SM90 (Hopper) FP4 MegaMoE fused kernel with swapAB small-batch path - #53
Open
qiushixiaoyu wants to merge 22 commits into
Open
SM90 (Hopper) FP4 MegaMoE fused kernel with swapAB small-batch path#53qiushixiaoyu wants to merge 22 commits into
qiushixiaoyu wants to merge 22 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>
Follow up deepseek-ai#63, which dropped the UE8M0 (power-of-two) activation scale on the SM90 FP8 path. The same reasoning applies to the FP4 path: SM90 WGMMA has no hardware block-scale operand (the SF is a plain FFMA in the epilogue), so rounding the scale to a power of two bought nothing and only cost precision. The L2 SF pool is already fp32 (`l2_acts_sf` is kFloat32), so this is byte/layout neutral. Only the activation scale changes. The FP4 weight SFB stays UE8M0 (DSV4 external format) and is still decoded through the LUT path. Also clamp amax before the reciprocal: padded rows have amax==0, and 448/0=inf would produce 0*inf=NaN. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Fast-math SiTU now computes the gate/up caps with __tanhf and the sigmoid via the 0.5 * (1 + tanh(x / 2)) identity instead of __expf + reciprocal. Non-fast-math and SwiGLU paths are unchanged. Verified diff=0.0000 on 3 Kimi scenarios and 2 DSV4Flash regressions; 4096 batch/rank latency improves 23.454 -> 22.830 ms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
GEMM1/GEMM2 FP32 activation scales were issued as one 1D TMA per K-group (4 for the SiTU 32/32 path). Extend the SF descriptor with a K-group box dimension and fetch all groups of a stage in one 2D TMA, cutting TMA instructions and barrier arrivals. expect_tx bytes are unchanged. Verified diff=0.0000 on Kimi and DSV4Flash scenarios; 4096 batch/rank latency improves to ~22.78 ms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The official Kimi-K3 per-32 activation-scale recipe costs ~15% at large batch on SM90: each K32 WGMMA must fully drain before its FP32 promotion, and the drain cannot be pipelined within the 640-thread register pool (ping-pong banks spill). Controlled experiments (doc 15.6-15.11) show the entire FP4-vs-DeepEP-normal gap comes from this, so SiTU now defaults to the legacy 128/64 recipe via get_act_sf_grans_for_mega_moe_sm90_fp4; DG_SM90_FP4_SITU_ACT_GRAN_128_64=0 switches back to per-32 for accuracy comparison. With 128/64 the prefill bundle is healthy again, but Kimi routing density (e = batch/7) wastes 43%/24% of BLOCK_M=128 rows at batch 1024/2048 where forced decode wins 20%/10%, while e=585 (batch 4096) favors prefill -- hence the SiTU e=512 boundary (per-32 keeps 2048). Sweep 1-8192: FP4 now beats DeepEP normal at every point (1.76-1.87x for batch 1-256, 1.26x/1.13x at 1024/2048, 1.015x at 4096/8192). Kimi 3 scenarios, per-32 switchback, and DSV4 regressions all diff=0.0000. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The swapAB decode path was gated on 'not use_situ' because its data layout only covers 128/64 activation scales, which SiTU's original per-32 recipe violated. Now that SiTU defaults to 128/64, gate on the granularity instead, and unify the L1 epilogue activation: the swapAB path previously hard-coded clamp+silu and now shares the same gated_activation lambda as the regular epilogue (bit-identical for SwiGLU, adds SiTU support). DG_SM90_FP4_SWAP_AB_MAX_E overrides the e<16 crossover for tuning; A/B on Kimi (896 experts, top-k 16) confirms the default 16 stands -- e=36.6/73 lose 26%/13% on the swapped path, matching the original swiglu tuning. Kimi decode batches 1-64 (e<16) gain 13-15%: batch 1 now 403us vs DeepEP normal 801us (1.99x), batch 4/16 reach 2.19x. Kimi 3 scenarios (swapAB specialization exercised), and DSV4 2 scenarios all diff=0.0000. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The valley around batch 512 on Kimi (896 experts, top-k 16, e = batch/7) is a tile-quantization band, with its true bottom at batch 640 (0.978x vs DeepEP normal). For e in (80, 112] a BLOCK_M=128 tile pads M no worse than two BLOCK_M=64 tiles while halving per-expert B decodes, and this is exactly the band where B-decode reuse is lowest: prefill config measures +4.1%/+2.5% at batch 640/768. Near e~128 routing spread makes single 128-tiles overflow into a second one, so decode wins big again (batch 896 decode is 15% faster), and batch 512 itself (e=73) ties -- it sits at the ceil(73/64)*64 = 128-row tile-math floor. Encode this as a non-monotonic is_fp4_sm90_prefill_band predicate shared by the block-config heuristics and the API feature defaults so the BLOCK_M choice and the early_b_decode/ss_nsplit/wide_load_decode bundle always switch together. DG_SM90_FP4_SITU_MIDBAND_LO/HI override the band. After: 512-1024 all >= 1.018x (640: 0.978->1.022, 768: 1.002->1.026); 512/896/1024 unchanged. Kimi 3 scenarios and DSV4 2 scenarios diff=0.0000. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Boundary A/B at batch 512/544/576/608 (e=73.1/77.7/82.3/86.9) locates the decode->prefill crossover between e=73 and e=78: prefill already wins +2.6% at e=77.7 and its latency is nearly flat across the band (one 128-tile per expert regardless of batch), while decode climbs as the share of experts with <=64 tokens vanishes. The previous lo=80 left batch 544 on the wrong side. Batch 512 (e=73.1) still favors decode by 0.8% -- the single-64-tile expert share (~16%) outweighs halved B decodes there, so it stays outside the band. Kimi 3 accuracy scenarios diff=0.0000. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Boundary scan on Kimi SiTU 128/64 at batch 128/160/192/224 (e=18.3/22.9/27.4/32.0): swapAB still wins +3.1% at e=18.3 and loses -6..-18% from e~23, so the crossover sits in (18.3, 22.9). The old bound 16 came from the original swiglu tuning and left batch 128 on the wrong side. Kimi 3 accuracy scenarios diff=0.0000. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the hand-tuned swapAB/decode/prefill boundaries with a per-expert expected-cost argmin. With X ~ Poisson(e) tokens per expert, in units of one regular-mainloop row: R(M) = E[ceil(X/M)]*M, T(M) = E[ceil(X/M)] cost_decode = R(64) + 6.29*T(64) cost_prefill = (1-0.003)*R(128) + 6.29*T(128) cost_swap = 1.335*R(8) + 40.6*P(X>0) Constants fitted offline (scratchpad fit_cost_model.py) on 21 same-shape A/B measurements; all 21 classify correctly and the model reproduces every hand-tuned boundary (swap<19 vs 20, prefill (76,115.5] vs (76,112], large-batch prefill from 578 vs 512). It additionally predicted prefill wins inside the former decode band -- confirmed +3.8%/+2.8% at batch 1540/2450 (e=220/350), which the hand bands were leaving on the table. The row/decode cost ratio is shape-independent to first order, so the boundaries transfer to other expert shapes / EP degrees without rescans. The block-config heuristic and the API feature defaults share the same argmin so BLOCK_M and the feature bundle never mix. Scope: SiTU 128/64 only; swiglu and per-32 keep their hand-tuned paths. DG_SM90_FP4_COST_MODEL=0 falls back to the hand-tuned bands; DG_SM90_FP4_SWAP_AB=0 still disables swapAB inside the model. Regression: Kimi 3 + DSV4 2 scenarios diff=0.0000; nine batch points (128..4096) all land on their measured-optimal config within 0.4%. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DG_SM90_FP4_MATH_WG_DECODE, DG_SM90_FP4_EARLY_B_DECODE and DG_SM90_FP4_STAGE_CAP expose the decode-participation, early-B-decode and pipeline-depth knobs for experiments; all default to current behavior. Small-batch tile-loop probing used them to rule out decode throughput (math-WG decode is 15-40% WORSE), decode/TMA overlap (early B-decode is neutral) and pipeline depth (4 vs 5 stages is neutral) as the source of the ~600ns/K-stage steady-state (doc 15.17). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The Linear1/Linear2 branch and the swapAB n_swap dispatch chain (4-5 compares + generic-lambda selection) executed on every K-stage of the math warpgroup -- the warp whose per-stage service time sets the whole pipeline's steady beat (~525ns, doc 15.17). Both decisions are tile constants: wrap the K-stage loop in a run_k_stages<kIsL1, kNSwap> template and dispatch once per tile instead. Each instantiation is a branch-free specialized loop body, which also unlocks compiler loop optimizations across stages. Small-batch gains (Kimi SiTU 128/64, 30-rep us_max): batch 1 409->399us (2.01x vs DeepEP normal, first time above 2x), batch 16 1627->1595us, batch 64 1812->1752us, batch 128 2184->2048us (+6%, the deeper N24 dispatch chain saved the most). batch 512/4096 unchanged. Kimi 3 + DSV4 2 accuracy scenarios diff=0.0000. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DG_SM90_FP4_EXPERTS_PER_WAVE overrides the wave size for sensitivity
probing; default behavior unchanged. The 512+ investigation swept it
({112,56,28,14}: +-1% flat) alongside stage caps (flat), BLOCK_N=256
(rejected on paper: B-decode volume is N-tiling-invariant and n256
accums blow registers) and L2-arrival refinement (measured wait is only
1-2% of kernel time). Phase-aligned batches (multiples of 896, e
divisible by 128) confirm the mid-range valleys are tile-phase effects
recoverable at the deployment layer, and at aligned large batches the
FP4 kernel reaches 80% of FP8 peak vs GLM's 79% -- the remaining
speedup gap vs GLM is normal-baseline strength, not kernel efficiency
(doc 15.22).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--profile-breakdown now also emits the five normal-baseline stages (dispatch / L1 GEMM / act+quant / L2 GEMM / combine). At batch 7168 (phase-aligned) this pins the comparison the fused kernel must answer to: normal pure GEMM = 29.36ms = 87.1% of FP8 peak vs fused = 32.0ms = 80%, i.e. the fused GEMM runs ~7 points below DeepEP normal's grouped GEMM and eats two thirds of the 3.5ms hideable-communication dividend. DG_SM90_FP4_WIDE_LOAD_DECODE probes wide-load decode outside the decode band: flat at 4096/7168, so the historical band gating stands. Doc 15.23. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Drop the five probes whose experiments concluded with no benefit: DG_SM90_FP4_MATH_WG_DECODE (15-40% worse), DG_SM90_FP4_EARLY_B_DECODE (neutral), DG_SM90_FP4_STAGE_CAP (depth-insensitive), DG_SM90_FP4_EXPERTS_PER_WAVE (flat) and DG_SM90_FP4_WIDE_LOAD_DECODE (flat outside the decode band). Behavior is identical to the probed defaults; conclusions live in docs 15.17/15.22/15.23. The functional switches (recipe/cost-model/swapAB kill-switch and the hand-tuned fallback parameters) stay. Co-Authored-By: Claude Fable 5 <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