From deb56a392665371a89b5a56fc4985bd4b42b4682 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Tue, 2 Dec 2025 03:13:19 +0000 Subject: [PATCH] Fix TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS for MTP/EAGLE Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- tensorrt_llm/_torch/speculative/eagle3.py | 7 ++++--- tensorrt_llm/_torch/speculative/mtp.py | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 44fed8fddade..89b1ff0ff1dc 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -530,9 +530,10 @@ def sample_and_accept_draft_tokens( dim=-1).sum(1) # Check for environment variable override if self.force_num_accepted_tokens != 0: - force_num_accepted_tokens = min(self.force_num_accepted_tokens, - self.max_draft_len + 1) - num_accepted_tokens[num_contexts:] = force_num_accepted_tokens + # total tokens per iteration = accepted draft tokens + 1 target token + force_total_tokens = min(self.force_num_accepted_tokens + 1, + self.max_draft_len + 1) + num_accepted_tokens[num_contexts:] = force_total_tokens return accepted_tokens, num_accepted_tokens def draft_decoder( diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 99f06c122ac7..29917d820a49 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -908,9 +908,10 @@ def sample_and_accept_draft_tokens( # Check for environment variable override if self.force_num_accepted_tokens != 0: - force_num_accepted_tokens = min(self.force_num_accepted_tokens, - mtp_num_modules + 1) - num_accepted_tokens[num_contexts:] = force_num_accepted_tokens + # total tokens per iteration = accepted draft tokens + 1 target token + force_total_tokens = min(self.force_num_accepted_tokens + 1, + mtp_num_modules + 1) + num_accepted_tokens[num_contexts:] = force_total_tokens return accepted_tokens, num_accepted_tokens