[None][fix] Fix one-model MTP KV cache accounting

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -4406,11 +4406,23 @@ def append_cross_attention_state(request: LlmRequest,
                 previous_pos_indices.extend([previous_batch_idx] *
                                             runtime_tokens_per_gen_step)
 
+                cached_token_num = (past_seen_token_num +
+                                    runtime_tokens_per_gen_step)
+                # The first generation batch overlaps the context sampler, so there is a previous
+                # token tensor but no previous speculative target forward in the KV cache.
+                # Backends without dynamic KV lengths cannot apply the runtime acceptance-length
+                # correction in _preprocess_inputs and must start at the prompt boundary.
+                # py_decoding_iter is a proxy for "the previous forward for this request was its
+                # last context chunk": _update_requests lags _prepare_and_schedule_batch by exactly
+                # one iteration, so the sampler's first increment has not landed yet at this point
+                # and only here. This branch already assumes that same batch-to-batch continuity
+                # (previous_batch_idx indexes the immediately preceding batch's device tensors).
+                if (request.py_decoding_iter == 0
+                        and not hasattr(attn_metadata, "kv_lens_cuda")):
+                    cached_token_num = request.max_beam_num_tokens
                 num_cached_tokens_per_seq.append(
-                    past_seen_token_num + runtime_tokens_per_gen_step -
-                    request.py_num_compressed_tokens)
-                request.cached_tokens = (past_seen_token_num +
-                                         runtime_tokens_per_gen_step)
+                    cached_token_num - request.py_num_compressed_tokens)
+                request.cached_tokens = cached_token_num
                 if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
                         self.attn_backend) and spec_config.is_linear_tree:
                     prompt_lengths.append(runtime_tokens_per_gen_step)

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3676,6 +3676,15 @@ def _prepare_and_schedule_batch(self):
                         LlmRequestState.GENERATION_IN_PROGRESS,
                         LlmRequestState.DISAGG_GENERATION_INIT):
                     continue
+                # Only fill in a placeholder when the Python-side list is empty (e.g. a
+                # DISAGG_GENERATION_INIT request, which never gets a draft snapshot). Overwriting it
+                # unconditionally would clobber the real draft tokens the one-model spec sampler
+                # wrote at the end of the previous iteration - which the overlap-disabled path
+                # reads back in `_prepare_tp_inputs` - and would also suppress the
+                # `current_num_draft_tokens == 0` signal that `_handle_dynamic_draft_len` uses to
+                # request a one-hot draft-probs placeholder under rejection sampling.
+                if not request.py_draft_tokens:
+                    request.py_draft_tokens = [0] * self.max_total_draft_tokens
                 request.draft_tokens = [0] * self.max_total_draft_tokens
 
         scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(

tests/unittest/_torch/executor/test_py_executor.py

@@ -1920,12 +1920,17 @@ class TestOneModelMTPDraftTokenScheduling:
     forward then builds a uniform ``1 + runtime_draft_len`` per gen request and
     overshoots ``max_num_tokens`` (``total_num_tokens > max_num_tokens``).
 
-    The fix populates ``request.draft_tokens = [0] * max_total_draft_tokens``
-    on every in-progress generation request so scheduling reserves the correct
-    token budget. This test drives ``_prepare_and_schedule_batch`` for a
-    one-model-MTP executor and asserts generation requests get
-    ``num_draft_tokens == max_total_draft_tokens`` while context requests are
-    left untouched.
+    The fix populates both the Python and C++ draft-token representations on
+    every in-progress generation request so both schedulers reserve the
+    correct token budget. This test drives ``_prepare_and_schedule_batch`` for
+    a one-model-MTP executor and asserts generation requests get the full
+    draft-token budget while context requests are left untouched.
+
+    The Python-side fill is placeholder-only: with the overlap scheduler
+    disabled, `_prepare_tp_inputs` sources a generation request's draft
+    tokens from `py_draft_tokens`, which the one-model spec sampler wrote at
+    the end of the previous iteration. Overwriting a populated list here would
+    feed zeros to the target model and collapse the acceptance rate.
 
     NOTE: Like ``test_fetch_called_once_even_in_benchmark_disagg`` in
     ``test_benchmark_disagg.py``, this uses ``object.__new__(PyExecutor)`` to
@@ -1999,6 +2004,8 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self):
         # Precondition: no draft tokens reserved yet on either gen request.
         assert gen.num_draft_tokens == 0
         assert disagg_gen.num_draft_tokens == 0
+        assert gen.py_draft_tokens == []
+        assert disagg_gen.py_draft_tokens == []
 
         ex = self._make_one_model_mtp_executor([gen, disagg_gen, ctx])
         scheduled_batch, _ = ex._prepare_and_schedule_batch()
