[TRTLLM-14093][feat] Eagle3 support for MiniMax-M3 - #18872
Conversation
zheyuf
changed the title
0329bf9 to
01add1a
Compare
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #72808 [ run ] triggered by Bot. Commit: |
|
PR_Github #72808 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #72969 [ run ] triggered by Bot. Commit: |
|
PR_Github #72969 [ run ] completed with state
|
Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Resolve the conflicts with NVIDIA#18611, which moved the MSA generation phase off fmha_sm100 onto MiniMax-M3's dedicated decode kernels (CuTe DSL indexer scorer, Triton sparse decode, trtllm-gen dense decode). Main's structure is kept and the Eagle3 pieces are re-targeted at it: - cache_manager.py: both sides had extracted _kv_slot_geometry; only the docstring wording differed. - msa_backend.py: the decode span accepts the uniform 1 + draft_len query length and the CuTe DSL scorer is validated up front for it. on_update_kv_lens now patches msa_seq_lens_cuda, which all three decode kernels read, plus the write slots and the per-token valid-block counts; the fmha_sm100 plan patching is dropped since the plans cover context rows only. The write slots are re-derived from the staged msa_block_table, so the per-token msa_req_to_token buffer and its per-step H2D copy go away. The num_extra_kv_tokens correction moves into _stage_host_lengths so every staged length, msa_seq_lens_cuda included, is the attended one. The pure-decode valid-block staging is per query token, and the proxy scratch keeps its token-based sizing. - test_msa_backend.py: keep both sides and rewrite the on_update_kv_lens tests against the new buffers. Fix the msa_utils import paths moved under kernels/ here and in test_llm_api_pytorch.py. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
f9f7e97 to
7dec97c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #73186 [ run ] triggered by Bot. Commit: |
|
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. |
… tables A drafter with a different K/V page stride puts KVCacheManagerV2 into per-layer page tables, where every layer already has its own attention-op pool and pool ids are layer ids. Appending pools at num_pools + i there aliased real layers' pools and indexed the host block-offset table with a layer id. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Only tree acceptance relocates accepted draft KV, and neither the MSA decode kernels nor M3's coalesced pool layout support it. Reject tree configs at startup and any relocation request in the manager. Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #73236 [ run ] triggered by Bot. Commit: |
|
PR_Github #73186 [ run ] completed with state |
|
PR_Github #73236 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #73300 [ run ] triggered by Bot. Commit: |
|
PR_Github #73300 [ run ] completed with state |
13 participants
Supersedes #16021
#16021 was the initial eagle support from Minimax-m3 side branch based on July design but many fixes are added on the side branch afterwards (#17341, #17457, #17656, #18066), so we decided to not merge #16021, which is pretty stale.
This current PR (#18872) contains the all eagle related PRs/changes in Minimax-m3 side branch. Also, I tried to minimize the changes in hot path that is not specifically for Minimax-m3 to make this PR easier to merge. Now there are only 8 lines changes in flashinfer_trtllm_gen.py and 4 lines changes in py_executor_creator.py in hot path, while all others should be contained in Minimax-m3 specific file/path.
What this PR mainly contains
Minimax-m3 Eagle performance on AgentX (tested on Minimax-m3 side branch)
It's a controlled ablation on one of Minimax-m3 AgentX submission stack: SemiAnalysisAI/InferenceX#2658, not chronological.
Grey curve is without spec dec. Red curve shows the performance with eagle.
PR Checklist
[JIRA/NVBUG/None][type]format/bot run)Dev Engineer Review
PR_Github#72108and `PR_Github `#72242, with relatedL0_MergeRequest_PRpipelines, failed. Review failures before merge.QA Engineer Review
TestMiniMaxM3::test_nvfp4_eagle3coverage for TP4/EP4, attention DP, disaggregated serving, and overlap scheduling.tests/integration/test_lists/qa/llm_function_core.txtandtests/integration/test_lists/test-db/l0_dgx_b200.yml. Both lists retain the Eagle3 selectors and 180-second timeouts.Per-File QA Perspective
docs/source/models/supported-models.md: Marks MiniMax-M3 EAGLE-3 Linear support as available. Verify the matrix matches runtime support.tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py: Removes configurable MLA backend selection and CuTeDSL execution. Verify TRTLLM-Gen validation and cache-manager token sizes.tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py: Rejects NVFP4 shared draft-layer pools and SWA scratch reuse. Verify pointers, offsets, layouts, and error messages.tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py: Enables multi-token speculative decode and device-side metadata correction. Verify graph buffers, slots, block counts, plan lengths, and scratch limits.tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py: Adds device-side cache-slot derivation and KV-length updates. Verify prefill, decode, persistent buffers, and multi-token routing.tensorrt_llm/_torch/models/modeling_minimaxm3.py: Restricts dense decode to one query token and changes output handling. Verify masks, output layout, and single-query assumptions.tensorrt_llm/_torch/pyexecutor/py_executor_creator.py: Changes draft-cache behavior and rejects Triton CUDA graphs for MiniMax-M3 Eagle3. Verify configuration errors and unaffected backends.tests/integration/defs/accuracy/references/gsm8k.yaml: Adds the MiniMax-M3 Eagle3 GSM8K reference. Verify selectors and expected accuracy.tests/integration/defs/accuracy/references/mmlu.yaml: Adds the MiniMax-M3 Eagle3 MMLU reference. Verify selectors and expected accuracy.tests/integration/defs/accuracy/test_disaggregated_serving.py: Addsrouter_urltoDuckLLM. Verify resolved URLs and default compatibility.tests/integration/defs/accuracy/test_llm_api_pytorch.py: Adds Eagle3 tests for local and disaggregated configurations. Verify acceptance metrics, overlap scheduling, attention-DP variants, and CUDA-graph settings.tests/integration/test_lists/qa/llm_function_core.txt: Updates two MiniMax-M3 Eagle3 manual-QA entries. Verify selectors match the test parameterization.tests/integration/test_lists/test-db/l0_dgx_b200.yml: Updates the same two MiniMax-M3 Eagle3 CI entries. Verify IDs, paths, selectors, and timeouts.tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py: Tests virtual pool creation, block copying, layout detection, fallback copying, and draft-layer counting. It is not listed in an integration test list.tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py: Tests token-based scratch sizing, multi-token block counts, KV-length correction, plan patching, and inactive-speculation behavior. It is not listed in an integration test list.tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py: Initializesspec_configfor manually constructed models. Verify checkpoint loading compatibility. It is not listed in an integration test list.