diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 03e4ea87..5c1243e8 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -581,7 +581,7 @@ dependencies = [ [[package]] name = "wh-iron" version = "0.1.0" -source = "git+https://github.com/thewafflehaus/iron?branch=dev#ee009c49bb60f87dd00ed6b34a559c7901d1690a" +source = "git+https://github.com/thewafflehaus/iron?branch=dev#131b1993ce4dce72c176af9368efe1d7263c7e97" dependencies = [ "bytemuck", "inventory", @@ -600,7 +600,7 @@ dependencies = [ [[package]] name = "wh-iron-codegen" version = "0.1.0" -source = "git+https://github.com/thewafflehaus/iron?branch=dev#ee009c49bb60f87dd00ed6b34a559c7901d1690a" +source = "git+https://github.com/thewafflehaus/iron?branch=dev#131b1993ce4dce72c176af9368efe1d7263c7e97" dependencies = [ "inventory", "rustc-hash", @@ -615,7 +615,7 @@ dependencies = [ [[package]] name = "wh-iron-core" version = "0.1.0" -source = "git+https://github.com/thewafflehaus/iron?branch=dev#ee009c49bb60f87dd00ed6b34a559c7901d1690a" +source = "git+https://github.com/thewafflehaus/iron?branch=dev#131b1993ce4dce72c176af9368efe1d7263c7e97" dependencies = [ "inventory", "rustc-hash", @@ -629,7 +629,7 @@ dependencies = [ [[package]] name = "wh-iron-macros" version = "0.1.0" -source = "git+https://github.com/thewafflehaus/iron?branch=dev#ee009c49bb60f87dd00ed6b34a559c7901d1690a" +source = "git+https://github.com/thewafflehaus/iron?branch=dev#131b1993ce4dce72c176af9368efe1d7263c7e97" dependencies = [ "proc-macro2", "quote", @@ -639,7 +639,7 @@ dependencies = [ [[package]] name = "wh-iron-runtime" version = "0.1.0" -source = "git+https://github.com/thewafflehaus/iron?branch=dev#ee009c49bb60f87dd00ed6b34a559c7901d1690a" +source = "git+https://github.com/thewafflehaus/iron?branch=dev#131b1993ce4dce72c176af9368efe1d7263c7e97" dependencies = [ "objc2", "objc2-foundation", @@ -656,7 +656,7 @@ dependencies = [ [[package]] name = "wh-iron-std" version = "0.1.0" -source = "git+https://github.com/thewafflehaus/iron?branch=dev#ee009c49bb60f87dd00ed6b34a559c7901d1690a" +source = "git+https://github.com/thewafflehaus/iron?branch=dev#131b1993ce4dce72c176af9368efe1d7263c7e97" dependencies = [ "half", "rustc-hash", diff --git a/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs b/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs index e4e6e418..732f327c 100644 --- a/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs +++ b/rust/crates/backends/wh-butter-cuda/tests/moe_grouped_mma_test.rs @@ -535,6 +535,142 @@ fn moe_q4_grouped_laguna_sched_ab() { eprintln!("moe_q4_grouped_laguna_sched_ab: PASS (default vs BUTTER_LAGUNA_SCHED=1 agree within 1e-3 on skewed groups)"); } +// Round 14: BUTTER_LAGUNA_MOE_NSPLIT=1 A/B correctness for the +// warp-count-doubled moe_q4_grouped_mma_nsplit redesign (see +// MOE_Q4_GROUPED_MMA_NSPLIT_SRC / laguna_moe_nsplit_on in wh-butter-ops). +// Covers uniform, ragged/uneven non-power-of-two, and zero-count-expert +// cases per the campaign's "a previously-shipped kernel bug was caught +// exactly by the zero-count-expert case" lesson. For each case: (1) a +// loose sanity check of the DEFAULT path's Q4 output against a host f32 +// reference (same tolerance style as moe_q4_grouped_laguna_sched_ab, +// confirms the Q4 pipeline itself is producing a real signal), then (2) +// the actual gate: nsplit-vs-default max relative diff on BOTH the +// host-descriptor (moe_q4_grouped_mma) and fully-on-device-descriptor +// (moe_q4_grouped_mma_dev) call paths. Both dispatch identical per-tile +// dequant+mma math (same K-reduction order within a tile, same tile -> +// output-element mapping, disjoint writes per tile) so nsplit and +// default should be near bit-exact -- 1e-3 relative is generous +// headroom, not a quantization-noise budget (mirrors sched_ab's framing). +fn run_nsplit_ab_case(d: &dyn Device, case: &str, n_exp: usize, k: usize, n: usize, groups: &[(usize, usize)]) { + let mut g_starts = vec![0usize]; + let mut expert_ids = vec![]; + for &(e, r) in groups { expert_ids.push(e); g_starts.push(g_starts.last().unwrap() + r); } + let mt = *g_starts.last().unwrap(); + + let a_f = rng(mt * k, 41); + let w_f = rng(n_exp * n * k, 43); + let up = |v: &[f32], sh: Vec| -> Tensor { + Tensor::new(d.upload(&tb(v)).unwrap(), vec![v.len()], DType::F32).reshaped(sh) }; + let a_dev = cast_f32_f16(d, &up(&a_f, vec![mt, k])).unwrap().reshaped(vec![mt, k]); + let (qs_v, sc_v) = quantize_q4(&w_f, n_exp * n, k); + let qs_b: Vec = qs_v.iter().flat_map(|x| x.to_le_bytes()).collect(); + let qs = Tensor::new(d.upload(&qs_b).unwrap(), vec![qs_v.len()], DType::U32); + let sc = cast_f32_f16(d, &up(&sc_v, vec![sc_v.len()])).unwrap(); + + // host f32 reference (loose sanity check only, see comment above) + let mut exp = vec![0f32; mt * n]; + for g in 0..groups.len() { + let (eid, _) = groups[g]; + for t in g_starts[g]..g_starts[g + 1] { + for nn in 0..n { + let mut acc = 0f32; + for kk in 0..k { acc += a_f[t * k + kk] * w_f[(eid * n + nn) * k + kk]; } + exp[t * n + nn] = acc; + } + } + } + if mt > 0 { + let sanity_cos = cos(&exp, &dl(d, &cast_f16_f32(d, &moe_q4_grouped_mma( + d, &a_dev, &qs, &sc, &g_starts, &expert_ids, n, k).unwrap()).unwrap())); + eprintln!("moe_q4_grouped_nsplit_ab[{case}]: sanity cosine (Q4 default vs f32 host ref) = {sanity_cos:.6}"); + assert!(sanity_cos > 0.99, "[{case}] moe_q4_grouped_mma default path looks broken, cosine {sanity_cos:.6}"); + } + + let rel_diff = |a: &[f32], b: &[f32]| -> f32 { + let mut maxr = 0f32; + for (x, y) in a.iter().zip(b.iter()) { + let rel = (x - y).abs() / x.abs().max(1e-2); + if rel > maxr { maxr = rel; } + } + maxr + }; + + // --- host-descriptor path (moe_q4_grouped_mma) --- + let out_default = dl(d, &cast_f16_f32(d, &moe_q4_grouped_mma( + d, &a_dev, &qs, &sc, &g_starts, &expert_ids, n, k).unwrap()).unwrap()); + unsafe { std::env::set_var("BUTTER_LAGUNA_MOE_NSPLIT", "1"); } + let out_nsplit = dl(d, &cast_f16_f32(d, &moe_q4_grouped_mma( + d, &a_dev, &qs, &sc, &g_starts, &expert_ids, n, k).unwrap()).unwrap()); + unsafe { std::env::remove_var("BUTTER_LAGUNA_MOE_NSPLIT"); } + let r_host = rel_diff(&out_default, &out_nsplit); + eprintln!("moe_q4_grouped_nsplit_ab[{case}] (host path): nsplit-vs-default max_rel={r_host:.2e}"); + assert!(r_host < 1e-3, "[{case}] host path: BUTTER_LAGUNA_MOE_NSPLIT=1 vs default max_rel {r_host:.2e} > 1e-3"); + + // --- fully-on-device descriptor path (moe_q4_grouped_mma_dev) --- + let offs: Vec = g_starts.iter().flat_map(|&s| (s as u32).to_le_bytes()).collect(); + let off = Tensor::new(d.upload(&offs).unwrap(), vec![g_starts.len()], DType::U32); + let out_dev_default = dl(d, &cast_f16_f32(d, &moe_q4_grouped_mma_dev( + d, &a_dev, &qs, &sc, &off, n_exp, mt, n, k).unwrap()).unwrap()); + unsafe { std::env::set_var("BUTTER_LAGUNA_MOE_NSPLIT", "1"); } + let out_dev_nsplit = dl(d, &cast_f16_f32(d, &moe_q4_grouped_mma_dev( + d, &a_dev, &qs, &sc, &off, n_exp, mt, n, k).unwrap()).unwrap()); + unsafe { std::env::remove_var("BUTTER_LAGUNA_MOE_NSPLIT"); } + let r_dev = rel_diff(&out_dev_default, &out_dev_nsplit); + eprintln!("moe_q4_grouped_nsplit_ab[{case}] (dev path): nsplit-vs-default max_rel={r_dev:.2e}"); + assert!(r_dev < 1e-3, "[{case}] dev path: BUTTER_LAGUNA_MOE_NSPLIT=1 vs default max_rel {r_dev:.2e} > 1e-3"); + + // dev-descriptor builder must also agree with the host-descriptor + // builder for both nsplit and default (cross-check, not the main gate). + let r_cross_default = rel_diff(&out_default, &out_dev_default); + let r_cross_nsplit = rel_diff(&out_nsplit, &out_dev_nsplit); + eprintln!("moe_q4_grouped_nsplit_ab[{case}]: host-vs-dev max_rel default={r_cross_default:.2e} nsplit={r_cross_nsplit:.2e}"); + assert!(r_cross_default < 1e-3 && r_cross_nsplit < 1e-3, "[{case}] host vs dev descriptor builders disagree"); + + eprintln!("moe_q4_grouped_nsplit_ab[{case}]: PASS (mt={mt} n_exp={n_exp} N={n} K={k})"); +} + +#[test] +fn moe_q4_grouped_nsplit_ab_uniform() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + // Uniform: several experts, identical row counts, exercises the common + // multi-full-tile (BM=64) case cleanly. + run_nsplit_ab_case(d.as_ref(), "uniform", 6, 128, 128, &[(0, 64), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64)]); +} + +#[test] +fn moe_q4_grouped_nsplit_ab_ragged() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + // Ragged/uneven, non-power-of-two: mixed group sizes (incl m=1) that + // don't divide BM=64 evenly, N/K are the real Nemotron-shape-derived + // dense proj dims (multiples of 64 per the kernel's constraint, but + // themselves not powers of two), and duplicate-adjacent small groups. + let sizes = [1usize, 2, 7, 13, 31, 64, 96, 128, 200, 17, 40, 55, 88, 3, 9, 3]; + let groups: Vec<(usize, usize)> = sizes.iter().enumerate().map(|(e, &r)| (e, r)).collect(); + run_nsplit_ab_case(d.as_ref(), "ragged", sizes.len(), 2688, 1856, &groups); +} + +#[test] +fn moe_q4_grouped_nsplit_ab_zero_count_expert() { + let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; + // Zero-count-expert: several experts selected zero times (empty + // groups, exercising the gend=0/padding-tile path), mixed with one + // big skewed group. One-group-per-expert-id (distinct sequential + // ids matching moe_q4_grouped_mma_dev's offsets[n_exp+1] contract -- + // it represents the real MoE-routing CSR layout, one contiguous + // range per expert index, so it cannot itself express a duplicate + // expert id across two separate groups; that's not a gap in this + // test, it's the on-device descriptor format's actual contract). + // "Duplicate expert id across tiles" is still exercised organically + // by every other case here whose groups exceed BM=64 rows: the host + // path's own tiling loop (`while t < e { ...eid.push(g's id)... + // t += 64 }`) tags MULTIPLE tiles with the same eid for any + // >64-row group (e.g. the ragged case's 200-row and 128-row groups, + // and this case's own 200-row group below), which is the actual + // "duplicate index" scenario this kernel encounters in production. + let groups: [(usize, usize); 8] = [(0, 200), (1, 0), (2, 3), (3, 0), (4, 37), (5, 0), (6, 3), (7, 0)]; + run_nsplit_ab_case(d.as_ref(), "zero_count_expert", 8, 128, 128, &groups); +} + #[test] fn moe_cutlass_f16_bench() { let Some(d) = CudaDevice::create().expect("cuda") else { eprintln!("no CUDA"); return; }; diff --git a/rust/crates/wh-butter-models/src/laguna.rs b/rust/crates/wh-butter-models/src/laguna.rs index 539526f1..f206cb1c 100644 --- a/rust/crates/wh-butter-models/src/laguna.rs +++ b/rust/crates/wh-butter-models/src/laguna.rs @@ -2615,8 +2615,33 @@ fn prefill_layer( Ok(x2) } +/// Size-adaptive default chunk policy for [`prefill`] (used when the +/// caller passes `chunk == 0`). Empirically measured on GB10/sm_121 +/// (laguna prefill campaign, round 4): a single monolithic chunk +/// (`chunk == total`) is fastest for `total <= 2048` -- at those sizes +/// `chunk.min(total)` already collapses to one chunk regardless, so this +/// is really just "don't force a smaller chunk than necessary" -- but +/// becomes UNSTABLE and often much SLOWER once `total` exceeds ~2048: the +/// batched per-layer activation tensors (`[T, hidden]` and friends) stop +/// fitting comfortably in cache/working-set at very large `T`, and this +/// hardware's boost clock/allocator behavior under one huge sustained +/// kernel sequence was measured to regress hard and unpredictably +/// (single monolithic n=8192 chunk measured between 191 and 364 tok/s +/// across repeated cold-start runs, vs a stable ~460-540 tok/s with a +/// fixed 2048 chunk at the same `n`). A fixed 2048-token chunk was +/// measured (isolated per-size runs, cold GPU each time) to beat BOTH +/// monolithic AND the previously-standard chunk=1024 at every size from +/// 2048 through 32768 (+1-15%, largest gains at the sizes where +/// monolithic was worst). `min(total, 2048)` therefore covers both +/// regimes with one formula: it IS monolithic chunking for `total <= +/// 2048` and a fixed 2048 chunk above that. +pub fn default_prefill_chunk(total: usize) -> usize { + total.min(2048).max(1) +} + /// Batched (multi-token) **prefill**: process the whole `tokens` prompt -/// `chunk` tokens at a time (default recommendation: 1024), writing every +/// `chunk` tokens at a time (pass 0 for the adaptive default, see +/// [`default_prefill_chunk`]), writing every /// token's K/V into `kv`/`scratch`, and return the LAST token's next-token /// logits, the prefill counterpart of [`decode_step`]'s per-token forward. /// Each chunk's `T` rows run through every layer as batched GEMMs @@ -2637,7 +2662,11 @@ fn prefill_layer( /// /// `chunk` only bounds how many prompt tokens are embedded/projected/attended /// in one forward pass (memory for the batched intermediates scales with -/// it); correctness is independent of the choice. +/// it); correctness is independent of the choice. **`chunk == 0` is a +/// sentinel for "auto"**: picks [`default_prefill_chunk`]'s adaptive policy +/// instead of a caller-fixed size. Any nonzero value is used verbatim +/// (existing explicit-override callers, e.g. `BUTTER_LAGUNA_PREFILL_CHUNK`, +/// are unaffected). pub fn prefill( dev: &dyn Device, model: &LagunaModel, @@ -2656,7 +2685,7 @@ pub fn prefill( "prefill: prompt len {total} exceeds scratch prompt_cap {}", scratch.prompt_cap ))); } - let chunk = chunk.max(1); + let chunk = if chunk == 0 { default_prefill_chunk(total) } else { chunk.max(1) }; // BUTTER_LAGUNA_MARLIN=1: ONE persistent all-zero `u32` indices buffer, // sized to `chunk` (every chunk's row count `t = chunk.min(total-start)` diff --git a/rust/crates/wh-butter-modeltests/src/lib.rs b/rust/crates/wh-butter-modeltests/src/lib.rs index ccca5304..644bdb77 100644 --- a/rust/crates/wh-butter-modeltests/src/lib.rs +++ b/rust/crates/wh-butter-modeltests/src/lib.rs @@ -588,20 +588,27 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { // BUTTER_LAGUNA_PP=: standalone batched-prefill throughput bench over a // synthetic n-token prompt (token id sequence i % 50000, any valid ids // exercise the same code path; this is a speed bench, not a semantic - // eval). BUTTER_LAGUNA_PREFILL_CHUNK overrides the chunk size (default - // 1024). Prints "pp tok/s" and returns, no generation, no decode. + // eval). BUTTER_LAGUNA_PREFILL_CHUNK overrides the chunk size; default is + // 0 == "auto" (laguna::default_prefill_chunk(n), see that fn's doc: a + // size-adaptive min(n, 2048) policy, measured round-4 to beat both + // monolithic and the old fixed-1024 convention at every size). Prints + // "pp tok/s" and returns, no generation, no decode. if let Ok(pp_s) = std::env::var("BUTTER_LAGUNA_PP") { // Comma-separated sizes bench the whole ladder on ONE model load // (the 71G requant load dominates a cycle). let sizes: Vec = pp_s.split(',').filter_map(|v| v.trim().parse().ok()).collect(); // BUTTER_LAGUNA_PREFILL_CHUNK also accepts a comma list: the harness // benches every (size, chunk) combo on the single model load. + // Default is `[0]` (auto): `laguna::prefill` resolves 0 into its + // own size-adaptive default per call, so different `n` in the same + // comma-separated `sizes` list each get their own right-sized + // chunk instead of one chunk value forced across all of them. let chunks: Vec = std::env::var("BUTTER_LAGUNA_PREFILL_CHUNK") .ok().map(|v| v.split(',').filter_map(|c| c.trim().parse().ok()).collect()) .filter(|v: &Vec| !v.is_empty()) - .unwrap_or_else(|| vec![1024]); + .unwrap_or_else(|| vec![0]); for &(mut chunk) in chunks.iter() { - chunk = chunk.max(1); + if chunk != 0 { chunk = chunk.max(1); } for &n in sizes.iter().filter(|&&n| n > 0) { let synth_tokens: Vec = (0..n).map(|i| (i % 50000) as u32).collect(); let Ok(kv_pp) = laguna::LagunaKvCache::new(d, &model.cfg, n + 1) else { @@ -635,8 +642,16 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { let dt = t0.elapsed().as_secs_f64(); let tps = n as f64 / dt; let mut lb = vec![0u8; model.cfg.vocab * 4]; - let _ = d.download(logits.buffer.as_ref(), &mut lb); // liveness check on the returned tensor - eprintln!("Laguna prefill bench (n={n}, chunk={chunk}, {plat}): {tps:.2} pp tok/s"); + // Also serves as a cross-chunk-size correctness cross-check: with the + // synthetic (i % 50000) token sequence this is deterministic, so running + // the SAME n with different chunk values (e.g. auto vs an explicit + // override) and comparing this argmax is a quick chunk-invariance sanity + // check without needing the separate BUTTER_LAGUNA_PREFILL tokenizer oracle. + let _ = d.download(logits.buffer.as_ref(), &mut lb); + let lf: Vec = lb.chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + let argmax = wh_butter_runtime::argmax(&lf); + let eff_chunk = if chunk == 0 { laguna::default_prefill_chunk(n) } else { chunk }; + eprintln!("Laguna prefill bench (n={n}, chunk={chunk} eff={eff_chunk}, {plat}): {tps:.2} pp tok/s, last-token argmax={argmax}"); laguna::prof_report(); } Err(e) => eprintln!("Laguna PP bench: timed prefill failed: {e:?}"), @@ -774,10 +789,15 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) { // correctness gate for the batched-prefill work: the greedy continuation // must come out byte-identical to the decode-only path (same prompt, // same seed token handed to decode at `pos = tokens.len()`). - // BUTTER_LAGUNA_PREFILL_CHUNK overrides the chunk size (default 1024). + // BUTTER_LAGUNA_PREFILL_CHUNK overrides the chunk size; default is 0 == + // "auto" (laguna::default_prefill_chunk, see its doc -- a size-adaptive + // min(n, 2048) policy). Correctness is chunk-invariant by construction + // (see `prefill`'s doc), so this default change carries no correctness + // risk; it only changes which chunk size the oracle happens to exercise + // when the caller doesn't pin one. let prefill_mode = std::env::var("BUTTER_LAGUNA_PREFILL").map(|v| v == "1").unwrap_or(false); let prefill_chunk: usize = std::env::var("BUTTER_LAGUNA_PREFILL_CHUNK") - .ok().and_then(|v| v.parse().ok()).unwrap_or(2048); + .ok().and_then(|v| v.parse().ok()).unwrap_or(0); let start_pos = if prefill_mode { let Ok(scratch) = laguna::LagunaPrefillScratch::new(d, &model.cfg, tokens.len()) else { eprintln!("Laguna: prefill scratch alloc failed, skipping"); diff --git a/rust/crates/wh-butter-ops/src/lib.rs b/rust/crates/wh-butter-ops/src/lib.rs index 24887eca..7487d8b9 100644 --- a/rust/crates/wh-butter-ops/src/lib.rs +++ b/rust/crates/wh-butter-ops/src/lib.rs @@ -1358,13 +1358,15 @@ pub fn argmax_f32_device(dev: &dyn Device, x: &Tensor) -> Result { /// The SwiGLU MoE path calls this for the gate and up stacks, then applies /// the activation elementwise before the down gather. F16 per-block scales. /// -/// TODO(rebrand-merge): blocked — dispatches `iron_moe_gather_q4` via -/// `wh_iron_std`'s kernel IR, but the pinned `thewafflehaus/iron@dev` kernel -/// library only ships `iron_moe_gather_q4_relu2/_down/_down_accum` — the -/// plain (non-relu2) variant this needs is not present upstream yet. Kept as -/// an always-error stub until that kernel lands in the iron repo. -pub fn moe_gather_q4(_dev: &dyn Device, _qs: &Tensor, _scales: &Tensor, _x: &Tensor, _idx: &Tensor, _top_k: usize, _inter: usize, _hid: usize) -> Result { - Err(Error::Msg("moe_gather_q4: blocked — iron_moe_gather_q4 kernel not yet in thewafflehaus/iron@dev (only _relu2/_down/_down_accum variants exist)".into())) +pub fn moe_gather_q4(dev: &dyn Device, qs: &Tensor, scales: &Tensor, x: &Tensor, idx: &Tensor, top_k: usize, inter: usize, hid: usize) -> Result { + let k = cached_ir("iron_moe_gather_q4", x.dtype, || { let mut k = wh_iron_std::kernels::moe::moe_gather_q4::iron_moe_gather_q4::kernel_ir_for(x.dtype); k.mode = wh_iron_core::ir::KernelMode::Reduction; k }); + let out = Tensor::empty(dev, vec![top_k * inter], x.dtype)?; + let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + // Default rpt=2: MoE-gather kernels are big + latency-bound; 2 warps/row hides + // global-load latency (matches the _relu2/_down siblings' default). + let rpt: u32 = std::env::var("BUTTER_MOE_RPT").ok().and_then(|v| v.parse().ok()).filter(|&r| r >= 1).unwrap_or(2); + dev.dispatch(&k, &[Binding::Buffer(qs.buffer.clone()), Binding::Buffer(scales.buffer.clone()), Binding::Buffer(x.buffer.clone()), Binding::Buffer(idx.buffer.clone()), Binding::Buffer(out.buffer.clone()), u(hid as u32), u(inter as u32), u(rpt)], Grid { grid: [(inter as u32).div_ceil(rpt), top_k as u32, 1], block: [32 * rpt, 1, 1] })?; + Ok(out) } /// Batched MoE GATE+UP gather with FUSED inline SwiGLU: reads the SAME @@ -1374,11 +1376,13 @@ pub fn moe_gather_q4(_dev: &dyn Device, _qs: &Tensor, _scales: &Tensor, _x: &Ten /// up stack) plus a host-side elementwise SwiGLU pass with one dispatch. /// F16 per-block scales on both `gate_scales` and `up_scales`. #[allow(clippy::too_many_arguments)] -/// TODO(rebrand-merge): blocked — same iron-kernel gap as `moe_gather_q4` -/// above (`iron_moe_gather_q4_swiglu` not in `thewafflehaus/iron@dev` yet). -#[allow(clippy::too_many_arguments)] -pub fn moe_gather_q4_swiglu(_dev: &dyn Device, _gate_qs: &Tensor, _gate_scales: &Tensor, _up_qs: &Tensor, _up_scales: &Tensor, _x: &Tensor, _idx: &Tensor, _top_k: usize, _inter: usize, _hid: usize) -> Result { - Err(Error::Msg("moe_gather_q4_swiglu: blocked — iron_moe_gather_q4_swiglu kernel not yet in thewafflehaus/iron@dev".into())) +pub fn moe_gather_q4_swiglu(dev: &dyn Device, gate_qs: &Tensor, gate_scales: &Tensor, up_qs: &Tensor, up_scales: &Tensor, x: &Tensor, idx: &Tensor, top_k: usize, inter: usize, hid: usize) -> Result { + let k = cached_ir("iron_moe_gather_q4_swiglu", x.dtype, || { let mut k = wh_iron_std::kernels::moe::moe_gather_q4::iron_moe_gather_q4_swiglu::kernel_ir_for(x.dtype); k.mode = wh_iron_core::ir::KernelMode::Reduction; k }); + let out = Tensor::empty(dev, vec![top_k * inter], x.dtype)?; + let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + let rpt: u32 = std::env::var("BUTTER_MOE_RPT").ok().and_then(|v| v.parse().ok()).filter(|&r| r >= 1).unwrap_or(2); + dev.dispatch(&k, &[Binding::Buffer(gate_qs.buffer.clone()), Binding::Buffer(gate_scales.buffer.clone()), Binding::Buffer(up_qs.buffer.clone()), Binding::Buffer(up_scales.buffer.clone()), Binding::Buffer(x.buffer.clone()), Binding::Buffer(idx.buffer.clone()), Binding::Buffer(out.buffer.clone()), u(hid as u32), u(inter as u32), u(rpt)], Grid { grid: [(inter as u32).div_ceil(rpt), top_k as u32, 1], block: [32 * rpt, 1, 1] })?; + Ok(out) } /// Batched MoE down + router-weighted accumulate into `acc[hid]`. `qs` is the @@ -2626,7 +2630,7 @@ extern "C" __global__ void moe_q4_grouped_mma( __half* __restrict__ Out, int N, int K, int band_n) { - const int BM=64,BN=128,BK=32,NS=16,NKS=2,GS=32,WPR=BK/8,GPR=BK/GS; // BK32 occ-tune (smem 25->12.5KB) + const int BM=64,BN=128,BK=32,NS=16,NKS=2,GS=32,WPR=BK/8,GPR=BK/GS,STAGES=3; // BK32 occ-tune (smem 25->12.5KB); STAGES=3 cp.async pipeline depth (deeper latency-hiding, occupancy-neutral -- see moe_q4_grouped_mma's doc) // band_n=0 (default): grid=[tile,nblk] with tile fast-varying -> a full // tile sweep (crossing experts) happens before the N-block advances. // band_n=1 (BUTTER_LAGUNA_SCHED=1): grid=[nblk,tile] with nblk fast-varying -> @@ -2638,9 +2642,9 @@ extern "C" __global__ void moe_q4_grouped_mma( int by = band_n ? blockIdx.x : blockIdx.y; int tok0=tok0_arr[bx], gend=gend_arr[bx], eid=eid_arr[bx], n0=by*BN; int tid=threadIdx.x,warp=tid>>5,lane=tid&31,gid=lane>>2,tid4=lane&3,wrow=warp*16; - __shared__ __half As[2][BM][BK]; - __shared__ unsigned Wqs[2][BN][WPR]; - __shared__ __half Scs[2][BN][GPR]; + __shared__ __half As[3][BM][BK]; + __shared__ unsigned Wqs[3][BN][WPR]; + __shared__ __half Scs[3][BN][GPR]; float oc[NS][4]; for(int s=0;s> 3) & 1) * 8; int acol = (lane >> 4) * 8; @@ -2689,6 +2706,132 @@ extern "C" __global__ void moe_q4_grouped_mma( } "#; +// F-round-14: warp-count-doubled variant of moe_q4_grouped_mma. Round 13 +// scoped this via standalone nvcc (zero GPU cost): splitting the NS=16 +// N-slots across an extra warp axis (4 warps->8 warps, 128->256 threads) +// halves the per-thread f32 accumulator (oc[16][4]=64 regs -> oc[8][4]=32 +// regs) while leaving BM/BN/the A-tile identical (so A-tile reuse and +// arithmetic intensity are UNCHANGED -- this is NOT the closed round-9/12 +// "BN/NS shrink", which shrunk the output tile itself and lost -9.6%/-6.5% +// from doubled A rereads). Under `__launch_bounds__(256,4)` this compiles +// to 64 registers with ZERO spill (verified via `nvcc -Xptxas -v`, +// ~/laguna-bench/variant_nsplit_lb4.cu), raising the reg-limited occupancy +// ceiling from 5 blocks/SM * 4 warps = 20 warps/SM (41.7%) to 4 blocks/SM * +// 8 warps = 32 warps/SM (66.7%) -- smem (19200B/block, unchanged) still +// allows 5 blocks/SM so the ceiling here is register-, not smem-, bound. +// `__launch_bounds__(256,5)` (targeting the smem-implied 5th block) SPILLS +// (48 regs but 72B stack + 96B spill-store + 88B spill-load per round 13), +// so (256,4)/64-reg/66.7% is the practical ceiling for this redesign. +// +// warp_m = warp>>1 in [0,4) selects the M-row block (same meaning as the +// baseline's `wrow=warp*16`); warp_n = warp&1 in [0,2) selects which half +// of the 16 N-slots (s in [warp_n*8, warp_n*8+8)) this warp accumulates. +// Both warp_n halves of a given warp_m redundantly ldmatrix the SAME +// shared As[] tile (cheap shared-mem reads, not extra global traffic) so +// the block's total output tile (BM=64 x BN=128) and A-tile reuse are +// unchanged vs the baseline. The LOADC cooperative-load distribution loops +// are pure strided covers (`for e=tid; e +__device__ __forceinline__ unsigned dq2(unsigned w, int sh, __half sc){ + unsigned two = (w >> sh) & 0xFFu; + unsigned packed = ((((two & 0xf0u) << 12) | (two & 0xfu)) ^ 0x00080008u) | 0x64006400u; + __half2 h = *reinterpret_cast<__half2*>(&packed); + h = __hmul2(__hsub2(h, __float2half2_rn(1032.0f)), __half2half2(sc)); + return *reinterpret_cast(&h); +} +#define CPA(SD,SP) asm volatile("cp.async.cg.shared.global [%0],[%1],16,16;\n" :: "r"(SD), "l"(SP)) +extern "C" __global__ void __launch_bounds__(256,4) moe_q4_grouped_mma_nsplit( + const __half* __restrict__ A, + const unsigned* __restrict__ Qs, + const __half* __restrict__ Sc, + const int* __restrict__ tok0_arr, const int* __restrict__ eid_arr, const int* __restrict__ gend_arr, + __half* __restrict__ Out, + int N, int K, int band_n) +{ + const int BM=64,BN=128,BK=32,NS_H=8,NKS=2,GS=32,WPR=BK/8,GPR=BK/GS,STAGES=3; + const int NTHREADS=256; + int bx = band_n ? blockIdx.y : blockIdx.x; + int by = band_n ? blockIdx.x : blockIdx.y; + int tok0=tok0_arr[bx], gend=gend_arr[bx], eid=eid_arr[bx], n0=by*BN; + int tid=threadIdx.x,warp=tid>>5,lane=tid&31,gid=lane>>2,tid4=lane&3; + int warp_m = warp >> 1; // 0..3 : M-row block (same meaning as baseline's warp*16) + int warp_n = warp & 1; // 0..1 : which half of the 16 N-slots + int wrow = warp_m*16; + int s_lo = warp_n*NS_H; + __shared__ __half As[3][BM][BK]; + __shared__ unsigned Wqs[3][BN][WPR]; + __shared__ __half Scs[3][BN][GPR]; + float oc[NS_H][4]; for(int s=0;s> 3) & 1) * 8; + int acol = (lane >> 4) * 8; + unsigned aaddr = (unsigned)__cvta_generic_to_shared(&As[buf][wrow+arow][ko+acol]); + unsigned a0,a1,a2,a3; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(a0),"=r"(a1),"=r"(a2),"=r"(a3) : "r"(aaddr)); + #pragma unroll + for(int s=0;s bool { + std::env::var("BUTTER_LAGUNA_MOE_NSPLIT").as_deref() == Ok("1") +} + // BUTTER_LAGUNA_SCHED=1: A/B switch for the two Laguna prefill CTA-scheduling // levers on the Q4 grouped GEMM (see moe_q4_grouped_mma / _dev below). Default // off = byte-identical launch geometry to before this change. @@ -2758,13 +2901,19 @@ pub fn moe_q4_grouped_mma( } else { [n_mtiles as u32, (n as u32).div_ceil(128), 1] }; + let nsplit = laguna_moe_nsplit_on(); + let (moe_src, moe_kname, moe_block): (&str, &str, [u32; 3]) = if nsplit { + (MOE_Q4_GROUPED_MMA_NSPLIT_SRC, "moe_q4_grouped_mma_nsplit", [256, 1, 1]) + } else { + (MOE_Q4_GROUPED_MMA_SRC, "moe_q4_grouped_mma", [128, 1, 1]) + }; dev.dispatch_raw_cuda( - MOE_Q4_GROUPED_MMA_SRC, "moe_q4_grouped_mma.cu", "moe_q4_grouped_mma", + moe_src, "moe_q4_grouped_mma.cu", moe_kname, &[(a.buffer.as_ref(), a.offset), (qs.buffer.as_ref(), qs.offset), (sc.buffer.as_ref(), sc.offset), (t0d.buffer.as_ref(), 0), (eidd.buffer.as_ref(), 0), (ged.buffer.as_ref(), 0), (out.buffer.as_ref(), 0)], &[i(n as i32), i(k as i32), i(band_n as i32)], - grid, [128, 1, 1], 0, false)?; + grid, moe_block, 0, false)?; Ok(out) } @@ -2790,10 +2939,42 @@ extern "C" __global__ void moe_build_tiles( int* order = smem + n_exp; // [n_exp] if (desc) { for (int e = 0; e < n_exp; e++) { cnt[e] = (int)(offsets[e+1] - offsets[e]); order[e] = e; } - for (int i2 = 0; i2 < n_exp; i2++) { - int best = i2; - for (int j = i2 + 1; j < n_exp; j++) if (cnt[order[j]] > cnt[order[best]]) best = j; - if (best != i2) { int tmp = order[i2]; order[i2] = order[best]; order[best] = tmp; } + // round7: O(n log n) heapsort replacing the prior O(n^2) selection sort + // (moe_build_tiles measured ~1.34ms/call via ncu at n_exp=256, single + // thread -- the O(n^2)=65536-compare sort was the dominant cost). + // Correctness is insensitive to exact tie-break order here: any + // permutation of `order[]` yields identical GEMM output since tiles + // write disjoint output regions regardless of emission order (see + // moe_q4_grouped_laguna_sched_ab, which A/B-tests exactly this + // desc-mode path against the default order on skewed/zero-count + // groups and requires bit-exact agreement). Build a max-heap over + // order[] keyed by cnt[order[i]], then extract-max repeatedly + // (ascending), then reverse in place for descending final order. + for (int start = n_exp / 2 - 1; start >= 0; start--) { + int root = start; + while (true) { + int child = 2 * root + 1; + if (child >= n_exp) break; + if (child + 1 < n_exp && cnt[order[child + 1]] > cnt[order[child]]) child++; + if (cnt[order[root]] >= cnt[order[child]]) break; + int tmp = order[root]; order[root] = order[child]; order[child] = tmp; + root = child; + } + } + for (int end = n_exp - 1; end > 0; end--) { + int tmp = order[0]; order[0] = order[end]; order[end] = tmp; + int root = 0; + while (true) { + int child = 2 * root + 1; + if (child >= end) break; + if (child + 1 < end && cnt[order[child + 1]] > cnt[order[child]]) child++; + if (cnt[order[root]] >= cnt[order[child]]) break; + tmp = order[root]; order[root] = order[child]; order[child] = tmp; + root = child; + } + } + for (int i = 0, j = n_exp - 1; i < j; i++, j--) { + int tmp = order[i]; order[i] = order[j]; order[j] = tmp; } } int t = 0; @@ -2853,13 +3034,19 @@ pub fn moe_q4_grouped_mma_dev( } else { [maxt as u32, (n as u32).div_ceil(128), 1] }; + let nsplit = laguna_moe_nsplit_on(); + let (moe_src, moe_kname, moe_block): (&str, &str, [u32; 3]) = if nsplit { + (MOE_Q4_GROUPED_MMA_NSPLIT_SRC, "moe_q4_grouped_mma_nsplit", [256, 1, 1]) + } else { + (MOE_Q4_GROUPED_MMA_SRC, "moe_q4_grouped_mma", [128, 1, 1]) + }; dev.dispatch_raw_cuda( - MOE_Q4_GROUPED_MMA_SRC, "moe_q4_grouped_mma.cu", "moe_q4_grouped_mma", + moe_src, "moe_q4_grouped_mma.cu", moe_kname, &[(a.buffer.as_ref(), a.offset), (qs.buffer.as_ref(), qs.offset), (sc.buffer.as_ref(), sc.offset), (t0d.buffer.as_ref(), 0), (eidd.buffer.as_ref(), 0), (ged.buffer.as_ref(), 0), (out.buffer.as_ref(), 0)], &[i(n as i32), i(k as i32), i(band_n as i32)], - grid, [128, 1, 1], 0, false)?; + grid, moe_block, 0, false)?; Ok(out) } @@ -2893,17 +3080,17 @@ const SDPA_TC_KPREP_SRC: &str = r#" #include extern "C" __global__ void sdpa_tc_kprep( const void* __restrict__ k_in, // [nkv, kv_stride, hd] - __half* __restrict__ k_out, // [nkv, n_kv, hd] f16 - int nkv, int kv_stride, int n_kv, int hd, int in_f32) + __half* __restrict__ k_out, // [nkv, blk, hd] f16, THIS BLOCK ONLY + int nkv, int kv_stride, int kb0, int blk, int hd, int in_f32) { long i = (long)blockIdx.x * blockDim.x + threadIdx.x; - long total = (long)nkv * n_kv * hd; + long total = (long)nkv * blk * hd; if (i >= total) return; int d = (int)(i % hd); long t = i / hd; - int p = (int)(t % n_kv); // kv position - int kh = (int)(t / n_kv); // kv head - long si = ((long)kh * kv_stride + p) * hd + d; + int p = (int)(t % blk); // position within this block + int kh = (int)(t / blk); // kv head + long si = ((long)kh * kv_stride + (kb0 + p)) * hd + d; float val = in_f32 ? ((const float*)k_in)[si] : __half2float(((const __half*)k_in)[si]); k_out[i] = __float2half(val); @@ -4148,7 +4335,6 @@ pub fn sdpa_multi_tc( // ── prep buffers (all f16) ────────────────────────────────────────────── let qh = Tensor::empty(dev, vec![nq, sq, hd], DType::F16)?; // [nq,Sq,hd] scaled - let kh = Tensor::empty(dev, vec![nkv, n_kv, hd], DType::F16)?; // [nkv,n_kv,hd] let i = |x: i32| x.to_le_bytes().to_vec(); let f = |x: f32| x.to_le_bytes().to_vec(); @@ -4158,10 +4344,6 @@ pub fn sdpa_multi_tc( &[(q.buffer.as_ref(), q.offset), (qh.buffer.as_ref(), 0)], &[i(sq as i32), i(nq as i32), i(hd as i32), f(scale), i(in_f32)], blk256(nq*sq*hd), [256,1,1], 0, false)?; - dev.dispatch_raw_cuda(SDPA_TC_KPREP_SRC, "sdpa_tc_kprep.cu", "sdpa_tc_kprep", - &[(k.buffer.as_ref(), k.offset), (kh.buffer.as_ref(), 0)], - &[i(nkv as i32), i(kv_stride as i32), i(n_kv as i32), i(hd as i32), i(in_f32)], - blk256(nkv*n_kv*hd), [256,1,1], 0, false)?; // ── running FlashAttention state (init via upload: o=0, l=0, m=-inf) ───── let o_run = Tensor::new(dev.alloc_zeroed(nq*sq*hd*4)?, vec![nq, sq, hd], DType::F32); @@ -4178,17 +4360,26 @@ pub fn sdpa_multi_tc( let bl = Tensor::empty(dev, vec![nq, sq], DType::F32)?; let o_blk = Tensor::empty(dev, vec![nq, sq, hd], DType::F16)?; let vt = Tensor::empty(dev, vec![nkv, hd, bk], DType::F16)?; // per-block V^T + let kh = Tensor::empty(dev, vec![nkv, bk, hd], DType::F16)?; // per-block K, f16 (was whole-prefix) let el = 2i64; // f16 bytes let mut kb0 = 0usize; while kb0 < n_kv { let blk = bk.min(n_kv - kb0); + // ── K prep: cast+copy THIS KV block only → kh[nkv,blk,hd] ─────────── + // (was one whole-causal-prefix conversion before the loop; see the + // PERF FIX note in `sdpa_multi_tc_varlen` below for the rationale.) + dev.dispatch_raw_cuda(SDPA_TC_KPREP_SRC, "sdpa_tc_kprep.cu", "sdpa_tc_kprep", + &[(k.buffer.as_ref(), k.offset), (kh.buffer.as_ref(), 0)], + &[i(nkv as i32), i(kv_stride as i32), i(kb0 as i32), i(blk as i32), i(hd as i32), i(in_f32)], + blk256(nkv*blk*hd), [256,1,1], 0, false)?; + // ── QKᵀ: per KV group (16 q-heads share one KV head) ──────────────── // C[Sq, blk] = Qh[Sq,hd] · Kh_blk[blk,hd]ᵀ ; batch over the hpg q-heads. for g in 0..nkv { let q_off = g * hpg * sq * hd * 2; // bytes into qh - let k_off = (g * n_kv + kb0) * hd * 2; // bytes into kh (this block) + let k_off = g * blk * hd * 2; // bytes into kh (buffer-relative now) let s_off = g * hpg * sq * blk * 2; // bytes into scores dev.gemm_strided_batched_off( qh.buffer.as_ref(), q_off, (sq*hd) as i64 * el, // X stride = one q-head @@ -4304,7 +4495,6 @@ pub fn sdpa_multi_tc_varlen( other => return Err(Error::Msg(format!("sdpa_multi_tc_varlen: unsupported dtype {other:?}"))) }; let qh = Tensor::empty(dev, vec![nq, sq, hd], DType::F16)?; - let kh = Tensor::empty(dev, vec![nkv, n_kv, hd], DType::F16)?; let i = |x: i32| x.to_le_bytes().to_vec(); let f = |x: f32| x.to_le_bytes().to_vec(); @@ -4314,23 +4504,58 @@ pub fn sdpa_multi_tc_varlen( &[(q.buffer.as_ref(), q.offset), (qh.buffer.as_ref(), 0)], &[i(sq as i32), i(nq as i32), i(hd as i32), f(scale), i(in_f32)], blk256(nq*sq*hd), [256,1,1], 0, false)?; - dev.dispatch_raw_cuda(SDPA_TC_KPREP_SRC, "sdpa_tc_kprep.cu", "sdpa_tc_kprep", - &[(k.buffer.as_ref(), k.offset), (kh.buffer.as_ref(), 0)], - &[i(nkv as i32), i(kv_stride as i32), i(n_kv as i32), i(hd as i32), i(in_f32)], - blk256(nkv*n_kv*hd), [256,1,1], 0, false)?; let o_run = Tensor::new(dev.alloc_zeroed(nq*sq*hd*4)?, vec![nq, sq, hd], DType::F32); let l_run = Tensor::new(dev.alloc_zeroed(nq*sq*4)?, vec![nq, sq], DType::F32); let neg: Vec = (0..nq*sq).flat_map(|_| (-3.4e38f32).to_le_bytes()).collect(); let m_run = Tensor::new(dev.upload(&neg)?, vec![nq, sq], DType::F32); - let bk = 2048usize.min(n_kv); + // KV block size. Full-attention layers (win==0) keep 2048 (bounds the + // scores buffer, see sibling `sdpa_multi_tc`'s doc). Sliding-window + // layers (win!=0): the block-skip above only omits blocks with + // q_cnt==0 ENTIRELY -- a block that IS touched still pays a full + // bk-wide QK^T/PV GEMM + softmax/merge even though the window (512 for + // Laguna, far smaller than 2048) only needs a fraction of that width. + // Shrinking bk for win!=0 tightens the per-block reachable-range union + // (`[kb0, kb0+blk+win)`, see the q_lo/q_cnt derivation below) at the + // cost of more loop iterations (more kprep/QK/softmax/merge launches). + // BUTTER_LAGUNA_SDPA_WIN_BK overrides the windowed-path block size for + // A/B sweeps. Default 1024, NOT the 512 window width itself: measured + // bk=512 vs bk=1024 vs bk=2048(orig) on Laguna-S-2.1/GB10 (chunk=1024, + // MARLIN+MOEFUSE) -- bk=512 regresses n=1024 prefill -4% (extra + // loop/launch overhead from doubling block count outweighs the reduced + // overcompute there), while bk=1024 nets +1-1.5% at n=8192/32768 with + // no regression anywhere tested. Full-attention layers (win==0) are + // never affected by this env var. + let win_bk: usize = std::env::var("BUTTER_LAGUNA_SDPA_WIN_BK") + .ok().and_then(|v| v.parse().ok()).filter(|&b: &usize| b >= 1).unwrap_or(1024); + let bk = if win != 0 { win_bk.min(n_kv) } else { 2048usize.min(n_kv) }; let scores = Tensor::empty(dev, vec![nq, sq, bk], DType::F16)?; let p_blk = Tensor::empty(dev, vec![nq, sq, bk], DType::F16)?; let bm = Tensor::empty(dev, vec![nq, sq], DType::F32)?; let bl = Tensor::empty(dev, vec![nq, sq], DType::F32)?; let o_blk = Tensor::empty(dev, vec![nq, sq, hd], DType::F16)?; let vt = Tensor::empty(dev, vec![nkv, hd, bk], DType::F16)?; + // PERF FIX (was `[nkv, n_kv, hd]`, one whole-causal-prefix f32->f16 + // conversion upfront every call): for sliding-window layers (`win != 0`) + // the vast majority of KV blocks have `q_cnt == 0` below (no row in this + // chunk can reach them) -- the QKᵀ/PV GEMMs were already correctly + // skipped for those blocks, but `kh` was still sized (and fully + // converted) to the ENTIRE causal prefix on every prefill chunk, and the + // softmax/merge dispatches ran unconditionally per block regardless of + // `q_cnt`. For a windowed layer deep into a long prefill (`base` large) + // that is an O(base) cost paid by EVERY chunk on every windowed layer + // (Laguna: 36 of 48 layers are sliding-window) -- the root cause of this + // engine's prefill throughput decaying with sequence length instead of + // staying roughly flat. Fix: convert K per-block into a `bk`-sized + // rolling buffer (mirroring `vt`/`SDPA_TC_VPREP_SRC`, which already + // worked this way), and gate kprep+softmax+merge on the same `q_cnt>0` + // check as the GEMMs -- a block no row can reach contributes nothing + // either way (softmax would derive `bm=-inf`/`bl=0` for every row via + // its own positional mask; `merge`'s `l_b <= 0` early-return already + // no-ops on that combination), so skipping all four dispatches together + // is bit-identical to running them, just without the wasted work. + let kh = Tensor::empty(dev, vec![nkv, bk, hd], DType::F16)?; let el = 2i64; let mut kb0 = 0usize; @@ -4372,15 +4597,18 @@ pub fn sdpa_multi_tc_varlen( }; // q_cnt == 0: no row in this chunk can reach this block at all (the // derivation above is a superset of the true reachable set, so an - // empty computed range proves the true set is also empty). Skip the - // QKᵀ/PV GEMMs and vprep entirely; softmax/merge stay unconditional - // below (softmax re-derives bm/bl fresh every iteration straight - // from the mask, and merge's `l_b <= 0` early-return means it never - // reads the skipped o_blk/vt buffers, so this is correctness-safe). + // empty computed range proves the true set is also empty). Skip + // kprep, the QKᵀ/PV GEMMs, vprep, softmax AND merge entirely (all + // four moved inside this same guard -- see the PERF FIX note above + // `kh`'s allocation for why this is bit-identical, not just faster). if q_cnt > 0 { + dev.dispatch_raw_cuda(SDPA_TC_KPREP_SRC, "sdpa_tc_kprep.cu", "sdpa_tc_kprep", + &[(k.buffer.as_ref(), k.offset), (kh.buffer.as_ref(), 0)], + &[i(nkv as i32), i(kv_stride as i32), i(kb0 as i32), i(blk as i32), i(hd as i32), i(in_f32)], + blk256(nkv*blk*hd), [256,1,1], 0, false)?; for g in 0..nkv { let q_off = g * hpg * sq * hd * 2 + q_lo * hd * 2; - let k_off = (g * n_kv + kb0) * hd * 2; + let k_off = g * blk * hd * 2; let s_off = g * hpg * sq * blk * 2 + q_lo * blk * 2; dev.gemm_strided_batched_off( qh.buffer.as_ref(), q_off, (sq*hd) as i64 * el, @@ -4388,15 +4616,13 @@ pub fn sdpa_multi_tc_varlen( scores.buffer.as_ref(), s_off, (sq*blk) as i64 * el, q_cnt, blk, hd, hpg, DType::F16)?; } - } - // varlen softmax: causal upper (base+r) AND segment lower (seg_lo[r]). - dev.dispatch_raw_cuda(SDPA_TC_SOFTMAX_VARLEN_SRC, "sdpa_tc_softmax_varlen.cu", "sdpa_tc_softmax_varlen", - &[(scores.buffer.as_ref(), 0), (p_blk.buffer.as_ref(), 0), - (bm.buffer.as_ref(), 0), (bl.buffer.as_ref(), 0), - (seg_lo.buffer.as_ref(), seg_lo.offset)], - &[i(nq as i32), i(sq as i32), i(blk as i32), i(kb0 as i32), i(base as i32)], - [(nq*sq) as u32, 1, 1], [256,1,1], 0, false)?; - if q_cnt > 0 { + // varlen softmax: causal upper (base+r) AND segment lower (seg_lo[r]). + dev.dispatch_raw_cuda(SDPA_TC_SOFTMAX_VARLEN_SRC, "sdpa_tc_softmax_varlen.cu", "sdpa_tc_softmax_varlen", + &[(scores.buffer.as_ref(), 0), (p_blk.buffer.as_ref(), 0), + (bm.buffer.as_ref(), 0), (bl.buffer.as_ref(), 0), + (seg_lo.buffer.as_ref(), seg_lo.offset)], + &[i(nq as i32), i(sq as i32), i(blk as i32), i(kb0 as i32), i(base as i32)], + [(nq*sq) as u32, 1, 1], [256,1,1], 0, false)?; dev.dispatch_raw_cuda(SDPA_TC_VPREP_SRC, "sdpa_tc_vprep.cu", "sdpa_tc_vprep", &[(v.buffer.as_ref(), v.offset), (vt.buffer.as_ref(), 0)], &[i(nkv as i32), i(kv_stride as i32), i(hd as i32), i(kb0 as i32), i(blk as i32), i(in_f32)], @@ -4411,12 +4637,12 @@ pub fn sdpa_multi_tc_varlen( o_blk.buffer.as_ref(), o_off, (sq*hd) as i64 * el, q_cnt, hd, blk, hpg, DType::F16)?; } + dev.dispatch_raw_cuda(SDPA_TC_MERGE_SRC, "sdpa_tc_merge.cu", "sdpa_tc_merge", + &[(o_blk.buffer.as_ref(), 0), (bm.buffer.as_ref(), 0), (bl.buffer.as_ref(), 0), + (o_run.buffer.as_ref(), 0), (m_run.buffer.as_ref(), 0), (l_run.buffer.as_ref(), 0)], + &[i(nq as i32), i(sq as i32), i(hd as i32)], + [(nq*sq) as u32, 1, 1], [128,1,1], 0, false)?; } - dev.dispatch_raw_cuda(SDPA_TC_MERGE_SRC, "sdpa_tc_merge.cu", "sdpa_tc_merge", - &[(o_blk.buffer.as_ref(), 0), (bm.buffer.as_ref(), 0), (bl.buffer.as_ref(), 0), - (o_run.buffer.as_ref(), 0), (m_run.buffer.as_ref(), 0), (l_run.buffer.as_ref(), 0)], - &[i(nq as i32), i(sq as i32), i(hd as i32)], - [(nq*sq) as u32, 1, 1], [128,1,1], 0, false)?; kb0 += blk; } let out = Tensor::empty(dev, vec![sq, nq, hd], q.dtype)?; @@ -6885,19 +7111,45 @@ pub fn mamba_split_conv( // CONV_DEVICE prefill path (TODO comment at line ~1543 of modeltests/lib.rs). /// RMSNorm with f16 output only. Equivalent to `cast_f32_f16(rms_norm(...))`, /// but avoids the intermediate f32 norm tensor and the separate cast launch. -/// TODO(rebrand-merge): the fast-path fused kernel (`iron_rms_norm_f16out`) -/// isn't in `thewafflehaus/iron@dev` yet, so this always takes the -/// `rms_norm` + `cast_f32_f16` fallback below (same numerics, one extra -/// dispatch) instead of the single-pass fused kernel. +/// Dispatches the fused `iron_rms_norm_f16out` kernel on the same fast-path +/// shape restriction as plain `rms_norm`'s fast path (f32 input, n a +/// multiple of 128, n <= 4096 — the 16-bit vec4 codegen restriction that +/// forces `rms_norm` onto `iron_rms_norm_wide` for f16/bf16 applies here +/// too, since this is the same 4-elements/thread codegen); anything outside +/// that falls back to `rms_norm` + `cast_f32_f16`. pub fn rms_norm_f16out(dev: &dyn Device, x: &Tensor, weight: &Tensor, eps: f32) -> Result { + let n = *x.shape.last().ok_or_else(|| Error::Msg("rms_norm_f16out: scalar input".into()))?; + if x.dtype == DType::F32 && n % 128 == 0 && n <= 4096 { + let rows = x.elem_count() / n; + let out = Tensor::empty(dev, x.shape.clone(), DType::F16)?; + let eps_buf = scalar_buf(dev, eps.to_bits())?; + let k = cached_ir("iron_rms_norm_f16out", x.dtype, || { + let mut k = wh_iron_std::kernels::norm::rms_norm::iron_rms_norm_f16out::kernel_ir_for(x.dtype); + k.mode = wh_iron_core::ir::KernelMode::Reduction; + k + }); + dev.dispatch( + &k, + &[ + Binding::Buffer(x.buffer.clone()), + Binding::Buffer(weight.buffer.clone()), + Binding::Buffer(out.buffer.clone()), + Binding::Buffer(eps_buf), + Binding::Scalar((n as u32).to_le_bytes().to_vec()), + ], + Grid { grid: [rows as u32, 1, 1], block: [(n / 4) as u32, 1, 1] }, + )?; + return Ok(out); + } let out = rms_norm(dev, x, weight, eps)?; cast_f32_f16(dev, &out) } -/// TODO(rebrand-merge): the fused kernel (`iron_gated_group_rmsnorm_batched_f16out`) -/// isn't in `thewafflehaus/iron@dev` yet (only the plain, non-f16out -/// `iron_gated_group_rmsnorm_batched` is) — always takes the -/// `gated_group_rmsnorm_batched` + `cast_f32_f16` fallback below. +/// Fused single-pass variant of `gated_group_rmsnorm_batched` that stores +/// directly to f16 (skips the intermediate f32 tensor + separate cast +/// launch). Dispatches `iron_gated_group_rmsnorm_batched_f16out`, same +/// grid/block scheme as the plain kernel (one threadgroup per (token, +/// norm-group) pair, 4 elements/thread). pub fn gated_group_rmsnorm_batched_f16out( dev: &dyn Device, y: &Tensor, @@ -6908,8 +7160,30 @@ pub fn gated_group_rmsnorm_batched_f16out( di: usize, gs: usize, ) -> Result { - let out = gated_group_rmsnorm_batched(dev, y, z, w, eps, s, di, gs)?; - cast_f32_f16(dev, &out) + let ng = di / gs; + let kernel = cached_ir("iron_gated_group_rmsnorm_batched_f16out", DType::F32, || { + use wh_iron_core::ir::KernelMode; + let mut k = wh_iron_std::kernels::ssm::scan::iron_gated_group_rmsnorm_batched_f16out::kernel_ir_for(); + k.mode = KernelMode::Reduction; + k + }); + let out = Tensor::empty(dev, vec![s * di], DType::F16)?; + let eps_buf = Tensor::new(scalar_buf(dev, eps.to_bits())?, vec![1], DType::F32); + let u = |v: u32| Binding::Scalar(v.to_le_bytes().to_vec()); + dev.dispatch( + &kernel, + &[ + Binding::Buffer(y.buffer.clone()), + Binding::Buffer(z.buffer.clone()), + Binding::Buffer(w.buffer.clone()), + Binding::Buffer(out.buffer.clone()), + Binding::Buffer(eps_buf.buffer.clone()), + u(gs as u32), + u(ng as u32), + ], + Grid { grid: [(s * ng) as u32, 1, 1], block: [(gs / 4) as u32, 1, 1] }, + )?; + Ok(out) } pub fn gated_group_rmsnorm_batched( @@ -7534,17 +7808,41 @@ using namespace nvcuda; // eliminating __half2float() inside the per-nibble dequant loop. // 3. Double-buffered X and WT staging: next tile's loads overlap // current tile's WMMA compute (ping-pong smem slots). +// 4. Padded smem_WT row stride (128->136 halfwords) to break the +// power-of-two bank-conflict alignment: a plain 128-halfword +// (256-byte) row stride is an exact multiple of the 32-bank×4-byte +// conflict period, so a warp's 32 threads (which vary the ROW index +// while the column/bank-determining index stays fixed within an +// unrolled dequant iteration) ALL alias the same bank on every +// LOAD_WT store — confirmed via `ncu` (l1tex bank-conflict counters +// vs wavefront counters): ~18x store-side overhead vs the +// conflict-free ideal on this kernel. 136 is the smallest multiple +// of 8 (WMMA's row-major `ldm` alignment requirement for __half) +// that reduces the conflict: the achievable minimum under that +// constraint is a 4-way conflict (down from 32-way); modular +// arithmetic (row_wt/2 mod 32) shows 8-alignment forbids getting to +// fully conflict-free for this 32-row/128-bank shape. smem_X (the +// other bank-conflict offender ncu found, ~3.5x overhead, load-side) +// is NOT padded: this device's actual opt-in dynamic-shared-memory +// cap is 101376 bytes (queried via +// CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 101376, +// NOT the ~228KB this comment previously assumed by analogy to a +// different GPU) and there is no budget left to also pad X (16384 +// bytes at the padded WT size already leaves only ~1.5KB of +// headroom) without exceeding it and failing `cuFuncSetAttribute` +// at launch. Left as a documented follow-up if the tile/accumulator +// layout is ever revisited to free up smem headroom. // // qs layout: Marlin tile-major, same 64-row tiles as moe_w4a16_marlin. // A 128-N-tile block straddles two consecutive 64-N Marlin tiles (nt0, nt0+1). // 256 threads × 1 u32/tile → both tiles load in a single fully-coalesced pass. // -// smem breakdown (per block): -// smem_X [2][128*32] f16 = 16384 B (double-buffered X) -// smem_WT[2][32*128] f16 = 16384 B (double-buffered W^T, already dequanted) +// smem breakdown (per block; WT row PADDED to 136 halfwords, X unpadded): +// smem_X [2][128*32] f16 = 16384 B (double-buffered X, unpadded) +// smem_WT[2][32*136] f16 = 17408 B (double-buffered W^T, padded) // smem_SC[128] f32 = 512 B (preloaded scales for this K-step) // smem_C [128*128] f32 = 65536 B (accumulator) -// Total = 98816 B (~96 KB; GB10 has 228 KB/SM) +// Total = 99840 B (~97.5 KB; GB10 opt-in cap is 101376 B) // // n_out % 128 == 0 and k_in % 32 == 0 required. extern "C" __global__ void moe_w4a16_marlin128( @@ -7564,16 +7862,16 @@ extern "C" __global__ void moe_w4a16_marlin128( const int warp_m_base = (warp_id >> 2) * 64; // 0 or 64 const int warp_n_base = (warp_id & 3) * 32; // 0, 32, 64, or 96 - // smem layout (all offsets in bytes): + // smem layout (all offsets in bytes; WT row PADDED to 136, X unpadded): // [0 .. 16384) smem_X [2][128*32] __half - // [16384 .. 32768) smem_WT[2][32*128] __half - // [32768 .. 33280) smem_SC[128] float - // [33280 .. 98816) smem_C [128*128] float + // [16384 .. 33792) smem_WT[2][32*136] __half + // [33792 .. 34304) smem_SC[128] float + // [34304 .. 99840) smem_C [128*128] float extern __shared__ char smem_raw[]; __half* smem_X = (__half*)(smem_raw); __half* smem_WT = (__half*)(smem_raw + 16384); - float* smem_SC = (float* )(smem_raw + 32768); - float* smem_C = (float* )(smem_raw + 33280); + float* smem_SC = (float* )(smem_raw + 33792); + float* smem_C = (float* )(smem_raw + 34304); const int bpr = k_in / 32; const int n_tiles64 = n_out / 64; @@ -7623,7 +7921,7 @@ extern "C" __global__ void moe_w4a16_marlin128( int nl = rg * 8 + _i; \ unsigned nib = (w0 >> (_i*4)) & 0xf; \ int qs_val = (int)(nib >= 8u ? nib - 16u : nib); \ - smem_WT[(slot_)*(32*128) + kl*128 + nl] = \ + smem_WT[(slot_)*(32*136) + kl*136 + nl] = \ __float2half((float)qs_val * smem_SC[nl]); \ } \ } \ @@ -7636,7 +7934,7 @@ extern "C" __global__ void moe_w4a16_marlin128( int nl = 64 + rg * 8 + _i; \ unsigned nib = (w1 >> (_i*4)) & 0xf; \ int qs_val = (int)(nib >= 8u ? nib - 16u : nib); \ - smem_WT[(slot_)*(32*128) + kl*128 + nl] = \ + smem_WT[(slot_)*(32*136) + kl*136 + nl] = \ __float2half((float)qs_val * smem_SC[nl]); \ } \ } \ @@ -7700,8 +7998,8 @@ extern "C" __global__ void moe_w4a16_marlin128( #pragma unroll for (int ni = 0; ni < 2; ni++) { wmma::load_matrix_sync(b_frag, - smem_WT + buf * (32*128) + k_off * 128 + (warp_n_base + ni*16), - 128); + smem_WT + buf * (32*136) + k_off * 136 + (warp_n_base + ni*16), + 136); wmma::mma_sync(c_frag[mi][ni], a_frag, b_frag, c_frag[mi][ni]); } } @@ -7804,8 +8102,12 @@ pub fn moe_w4a16_marlin( ], [(n_out as u32) / 128, (m_total as u32).div_ceil(128), 1], [256, 1, 1], - // smem: X_db(16384) + WT_db(16384) + SC(512) + C(65536) = 98816 bytes (~96 KB) - 98816, + // smem: X_db(16384, unpadded) + WT_db(17408, padded 128->136) + SC(512) + // + C(65536) = 99840 bytes (~97.5 KB). GB10's opt-in dynamic-smem cap is + // 101376 bytes (queried via CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN), + // so this is the largest WT padding this device's budget allows without also + // shrinking something else; see the kernel doc for the full rationale. + 99840, false, )?; } else {