[https://nvbugs/6460072][fix] Split the caller — for `use_lm_head_tp_in_adp and is_all_greedy_sample`, keep… by trtllm-agent · Pull Request #16440 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-15T16:03:23Z

Summary

Root cause: MTP-Eagle greedy path under ADP+LM-head-TP was trimming vocab-sharded logits to token_count and then plain-argmaxing per rank, so every rank picked a local-vocab-shard token (not the global argmax) and rows [:token_count] weren't even this rank's real requests (allgather-along-dim=0 order).
Fix: Split the caller — for use_lm_head_tp_in_adp and is_all_greedy_sample, keep padded logits, call self.draft_sampler(logits, draft_model.mtp_layers[0].shared_head.mapping_lm_head_tp) (which all-gathers local argmax/max across mapping_lm_head_tp ranks and picks global argmax), then trim result to token_count. Non-greedy path unchanged.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6460072

Summary by CodeRabbit

Bug Fixes
- Improved greedy token selection for speculative decoding when tensor-parallel language-head processing is enabled.
- Ensured draft tokens are sampled from correctly sharded logits, improving sampling correctness and consistency.
Refactor
- Clarified internal handling and documentation for greedy sampling across parallel decoding modes.

coderabbitai · 2026-07-15T16:06:14Z

📝 Walkthrough

Walkthrough

Eagle3 now routes all-greedy LM-head-TP-in-ADP draft-token selection through the TP-aware sampler, slices results to token_count, and retains the existing decoder path for other cases. Decoder comments clarify the upstream handling distinction.

Changes

Eagle3 draft sampling

Layer / File(s)	Summary
Greedy sampling routing and decoder semantics `tensorrt_llm/_torch/speculative/eagle3.py`	`_forward_linear_draft_loop` uses `draft_sampler` with `mapping_lm_head_tp` for all-greedy LM-head-TP-in-ADP execution, slices to `token_count`, and preserves `draft_decoder` for other cases. Comments document the corresponding decoder routing behavior.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: asfiyab-nvidia, zhaoyangwang-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 is specific and matches the main fix, including the NVBugs ID and fix type.
Description check	✅ Passed	The description covers the issue, fix, test plan, and bug link, with only the template's PR checklist section missing.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 1

🤖 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/speculative/eagle3.py`:
- Around line 1254-1257: Update the comment near the plain-TP and
LM-head-TP-in-ADP routing explanation to reference `_forward_linear_draft_loop`
instead of `_sample_draft_tokens_common`, preserving the existing description of
pre-trim padded logits and `mapping_lm_head_tp`.

🪄 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: f2d42e53-7ade-4ba6-ba37-8e0ef615b6dd

📥 Commits

Reviewing files that changed from the base of the PR and between 0623958 and 8d64e31.

📒 Files selected for processing (1)

tensorrt_llm/_torch/speculative/eagle3.py

Tabrizian · 2026-07-16T00:43:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T00:49:56Z

PR_Github #59562 [ run ] triggered by Bot. Commit: 8d64e31 Link to invocation

tensorrt-cicd · 2026-07-16T13:43:52Z

PR_Github #59562 [ run ] completed with state SUCCESS. Commit: 8d64e31
/LLM/main/L0_MergeRequest_PR pipeline #48007 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

zhaoyangwang-nvidia · 2026-07-17T02:49:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-17T02:55:10Z

PR_Github #59846 [ run ] triggered by Bot. Commit: a28ed86 Link to invocation

…-Eagle greedy draft under attention DP The draft_sampler consolidation dropped the ADP+LM-head-TP branch: under enable_lm_head_tp_in_adp the MTP shared head returns the LM-head-TP group's row-stacked, vocab-sharded logits, but the worker trimmed rows [:token_count] (group rank 0's requests, not this rank's) and greedy sampling fell through to a plain per-rank argmax over a local vocab shard, producing wrong draft tokens on every rank. Restore the group-wide distributed argmax in greedy_sample_draft_with_tp_gather: combine each rank's local (global index, max) over mapping_lm_head_tp, slice this rank's own row segment at tp_rank * max_num_requests, and only then trim the padding to token_count in the caller. Advanced (rejection) draft sampling asserts it never sees this layout, as rejection is config-gated off under attention DP. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

zhaoyangwang-nvidia · 2026-07-17T05:28:00Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-17T05:32:26Z

PR_Github #59846 [ run ] completed with state SUCCESS. Commit: a28ed86
/LLM/main/L0_MergeRequest_PR pipeline #48252 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

tensorrt-cicd · 2026-07-17T05:33:20Z

PR_Github #59888 [ run ] triggered by Bot. Commit: 3a94f75 Link to invocation

tensorrt-cicd · 2026-07-17T09:17:32Z

PR_Github #59888 [ run ] completed with state SUCCESS. Commit: 3a94f75
/LLM/main/L0_MergeRequest_PR pipeline #48289 completed with status: 'SUCCESS'

CI Report

Link to invocation

trtllm-agent requested a review from a team as a code owner July 15, 2026 16:03

trtllm-agent requested review from asfiyab-nvidia and zhaoyangwang-nvidia July 15, 2026 16:03

github-actions Bot assigned trtllm-agent Jul 15, 2026

coderabbitai Bot reviewed Jul 15, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/speculative/eagle3.py Outdated

Tabrizian requested a review from xwang233 July 15, 2026 16:32

longlee0622 changed the title ~~[https://nvbugs/6460072][fix] Split the caller — for use_lm_head_tp_in_adp and is_all_greedy_sample, keep…~~ [None][fix] Split the caller — for use_lm_head_tp_in_adp and is_all_greedy_sample, keep… Jul 16, 2026

longlee0622 changed the title ~~[None][fix] Split the caller — for use_lm_head_tp_in_adp and is_all_greedy_sample, keep…~~ [https://nvbugs/6460072][fix] Split the caller — for use_lm_head_tp_in_adp and is_all_greedy_sample, keep… Jul 16, 2026

xwang233 reviewed Jul 16, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/speculative/eagle3.py Outdated

zhaoyangwang-nvidia self-assigned this Jul 16, 2026

zhaoyangwang-nvidia force-pushed the repair-bot-bug6460072 branch 2 times, most recently from 46f8aae to a28ed86 Compare July 17, 2026 02:37

qiaoxj07 mentioned this pull request Jul 17, 2026

[None][test] Assert MTP acceptance length in ADP + LM-head-TP accuracy tests #16521

Merged

zhaoyangwang-nvidia force-pushed the repair-bot-bug6460072 branch from a28ed86 to 3a94f75 Compare July 17, 2026 03:49

zhaoyangwang-nvidia approved these changes Jul 17, 2026

View reviewed changes

zhaoyangwang-nvidia enabled auto-merge (squash) July 17, 2026 09:09

zhaoyangwang-nvidia merged commit 28be4ae into NVIDIA:main Jul 17, 2026
8 checks passed

longlee0622 deleted the repair-bot-bug6460072 branch July 17, 2026 09:24

zhaoyangwang-nvidia mentioned this pull request Jul 17, 2026

[None][feat] Support rejection sampling under attention DP (incl. LM-head TP) #16544

Merged

1 task

Uh oh!

Conversation

trtllm-agent commented Jul 15, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Tabrizian commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

zhaoyangwang-nvidia commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

zhaoyangwang-nvidia commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

tensorrt-cicd commented Jul 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

trtllm-agent commented Jul 15, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 15, 2026 •

edited

Loading