[SM90] Optimize FP8 MegaMoE across L2 epilogue, 4WG launch heuristic and decode swap-AB - #69
Open
Helicoplacoidea wants to merge 8 commits into
Open
[SM90] Optimize FP8 MegaMoE across L2 epilogue, 4WG launch heuristic and decode swap-AB#69Helicoplacoidea wants to merge 8 commits into
Helicoplacoidea wants to merge 8 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>
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.
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 4WG (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 2WG (only launch 4WG when>512tokens/expert); on H200 keep the original>=64threshold and stay 4WG-first.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 sees ~2%/8% in the decode range (t16/t64). t256 sits on the threshold-switch boundary; t2048 still gains ~10% (Pro), while t4096+ is essentially flat near the bandwidth/compute ceiling.
Accuracy
All three changes are performance-only and do not change numerical results; accuracy is consistent with the dev branch:
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.