[None][feat] Add environment variable to force spec-dec number of accepted tokens (#9371)

Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

Author: achartier

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -56,6 +56,7 @@
 from tensorrt_llm.sampling_params import SamplingParams
 
 from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
+from ..speculative.interface import get_force_num_accepted_tokens
 from ..speculative.spec_tree_manager import SpecTreeManager
 from .finish_reason import FinishedState
 from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length
@@ -662,6 +663,9 @@ def __init__(self, args: Args):
         self._global_seed = 42
         self._generator = None
 
+        # Force number of accepted tokens for speculative decoding testing
+        self._force_num_accepted_tokens = get_force_num_accepted_tokens()
+
     def get_generator(self, device: torch.device) -> torch.Generator:
         """Get a deterministic generator for the specified device.
 
@@ -784,15 +788,24 @@ def _process_draft_tokens_greedy(
             return 0
         num_accepted = 0
 
-        for draft_token in request.py_draft_tokens:
-            if draft_token != new_token:
-                # Reject.
-                break
-
-            num_accepted += 1
-            new_token = add_token(request, new_tokens, beam=BEAM, step=num_accepted)
-            if self.finish_if_reason(request, finish_reasons, step=num_accepted):
-                break
+        if self._force_num_accepted_tokens != 0:
+            # Force acceptance of up to force_num_accepted_tokens draft tokens
+            force_limit = min(self._force_num_accepted_tokens, len(request.py_draft_tokens))
+            for _ in request.py_draft_tokens[:force_limit]:
+                num_accepted += 1
+                new_token = add_token(request, new_tokens, beam=BEAM, step=num_accepted)
+                if self.finish_if_reason(request, finish_reasons, step=num_accepted):
+                    break
+        else:
+            for draft_token in request.py_draft_tokens:
+                if draft_token != new_token:
+                    # Reject.
+                    break
+
+                num_accepted += 1
+                new_token = add_token(request, new_tokens, beam=BEAM, step=num_accepted)
+                if self.finish_if_reason(request, finish_reasons, step=num_accepted):
+                    break
         return num_accepted
 
     def _process_draft_tokens_tree(

tensorrt_llm/_torch/speculative/eagle3.py

@@ -12,7 +12,7 @@
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.scheduler import ScheduledRequests
-from .interface import SpecMetadata
+from .interface import SpecMetadata, get_force_num_accepted_tokens
 from .mtp import MTPSampler
 from .spec_tree_manager import SpecTreeManager
 
@@ -365,6 +365,7 @@ def __init__(self, spec_config: "EagleDecodingConfig", mapping: Mapping):
         self.max_draft_len = self.spec_config.max_draft_len
         self.mapping = mapping
         self.guided_decoder: Optional[CapturableGuidedDecoder] = None
+        self.force_num_accepted_tokens = get_force_num_accepted_tokens()
 
     # Skip torch.compile for now since current Torch is not compatible with Triton 3.4
     # @torch.compile(options={"max-autotune": True})
@@ -527,6 +528,11 @@ def sample_and_accept_draft_tokens(
         num_accepted_tokens[num_contexts:] += torch.cumprod(
             (draft_tokens == gen_target_tokens[:, :self.max_draft_len]).int(),
             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
         return accepted_tokens, num_accepted_tokens
 
     def draft_decoder(

tensorrt_llm/_torch/speculative/interface.py

@@ -1,14 +1,37 @@
 import copy
+import os
 from dataclasses import dataclass, field
 from enum import IntEnum, auto
 from typing import List, Optional, Type
 
 import torch
 
+from tensorrt_llm.logger import logger
+
 from ..._utils import get_sm_version
 from ..attention_backend.trtllm import AttentionBackend, TrtllmAttention
 from ..pyexecutor.resource_manager import BaseResourceManager
 
+# Environment variable name for forcing the number of accepted tokens in speculative decoding
+FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR = "TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS"
+
+
+def get_force_num_accepted_tokens() -> int:
+    """
+    Read and parse the TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS environment variable.
+
+    Returns:
+        int: The forced number of accepted tokens, or 0 if not set or invalid.
+    """
+    env_value = os.environ.get(FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, "0")
+    try:
+        return int(env_value)
+    except ValueError:
+        logger.warning(
+            f"{FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR} must be a valid integer, "
+            f"got '{env_value}'. Using default value 0.")
+        return 0
+
 
 class SpeculativeDecodingMode(IntEnum):
     MTP = auto()

tensorrt_llm/_torch/speculative/mtp.py

@@ -17,7 +17,7 @@
                                   SampleStateTensors, TorchSampler, add_token,
                                   int_tensor)
 from ..pyexecutor.scheduler import ScheduledRequests
-from .interface import SpecMetadata
+from .interface import SpecMetadata, get_force_num_accepted_tokens
 
 if TYPE_CHECKING:
     from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
@@ -347,6 +347,7 @@ def __init__(self, spec_config: "MTPDecodingConfig", model_config=None):
         self.model_config = model_config
         self.is_thop = False
         self.guided_decoder: Optional[CapturableGuidedDecoder] = None
+        self.force_num_accepted_tokens = get_force_num_accepted_tokens()
 
     def forward(
         self,
@@ -895,6 +896,12 @@ def sample_and_accept_draft_tokens(
                      ).int(),
                     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,
+                                            mtp_num_modules + 1)
+            num_accepted_tokens[num_contexts:] = force_num_accepted_tokens
+
         return accepted_tokens, num_accepted_tokens
 
     def change_attn_metadata(self, num_accepted_tokens: torch.Tensor,