fix: Align bitmask batch size with logits tensor dimension inside GuidedDecoder

[Priyanshu31102003]

Description

Fixes a RuntimeError: bitmask must have the same batch size as logits crash during CUDA graph capture warmup when draft_len_schedule drops to 0.

Changes

• Implemented dynamic runtime shape alignment checks inside _apply_bitmask (guided_decoder.py) to synchronize num_bitmask_tokens with the actual batch size of the logits tensor when speculative generation switches modes.
• Added strict out-of-bounds boundary protection against self.bitmask.size(0) to structurally prevent any potential buffer overflows, safely addressing the architectural edge-case missing in other temporary patches.

Closes #15022

Summary by CodeRabbit

• New Features

  • Enabled deployment of NVIDIA Llama-3.1-Nemotron-Ultra-253B models (standard and FP8 variants).

• Bug Fixes

  • Improved stability of guided decoding during inference to prevent potential indexing errors.

[coderabbitai]

Complex PR? Review this PR in Change Stack to move by importance, not file order.

📝 Walkthrough

Walkthrough

This PR addresses a CUDA graph capture crash in guided decoding by defensively computing and clamping num_bitmask_tokens in GuidedDecoder._apply_bitmask, and enables two Nemotron Ultra-253B model configurations in the deployment registry.

Changes

Guided Decoding and Model Configuration

Layer / File(s)	Summary
Guided decoder bitmask defensive clamping tensorrt_llm/_torch/pyexecutor/guided_decoder.py	_apply_bitmask derives num_bitmask_tokens from requests when available, falls back to logits.size(0), and clamps to both logits batch/token dimension and allocated self.bitmask size to prevent out-of-bounds indexing during CUDA graph capture.
Model registry configuration update examples/auto_deploy/model_registry/models.yaml	Two previously commented-out nvidia/Llama-3_1-Nemotron-Ultra-253B-v1 entries (standard and FP8) are enabled with config_id: simple_shard_only and matching yaml_extra composition.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14686: Both PRs update model registry entries in examples/auto_deploy/model_registry/models.yaml to enable/adjust model configurations.

Suggested reviewers

• bmarimuthu-nv
• QiJune

🚥 Pre-merge checks | ✅ 2 | ❌ 3

❌ Failed checks (2 warnings, 1 inconclusive)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The changes include enabling NVIDIA Llama/Nemotron model registry entries in models.yaml, which is unrelated to the bitmask-logits alignment fix described in the PR objectives.	Remove the models.yaml changes (uncommenting model entries) from this PR as they are out of scope; create a separate PR for model registry updates.
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	The PR description explains the issue and changes clearly, but lacks explicit test coverage and PR checklist information as specified in the template.	Add a 'Test Coverage' section listing relevant tests that verify the bitmask-logits alignment fix and confirm CUDA-graph capture succeeds with draft_len_schedule including 0.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main fix: aligning bitmask batch size with logits tensor dimension in GuidedDecoder, which directly addresses the core issue.
Linked Issues check	✅ Passed	The PR addresses the core objectives from #15022: implements dynamic runtime shape alignment checks in _apply_bitmask and adds out-of-bounds protection against self.bitmask.size(0).

_{✏️ 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: 1

🧹 Nitpick comments (1)

tensorrt_llm/_torch/pyexecutor/guided_decoder.py (1)

309-330: ⚡ Quick win

Consider consolidating the clamping logic into a single expression.

After fixing the bug at lines 311-312, the three-step clamping pattern can be simplified for clarity:

♻️ Proposed refactor

 if num_bitmask_tokens is None:
     num_bitmask_tokens = requests.num_bitmask_tokens if requests is not None else logits.size(0)

-if logits.size(0) > num_bitmask_tokens:
-    num_bitmask_tokens = logits.size(0)
-
 vocab_size_padded = self.vocab_size_padded if d2t is None else d2t.size(
     0)
 assert vocab_size_padded % logits.size(1) == 0
 tp_size = vocab_size_padded // logits.size(1)
 assert self.bitmask_size % tp_size == 0
 tp_rank = self.rank % tp_size
 bitmask_start = tp_rank * self.bitmask_size // tp_size
 bitmask_end = bitmask_start + self.bitmask_size // tp_size

 if d2t is not None:
     d2t_start = tp_rank * vocab_size_padded // tp_size
     d2t_end = d2t_start + vocab_size_padded // tp_size
     d2t = d2t[d2t_start:d2t_end]

-if self.bitmask.size(0) < num_bitmask_tokens:
-    num_bitmask_tokens = self.bitmask.size(0)
-
+num_bitmask_tokens = min(num_bitmask_tokens, logits.size(0), self.bitmask.size(0))
+
 torch.ops.trtllm.logits_bitmask(

🤖 Prompt for 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.

In `@tensorrt_llm/_torch/pyexecutor/guided_decoder.py` around lines 309 - 330,
Replace the three-step clamping of num_bitmask_tokens with a single expression
that computes the initial value (requests.num_bitmask_tokens if requests is not
None else logits.size(0)), then clamps it to be at least logits.size(0) and at
most self.bitmask.size(0); i.e. set num_bitmask_tokens =
min(self.bitmask.size(0), max(logits.size(0), initial_value)). Apply this change
where num_bitmask_tokens is defined/updated in guided_decoder.py so the later
d2t handling and other logic using num_bitmask_tokens remain unchanged.

🤖 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 311-312: The comparison is reversed causing num_bitmask_tokens to
grow instead of being clamped to the logits batch size; change the logic around
logits and num_bitmask_tokens so that if num_bitmask_tokens is greater than
logits.size(0) you set num_bitmask_tokens = logits.size(0). Update the block
using logits, num_bitmask_tokens, logits_bitmask and self.bitmask in
guided_decoder.py to ensure you slice logits[:num_bitmask_tokens] and
self.bitmask[:num_bitmask_tokens, ...] with matching batch sizes (i.e., clamp
down num_bitmask_tokens to logits.size(0) before calling logits_bitmask).

---

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/guided_decoder.py`:
- Around line 309-330: Replace the three-step clamping of num_bitmask_tokens
with a single expression that computes the initial value
(requests.num_bitmask_tokens if requests is not None else logits.size(0)), then
clamps it to be at least logits.size(0) and at most self.bitmask.size(0); i.e.
set num_bitmask_tokens = min(self.bitmask.size(0), max(logits.size(0),
initial_value)). Apply this change where num_bitmask_tokens is defined/updated
in guided_decoder.py so the later d2t handling and other logic using
num_bitmask_tokens remain unchanged.

🪄 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: d3abe478-c84f-4506-b295-944adde63110

📥 Commits

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

📒 Files selected for processing (2)

• examples/auto_deploy/model_registry/models.yaml
• tensorrt_llm/_torch/pyexecutor/guided_decoder.py

[achartier]

Is this still needed now that #15023 is merged?

[karljang]

@Priyanshu31102003 ,
Thanks for your contribution!

I think the guided-decoding fix here has been superseded by #15023, which has already merged. #15023 addresses the root cause by making the guided-decoding bitmask layout follow the runtime draft length and passing runtime_draft_len through from model_engine.py.

Also, the examples/auto_deploy/model_registry/models.yaml change seems unrelated to the guided-decoding fix. If that registry update is still needed, please split it into a separate PR.

Given #15023 is merged, I think we can close this PR as superseded. Please feel free to open a new one if the registry update is still needed.

Thank you!