@@ -2008,7 +2015,33 @@ def test_one_model_mtp_populates_draft_tokens_for_scheduling(self):
         # full draft-token budget so the micro-batch scheduler reserves
         # beam + max_total_draft_tokens.
         assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
+        assert gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS
         # Disaggregated case: decode-worker request awaiting KV also normalized.
         assert disagg_gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS
+        assert disagg_gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS
         # Context requests are not generation requests and must be left alone.
         assert ctx.num_draft_tokens == 0
+        assert ctx.py_draft_tokens == []
+
+    def test_one_model_mtp_preserves_sampler_draft_tokens(self):
+        """Normalization must not clobber real draft tokens.
+
+        With `disable_overlap_scheduler=True` the one-model spec sampler writes the next iteration's
+        draft tokens into `py_draft_tokens`, and `_prepare_tp_inputs` reads them straight back into
+        `input_ids` / `draft_tokens_cuda` (there is no previous-iteration device tensor to source
+        them from). Overwriting them with the zero placeholder leaves the target model verifying
+        token id 0, silently dropping the acceptance rate to chance.
+        Only the C++ count needs unconditional normalization.
+        """
+        sampler_drafts = [7, 8, 9]
+        assert len(sampler_drafts) == self.MAX_TOTAL_DRAFT_TOKENS
+
+        gen = self._make_llm_request(0, LlmRequestState.GENERATION_IN_PROGRESS)
+        gen.py_draft_tokens = list(sampler_drafts)
+
+        ex = self._make_one_model_mtp_executor([gen])
+        ex._prepare_and_schedule_batch()
+
+        assert gen.py_draft_tokens == sampler_drafts
+        # The C++ count is still normalized for the micro-batch scheduler.
+        assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS

tests/unittest/_torch/executor/test_pytorch_model_engine.py

@@ -1172,6 +1172,59 @@ def set_attn_max_seq_len(self, max_seq_len: int) -> None:
                                  (encoder_batch_size, encoder_max_num_tokens))
                 self.assertEqual(encoder.max_seq_len, expected_max_seq_len)
 
+    def test_first_speculative_generation_uses_prompt_kv_boundary(self) -> None:
+        spec_config = SADecodingConfig(
+            max_draft_len=3,
+            draft_len_schedule={1: 3},
+        )
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
+            spec_config=spec_config)
+        model_engine.runtime_draft_len = 3
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+        attn_metadata = AttentionMetadata(max_num_requests=4,
+                                          max_num_tokens=32,
+                                          kv_cache_manager=kv_cache_manager)
+        attn_metadata.is_cuda_graph = False
+
+        generation = _create_request_with_tokens([50, 51, 52, 53, 54], 1)
+        generation.py_seq_slot = 0
+        # The final context batch populated this overlap slot, but it did not
+        # run a speculative target step.
+        generation.py_batch_idx = 0
+        generation.py_draft_tokens = [0, 0, 0]
+        self.assertEqual(generation.py_decoding_iter, 0)
+
+        graph_batch = ScheduledRequests()
+        graph_batch.generation_requests = [generation]
+        overlap_state = SampleStateTensorsSpec(
+            new_tokens=torch.zeros((4, 4, 1), dtype=torch.int32, device="cuda"),
+            new_tokens_lens=torch.ones(4, dtype=torch.int32, device="cuda"),
+            next_draft_tokens=torch.zeros((4, 3),
+                                          dtype=torch.int32,
+                                          device="cuda"),
+        )
+        spec_metadata = Mock(_force_non_greedy_for_capture=False)
+
+        inputs, _ = model_engine._prepare_tp_inputs(
+            scheduled_requests=graph_batch,
+            kv_cache_manager=kv_cache_manager,
+            attn_metadata=attn_metadata,
+            spec_metadata=spec_metadata,
+            new_tensors_device=overlap_state,
+            resource_manager=resource_manager,
+        )
+
+        prompt_len = generation.max_beam_num_tokens
+        self.assertEqual(
+            attn_metadata.kv_cache_params.num_cached_tokens_per_seq,
+            [prompt_len])
+        self.assertEqual(generation.cached_tokens, prompt_len)
+        model_engine._preprocess_inputs(inputs)
+        self.assertEqual(inputs["position_ids"][0, :4].cpu().tolist(),
+                         list(range(prompt_len, prompt_len + 4)))
+        kv_cache_manager.shutdown()
+
     def test_pad_generation_requests(self) -> None:
         model_engine, kv_cache_manager = create_model_engine_and_kvcache()
         resource_manager = ResourceManager(