perf(moe): warp-split moe_q4_grouped_mma variant, opt-in (BUTTER_LAGUNA_MOE_NSPLIT=1) - #75
Merged
Conversation
sdpa_multi_tc_varlen (and its non-varlen sibling sdpa_multi_tc, which shares the SDPA_TC_KPREP_SRC kernel) converted the WHOLE causal-prefix K range into the kh scratch buffer upfront, once per call, and then dispatched the softmax/merge kernels unconditionally for every 2048-wide KV block in that range -- even though the QK^T/PV GEMMs were already correctly skipped via q_cnt==0 for blocks outside a sliding-window layer's reachable range. For Laguna-S-2.1 (36 of 48 layers are sliding-window, window=512) this meant every prefill chunk on every windowed layer paid an O(prefix) cost (kh sized to base+chunk_len, kprep dispatched over that whole range, softmax+merge run per block) instead of O(window+chunk_len) -- the cost grows with how deep into the prompt the chunk sits. Fix: make K-prep per-block (mirroring V-prep/SDPA_TC_VPREP_SRC, which already worked this way) into a fixed bk-sized rolling buffer instead of a whole-prefix one, and gate kprep+softmax+merge on the same q_cnt>0 check the GEMMs already use. A block no row can reach contributes nothing either way (softmax's own positional mask already forces bm=-inf/bl=0 for such rows; merge's l_b<=0 early-return already no-ops on that combination) -- bit-identical output, just without the wasted dispatches. Side effect: kh's allocation is now a fixed size across calls instead of growing with prefix depth, which should also help the CUDA caching allocator reuse pooled buffers instead of re-allocating on every larger chunk. Ported from the pre-rename ffai-rust prefill-optimization campaign (source commit 8f6eb3c); gated there against the full laguna test suite (10/10) and the prefill-then-decode token-identity oracle (byte-identical greedy continuation, last-token argmax=290).
The per-block skip in sdpa_multi_tc_varlen omits whole KV blocks a sliding-window layer's rows can't reach at all (q_cnt==0), but a block that IS touched still runs the full bk-wide QK^T/PV GEMM + softmax/merge, even though Laguna's sliding_window (512) needs only a fraction of a 2048-wide block. Make the KV block size conditional: full-attention layers (win==0) keep bk=2048 unchanged (bounds the scores buffer, no window to exploit). Sliding-window layers (win!=0) now use a separate, independently-tunable BUTTER_LAGUNA_SDPA_WIN_BK (default 1024) -- picked over the window width itself (512) because at 512 the extra loop/launch overhead from doubling block count outweighs the reduced overcompute at some sequence lengths, while 1024 is flat-to-positive everywhere measured. Ported from the pre-rename ffai-rust campaign (source commit 6fece8b); gated there against the full laguna suite (10/10) plus the prefill-then-decode token-identity oracle (argmax=290, byte-identical greedy continuation).
BUTTER_LAGUNA_PREFILL_CHUNK still overrides explicitly; unset/0 now resolves through laguna::default_prefill_chunk(total) instead of a fixed 1024/2048 constant baked into the two test-harness call sites. A fixed 2048-token chunk (isolated, cold-GPU-per-run measurements) beats BOTH true monolithic (chunk==n) and the previous chunk=1024 convention at every n from 2048 through 32768. Monolithic is fine (and equivalent, since chunk.min(n) collapses either way) up to n=2048, but is unstable and often much slower above that: a single huge batched activation tensor stops fitting the working set (single monolithic n=8192 chunk measured between 191 and 364 tok/s across repeated cold-start runs, vs a stable ~460-540 tok/s at chunk=2048). Also adds an argmax cross-check to the BUTTER_LAGUNA_PP throughput bench: the synthetic token sequence is deterministic, so running the same n at different chunk sizes (auto vs an explicit override) and comparing the printed last-token argmax is a quick chunk-invariance sanity check alongside the separate BUTTER_LAGUNA_PREFILL tokenizer oracle. Ported from the pre-rename ffai-rust campaign (source commit 6c301a1); gated there against the full laguna suite (10/10), the tokenizer argmax=290 oracle (byte-identical greedy continuation -- chunk choice does not affect correctness by construction), and a chunk=0 (auto) vs explicit-override argmax cross-check at several prompt lengths.
ncu-confirmed (l1tex bank-conflict vs wavefront counters, moe_w4a16_marlin128): the dequant STORE into smem_WT was hitting ~18x overhead vs the conflict-free ideal. Root cause: smem_WT's row stride was exactly 128 halfwords (256 bytes), a multiple of the 32-bank x 4-byte conflict period, so a warp's 32 threads (which vary the ROW index while the bank-determining column index stays fixed within an unrolled dequant iteration) all aliased the SAME bank on every store. Padded the row stride to 136 halfwords -- the smallest multiple of 8 (WMMA's row-major ldm alignment requirement for __half) that reduces the conflict; modular arithmetic (row_wt/2 mod 32) shows 8-alignment makes a fully conflict-free stride impossible for this 32-row/128-bank shape, so 136 is the best achievable (32-way -> 4-way conflict, an 8x reduction). smem_X (the ~3.5x load-side offender) is NOT padded this round: GB10's actual opt-in dynamic-shared-memory cap is 101376 bytes (queried via CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN -- not the ~228KB the kernel's old doc comment assumed by analogy to a different GPU), and the padded-WT kernel already uses 99840 bytes, leaving no room to also pad X without exceeding the cap and failing cuFuncSetAttribute at launch. Ported from the pre-rename ffai-rust campaign (source commit 24da67e); gated there against the full laguna suite (10/10, with BUTTER_LAGUNA_MARLIN=1 actually set -- the bare default suite doesn't exercise this kernel), the tokenizer argmax=290 oracle with MARLIN=1 (byte-identical greedy continuation), and PP-bench argmax cross-checks at several sizes against pre-padding values (all exact matches).
…election sort moe_build_tiles (single-thread pre-pass kernel building per-tile descriptors for moe_q4_grouped_mma_dev, used by prefill_layer's two grouped-GEMM calls per layer under BUTTER_LAGUNA_SCHED=1 default-on desc=true mode) used an O(n_exp^2) selection sort to order experts by descending token count. At Laguna's n_expert=256 that's 32768 compares on a SINGLE GPU thread (the kernel is intentionally single-threaded: only blockIdx.x==0 && threadIdx.x==0 runs). ncu measured ~1.34ms/call average at pp2048 -- roughly 96 calls per prefill pass (48 layers x 2 dev-path grouped-GEMM calls) -- a real cost hidden inside the moe_gateup_gemm/moe_down_gemm profiling buckets. Replaced the selection sort with a standard heapsort (build max-heap keyed by cnt[order[i]], repeated extract-max, reverse for descending), verified correct via a standalone Python simulation (200+ random cases + edge cases: n=0, n=1, all-equal, all-zero, both sort directions, n=256) before writing any CUDA. ncu-confirmed the fix: moe_build_tiles average duration dropped 1,338,549ns -> 245,584ns (5.45x) at pp2048, same launch-skip/count sample. Correctness of the caller is provably insensitive to the exact tie-break order this sort produces (tiles write disjoint output regions regardless of emission order -- moe_q4_grouped_laguna_sched_ab already A/B-tests desc-mode vs default order on skewed/zero-count-expert groups and requires bit-exact agreement). Ported from the pre-rename ffai-rust campaign (source commit c8c76da); gated there against the full laguna suite (10/10, BUTTER_LAGUNA_MARLIN=1 BUTTER_LAGUNA_MOEFUSE=1), the tokenizer argmax oracle (byte-identical greedy continuation), and per-size argmax cross-checks matching known values at n=512/1024/2048/4096/8192/32768.
Prior rounds of profiling (tile-size, launch_bounds, launch-geometry, smem-budget, pipe-competition) converged on this kernel's ~31% tensor-pipe utilization being latency/occupancy-bound, with occupancy pinned at 5 blocks/SM by a 92-95 reg/thread register-file ceiling. This deepens the software-pipelined cp.async prefetch depth (the remaining latency-hiding lever) from 2 buffers to 3. Standalone nvcc -Xptxas -v confirms zero register-count change (95->95 regs, 0 spill) -- only smem grows (12800->19200 bytes), which still fits 5 blocks/SM under the GB10 101376-byte opt-in cap (96000 bytes), so occupancy is unaffected by construction, not just by measurement. A live same-session A/B (patched vs reverted baseline, both run back-to-back to control for this box's substantial run-to-run/day-to-day throughput drift) measured 3-stage as +6.4% @2048, +2.3% @4096, +1.2% @8192, and flat (within noise) at 512/1024/16384/32768 -- modest but real, consistent with occupancy remaining the dominant cost and pipeline depth only recovering part of the gap. Ported from the pre-rename ffai-rust campaign (source commit eb51c29); gated there against the full laguna suite (10/10, byte-identical greedy continuation, decode-step-1 argmax=290 preserved), moe_grouped_mma_test (15/21 pass, same 6 known unrelated failures as the documented baseline), and moe_q4_grouped_laguna_sched_ab (ragged/zero-count-expert case).
…msnorm_batched_f16out to iron@dev iron PR #75 landed iron_moe_gather_q4, iron_moe_gather_q4_swiglu, iron_rms_norm_f16out, and iron_gated_group_rmsnorm_batched_f16out on thewafflehaus/iron@dev. Replace the always-error stubs and the rms_norm+cast/gated_group_rmsnorm_batched+cast fallbacks with real dispatches, mirroring the existing _relu2/_down sibling patterns (cached_ir + kernel_ir_for). rms_norm_f16out keeps the f32-only, n%128==0 && n<=4096 fast-path guard (same 16-bit vec4 codegen restriction as plain rms_norm) and falls back otherwise. Bump the iron git dependency lock to pick up the merged kernels.
…fix) Picks up the pool-cap fix landed in thewafflehaus/iron#76 on top of the #75 kernel port already locked in the prior commit. Confirmed in the gate below: the single-process comma-list ladder now matches isolated-per-config within normal run-to-run noise (no more mid-ladder degradation from the parking-cap pressure #76 fixes).
…A_MOE_NSPLIT=1) Splits the NS=16 N-slots across an extra warp axis (4->8 warps, 128->256 threads) instead of shrinking BN/the output tile, halving the per-thread f32 accumulator (oc[16][4]=64 regs -> oc[8][4]=32 regs) while leaving the A-tile and its reuse unchanged. Under __launch_bounds__(256,4) this compiles to 64 registers with zero spill, raising the register-limited occupancy ceiling from 41.7% (5 blocks/SM x 4 warps) to 66.7% (4 blocks/SM x 8 warps) -- confirmed empirically via ncu (66.07-66.33% achieved vs the kernel's own 66.67% ceiling), not just via standalone nvcc projection. Correctness: nsplit vs default max_rel=0.00e0 (bit-exact) across uniform, ragged/non-power-of-two, and zero-count-expert cases, on both the host-descriptor (moe_q4_grouped_mma) and on-device-descriptor (moe_q4_grouped_mma_dev) call paths. Full-model gate byte-identical (argmax=290, decode steps 2-17, 16-token greedy continuation) with the same env stack used for benchmarking. Isolated-cold ladder (2 reps/config, BUTTER_LAGUNA_MARLIN=1 BUTTER_LAGUNA_MOEFUSE=1): nsplit beats default at every size by a small, consistent margin (+0.2% to +1.2%, mean of 2 reps) -- 512: 458.7->464.0, 1024: 593.9->595.2, 2048: 669.9->676.8, 4096: 648.7->654.3, 8192: 619.1->625.9 tok/s. This is a real, directionally consistent win (9/10 paired samples favor nsplit) but sits at the edge of this campaign's documented 1-3% run-to-run noise band, well short of the occupancy jump's theoretical ceiling -- consistent with a latency-bound kernel that is only ~30-38% of prefill wall time. Gated default-off pending a longer confirmation run before considering flipping it on.
TheTom
marked this pull request as ready for review
July 29, 2026 17:53
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.
What
Adds
moe_q4_grouped_mma_nsplit: doubles the warp count (4→8, 128→256 threads) and splits the NS=16 N-slots across the new warp axis, halving the per-thread f32 accumulator (64→32 regs). BM/BN and the A-tile are unchanged, so A-tile reuse is preserved — this is a restructure, not the tile-parameter shrink that was measured negative earlier in the campaign.Default OFF, behind
BUTTER_LAGUNA_MOE_NSPLIT=1.Result — real but small, hence opt-in
__launch_bounds__(256,4)(verified on the shipped kernel source, not just a scoping copy)max_rel = 0.00e0) vs the default kernel across uniform, ragged/non-power-of-two, and zero-count-expert cases, on both host- and device-descriptor pathsWhy it's small, and why that's the interesting part
Profiling established this kernel is latency-bound with no pipe saturated, at ~30-38% of prefill wall time. A large occupancy gain therefore hides latency better but cannot move the total much — which is what we measure. That is a useful negative result for the campaign: occupancy was not the binding constraint.
Left default-OFF because a ~1% gain sits at the edge of this box's documented 1-3% run-to-run noise; flipping the default deserves a larger measurement budget than one afternoon.
FYI @ekryski.