[DCP] Bound a request by the aggregate KV pool, not one rank's share (#33448)

This commit is contained in:
Khoa Pham
2026-08-03 19:36:25 -07:00
committed by GitHub
parent c113ead98a
commit 91fae8a72c
2 changed files with 10 additions and 3 deletions
+4 -1
View File
@@ -2080,7 +2080,10 @@ class Scheduler(
min(
max_new_tokens,
self.max_req_len - input_len - 1,
self.max_total_num_tokens - paged_input_len - self.page_size - 1,
self.max_total_num_tokens * self.server_args.dcp_size
- paged_input_len
- self.page_size
- 1,
),
)
# Clipping above can push max_new_tokens below min_new_tokens, which
+6 -2
View File
@@ -385,7 +385,9 @@ class TpModelWorker(BaseTpWorker):
assert self.model_runner.max_running_requests > 0, "max_running_request is zero"
max_req_len = min(
self.model_config.context_len - 1,
self.model_runner.effective_max_total_num_tokens - 1,
self.model_runner.effective_max_total_num_tokens
* self.model_runner.dcp_size
- 1,
)
assert max_req_len > 0, "Memory pool size is too small"
@@ -485,7 +487,9 @@ class TpModelWorker(BaseTpWorker):
def get_worker_info(self):
max_req_len = min(
self.model_config.context_len - 1,
self.model_runner.effective_max_total_num_tokens - 1,
self.model_runner.effective_max_total_num_tokens
* self.model_runner.dcp_size
- 1,
)
return (
self.model_runner.max_total_num_tokens,