[SM90] Optimize FP8 MegaMoE across L2 epilogue, 4WG launch heuristic and decode swap-AB - #69
Open
Helicoplacoidea wants to merge 9 commits into
Open
[SM90] Optimize FP8 MegaMoE across L2 epilogue, 4WG launch heuristic and decode swap-AB#69Helicoplacoidea wants to merge 9 commits into
Helicoplacoidea wants to merge 9 commits into
Conversation
…params; keep original 4-WG path on H200
…l lambda Replace the inlined swizzle computation in the swapAB STS write path with calls to the existing swiz_col(row, col) lambda, matching the non-swap STS write and LDS scatter read paths. Removes the swiz_col_0/swiz_col_1 locals that collided in naming with the lambda. Computation is bitwise-identical; no numerical or behavior change. Co-Authored-By: Claude <noreply@anthropic.com>
…20 flash On H20 + flash (intermediate_hidden <= 2048), 2-WG + BLOCK_N=256 beats 4-WG across part of the flash batch range. Raise the 4-WG launch bar to expected_tokens_per_expert > 576 for that regime (pro stays > 512, other cards and the H20 mid-range keep >= 64). Also drop the intermediate_hidden >= 3072 gate on decode_use_block_n_256 so the 2-WG split-N path picks BLOCK_N=256 for flash too (was forced to 128). Co-Authored-By: Claude <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
1 participant
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.
[SM90] Optimize FP8 MegaMoE decode path with swap-AB, 4WG heuristic and L2 bank swizzle
Summary
Relative to the dev branch, this PR applies three groups of performance optimizations to the SM90 FP8 MegaMoE kernel, covering the L2 epilogue, the 4-WG (split-MN) launch heuristic, and the decode swap-AB path.
Main changes:
intermediate_hiddenand GPU model (SM count). On H20 + pro (intermediate_hidden >= 3072), where the smaller SM count makes register spilling costly, prefer 2-WG (only launch 4-WG when>512tokens/expert); on H20 + flash (intermediate_hidden <= 2048), where 2-WG +BLOCK_N=256beats 4-WG across part of the batch range, prefer 2-WG (4-WG only when>576); on H200/H100 and the H20 mid-range keep the original>=64threshold and stay 4-WG-first. TheBLOCK_N=256selection for the 2-WG split-N path is also un-gated fromintermediate_hidden >= 3072, so flash picksn=256instead of being forced ton=128.BLOCK_N=128toBLOCK_N=256(2-WG N-split, each WG driving twom64n64slabs), and fix the L1-output TMA descriptor to cover the full tile rather than degrading towg_l1_out_block_n = 64.Performance
Notes:
speedup = latency(baseline) / latency(this PR).Flash shape (256 experts)
Pro shape (384 experts)
Pro sees a ~7–18% latency reduction across t16–t1024 (1.22× peak at t512). Flash also gains across the prefill range: t256 1.09× (BLOCK_N=256 un-gated, was 1.00×), t1024 1.05× and t2048 1.04× (4-WG → 2-WG + n256, t2048 flipping from 0.99× to a gain); t512 dips to 0.98× on the 2-WG/4-WG boundary, where 4-WG still wins — consistent with the 2-WG + n256 advantage being shape-dependent. The decode range keeps ~1.01×/1.08× (t16/t64). t4096+ stays 4-WG and is essentially flat near the FP8 compute ceiling (251 TFLOPS at t8192, ~85% of H20 FP8 peak).
Accuracy
All three changes are performance-only and do not change numerical results; accuracy is consistent with the dev branch:
m64n128WGMMA halves sharing the A load) vs 4-WG, both tiling the same C block. Each shape independently accumulates the full-tile GEMM, so the result is mathematically equivalent.BLOCK_N=256: performs the same GEMM with FP32 accumulation and unchanged K-order, splitting the originalm64n128into twom64n64slabs.Single-operator and model-level accuracy validations match the dev branch, with no accuracy regression.