perf(moe): warp-split moe_q4_grouped_mma variant, opt-in (BUTTER_LAGUNA_MOE_NSPLIT=1)

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<usize>| -> 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<u8> = 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<u8> = 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; };

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)`

rust/crates/wh-butter-modeltests/src/lib.rs

@@ -588,20 +588,27 @@ pub fn verify_laguna(d: &dyn Device, plat: &str) {
     // BUTTER_LAGUNA_PP=<n>: 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<usize> = 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<usize> = std::env::var("BUTTER_LAGUNA_PREFILL_CHUNK")
             .ok().map(|v| v.split(',').filter_map(|c| c.trim().parse().ok()).collect())
             .filter(|v: &Vec<usize>| !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<u32> = (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<f32> = 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");