From edd543128e5b22c33079db58418e27e157626ae0 Mon Sep 17 00:00:00 2001 From: Chung-En Ho <75154854+chungen04@users.noreply.github.com> Date: Thu, 11 Jun 2026 16:10:47 -0700 Subject: [PATCH 1/2] [#15022][fix] Guided decoding (xgrammar) + EAGLE-3 + draft_len_schedule reaching 0 crashes during CUDA graph capture, "bitmask must have the same batch size as logits" (#15023) Signed-off-by: chungen04 --- .../_torch/pyexecutor/guided_decoder.py | 25 ++++++++++++++----- .../_torch/pyexecutor/model_engine.py | 6 +++-- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index a6b23d0a45ca..8b1d1b221eae 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/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)) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 7599e1e4e6ca..4210748326b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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, From 1fadaba6065e4babe6964d83bacb849f97c52eb9 Mon Sep 17 00:00:00 2001 From: Sihan Wang Date: Wed, 17 Jun 2026 22:02:46 +0000 Subject: [PATCH 2/2] [None][fix] spec-decode NaN guard: check ~isnan, not isfinite/any (guided-decoding -inf is legal) Standardize the spec-decode target-logits guard on (~torch.isnan(logits)).all() so guided-decoding -inf (legitimate grammar masking) never trips it, while real NaN (argmax->PAD corruption) still does. Aligns deep-main with the priority build branch and the shipped rc16 images. Signed-off-by: Sihan Wang --- tensorrt_llm/_torch/speculative/interface.py | 2 +- tensorrt_llm/_torch/speculative/mtp.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 1c4f8e63a376..ce78ff34d21b 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1306,7 +1306,7 @@ def _sample_tokens_for_batch( # in SpecSamplerBase.update_requests where there's already a sync. # See mtp.MTPWorker.sample_and_accept_draft_tokens for the analogous # check on the MTP-relaxed and MTP-strict-thop paths. - self.logits_finite = torch.isfinite(logits).any(dim=-1).all() + self.logits_finite = (~torch.isnan(logits)).all() # Counts for the failure diagnostic; only inspected on the crash # path. Cheap GPU reductions, no host sync, no fresh allocations # (safe inside CUDA graph capture). diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 17f953ec503d..33c9db23eabc 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -794,7 +794,7 @@ def sample_and_accept_draft_tokens( # -inf values from guided-decoding masks are legitimate; rows with # a mix of -inf and finite still have a sampleable max and must not # trip this guard. - self.logits_finite = torch.isfinite(logits).any(dim=-1).all() + self.logits_finite = (~torch.isnan(logits)).all() # Counts for the failure diagnostic; only inspected on the crash # path. Cheap GPU reductions, no host sync, no fresh allocations # (safe inside CUDA graph capture).