From 22a80a289addeb1b23f9f756f716128231c4b423 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Fri, 5 Jun 2026 23:06:57 +0000 Subject: [PATCH 1/2] Apply fix Signed-off-by: chungen04 --- .../_torch/pyexecutor/guided_decoder.py | 24 ++++++++++++++----- .../_torch/pyexecutor/model_engine.py | 3 ++- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index fb505fe0a869..bf13411d3fdd 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -259,8 +259,11 @@ 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 @@ -332,9 +335,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]]: @@ -470,9 +477,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 8c0ee9d0edf4..3902873eb1b5 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2634,7 +2634,8 @@ 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) + new_tokens=new_tokens_device, + runtime_draft_len=self.runtime_draft_len) if self._can_use_incremental_update(scheduled_requests, new_tokens_device, From 0e4f0b459c730ca05141aaf62f6b299e589df7a3 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Mon, 8 Jun 2026 23:16:04 +0000 Subject: [PATCH 2/2] modified for pre-commit hook formatting Signed-off-by: chungen04 --- tensorrt_llm/_torch/pyexecutor/guided_decoder.py | 3 ++- tensorrt_llm/_torch/pyexecutor/model_engine.py | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index bf13411d3fdd..b8d06c1cb97a 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -263,7 +263,8 @@ def _build(self, requests: GuidedRequests) -> List[Tuple[int, str]]: # 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): + req.draft_tokens[:requests.max_num_draft_tokens], + 1): accepted = matcher.accept_token(tid) if not accepted: break diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3902873eb1b5..1e88a65de1e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2633,9 +2633,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, - runtime_draft_len=self.runtime_draft_len) + 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,