perf(laguna): prefill optimization campaign — sliding-window SDPA, marlin smem padding, moe heapsort + pipeline depth - #74
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).
TheTom
marked this pull request as ready for review
July 25, 2026 23:27
Contributor
Author
|
Now fully gated end-to-end on GB10 — decode included. Two commits added: the four Decode was previously blocked at step 2 on the Ladder (isolated-per-config): 461 / 594 / 668 / 650 / 622 tok/s @ 512/1K/2K/4K/8K. Notably, single-process now measures the same within ~1% — previously it read 10-32% lower at mid sizes. That's iron #76 (allocator pool cap 4→8 GiB) confirmed in the wild, not just in an A/B. |
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.
Summary
Ports 6 commits from an overnight Laguna prefill-optimization campaign (run on a pre-rename
ffai-rustworking tree) onto this repo's post-renamedevbranch. Each source commit was individually located, re-derived against this tree's current structure (crate names,BUTTER_LAGUNA_*env convention,wh-butter-*/wh-iron-*paths), and re-verified to build here — see the per-commit list below.Verified ladder (isolated-per-config methodology — one process per size, cold GPU each run — measured on the source tree at HEAD
eb51c29):vs. the campaign's starting point (559 @2048 / 504 @8192 / 317 @32768) and the reference C++ fork's stock numbers (663 / 660 / 631 @2048/8192/32768) — we now beat fork-stock at 2048.
Every change was gated (on the source tree, GB10) against byte-identical greedy continuation and the last-token argmax=290 oracle vs. the reference implementation.
The 6 fixes
fix(sdpa)8f6eb3c→f7c0ab2— sliding-window attention layers were K-prepping and running softmax/merge over the entire causal prefix instead of just the touched window; fixed to per-block K-prep gated on the sameq_cnt>0check the QK/PV GEMMs already used (bit-identical output, removes wasted dispatches — the root cause of prefill throughput decaying with sequence length).perf(sdpa)6fece8b→0425f8b— windowed-layer KV block granularity shrunk 2048→1024 (newBUTTER_LAGUNA_SDPA_WIN_BKenv, full-attention layers unaffected) — a touched block was still 4x wider than the 512 sliding window even after fix Release v0.1.0 #1.perf(prefill)6c301a1→f0106ab— chunk size now resolves tomin(total, 2048)when unset (chunk==0sentinel), replacing a fixed 1024 default; monolithic chunking is unstable and often much slower above ~2048 tokens (single monolithic n=8192 chunk measured 191-364 tok/s across cold-start runs vs. a stable 460-540 at chunk=2048).BUTTER_LAGUNA_PREFILL_CHUNKstill overrides explicitly.perf(marlin)24da67e→bb548cc— paddedmoe_w4a16_marlin128'ssmem_WTrow stride 128→136 halfwords; ncu showed the unpadded stride caused a 32-way shared-memory bank conflict on every dequant store (~18x overhead), 136 is the best achievable under WMMA's 8-halfword alignment requirement and the device's 101376-byte opt-in smem cap.perf(moe)c8c76da→b2ffe9d— replaced an O(n²) single-thread selection sort inmoe_build_tileswith an O(n log n) heapsort; at n_expert=256 this was ~1.34ms/call (ncu), dropping to ~0.246ms (5.45x). Correctness is insensitive to tie-break order (tiles write disjoint regions regardless of emission order).perf(moe)eb51c29→2ffc2b6— deepenedmoe_q4_grouped_mma's cp.async software pipeline from 2 to 3 stages; register count is unchanged (95→95, 0 spill) so occupancy (5 blocks/SM) is unaffected by construction, only smem grows (12800→19200 bytes, still under cap). Modest but real gain from deeper latency-hiding; occupancy remains the kernel's dominant cost.This port has NOT yet been re-validated on the GB10 — the box is in use by another job concurrently with this port. Every source commit was individually gated on the GB10 (full
lagunatest suite 10/10, byte-identical greedy continuation, argmax oracles) on the pre-rename tree, but this re-derived port (adapted to the post-rename crate/env names and this tree's current state, which has diverged ~900 lines from what was previously merged) needs its own confirming run — fulllagunasuite + the prefill/decode argmax oracle — before merge.Gates run (Mac, no CUDA hardware — this is the ceiling available in this environment)
cargo build --workspace— clean (pre-existing warnings only, no errors)cargo check -p wh-butter-cuda --features cuda— clean, typechecks the CUDA backendrustfmt --checkon the 3 touched files — fails, but this is pre-existing/repo-wide:origin/dev's versions of these same files already failrustfmt --checkbefore this port (norustfmt.toml, no CI fmt job;wh-butter-ops/src/lib.rsuses a dense hand-formatted convention for CUDA dispatch calls throughout). This port's diff follows the same established convention as the surrounding code.cargo clippyon the 3 touched crates — no new errors; one new trivial warning (clippy::manual_clampondefault_prefill_chunk'stotal.min(2048).max(1), kept literal to match the ported source commit) alongside ~46 pre-existing warnings elsewhere inwh-butter-modeltestsunrelated to this change.grepfor leftoverffai/metaltile/FFAI_naming across the full diff — zero matches.Files touched
rust/crates/wh-butter-ops/src/lib.rs(commits 1, 2, 4, 5, 6)rust/crates/wh-butter-models/src/laguna.rs(commit 3)rust/crates/wh-butter-modeltests/src/lib.rs(commit 3)Test plan
lagunatest suite (10/10) on GB10 against this branchBUTTER_LAGUNA_PREFILL=1argmax=290 oracle, byte-identical greedy continuationBUTTER_LAGUNA_MARLIN=1variant of the same oracle (exercises fix ci: PRs run unit only; release runs unit + integration #4)moe_grouped_mma_test(expect the same 6 known-unrelated failures as documented baseline)