[None][test] Assert MTP acceptance length in ADP + LM-head-TP accuracy tests by qiaoxj07 · Pull Request #16521 · NVIDIA/TensorRT-LLM

qiaoxj07 · 2026-07-17T03:36:13Z

Background

https://nvbugs/6460072 (fix in #16440): on main, the ADP + enable_lm_head_tp_in_adp + greedy draft-sampling path produced wrong MTP draft tokens (per-rank argmax on the wrong all-gather rows / a local vocab shard).

The bug shipped through green CI even though three parametrizations of TestDeepSeekR1::test_nvfp4_multi_gpus (latency_adp_lmtp, latency_trtllmgen_adp_lmtp, latency_adp_lmtp_tp4) exercise exactly this configuration — because accuracy scores cannot see this bug class: rejected drafts are regenerated by the target model, so outputs (and GSM8K/MMLU scores) stay correct and the only observable symptom is the acceptance length collapsing to ~1.0, i.e. a pure perf regression. In one production benchmark this manifested as a hidden +13% gen-only iteration-time regression (force-accepted wrong drafts degrading MoE routing balance).

Summary

Add an acceptance-length guard to the three ADP + LM-head-TP parametrizations:

enable enable_iter_perf_stats for those cases only (all other parametrizations keep the default False — zero behavior change),
after the accuracy tasks, read llm.get_stats() and assert the mean specDecodingStats.acceptanceLength over the retained iterations is >= 1.5.

Healthy DeepSeek-R1 + MTP3 greedy acceptance length is ~2.5–3.0; a broken draft path gives ~1.0, so the 1.5 threshold has wide margins on both sides. The stats-reading pattern follows the existing test_eagle3_spec_decoding_stats unit test.

With this assertion, a regression of the same class as https://nvbugs/6460072 fails these tests instead of silently passing.

Impact

Test-only change. No runtime behavior change for any other parametrization.

Related PRs

[https://nvbugs/6460072][fix] Split the caller — for use_lm_head_tp_in_adp and is_all_greedy_sample, keep… #16440 (the fix for the motivating bug)

Summary by CodeRabbit

Tests
- Enhanced validation for multi-GPU speculative decoding scenarios.
- Added runtime checks for draft-token acceptance, including a minimum average acceptance threshold.
- Improved diagnostics when speculative decoding produces unexpectedly low-quality drafts.

…y tests Wrong MTP draft tokens are invisible to accuracy scores: the target model rejects and regenerates them, so GSM8K/MMLU stay intact and the only observable symptom is a collapsed acceptance length (~1.0). This is how https://nvbugs/6460072 (wrong draft argmax in the ADP + LM-head-TP greedy path) shipped through green CI. Enable iteration stats for the latency_adp_lmtp, latency_trtllmgen_adp_lmtp and latency_adp_lmtp_tp4 parametrizations of TestDeepSeekR1::test_nvfp4_multi_gpus and assert the mean acceptance length over the profiled iterations stays >= 1.5 (healthy MTP3 is ~2.5, a broken draft path gives ~1.0). Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

qiaoxj07 · 2026-07-17T03:37:09Z

Merge-order dependency: #16440 must land first. The motivating bug is still live on main, so on the current tip these assertions will (correctly) fail with an acceptance length of ~1.0 — which is exactly the detection this PR adds. Will trigger /bot run after #16440 merges and this branch is rebased.

coderabbitai · 2026-07-17T03:39:25Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 156f1b2e-94cf-4073-a010-597155294864

📥 Commits

Reviewing files that changed from the base of the PR and between 1228fed and 21b9116.

📒 Files selected for processing (1)

tests/integration/defs/accuracy/test_llm_api_pytorch.py

📝 Walkthrough

Walkthrough

The integration test now conditionally enables iteration statistics and validates MTP speculative-decoding acceptance length for configurations using MTP with LM-head tensor parallelism in ADP.

Changes

MTP acceptance validation

Layer / File(s)	Summary
Conditional statistics and acceptance assertion `tests/integration/defs/accuracy/test_llm_api_pytorch.py`	Determines when acceptance validation applies, enables iteration performance statistics, computes average `acceptanceLength` from speculative-decoding stats, and requires an average of at least `1.5`.

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

Possibly related PRs

NVIDIA/TensorRT-LLM#15775: Adds related MTP speculative-decoding acceptance-length validation and compare/accept behavior.

Suggested reviewers: nv-guomingz

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the change: adding an MTP acceptance-length assertion to ADP + LM-head-TP accuracy tests.
Description check	✅ Passed	The description explains the bug, fix, impact, and related PR, but it omits the template's Test Coverage and PR Checklist sections.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

longlee0622 · 2026-07-17T05:22:43Z

+        # stays intact and the only observable symptom is a collapsed
+        # acceptance length (https://nvbugs/6460072). Assert on it for the
+        # ADP + LM-head-TP cases, whose draft sampling takes a dedicated path.
+        check_acceptance = mtp_nextn > 0 and enable_lm_head_tp_in_adp


@mikeiovine for viz. @qiaoxj07 contributed one such AL protection test. Is it suitable for test suite? I'm not sure if this is connected to any CI testing system already.

BowenFu

LGTM — test-only; the MTP acceptance-length assertion and iter-perf-stats are scoped to the 3 ADP+LM-head-TP parametrizations, no behavior change elsewhere.

qiaoxj07 requested review from a team as code owners July 17, 2026 03:36

qiaoxj07 requested review from ZhanruiSunCh, brnguyen2, fredricz-20070104, mlefeb01 and schetlur-nv July 17, 2026 03:36

github-actions Bot assigned qiaoxj07 Jul 17, 2026

longlee0622 reviewed Jul 17, 2026

View reviewed changes

BowenFu approved these changes Jul 17, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

[None][test] Assert MTP acceptance length in ADP + LM-head-TP accuracy tests#16521
qiaoxj07 wants to merge 1 commit into
NVIDIA:mainfrom
qiaoxj07:test/mtp-adp-lmtp-acceptance-guard

qiaoxj07 commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

qiaoxj07 commented Jul 17, 2026

Uh oh!

coderabbitai Bot commented Jul 17, 2026

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

longlee0622 Jul 17, 2026

Uh oh!

BowenFu left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

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

Uh oh!

Background

Summary

Impact

Related PRs

Summary by CodeRabbit

Uh oh!

qiaoxj07 commented Jul 17, 2026

Uh oh!

coderabbitai Bot commented Jul 17, 2026

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

longlee0622 Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

BowenFu left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

qiaoxj07 commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading