From bfe9d463f60571b6f4a37a9174d0fe1c2591e26f Mon Sep 17 00:00:00 2001 From: William Zhang Date: Mon, 17 Aug 2026 11:10:04 -0700 Subject: [PATCH] [None][fix] Synchronize one-model MTP scheduler draft tokens Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++- .../_torch/executor/test_py_executor.py | 72 ++++++++++++++++--- 2 files changed, 73 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4770fe2e6722..c6cc14c4304c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3430,7 +3430,7 @@ def _handle_dynamic_draft_len(self, # so next iter's rejection kernel would read a stale draft_probs # row. Mark it (pre-pad signal) so _prepare_tp_inputs writes a # one-hot placeholder row after spec_metadata.prepare(). - request.py_needs_onehot_draft_probs = ( + request.py_needs_onehot_draft_probs |= ( rejection_on and current_num_draft_tokens == 0) if spec_dec_mode.is_pard(): # special case: PARD carries 2K-1 draft tokens per request @@ -3874,11 +3874,19 @@ def _prepare_and_schedule_batch(self): # 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. + rejection_on = ( + self.model_engine.spec_config is not None + and self.model_engine.spec_config.use_rejection_sampling) for request in self.active_requests: if request.state not in ( LlmRequestState.GENERATION_IN_PROGRESS, LlmRequestState.DISAGG_GENERATION_INIT): continue + # Populate the Python-side count only when no real draft tokens + # are available from the previous iteration. + if not request.py_draft_tokens: + request.py_needs_onehot_draft_probs |= rejection_on + 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 6ad05fc601bc..186af2cb6243 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -49,6 +49,7 @@ ScheduledRequests, SerializableSchedulerOutput, ) +from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError pytestmark = pytest.mark.cpu_only @@ -2876,12 +2877,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 @@ -2908,7 +2914,11 @@ def _make_llm_request(request_id: int, state: LlmRequestState) -> LlmRequest: return req @classmethod - def _make_one_model_mtp_executor(cls, active_requests): + def _make_one_model_mtp_executor( + cls, + active_requests: list[LlmRequest], + use_rejection_sampling: bool = False, + ) -> PyExecutor: """Construct a partially-initialised one-model-MTP PyExecutor. drafter is None (one-model MTP has no separate drafter) and @@ -2921,11 +2931,23 @@ def _make_one_model_mtp_executor(cls, active_requests): 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) + spec_config = MTPDecodingConfig( + max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS, + mtp_eagle_one_model=True, + use_rejection_sampling=use_rejection_sampling, + draft_len_schedule={1: cls.MAX_TOTAL_DRAFT_TOKENS}, + ) + ex.model_engine = Mock( + is_spec_decode=True, + spec_config=spec_config, + max_draft_len=cls.MAX_TOTAL_DRAFT_TOKENS, + max_total_draft_tokens=cls.MAX_TOTAL_DRAFT_TOKENS, + ) ex.kv_cache_transceiver = None ex.is_shutdown = False ex.enable_iter_perf_stats = False ex.enable_attention_dp = False + ex.speculation_permanently_disabled = False ex.active_requests = active_requests ex.waiting_queue = [] @@ -2958,6 +2980,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() @@ -2967,7 +2991,39 @@ 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_zero_proposal_signal_for_rejection(self) -> None: + gen = self._make_llm_request(0, LlmRequestState.GENERATION_IN_PROGRESS) + ex = self._make_one_model_mtp_executor([gen], use_rejection_sampling=True) + + ex._prepare_and_schedule_batch() + + assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS + assert gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS + assert gen.py_needs_onehot_draft_probs + + batch = ScheduledRequests() + batch.append_generation_request(gen) + ex._handle_dynamic_draft_len(batch) + + assert ex.model_engine.runtime_draft_len == self.MAX_TOTAL_DRAFT_TOKENS + assert gen.py_needs_onehot_draft_probs + assert gen.py_draft_tokens == [0] * self.MAX_TOTAL_DRAFT_TOKENS + + def test_one_model_mtp_preserves_sampler_draft_tokens(self) -> None: + sampler_drafts = [7, 8] + 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 + assert gen.num_draft_tokens == self.MAX_TOTAL_DRAFT_TOKENS