Repository navigation
[None][perf] MiniMax-M3: native P128 draft KV view for the NVFP4 hybrid cache - #19731
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review. WalkthroughHybrid NVFP4 shared Eagle draft layers retain FP8/P128 pages and expose them through a draft KV cache view. Speculative lookup and dense-decode page handling use the revised view. The P32 subpage view and scatter path are removed. ChangesHybrid draft cache view
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant SpecWorkerBase
participant get_draft_kv_cache_manager
participant MiniMaxM3KVCacheManagerV2
participant MiniMaxM3DraftKVCacheView
participant copy_batch_block_offsets_to_device
SpecWorkerBase->>get_draft_kv_cache_manager: request shared draft KV view
get_draft_kv_cache_manager->>MiniMaxM3KVCacheManagerV2: call get_draft_kv_cache_view
MiniMaxM3KVCacheManagerV2->>MiniMaxM3DraftKVCacheView: create and cache the view
MiniMaxM3DraftKVCacheView->>copy_batch_block_offsets_to_device: convert and copy block offsets
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The inspected draft-cache paths preserve the expected P128 page mapping. No actionable issue remains from the available evidence; merge after normal checks. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
…id cache With NVFP4 KV, the shared Eagle3 draft layer used a physical-P32 sub-page view (MiniMaxM3DraftSubpageView). Its copy_batch_block_offsets rebuilt the draft block table in Python/numpy on every decode prepare(), over each request's full padded page row (max_blocks_per_seq, 8196 for max_seq_len 1M), then expanded it x4 on the CPU and copied it H2D. That costs ~0.5 ms per request per decode step on the host and makes large-batch NVFP4 decode host-bound (B300 TP4, 128k context, batch 32: 38.4 ms/step vs 19.9 ms with FP8 KV). Run the Eagle layer at native FP8/P128 instead, as feat/m3_with_msa does since NVIDIA#18066: - _build_cache_config only strips the NVFP4 block-scale buffers from the dense/Eagle layers; no tokens_per_block_override for the Eagle layer. - New MiniMaxM3DraftKVCacheView: FP8, the trtllm-gen P128 opt-in, one pool rooted at the draft layer's K page, and a block table filled by a single copy_batch_block_offsets_to_device call with the view's own index_scales/kv_offset over the draft layer's lifecycle-pool slot rows. - Remove the P32 sub-page machinery (MiniMaxM3DraftSubpageView, DRAFT_PHYSICAL_TOKENS_PER_BLOCK, the P32 page-index override and helpers, the unused sub-paged scatter kernel) and inline the SWA-scratch expansion check back into KVCacheManagerV2._prepare_swa_scratch_copy_tensors, dropping the NVIDIA#19422-only hook. - Rename get_draft_subpage_view -> get_draft_kv_cache_view and is_fp8_subpaged_layer -> is_hybrid_draft_layer; update the unit tests. B300 (SM103), TP only, max_draft_len=3, measured on NVIDIA#19422 rev 4d7e78f plus this patch: batch-32 decode 38.4 -> 20.9 ms/step (FP8 KV 19.9). InferenceX AgentX P90 interactivity with forced acceptance: TP4 C40 79.5 -> 100.3, TP2 C45 44.6 -> 50.6. Rerun on main 00bdcb6 (interleaved): TP4 C40 81.7 -> 100.4, TP2 C45 46.0 -> 50.4. Draft acceptance (24-prompt greedy probe) identical; GSM8K accuracy unchanged within run-to-run noise. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
b91bad7 to
a436008
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #76196 [ run ] triggered by Bot. Commit: |
|
PR_Github #76196 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #76205 [ run ] triggered by Bot. Commit: |
|
PR_Github #76205 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #76220 [ run ] triggered by Bot. Commit: |
|
PR_Github #76220 [ run ] completed with state |
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp on behalf of runtime devs, delegating proper review to @NVIDIA/trt-llm-kv-cache-manager-devs; please ping me if you think this is not accurate
pranav-nvidia
left a comment
pranav-nvidia
left a comment
There was a problem hiding this comment.
LGTM - The P128 pool geometry, block-table conversion, speculative metadata swapping, and updated callers are consistent. Local syntax and diff checks pass; the latest reported L0 run passed.
Please ensure the four-GPU NVFP4 Eagle3 accuracy tests are validated on CI. The files changed may not automatically trigger these test stages.
brb-nv
left a comment
brb-nv
left a comment
There was a problem hiding this comment.
Minor comment about tests.
farazkh80
left a comment
farazkh80
left a comment
There was a problem hiding this comment.
Approved from attention perspective, with a couple questions (non-blocking).
|
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. |
…view tests - MiniMaxM3DraftKVCacheView.copy_batch_block_offsets now asserts beam_width == 1 and one copy_idx entry per sequence, as KVCacheManagerV2.copy_batch_block_offsets does. The view writes one draft block-table row per copy_idx entry, so a mismatch would otherwise leave stale rows without an error. - Mark test_draft_kv_cache_view_accessor_respects_local_layers and test_draft_kv_cache_view_rejects_multiple_local_layers cpu_only: both build the manager with __new__ and either mock the view or fail its single-layer check before any tensor is created. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…anager "Not for every shape" read as if the draft shape itself might lack P128 kernels. It means other shapes in the trtllm-gen artifact are P32-only, which is why P128 is opted in per manager rather than globally. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #76274 [ run ] triggered by Bot. Commit: |
|
PR_Github #76274 [ run ] completed with state |
Description
The problem. MiniMax-M3 support was developed on the
feat/m3_with_msaside branch, and dozens of PRs have been ported to main. When aligning perf between the two, we found a regression on main with an NVFP4 KV cache: MiniMax-M3's shared Eagle3 draft layer still used a P32 sub-page path that the side branch had already replaced with native P128 (#18066). That path rebuilds the draft block table in Python for every request on every decode step (~0.5 ms per request), so high-concurrency decode becomes host-bound.What this PR does. It stores the Eagle3 draft layer at native P128 pages when the KV cache is NVFP4, and fills the draft block table with the same single C++ device copy the main KV cache uses. Removing the P32 path also simplifies the code: the PR is a net ~400-line reduction (+181 / −590).
Results on AgentX two high throughput pareto points (B300, InferenceX AgentX; x90 = 1 / P90 inter-token latency in tok/s/user, y = total tok/s/GPU):
With this PR, main recovers perf and beats the side branch (+3% / +2% at TP4 C40, +6% / +4% at TP2 C45).
Test Coverage
tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py: the GPU test builds the view on a real NVFP4hybrid manager and checks its pool pointer, slot-id source rows, the single device copy and the resulting page indices; the
P32-only tests are removed. Local B300 run of the M3 sparse unit tests: 149 passed (151 before; two P32 tests removed).
accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3in pre-mergel0_dgx_b200, which thisPR's files do not trigger automatically (needs
/bot run --add-multi-gpu-test).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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.🤖 Generated with Claude Code
Dev Engineer Review
The NVFP4 hybrid-cache path now exposes shared Eagle draft KV data through native FP8/P128 pages. It removes the P32 subpage view and scatter helper, and uses the device-copy path to populate the draft block table. The copy path asserts that the copy-index count matches the sequence count. Confirm that callers use the new view API and that no runtime path depends on the removed subpage writer.
Speculative-cache lookup now requests a draft KV view. Dense decode excludes hybrid draft layers when it calculates the uniform subpage factor. SWA scratch-copy preparation reads page-index converter values directly and rejects expansions other than 1.
QA Engineer Review
The shared-draft-layer unit test changed. It covers native P128 page-table conversion, pool pointers, K/V offsets, dense-decode sizing, view access, and rejected configurations. The unit-test file does not appear in the integration test lists. The MiniMax-M3
test_nvfp4_eagle3selectors are present in thel0_dgx_b200CI list; the QA list has FP8 selectors but no NVFP4 selector. Current test-run status is unavailable.Coverage verdict: needs follow-up. The supplied evidence describes broad unit-test coverage and CI selectors, but does not establish current test results.
Per-File QA Perspective
tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py— Verify native P128 view construction, block-table copying, and supported draft-layer configurations. The P32 subpage view is removed.tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py— The subpaged scatter helper is removed. Verify that no runtime path depends on it.tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py— Uniform subpage-factor selection now excludes hybrid draft layers. Verify dense decode with mixed cache roles.tensorrt_llm/_torch/models/modeling_minimaxm3.py— A comment clarifies the missing trtllm-gen cubin. No runtime behavior change is reported.tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py— SWA scratch-copy preparation now reads converter values directly and rejects expansions other than 1. Verify supported layouts and the rejection path.tensorrt_llm/_torch/speculative/interface.py— Draft-cache lookup now usesget_draft_kv_cache_view. Verify API consistency with the target cache manager.tensorrt_llm/_torch/speculative/utils.py— Draft-cache lookup requests the shared view when no explicit draft manager exists. Verify fallback behavior when the view is unavailable.tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py— Tests native P128 shared-draft behavior, view access, and rejected configurations. The file is not listed in the integration test lists.