diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index fb505fe0a869..b8d06c1cb97a 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -259,8 +259,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 @@ -332,9 +336,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 +478,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..1e88a65de1e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2633,8 +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) + 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,