SM90 (Hopper) FP4 MegaMoE fused kernel with swapAB small-batch path#53
Open
qiushixiaoyu wants to merge 6 commits into
Open
SM90 (Hopper) FP4 MegaMoE fused kernel with swapAB small-batch path#53qiushixiaoyu wants to merge 6 commits into
qiushixiaoyu wants to merge 6 commits into
Conversation
Add the SM90 FP8xFP4 MegaMoE runtime, kernel path, Python API, Hopper correctness and benchmark coverage, tuned runtime decode heuristics, swapAB support, synchronization/spill fixes, and the SM90 MegaMoE alignment export.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The FP4 path had ~49 tuned conditions (17 shape x e-band API-default bands, 9+12 wave-table rules, 11 stage-cap rules) fitted point-by-point on benchmark batches; on real shapes most bands collapsed to constants or slivers. Retire them: - API defaults: constants + one decode/prefill split + a single swapAB threshold (measured crossover on H20: swapAB wins at e<=12, ties at e~16, loses beyond -> bound 16). wide_load_decode on for all decode. - Wave scheduling: delegate to the FP8 rule (get_num_experts_per_wave_for_mega_moe_sm90). - Stage cap: removed; always use as many pipeline stages as SMEM allows. - CPU heuristic mirrors in the test synced to the same closed forms. Verified on 8xH20 EP8: accuracy suites pass (diff <= 0.0006, tol 0.05); GLM5.2 flat-or-better (b2 -12%, b96 -9%, prefill bit-identical); DSV4-Flash within noise; DSV4-Pro avg ~+2%, worst b1 +8.6%. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The decode config (BLOCK_M=64) and prefill config (BLOCK_M=128) flipped at a hardcoded e>=64. Measured on 8xH20 both modes across the boundary: decode wins for e in [64, 80) (its first m-block is exactly full while prefill's 128-row block runs half empty, -13% at GLM5.2 b256); prefill wins from e=80 up (decode's second m-block is mostly empty, +2~9%). Parameterize the four boundary checks (block config auto_split_mn, split-N band, 2-WG thread band, API prefill band) through a single DG_SM90_FP4_PREFILL_E knob (default 80) plus a PREFILL_E mirror constant in the test. GLM5.2 b256: 1277 -> 1123 us; vs FP8 low-latency 0.98x -> 1.11x (cliff fixed, b245->b320 now monotone). DSV4-Flash e=72 also improves -1.8%. Accuracy: forced-decode reference checks at e=64/96/128 pass (diff ~ 6e-4, tol 0.1). Co-Authored-By: Claude Opus 4.8 <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.
This PR adds an FP4-weight MegaMoE fused kernel that:
tokens on N), which is more efficient when tokens-per-rank is small.
Changes
sm90_fp8_fp4_mega_moe_impl(
deep_gemm.fp8_fp4_mega_moe): FP8 (E4M3) activations × packed FP4 (E2M1)expert weights with per-32-K UE8M0 weight scales folded into the FP4→E4M3
dequant; fused L1 GEMM → SwiGLU → per-token FP8 requant → L2 GEMM → combine.
runs with A/B swapped, selected by the L1/L2 dispatch ladders.
Accuracy (DeepSeek-V4-Flash, 8×H20, swapAB on)
sgl-eval run gpqa
--n-repeats 16 --max-tokens 200000
--temperature 1.0 --top-p 1.0 --thinking
--out-dir /sgl-workspace/logs
--base-url http://localhost:30000/v1
2>&1 | tee /sgl-workspace/logs/gpqa_$(date +%Y%m%d_%H%M%S).console.log
== gpqa ==
198 examples x 16 repeats | 11789.7s | 2895 tok/s | 34.1M tokens
pass@16 = 96.46%
majority@16 = 90.15%
no_answer = 0.00%
stop_rate = 100.00%
truncated_rate = 0.00%
error_rate = 0.00%
Performance (single-op MegaMoE kernel, 8×H20, bench_kineto)
DeepSeekV4Flash
DeepSeekV4Pro