From 8eb0aca4ecb47477402e75b4b35f7d8c663b0927 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 cbf4004d2052b0554f6ca2d1cf327a692dd29121 Mon Sep 17 00:00:00 2001 From: Sihan Wang Date: Wed, 17 Jun 2026 22:02:12 +0000 Subject: [PATCH 2/2] [None][fix] spec-decode NaN guard: check ~isnan, not isfinite (guided-decoding -inf is legal) The spec-decode target-logits guard used torch.isfinite(logits).all(), which is False for -inf as well as NaN. Guided decoding (xgrammar) writes -inf into masked-token logits as normal grammar masking, so the guard false-positively terminated the engine (os._exit) on any structured-output request under speculative decoding (Eagle3/MTP). Check NaN only -- the actual argmax->PAD corruption the guard targets; -inf is legitimate and handled by softmax/argmax. Restores the NaN-only intent of the original guard. 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 477a9123d7db..e7c828919f05 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/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 diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 3f60b11bd1a4..90417cb5bb12 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/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: