From 9cb549325c1624ad8901f28d60c724ff18ddfc81 Mon Sep 17 00:00:00 2001 From: Jhao-Ting Chen Date: Wed, 8 Jul 2026 17:06:35 -0700 Subject: [PATCH] [https://nvbugs/6427240][fix] Reserve MTP draft tokens in scheduler for one-model speculative decoding One-model speculative decoding (MTP / Eagle3 one-model) has no separate drafter, so get_spec_drafter() returns None and the `if self.drafter is not None` block in _prepare_and_schedule_batch is skipped -- including the loop that sets `request.draft_tokens = [0] * max_total_draft_tokens` for generation requests. That loop is what makes the C++ micro-batch scheduler budget each generation request as beam_width + getNumDraftTokens(). Without it, under enable_chunked_prefill + the overlap scheduler, the scheduler under-reserves and the forward's uniform (1 + runtime_draft_len) build overshoots max_num_tokens, aborting requests with an AssertionError (total_num_tokens > max_num_tokens) in _prepare_tp_inputs. Add an elif branch mirroring the two-model draft-token normalization for the one-model path so scheduling reserves the correct budget. Using self.max_total_draft_tokens keeps it consistent with dynamic draft_len_schedule. model_engine is guarded (getattr ... is not None) first so partially- constructed executors in unit tests do not raise AttributeError. The state filter {GENERATION_IN_PROGRESS, DISAGG_GENERATION_INIT} matches the two-model normalization, so this single fix covers both aggregated and disaggregated (decode-worker) serving. Add a regression test (test_py_executor.py) that drives _prepare_and_schedule_batch on a one-model-MTP executor with real LlmRequests in GENERATION_IN_PROGRESS (aggregated) and DISAGG_GENERATION_INIT (disagg) states, asserting both are normalized to max_total_draft_tokens while CONTEXT_INIT is left untouched. The test fails without the fix and passes with it. Verified on Qwen3.6-35B-A3B-NVFP4 (MTP num_nextn_predict_layers=3, chunked prefill, overlap scheduler on): a concurrency sweep that previously produced 130 assertions at concurrency 32 now completes with zero assertions and zero request errors. Signed-off-by: Jhao-Ting Chen --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 19 +++ .../_torch/executor/test_py_executor.py | 109 +++++++++++++++++- 2 files changed, 127 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 0fbf170aaae6..e16422af232c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3418,6 +3418,25 @@ def _prepare_and_schedule_batch(self): # with dummy draft tokens to make the scheduler aware of the fact # that speculation is about to happen. self._prepare_draft_requests() + elif (getattr(self, "model_engine", None) is not None + and self.model_engine.is_spec_decode + and self.max_total_draft_tokens > 0): + # One-model speculative decoding (MTP / Eagle3 one-model) has no separate + # drafter, so the block above is skipped -- but the model still verifies + # max_total_draft_tokens draft tokens per generation step. The C++ micro-batch + # scheduler budgets each gen request as beam_width + getNumDraftTokens(); without + # populating the draft-token count here it under-reserves and, under chunked + # prefill + overlap scheduler, the forward's uniform (1 + runtime_draft_len) + # build overshoots max_num_tokens (total_num_tokens > max_num_tokens). Mirror the + # two-model normalization so scheduling reserves the correct token budget. + # model_engine is guarded first so partially-constructed executors in unit tests + # (which may not set model_engine) do not raise AttributeError. + for request in self.active_requests: + if request.state not in ( + LlmRequestState.GENERATION_IN_PROGRESS, + LlmRequestState.DISAGG_GENERATION_INIT): + continue + request.draft_tokens = [0] * self.max_total_draft_tokens scheduled_batch, 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 29b04404bbf1..352ec3347679 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -22,7 +22,7 @@ SHUTDOWN_REQUEST_ID, RequestQueueItem, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( @@ -1307,3 +1307,110 @@ def test_handles_error_on_single_rank(self): stub = _make_disagg_err_stub(world_size=1, active_requests=[err]) stub._check_cache_transfer_errors("gen") assert len(stub.handle_errors_calls) == 1 + + +class TestOneModelMTPDraftTokenScheduling: + """Regression tests for the one-model MTP over-scheduling bug (#16101). + + One-model MTP (``mtp_eagle_one_model``) has no separate drafter, so + ``get_spec_drafter()`` returns None and the ``if self.drafter is not None`` + draft-token normalization block in ``_prepare_and_schedule_batch`` is + skipped. Without the ``elif`` fallback that mirrors it, generation requests + keep ``num_draft_tokens == 0`` and the C++ micro-batch scheduler + under-reserves each gen request (it budgets ``beam_width + + getNumDraftTokens()``). Under chunked prefill + overlap scheduler the + 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. + + NOTE: Like ``test_fetch_called_once_even_in_benchmark_disagg`` in + ``test_benchmark_disagg.py``, this uses ``object.__new__(PyExecutor)`` to + bypass ``__init__`` and sets internal attributes by hand. Real + ``LlmRequest`` objects (not Mocks) are used so ``draft_tokens = [0] * N`` + actually updates the C++-backed ``num_draft_tokens`` count. + """ + + MAX_TOTAL_DRAFT_TOKENS = 3 + + @staticmethod + def _make_llm_request(request_id: int, state: LlmRequestState) -> LlmRequest: + """Build a real LlmRequest in the given state (mirrors the helper in + test_py_scheduler.py::make_generation_request).""" + req = LlmRequest( + request_id=request_id, + max_new_tokens=10, + input_tokens=list(range(10)), + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=None, + ) + req.state = state + return req + + @classmethod + def _make_one_model_mtp_executor(cls, active_requests): + """Construct a partially-initialised one-model-MTP PyExecutor. + + drafter is None (one-model MTP has no separate drafter) and + model_engine.is_spec_decode is True, so _prepare_and_schedule_batch + takes the elif draft-token normalization branch. kv_cache_transceiver + is None to keep the test hermetic (skips the disagg blocks). + """ + ex = object.__new__(PyExecutor) + ex.drafter = None + ex.max_total_draft_tokens = cls.MAX_TOTAL_DRAFT_TOKENS + ex.model_engine = Mock(is_spec_decode=True) + ex.kv_cache_transceiver = None + ex.is_shutdown = False + ex.enable_iter_perf_stats = False + ex.active_requests = active_requests + ex.waiting_queue = [] + + ex._fetch_and_activate_new_requests = Mock(return_value=[]) + ex._check_disagg_ctx_schedulable_status = Mock() + ex._check_disagg_gen_transfer_status = Mock() + ex._check_kv_transfer_timeout = Mock() + ex._check_disagg_ctx_cache_transfer_status = Mock() + ex._pad_attention_dp_dummy_request = Mock() + ex._prefetch_for_context_requests = Mock() + ex._prepare_disagg_gen_init = Mock() + ex._schedule = Mock(return_value=(ScheduledRequests(), [], 0)) + return ex + + def test_one_model_mtp_populates_draft_tokens_for_scheduling(self): + """The draft-token normalization must cover BOTH aggregated and + disaggregated serving in one shot. + + The fix's state filter is {GENERATION_IN_PROGRESS, + DISAGG_GENERATION_INIT}, mirroring the two-model normalization, so a + single fix covers the aggregated decode path (GENERATION_IN_PROGRESS) + and the disagg decode-worker path (DISAGG_GENERATION_INIT). Context + requests (CONTEXT_INIT) are not generation requests and must be left + untouched. + """ + gen = self._make_llm_request(0, LlmRequestState.GENERATION_IN_PROGRESS) + disagg_gen = self._make_llm_request(1, LlmRequestState.DISAGG_GENERATION_INIT) + ctx = self._make_llm_request(2, LlmRequestState.CONTEXT_INIT) + + # Precondition: no draft tokens reserved yet on either gen request. + assert gen.num_draft_tokens == 0 + assert disagg_gen.num_draft_tokens == 0 + + ex = self._make_one_model_mtp_executor([gen, disagg_gen, ctx]) + scheduled_batch, _ = ex._prepare_and_schedule_batch() + + assert scheduled_batch is not None + # Aggregated case: in-progress generation request is normalized to the + # 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 + # Disaggregated case: decode-worker request awaiting KV also normalized. + assert disagg_gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS + # Context requests are not generation requests and must be left alone. + assert ctx.num_draft_tokens == 0