From c73707bfb1a8c71bb1c090735eeeac39ea32bc85 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 28 Jul 2026 05:00:37 -0700 Subject: [PATCH 1/3] [None][fix] reserve speculative KV for ADP dummies Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 28 +++++++++++++++++-- .../_torch/executor/test_py_executor.py | 21 ++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2efc017e0e4c..b7551e6708b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5775,7 +5775,30 @@ def _pad_attention_dp_dummy_request(self): token_nums = None if (not self._adp_dummy_is_gen and self.kv_cache_transceiver is not None and self.max_num_tokens is not None): - token_nums = [self.max_num_tokens] + # max_num_tokens is the aggregate per-iteration budget and can + # exceed the legal capacity of one sequence. + token_num = min( + self.max_num_tokens, + self.model_engine.max_num_tokens, + self.model_engine.max_seq_len, + self.kv_cache_manager.max_seq_len, + ) + # One-engine speculative decoding appends extra KV tokens after + # add_sequence_batch(). Keep them in the same block count that + # the capacity scheduler reserves from the prompt length. + extra_kv_tokens = self.kv_cache_manager.num_extra_kv_tokens + tokens_per_block = self.kv_cache_manager.tokens_per_block + block_capacity = ((token_num + tokens_per_block - 1) // + tokens_per_block) * tokens_per_block + token_num = max(1, min(token_num, block_capacity - extra_kv_tokens)) + token_nums = [token_num] + + if not self._has_adp_dummy_kv_capacity(token_nums): + logger.warning_once( + "Unable to fit the complete attention-DP dummy KV allocation; " + "skipping this forward iteration and retrying", + key="attention_dp_dummy_insufficient_kv_capacity") + return if (not self._enable_dsv4_adp_dummy_fixes or self.kv_cache_transceiver is None): @@ -5798,8 +5821,7 @@ def _pad_attention_dp_dummy_request(self): has_live_adp_dummy = any( request.py_request_id == ATTENTION_DP_DUMMY_REQUEST_ID for request in self.active_requests) - if has_live_adp_dummy or not self._has_adp_dummy_kv_capacity( - token_nums): + if has_live_adp_dummy: return try: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 101edffbe155..50f33d321c2f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1300,6 +1300,8 @@ def __init__( enable_attention_dp=True, kv_cache_transceiver=object(), max_num_tokens=8192, + max_seq_len=8192, + kv_manager_max_seq_len=None, is_warmup=False, benchmark_req_queues_size=0, enable_dsv4_adp_dummy_fixes=True, @@ -1318,6 +1320,7 @@ def __init__( self._pending_adp_dummy_request = None self._enable_dsv4_adp_dummy_fixes = enable_dsv4_adp_dummy_fixes self.add_dummy_calls = [] + self.model_engine = Mock(max_num_tokens=max_num_tokens, max_seq_len=max_seq_len) self.dist = Mock() self.dist.tp_size = 1 @@ -1326,6 +1329,11 @@ def __init__( kv_cache_manager = Mock() kv_cache_manager.mapping.has_cp_helix.return_value = False kv_cache_manager.get_num_available_tokens.return_value = 1 << 30 + kv_cache_manager.max_seq_len = ( + max_seq_len if kv_manager_max_seq_len is None else kv_manager_max_seq_len + ) + kv_cache_manager.num_extra_kv_tokens = 0 + kv_cache_manager.tokens_per_block = 128 def _add_dummy(**kwargs): self.add_dummy_calls.append(kwargs) @@ -1496,6 +1504,19 @@ def test_pad_dummy_allocation_failure_skips_padding(): assert not any(r.is_attention_dp_dummy for r in stub.active_requests) +def test_disabled_dsv4_gate_checks_full_generation_capacity(): + stub = _StubADPExecutor(enable_dsv4_adp_dummy_fixes=False) + stub.max_total_draft_tokens = 4 + stub.kv_cache_manager.get_num_available_tokens.return_value = 4 + + _run_pad(stub) + + stub.kv_cache_manager.get_num_available_tokens.assert_called_once_with( + token_num_upper_bound=5, max_num_draft_tokens=4 + ) + stub.kv_cache_manager.add_dummy_requests.assert_not_called() + + def test_dsv4_pad_dummy_checks_full_context_capacity(): stub = _StubADPExecutor(max_num_tokens=4096) stub._adp_dummy_is_gen = False From 76ced54b1335dc1c2ce5f93b05f3d26a7c02e82e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 28 Jul 2026 06:10:20 -0700 Subject: [PATCH 2/3] [None][fix] reserve draft KV at disagg transition Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 15 +++---- .../test_kv_cache_v2_capacity_only.py | 40 +++++++++++++++++++ 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 4b7bde587dbf..56c6c00e0218 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2154,15 +2154,10 @@ def is_request_active(self, request_id: int) -> bool: def _effective_draft_len(self, req: LlmRequest) -> int: """Draft token length to use for next-step KV capacity calculation. - For a disagg gen request whose KV transmission just completed - (state == DISAGG_GENERATION_TRANS_COMPLETE), py_draft_tokens is - still [] when the scheduler asks for capacity, because it gets - mirrored from context_phase_params.draft_tokens later in - _prepare_disagg_gen_transmission_complete (which runs AFTER the - scheduler in the executor loop). Without compensating here, the - first gen forward writes 1 + len(ctx_draft_tokens) tokens into - KV cache but only +1 was reserved, OOB-ing the KV block table at - the next tokens_per_block-aligned boundary. + During the context-to-generation transition, ``py_draft_tokens`` is + still empty when the scheduler asks for capacity. Prefer transferred + context draft tokens; if there are none, reserve the generation + worker's configured draft capacity before its first forward pass. """ draft_len = get_draft_token_length(req) if ( @@ -2173,6 +2168,8 @@ def _effective_draft_len(self, req: LlmRequest) -> int: ctx_draft_tokens = req.context_phase_params.draft_tokens if ctx_draft_tokens is not None: draft_len = len(ctx_draft_tokens) + if draft_len == 0 and not self.is_draft and not req.py_disable_speculative_decoding: + draft_len = self.max_total_draft_tokens return draft_len def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py index a37288a80895..1cd615c4d648 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -182,3 +182,43 @@ def test_llm_request_has_no_compression_consumer_marker() -> None: assert "py_kv_cache_kv_compression_manages_history" not in vars(request) assert "py_kv_cache_compaction" not in vars(request) + + +def test_disagg_gen_transition_reserves_target_drafts_without_context_drafts(): + manager = _manager(is_draft=False) + manager.max_total_draft_tokens = 4 + request = SimpleNamespace( + py_draft_tokens=[], + is_disagg_generation_transmission_complete=True, + context_phase_params=SimpleNamespace(draft_tokens=None), + py_disable_speculative_decoding=False, + ) + + assert manager._effective_draft_len(request) == 4 + assert manager._required_gen_capacity(request, 128) == 137 + + +def test_disagg_gen_transition_does_not_reserve_disabled_speculation(): + manager = _manager(is_draft=False) + manager.max_total_draft_tokens = 4 + request = SimpleNamespace( + py_draft_tokens=[], + is_disagg_generation_transmission_complete=True, + context_phase_params=SimpleNamespace(draft_tokens=None), + py_disable_speculative_decoding=True, + ) + + assert manager._effective_draft_len(request) == 0 + + +def test_disagg_gen_transition_prefers_context_drafts(): + manager = _manager(is_draft=False) + manager.max_total_draft_tokens = 4 + request = SimpleNamespace( + py_draft_tokens=[], + is_disagg_generation_transmission_complete=True, + context_phase_params=SimpleNamespace(draft_tokens=[1, 2]), + py_disable_speculative_decoding=False, + ) + + assert manager._effective_draft_len(request) == 2 From d463494a9ca1dd4d719f751d680e01556e80abb9 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:12:59 -0700 Subject: [PATCH 3/3] fix ci failure Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../unittest/_torch/executor/test_kv_cache_v2_capacity_only.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py index 1cd615c4d648..a67cafdc2b01 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -195,7 +195,7 @@ def test_disagg_gen_transition_reserves_target_drafts_without_context_drafts(): ) assert manager._effective_draft_len(request) == 4 - assert manager._required_gen_capacity(request, 128) == 137 + assert manager._required_gen_capacity(request, 128) == 133 def test_disagg_gen_transition_does_not_reserve_disabled_speculation():