AR based speculation off

Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>

Author: zheyuf

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -328,6 +328,7 @@ def __init__(
         self.max_beam_width = max_beam_width
         self.max_draft_len = max_draft_len
         self.max_total_draft_tokens = max_total_draft_tokens
+        self.use_spec_decode = self.model_engine.enable_spec_decode
         self.llm_args = self.model_engine.llm_args
         self.max_stats_len = max(self.llm_args.max_stats_len, 1)
         self.max_num_tokens = self.llm_args.max_num_tokens
@@ -352,19 +353,16 @@ def __init__(
         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
+        # Rolling true acceptance-rate tracking for permanent speculation
+        # disable.
         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)
+        spec_config = getattr(self.model_engine, 'spec_config', None)
+        if spec_config is not None:
+            window = getattr(spec_config, 'acceptance_rate_window_size', None)
+            threshold = getattr(spec_config, 'acceptance_rate_threshold', None)
+            if window and threshold is not None:
+                self.speculation_gate = SpeculationGate(window, threshold)
 
         # response used data
         self.response_lock = threading.Lock()
@@ -1185,6 +1183,42 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
                 draft_latency_ms) / float(iter_latency_ms)
         return stats
 
+    def _record_batch_acceptance_rate(
+            self,
+            scheduled_batch: ScheduledRequests,
+            sample_state: SampleState,
+            iteration_id: Optional[int] = None) -> Tuple[bool, Optional[float]]:
+        if (self.speculation_gate is None
+                or self.speculation_permanently_disabled or self.is_warmup):
+            return False, None
+
+        if (getattr(self.dist, 'has_pp', False)
+                and not self.dist.is_last_pp_rank):
+            return False, None
+        new_tokens_lens = getattr(sample_state.host, 'new_tokens_lens', None)
+        if new_tokens_lens is None:
+            return False, None
+        new_tokens_lens_list = (new_tokens_lens.tolist() if hasattr(
+            new_tokens_lens, 'tolist') else list(new_tokens_lens))
+        total_draft_tokens = 0
+        total_accepted_tokens = 0
+        for request in scheduled_batch.generation_requests:
+            draft_len = request.num_draft_tokens
+            if draft_len <= 0 or request.is_dummy:
+                continue
+            total_draft_tokens += draft_len
+            total_accepted_tokens += request.py_num_accepted_draft_tokens
+
+        if total_draft_tokens <= 0:
+            return False, None
+
+        acceptance_rate = total_accepted_tokens / total_draft_tokens
+        disabled_now, avg = self.speculation_gate.record_acceptance_rate(
+            acceptance_rate, sample_id=iteration_id)
+        if disabled_now:
+            self.speculation_permanently_disabled = True
+        return disabled_now, avg
+
     def _append_iter_stats(self,
                            stats: IterationStats,
                            req_stats: Optional[List[RequestStats]] = None):
@@ -1679,6 +1713,10 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]):
         if executed_batch is not None:
             with torch.cuda.nvtx.range("_handle_executed_batch_pp"):
                 self._update_requests(executed_batch.sample_state)
+                self._record_batch_acceptance_rate(
+                    executed_batch.scheduled_requests,
+                    executed_batch.sample_state,
+                    iteration_id=self.iter_counter)
 
                 scheduled_requests = executed_batch.scheduled_requests
                 if self.kv_cache_transceiver:
@@ -1745,6 +1783,12 @@ def _handle_dynamic_draft_len(self,
         if not hasattr(self.model_engine, 'max_draft_len'):
             return
 
+        if self.speculation_permanently_disabled:
+            for request in scheduled_batch.generation_requests:
+                request.py_draft_tokens = []
+            self.model_engine.runtime_draft_len = 0
+            return
+
         if (self.model_engine.spec_config is not None
                 and self.model_engine.spec_config.draft_len_schedule is not None
                 and self.model_engine.spec_config.spec_dec_mode.
@@ -1846,7 +1890,6 @@ 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()
-
         scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
         )
 
@@ -2134,6 +2177,10 @@ def _executor_loop(self):
 
                     self._update_request_states(scheduled_batch)
                     self._update_requests(sample_state, self.resource_manager)
+                    self._record_batch_acceptance_rate(
+                        scheduled_batch,
+                        sample_state,
+                        iteration_id=self.iter_counter)
 
                     self._send_kv_async(scheduled_batch.all_requests())
 
@@ -2376,6 +2423,10 @@ def _executor_loop_overlap(self):
 
                 if self.previous_batch is not None and should_process_previous_batch:
                     self._update_requests(self.previous_batch.sample_state)
+                    self._record_batch_acceptance_rate(
+                        self.previous_batch.scheduled_requests,
+                        self.previous_batch.sample_state,
+                        iteration_id=self.iter_counter)
 
                     self._send_kv_async(
                         self.previous_batch.scheduled_requests.all_requests())
@@ -3685,31 +3736,6 @@ def _handle_responses(self):
                     new_responses.append((req_id, response))
 
             if request_done:
-                if (self.drafter is not None and getattr(
-                        self.model_engine, 'enable_spec_decode', False)
-                        and not self.speculation_permanently_disabled
-                        and not request.is_dummy and not self.is_warmup):
-                    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)
-                            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 partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1,
                 # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse
                 # path is fixed.

tensorrt_llm/_torch/speculative/speculation_gate.py

@@ -1,77 +1,82 @@
 from collections import deque
