perf(sm120): wave+residue tile-selection for plain groupwise MoE GEMM - #4318
Merged
Merged
Conversation
Plain (unfused) groupwise MoE GEMM tile-selection used bare per-expert-M
thresholds and dispatched tile_m=128 across the M-residue band where a
2-M-tile layout wastes much of the second tile, while tile_m=64 packs
better. Add select_{mxfp8,fp8}_plain_moe_tile_m sharing
select_plain_m64_or_m128 (long-K M64<->M128 chosen by
cost=ceil(tiles/num_sms)*(tile_m+overhead), tie->128; small-M and short-K
stay bit-identical to the previous thresholds), and drop the fp8 fused
m64_tiles<=num_sms*8 guard so the fused selector matches the wave-aware
reference. Correctness unchanged; benchmark shows net wins with worst-case
within noise.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
CarstyYou
requested review from
Aneureka,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
August 3, 2026 05:01
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughSM120 FP8 and MXFP8 MoE paths now support tuned tile tactics, workload-based tile selection, tactic validation, and autotuner integration. Untuned entry points remain available. Tests cover cache behavior, profiling, eligibility, and forced kernel execution. ChangesSM120 MoE tuning and dispatch
Estimated code review effort: 5 (Critical) | ~100 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…ributions
The plain groupwise MoE GEMM tile selector estimated the tile count as
num_experts * ceil(mean_m / tile_m) (ceil-of-mean), but the scheduler launches
sum_e ceil(m_e / tile_m). Because ceil is convex, ceil-of-mean underestimates
the tile count when the per-expert mean sits near a tile boundary, so the
selector over-picks tile_m=64. On SM120 (110 SM) this is a measured 3-7%
regression on realistic non-uniform routing (multinomial, topk=8,
num_experts=128/256), where the optimum flips from tile_m=64 (uniform) to
tile_m=128 (e.g. per-expert-mean 129/192).
Per-expert row counts (token_offset) are device-side and the host selector only
has the mean, so instead of a D2H sync this uses the expected per-expert tile
count under Poisson(mean) routing noise, normal-approximated:
E[ceil(X/t)] = sum_{j>=0} P(X > j*t)
~ sum_{j>=0} 0.5 * (1 + erf((mean - j*t) / sqrt(2*mean))).
Only boundary-straddling terms contribute; away from a boundary it reduces to
ceil-of-mean, so shapes not near a tile boundary are unchanged. Applied
symmetrically to the fp8 and mxfp8 plain runners; host per-dispatch cost only,
no kernel hot-path change. Validated on SM120 (0-spill, correctness unchanged,
forced/natural tile picks match the measured-faster tile).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Remove the algorithm-derivation comment added in the previous commit; keep the code unchanged. Rationale lives in the commit message / PR description instead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu`:
- Around line 165-166: Update the call to select_fp8_fused_moe_tile_m near the
out_n declaration to pass out_n as the logical N dimension instead of shape_n.
Leave the packed weight handling and other selector arguments unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 76ff16b1-9a6a-425e-9a82-2ba5e1616b6c
📒 Files selected for processing (3)
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_op.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu
🚧 Files skipped from review as they are similar to previous changes (1)
- csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu (1)
51-54: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy liftFused-MoE selection reuses a plain-path cost model with the wrong TileN.
select_fp8_fused_moe_tile_mnow delegates toselect_fp8_plain_moe_tile_m. That selector assumes TileN=128 for every candidate, becauseselect_plain_m64_or_m128and the SWAPAB gate both divideshape_nby 128.The fused kernels do not use those tile shapes.
fused_moe_fp8_nt_groupwise_implinstantiatesSM120BlockScalingFusedMoeBuilder<128, 64, 128, 2>, so the fused M128 kernel has TileN=64, not 128.The caller at Line 166 also passes packed
shape_n, whilelaunch_fused_moereceivesout_n = shape_n / 2as the logical N. The two deviations interact:
- For the M128 candidate,
ceil(2*out_n / 128)equals the true N-tile countceil(out_n / 64). The estimate is correct by coincidence.- For the M64, M32, and SWAPAB candidates, the real TileN is 128, so the true count is
ceil(out_n / 128). The selector computes twice that value.The result is a systematic 2x cost inflation for the smaller-M fused candidates, which biases the fused path toward M128 and toward skipping the SWAPAB branch. Pass the logical N and model the per-candidate TileN, or keep a separate fused selector.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu` around lines 51 - 54, Update select_fp8_fused_moe_tile_m instead of delegating to select_fp8_plain_moe_tile_m: pass logical N (out_n, not packed shape_n) and compute candidate costs using each fused kernel’s actual TileN, including TileN=64 for M128 and TileN=128 for M64, M32, and SWAPAB. Preserve the existing fused candidate-selection behavior while eliminating the 2x inflation for smaller-M candidates.
🧹 Nitpick comments (4)
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu (1)
86-96: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
select_plain_m64_or_m128is duplicated across the FP8 and MXFP8 runners.This function is byte-identical to
select_plain_m64_or_m128incsrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cuat Lines 56-66, including thekPlainTileOverhead = 48constant.Both copies are
static, so there is no linkage conflict. The risk is drift: a future retune of the overhead constant or the wave model must be applied in both files. Consider moving the helper into a shared header undercsrc/cute_sm120_mxfp8_groupwise/sm120_common/.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu` around lines 86 - 96, Move the duplicated select_plain_m64_or_m128 helper and its kPlainTileOverhead wave-cost logic into a shared header under sm120_common, then include and reuse that definition from both the FP8 and MXFP8 runners. Remove the separate static implementations while preserving the existing 64-versus-128 selection behavior.flashinfer/grouped_mm/cute_sm120_mxfp8_groupwise/core.py (1)
40-42: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse iterable unpacking instead of tuple concatenation.
Ruff reports RUF005 here. Unpacking avoids the intermediate tuple and matches the rule.
♻️ Proposed change
-_MXFP8_MOE_TACTICS_GRANK128 = _MXFP8_MOE_TACTICS + ( - (_MXFP8_MOE_TACTIC_SCHEMA_VERSION, 128, 128), -) +_MXFP8_MOE_TACTICS_GRANK128 = ( + *_MXFP8_MOE_TACTICS, + (_MXFP8_MOE_TACTIC_SCHEMA_VERSION, 128, 128), +)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/grouped_mm/cute_sm120_mxfp8_groupwise/core.py` around lines 40 - 42, Update the _MXFP8_MOE_TACTICS_GRANK128 definition to use iterable unpacking when adding the extra tactic entry instead of concatenating tuples, preserving the existing tactic contents and order.Source: Linters/SAST tools
tests/autotuner/test_cute_sm120_fp8_groupwise.py (1)
114-154: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a positive cache-hit assertion.
This test asserts only that invalid tactics miss. The assertions still pass if
search_cachealways returns a miss. Add one case that stores a valid tactic underkeyand asserts a hit with the same tactic. That proves the validation gate rejects only invalid entries.💚 Proposed addition
tuner.profiling_cache.clear() tuner._file_configs[key.file_key] = (runner.__class__.__name__, invalid) assert not tuner.search_cache( "cute_sm120_fp8_groupwise_moe", [runner], shapes, fp8_core._FP8_MOE_TUNING_CONFIG, inputs, )[0] + + tuner.clear_cache() + valid = fp8_core._FP8_MOE_PLAIN_TACTICS[0] + tuner.profiling_cache[key] = (valid, None) + hit, _, cached_tactic, _ = tuner.search_cache( + "cute_sm120_fp8_groupwise_moe", + [runner], + shapes, + fp8_core._FP8_MOE_TUNING_CONFIG, + inputs, + ) + assert hit + assert cached_tactic == valid🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/autotuner/test_cute_sm120_fp8_groupwise.py` around lines 114 - 154, Add a positive cache-hit case to test_invalid_memory_and_loaded_tactics_are_cache_misses: after the invalid-entry assertions, store a valid tactic under the same cache key and assert search_cache returns a hit with that tactic. Reuse the runner/configuration context and obtain a tactic known to pass validation, ensuring the test proves valid entries are accepted while invalid entries are rejected.tests/grouped_mm/test_cute_sm120_mxfp8.py (1)
440-450: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConfirm the parametrization matrix runtime is acceptable.
This matrix expands to 2 gated modes × 11
(k_gran, tactic)pairs × 10 shape combinations, which is 220 GPU cases in one test. Each case builds reference tensors with Python-eager per-expert loops. Confirm the SM120 CI job tolerates the added wall-clock time, or reduce the shape list for the less informative tactics.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/grouped_mm/test_cute_sm120_mxfp8.py` around lines 440 - 450, Reduce the parametrization matrix used by the test to keep SM120 CI runtime acceptable, prioritizing coverage from the most informative tactics and retaining the gated-mode and key shape cases. Update the `k_gran,tactic` and/or `expert_rows,physical_n,k` parametrization near `test_cute_sm120_mxfp8` without changing the test’s functional assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu`:
- Around line 242-249: Ensure both tuned dispatch implementations,
fused_moe_mxfp8_nt_groupwise_tuned_impl at
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu:242-249 and
moe_gemm_mxfp8_nt_groupwise_tuned_impl at
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu:465-472, terminate
their GranK dispatch chains with an else that raises TVM_FFI_ICHECK(false) and
reports the tactic and GranK values when no supported tactic matches.
In `@tests/grouped_mm/test_cute_sm120_mxfp8.py`:
- Around line 510-518: Update the output initialization in the test around
_CuteSm120Mxfp8MoeRunner to use torch.zeros_like(ref) instead of
torch.empty_like(ref), ensuring unwritten rows produce deterministic
cosine-similarity results.
---
Outside diff comments:
In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cu`:
- Around line 51-54: Update select_fp8_fused_moe_tile_m instead of delegating to
select_fp8_plain_moe_tile_m: pass logical N (out_n, not packed shape_n) and
compute candidate costs using each fused kernel’s actual TileN, including
TileN=64 for M128 and TileN=128 for M64, M32, and SWAPAB. Preserve the existing
fused candidate-selection behavior while eliminating the 2x inflation for
smaller-M candidates.
---
Nitpick comments:
In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu`:
- Around line 86-96: Move the duplicated select_plain_m64_or_m128 helper and its
kPlainTileOverhead wave-cost logic into a shared header under sm120_common, then
include and reuse that definition from both the FP8 and MXFP8 runners. Remove
the separate static implementations while preserving the existing 64-versus-128
selection behavior.
In `@flashinfer/grouped_mm/cute_sm120_mxfp8_groupwise/core.py`:
- Around line 40-42: Update the _MXFP8_MOE_TACTICS_GRANK128 definition to use
iterable unpacking when adding the extra tactic entry instead of concatenating
tuples, preserving the existing tactic contents and order.
In `@tests/autotuner/test_cute_sm120_fp8_groupwise.py`:
- Around line 114-154: Add a positive cache-hit case to
test_invalid_memory_and_loaded_tactics_are_cache_misses: after the invalid-entry
assertions, store a valid tactic under the same cache key and assert
search_cache returns a hit with that tactic. Reuse the runner/configuration
context and obtain a tactic known to pass validation, ensuring the test proves
valid entries are accepted while invalid entries are rejected.
In `@tests/grouped_mm/test_cute_sm120_mxfp8.py`:
- Around line 440-450: Reduce the parametrization matrix used by the test to
keep SM120 CI runtime acceptable, prioritizing coverage from the most
informative tactics and retaining the gated-mode and key shape cases. Update the
`k_gran,tactic` and/or `expert_rows,physical_n,k` parametrization near
`test_cute_sm120_mxfp8` without changing the test’s functional assertions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: cbc85946-9f41-4ccf-8c4b-8bc95d1c433f
📒 Files selected for processing (20)
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_op.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_op_jit_binding.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_fp8_runner.hcsrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_op.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_op_jit_binding.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cucsrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.hcsrc/cute_sm120_mxfp8_groupwise/sm120_common/epilogue.cuhcsrc/cute_sm120_mxfp8_groupwise/sm120_fused_moe/fp8_builder.cuhcsrc/cute_sm120_mxfp8_groupwise/sm120_fused_moe/fp8_kernel_impl.cuhflashinfer/autotuner/autotuner.pyflashinfer/grouped_mm/_sm120_moe_autotune.pyflashinfer/grouped_mm/cute_sm120_fp8_groupwise/core.pyflashinfer/grouped_mm/cute_sm120_mxfp8_groupwise/core.pytests/autotuner/test_autotuner_core.pytests/autotuner/test_cute_sm120_fp8_groupwise.pytests/autotuner/test_cute_sm120_mxfp8_groupwise.pytests/grouped_mm/test_cute_sm120_fp8.pytests/grouped_mm/test_cute_sm120_mxfp8.py
Comment on lines
+242
to
+249
| } else if constexpr (GranK == 128) { | ||
| using KT_M128_N128 = | ||
| sm120_blockscaled::SM120BlockScaledFusedMoeBuilder<128, 128, 64, 4, GranK>; | ||
| sm120_blockscaled::launch_fused_moe<KT_M128_N128>(ptr_A, ptr_B, ptr_SFA, ptr_SFB, ptr_D, | ||
| total_rows, shape_n, shape_k, num_experts, | ||
| token_offset, num_sms, stream); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Both MXFP8 tuned dispatch chains can fall through without launching a kernel. Each chain terminates in else if constexpr (GranK == 128) with no final else. When GranK == 32, that branch is discarded and an unmatched tactic launches nothing, leaving the output buffer D untouched with no error. The (128, 128) guard in moe_gemm_mxfp8_nt_groupwise_tuned at Lines 342-343 makes both paths unreachable today, but neither dispatch carries its own guard.
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu#L242-L249: add a terminatingelsetofused_moe_mxfp8_nt_groupwise_tuned_implthat raisesTVM_FFI_ICHECK(false)with the tactic andGranKvalues.csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu#L465-L472: add the same terminatingelsetomoe_gemm_mxfp8_nt_groupwise_tuned_impl.
📍 Affects 1 file
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu#L242-L249(this comment)csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu#L465-L472
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu` around lines 242
- 249, Ensure both tuned dispatch implementations,
fused_moe_mxfp8_nt_groupwise_tuned_impl at
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu:242-249 and
moe_gemm_mxfp8_nt_groupwise_tuned_impl at
csrc/cute_sm120_mxfp8_groupwise/cute_sm120_mxfp8_runner.cu:465-472, terminate
their GranK dispatch chains with an else that raises TVM_FFI_ICHECK(false) and
reports the tactic and GranK values when no supported tactic matches.
Comment on lines
+510
to
+518
| out = torch.empty_like(ref) | ||
| runner = _CuteSm120Mxfp8MoeRunner(out, is_gated, (1, 1, k_gran), "MN") | ||
| runner([a_fp8, b_fp8, a_sf, b_sf, m_indptr], tactic=tactic) | ||
| cos_sim = F.cosine_similarity( | ||
| out.reshape(-1).float(), ref.reshape(-1).float(), dim=0 | ||
| ).item() | ||
| assert cos_sim > COS_SIM_THRESHOLD, ( | ||
| f"forced tactic cos_sim={cos_sim:.4f} < {COS_SIM_THRESHOLD}" | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Initialize out deterministically.
torch.empty_like(ref) leaves out uninitialized. If a tactic fails to write some rows, the comparison reads stale memory and the cosine similarity becomes non-deterministic. Use torch.zeros_like(ref) so a partial write produces a stable, reproducible failure.
💚 Proposed fix
- out = torch.empty_like(ref)
+ out = torch.zeros_like(ref)
runner = _CuteSm120Mxfp8MoeRunner(out, is_gated, (1, 1, k_gran), "MN")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| out = torch.empty_like(ref) | |
| runner = _CuteSm120Mxfp8MoeRunner(out, is_gated, (1, 1, k_gran), "MN") | |
| runner([a_fp8, b_fp8, a_sf, b_sf, m_indptr], tactic=tactic) | |
| cos_sim = F.cosine_similarity( | |
| out.reshape(-1).float(), ref.reshape(-1).float(), dim=0 | |
| ).item() | |
| assert cos_sim > COS_SIM_THRESHOLD, ( | |
| f"forced tactic cos_sim={cos_sim:.4f} < {COS_SIM_THRESHOLD}" | |
| ) | |
| out = torch.zeros_like(ref) | |
| runner = _CuteSm120Mxfp8MoeRunner(out, is_gated, (1, 1, k_gran), "MN") | |
| runner([a_fp8, b_fp8, a_sf, b_sf, m_indptr], tactic=tactic) | |
| cos_sim = F.cosine_similarity( | |
| out.reshape(-1).float(), ref.reshape(-1).float(), dim=0 | |
| ).item() | |
| assert cos_sim > COS_SIM_THRESHOLD, ( | |
| f"forced tactic cos_sim={cos_sim:.4f} < {COS_SIM_THRESHOLD}" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/grouped_mm/test_cute_sm120_mxfp8.py` around lines 510 - 518, Update the
output initialization in the test around _CuteSm120Mxfp8MoeRunner to use
torch.zeros_like(ref) instead of torch.empty_like(ref), ensuring unwritten rows
produce deterministic cosine-similarity results.
CarstyYou
force-pushed
the
sm120_fused_moe
branch
2 times, most recently
from
August 8, 2026 02:32
10914f2 to
a53e94b
Compare
jiahanc
reviewed
Aug 10, 2026
Collaborator
There was a problem hiding this comment.
we dont need helper tests files under autotuner
Collaborator
|
/bot run tests/grouped_mm |
Collaborator
Contributor
Author
|
Hi @jiahanc, I pushed e5a6d68 to address the selector P1s:
Could you please take another look when you have a chance? |
Collaborator
|
[SUCCESS] Pipeline #61895071: 18/18 executed test jobs passed |
1 task
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…flashinfer-ai#4318) ## Summary The plain (unfused) SM120 groupwise MoE GEMM tile-selection used bare per-expert-M thresholds and dispatched `tile_m=128` across the per-expert-M residue band, where a 2-M-tile layout wastes much of the second tile while `tile_m=64` packs better. This PR replaces the plain-path selection with a wave+residue cost model and aligns the fp8 fused selector with the wave-aware reference. ## Changes - Add `select_{mxfp8,fp8}_plain_moe_tile_m` sharing `select_plain_m64_or_m128`. For long-K, M64↔M128 is chosen by `cost = ceil(tiles / num_sms) * (tile_m + overhead)`, tie → 128. Small-M and short-K dispatch stay bit-identical to the previous thresholds. - Drop the `m64_tiles <= num_sms * 8` guard in `select_fp8_fused_moe_tile_m` that suppressed `tile_m=64` for large `num_experts` (E ≥ 64), so the fused selector matches the wave-aware reference. - Runtime `num_sms`-adaptive; validated at `num_sms = 110`. ## Performance Measured on 5K Pro (SM120 Blackwell, `num_sms = 110`): a forced tile_m sweep plus representative end-to-end grouped-MoE shapes. Speedup = baseline / candidate − 1. The previous bare-threshold selector was systematically suboptimal in the per-expert-M residue band (~129–207) and for some deep-K per-expert-M 72–80, leaving up to +38% (mxfp8 plain) / +17% (fp8 fused, E ≥ 64) / +5% (fp8 plain) speedup over its choice vs the measured-optimal tile_m at those points. The new selector eliminates all of the large residue-band regressions and is faster on ~100 of the swept shapes. Worst case is −4.7% (a slowdown) in a small-N / per-expert-M ≈ 72 occupancy corner where the tile_m choice is a noise-level near-tie (±1–3%, no clean threshold). Representative end-to-end speedups: | case | before | after | speedup | | ------------------------------------- | ------- | ------- | ------- | | mxfp8 grouped MoE, per-expert-M ≈ 192 | 1411 µs | 1292 µs | +9.2% | | fp8 fused MoE, per-expert-M ≈ 192 | 1507 µs | 1297 µs | +16.2% | (per-expert-M = 256 unchanged — already optimal.) ## Correctness Unchanged: mxfp8 and fp8 grouped-GEMM tests pass. Small-M and short-K dispatch is bit-identical to the previous thresholds, so only the long-K M64↔M128 boundary moves. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance Improvements** * Improved plain and fused MoE workload scheduling for uneven expert workloads. * Enhanced tile-size selection across expert sizes, shapes, and workload sizes. * Added autotuning support for SM120 FP8 and MXFP8 groupwise MoE operations. * Refined dispatch decisions to improve compute utilization and performance. * **Bug Fixes** * Invalid or outdated tuning choices are now rejected and replaced through fallback profiling. * Improved handling of gated and ungated workloads, including non-aligned dimensions and empty experts. * **New Features** * Added tuned MoE execution options for selecting tile dimensions. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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
3 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
The plain (unfused) SM120 groupwise MoE GEMM tile-selection used bare
per-expert-M thresholds and dispatched
tile_m=128across the per-expert-Mresidue band, where a 2-M-tile layout wastes much of the second tile while
tile_m=64packs better. This PR replaces the plain-path selection with awave+residue cost model and aligns the fp8 fused selector with the wave-aware
reference.
Changes
select_{mxfp8,fp8}_plain_moe_tile_msharingselect_plain_m64_or_m128.For long-K, M64↔M128 is chosen by
cost = ceil(tiles / num_sms) * (tile_m + overhead), tie → 128. Small-M andshort-K dispatch stay bit-identical to the previous thresholds.
m64_tiles <= num_sms * 8guard inselect_fp8_fused_moe_tile_mthat suppressed
tile_m=64for largenum_experts(E ≥ 64), so the fusedselector matches the wave-aware reference.
num_sms-adaptive; validated atnum_sms = 110.Performance
Measured on 5K Pro (SM120 Blackwell,
num_sms = 110): a forcedtile_m sweep plus representative end-to-end grouped-MoE shapes.
Speedup = baseline / candidate − 1.
The previous bare-threshold selector was systematically suboptimal in the
per-expert-M residue band (~129–207) and for some deep-K per-expert-M 72–80,
leaving up to +38% (mxfp8 plain) / +17% (fp8 fused, E ≥ 64) / +5% (fp8 plain)
speedup over its choice vs the measured-optimal tile_m at those points.
The new selector eliminates all of the large residue-band regressions and is
faster on ~100 of the swept shapes. Worst case is −4.7% (a slowdown) in a
small-N / per-expert-M ≈ 72 occupancy corner where the tile_m choice is a
noise-level near-tie (±1–3%, no clean threshold).
Representative end-to-end speedups:
(per-expert-M = 256 unchanged — already optimal.)
Correctness
Unchanged: mxfp8 and fp8 grouped-GEMM tests pass. Small-M and short-K dispatch
is bit-identical to the previous thresholds, so only the long-K M64↔M128
boundary moves.
Summary by CodeRabbit