[None][fix] reserve draft KV at disagg transition for gen-only-dspark case

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:

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:

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) == 133
+
+
+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

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