From 9ff8de85849756752760854ab6a7c4f22197dac2 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Tue, 4 Aug 2026 12:38:21 -0700 Subject: [PATCH] [None][fix] Fix one-model MTP KV cache accounting * Why? The first overlapped MTP generation step could overstate its cached length and access an invalid FlashInfer page, causing a CUDA illegal memory access. The KV scheduler also lacked the Python-side draft-token count needed for correct capacity planning. * What? Keep the Python and C++ draft-token state synchronized, and use the prompt KV boundary when generation immediately follows context for backends without dynamic KV-length correction. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 20 +++++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 ++++ .../_torch/executor/test_py_executor.py | 45 +++++++++++++--- .../executor/test_pytorch_model_engine.py | 53 +++++++++++++++++++ 4 files changed, 117 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3fdd8190dd39..ceb1d990b11c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6e3252b48688..fa8f1095ddfb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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( diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c9c05b39137f..d068776c08e9 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/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 diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 9a1242441523..b072210dbd9e 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/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(