[#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"

[chungen04]

@coderabbitai summary

Description

Addressing Issue #15022.

With dynamic draft length (draft_len_schedule), the model engine tracks this as runtime_draft_len:

py_executor.py:2346-2362 per iteration, from the schedule

runtime_draft_len = get_draft_len_for_batch_size(... batch_size ...)
self.model_engine.runtime_draft_len = runtime_draft_len

When runtime_draft_len is, say, 1, the target model emits only 1 + 1 = 2 logit rows per generation request. But now, the guided decoder always laid out its bitmask for the static max
guided_decoder.py:335-337

def add_batch(self, scheduled_requests):
    self.requests = GuidedRequests.from_scheduled_requests(
        scheduled_requests, self.max_num_draft_tokens)   # <-- always the static max

So the guided decoder expected max_num_draft_tokens + 1 rows per request, while the model produced only runtime_draft_len + 1.

That row count flows into num_bitmask_tokens:

guided_decoder.py:110-116

def num_bitmask_tokens(self):
    return self.num_contexts + self.num_generations * (self.max_num_draft_tokens + 1)

and is then used to index both the bitmask and the logits tensor:

guided_decoder.py:328-332 (_apply_bitmask)

torch.ops.trtllm.logits_bitmask(
    logits[:num_bitmask_tokens],                    # logits only has runtime_draft_len+1 rows/req
    self.bitmask[:num_bitmask_tokens, ...],
    token_mask=self.token_mask[:num_bitmask_tokens], ...)

Because num_bitmask_tokens (computed from the static max) is larger than the actual number of logit rows the model produced, the indexing reads past the end of logits and crash.

At #10860, it let runtime_draft_len drop below the static max per iteration, but guided_decoder.py was not updated. So after #10860, the guided decoder kept laying out its bitmask for max_num_draft_tokens while the target emitted only runtime_draft_len + 1 rows per request.

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

Ready to act? Review this PR in Change Stack to turn feedback into patch suggestions you can inspect and refine.

📝 Walkthrough

Walkthrough

The PR threads a runtime draft length parameter through guided decoder APIs to bound draft-token processing. GuidedDecoder and CapturableGuidedDecoder now accept an optional runtime_draft_len parameter, which is used to compute the effective draft-token count for each batch. The draft-token loop in the _build method is bounded to iterate only over tokens within the runtime-reserved range. ModelEngine propagates its runtime draft length to the guided decoder, completing the wiring.

Changes

Runtime Draft Length Parameter

Layer / File(s)	Summary
Add runtime_draft_len parameter to guided decoder API tensorrt_llm/_torch/pyexecutor/guided_decoder.py	GuidedDecoder.add_batch and CapturableGuidedDecoder.add_batch accept optional runtime_draft_len parameter; when set, it overrides self.max_num_draft_tokens to compute num_draft_tokens for GuidedRequests construction.
Bound draft-token processing by runtime length tensorrt_llm/_torch/pyexecutor/guided_decoder.py	_build method iterates over req.draft_tokens[:requests.max_num_draft_tokens] to limit bitmask filling and token-mask marking within the runtime-reserved draft-token slots.
Propagate runtime_draft_len from model engine tensorrt_llm/_torch/pyexecutor/model_engine.py	PyTorchModelEngine._prepare_tp_inputs passes runtime_draft_len=self.runtime_draft_len to guided_decoder.add_batch, connecting engine runtime state to guided decoder bounding logic.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related issues

• [Bug]: Guided decoding (xgrammar) + EAGLE-3 + draft_len_schedule reaching 0 crashes during CUDA graph capture, "bitmask must have the same batch size as logits" #15022: The runtime_draft_len parameter and bounded draft-token handling address the guided_decoder bitmask/logits batch-size mismatch when draft_len=0.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	PR description explains the issue and root cause clearly, but Test Coverage section is empty and not all checklist items are addressed.	Complete the Test Coverage section by listing relevant tests that validate the fix, and ensure all PR checklist items are properly reviewed and addressed.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: fixing a crash when guided decoding is combined with EAGLE-3 and draft_len_schedule reaches 0, with specific error message context.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/guided_decoder.py`:
- Around line 341-344: Validate that runtime_draft_len does not exceed
self.max_num_draft_tokens before constructing GuidedRequests: if
runtime_draft_len is not None and runtime_draft_len > self.max_num_draft_tokens,
raise a ValueError (or clamp and log) so requests =
GuidedRequests.from_scheduled_requests(...) cannot create
requests.max_num_draft_tokens larger than the buffers allocated in __init__;
this prevents buffer overflows in _copy_bitmask and _apply_bitmask.
- Around line 484-487: Validate that runtime_draft_len does not exceed the
preallocated capacity before using it: inside GuidedDecoder (the constructor
block where runtime_draft_len is used to set num_draft_tokens and call
GuidedRequests.from_scheduled_requests) add a check comparing runtime_draft_len
(when not None) against self.max_num_draft_tokens and raise a clear ValueError
(or clamp/handle per project convention) if runtime_draft_len >
self.max_num_draft_tokens; mirror the same validation logic used in
GuidedDecoder.add_batch to prevent buffer overflows during
_copy_bitmask/_apply_bitmask.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6e2a708e-3128-4228-9570-48ff238ea55a

📥 Commits

Reviewing files that changed from the base of the PR and between 86f9602 and e6ffe66.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/guided_decoder.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py

[Priyanshu31102003]

Reviewing the current patch, the concern raised by @coderabbitai regarding potential buffer overflows if runtime_draft_len exceeds allocated capacity is highly valid.

Instead of overriding num_draft_tokens early during scheduling, a cleaner and defensively safer approach to align batch dimensions right before the custom kernel call inside _apply_bitmask is to validate shapes dynamically and enforce a hard boundary protection check against self.bitmask.size(0).

I have submitted a bullet-proof fix addressing this edge-case safely here: #15026

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
gh pr view 15026 --repo NVIDIA/TensorRT-LLM --json title,body,files,additions,deletions,state
gh pr diff 15026 --repo NVIDIA/TensorRT-LLM

Length of output: 5928

---

---

@Priyanshu31102003 Thanks for the detailed review and for referencing #15026. After examining the proposed diff, there are a few concerns worth raising with the PR author:

1. Silent expansion risk (Line if logits.size(0) > num_bitmask_tokens)
   When logits has more rows than num_bitmask_tokens, the proposed code silently expands num_bitmask_tokens to logits.size(0). The subsequent self.bitmask.size(0) clamp prevents an out-of-bounds kernel call, but any bitmask rows between the original num_bitmask_tokens and the expanded value may hold stale or uninitialized data from a previous iteration. That means the kernel could silently apply incorrect masks to those logit rows — arguably worse than an explicit crash, since it produces wrong outputs without any signal.

2. Unrelated change bundled in
   The examples/auto_deploy/model_registry/models.yaml change (enabling Nemotron Ultra 253B models) is unrelated to the bitmask crash fix and should ideally be a separate PR to keep the fix bisectable and reviewable on its own merits.

3. Indentation inconsistency
   The parameter indentation in _apply_bitmask is changed from the surrounding codebase style, which is a cosmetic noise addition to an otherwise focused bug fix.

The approach in the current PR (#15023) — adjusting num_draft_tokens upstream during GuidedRequests construction — ensures the bitmask and logits are consistently sized before any kernel call is made, which is architecturally cleaner. The key question for both approaches is whether all callers of add_batch correctly propagate runtime_draft_len; that's a valid concern to verify on #15023 itself.

[chungen04]

Despite there are no explicit guarding and clamping of runtime_draft_len, it is currently safe that runtime_draft_len do not exceed max_num_draft_tokens, justification:

• The caller with explicit runtime_draft_len assignment is at model_engine.py
• Assignments of runtime_draft_len includes:
  • ref1 -> assigned to self.max_total_draft_tokens.
  • ref2 (init), ref3 (init) -> assigned to self.max_total_draft_tokens or 0.
  • ref4 (capture) -> From _get_graphs_to_capture.
    • Case 1: deprecating soon.
    • Case 3: draft_lengths = [max_total_draft_tokens], optionally + [0].
    • Case 2: draft_len ranges over self._dynamic_draft_len_mapping.values(), built by _compute_dynamic_draft_len_mapping: for each graph batch size it picks schedule schedule[schedule_keys[key_idx]], else 0. Schedule values are Pydantic-validated max_draft_len src, and max_draft_len <= max_total_draft_tokens. max_total_draft_tokens == max_num_draft_tokens (src, src).

[achartier]

Could you clarify the intended order between this PR and #15026?

[chungen04]

@achartier I am not sure about the content and author in #15026 , looks like he proposed a PR quickly soon after my issue & PR, but there are some design choices to be justified there.

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #52859 [ run ] triggered by Bot. Commit: e760fa2 Link to invocation

[tensorrt-cicd]

PR_Github #52859 [ run ] completed with state FAILURE. Commit: e760fa2
/LLM/main/L0_MergeRequest_PR pipeline #42111 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

Please follow these steps to setup pre-commit and run it on this PR: https://github.com/NVIDIA/TensorRT-LLM/blob/main/CONTRIBUTING.md

[chungen04]

@achartier Done, please check

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #52871 [ run ] triggered by Bot. Commit: f82e6eb Link to invocation

[tensorrt-cicd]

PR_Github #52871 [ run ] completed with state FAILURE. Commit: f82e6eb
/LLM/main/L0_MergeRequest_PR pipeline #42123 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

@chungen04 Could you double check? yapf is still failing on the last version.

[chungen04]

@achartier Done, please check. Looks like I missed a message saying yapf was installing on my end.

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #53089 [ run ] completed with state FAILURE. Commit: 7b35bde
/LLM/main/L0_MergeRequest_PR pipeline #42299 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[achartier]

@achartier Should I do anything regarding the CI failure?

Infra issue. I rebased and relaunched the pipeline.

[tensorrt-cicd]

PR_Github #53121 [ run ] triggered by Bot. Commit: 819227e Link to invocation

[tensorrt-cicd]

PR_Github #53121 [ run ] completed with state FAILURE. Commit: 819227e
/LLM/main/L0_MergeRequest_PR pipeline #42327 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #53239 [ run ] triggered by Bot. Commit: 0e4f0b4 Link to invocation

[tensorrt-cicd]

PR_Github #53239 [ run ] completed with state FAILURE. Commit: 0e4f0b4
/LLM/main/L0_MergeRequest_PR pipeline #42434 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #53371 [ run ] triggered by Bot. Commit: 0e4f0b4 Link to invocation

[tensorrt-cicd]

PR_Github #53371 [ run ] completed with state FAILURE. Commit: 0e4f0b4
/LLM/main/L0_MergeRequest_PR pipeline #42549 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #53438 [ run ] triggered by Bot. Commit: 0e4f0b4 Link to invocation

[tensorrt-cicd]

PR_Github #53438 [ run ] completed with state FAILURE. Commit: 0e4f0b4
/LLM/main/L0_MergeRequest_PR pipeline #42606 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #53599 [ run ] triggered by Bot. Commit: 0e4f0b4 Link to invocation

[tensorrt-cicd]

PR_Github #53599 [ run ] completed with state SUCCESS. Commit: 0e4f0b4
/LLM/main/L0_MergeRequest_PR pipeline #42745 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #53655 [ run ] triggered by Bot. Commit: 0e4f0b4 Link to invocation

[tensorrt-cicd]

PR_Github #53655 [ run ] completed with state SUCCESS. Commit: 0e4f0b4
/LLM/main/L0_MergeRequest_PR pipeline #42797 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chungen04]

Oh, here we go. Thank you @achartier for the support!