Repository navigation
[https://nvbugs/6948129][fix] Test H=96 single-token MLA routing at its measured batch floor - #20652
Conversation
…ts batch floor NVIDIA#18599 raised the measured FP8 perf-table entry for single-token H=96 MLA decode, _PERF_MIN_BATCH_FP8[(96, 1)], from 1 to 8. Below batch 8, CuteDslMlaFmha now declines that shape and selection falls through to TRTLLM-Gen, as intended. NVIDIA#19960 then added test_mla_generation_routes_wide_head_fp8_verify_away_from_trtllm_gen, whose helper hard-codes batch_size=4 and whose [103-96-1-True-False] row still expects CuTe DSL under the old entry of 1. Pass the batch size to the helper. The single-token row now runs at the batch-8 floor, where its perf-table entry admits CuTe DSL. The multi-token rows stay at batch 4, below the floor, so they keep checking that the verify-group carve-outs do not depend on the batch size. Add a batch-4 single-token case that pins the below-floor fall-through to TRTLLM-Gen, and remove this bug's waiver. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
WalkthroughThe MLA selection helper now receives batch size from its callers. Routing tests cover batch sizes 8 and 4, including TRTLLM-Gen selection for single-token FP8 H=96 requests at batch size 4. The associated skip entry is removed. ChangesMLA routing coverage
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~8 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title clearly identifies the NVBug, fix type, and primary change: testing H=96 single-token MLA routing at the measured batch floor. Description check The description explains the root cause, fix, affected test cases, waiver removal, reproduction result, and test plan. It does not include a separate PR Checklist section, but the key required informa… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Full details: Description check
-
Explanation
The description explains the root cause, fix, affected test cases, waiver removal, reproduction result, and test plan. It does not include a separate PR Checklist section, but the key required information is present.
-
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
@PerkzZheng, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~ Assignment explanation
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once. |
|
/bot run --disable-fail-fast |
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
|
PR_Github #79131 [ run ] triggered by Bot. Commit: |
|
PR_Github #79131 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79215 [ run ] triggered by Bot. Commit: |
|
PR_Github #79215 [ run ] completed with state |
Check the FP8 batch floor of single-token H=96 decode at 4 / 7 / 8 / 9 and that BF16 KV is not gated, and mark the routing tests cpu_only so they run in the CPU-Generic stages. Signed-off-by: Vismay Shrouty <vismayshrouty@rediffmail.com> [Zihua Wu: rebased onto main after NVIDIA#20652, which added the batch_size argument to the routing helper; the threshold cases now use it directly.] Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Summary
_PERF_MIN_BATCH_FP8[(96, 1)]from 1 to 8, but the [None][feat] Kimi K3 optimizations on Rubin #19960 test helper hard-codes batch 4, so the single-token H=96 FP8 row correctly routes to TRTLLM-Gen while still expecting CuTe DSL.batch_sizeas a required keyword; only the single-token row uses batch 8, the nine multi-token rows stay at 4, a new exact test checks that batch 4 falls through to TRTLLM-Gen, and this bug's waiver is removed.bash .repair-bot/handoff/run_repro.shTest plan
Links
Reproduction comparison
The failure on ToT is the one the report describes, so the bug is still present on ToT.
Dev Engineer Review
The test helper now requires an explicit
batch_size. H=96 FP8 single-token tests use batch 8, while multi-token tests use batch 4. A new batch-4 test verifies fallback from Cute DSL MLA to TRTLLM-Gen below the performance floor. No production source changes are present.QA Engineer Review
Modified
tests/unittest/_torch/attention/test_cute_dsl_mla.pyand removed the NVBug 6948129 waiver fromtests/integration/test_lists/waives.txt. Coverage includes H=96 FP8 routing above and below the batch threshold. Coverage verdict: needs follow-up, because test execution results and CI/manual-QA list membership were not supplied.Per-File QA Perspective
tests/unittest/_torch/attention/test_cute_dsl_mla.py: Covers H=96 FP8 single-token routing at batch 8, multi-token routing at batch 4, and TRTLLM-Gen fallback at batch 4. CI or manual-QA list membership is not confirmed.tests/integration/test_lists/waives.txt: Removes the waiver for the MLA routing test. QA should confirm that the test passes without the waiver.