[None][fix] spec-decode NaN guard (~isnan, not isfinite) + draft_len_schedule CUDA-graph crash

tensorrt_llm/_torch/pyexecutor/guided_decoder.py

@@ -260,8 +260,12 @@ def _build(self, requests: GuidedRequests) -> List[Tuple[int, str]]:
                     matcher.fill_next_token_bitmask(self.bitmask_host, offset)
                     self.token_mask_host[offset] = 1
                     self.num_guided_tokens[slot] += 1
-                    # Process draft tokens
-                    for i, tid in enumerate(req.draft_tokens, 1):
+                    # Process draft tokens. Bound by the layout's draft length:
+                    # the new_tokens buffer always holds the static max, but only
+                    # `max_num_draft_tokens` slots are reserved this iteration.
+                    for i, tid in enumerate(
+                            req.draft_tokens[:requests.max_num_draft_tokens],
+                            1):
                         accepted = matcher.accept_token(tid)
                         if not accepted:
                             break
@@ -333,9 +337,13 @@ def _apply_bitmask(self,
             d2t=d2t)
 
     @nvtx_range("GuidedDecoder.add_batch")
-    def add_batch(self, scheduled_requests: ScheduledRequests) -> None:
+    def add_batch(self,
+                  scheduled_requests: ScheduledRequests,
+                  runtime_draft_len: Optional[int] = None) -> None:
+        num_draft_tokens = (self.max_num_draft_tokens
+                            if runtime_draft_len is None else runtime_draft_len)
         self.requests = GuidedRequests.from_scheduled_requests(
-            scheduled_requests, self.max_num_draft_tokens)
+            scheduled_requests, num_draft_tokens)
 
     @nvtx_range("GuideDecoder.build")
     def build(self) -> List[Tuple[int, str]]:
@@ -471,9 +479,14 @@ def __init__(self,
     @nvtx_range("GuidedDecoder.add_batch")
     def add_batch(self,
                   scheduled_requests: ScheduledRequests,
-                  new_tokens: Optional[torch.Tensor] = None) -> None:
+                  new_tokens: Optional[torch.Tensor] = None,
+                  runtime_draft_len: Optional[int] = None) -> None:
+        # See GuidedDecoder.add_batch: the layout must follow the runtime draft
+        # length so the captured graph's bitmask matches the target logits.
+        num_draft_tokens = (self.max_num_draft_tokens
+                            if runtime_draft_len is None else runtime_draft_len)
         self.requests = GuidedRequests.from_scheduled_requests(
-            scheduled_requests, self.max_num_draft_tokens)
+            scheduled_requests, num_draft_tokens)
         if new_tokens is not None:
             self.new_tokens.copy_(new_tokens.squeeze(-1), non_blocking=True)
         self.queue.put((self.requests, new_tokens is not None))

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -2339,8 +2339,10 @@ def _prepare_tp_inputs(
 
         # Must be before the update of py_batch_idx
         if self.guided_decoder is not None:
-            self.guided_decoder.add_batch(scheduled_requests,
-                                          new_tokens=new_tokens_device)
+            self.guided_decoder.add_batch(
+                scheduled_requests,
+                new_tokens=new_tokens_device,
+                runtime_draft_len=self.runtime_draft_len)
 
         if self._can_use_incremental_update(scheduled_requests,
                                             new_tokens_device,

tensorrt_llm/_torch/speculative/interface.py

@@ -1299,7 +1299,7 @@ def _sample_tokens_for_batch(
         # already a sync. See mtp.MTPWorker.sample_and_accept_draft_tokens
         # for the analogous check on the MTP-relaxed and MTP-strict-thop
         # paths that bypass this method.
-        self.logits_finite = torch.isfinite(logits).all()
+        self.logits_finite = (~torch.isnan(logits)).all()
 
         if spec_metadata.allow_advanced_sampling:
             num_gens = batch_size - num_contexts

tensorrt_llm/_torch/speculative/mtp.py

@@ -788,7 +788,7 @@ def sample_and_accept_draft_tokens(
         # propagate this bool through the worker forward output dict so the
         # natural sync in SpecSamplerBase.update_requests can check it without
         # an extra host sync — safe under CUDA graph capture.
-        self.logits_finite = torch.isfinite(logits).all()
+        self.logits_finite = (~torch.isnan(logits)).all()
 
         # The return buffer
         if self.spec_config.use_relaxed_acceptance_for_thinking or not self.is_thop: