[None][perf] Fast-path the MiniMax-M3 MSA plan_kernel - #17097
Closed
zheyuf wants to merge 3 commits into
Closed
Conversation
zheyuf
force-pushed
the
perf/m3-planner-greedy-fastpaths
branch
from
July 31, 2026 02:59
ea9dfef to
74be5ed
Compare
The fmha_sm100 plan_kernel builds three attention work plans per engine step on a single CTA (1 of 148 SMs) with zero overlap, costing 3.2 ms of every conc256 decode step and 5.9 ms of every mixed step. Three exact fast paths in direct_greedy remove 93-96% of that without changing one bit of plan output: 1. Uniform cost: with equal tile costs the greedy reduces to round-robin (task t -> bucket t%N), so every output becomes closed-form. Covers the sparse/gqa decode plan (kv extent pinned to topk*page). 94x. 2. Single head: with num_heads==1 each round needs only the cheapest bucket; min(pack_cost_idx(cost,idx)) via the existing warp-shuffle reduction replaces the O(num_buckets) per-thread rank scan, and memoising pass 1's placements turns pass 2 into a parallel scatter. Covers the proxy and dense decode plans. 11x. 3. Row compaction: mixed steps visit max_qo_tiles x batch (qt,b) pairs twice while ~96% are inactive; compacting active rows once via cub::BlockScan and walking the list removes the dead iteration space. Covers prefill/mixed steps. 2.8x. Each path is guarded and falls back to the original loop unchanged. Applied as a downstream MSA patch; apply_msa_patch now applies every 3rdparty/patches/msa_*.patch in sorted order instead of one hardcoded file. Measured on 4x B300 TP4/EP1, InferenceMAX lengths (noise floors 0.05-0.12%, all deltas with non-overlapping ranges): - decode fast paths: +2.85% e2e at conc8, +3.30% at conc256 - row compaction: +1.14% further at conc256 (+4.42% cumulative) - plan output bit-identical to stock on 386 constructed cases including MAX_UNSPLIT truncation, ki==0 rows, bucket-count boundaries, and Eagle3 spec-dec shapes with variable acceptance - Eagle3 InferenceMAX GSM8K: 95.451 vs 95.490 reference (band 95.15-96.06), acceptance rate 0.838, mean acceptance length 3.513 Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
force-pushed
the
perf/m3-planner-greedy-fastpaths
branch
from
July 31, 2026 06:04
74be5ed to
de0c1e0
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63001 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63001 [ run ] completed with state |
Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
zheyuf
marked this pull request as ready for review
August 3, 2026 08:37
zheyuf
requested review from
BowenFu,
QiJune,
dpitman-nvda,
kris1025,
yiqingy0 and
yuxianq
and removed request for
a team,
BowenFu,
QiJune,
dpitman-nvda,
kris1025,
yiqingy0 and
yuxianq
August 3, 2026 08:37
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63400 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63400 [ run ] completed with state |
zheyuf
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR provides three fast path to planner. Under conc=4096, it is able to save~93% overhead in planner and have ~4.9% gain in tok/s/user while maintaining tok/s/GPU. It also shows gains for conc=64, 256, 1024.
There will be less gain after #17171 get merged. This is because
Fast path 1 — uniform cost ⇒ closed form (sparse/GQA decode, 94×)
When every task has the same cost, the greedy assignment is deterministic round-robin: task t → SM (t % N). We replace both greedy passes with parallel closed-form calculations. A min == max reduction guards the path; non-uniform cases use the original algorithm.
Gain: 1080 → 11.7 µs at batch 256 (94×; 113× at batch 2048).
Fast path 2 — single head ⇒ warp argmin + cached scatter (dense/proxy decode, 11×)
With num_heads == 1, each round assigns only one task, so the full per-thread rank scan is unnecessary. We use the existing warp-shuffle argmin, preserving the lowest-index tie-break, and cache pass-1 placements so pass 2 becomes a parallel scatter.
Gain: 1075 → 101 µs at batch 256 (11×).
Fast path 3 — active-row compaction (prefill/mixed, 2.8×)
Mixed steps scan the full max_qo_tiles × batch rectangle twice, although most entries are inactive—96% at concurrency 256. We compact active rows once with cub::BlockScan, preserving stock order, then run both passes over the dense list. Oversized cases fall back to the original path.
Gain: 2939 → 1062 µs per mixed call (2.8×).
Delivered as a downstream MSA patch. apply_msa_patch now applies all msa_*.patch files in sorted order using the existing patch mechanism.