Repository navigation
[None][feat] MiniMax-M3 NVFP4 low latency decode megakernel for Rubin (sm_107) TP2 / TP4 / TP8 in modeling_v2 - #20623
Conversation
|
/bot run --disable-fail-fast |
cd5f671 to
5d24e81
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79103 [ run ] triggered by Bot. Commit: |
|
PR_Github #79103 [ run ] completed with state
|
5d24e81 to
4d16da2
Compare
|
/bot run --disable-fail-fast |
4d16da2 to
7333ac8
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #79291 [ run ] triggered by Bot. Commit: |
|
PR_Github #79291 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79385 [ run ] triggered by Bot. Commit: |
|
PR_Github #79385 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #79616 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
Brings in NVIDIA#20623 at ffa76ab (the NVIDIA#20306 head 300c1c1 merged in, sm107 target kernels under ops). No change to the context path. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
PR_Github #79695 [ run ] triggered by Bot. Commit: |
|
PR_Github #79616 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #82111 [ run ] triggered by Bot. Commit: |
|
PR_Github #82111 [ run ] completed with state
|
|
/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 #82335 [ run ] triggered by Bot. Commit: |
|
PR_Github #82335 [ run ] completed with state
|
Retain upstream single KV-metadata invalidation after overlap corrections, including the model-specific fused offset path. Align the existing CPU regression with the upstream behavior for both legacy hook settings. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #82359 [ run ] triggered by Bot. Commit: |
|
PR_Github #82359 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #82417 [ run ] triggered by Bot. Commit: |
|
PR_Github #82417 [ run ] completed with state |
The sm_107 targets (NVIDIA#20623) issued their step's copies through step_inputs / h2d_cache with RuntimeHooks switches. They now declare a CopyPlan like the sm_103 targets. * CopyPlan.parallel_scatter: an order-free segment is scattered with one program per copy instead of the serial one-CTA scatter. * CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every decode graph is captured with a prologue that issues a pinned window, and the final flush before a replay only fills it (early_copy: uploaded on a side copy stream; prologue_op: the model's one-launch prologue; step_epoch: a device counter bumped once per replay, read through device_copy.step_epoch()). Barrier segments stay eager: only the step's last segment rides the window, and only when it fits and is order-free. No prologue is captured when the graph runner has pre-replay hooks; a replay outside the copier's context gets an empty window. * CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True): a changed merged 2-D upload stages only each row's changed span (MSA's block table on a page crossing). * M3SpecCopySite: the spec fold's arming uploads, owned by DeviceCopyBuffers and merged by the sm_107 plans together with the speculative RNG sites. * The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode shortcuts and the sampler's store copies are those components' own. * MSA's in-forward paths are its fixed behaviour: the sub-page table is always expanded inside the forward, the block table always keeps its owner-side record, and the out_cache_loc tail is always filled in the forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the eager tail fill are removed; nothing reads a merged destination before the replay, so no M3 plan needs a flush_after barrier. * The step prologue ops read device_copy's graph scatter as reference. * Ported tests: test_step_inputs' parallel-scatter and in-graph tests (test_device_copy_in_graph.py), test_h2d_changed_span (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph runner's prologue test, the MSA skip-routing scan, the sm_107 contract scans and step-prologue tests; test_copy_plans.py checks every M3 target's plan. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The sm_107 targets (NVIDIA#20623) issued their step's copies through step_inputs / h2d_cache with RuntimeHooks switches. They now declare a CopyPlan like the sm_103 targets. * CopyPlan.parallel_scatter: an order-free segment is scattered with one program per copy instead of the serial one-CTA scatter. * CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every decode graph is captured with a prologue that issues a pinned window, and the final flush before a replay only fills it (early_copy: uploaded on a side copy stream; prologue_op: the model's one-launch prologue; step_epoch: a device counter bumped once per replay, read through device_copy.step_epoch()). Barrier segments stay eager: only the step's last segment rides the window, and only when it fits and is order-free. No prologue is captured when the graph runner has pre-replay hooks; a replay outside the copier's context gets an empty window. * CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True): a changed merged 2-D upload stages only each row's changed span (MSA's block table on a page crossing). * M3SpecCopySite: the spec fold's arming uploads, owned by DeviceCopyBuffers and merged by the sm_107 plans together with the speculative RNG sites. * The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode shortcuts and the sampler's store copies are those components' own. * MSA's in-forward paths are its fixed behaviour: the sub-page table is always expanded inside the forward, the block table always keeps its owner-side record, and the out_cache_loc tail is always filled in the forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the eager tail fill are removed; nothing reads a merged destination before the replay, so no M3 plan needs a flush_after barrier. * The step prologue ops read device_copy's graph scatter as reference. * Ported tests: test_step_inputs' parallel-scatter and in-graph tests (test_device_copy_in_graph.py), test_h2d_changed_span (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph runner's prologue test, the MSA skip-routing scan, the sm_107 contract scans and step-prologue tests; test_copy_plans.py checks every M3 target's plan. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The sm_107 targets (NVIDIA#20623) issued their step's copies through step_inputs / h2d_cache with RuntimeHooks switches. They now declare a CopyPlan like the sm_103 targets. * CopyPlan.parallel_scatter: an order-free segment is scattered with one program per copy instead of the serial one-CTA scatter. * CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every decode graph is captured with a prologue that issues a pinned window, and the final flush before a replay only fills it (early_copy: uploaded on a side copy stream; prologue_op: the model's one-launch prologue; step_epoch: a device counter bumped once per replay, read through device_copy.step_epoch()). Barrier segments stay eager: only the step's last segment rides the window, and only when it fits and is order-free. No prologue is captured when the graph runner has pre-replay hooks; a replay outside the copier's context gets an empty window. * CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True): a changed merged 2-D upload stages only each row's changed span (MSA's block table on a page crossing). * M3SpecCopySite: the spec fold's arming uploads, owned by DeviceCopyBuffers and merged by the sm_107 plans together with the speculative RNG sites. * The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode shortcuts and the sampler's store copies are those components' own. * MSA's in-forward paths are its fixed behaviour: the sub-page table is always expanded inside the forward, the block table always keeps its owner-side record, and the out_cache_loc tail is always filled in the forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the eager tail fill are removed; nothing reads a merged destination before the replay, so no M3 plan needs a flush_after barrier. * The step prologue ops read device_copy's graph scatter as reference. * Ported tests: test_step_inputs' parallel-scatter and in-graph tests (test_device_copy_in_graph.py), test_h2d_changed_span (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph runner's prologue test, the MSA skip-routing scan, the sm_107 contract scans and step-prologue tests; test_copy_plans.py checks every M3 target's plan. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The sm_107 targets (NVIDIA#20623) issued their step's copies through step_inputs / h2d_cache with RuntimeHooks switches. They now declare a CopyPlan like the sm_103 targets. * CopyPlan.parallel_scatter: an order-free segment is scattered with one program per copy instead of the serial one-CTA scatter. * CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every decode graph is captured with a prologue that issues a pinned window, and the final flush before a replay only fills it (early_copy: uploaded on a side copy stream; prologue_op: the model's one-launch prologue; step_epoch: a device counter bumped once per replay, read through device_copy.step_epoch()). Barrier segments stay eager: only the step's last segment rides the window, and only when it fits and is order-free. No prologue is captured when the graph runner has pre-replay hooks; a replay outside the copier's context gets an empty window. * CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True): a changed merged 2-D upload stages only each row's changed span (MSA's block table on a page crossing). * M3SpecCopySite: the spec fold's arming uploads, owned by DeviceCopyBuffers and merged by the sm_107 plans together with the speculative RNG sites. * The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode shortcuts and the sampler's store copies are those components' own. * MSA's in-forward paths are its fixed behaviour: the sub-page table is always expanded inside the forward, the block table always keeps its owner-side record, and the out_cache_loc tail is always filled in the forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the eager tail fill are removed; nothing reads a merged destination before the replay, so no M3 plan needs a flush_after barrier. * The step prologue ops read device_copy's graph scatter as reference. * Ported tests: test_step_inputs' parallel-scatter and in-graph tests (test_device_copy_in_graph.py), test_h2d_changed_span (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph runner's prologue test, the MSA skip-routing scan, the sm_107 contract scans and step-prologue tests; test_copy_plans.py checks every M3 target's plan. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The sm_107 targets (NVIDIA#20623) issued their step's copies through step_inputs / h2d_cache with RuntimeHooks switches. They now declare a CopyPlan like the sm_103 targets. * CopyPlan.parallel_scatter: an order-free segment is scattered with one program per copy instead of the serial one-CTA scatter. * CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every decode graph is captured with a prologue that issues a pinned window, and the final flush before a replay only fills it (early_copy: uploaded on a side copy stream; prologue_op: the model's one-launch prologue; step_epoch: a device counter bumped once per replay, read through device_copy.step_epoch()). Barrier segments stay eager: only the step's last segment rides the window, and only when it fits and is order-free. No prologue is captured when the graph runner has pre-replay hooks; a replay outside the copier's context gets an empty window. * CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True): a changed merged 2-D upload stages only each row's changed span (MSA's block table on a page crossing). * M3SpecCopySite: the spec fold's arming uploads, owned by DeviceCopyBuffers and merged by the sm_107 plans together with the speculative RNG sites. * The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode shortcuts and the sampler's store copies are those components' own. * MSA's in-forward paths are its fixed behaviour: the sub-page table is always expanded inside the forward, the block table always keeps its owner-side record, and the out_cache_loc tail is always filled in the forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the eager tail fill are removed; nothing reads a merged destination before the replay, so no M3 plan needs a flush_after barrier. * The step prologue ops read device_copy's graph scatter as reference. * Ported tests: test_step_inputs' parallel-scatter and in-graph tests (test_device_copy_in_graph.py), test_h2d_changed_span (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph runner's prologue test, the MSA skip-routing scan, the sm_107 contract scans and step-prologue tests; test_copy_plans.py checks every M3 target's plan. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…ngine's environment overrides LLM writes env_overrides into the calling process's os.environ and does not restore them, so the modeling_v2 accuracy tests open each engine inside scoped_env_overrides, and test_modeling_v2_env_scope.py checks that every LLM(...) call under accuracy/modeling_v2 is built that way. TestModelingV2Glm52Nvfp4Sm103Tp8 passed env_overrides to an engine built outside such a scope, so TRTLLM_MODELING_V2_VALIDATE and TRTLLM_DSA_INDEXER_BF16 stayed set in the pytest process after the test, and the check fails on glm_moe_dsa/glm_5_2_nvfp4__sm_103__tp8/test_accuracy.py. The file-level waiver of this bug hid that failure; it was added when the check failed on test_minimax_m3.py, which NVIDIA#20623 has since scoped. Open the engine inside scoped_env_overrides and pass it the yielded mapping, as _env.py documents, and remove the waiver. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
The sm_107 targets (NVIDIA#20623) issued their step's copies through step_inputs / h2d_cache with RuntimeHooks switches. They now declare a CopyPlan like the sm_103 targets. * CopyPlan.parallel_scatter: an order-free segment is scattered with one program per copy instead of the serial one-CTA scatter. * CopyPlan.final_flush=InGraph(early_copy, prologue_op, step_epoch): every decode graph is captured with a prologue that issues a pinned window, and the final flush before a replay only fills it (early_copy: uploaded on a side copy stream; prologue_op: the model's one-launch prologue; step_epoch: a device counter bumped once per replay, read through device_copy.step_epoch()). Barrier segments stay eager: only the step's last segment rides the window, and only when it fits and is order-free. No prologue is captured when the graph runner has pre-replay hooks; a replay outside the copier's context gets an empty window. * CopyPlan.changed_span and DeviceCopyBuffer.update(..., changed_span=True): a changed merged 2-D upload stages only each row's changed span (MSA's block table on a page crossing). * M3SpecCopySite: the spec fold's arming uploads, owned by DeviceCopyBuffers and merged by the sm_107 plans together with the speculative RNG sites. * The sm_107 RuntimeHooks keep only step-loop choices; MSA's decode shortcuts and the sampler's store copies are those components' own. * MSA's in-forward paths are its fixed behaviour: the sub-page table is always expanded inside the forward, the block table always keeps its owner-side record, and the out_cache_loc tail is always filled in the forward. The SUBPAGE_IN_GRAPH, SKIP_UNCHANGED_BLOCK_TABLE and TAIL_FILL_IN_FORWARD switches, the prepare()-time expansion and the eager tail fill are removed; nothing reads a merged destination before the replay, so no M3 plan needs a flush_after barrier. * The step prologue ops read device_copy's graph scatter as reference. * Ported tests: test_step_inputs' parallel-scatter and in-graph tests (test_device_copy_in_graph.py), test_h2d_changed_span (test_device_copy_changed_span.py), test_h2d_cache's pool-shared tests (test_device_copy_pool_shared.py), test_rng_upload_pack, the graph runner's prologue test, the MSA skip-routing scan, the sm_107 contract scans and step-prologue tests; test_copy_plans.py checks every M3 target's plan. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
@coderabbitai summary
Description
Depends on #20306 (flat staircase layout of the sm_103 targets). This PR adds only the last commit; please review that commit.
This PR brings the MiniMax-M3 NVFP4 decode targets to Rubin (sm_107):
minimax_m3_nvfp4__sm_107__tp2,__tp4and__tp8. They follow the flat layout of the sm_103 targets: onemodeling.pyper target, aweights.pyweight table, and anops/catalog entry with certified cells and a test for every kernel the model calls. The router picks a target only for the matching checkpoint, SM and width, andmodeling_v2: requiremakes it mandatory.Performance
On Rubin (VR200), MiniMax-M3 NVFP4 AgentX with Eagle3, against the stock decode path (
modeling_v2off):sm_107 enablement
Arch.SM_107, and routing entries for (10, 7) at 2, 4 and 8 ranks, with the same drafter and runtime declines as sm_103.minimaxM3Sm107Tp{2,4,8}with theirtrtllm::minimax_m3_sm107*thop ops, plus two shared GEN kernels: the TF32 tensor-core router GEMV and the spec-tail draft-input step.RuntimeHooksturns it on, so other models are unchanged.WorkTileInforebind to its own compiles, so other CuTe DSL kernels (e.g. FA4) are unaffected.RuntimeHooks; the MiniMax-M3 sm_107 targets do, and other models are unchanged. The AllReduce tuner leavesNCCL_SYMMETRICout of its candidates only withTRTLLM_AR_TUNE_NCCL_SYMMETRIC=0(opt-in; the default keeps the old candidates).The targets
All three share one structure; they differ in rank geometry and row ranges.
m3_sparse_decode_lcm3_sparse_decode_lcm3_sparse_decode_lc8TRTLLM_MODELING_V2_VALUE_DUMP), as in the sm_103 targets. Every fast path is on by default and has one module constant that turns it off. Capacity bounds carry a# capacity:annotation.Aggregated (wide) decode at TP2
Numerics. No lossy steps. Most kernels are bitwise equal to the stock op they replace; the rest reduce in a fixed order, are deterministic in eager and graph mode, and stay within the stock kernel's error band. One known divergence: the TP4 drafter all-reduce + RMSNorm always sums like ONESHOT, so under a forced NCCL strategy its rounding differs from stock NCCL. Against the earlier wide TP2 reference, target tokens are identical: at 12+ concurrent requests under Eagle3 some draft tokens can resolve near-ties differently (target logits bitwise), and without speculation (not a production configuration) decode steps of 9..48 rows are bitwise on the first such step and differ after it, with identical output tokens. Steps of 1..8 rows are bitwise for prompts up to about 12K tokens, and at TP8 also for a single 83K-token request (prefill and decode).
Not in this PR: batch 4/8 decode at TP4 / TP8 (follow-up). The prefill (context) optimizations are in #20629.
Known behaviours: The pinned SM107 MXFP8 GEMM tactics are armed process-wide by an executor whose model opts in, so another model in the same process at the same GEMM shape also uses them; its outputs are bitwise unchanged (every pinned tactic equals the default), and the LLM API runs one executor per worker process.
Test Coverage
new_tokensin Eagle3 context steps. Re-run on the final commit at TP2 and TP4 (no speculation and Eagle3): bitwise at every step of 1..8 rows; at TP2 the first difference is a 45-row context step, which this PR moves to the wide cells.executor/,speculative/, attention): 255 pass, 1 skipped.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.