perf(laguna): prefill optimization campaign — sliding-window SDPA, marlin smem padding, moe heapsort + pipeline depth

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");