SM120: MoE grouped GEMM decode optimizations: skip padding I/O, BLOCK_M=32#59
Open
leavelet wants to merge 5 commits into
Open
SM120: MoE grouped GEMM decode optimizations: skip padding I/O, BLOCK_M=32#59leavelet wants to merge 5 commits into
leavelet wants to merge 5 commits into
Conversation
Two decode-focused optimizations for MoE grouped GEMM (contiguous + masked):
1. Skip D stores of M-padding rows. M-grouped 1D1D now always uses the
direct-store epilogue (swizzle_cd = 0) and predicates row stores on
validity (contiguous: m_indices >= 0; masked: local row < masked_m).
Padding rows of D become unspecified -- they were already
caller-dependent (A padding is not zeroed by the API) and no framework
reads them. Fewer DRAM writes helps decode AND prefill.
tests/test_fp8_fp4.py compares valid rows only (masked/psum already did).
2. BLOCK_M=32 via a 2x4 warp layout: kNWarps = (BLOCK_M < 64) ? 4 : 2
(the BF16 kernel already had the template parameter). SM120
contiguous-layout alignment drops to {128, 64, 32} keyed on per-expert
expected_m (96 rounds up to 128); BLOCK_M=32 is only picked for
m-grouped layouts at alignment 32, with BLOCK_M | alignment and
BLOCK_N % 64 == 0 (even kNTilesPerWarp for ldmatrix.x2) enforced.
Same-methodology A/B (bench_kineto, L2-flushed), M=32 decode:
W1 (128E topk2, 5120/2048 + 2048/2560): gate_up 421->397, down 244->234 us
W3 (256E topk8, 512/6144 + 6144/256): gate_up 444->412, down 249->212 us
(M=64 down: 357->288 us, -19.6%)
W2 (128E topk6 FP8xFP4): gate_up 716->656, down 389->355 us
Warm-loop (flush_l2=False): W1 down 200.4 us, W3 down 181.2 us -- the
latter beats the CUTLASS ptr-array blockwise baseline (Ping_64x128, 190.5 us).
Full test_fp8_fp4.py (444 cases incl. FP4/mixed) and test_bf16.py (220)
suites pass on RTX PRO 6000.
BN=64 satisfies the even-kNTilesPerWarp (ldmatrix.x2) constraint and measures never worse than BN=128 across W1/W2/W3 decode shapes, -1.6% on thin-N down layers, and doubles SM fill at tiny M: M=1 gate_up drops 30.6 -> 8.3 us (warm), now ahead of the Marlin fused-MoE kernel (9.2 us); M=8 reaches parity.
For m-grouped 1D1D FP8 layouts (BLOCK_K=128, SW128 A), the leader warp loads only the valid-prefix A rows of each tile with cp.async.cg (Swizzle<3,4,3> matches the TMA SW128 layout); padding rows keep stale SMEM — their outputs are never stored. Completion is tracked by the full barrier itself via cp.async.mbarrier.arrive.noinc (expected count 33 = 32 lanes + the TMA producer arrival), so the producer never stalls. TMA issue and barrier ops stay on lane 0; scheduler math is lane-uniform. The BK=64 candidate is dropped for m-grouped alignment-32 so the SW128 path applies at decode. Also fixes m-grouped launchers passing gmem_a_ptr = nullptr. Equal-conditions vs the CUTLASS ptr-array baselines (ncu locked clocks, M=32 decode, RTX PRO 6000) — DeepGEMM now leads on every target shape: W1 gate_up 350.2 vs 357.7 us W1 down 182.5 vs 194.4 us W3 gate_up 330.8 vs 346.7 us W3 down 171.1 vs 183.1 us W2 (FP4 wgt, warm) 645 vs 678, 336 vs 366 us (mxf8-mxf4 grouped) Read rate 1.48-1.49 TB/s = 98-99% of the measured pure-read peak. Full test_fp8_fp4.py (444) and test_bf16.py (220) suites pass.
Review findings on the decode optimizations, addressed: 1. Partial-K (k % 128 != 0, accepted for pure FP8): the cp.async A path would read past the row tail where TMA zero-fills. The cp.async path is now gated host-side (new kACpAsync stub flag) on contiguous A and k % BLOCK_K == 0; anything else falls back to the TMA A path. 2. Strided A (views are accepted; the TMA descriptor carries the real stride): same host gate covers it. 3. Odd N in the direct-store pair path: not reachable in practice (BF16 D requires n % 8 == 0 or the TMA descriptor creation fails), but the pair store now writes a single-element tail instead of dropping the last column, as defense in depth. 4. Small-N (n <= 32) m-grouped at alignment 32 left zero layout candidates (opaque assert): now rejected with an explicit message (no valid BLOCK_M=32 warp layout exists for N <= 32 tiles). Edge tests: partial-K/strided-A/N%64!=0 all produce correct results via fallback (diff ~7e-4); odd-N rejected at descriptor creation; small-N raises the explicit error. Main-path perf unchanged (W3 gate_up 348.6, W1 gate_up 381.4 us warm).
Mirrors the tests/test_fp8_fp4.py change: M-padding rows of D are unspecified with the padding-store skip; masked/psum checks already slice valid rows.
67 tasks
Fridge003
reviewed
Jul 15, 2026
| check_fp8_fp4_psum_zero_padding(a, d, grouped_layout) | ||
| else: | ||
| diff = calc_diff(d, ref_d) | ||
| # M-padding rows of `d` are unspecified; compare valid rows only |
Collaborator
There was a problem hiding this comment.
Revert change to this file
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
2 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.
Picked deepseek-ai#380
Motivation
MoE grouped GEMM at decode batch sizes is weight-bandwidth-bound on SM120: almost all DRAM time goes to streaming expert weights, while the M-padding of the contiguous/masked layouts adds pure-waste traffic (padded A-row reads, padded D-row writes) and the BLOCK_M=64 floor inflates padded compute.
This PR removes the padding waste and lifts the BLOCK_M floor, after it, DeepGEMM leads CUTLASS
on all tested MoE shapes (see numbers below) and streams weights at 98%~99% of
the device's measured pure-read DRAM peak.
What's in here (3 commits)
e18f68b).e18f68b,5c018f3).cp.async(7fdc998).Performance
RTX PRO 6000 Blackwell (188 SMs), M=32 decode tokens, multinomial top-k
routing.