SM120: MoE grouped GEMM decode optimizations: skip padding I/O, BLOCK_M=32 - #380
SM120: MoE grouped GEMM decode optimizations: skip padding I/O, BLOCK_M=32#380leavelet wants to merge 4 commits into
Conversation
| uint32_t offset = cute::Swizzle<3, 4, 3>::apply(item * 16); | ||
| const auto dst = static_cast<uint32_t>(__cvta_generic_to_shared(smem_dst + offset)); | ||
| const char* src = gmem_src + row * static_cast<uint64_t>(gmem_row_stride) + (item % kChunks) * 16; | ||
| asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(dst), "l"(src) : "memory"); |
There was a problem hiding this comment.
🔴 critical: Zero-fill the partial cp.async K tile: When shape_k is not divisible by 128, the final K iteration still issues all eight 16-byte copies per row. Unlike the replaced TMA load, these copies do not zero-fill out-of-bounds K elements and instead read into the next row or beyond the allocation, producing incorrect accumulations for inputs that the API currently accepts.
🤖 v6
| // stores (padding rows of D are unspecified); faster at decode AND prefill (fewer bytes). | ||
| const bool grouped_skip_padding_store = desc.kernel_type == KernelType::Kernel1D1D | ||
| and (desc.gemm_type == GemmType::MGroupedContiguous or desc.gemm_type == GemmType::MGroupedMasked); | ||
| const auto swizzle_mode_cd = (not grouped_skip_padding_store and desc.cd_n_contiguous and layout.block_n * cd_size >= 128) ? 128 : 0; |
There was a problem hiding this comment.
🔴 critical: Preserve the final column for odd N: For an SM120 grouped 1D1D GEMM with odd N, this forces the direct-store epilogue, while contiguous outputs take its paired-store branch and only store when col + 1 < shape_n. The final unpaired column is therefore never written, even though the public grouped APIs only require n > 0; retain TMA storage for this case or add a scalar tail store.
🤖 v6
| if (block_m < 64 and block_n != 64) | ||
| continue; |
There was a problem hiding this comment.
🔴 critical: Keep a layout candidate for small-N decode: When the SM120 grouped alignment is 32 and N <= 32, block_m_candidates contains only 32 while block_n_candidates contains only 16 and optionally 32. This condition rejects every combination because none has BLOCK_N == 64, causing the later non-empty-candidate assertion to fail for valid small-N grouped calls.
🤖 v6
| smem_a[s], reinterpret_cast<const char*>(gmem_a_ptr) + static_cast<uint64_t>(m_idx) * shape_k + k_idx, | ||
| valid_rows, shape_k, lane_idx); |
There was a problem hiding this comment.
🟡 warning: Pass the actual A row stride to cp.async: For a grouped-contiguous K-major A view with stride(-2) > K, which passes the existing layout checks, this computes row addresses and passes a row stride using shape_k. The replaced TMA descriptor used a.stride(-2), so the new path reads subsequent rows from incorrect offsets; the actual stride must be carried into the kernel or this path must be restricted to compact inputs.
🤖 v6
| // to prevent tiles from straddling expert boundaries. | ||
| if (is_m_grouped_contiguous(desc.gemm_type) and block_m > runtime_align) | ||
| if (is_m_grouped_contiguous(desc.gemm_type) | ||
| and (block_m > runtime_align or runtime_align % block_m != 0)) |
There was a problem hiding this comment.
🔴 critical: BLOCK_M=32 is pushed for EVERY m-grouped/runtime_align==32 GEMM, but the device kernel (kNWarps=4) only supports BLOCK_N in {64,128} and the heuristic enforces block_n == 64 for block_m < 64. For an m-grouped decode with small N (n <= 32, the is_small_n branch) this filters out ALL block_n candidates and trips DG_HOST_ASSERT(not candidates.empty()), regressing a shape that previously used BM=64/BN=16/32. Gate the BM=32 candidate on not is_small_n (or fall back to BM=64 when no valid BLOCK_N exists).
🤖 v4
| @@ -145,7 +157,11 @@ struct SM120ArchSpec { | |||
|
|
|||
There was a problem hiding this comment.
🟡 warning: grouped_skip_padding_store forces swizzle_mode_cd = 0 for ALL m-grouped 1D1D, including the BF16 (KernelNoSF) dispatches whose kernel has NEITHER the cp.async A-skip NOR the row-is-valid store skip. It silently switches those shapes from the TMA-store epilogue to the direct-store epilogue (correct result, different codegen/perf) purely as a side effect, contradicting the comment's 'so the kernel can skip M-padding row stores'. The store/A-skip capability is FP8-only; the BF16 cd-store path should not be toggled by it.
🤖 v4
| } | ||
| } | ||
|
|
||
| CUTLASS_DEVICE void cpasync_commit() { |
There was a problem hiding this comment.
🔵 suggestion: cpasync_commit() and cpasync_wait<N>() are defined but never called anywhere in the codebase — the completion path uses cp.async.mbarrier.arrive.noinc directly (which tracks all prior uncommitted cp.async of the calling thread), so commit_group/wait_group are unnecessary. Recommend removing the two dead helpers to avoid implying an unused commit/wait protocol, or add a comment noting they are intentionally provided for future use. Not a correctness issue.
🤖 v3
| constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); | ||
| const uint32_t batch_idx = kIsBatchedMM ? scheduler.current_group_idx : 0; | ||
|
|
||
| // Valid (non-padding) A rows in this tile; padding is a per-expert suffix |
There was a problem hiding this comment.
🔵 suggestion: The cp.async prefix load (cpasync_load_rows on rows [0, valid_rows)) and the valid_rows computation both assume valid rows are a contiguous prefix and padding a contiguous suffix within each BLOCK_M tile. This is correct for the current contiguous/masked layouts (per-expert m is aligned up to a multiple of the alignment, which is a multiple of BLOCK_M, so tiles never straddle a valid/padding boundary), but it is an undocumented invariant in this file. Consider adding a brief comment stating that padding is guaranteed to be a per-expert suffix so the prefix-count is equivalent to the epilogue's per-row row_is_valid predicate; if a future layout ever interleaves -1 markers, the producer would load wrong rows while the epilogue would still store the right ones, silently corrupting results.
🤖 v3
| if constexpr (kUseCpAsyncA) { | ||
| if constexpr (kGemmType == GemmType::MGroupedContiguous) { | ||
| uint32_t cnt = 0; | ||
| for (uint32_t r = lane_idx; r < BLOCK_M; r += 32) |
There was a problem hiding this comment.
🔵 suggestion: For the MGroupedContiguous branch, valid_rows is computed by every lane via a warp loop + __reduce_add_sync, but m_idx used for the A base pointer is derived with get_global_idx<kAGroupOffset=false> (= m_block_idx*BLOCK_M) while the count reads grouped_layout + m_idx + r. Since grouped_layout is indexed by absolute row and m_idx equals the absolute tile start here, this is correct — but the double meaning of m_idx (row index for grouped_layout vs A element offset via m_idx * shape_k) is subtle. A one-line comment clarifying that m_idx is the absolute starting row of the tile would help future readers avoid an off-by-group error if the offset convention ever changes.
🤖 v3
| if (runtime_align <= kMinBlockM) | ||
| const bool is_m_grouped = is_m_grouped_contiguous(desc.gemm_type) | ||
| or desc.gemm_type == GemmType::MGroupedMasked; | ||
| if (runtime_align == 32 and is_m_grouped) |
There was a problem hiding this comment.
🔵 suggestion: is_m_grouped here (used to allow BLOCK_M=32 and to drop the BK=64 candidate) includes MGroupedContiguousWithPsumLayout via is_m_grouped_contiguous, whereas the padding-skip store gate (grouped_skip_padding_store in get_storage_config) and the kernel's kUseCpAsyncA deliberately exclude the psum layout. The net effect is that the psum layout can now be dispatched at BLOCK_M=32 / BLOCK_N=64 through the TMA-store epilogue path (no cp.async, no store skip). This is internally consistent and presumably covered by the psum test cases, but it is an easy-to-miss asymmetry: BLOCK_M selection includes psum while the two decode optimizations exclude it. Worth a short comment noting the psum layout intentionally rides the BLOCK_M=32 tile but keeps the standard TMA-store/TMA-A path.
🤖 v3
🤖 ds-review-bot Code Reviewv6The new cp.async and direct-store paths produce incorrect results for accepted partial-K, odd-N, and strided-A inputs. The new block layout filtering also leaves valid small-N decode configurations with no candidate. v4The PR is already fully committed/clean at HEAD 7fdc998 (base a6b593d). The three optimizations (skip padding row stores; BLOCK_M=32 via a 2x4 warp layout; skip padding A reads via warp-cooperative cp.async) are all present. I read the full kernel/heuristics/host-dispatch stack and verified the central correctness invariants; everything that is actually exercised by the tested decode shapes (large N, SW128/BK=128 FP8-A, no split-K, no accumulation) is internally consistent. The remaining concerns are robustness/divergence gaps that the current per-shape gating leaves unprotected, and a brittle (currently correct) producer warp/barrier model around the new cp.async path. v3This PR delivers three decode-focused optimizations for SM120 MoE grouped GEMM (m-grouped contiguous + masked, 1D1D FP8/FP4): (1) skipping D stores of M-padding rows via the direct-store epilogue and per-row validity predication, (2) enabling BLOCK_M=32 through a 2x4 warp layout (kNWarps=4) with BLOCK_N=64, and (3) skipping M-padding A reads via a warp-cooperative cp.async prefix loader whose completion is tracked by the full mbarrier itself (expected count 33 = 32 lanes + 1 TMA producer arrival). The implementation is coherent and the gating is consistent across the three layers that must agree: the storage-config gate (grouped_skip_padding_store), the runtime launcher gate (should_skip_padding_store), and the kernel template flag (kSkipPaddingStore / kUseCpAsyncA) all key on the same {MGroupedContiguous, MGroupedMasked} set, correctly excluding MGroupedContiguousWithPsumLayout (which keeps the TMA-store epilogue). The producer barrier init is bumped to 33 only when kUseCpAsyncA, the transaction-byte count is correctly reduced by TMA_A_BYTES (SMEM_TMA_BYTES_PRODUCER), and the m-grouped launchers now pass a real gmem_a_ptr instead of nullptr. Correctness of the cp.async prefix load rests on the invariant that valid rows form a contiguous prefix and padding a contiguous suffix within every BLOCK_M tile; this holds because per-expert row counts are aligned up to a multiple of the contiguous-layout alignment (itself a multiple of BLOCK_M), so no tile straddles a valid/padding transition mid-expert. The producer's __reduce_add_sync count over m_indices>=0 is therefore equivalent to the epilogue's per-row row_is_valid check. The masked path avoids unsigned underflow because the scheduler only yields blocks with m_block_idx*BLOCK_M < masked_m. Test coverage is claimed to pass (test_fp8_fp4.py 444 cases, test_bf16.py 220) and the epilogue-store padding-skip is matched by the test comparing valid rows only. Overall this is a well-scoped, well-commented, low-risk change with strong measured wins. The findings below are minor. Files reviewed: 7 |
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.
validation compares valid rows only (masked/psum already did); the
SM120-aware test copy lives in the sglang fork's sgl_deep_gemm/tests.
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 m-grouped FP8/FP4/mixed and BF16 suites pass on RTX PRO 6000 with
the valid-rows comparison (in the sglang fork's test copy; this repo's
tests/ is kept unchanged).
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).
2 participants
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.