Address Mike's comments

Signed-off-by: Zheyu Fu <[email protected]>

Author: zheyuf

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -297,19 +297,6 @@ def __init__(
         self.spec_config = spec_config
         self.is_spec_decode = spec_config is not None
         self.enable_spec_decode = self.is_spec_decode
-        # Rolling acceptance tracking
-        self.acceptance_window = getattr(
-            spec_config, 'acceptance_window',
-            None) if spec_config is not None else None
-        self.acceptance_length_threshold = getattr(
-            spec_config, 'acceptance_length_threshold',
-            None) if spec_config is not None else None
-        # Initialize speculation gate early since it only depends on config
-        self.speculation_permanently_disabled = False
-        self.speculation_gate = None
-        if self.acceptance_window and self.acceptance_length_threshold is not None:
-            self.speculation_gate = SpeculationGate(
-                self.acceptance_window, self.acceptance_length_threshold)
         self.is_draft_model = is_draft_model
 
         self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures(

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -38,6 +38,7 @@
 from ..models.modeling_utils import DecoderModelForCausalLM
 from ..modules.decoder_layer import DecoderLayer
 from ..speculative.drafter import Drafter
+from ..speculative.speculation_gate import SpeculationGate
 from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem
 from .guided_decoder import GuidedDecoder
 from .handle_logits import HandleLogits
@@ -207,6 +208,20 @@ def __init__(self,
         self.num_fetch_requests = 0
         self.shutdown_event = threading.Event()
 
+        # Rolling acceptance tracking for spec decode (disable speculation if rolling acceptance is below threshold)
+        spec_config = getattr(self.model_engine, 'spec_config', None)
+        self.acceptance_window = getattr(
+            spec_config, 'acceptance_window',
+            None) if spec_config is not None else None
+        self.acceptance_length_threshold = getattr(
+            spec_config, 'acceptance_length_threshold',
+            None) if spec_config is not None else None
+        self.speculation_permanently_disabled = False
+        self.speculation_gate = None
+        if self.acceptance_window and self.acceptance_length_threshold is not None:
+            self.speculation_gate = SpeculationGate(
+                self.acceptance_window, self.acceptance_length_threshold)
+
         # response used data
         self.response_lock = threading.Lock()
         self.response_cv = threading.Condition(self.response_lock)
@@ -969,15 +984,14 @@ def _prepare_and_schedule_batch(self):
 
         if self.drafter is not None:
             # Honor permanent disable flag based on rolling acceptance first
-            if getattr(self.model_engine, 'speculation_permanently_disabled',
-                       False):
+            if getattr(self, 'speculation_permanently_disabled', False):
                 self.use_spec_decode = False
             else:
                 self.use_spec_decode = self.drafter.should_use_spec_decode(
                     self.active_requests, self.max_batch_size,
                     self.model_engine.max_num_tokens,
                     self.model_engine.spec_config.max_draft_len)
-
+            logger.debug(f"Use spec decode: {self.use_spec_decode}")
             self.model_engine.enable_spec_decode = self.use_spec_decode
 
             # When overlap scheduler is enabled, and we already prepared the draft tokens in the previous batch,
@@ -1911,24 +1925,29 @@ def _handle_responses(self):
                     new_responses.append((req_id, response))
 
             if request_done:
-                if (self.model_engine.enable_spec_decode and
-                        not self.model_engine.speculation_permanently_disabled
+                if (self.model_engine.enable_spec_decode
+                        and not self.speculation_permanently_disabled
                         and not request.is_dummy and not self.is_warmup):
-                    if self.model_engine.speculation_gate is not None:
+                    if self.speculation_gate is not None:
                         # Response handling runs on multiple PP ranks. Only the last PP rank performs
                         # sampling; restrict rolling stat updates to it to avoid overcounting.
                         if (not getattr(self.dist, 'has_pp',
                                         False)) or self.dist.is_last_pp_rank:
                             avg_decoded = getattr(
                                 request, 'avg_decoded_tokens_per_iter', None)
-                            disabled_now, _ = self.model_engine.speculation_gate.record_avg_decoded(
-                                avg_decoded,
-                                request_id=getattr(request, 'py_request_id',
-                                                   None))
-                            if disabled_now:
-                                # disable speculation permanently
-                                # starting from next iteration, _prepare_and_schedule_batch will set self.use_spec_decode to False
-                                self.model_engine.speculation_permanently_disabled = True
+                            if avg_decoded is not None:
+                                disabled_now, _ = self.speculation_gate.record_avg_decoded(
+                                    avg_decoded,
+                                    request_id=getattr(request, 'py_request_id',
+                                                       None))
+                                if disabled_now:
+                                    # disable speculation permanently
+                                    # starting from next iteration, _prepare_and_schedule_batch will set self.use_spec_decode to False
+                                    self.speculation_permanently_disabled = True
+                            else:
+                                logger.debug(
+                                    f"Request {request.py_request_id} has no avg_decoded_tokens_per_iter"
+                                )
                 if request.is_disagg_context_transmission_state:
                     self.ctx_in_transmission_requests.append(request)
                 else:

tensorrt_llm/_torch/speculative/speculation_gate.py

@@ -29,18 +29,21 @@ def reset(self) -> None:
 
     def record_avg_decoded(
             self,
-            avg_decoded_tokens_per_iter: Optional[float],
+            avg_decoded_tokens_per_iter: float,
             request_id: Optional[int] = None) -> Tuple[bool, Optional[float]]:
         """
-		Record a completed request's avg_decoded_tokens_per_iter.
-		Returns (disabled_now, current_avg_accept) where disabled_now is True only when the call causes disable.
-		"""
+        Record a completed request's avg_decoded_tokens_per_iter.
+        Returns (disabled_now, current_avg_accept) where disabled_now is True only when the call causes disable.
+        """
         if self.disabled or self.window is None or self.window <= 0 or self.threshold is None:
             return False, None
 
+        # Extra Guard: if caller passed None, skip updating the rolling stats
+        if avg_decoded_tokens_per_iter is None:
+            return False, None
+
         accepted_len = 0.0
-        if avg_decoded_tokens_per_iter is not None:
-            accepted_len = max(0.0, float(avg_decoded_tokens_per_iter) - 1.0)
+        accepted_len = max(0.0, float(avg_decoded_tokens_per_iter) - 1.0)
 
         # Log per-request completion for debug
         if request_id is not None:
@@ -50,6 +53,8 @@ def record_avg_decoded(
 
         # O(1) rolling update
         self.acceptance_history.append(accepted_len)
+        logger.debug(
+            f"[SpeculationGate] Acceptance history: {self.acceptance_history}")
         self.acceptance_sum += accepted_len
         if len(self.acceptance_history) > self.window:
             removed = self.acceptance_history.popleft()

tensorrt_llm/llmapi/llm_args.py

@@ -367,9 +367,6 @@ class DecodingBaseConfig(StrictBaseModel):
     # (N = acceptance_window) drops below this value.
     acceptance_length_threshold: Optional[float] = None
 
-    # Upper bound to avoid accidental huge windows
-    MAX_ACCEPTANCE_WINDOW: ClassVar[int] = 100000
-
     # Validate acceptance controls at field level so they run on model creation
     @field_validator('acceptance_window')
     @classmethod
@@ -379,10 +376,6 @@ def _validate_acceptance_window(cls, v: Optional[int]):
         if v < 0:
             raise ValueError(
                 f"acceptance_window must be >= 0 (0 disables), got {v}")
-        if v > cls.MAX_ACCEPTANCE_WINDOW:
-            raise ValueError(
-                f"acceptance_window must be <= {cls.MAX_ACCEPTANCE_WINDOW}, got {v}"
-            )
         return v
 
     @field_validator('acceptance_length_threshold')