-from typing import Optional, Tuple
+from typing import Deque, Optional, Tuple
 
 from tensorrt_llm.logger import logger
 
 
 class SpeculationGate:
     """
-	Tracks rolling average of accepted draft tokens per iteration over the last N completed requests.
-	Permanently disables speculation when average falls below a threshold.
-	"""
+    Tracks a rolling average of true acceptance-rate samples over the last N
+    speculation-enabled decoding iterations.
+
+    Permanently disables speculation when the rolling average falls below the
+    configured threshold.
+    """
 
     def __init__(self, window: int, threshold: float):
         self.window = window
         self.threshold = threshold
-        self.acceptance_history: Deque[float] = deque()
-        self.acceptance_sum: float = 0.0
-        self.num_completed_for_acceptance = 0
+        self.acceptance_rate_history: Deque[float] = deque()
+        self.acceptance_rate_sum: float = 0.0
+        self.num_recorded_samples = 0
         self.disabled = False
         logger.debug(
             f"[SpeculationGate] SpeculationGate initialized with window={self.window}, threshold={self.threshold}"
         )
 
     def reset(self) -> None:
-        self.acceptance_history.clear()
-        self.acceptance_sum = 0.0
-        self.num_completed_for_acceptance = 0
+        self.acceptance_rate_history.clear()
+        self.acceptance_rate_sum = 0.0
+        self.num_recorded_samples = 0
         self.disabled = False
 
-    def record_avg_decoded(
+    def record_acceptance_rate(
             self,
-            avg_decoded_tokens_per_iter: float,
-            request_id: Optional[int] = None) -> Tuple[bool, Optional[float]]:
+            acceptance_rate: float,
+            sample_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 one speculation-enabled iteration's true acceptance rate.
+
+        Returns (disabled_now, current_avg_acceptance_rate) where
+        disabled_now is True only when this call causes permanent 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:
+        if acceptance_rate is None:
             return False, None
 
-        accepted_len = 0.0
-        accepted_len = max(0.0, float(avg_decoded_tokens_per_iter) - 1.0)
+        acceptance_rate = float(acceptance_rate)
+        if acceptance_rate < 0.0 or acceptance_rate > 1.0:
+            raise ValueError("acceptance_rate must be in the range [0.0, 1.0], "
+                             f"got {acceptance_rate}")
 
-        # Log per-request completion for debug
-        if request_id is not None:
-            logger.debug(
-                f"[SpeculationGate] Request {request_id} completed: avg_decoded={avg_decoded_tokens_per_iter if avg_decoded_tokens_per_iter is not None else 'None'}, accepted_len={accepted_len:.3f}"
-            )
+        if sample_id is not None:
+            logger.debug(f"[SpeculationGate] Iteration {sample_id} recorded "
+                         f"acceptance_rate={acceptance_rate:.3f}")
 
         # 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()
-            self.acceptance_sum -= removed
+        self.acceptance_rate_history.append(acceptance_rate)
+        logger.debug(f"[SpeculationGate] Acceptance-rate history: "
+                     f"{self.acceptance_rate_history}")
+        self.acceptance_rate_sum += acceptance_rate
+        if len(self.acceptance_rate_history) > self.window:
+            removed = self.acceptance_rate_history.popleft()
+            self.acceptance_rate_sum -= removed
 
-        self.num_completed_for_acceptance += 1
+        self.num_recorded_samples += 1
 
-        if self.num_completed_for_acceptance >= self.window:
-            avg_accept = self.acceptance_sum / len(self.acceptance_history)
-            if avg_accept < self.threshold:
+        if self.num_recorded_samples >= self.window:
+            avg_acceptance_rate = (self.acceptance_rate_sum /
+                                   len(self.acceptance_rate_history))
+            if avg_acceptance_rate < self.threshold:
                 self.disabled = True
                 logger.info(
-                    f"[SpeculationGate] Speculative decoding disabled: rolling acceptance avg {avg_accept:.3f} < threshold {self.threshold} over last {self.window} requests"
-                )
-                return True, avg_accept
-            else:
-                # speculation is still enabled
-                return False, avg_accept
+                    "[SpeculationGate] Speculative decoding disabled: "
+                    f"rolling acceptance rate avg "
+                    f"{avg_acceptance_rate:.3f} < threshold "
+                    f"{self.threshold} over last {self.window} iterations")
+                return True, avg_acceptance_rate
+            return False, avg_acceptance_rate
 
         return False, None

tensorrt_llm/llmapi/llm_args.py

@@ -814,19 +814,23 @@ class DecodingBaseConfig(StrictBaseModel):
     load_format: Optional[str] = Field(
         default=None, description="The load format of the speculative model.")
 
-    acceptance_window: Optional[NonNegativeInt] = Field(
+    acceptance_rate_window_size: Optional[NonNegativeInt] = Field(
         default=None,
         description=
-        "The rolling average window size (N) for acceptance length across completed requests. "
+        "The rolling average window size (N) for acceptance rate across "
+        "speculation-enabled decoding iterations. "
         "If not set or set to 0, the feature is disabled. PyTorch backend only."
     )
 
-    acceptance_length_threshold: Optional[NonNegativeFloat] = Field(
+    acceptance_rate_threshold: Optional[float] = Field(
         default=None,
-        description=
-        "The threshold for average acceptance length; speculation will be disabled permanently once the "
-        "rolling average over the last N completed requests (N = acceptance_window) drops below this value. "
-        "PyTorch backend only.")
+        ge=0.0,
+        le=1.0,
+        description="The threshold for average true acceptance rate "
+        "(accepted_draft_tokens / drafted_tokens); speculation will be "
+        "disabled permanently once the rolling average over the last N "
+        "speculation-enabled decoding iterations "
+        "(N = acceptance_rate_window_size) drops below this value. ")
 
     allow_advanced_sampling: bool = Field(
         default=False,