memra_engine/decode_batch.rs
1//! Batched decode step — B sequences share one fused pass (ARCHITECTURE-H100.md §3 B2').
2//!
3//! The bandwidth thesis: decode is weight-stream-bound, so every projection at m=B rows
4//! amortizes one weight read across B sequences. Row-parallel ops (norm/rope/quantize/
5//! activation) batch trivially — they are the SAME kernels prefill already runs at T rows.
6//! Only truly per-sequence state stays in a loop: KV append + fa_decode over each cache,
7//! and the GDN/conv recurrent step (v1: per-seq loop via the existing single-seq path;
8//! a blockIdx.z-batched GDN state kernel is the v2 fusion).
9//!
10//! EXACTNESS CONTRACT (the law this module lives under):
11//! - B == 1 must be BIT-IDENTICAL to `decode_step_h` (gate: decode-batch-gate).
12//! - 2 <= B <= 8: each row rides the m=2..9 verify-tier mmvq kernels, which are per-row
13//! bit-identical to m=1 (the spec-exactness machinery decode_step_t relies on). Each
14//! sequence's token stream must equal its isolated single-seq run (worker.rs contract:
15//! "byte-identical to isolated").
16//! - 9 <= B <= 16 (the EXACT-16 tier, inc3 2026-08-01): admitted iff
17//! `decode_batch_exact16_ok` — every matmul rides the b16 batched-mmvq class
18//! (bit-identical per (token,row) to m=1; Q8_0 needs the q8rp mirror) under a
19//! verify_exact scope that disables the m>=16 GEMM/MMQ arms. gate2 bit-strength
20//! PASS at B=12/16 (research/batched-tick-inc3-20260801). Refused otherwise.
21//! - B > 16 crosses into GEMM/dp4a-tail numeric configs with NO exact kernel class —
22//! refused (MEMRA_DECODE_BATCH_CAP stays a measurement door).
23//!
24//! v1 scope: the hybrid (Qwen3.5-class) non-gemma4 trunk. Fused m=1 micro-launches
25//! (fused3 QKV, cross-layer add+norm+q8 chain) are NOT used — the unfused sequence is
26//! bit-identical (kernel_check: add_rms_norm == add;rms_norm; _q8_1 == +quantize_q8_1)
27//! and keeps the batched path simple. Batched fusions are tuning work, not correctness.
28
29use crate::Engine;
30use crate::cache::Cache;
31use crate::hybrid::{HybridModel, Mixer};
32use cudarc::driver::{CudaEvent, CudaSlice};
33use memra_gguf::config::Arch;
34
35type DualPpCudaSpan = Option<(CudaEvent, CudaEvent)>;
36
37fn dual_pp_timing_event(e: &Engine, context: &str) -> Option<CudaEvent> {
38 if !crate::pp::dual_pp_timing_on() {
39 return None;
40 }
41 match e
42 .stream()
43 .record_event(Some(cudarc::driver::sys::CUevent_flags::CU_EVENT_DEFAULT))
44 {
45 Ok(event) => Some(event),
46 Err(err) => {
47 crate::pp::record_dual_pp_timing_drop(context, &err);
48 None
49 }
50 }
51}
52
53/// Per-step, per-LAYER-RANGE invariants the batched trunk needs: the device state-pointer
54/// table for the range's layers, the arm picks, and the per-row `t_kv` snapshot. Built once
55/// per step per range by `HybridModel::batch_layer_ctx`, consumed by `decode_batch_layers`.
56///
57/// WHY IT IS RANGE-SCOPED AND NOT STEP-SCOPED (this is the whole point of the struct):
58/// `ptr_table` is a `CudaSlice<u64>` of DEVICE ADDRESSES, uploaded through `e` — so it lives
59/// on `e`'s device, and its entries are pointers into caches that live on the device that
60/// OWNS those layers. Under a pp stage split, stage s runs layers [fence[s], fence[s+1])
61/// whose cache state was allocated by stage s's engine (`pp::new_cache` -> `Cache::new_ppn`),
62/// so stage s must build its OWN table through its OWN engine. One step-wide table built on
63/// the primary would put every stage's kernel arguments in stage-0's HBM — a peer read per
64/// pointer fetch, which is the exact cliff `pp::refuse_unsplit_if_remote` exists to stop.
65/// `lo`/`hi` are recorded so the consumer can assert the ctx it was handed matches the range
66/// it was asked to run (the offsets in `lin_base`/`attn_base` are only valid for that range).
67pub(crate) struct BatchLayerCtx {
68 /// Offset into `ptr_table` of layer il's [conv x B][ssm_in x B][ssm_out x B] block
69 /// (linear-attn layers only). Indexed by ABSOLUTE layer id; `None` off-range.
70 lin_base: Vec<Option<usize>>,
71 /// Offset into `ptr_table` of layer il's [k0,v0,k1,v1,..] block (full-attn layers only).
72 /// Indexed by ABSOLUTE layer id; `None` off-range.
73 attn_base: Vec<Option<usize>>,
74 ptr_table: Option<CudaSlice<u64>>,
75 /// Per-row `pos + 1` — the t_kv each sequence attends at this step. Layer-invariant
76 /// within a step, so the arm picks below are decided once.
77 t_kvs: Vec<usize>,
78 t_kv_max: usize,
79 /// The single `fa_split_keys` rung every row shares (the rows-twins straddle law).
80 sp0: usize,
81 seqs_append: bool,
82 seqs_fa: bool,
83 lo: usize,
84 hi: usize,
85}
86
87// ---- MEMRA_BATCH_PHASE=1 (diagnostics): sync-bounded per-phase accumulators for the batched
88// tick. Each boundary syncs the stream, so the TOTAL inflates (launch pipelining is destroyed);
89// the value is the RANKING/shares, not absolute ms. Read via `batch_phase_report()`.
90pub(crate) static BATCH_PHASE: std::sync::Mutex<[f64; 12]> = std::sync::Mutex::new([0.0; 12]);
91/// Device-sample request for one batched row: (temp, seed, ctr, top_k, top_p, min_p).
92/// `top_k=0` / `top_p>=1.0` / `min_p<=0.0` = that filter off. Greedy = temp<=0 (device
93/// argmax); pure temperature = seeded gumbel; any filter on = filter_stats floor + the
94/// filtered gumbel draw. Penalty configs never reach device sampling (worker eligibility).
95pub type DevSamp = (f32, u64, u32, i32, f32, f32);
96
97pub const BATCH_PHASE_NAMES: [&str; 12] = [
98 "setup(ptrs+embed H2D)",
99 "attn batched pre (norm/qkv/rope)",
100 "attn per-seq: kv append",
101 "attn per-seq: q/a dtod copies",
102 "attn per-seq: fa_decode",
103 "attn post (gate+o-proj)",
104 "gdn batched projections",
105 "gdn state ops (conv/prep/scan)",
106 "gdn out (gated norm+proj)",
107 "ffn (add/norm/gate/up/act/down)",
108 "lm_head (norm+matmul)",
109 "logits D2H + host split",
110];
111pub fn batch_phase_on() -> bool {
112 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
113 *ON.get_or_init(|| std::env::var("MEMRA_BATCH_PHASE").as_deref() == Ok("1"))
114}
115/// Accumulate the elapsed time since `last` into phase slot `slot` and re-stamp `last`.
116/// No-op unless `MEMRA_BATCH_PHASE=1`. Syncs the ambient stream first, so under a pp stage
117/// scope this bounds the STAGE's stream, which is what the caller is timing.
118///
119/// A free fn rather than the closure it replaced: `decode_batch_layers` (the pp stage seam)
120/// runs the instrumented layer loop, so the marker has to be callable from both the seam
121/// and its caller's epilogue. `batch_phase_on()` is a `OnceLock` memo, so per-call cost is
122/// the same atomic load the hoisted `ph_on` local was.
123fn ph_mark(
124 e: &Engine,
125 slot: usize,
126 last: &mut std::time::Instant,
127) -> Result<(), Box<dyn std::error::Error>> {
128 if batch_phase_on() {
129 e.stream().synchronize()?;
130 let now = std::time::Instant::now();
131 BATCH_PHASE.lock().unwrap()[slot] += (now - *last).as_secs_f64();
132 *last = now;
133 }
134 Ok(())
135}
136
137pub fn batch_phase_report() -> String {
138 let ph = BATCH_PHASE.lock().unwrap();
139 let tot: f64 = ph.iter().sum();
140 let mut rows: Vec<(usize, f64)> = ph.iter().copied().enumerate().collect();
141 rows.sort_by(|a, b| b.1.total_cmp(&a.1));
142 let mut s = format!(
143 "[batch-phase] total {:.1} ms (sync-bounded; shares rank, not walltime)\n",
144 tot * 1e3
145 );
146 for (i, v) in rows {
147 s += &format!(
148 " {:>6.1} ms {:>5.1}% {}\n",
149 v * 1e3,
150 v / tot * 100.0,
151 BATCH_PHASE_NAMES[i]
152 );
153 }
154 s
155}
156
157impl HybridModel {
158 /// Batched-decode width cap. 8 = the exactness-tier default (see the assert below);
159 /// MEMRA_DECODE_BATCH_CAP overrides for tier-probe measurement, clamped to 32.
160 pub fn decode_batch_cap() -> usize {
161 use std::sync::OnceLock;
162 static CAP: OnceLock<usize> = OnceLock::new();
163 *CAP.get_or_init(|| {
164 std::env::var("MEMRA_DECODE_BATCH_CAP")
165 .ok()
166 .and_then(|v| v.parse().ok())
167 .map(|c: usize| c.clamp(1, 32))
168 .unwrap_or(8)
169 })
170 }
171
172 /// EXACT-16 TIER admission (increment 3a, 2026-08-01, 5090 receipts
173 /// research/batched-tick-inc3-20260801): true iff EVERY matmul the batched decode step
174 /// runs has a per-(token,row) bit-exact kernel class at m=9..16 under the verify_exact
175 /// scope — i.e. the batched-mmvq b16 family (32-thread warp reduce, the exact m=1 mmvq
176 /// program per column) or the e4m3 grid.y=m mmvq catch-all. Q8_0 qualifies only with
177 /// the split-plane mirror (rp4, MEMRA_Q8RP): its b16 kernel exists only as the _rp twin.
178 /// Float matmuls (cuBLASLt, n-dependent reductions) and MoE FFNs disqualify the model.
179 /// Measured attribution for WHY the naked m=16 tier is not exact: the m>=16 arms
180 /// (MMQ int8-MMA `mul_mat_q` — MEMRA_PP_Q8MMQ default-on — and `qmatvec_gemm`, both
181 /// block-scale f32) and the m=9..15 dp4a tail (128-thread two-level reduce) all break
182 /// per-row bit-identity vs isolated decode (gate2 step-0 bit-diffs, maxdiff ~1.3-2.3e-1).
183 pub fn decode_batch_exact16_ok(&self) -> bool {
184 fn ok(w: &crate::model::GpuTensor) -> bool {
185 match w {
186 crate::model::GpuTensor::Quant { qtype, .. } => {
187 *qtype == crate::QT_Q4_0 || *qtype == crate::QT_Q6_K
188 || *qtype == crate::QT_F8_E4M3
189 // BLOCK-128 FP8-ST (lane/rp-on-st, 2026-08-06): admitted now that the class
190 // has a b16 batched kernel (`qmatvec_e4m3_blk_mmvq_b16`), bit-identical per
191 // (token,row) to its m=1 launch. Before that kernel existed this class fell to
192 // the grid.y=m form at every width — still EXACT, so the tier's correctness
193 // bar was met, but it re-read the weight m times, which is why admitting it
194 // without the kernel would have been a throughput trap rather than a win.
195 || *qtype == crate::QT_F8_E4M3_BLK
196 // NVFP4 (lane/rp-on-st, 2026-08-06) — THE blocker this lane measured. The
197 // mixed FP8-ST 27B is 193 NVFP4 dense-MLP tensors, and this predicate is an
198 // ALL over every matmul, so NVFP4's missing b16 refused the whole checkpoint
199 // (`B=16 > cap 8 with no exact tier ... refused`) even with both e4m3 classes
200 // admitted. It now has base + _rp b16 twins off its existing batched template
201 // (bit-identical per (token,row) to the m=1 mmvq: same nibble decode, dp4a
202 // order, ue4m3 scale, warp reduce). This also opens the tier for pure-NVFP4
203 // GGUF models, which is a behavior change on the primary format — hence the
204 // full decode-batch config+strict battery on both.
205 || *qtype == crate::QT_NVFP4
206 // Q4_K (lane/rp-on-st): named by MEMRA_EXACT16_WHY as the 9B NVFP4 GGUF's
207 // refusing class (`L0.wqkv qtype=1`) — mixed NVFP4 checkpoints keep Q4_K
208 // attention. Now has base + _rp b16.
209 || *qtype == crate::QT_Q4_K
210 // Q5_K (lane/rp-on-st): the FOURTH class the diagnostic named on the same 9B
211 // GGUF (`L0.wqkv_gate qtype=3`). A shipped mixed checkpoint spreads ~500
212 // matmuls over four/five classes, and this predicate is an ALL — so chunk 16
213 // was unreachable for every real artifact until every class had a b16.
214 || *qtype == crate::QT_Q5_K
215 // Q8_0 NO LONGER requires the mirror (rp4): it has a base b16 too, so the
216 // tier is reachable at zero VRAM. Named by the diagnostic as the FP8-ST
217 // refusal — `L0.ssm_beta qtype=0 rp4=false`, a 23.9 MiB residual class that
218 // was gating chunk 16 for a 16.4 GiB checkpoint.
219 || *qtype == crate::QT_Q8_0
220 }
221 _ => false,
222 }
223 }
224 // WHY-NOT DIAGNOSTIC (lane/rp-on-st, 2026-08-06): this predicate is a bare bool over
225 // ~500 tensors, so a refusal produced only `B=16 > cap 8 with no exact tier ... refused`
226 // with no way to tell WHICH class refused. That cost this lane two wrong hypotheses (the
227 // rp mirror, then e4m3-only) before the NVFP4 gap was found. MEMRA_EXACT16_WHY=1 names
228 // the first refusing tensor + its qtype. Diagnostic-only per flags doctrine; default off,
229 // zero cost when unread.
230 let why = std::env::var("MEMRA_EXACT16_WHY").is_ok();
231 macro_rules! chk {
232 ($t:expr, $label:expr) => {{
233 let r = ok($t);
234 if !r && why {
235 // qtype = -1 means the tensor is NOT Quant at all (a float/BF16/F16
236 // container), which the tier can never admit — a distinct diagnosis from
237 // "quantized, but in a class with no b16 kernel".
238 let (qt, rp4) = match $t {
239 crate::model::GpuTensor::Quant { qtype, rp4, .. } => {
240 (*qtype, rp4.is_some())
241 }
242 _ => (-1, false),
243 };
244 eprintln!("[exact16] REFUSED by {} qtype={qt} rp4={rp4}", $label);
245 }
246 r
247 }};
248 }
249 if self.cfg.m3.is_some() || self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
250 if why {
251 eprintln!("[exact16] REFUSED by architecture (m3/gemma4)");
252 }
253 return false;
254 }
255 self.layers.iter().enumerate().all(|(li, l)| {
256 let mix_ok = match &l.mixer {
257 Mixer::Full(fa) => {
258 chk!(&fa.wq, format!("L{li}.wq"))
259 && chk!(&fa.wk, format!("L{li}.wk"))
260 && chk!(&fa.wv, format!("L{li}.wv"))
261 && chk!(&fa.wo, format!("L{li}.wo"))
262 }
263 Mixer::Linear(la) => {
264 chk!(&la.wqkv, format!("L{li}.wqkv"))
265 && chk!(&la.wqkv_gate, format!("L{li}.wqkv_gate"))
266 && chk!(&la.ssm_beta, format!("L{li}.ssm_beta"))
267 && chk!(&la.ssm_alpha, format!("L{li}.ssm_alpha"))
268 && chk!(&la.ssm_out, format!("L{li}.ssm_out"))
269 }
270 // MLA rides its own increment-4 arm; never admitted to the exact-16 tier here.
271 Mixer::Mla(_) => {
272 if why {
273 eprintln!("[exact16] REFUSED by L{li} MLA mixer");
274 }
275 false
276 }
277 };
278 let ffn_ok = match &l.ffn {
279 crate::hybrid::Ffn::Dense {
280 ffn_gate,
281 ffn_up,
282 ffn_down,
283 } => {
284 chk!(ffn_gate, format!("L{li}.ffn_gate"))
285 && chk!(ffn_up, format!("L{li}.ffn_up"))
286 && chk!(ffn_down, format!("L{li}.ffn_down"))
287 }
288 crate::hybrid::Ffn::Moe(_) => {
289 if why {
290 eprintln!("[exact16] REFUSED by L{li} MoE ffn");
291 }
292 false
293 }
294 };
295 mix_ok && ffn_ok
296 }) && chk!(&self.output, "output".to_string())
297 }
298
299 /// Opt-in/A-B seam for the eager B=1 fusion program. `MEMRA_SERVE_B1FAST=1` sends an
300 /// eligible solo tick through that program; unset/other values keep B=1 on the generic
301 /// batched body, the same numeric class used at B>=2.
302 ///
303 /// EXACTNESS, stated precisely (measured on-box 2026-08-05, sm_120 q9 NVFP4-MTP):
304 /// the fast path is BIT-IDENTICAL TO `decode_step_h` — decode-batch-gate's STRICT
305 /// gate1 (`--mode strict`) PASSes with it ON and FAILs with it OFF at maxdiff
306 /// 1.591e-1. It is deliberately NOT bit-identical to the batched body: the two
307 /// carry a decode-config FP-composition gap (same class gate1's config mode measures).
308 /// That gap became correctness-visible under live load: Step35, Q35-MoE, and finally
309 /// dense Q27 all produced load-history-dependent token streams, including early EOS,
310 /// when a request crossed between the two programs. The generic body is therefore the
311 /// correctness default; the eager program remains available only for fixed-solo A/Bs.
312 /// Historical token-stream/performance receipts:
313 /// research/servepath-p2-20260805 (greedy 150 ids + seeded-sampled identical to the
314 /// run-gen oracle AND cross-arm, so the gap is sub-token here as designed).
315 ///
316 /// Read fresh (an `AtomicU8` memo, not a `OnceLock`): decode-batch-gate flips this
317 /// seam BETWEEN gates in-process — gate1 needs the fast path ON to prove bit-identity,
318 /// gate2 needs it pinned OFF to keep testing the batched body. A latch-once read would
319 /// bake whichever gate ran first, so the gate could never test both sides. The memo
320 /// caches the parse but `set_b1_fast` invalidates it.
321 pub fn b1_fast_on() -> bool {
322 // 0 = unknown/invalidated, 1 = off, 2 = on
323 match Self::b1_fast_memo().load(std::sync::atomic::Ordering::Relaxed) {
324 1 => false,
325 2 => true,
326 _ => {
327 let value = std::env::var("MEMRA_SERVE_B1FAST").ok();
328 let on = b1_fast_env_on(value.as_deref());
329 Self::b1_fast_memo()
330 .store(if on { 2 } else { 1 }, std::sync::atomic::Ordering::Relaxed);
331 on
332 }
333 }
334 }
335
336 fn b1_fast_memo() -> &'static std::sync::atomic::AtomicU8 {
337 static MEMO: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
338 &MEMO
339 }
340
341 /// Test/gate seam: force the B=1 fast path on or off for the rest of the process,
342 /// overriding the env. Used by decode-batch-gate to exercise the opt-in eager arm and
343 /// pin gate2's default reference arm.
344 pub fn set_b1_fast(on: bool) {
345 Self::b1_fast_memo().store(if on { 2 } else { 1 }, std::sync::atomic::Ordering::Relaxed);
346 }
347
348 /// Whether this architecture may switch a live serving row onto the eager B=1 fusion
349 /// class. Qwen35-MoE must stay on the batched trunk at every width: its eager and batched
350 /// hybrid/MoE walks are each deterministic, but crossing B=1 -> B>=2 changes greedy token
351 /// ids and can introduce an early EOS (Q35 sellgate, 2026-08-12).
352 pub fn b1_fast_arch_eligible(&self) -> bool {
353 b1_fast_arch_eligible(&self.cfg.arch)
354 }
355
356 /// H3 body: the m=1 FUSED trunk (`decode_layers_eager` — shared verbatim with
357 /// `decode_step_h`/the ppN stages) plus the batched path's own serving epilogue
358 /// (grammar mask, device sample, lean-logits park). See the call-site comment in
359 /// `decode_step_batch_sampled_lean_masked` for why this is bit-identical.
360 fn decode_step_b1_fast(
361 &self,
362 e: &Engine,
363 token: u32,
364 caches: &mut [&mut Cache],
365 samp: &[Option<DevSamp>],
366 masks: &[Option<(&CudaSlice<u32>, usize)>],
367 lean: bool,
368 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
369 let n_embd = self.cfg.n_embd as usize;
370 let eps = self.cfg.rms_eps;
371 let pos = caches[0].pos;
372 let pos_d = e.htod_i32(&[pos as i32])?;
373 let x = e.htod(&self.embd.gather(n_embd, &[token]))?;
374 // the SHARED m=1 trunk: same function decode_step_h runs, so every m=1 fusion
375 // (cross-layer add+norm+q8_1, fused SwiGLU, lever 1's gate+up dual) fires here.
376 let x = self.decode_layers_eager(e, x, 0, self.layers.len(), &pos_d, pos, caches[0])?;
377 let mut hn = e.uninit(n_embd)?;
378 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
379 let logits = e.matmul(&self.output, &hn, 1)?;
380
381 // ---- epilogue: byte-for-byte the batched path's, at b_n=1 ----
382 let n_vocab = self.output.out_features();
383 let mut logits = logits;
384 let mut pristine: Option<CudaSlice<f32>> = None;
385 if let Some((mask, words)) = masks.first().copied().flatten() {
386 assert!(
387 samp.first().copied().flatten().is_some(),
388 "grammar-masked row 0 must request a device sample"
389 );
390 if lean {
391 let cache = &mut caches[0];
392 if cache
393 .last_logits_dev
394 .as_ref()
395 .map(|d| d.len() < n_vocab)
396 .unwrap_or(true)
397 {
398 cache.last_logits_dev = Some(e.uninit(n_vocab)?);
399 }
400 let dst = cache.last_logits_dev.as_mut().unwrap();
401 e.dtod_copy_view(&logits.slice(0..n_vocab), dst)?;
402 } else {
403 let mut p = e.uninit(n_vocab)?;
404 e.dtod_copy_view(&logits.slice(0..n_vocab), &mut p)?;
405 pristine = Some(p);
406 }
407 e.mask_logits_col(&mut logits, mask, 0, n_vocab, words)?;
408 }
409
410 let mut next: Vec<Option<u32>> = vec![None; 1];
411 if let Some((temp, seed, ctr, top_k, top_p, min_p)) = samp.first().copied().flatten() {
412 let mut toks = e.alloc_u32_zeroed(1)?;
413 // Filtered-greedy degenerates to plain argmax (the max always survives every
414 // truncation filter), so temp<=0 short-circuits regardless of filters.
415 let filtered = temp > 0.0 && (top_k > 0 || top_p < 1.0 || min_p > 0.0);
416 if temp <= 0.0 {
417 e.argmax_token_device_col(&logits, 0, n_vocab, &mut toks, 0)?;
418 } else if filtered {
419 let mut pb = e.zeros(n_vocab)?;
420 self.devsample_filtered_col(
421 e, &logits, 0, n_vocab, temp, seed, ctr, top_k, top_p, min_p, &mut pb,
422 &mut toks, 0,
423 )?;
424 } else {
425 let mut pb = e.zeros(n_vocab)?;
426 e.gumbel_perturb_col(&logits, 0, &mut pb, n_vocab, seed, ctr, temp)?;
427 e.argmax_token_device_col(&pb, 0, n_vocab, &mut toks, 0)?;
428 }
429 next[0] = Some(e.dtoh_u32(&toks)?[0]);
430 }
431
432 let sampled = samp.first().copied().flatten().is_some();
433 let rows: Vec<Vec<f32>> = if lean && sampled {
434 if masks.first().copied().flatten().is_none() {
435 let cache = &mut caches[0];
436 if cache
437 .last_logits_dev
438 .as_ref()
439 .map(|d| d.len() < n_vocab)
440 .unwrap_or(true)
441 {
442 cache.last_logits_dev = Some(e.uninit(n_vocab)?);
443 }
444 let dst = cache.last_logits_dev.as_mut().unwrap();
445 e.dtod_copy_view(&logits.slice(0..n_vocab), dst)?;
446 }
447 vec![Vec::new()]
448 } else if let Some(p) = pristine.as_ref() {
449 vec![e.dtoh(p)?]
450 } else {
451 vec![e.dtoh(&logits)?]
452 };
453 // decode_layers_eager does NOT advance cache.pos (decode_step_h advances it after
454 // the head); the batched path advances every cache at the tail — same here.
455 caches[0].pos += 1;
456 Ok((rows, next))
457 }
458
459 /// One filtered device draw for stacked-logits row `col`: `filter_stats` solves the
460 /// single unnormalized-prob floor that encodes top-k AND top-p AND min-p (block-internal
461 /// binary search, bit-stable), then the filtered gumbel perturb + argmax draws one token
462 /// from the truncated softmax into `toks[slot]`. All device-side — no stat D2H, no row
463 /// copy; the only host traffic stays the caller's one [B]-u32 token readback.
464 #[allow(clippy::too_many_arguments)]
465 fn devsample_filtered_col(
466 &self,
467 e: &Engine,
468 logits: &CudaSlice<f32>,
469 col: usize,
470 n_vocab: usize,
471 temp: f32,
472 seed: u64,
473 ctr: u32,
474 top_k: i32,
475 top_p: f32,
476 min_p: f32,
477 pb: &mut CudaSlice<f32>,
478 toks: &mut CudaSlice<u32>,
479 slot: usize,
480 ) -> Result<(), Box<dyn std::error::Error>> {
481 let rows = e.htod_i32(&[col as i32])?;
482 let mut th = e.zeros(1)?;
483 let mut z = e.zeros(1)?;
484 let mut mx = e.zeros(1)?;
485 e.filter_stats(
486 logits, n_vocab, &rows, &mut th, &mut z, &mut mx, n_vocab, 1, temp, top_k, top_p, min_p,
487 )?;
488 e.gumbel_perturb_filtered_col(logits, col, pb, n_vocab, seed, ctr, temp, &mx, &th, 0)?;
489 e.argmax_token_device_col(pb, 0, n_vocab, toks, slot)?;
490 Ok(())
491 }
492
493 /// One batched greedy-decode step over B independent sequences.
494 /// `tokens[b]` is sequence b's input token; `caches[b]` its private cache (position,
495 /// quantized KV, GDN/conv state). Returns the B logits rows (host, [n_vocab] each).
496 /// Each cache's pos/len advance exactly as `decode_step_h` would.
497 pub fn decode_step_batch(
498 &self,
499 e: &Engine,
500 tokens: &[u32],
501 caches: &mut [&mut Cache],
502 ) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
503 let (rows, _) = self.decode_step_batch_sampled(e, tokens, caches, &[])?;
504 Ok(rows)
505 }
506
507 /// `decode_step_batch` + DEVICE-SIDE SAMPLING for eligible rows (the batched-tick lever,
508 /// 2026-08-01): the host sampler's temp-path is O(n_vocab) with a full-vocab exp per row
509 /// (measured 1.36 ms/row at the 9B's 248320 vocab = 10.9 ms/tick at B=8 — the single
510 /// largest component of the serving tick). Here each requested row samples ON DEVICE
511 /// between the lm_head matmul and the logits D2H:
512 /// temp <= 0 (greedy): the 2-pass device argmax — bit-identical to host argmax
513 /// (argmax-gate contract, same kernels as the dc serving path).
514 /// temp > 0: gumbel_perturb(seed, ctr, temp) + the same argmax = ONE categorical draw
515 /// from softmax(logits/temp) — the sampled-spec Philox machinery. Deterministic per
516 /// (seed, ctr) and INDEPENDENT of batch composition (the isolation contract;
517 /// decode-batch-gate gate3). NOTE: the draw stream differs from the host sampler's
518 /// SplitMix64 (distribution-equal, seed-deterministic, NOT byte-equal to the old
519 /// host draws) — greedy rows are unchanged bit-exact.
520 /// `samp[bi] = Some((temp, seed, ctr))` requests a device sample for row bi; the full
521 /// logits rows are still returned (worker keeps last_logits semantics + fallback rows).
522 pub fn decode_step_batch_sampled(
523 &self,
524 e: &Engine,
525 tokens: &[u32],
526 caches: &mut [&mut Cache],
527 samp: &[Option<DevSamp>],
528 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
529 self.decode_step_batch_sampled_lean(e, tokens, caches, samp, false)
530 }
531
532 /// `decode_step_batch_sampled` + LEAN LOGITS (increment 2 component 3, 2026-08-01):
533 /// with `lean`, device-sampled rows SKIP the [n_vocab] logits D2H (9.4%/32.5% of the
534 /// pre-/post-inc2 tick profile) — their returned row is EMPTY. The audit-mapped
535 /// consumers: (a) the next tick's host sample — never fires, `device_next` carries the
536 /// token; (b) the graph-promotion argmax — reads only prefill logits (generated empty);
537 /// (c) the KV-reuse pool park at retire — the REAL consumer, served by a per-cache
538 /// device park: the row is dtod-copied into `cache.last_logits_dev` (device bandwidth)
539 /// and D2H'd ONCE at retire by the worker. Rows without a device sample keep a per-row
540 /// D2H. `lean=false` is bit-for-bit the previous method (gates + non-serving callers).
541 pub fn decode_step_batch_sampled_lean(
542 &self,
543 e: &Engine,
544 tokens: &[u32],
545 caches: &mut [&mut Cache],
546 samp: &[Option<DevSamp>],
547 lean: bool,
548 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
549 self.decode_step_batch_sampled_lean_masked(e, tokens, caches, samp, &[], lean)
550 }
551
552 /// `decode_step_batch_sampled_lean` + GRAMMAR MASKS (constrained decoding, 2026-08-03):
553 /// `masks[bi] = Some((packed_bitset, words))` bans every unset-bit vocab id on row bi
554 /// (mask_logits_f32, -FLT_MAX) BETWEEN the lm_head matmul and the device sampler, so a
555 /// constrained row rides the SAME device-sample/lean-logits tick as everyone else — no
556 /// full-row D2H, no host O(n_vocab) sample. Contract: a masked row must also request a
557 /// device sample. The row's PRISTINE logits are preserved for their consumers before the
558 /// in-place ban: lean rows park the unmasked row into `cache.last_logits_dev` (the
559 /// retire-time reuse-pool park stays unmasked — continuations resume grammar-free, the
560 /// v1 host-path contract), non-lean rows D2H the unmasked row. `masks = &[]` is
561 /// bit-for-bit the unmasked method.
562 pub fn decode_step_batch_sampled_lean_masked(
563 &self,
564 e: &Engine,
565 tokens: &[u32],
566 caches: &mut [&mut Cache],
567 samp: &[Option<DevSamp>],
568 masks: &[Option<(&CudaSlice<u32>, usize)>],
569 lean: bool,
570 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
571 self.decode_step_batch_sampled_lean_masked_schedule(
572 e, tokens, caches, samp, masks, lean, None,
573 )
574 }
575
576 /// Worker-scheduled twin of [`Self::decode_step_batch_sampled_lean_masked`]. The worker
577 /// supplies the balanced dual-wave boundary it used when forming this tick. Direct engine
578 /// callers keep the automatic midpoint above; the explicit seam makes scheduler chunking and
579 /// engine execution one checked contract instead of two coincident width calculations.
580 pub fn decode_step_batch_sampled_lean_masked_scheduled(
581 &self,
582 e: &Engine,
583 tokens: &[u32],
584 caches: &mut [&mut Cache],
585 samp: &[Option<DevSamp>],
586 masks: &[Option<(&CudaSlice<u32>, usize)>],
587 lean: bool,
588 dual_wave_mid: usize,
589 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
590 self.decode_step_batch_sampled_lean_masked_schedule(
591 e,
592 tokens,
593 caches,
594 samp,
595 masks,
596 lean,
597 Some(dual_wave_mid),
598 )
599 }
600
601 #[allow(clippy::too_many_arguments)]
602 fn decode_step_batch_sampled_lean_masked_schedule(
603 &self,
604 e: &Engine,
605 tokens: &[u32],
606 caches: &mut [&mut Cache],
607 samp: &[Option<DevSamp>],
608 masks: &[Option<(&CudaSlice<u32>, usize)>],
609 lean: bool,
610 scheduled_dual_mid: Option<usize>,
611 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
612 // NOTE (inc3 3c, 2026-08-01, KILLED ARM): a deferred-token-readback variant (all
613 // chunks of a tick writing device-sampled tokens into one shared buffer, ONE
614 // dtoh_u32 after the last chunk instead of one per chunk) measured FLAT at serve
615 // level on the 5090 (N=4 medians within +-0.7% at c=8/16/32 — 3 saved syncs
616 // against a ~100 ms weight-bound tick is ~0.1%, below resolution). Killed per the
617 // flags doctrine; receipts research/batched-tick-inc3-20260801 (serve-points.jsonl
618 // base vs defer arms) are the record. The per-chunk [B]-u32 readback below IS the
619 // tick's only steady-state D2H — one per chunk, none per seq.
620 let b_n = tokens.len();
621 assert!(
622 b_n >= 1 && b_n == caches.len(),
623 "tokens/caches length mismatch"
624 );
625 // ---- PP DOOR: THE BATCHED STAGE SPLIT (pp2-batch 2026-08-06) ----------------------
626 // Until this increment this body had NO pp arm: it walked lo=0..n_layers on the
627 // primary engine's stream, with no stage split, no boundary, and no `rt.enter()`. With
628 // the door open and a sharded cross-device placement, every projection for the remote
629 // stages' layers was read over PCIe, per step, silently — measured 7.4 vs 208.9 tok/s
630 // at B=1 (28x), 47.4 vs 657.0 at B=8 (13.9x) on a PRO 6000 pair over Gen5 x16 P2P.
631 // Nothing failed or warned, because peer reads return identical bytes and all three
632 // `decode-batch-gate` gates PASS on that config — the failure mode was performance,
633 // and a green exactness battery hid it. `pp2-hardening` made that regime FAIL CLOSED
634 // (research/pp2-hardening-20260806); this lane makes it legitimately split, so the
635 // refusal lifts for the batched path.
636 //
637 // `decode_step_batch_ppn` runs each stage's layer range through that stage's engine
638 // and stream with a [B, n_embd] boundary transfer between them, i.e. every stage
639 // touches only LOCAL weights and LOCAL cache state. The refusal below still guards
640 // the residue: the door open with `MEMRA_PP_STREAMS=0` (the same-stream rollback,
641 // which also disables the sharded loader, so nothing is remote — `pp_shard_off` and
642 // `pp2_streams_off` both make `pp_sharded_cross_device()` false) or a placement whose
643 // PpNRt fails to build. Keeping the call means a future path that reaches here in a
644 // remote regime still refuses instead of regressing 28x.
645 if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
646 if !crate::pp::pp2_streams_off() && crate::pp::batch_pp_on() {
647 // Auto (flipped default) routes dual only in the re-gated regime and
648 // degrades serially elsewhere; Forced keeps every ineligible placement on
649 // the refusing dual body so the binding negative cells stay reachable.
650 let route_dual = crate::pp::dual_pp_route(
651 crate::pp::dual_pp_mode(),
652 b_n,
653 fence.len() - 1,
654 crate::pp::pp2_overlap(),
655 crate::pp::pp_host_bounce_active(),
656 );
657 if route_dual {
658 let mid = scheduled_dual_mid
659 .or_else(|| crate::pp::dual_pp_wave_mid(b_n))
660 .expect("dual PP B>=2 must have a wave midpoint");
661 return self
662 .decode_step_batch_dual(e, tokens, caches, samp, masks, lean, &fence, mid);
663 }
664 return self.decode_step_batch_ppn(e, tokens, caches, samp, masks, lean, &fence);
665 }
666 }
667 if scheduled_dual_mid.is_some() {
668 return Err(
669 "decode_step_batch: worker supplied a dual-wave schedule but the PP-2 dual path is unavailable"
670 .into(),
671 );
672 }
673 crate::pp::refuse_unsplit_if_remote(
674 "decode_step_batch",
675 "drop MEMRA_PP_STREAMS=0 / MEMRA_BATCH_PP=0 so the batched path takes its OWN \
676 stage split (decode_step_batch_ppn), or serve single-stream over the eager pp \
677 arm (decode_step_h), which is also split",
678 )?;
679 // ---- H3: B=1 FAST-PATH (serve-path phase 2, 2026-08-05) ----------------------------
680 // At b_n==1 every projection below calls `matmul_pre(.., b_n)` with m=1, which is
681 // ALREADY the m=1 mmvq dispatch — so the m=1 *kernel family* was never the gap. What
682 // this body does NOT have is the m=1 *fusion chain* that `decode_step_h` carries:
683 // - the cross-layer add+norm+quantize fusion (`add_rms_norm_q8_1`: 3 launches -> 1),
684 // - the fused SwiGLU epilogue (`silu_mul_scaled_q8_1`: folds ffn_down's quantize
685 // into its producer) and, with it, `matmul_pre_dual_noscale`'s gate+up pair
686 // fusion — i.e. phase-1 LEVER 1.
687 // Routing b_n==1 through `decode_layers_eager` (the SHARED trunk `decode_step_h` and
688 // the ppN stages already use, lifted verbatim — not a copy) makes every present and
689 // future m=1 lever fire on the opt-in path automatically. The epilogue (grammar mask ->
690 // device sample -> lean logits park) stays exactly as the batched path runs it; the trunk's
691 // different FP composition is why this path cannot be a load-changing default.
692 // BIT-IDENTITY: the trunk is the same function `decode_step_h` calls, and every
693 // fusion it enables is kernel-check-pinned bit-identical to its unfused sequence
694 // (add_rms_norm == add;rms_norm | _q8_1 == +quantize_q8_1 | dual_noscale == two
695 // matmul_pre_noscale). Gate: decode-batch-gate B=1 vs decode_step_h + serve stream
696 // identity. MEMRA_SERVE_B1FAST=1 is the fixed-solo opt-in/A-B seam; the default
697 // stays on this function's generic body so batch-width changes cannot change the
698 // FP program mid-request.
699 if b_n == 1
700 && Self::b1_fast_on()
701 && self.b1_fast_arch_eligible()
702 && !self.is_gemma4_e4b()
703 && self.cfg.gemma4.is_none()
704 && self.cfg.m3.is_none()
705 && crate::pp::pp_cuts(self.layers.len()).is_none()
706 && !e.verify_exact_on()
707 {
708 return self.decode_step_b1_fast(e, tokens[0], caches, samp, masks, lean);
709 }
710 // MEMRA_DECODE_BATCH_CAP (experimental door, serving-lane tier probe 2026-08-01):
711 // default 8 keeps the v1 exactness policy — B=2..8 rides the verify-tier batched
712 // mmvq arms, per-row bit-identical to isolated m=1 decode. Values >8 are a
713 // MEASUREMENT DOOR ONLY: m=9..15 falls to the grid.y=m dp4a tail (m weight
714 // re-reads + a different reduce shape) and m>=16 crosses into the GEMM tier
715 // (block-scale f32 rounding) — BOTH break the "byte-identical to isolated"
716 // serving contract. Never default this above 8 without the batched-tier
717 // exactness policy landing.
718 let cap = Self::decode_batch_cap();
719 // EXACT-16 TIER (increment 3a): chunks of 9..=16 are admitted WITHOUT the env door
720 // when every matmul has a bit-exact b16-class kernel (see decode_batch_exact16_ok).
721 // The verify_exact scope below pins that dispatch for the whole step: it turns off
722 // the m>=16 GEMM arms (qmatvec_gemm + MMQ + fp8/f16/fp4 — all block-scale/foreign
723 // numeric configs) so every projection rides the batched-mmvq b16 tier, which is
724 // per-(token,row) bit-identical to isolated m=1 decode (gate2 bit-strength PASS at
725 // B=12/16, s32+s160, 5090 receipts research/batched-tick-inc3-20260801). Without
726 // the exact tier, B>cap stays refused; the env door (MEMRA_DECODE_BATCH_CAP) keeps
727 // its old meaning as the non-exact measurement probe.
728 let exact16 = b_n > 8 && b_n <= 16 && self.decode_batch_exact16_ok();
729 assert!(
730 b_n <= cap || exact16,
731 "decode_step_batch: B={b_n} > cap {cap} with no exact tier — refused. Either \
732 B>16 (there is NO exact kernel class above 16: m>16 crosses GEMM/dp4a numeric \
733 configs; the serve scheduler chunks wider concurrency into <=16 groups instead), \
734 or some matmul in this checkpoint has no bit-exact b16 kernel — run with \
735 MEMRA_EXACT16_WHY=1 to see which tensor and qtype refuses"
736 );
737 struct ExactScope<'a>(&'a Engine, bool);
738 impl Drop for ExactScope<'_> {
739 fn drop(&mut self) {
740 if self.1 {
741 self.0.set_verify_exact(false);
742 }
743 }
744 }
745 let _exact_scope = ExactScope(e, exact16);
746 if exact16 {
747 e.set_verify_exact(true);
748 }
749 // gemma4: NO batched arm at any B (per-layer SWA/global geometry, hd-512 MQA globals,
750 // weightless V-norm, softcapped head — none of it in the generic body below). This was
751 // an assert until 2026-08-07: one serve request panicked the worker, the respawn
752 // re-panicked on the queued request, and the process FATALed
753 // (research/gemma4-serve-20260807/raw/repro-panic-server-*.log). The worker now routes
754 // gemma4 sessions to the per-session eager loop and never calls here; this Err is the
755 // defense-in-depth backstop — a future path that reaches it refuses PER-REQUEST
756 // instead of killing the process. The eager arm (gemma4_decode_step_h) is the
757 // supported decode.
758 if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
759 // BATCHED ARM (lane/gemma-batched, 2026-08-16): the dense 31B gets its own
760 // per-session batched walk (gemma4_decode_batch) — DEFAULT ON since the owner
761 // flip (MEMRA_GEMMA4_BATCH=0 = the eager kill switch). Same shape law as
762 // step35: projections/norms/rope/FFN/head run at m=B (one weight stream, B
763 // rows — decode is weight-BW-bound), KV append + fa_decode stay a per-session
764 // loop (each session's own len drives its SWA/global view). E4B keeps its
765 // dedicated decode; it never enters here.
766 if self.cfg.gemma4.is_some() && !self.is_gemma4_e4b() && Self::gemma4_batch_on() {
767 return self.gemma4_decode_batch(e, tokens, caches, samp, masks, lean);
768 }
769 return Err(
770 "decode_step_batch has no gemma4 arm for this model class (per-layer \
771 swa/global geometry, softcapped head; the dense-31B batched arm is \
772 default-on, MEMRA_GEMMA4_BATCH=0 forces eager) — serve gemma4 on the \
773 eager per-session path"
774 .into(),
775 );
776 }
777 // step35 (lane/step35-batched-decode, 2026-08-08): its OWN batched walk. The generic
778 // body below is the uniform Full arm — global n_head, 128-dim rope on every layer, no
779 // SWA window, no head-wise gate — which on step35 produced HTTP-200 GARBAGE at c>1
780 // (research/step-sku-20260807/raw/b2ab-pre-*.log), so step35 NEVER enters it at any B.
781 // `step35_decode_batch_layers` carries the real geometry: per-layer n_head (64/96),
782 // partial rope (64 full / 128 SWA, dual base, rope_freqs on FULL only), per-SESSION
783 // SWA view offsets from each session's own kvl.len, the separate head-wise gate at
784 // m=B, and the sigmoid-router MoE via the same moe_ffn_il_zq8 the eager path uses.
785 // MEMRA_STEP35_BATCH=0 = the fail-closed rollback seam. The server caps chunks at
786 // B=1; on PP-N the B=1 correctness default also refuses the eager numeric class, while
787 // an unsplit deployment can still use its existing eager B=1 route.
788 if self.cfg.step35.is_some() {
789 if !Self::step35_batch_on() {
790 return Err(
791 "step35 batched decode is disabled (MEMRA_STEP35_BATCH=0) — \
792 only a non-PP eager B=1 route remains available"
793 .into(),
794 );
795 }
796 let n_embd = self.cfg.n_embd as usize;
797 let eps = self.cfg.rms_eps;
798 let mut ph_last = std::time::Instant::now();
799 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
800 let pos_d = e.htod_i32(&pos_v)?;
801 let x = e.htod(&self.embd.gather(n_embd, tokens))?;
802 ph_mark(e, 0, &mut ph_last)?;
803 let x = self.step35_decode_batch_layers(
804 e,
805 x,
806 caches,
807 &pos_d,
808 0,
809 self.layers.len(),
810 &mut ph_last,
811 )?;
812 let mut hn = e.uninit(b_n * n_embd)?;
813 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
814 let logits = e.matmul(&self.output, &hn, b_n)?;
815 ph_mark(e, 10, &mut ph_last)?;
816 return self.decode_batch_epilogue(
817 e,
818 caches,
819 samp,
820 masks,
821 lean,
822 logits,
823 b_n,
824 &mut ph_last,
825 );
826 }
827 let n_embd = self.cfg.n_embd as usize;
828 let eps = self.cfg.rms_eps;
829
830 // MEMRA_BATCH_PHASE=1: sync-bounded phase accumulation (diagnostics — see header note).
831 // Initialized BEFORE the tick-input assembly below so slot 0 covers the HOST side of
832 // setup (pos_v/ptr-table builds, embed gather) as well as the H2D sync — the audit-fix
833 // lane's Q6 instrumentation gap (research/audit-fixes2-20260805): the old placement
834 // started the clock after the assembly, so slot 0 under-reported setup.
835 let mut ph_last = std::time::Instant::now();
836
837 // Per-row rope positions (each sequence at its own depth).
838 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
839 let pos_d = e.htod_i32(&pos_v)?;
840
841 // Per-step, whole-trunk layer context: state pointer table + arm picks. Under a pp
842 // split this call is made once PER STAGE with that stage's engine and range instead
843 // (see `batch_layer_ctx`'s doc for why the table cannot be shared across devices).
844 let n_layers = self.layers.len();
845 let ctx = self.batch_layer_ctx(e, caches, 0, n_layers)?;
846
847 // Embed all B tokens -> x [B, n_embd] (host gather, one H2D).
848 let x = e.htod(&self.embd.gather(n_embd, tokens))?;
849 ph_mark(e, 0, &mut ph_last)?;
850
851 let x = self.decode_batch_layers(e, x, caches, &ctx, &pos_d, &mut ph_last)?;
852
853 // ---- output norm + lm_head at m=B, one D2H ----
854 let mut hn = e.uninit(b_n * n_embd)?;
855 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
856 let logits = e.matmul(&self.output, &hn, b_n)?;
857 ph_mark(e, 10, &mut ph_last)?;
858
859 self.decode_batch_epilogue(e, caches, samp, masks, lean, logits, b_n, &mut ph_last)
860 }
861
862 /// DUAL-ACTIVE PP-2 DECODE (increment 0): split one batch into wave A/B and drive
863 /// stage 0(B) from a scoped host walker while this thread drives stage 1(A). Step's
864 /// per-layer router readback synchronizes the host, so two CUDA streams issued by one
865 /// host thread would remain serial; this mirrors the proven prime PP-2 host schedule.
866 ///
867 /// This arm is the naked PP-2 default since the 2026-08-11 owner flip (`MEMRA_DUAL_PP`
868 /// unset = Auto; `0` is the serial rollback seam). It is fail-closed unless the
869 /// double-slot door is open, prewarms both slots, and uses `tx_pipelined` exclusively.
870 #[allow(clippy::too_many_arguments)]
871 fn decode_step_batch_dual(
872 &self,
873 e: &Engine,
874 tokens: &[u32],
875 caches: &mut [&mut Cache],
876 samp: &[Option<DevSamp>],
877 masks: &[Option<(&CudaSlice<u32>, usize)>],
878 lean: bool,
879 fence: &[usize],
880 mid: usize,
881 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
882 let b_n = tokens.len();
883 assert!(
884 b_n >= 1 && b_n == caches.len(),
885 "tokens/caches length mismatch"
886 );
887 let Some(expected_mid) = crate::pp::dual_pp_wave_mid(b_n) else {
888 return self.decode_step_batch_ppn(e, tokens, caches, samp, masks, lean, fence);
889 };
890 if mid != expected_mid {
891 return Err(format!(
892 "decode_step_batch_dual: worker midpoint {mid} is not the balanced midpoint {expected_mid} for B={b_n}"
893 ).into());
894 }
895 if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
896 return Err(
897 "decode_step_batch_dual has no gemma4 arm — serve gemma4 on the eager \
898 per-session path"
899 .into(),
900 );
901 }
902 assert!(
903 samp.is_empty() || samp.len() == b_n,
904 "decode_step_batch_dual: samp must be empty or have one entry per row"
905 );
906 assert!(
907 masks.is_empty() || masks.len() == b_n,
908 "decode_step_batch_dual: masks must be empty or have one entry per row"
909 );
910
911 let cap = Self::decode_batch_cap();
912 let max_wave = mid.max(b_n - mid);
913 let exact16 = max_wave > 8 && max_wave <= 16 && self.decode_batch_exact16_ok();
914 if max_wave > cap && !exact16 {
915 return Err(format!(
916 "decode_step_batch_dual: B={b_n} waves {mid}+{} exceed per-wave cap {cap} with no exact tier — refused",
917 b_n - mid,
918 ).into());
919 }
920 let n_st = fence.len() - 1;
921 crate::pp::dual_pp_eligibility(
922 n_st,
923 crate::pp::pp2_overlap(),
924 crate::pp::pp_host_bounce_active(),
925 )
926 .map_err(|msg| -> Box<dyn std::error::Error> { msg.into() })?;
927 let rt = crate::pp::PpNRt::get(e)?;
928 assert_eq!(
929 rt.n_stages(),
930 n_st,
931 "PpNRt stage count {} != fence stages {n_st}",
932 rt.n_stages()
933 );
934 let caller_stream = e.stream();
935 rt.fence_stages_behind(&caller_stream)?;
936
937 let n_embd = self.cfg.n_embd as usize;
938 let wave_cap = mid.max(b_n - mid) * n_embd;
939 rt.prepare_overlap_slots(0, wave_cap)?;
940
941 // EXACT-16 is a property of either scheduled wave, not the combined live width. Keep
942 // the scope live across both host walkers and set it on both stage-owned Engines.
943 struct ExactScopeN<'a>(Vec<&'a Engine>);
944 impl Drop for ExactScopeN<'_> {
945 fn drop(&mut self) {
946 for eng in &self.0 {
947 eng.set_verify_exact(false);
948 }
949 }
950 }
951 let _exact_scope = if exact16 {
952 let engines: Vec<&Engine> = (0..n_st).map(|s| rt.engine(s, e)).collect();
953 for eng in &engines {
954 eng.set_verify_exact(true);
955 }
956 Some(ExactScopeN(engines))
957 } else {
958 None
959 };
960
961 let step35_batched = self.cfg.step35.is_some();
962 if step35_batched && !Self::step35_batch_on() {
963 return Err(
964 "step35 batched decode is disabled (MEMRA_STEP35_BATCH=0) — \
965 dual-active PP-2 decode has no correct fallback trunk"
966 .into(),
967 );
968 }
969
970 let (tokens_a, tokens_b) = tokens.split_at(mid);
971 let (caches_a, caches_b) = caches.split_at_mut(mid);
972 let (samp_a, samp_b) = if samp.is_empty() {
973 (&[][..], &[][..])
974 } else {
975 samp.split_at(mid)
976 };
977 let (masks_a, masks_b) = if masks.is_empty() {
978 (&[][..], &[][..])
979 } else {
980 masks.split_at(mid)
981 };
982
983 let (slot_a, ph_a, span_a0) = self.decode_step_batch_dual_stage0(
984 e,
985 rt,
986 tokens_a,
987 caches_a,
988 fence,
989 step35_batched,
990 false,
991 )?;
992
993 static LOGGED: std::sync::Once = std::sync::Once::new();
994 LOGGED.call_once(|| {
995 eprintln!("[dual-pp] dual-active PP-2 decode engaged (naked default since 2026-08-11; two waves)");
996 });
997
998 let (out_a, out_b, span_b0, span_b1) = std::thread::scope(
999 |scope| -> Result<_, Box<dyn std::error::Error>> {
1000 let stage0_b = scope.spawn(move || {
1001 let staged = self
1002 .decode_step_batch_dual_stage0(
1003 e,
1004 rt,
1005 tokens_b,
1006 caches_b,
1007 fence,
1008 step35_batched,
1009 true,
1010 )
1011 .map_err(|err| err.to_string())?;
1012 Ok::<_, String>((staged, caches_b))
1013 });
1014
1015 let out_a = self.decode_step_batch_dual_stage1(
1016 e,
1017 rt,
1018 slot_a,
1019 caches_a,
1020 samp_a,
1021 masks_a,
1022 lean,
1023 fence,
1024 step35_batched,
1025 ph_a,
1026 true,
1027 )?;
1028 let ((slot_b, ph_b, span_b0), caches_b) = stage0_b
1029 .join()
1030 .map_err(|_| "dual PP stage-0 wave-B host walker panicked")?
1031 .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
1032 if !crate::pp::record_dual_pp_slot_pair(slot_a, slot_b) {
1033 return Err(format!(
1034 "decode_step_batch_dual: refused: wave A and B both selected boundary slot {slot_a}"
1035 ).into());
1036 }
1037 let (out_b, span_b1) = self.decode_step_batch_dual_stage1(
1038 e,
1039 rt,
1040 slot_b,
1041 caches_b,
1042 samp_b,
1043 masks_b,
1044 lean,
1045 fence,
1046 step35_batched,
1047 ph_b,
1048 false,
1049 )?;
1050 Ok((out_a, out_b, span_b0, span_b1))
1051 },
1052 )?;
1053
1054 // Wave B is the final producer. One event publishes all last-stage work back to the
1055 // caller after both epilogues, preserving the ordinary PP-N exit law.
1056 rt.publish_to(1, &caller_stream)?;
1057 let (out_a, span_a1) = out_a;
1058 for (stage, span) in [span_a0, span_a1, span_b0, span_b1].into_iter().enumerate() {
1059 if let Some((start, end)) = span {
1060 crate::pp::record_dual_pp_stage_result(stage, start.elapsed_ms(&end));
1061 }
1062 }
1063 let (mut rows, mut next) = out_a;
1064 rows.extend(out_b.0);
1065 next.extend(out_b.1);
1066 Ok((rows, next))
1067 }
1068
1069 #[allow(clippy::too_many_arguments)]
1070 fn decode_step_batch_dual_stage0(
1071 &self,
1072 e: &Engine,
1073 rt: &crate::pp::PpNRt,
1074 tokens: &[u32],
1075 caches: &mut [&mut Cache],
1076 fence: &[usize],
1077 step35_batched: bool,
1078 track_overlap: bool,
1079 ) -> Result<(usize, std::time::Instant, DualPpCudaSpan), Box<dyn std::error::Error>> {
1080 let b_n = tokens.len();
1081 let n_embd = self.cfg.n_embd as usize;
1082 let mut ph_last = std::time::Instant::now();
1083 rt.bind_stage(0)?;
1084 let _st0 = rt.enter(0);
1085 let e0 = rt.engine(0, e);
1086 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1087 let pos_d = e0.htod_i32(&pos_v)?;
1088 let x = e0.htod(&self.embd.gather(n_embd, tokens))?;
1089 ph_mark(e0, 0, &mut ph_last)?;
1090 let timing_start = dual_pp_timing_event(e0, "stage0 start event");
1091 let x = {
1092 let _overlap = track_overlap.then(crate::pp::enter_dual_pp_stage);
1093 if step35_batched {
1094 self.step35_decode_batch_layers(
1095 e0,
1096 x,
1097 caches,
1098 &pos_d,
1099 fence[0],
1100 fence[1],
1101 &mut ph_last,
1102 )?
1103 } else {
1104 let ctx = self.batch_layer_ctx(e0, caches, fence[0], fence[1])?;
1105 self.decode_batch_layers(e0, x, caches, &ctx, &pos_d, &mut ph_last)?
1106 }
1107 };
1108 let timing = timing_start
1109 .and_then(|start| dual_pp_timing_event(e0, "stage0 end event").map(|end| (start, end)));
1110 let slot = rt.tx_pipelined(0, &x, b_n * n_embd)?;
1111 Ok((slot, ph_last, timing))
1112 }
1113
1114 #[allow(clippy::too_many_arguments)]
1115 fn decode_step_batch_dual_stage1(
1116 &self,
1117 e: &Engine,
1118 rt: &crate::pp::PpNRt,
1119 slot: usize,
1120 caches: &mut [&mut Cache],
1121 samp: &[Option<DevSamp>],
1122 masks: &[Option<(&CudaSlice<u32>, usize)>],
1123 lean: bool,
1124 fence: &[usize],
1125 step35_batched: bool,
1126 mut ph_last: std::time::Instant,
1127 track_overlap: bool,
1128 ) -> Result<((Vec<Vec<f32>>, Vec<Option<u32>>), DualPpCudaSpan), Box<dyn std::error::Error>>
1129 {
1130 let b_n = caches.len();
1131 let n_embd = self.cfg.n_embd as usize;
1132 let eps = self.cfg.rms_eps;
1133 rt.bind_stage(1)?;
1134 let _st1 = rt.enter(1);
1135 let e1 = rt.engine(1, e);
1136 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1137 let pos_d = e1.htod_i32(&pos_v)?;
1138 let x = rt.rx(0, slot, b_n * n_embd)?;
1139 let timing_start = dual_pp_timing_event(e1, "stage1 start event");
1140 let x = {
1141 let _overlap = track_overlap.then(crate::pp::enter_dual_pp_stage);
1142 if step35_batched {
1143 self.step35_decode_batch_layers(
1144 e1,
1145 x,
1146 caches,
1147 &pos_d,
1148 fence[1],
1149 fence[2],
1150 &mut ph_last,
1151 )?
1152 } else {
1153 let ctx = self.batch_layer_ctx(e1, caches, fence[1], fence[2])?;
1154 self.decode_batch_layers(e1, x, caches, &ctx, &pos_d, &mut ph_last)?
1155 }
1156 };
1157 let timing = timing_start
1158 .and_then(|start| dual_pp_timing_event(e1, "stage1 end event").map(|end| (start, end)));
1159 let mut hn = e1.uninit(b_n * n_embd)?;
1160 e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
1161 let logits = e1.matmul(&self.output, &hn, b_n)?;
1162 ph_mark(e1, 10, &mut ph_last)?;
1163 Ok((
1164 self.decode_batch_epilogue(e1, caches, samp, masks, lean, logits, b_n, &mut ph_last)?,
1165 timing,
1166 ))
1167 }
1168
1169 /// THE BATCHED PP-N STEP (pp2-batch increment 2, 2026-08-06): the batched tick split
1170 /// across `fence.len()-1` stages, each stage running ONLY its own layer range through
1171 /// ITS OWN engine and stream, with a `[B, n_embd]` boundary activation between them.
1172 /// The batched twin of `decode_step_h_ppn`, and the #1 item on the PP-2 serving bill —
1173 /// without it a >VRAM SKU (Step-3.7-Flash: 105 GB, fits only across two cards) serves
1174 /// SINGLE-STREAM only, because the batched path was the one loop with no stage split.
1175 ///
1176 /// STRUCTURE (mirrors the eager arm exactly, so the two stay comparable):
1177 /// stage 0 `rt.enter(0)` -> per-stage pos_d + embed -> range -> `rt.tx`
1178 /// middle stages `rt.rx` -> per-stage pos_d -> range -> `rt.tx`
1179 /// last stage `rt.rx` -> per-stage pos_d -> range -> output_norm + lm_head ->
1180 /// the batched serving epilogue (masks, device sample, lean park)
1181 ///
1182 /// FOUR THINGS ARE PER-STAGE, and each is per-stage for a measured reason:
1183 ///
1184 /// 1. THE ENGINE (`rt.engine(s, e)`). Not just for the remote device: `Engine` owns
1185 /// lazily-grown stable-pointer scratch pools (`fa_part_pool`, `fa_vf16_scratch`,
1186 /// `argmax_partials`) that are single-stream-safe BY DESIGN. Two stage streams
1187 /// through one Engine is the shared-scratch race the pp2 lane hit (2026-08-02
1188 /// nondeterministic all-logits divergence, 35% flake). `PpNRt::build` already gives
1189 /// every stage s>0 its own Engine even on the primary device, so honouring
1190 /// `rt.engine(s, e)` here is what scopes the pools per stage — the batched path
1191 /// allocates MORE of that scratch than the eager one (fa at m=B), so this is the
1192 /// load-bearing half of the trap's mitigation, not an inherited nicety.
1193 ///
1194 /// 2. THE POINTER TABLE (`batch_layer_ctx(es, caches, lo, hi)`). See [`BatchLayerCtx`]:
1195 /// it holds DEVICE ADDRESSES of that range's cache state, uploaded through that
1196 /// stage's engine. One step-wide table on the primary would put every stage's kernel
1197 /// arguments in stage-0's HBM — a peer read per pointer fetch, the exact cliff this
1198 /// whole lane exists to remove.
1199 ///
1200 /// 3. `pos_d` (the M2 pipelining law, learned on the eager arm): each stage uploads its
1201 /// own copy of the step's per-row positions on ITS stream, so the buffer is
1202 /// allocated, consumed and freed on one stream. A shared stage-0 `pos_d` freed at fn
1203 /// return breaks under deferred readback — the free enqueues on stream 0 while later
1204 /// stages still dereference it.
1205 ///
1206 /// 4. THE HEAD + EPILOGUE run on the LAST stage: `output_norm`/`output` were uploaded
1207 /// through the last stage's engine by the sharded loader (`hybrid.rs`: `e_head =
1208 /// layer_engine(e, n_trunk, n_trunk-1)`), and `cache.last_logits_dev` must be
1209 /// allocated where the logits are.
1210 ///
1211 /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME
1212 /// bytes in the same order — the split only moves where the residual is materialized,
1213 /// and the boundary is a straight f32 copy (dtod same-device / `cudaMemcpyPeerAsync`
1214 /// cross-device, no conversion). So batched PP-N must be BIT-IDENTICAL to single-device
1215 /// batched at the same B, in both placement orders. Gate: `decode-batch-gate --mode
1216 /// pp` (logit-dump, both orders) — the batched analogue of the eager arm's 48 steps x
1217 /// 248,320 f32 logits with zero differing bits.
1218 ///
1219 /// The B=1 fast path is NOT taken here (its condition already excludes an open door):
1220 /// it routes through `decode_layers_eager` whole-trunk on one engine, which is exactly
1221 /// the unsplit walk. B=1 under the door rides this function's B=1 case instead — the
1222 /// same trade the eager arm's own ppn step makes, and the reason the pp2 lane measured
1223 /// B=1 door-open at 0.854x (the lost fusion chain), not a cliff.
1224 #[allow(clippy::too_many_arguments)]
1225 fn decode_step_batch_ppn(
1226 &self,
1227 e: &Engine,
1228 tokens: &[u32],
1229 caches: &mut [&mut Cache],
1230 samp: &[Option<DevSamp>],
1231 masks: &[Option<(&CudaSlice<u32>, usize)>],
1232 lean: bool,
1233 fence: &[usize],
1234 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1235 let b_n = tokens.len();
1236 assert!(
1237 b_n >= 1 && b_n == caches.len(),
1238 "tokens/caches length mismatch"
1239 );
1240 // gemma4: same no-arm refusal as the unsplit body (see decode_step_batch), Err not
1241 // assert — a request must never kill the worker process.
1242 if self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
1243 return Err(
1244 "decode_step_batch_ppn has no gemma4 arm — serve gemma4 on the eager \
1245 per-session path"
1246 .into(),
1247 );
1248 }
1249 // Same width policy as the unsplit body — the stage split changes WHERE kernels run,
1250 // never WHICH tier admits the width. Duplicated deliberately rather than hoisted:
1251 // the exact-16 scope must wrap the whole multi-stage walk (`set_verify_exact` is
1252 // per-Engine state read at dispatch on every stage), so it has to be established
1253 // here, and a shared helper returning a guard would have to own `e` plus the flag.
1254 let cap = Self::decode_batch_cap();
1255 let exact16 = b_n > 8 && b_n <= 16 && self.decode_batch_exact16_ok();
1256 assert!(
1257 b_n <= cap || exact16,
1258 "decode_step_batch_ppn: B={b_n} > cap {cap} with no exact tier — refused"
1259 );
1260 let rt = crate::pp::PpNRt::get(e)?;
1261 let n_st = fence.len() - 1;
1262 assert_eq!(
1263 rt.n_stages(),
1264 n_st,
1265 "PpNRt stage count {} != fence stages {n_st}",
1266 rt.n_stages()
1267 );
1268 // #87 REVERSE PUBLICATION (lane/pp2spec-crash): order every stage stream behind
1269 // the caller before this body's first stage allocation can reuse a pool block
1270 // whose queued primary-stream consumer has not read it yet. Anatomy:
1271 // `PpNRt::fence_stages_behind`. (This body dtoh+syncs its own logits, but its
1272 // PP-mode callers interleave with the spec verify's device-resident outputs in
1273 // the same worker, so the entry fence is the uniform law, not an optimization.)
1274 rt.fence_stages_behind(&e.stream())?;
1275 let n_embd = self.cfg.n_embd as usize;
1276 let eps = self.cfg.rms_eps;
1277 let payload = b_n * n_embd;
1278
1279 // EXACT-16 SCOPE, PER STAGE ENGINE: `verify_exact` is per-Engine state (an AtomicBool
1280 // on the Engine the dispatch reads), and each stage runs through a DIFFERENT Engine —
1281 // so setting it on the primary alone would leave stages 1..N-1 dispatching the m>=16
1282 // GEMM/MMQ arms while stage 0 used the exact b16 tier. That is a silent per-stage
1283 // numeric split (the failure this tier exists to prevent), so the flag is set on
1284 // every stage engine and cleared on all of them at scope exit.
1285 struct ExactScopeN<'a>(Vec<&'a Engine>);
1286 impl Drop for ExactScopeN<'_> {
1287 fn drop(&mut self) {
1288 for eng in &self.0 {
1289 eng.set_verify_exact(false);
1290 }
1291 }
1292 }
1293 let _exact_scope = if exact16 {
1294 let engines: Vec<&Engine> = (0..n_st).map(|s| rt.engine(s, e)).collect();
1295 for eng in &engines {
1296 eng.set_verify_exact(true);
1297 }
1298 Some(ExactScopeN(engines))
1299 } else {
1300 None
1301 };
1302
1303 let mut ph_last = std::time::Instant::now();
1304
1305 // B=1 PER-STAGE FAST PATH (measured 2026-08-06, PRO 6000 pair). The unsplit body's
1306 // b1_fast guard includes `pp_cuts().is_none()`, so opening the pp door dropped every
1307 // solo session off the m=1 FUSION chain (cross-layer add+norm+q8_1, fused SwiGLU,
1308 // lever 1's gate+up dual) and onto the batched m=1 walk. Cost, arm A vs arm C at B=1:
1309 // 208.5 vs 177.3 tok/s = -15.0% — and NOT a split cost, since arm B (stages=2 on ONE
1310 // card) pays the same 177, and the prior lane's `MEMRA_PP_SHARD=0` batched-body B=1
1311 // was 178.5. It was the fusion chain going missing, on the config the Step SKU serves
1312 // solo requests from.
1313 //
1314 // `decode_layers_eager(lo, hi)` is ALREADY range-scoped and is exactly what the eager
1315 // ppn arm (`decode_step_h_ppn`) calls per stage, so B=1 rides the same per-stage
1316 // structure: same engines, same streams, same [1, n_embd] boundary slots, same
1317 // stage-owned caches. Only the trunk kernels differ, and they differ identically to
1318 // how they differ off-door. Exactness is therefore the SAME accepted decode-config FP
1319 // class the unsplit b1_fast lever already carries (strict gate1 PASSes with it on,
1320 // FAILs with it off at maxdiff 1.591e-1) — which is why the pp gate pins
1321 // `set_b1_fast(false)`: with it on, the B=1 reference and the split arm would
1322 // legitimately sit on opposite sides of that gap and the bit-identity arm would
1323 // report a fake stage-split failure.
1324 //
1325 // Step3.5/Step3.7 are an exception (lane/cx-b1fix, 2026-08-10): their B>1 route is
1326 // `step35_decode_batch_layers`, and the live scheduler may move a session from B=1
1327 // to B>1. The eager/fused class and that batched class produce different greedy bytes,
1328 // so selecting the eager arm at B=1 made output depend on load history. Keep one
1329 // numeric class for this model family: Step35 always takes its stage-scoped batched
1330 // trunk at every width. The live transition gate in step35-b2-geometry-gate pins it.
1331 // Qwen35-MoE is the second exception (lane/cx-q35bug, 2026-08-12): on the Q35
1332 // sellgate workload the eager-B1 -> batched-B2 transition changed emitted token ids and
1333 // selected EOS at tokens 15/17/25. Keep that family on this generic batched trunk at B=1
1334 // too; dense Qwen35 retains the measured eager fast path.
1335 let b1_stage_fast = b_n == 1
1336 && Self::b1_fast_on()
1337 && self.b1_fast_arch_eligible()
1338 && !self.is_gemma4_e4b()
1339 && self.cfg.gemma4.is_none()
1340 && self.cfg.m3.is_none()
1341 && self.cfg.step35.is_none()
1342 && !e.verify_exact_on();
1343 // step35 (lane/step35-batched-decode, 2026-08-08): B>1 rides its OWN stage-scoped
1344 // batched walk (`step35_decode_batch_layers`) — the generic `decode_batch_layers`
1345 // remains OFF-LIMITS for this arch at every B (its uniform geometry produced the
1346 // b2ab HTTP-200 garbage: research/step-sku-20260807/raw/b2ab-pre-*.log). Since
1347 // lane/cx-b1fix, B=1 also takes this walk: a Step35 PP-N session must not change
1348 // numeric class when live decode width changes. The refusal below guards the
1349 // rollback residue; under PP-N, disabling the only correct trunk makes Step35
1350 // requests fail closed instead of falling back to the eager class.
1351 let step35_batched = self.cfg.step35.is_some();
1352 if step35_batched && !Self::step35_batch_on() {
1353 return Err(
1354 "step35 batched decode is disabled (MEMRA_STEP35_BATCH=0) — \
1355 PP-N Step35 decode is unavailable because eager B=1 is a different \
1356 numeric class"
1357 .into(),
1358 );
1359 }
1360 // Hoisted: `caches[0].pos` as a value argument alongside `caches[0]` as `&mut` in one
1361 // call is a borrow conflict; `pos` is Copy and the epilogue is what advances it.
1362 let pos0 = if b1_stage_fast { caches[0].pos } else { 0 };
1363
1364 // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
1365 let mut slot = {
1366 let _st0 = rt.enter(0);
1367 let e0 = rt.engine(0, e);
1368 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1369 let pos_d = e0.htod_i32(&pos_v)?;
1370 let x = e0.htod(&self.embd.gather(n_embd, tokens))?;
1371 ph_mark(e0, 0, &mut ph_last)?;
1372 let x = if b1_stage_fast {
1373 self.decode_layers_eager(e0, x, fence[0], fence[1], &pos_d, pos0, caches[0])?
1374 } else if step35_batched {
1375 self.step35_decode_batch_layers(
1376 e0,
1377 x,
1378 caches,
1379 &pos_d,
1380 fence[0],
1381 fence[1],
1382 &mut ph_last,
1383 )?
1384 } else {
1385 let ctx = self.batch_layer_ctx(e0, caches, fence[0], fence[1])?;
1386 self.decode_batch_layers(e0, x, caches, &ctx, &pos_d, &mut ph_last)?
1387 };
1388 rt.tx(0, &x, payload)?
1389 // x + pos_d + ctx.ptr_table drop here: freed stream-ordered on stage-0's stream.
1390 };
1391
1392 // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
1393 for s in 1..n_st - 1 {
1394 let _st = rt.enter(s);
1395 let es = rt.engine(s, e);
1396 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1397 let pos_d = es.htod_i32(&pos_v)?;
1398 let x = rt.rx(s - 1, slot, payload)?;
1399 let x = if b1_stage_fast {
1400 self.decode_layers_eager(es, x, fence[s], fence[s + 1], &pos_d, pos0, caches[0])?
1401 } else if step35_batched {
1402 self.step35_decode_batch_layers(
1403 es,
1404 x,
1405 caches,
1406 &pos_d,
1407 fence[s],
1408 fence[s + 1],
1409 &mut ph_last,
1410 )?
1411 } else {
1412 let ctx = self.batch_layer_ctx(es, caches, fence[s], fence[s + 1])?;
1413 self.decode_batch_layers(es, x, caches, &ctx, &pos_d, &mut ph_last)?
1414 };
1415 slot = rt.tx(s, &x, payload)?;
1416 }
1417
1418 // ---- LAST STAGE: RX + final range + head + the batched serving epilogue ----
1419 let _stl = rt.enter(n_st - 1);
1420 let el = rt.engine(n_st - 1, e);
1421 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1422 let pos_d = el.htod_i32(&pos_v)?;
1423 let x = rt.rx(n_st - 2, slot, payload)?;
1424 let x = if b1_stage_fast {
1425 self.decode_layers_eager(el, x, fence[n_st - 1], fence[n_st], &pos_d, pos0, caches[0])?
1426 } else if step35_batched {
1427 self.step35_decode_batch_layers(
1428 el,
1429 x,
1430 caches,
1431 &pos_d,
1432 fence[n_st - 1],
1433 fence[n_st],
1434 &mut ph_last,
1435 )?
1436 } else {
1437 let ctx = self.batch_layer_ctx(el, caches, fence[n_st - 1], fence[n_st])?;
1438 self.decode_batch_layers(el, x, caches, &ctx, &pos_d, &mut ph_last)?
1439 };
1440
1441 let mut hn = el.uninit(payload)?;
1442 el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
1443 let logits = el.matmul(&self.output, &hn, b_n)?;
1444 ph_mark(el, 10, &mut ph_last)?;
1445
1446 self.decode_batch_epilogue(el, caches, samp, masks, lean, logits, b_n, &mut ph_last)
1447 }
1448
1449 /// Build the per-step layer context for layers `[lo, hi)`: the device state-pointer
1450 /// table plus the step's arm picks. See [`BatchLayerCtx`] for why this is RANGE-scoped
1451 /// (the table holds device addresses and must be uploaded through the engine whose
1452 /// device runs those layers).
1453 ///
1454 /// Table layout is unchanged from the whole-trunk version — `lin_base`/`attn_base` are
1455 /// still indexed by ABSOLUTE layer id, so `decode_batch_layers`' body indexes them
1456 /// exactly as the old inline loop did. Only layers in `[lo, hi)` contribute entries; the
1457 /// rest stay `None`, which is a loud `expect` if a range ever reads outside its own.
1458 pub(crate) fn batch_layer_ctx(
1459 &self,
1460 e: &Engine,
1461 caches: &[&mut Cache],
1462 lo: usize,
1463 hi: usize,
1464 ) -> Result<BatchLayerCtx, Box<dyn std::error::Error>> {
1465 let cfg = &self.cfg;
1466 let head_dim = cfg.head_dim_k as usize;
1467 // Per-step STATE POINTER TABLE (one H2D): for every linear layer, [conv x B]
1468 // [ssm_in x B][ssm_out x B] device addresses. The batched state kernels read their
1469 // sequence's pointer from these arrays — states stay per-cache (no pooling refactor),
1470 // yet conv/prep/scan collapse from 3xB launches per layer to 3. Rebuilt every step
1471 // because the ssm ping-pong swaps pointers host-side after each scan.
1472 // INCREMENT 2 (2026-08-01): the SAME table now also carries, for every FULL-attn
1473 // layer, [k0,v0,k1,v1,...] cache base addresses — the z-batched seqs append and
1474 // seqs fa_decode kernels read their sequence's cache through it (the MoE
1475 // expert-table pattern), collapsing 2xB launches per attn layer to 2.
1476 let mut lin_base: Vec<Option<usize>> = vec![None; self.layers.len()];
1477 let mut attn_base: Vec<Option<usize>> = vec![None; self.layers.len()];
1478 let mut ptrs: Vec<u64> = Vec::new();
1479 {
1480 use cudarc::driver::DevicePtr;
1481 let s = &e.gpu.stream();
1482 for il in lo..hi {
1483 match &self.layers[il].mixer {
1484 Mixer::Linear(_) => {
1485 lin_base[il] = Some(ptrs.len());
1486 for c in caches.iter() {
1487 let rl = c.recur[il].as_ref().unwrap();
1488 let (p, _g) = rl.conv_state.device_ptr(s);
1489 ptrs.push(p as u64);
1490 }
1491 for c in caches.iter() {
1492 let rl = c.recur[il].as_ref().unwrap();
1493 let (p, _g) = rl.ssm_state.device_ptr(s);
1494 ptrs.push(p as u64);
1495 }
1496 for c in caches.iter() {
1497 let rl = c.recur[il].as_ref().unwrap();
1498 let (p, _g) = rl.ssm_state_alt.device_ptr(s);
1499 ptrs.push(p as u64);
1500 }
1501 }
1502 Mixer::Full(_) => {
1503 attn_base[il] = Some(ptrs.len());
1504 for c in caches.iter() {
1505 let kvl = c.kv[il].as_ref().unwrap();
1506 let (pk, _g) = kvl.k.device_ptr(s);
1507 let (pv, _g2) = kvl.v.device_ptr(s);
1508 ptrs.push(pk as u64);
1509 ptrs.push(pv as u64);
1510 }
1511 }
1512 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1513 }
1514 }
1515 }
1516 let ptr_table = if ptrs.is_empty() {
1517 None
1518 } else {
1519 Some(e.htod_u64(&ptrs)?)
1520 };
1521
1522 // INCREMENT 2 arm picks (per STEP — t_kv is layer-invariant within a tick):
1523 // - seqs APPEND: format-only condition (per-row program is t_kv-independent);
1524 // default flash module only (fp8-KV rides the per-seq g-module path).
1525 // - seqs FA: every row must take the v4 eager arm at ITS OWN t_kv AND all rows
1526 // must share ONE fa_split_keys rung (the rows-twins' straddle law) — a rung
1527 // crossing inside the batch keeps the per-seq loop for that step, so each
1528 // sequence always executes the exact program its isolated run would.
1529 // MEMRA_BATCH_APPEND=0 / MEMRA_BATCH_FA=0 are the rollback/A-B seams.
1530 //
1531 // The picks are t_kv-driven, and t_kv is layer-INVARIANT within a step, so every
1532 // stage of a pp split independently computes the SAME arms from the same `caches`
1533 // — a stage cannot silently take a different program than its unsplit self.
1534 let t_kvs: Vec<usize> = caches.iter().map(|c| c.pos + 1).collect();
1535 let t_kv_max = *t_kvs.iter().max().unwrap();
1536 let seqs_append = {
1537 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1538 *ON.get_or_init(|| std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0"))
1539 } && !Engine::kv_fp8_on();
1540 let sp0 = crate::fa_split_keys(t_kvs[0], cfg.n_head_kv as usize);
1541 let seqs_fa = {
1542 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1543 *ON.get_or_init(|| std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0"))
1544 } && t_kvs.iter().all(|&t| crate::fa_seqs_eligible(t, head_dim))
1545 && t_kvs
1546 .iter()
1547 .all(|&t| crate::fa_split_keys(t, cfg.n_head_kv as usize) == sp0);
1548
1549 Ok(BatchLayerCtx {
1550 lin_base,
1551 attn_base,
1552 ptr_table,
1553 t_kvs,
1554 t_kv_max,
1555 sp0,
1556 seqs_append,
1557 seqs_fa,
1558 lo,
1559 hi,
1560 })
1561 }
1562
1563 /// THE PP SEAM (pp2-batch increment 1, 2026-08-06): run the batched trunk over layers
1564 /// `[ctx.lo, ctx.hi)`, entering with a materialized `[B, n_embd]` residual and exiting
1565 /// with the range's final residual materialized. The batched twin of
1566 /// `decode_layers_eager` — the eager arm has had this seam since M1-PP2 and every ppN
1567 /// stage calls it; the batched body had no equivalent, which is why every later PP-2
1568 /// increment (and spec-over-PP2, whose verify is a batched T=K+1 forward) waited on this
1569 /// extraction (`research/pp2-hardening-20260806/PROGRESS.md` bill item 1).
1570 ///
1571 /// SINGLE-DEVICE SEMANTICS ARE UNCHANGED BY CONSTRUCTION: the body is the old
1572 /// `for (il, layer) in self.layers.iter().enumerate()` loop moved verbatim, with `for il
1573 /// in ctx.lo..ctx.hi` as the header and the per-step invariants (`ptr_table`, arm picks,
1574 /// `t_kv`) read from `ctx` instead of enclosing locals. At `lo=0, hi=n_layers` — every
1575 /// call today — the launch sequence is identical, so the exactness contract in this
1576 /// module's header carries over untouched rather than needing a re-proof.
1577 ///
1578 /// UNLIKE the eager seam, this one is NOT yet stage-callable: `caches` is `&mut [&mut
1579 /// Cache]` mutated in place (KV `len` bumps, ssm ping-pong swaps), and `pos_d`/`x` come
1580 /// from the caller's device. Wiring a stage split means per-stage `pos_d` + a boundary
1581 /// `[B, n_embd]` transfer around this call, which is the NEXT increment. The seam exists
1582 /// so that increment is a call-site change, not a 250-line surgery.
1583 #[allow(clippy::too_many_arguments)]
1584 pub(crate) fn decode_batch_layers(
1585 &self,
1586 e: &Engine,
1587 mut x: CudaSlice<f32>,
1588 caches: &mut [&mut Cache],
1589 ctx: &BatchLayerCtx,
1590 pos_d: &CudaSlice<i32>,
1591 ph_last: &mut std::time::Instant,
1592 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
1593 let b_n = caches.len();
1594 let cfg = &self.cfg;
1595 let n_embd = cfg.n_embd as usize;
1596 let eps = cfg.rms_eps;
1597 let (lin_base, attn_base) = (&ctx.lin_base, &ctx.attn_base);
1598 let ptr_table = &ctx.ptr_table;
1599 let (seqs_append, seqs_fa, sp0, t_kv_max) =
1600 (ctx.seqs_append, ctx.seqs_fa, ctx.sp0, ctx.t_kv_max);
1601 debug_assert_eq!(
1602 ctx.t_kvs.len(),
1603 b_n,
1604 "ctx built for a different batch width"
1605 );
1606
1607 for il in ctx.lo..ctx.hi {
1608 let layer = &self.layers[il];
1609 // ---- attn_norm + q8_1 quantize, batched (B rows) ----
1610 let anorm = layer.attn_norm.float_data();
1611 let mut xn = e.uninit(b_n * n_embd)?;
1612 e.rms_norm(&x, anorm, &mut xn, n_embd, b_n, eps)?;
1613 let (hq, hd) = e.quantize_q8_1(&xn, b_n, n_embd)?;
1614
1615 // ---- mixer ----
1616 let mixed: CudaSlice<f32> = match &layer.mixer {
1617 Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
1618 Mixer::Full(fa) => {
1619 let geometry = cfg.full_attention_geometry_at(il as u32);
1620 let n_head = geometry.n_head as usize;
1621 let n_head_kv = geometry.n_head_kv as usize;
1622 let head_dim = geometry.head_dim_k as usize;
1623 let rope_dims = geometry.n_rot as usize;
1624 let rope_base = geometry.rope_base;
1625 let scale = geometry.attention_scale();
1626 // Batched projections: one weight read serves all B rows. At B=1 the
1627 // QKV triple fuses into ONE launch (rig-native decode increment 1 —
1628 // bit-identical per (tensor,row), RIG-NATIVE-DECODE.md); B>1 and
1629 // non-NVFP4 trunks keep the three singles.
1630 let (qf, mut k, v) =
1631 match e.matmul_nvfp4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hd, b_n)? {
1632 Some(t) => t,
1633 None => (
1634 e.matmul_pre(&fa.wq, &hq, &hd, &xn, b_n)?,
1635 e.matmul_pre(&fa.wk, &hq, &hd, &xn, b_n)?,
1636 e.matmul_pre(&fa.wv, &hq, &hd, &xn, b_n)?,
1637 ),
1638 };
1639
1640 let gated =
1641 geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
1642 let (mut q, gate) = if gated {
1643 let mut qs = e.uninit(b_n * n_head * head_dim)?;
1644 let mut gs = e.uninit(b_n * n_head * head_dim)?;
1645 e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, b_n)?;
1646 (qs, Some(gs))
1647 } else {
1648 (qf, None)
1649 };
1650
1651 // QK-norm over B*n_head rows, rope with per-row positions.
1652 let mut qn = e.uninit(b_n * n_head * head_dim)?;
1653 e.rms_norm(
1654 &q,
1655 fa.q_norm.float_data(),
1656 &mut qn,
1657 head_dim,
1658 b_n * n_head,
1659 eps,
1660 )?;
1661 q = qn;
1662 let mut kn = e.uninit(b_n * n_head_kv * head_dim)?;
1663 e.rms_norm(
1664 &k,
1665 fa.k_norm.float_data(),
1666 &mut kn,
1667 head_dim,
1668 b_n * n_head_kv,
1669 eps,
1670 )?;
1671 k = kn;
1672 e.rope_neox(
1673 &mut q, &pos_d, head_dim, rope_dims, n_head, b_n, rope_base, 1.0,
1674 )?;
1675 e.rope_neox(
1676 &mut k, &pos_d, head_dim, rope_dims, n_head_kv, b_n, rope_base, 1.0,
1677 )?;
1678 ph_mark(e, 1, ph_last)?;
1679
1680 // INCREMENT 2 (2026-08-01): the per-seq (append, attend) launch train
1681 // becomes two phases. Phase A appends all B rows (one z-batched launch,
1682 // or the per-seq loop on the seam/fp8 path); phase B attends all B
1683 // sequences (one blockIdx.z launch + one combine on the batched arm —
1684 // which also reads q / writes attn at row offsets, killing the per-seq
1685 // q/a dtod copies — or the per-seq loop when any row is outside the v4
1686 // arm / a split rung crosses inside the batch). Caches are disjoint per
1687 // sequence, so the phase split leaves every row's math untouched.
1688 let q_dim = n_head * head_dim;
1689 let kv_dim = n_head_kv * head_dim;
1690 let mut attn = e.uninit(b_n * q_dim)?;
1691 // ---- phase A: KV append (all B rows) ----
1692 if seqs_append {
1693 let (kdk, kdv, ktb, vtb) = {
1694 let kvl = caches[0].kv[il].as_ref().unwrap();
1695 (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
1696 };
1697 let base = attn_base[il].expect("full layer missing from pointer table");
1698 let table = ptr_table.as_ref().expect("pointer table missing");
1699 let kv_view = table.slice(base..base + 2 * b_n);
1700 e.append_kv_quantized_seqs(
1701 &k, &v, &kv_view, &pos_d, b_n, kdk, kdv, ktb, vtb,
1702 )?;
1703 for cache in caches.iter_mut() {
1704 let kvl = cache.kv[il].as_mut().unwrap();
1705 debug_assert_eq!(kvl.len, cache.pos, "kv len / pos out of lockstep");
1706 kvl.len += 1;
1707 }
1708 } else {
1709 for (bi, cache) in caches.iter_mut().enumerate() {
1710 let kvl = cache.kv[il].as_mut().unwrap();
1711 let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
1712 let v_row = v.slice(bi * kv_dim..(bi + 1) * kv_dim);
1713 e.append_kv_quantized_view(
1714 &k_row,
1715 &v_row,
1716 &mut kvl.k,
1717 &mut kvl.v,
1718 kvl.len,
1719 kvl.kv_dim_k,
1720 kvl.kv_dim_v,
1721 kvl.k_tok_bytes,
1722 kvl.v_tok_bytes,
1723 Engine::kv_fp8_on(),
1724 )?;
1725 kvl.len += 1;
1726 }
1727 }
1728 ph_mark(e, 2, ph_last)?;
1729 // ---- phase B: attention (all B sequences) ----
1730 if seqs_fa {
1731 let (ktb, vtb) = {
1732 let kvl = caches[0].kv[il].as_ref().unwrap();
1733 (kvl.k_tok_bytes, kvl.v_tok_bytes)
1734 };
1735 let base = attn_base[il].expect("full layer missing from pointer table");
1736 let table = ptr_table.as_ref().expect("pointer table missing");
1737 let kv_view = table.slice(base..base + 2 * b_n);
1738 e.fa_decode_batch_seqs_v4(
1739 &q, &kv_view, &pos_d, &mut attn, head_dim, n_head, n_head_kv, b_n,
1740 t_kv_max, scale, sp0, ktb, vtb,
1741 )?;
1742 ph_mark(e, 4, ph_last)?;
1743 } else {
1744 for (bi, cache) in caches.iter_mut().enumerate() {
1745 let kvl = cache.kv[il].as_mut().unwrap();
1746 let t_kv = kvl.len;
1747 let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
1748 let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
1749 // The fallback keeps one FA launch per distinct KV view, but Q and
1750 // attention already live in packed row-major buffers. Pass those row
1751 // views directly; only the arithmetic-free materialization copies go.
1752 let q_row = q.slice(bi * q_dim..(bi + 1) * q_dim);
1753 let mut a_row = attn.slice_mut(bi * q_dim..(bi + 1) * q_dim);
1754 e.fa_decode_kvmod_view(
1755 &q_row,
1756 &k_view,
1757 &v_view,
1758 &mut a_row,
1759 head_dim,
1760 n_head,
1761 n_head_kv,
1762 t_kv,
1763 scale,
1764 kvl.k_tok_bytes,
1765 kvl.v_tok_bytes,
1766 Engine::kv_fp8_on(),
1767 )?;
1768 ph_mark(e, 4, ph_last)?;
1769 }
1770 }
1771
1772 // Output gate (element-wise — batches whole) + o-proj at m=B.
1773 let attn_g = match &gate {
1774 Some(g) => {
1775 let n = b_n * q_dim;
1776 let mut gsig = e.uninit(n)?;
1777 e.sigmoid(g, &mut gsig, n)?;
1778 let mut ag = e.uninit(n)?;
1779 e.mul(&attn, &gsig, &mut ag, n)?;
1780 ag
1781 }
1782 None => attn,
1783 };
1784 let o = e.matmul(&fa.wo, &attn_g, b_n)?;
1785 ph_mark(e, 5, ph_last)?;
1786 o
1787 }
1788 Mixer::Linear(la) => {
1789 // v2 (the B-scaling fix): the GDN mixer's PROJECTIONS carry the layer's
1790 // weight mass — batch them at m=B so wqkv/gate/beta/alpha/ssm_out stream
1791 // ONCE per step instead of once per sequence. Only the recurrent state ops
1792 // (fused conv ring, gdn prep, gdn scan) stay per-seq — they are state-bound
1793 // micro-kernels, not weight readers. Composition unchanged vs v1 (matmul_pre
1794 // == fused2 per (tensor,row); _bN mmvq per-row == m=1): same numeric config.
1795 let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
1796 let d_state = ssm.state_size as usize;
1797 let num_k = ssm.group_count as usize;
1798 let num_v = ssm.time_step_rank as usize;
1799 let d_conv = ssm.conv_kernel as usize;
1800 let key_dim = d_state * num_k;
1801 let value_dim = d_state * num_v;
1802 let conv_dim = key_dim * 2 + value_dim;
1803 let gdn_scale = 1.0 / (d_state as f32).sqrt();
1804
1805 // ---- batched projections (the weight win) ----
1806 // At B=1 the mixer quartet fuses into ONE launch (rig-native decode
1807 // increment 2 — bit-identical per (tensor,row), RIG-NATIVE-DECODE.md);
1808 // B>1 and non-NVFP4 trunks keep the four singles.
1809 let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_nvfp4_fused4(
1810 &la.wqkv,
1811 &la.wqkv_gate,
1812 &la.ssm_beta,
1813 &la.ssm_alpha,
1814 &hq,
1815 &hd,
1816 b_n,
1817 )? {
1818 Some(t) => t,
1819 None => (
1820 e.matmul_pre(&la.wqkv, &hq, &hd, &xn, b_n)?,
1821 e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, b_n)?,
1822 e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, b_n)?,
1823 e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, b_n)?,
1824 ),
1825 };
1826 ph_mark(e, 6, ph_last)?;
1827
1828 // ---- batched recurrent state ops (3 launches for all B sequences) ----
1829 let base = lin_base[il].expect("linear layer missing from pointer table");
1830 let table = ptr_table.as_ref().expect("pointer table missing");
1831 let conv_view = table.slice(base..base + b_n);
1832 let in_view = table.slice(base + b_n..base + 2 * b_n);
1833 let out_view = table.slice(base + 2 * b_n..base + 3 * b_n);
1834 let mut conv_outs = e.uninit(b_n * conv_dim)?;
1835 e.ssm_conv1d_fused_decode_b(
1836 &qkv_mixed,
1837 &conv_view,
1838 la.ssm_conv1d.float_data(),
1839 &mut conv_outs,
1840 conv_dim,
1841 d_conv,
1842 b_n,
1843 )?;
1844 let mut q_l2 = e.uninit(b_n * value_dim)?;
1845 let mut k_l2 = e.uninit(b_n * value_dim)?;
1846 let mut v_gd = e.uninit(b_n * value_dim)?;
1847 let mut beta_b = e.uninit(b_n * num_v)?;
1848 let mut g_log = e.uninit(b_n * num_v)?;
1849 e.gdn_prep_decode_b(
1850 &conv_outs,
1851 &beta_raw,
1852 &alpha,
1853 la.ssm_dt.float_data(),
1854 la.ssm_a.float_data(),
1855 &mut q_l2,
1856 &mut k_l2,
1857 &mut v_gd,
1858 &mut beta_b,
1859 &mut g_log,
1860 d_state,
1861 num_v,
1862 num_k,
1863 key_dim,
1864 eps,
1865 conv_dim,
1866 b_n,
1867 )?;
1868 let mut o_all = e.uninit(b_n * value_dim)?;
1869 e.gdn_scan_s128_batched(
1870 &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_all,
1871 num_v, b_n, gdn_scale,
1872 )?;
1873 // ping-pong: scan wrote each seq's alt buffer; swap host handles (the
1874 // NEXT step's table rebuild picks up the new canonical pointers).
1875 for cache in caches.iter_mut() {
1876 let rl = cache.recur[il].as_mut().unwrap();
1877 std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
1878 }
1879 ph_mark(e, 7, ph_last)?;
1880
1881 // ---- batched gated norm + out-projection ----
1882 let o = if e.uses_q8_1_fast(&la.ssm_out) {
1883 let (gq, gd) = e.gated_rmsnorm_q8_1(
1884 &o_all,
1885 la.ssm_norm.float_data(),
1886 &z,
1887 d_state,
1888 b_n * num_v,
1889 eps,
1890 )?;
1891 let g0 = e.zeros(0)?;
1892 e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, b_n)?
1893 } else {
1894 let mut gn = e.uninit(b_n * value_dim)?;
1895 e.gated_rmsnorm(
1896 &o_all,
1897 la.ssm_norm.float_data(),
1898 &z,
1899 &mut gn,
1900 d_state,
1901 b_n * num_v,
1902 eps,
1903 )?;
1904 e.matmul(&la.ssm_out, &gn, b_n)?
1905 };
1906 ph_mark(e, 8, ph_last)?;
1907 o
1908 }
1909 };
1910
1911 // ---- residual add + post_attn_norm + FFN, batched ----
1912 let pnorm = layer.post_attn_norm.float_data();
1913 let mut x1 = e.uninit(b_n * n_embd)?;
1914 let mut z = e.uninit(b_n * n_embd)?;
1915 e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut z, n_embd, b_n, eps)?;
1916 let ffn_out = match &layer.ffn {
1917 crate::hybrid::Ffn::Dense {
1918 ffn_gate,
1919 ffn_up,
1920 ffn_down,
1921 } => {
1922 // v1 covers the SiLU family; M3's swigluoai clamp rides a scaled epilogue
1923 // (m=1 fused tier) — batched M3 lands with the batched-fusion pass.
1924 assert!(
1925 self.cfg.m3.is_none(),
1926 "decode_step_batch v1: M3 swigluoai FFN not yet batched"
1927 );
1928 let n_ff = ffn_gate.out_features();
1929 let (zq, zd) = e.quantize_q8_1(&z, b_n, n_embd)?;
1930 // REFUTED ARM (lane/q27-deepdive, 2026-08-05): fusing this gate+up pair
1931 // into `matmul_q8_fused2_t` (the fused2_b8 tier) measured FLAT-TO-NEGATIVE
1932 // at the serving tick — bench c=8 213.1/213.8, 213.9/214.4, 214.4/213.5
1933 // (sign flips) and serve c=8 paired mean −0.20% over 3 passes. Mechanism:
1934 // unlike m=1 (where the pair is 128 of 1015 launches in a 7.67%-gap tick),
1935 // the c=8 tick is 73.2% one weight-bound kernel class with launch cost
1936 // already hidden — halving 128 launches of ~28k buys nothing. The m=1 arm
1937 // in `matmul_pre_dual_noscale` (+0.94%) stays; this call site keeps the two
1938 // launches. Kernel + fused2_b8 wrapper retained: kernel-check gates it at
1939 // m=5/8 and matmul_q8_fused2_t serves the verify tier. Receipts:
1940 // research/q27-deepdive-20260805/ (lever3-bench-*, serve-points.jsonl).
1941 let g = e.matmul_pre(ffn_gate, &zq, &zd, &z, b_n)?;
1942 let u = e.matmul_pre(ffn_up, &zq, &zd, &z, b_n)?;
1943 let mut act = e.uninit(b_n * n_ff)?;
1944 e.silu_mul(&g, &u, &mut act, b_n * n_ff)?;
1945 let (aq, ad) = e.quantize_q8_1(&act, b_n, n_ff)?;
1946 e.matmul_pre(ffn_down, &aq, &ad, &act, b_n)?
1947 }
1948 crate::hybrid::Ffn::Moe(m) => {
1949 self.moe_ffn_il_zq8(e, m, &z, None, b_n, il as u16)?
1950 }
1951 };
1952 // next-layer input x = x1 + ffn_out (batched element-wise add)
1953 let mut x2 = e.uninit(b_n * n_embd)?;
1954 e.add(&x1, &ffn_out, &mut x2, b_n * n_embd)?;
1955 x = x2;
1956 ph_mark(e, 9, ph_last)?;
1957 }
1958 Ok(x)
1959 }
1960
1961 /// Rollback seam for the step35 batched decode arm (lane/step35-batched-decode,
1962 /// 2026-08-08). Default ON; `MEMRA_STEP35_BATCH=0` caps serving at B=1 and makes the
1963 /// batched bodies return Err. Since lane/cx-b1fix, PP-N also refuses the eager B=1
1964 /// numeric class, so the seam disables PP-N Step35 decode rather than serving unstable
1965 /// bytes. Also the b2geo35 gate's CANARY seam — the live assertions must fail under it.
1966 pub fn step35_batch_on() -> bool {
1967 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1968 *ON.get_or_init(|| std::env::var("MEMRA_STEP35_BATCH").as_deref() != Ok("0"))
1969 }
1970
1971 /// THE step35 BATCHED LAYER WALK (lane/step35-batched-decode, 2026-08-08): B sequences
1972 /// share one pass over layers `[lo, hi)` with the REAL step35 geometry — the arm that
1973 /// kills the B=1 pin (34 tok/s aggregate FLAT across c=1..8, round-robin serialized;
1974 /// research/step-sku-20260807 §4) without re-opening the b2ab garbage hole (the generic
1975 /// `decode_batch_layers` ran uniform n_head/full-width rope/no window/no gate over
1976 /// step35 weights and returned HTTP-200 garbage at c>1).
1977 ///
1978 /// SHAPE — batched where the weights are, per-session where the state is:
1979 /// * attn_norm + quantize + wq/wk/wv/attn_gate projections + q/k norms + rope + head
1980 /// gate + wo + residual/post-norm + FFN all run at m=B: ONE weight stream serves B
1981 /// rows (decode is weight-BW-bound; this is the entire win).
1982 /// * KV append + fa_decode stay a per-session loop — the SWA window makes each
1983 /// session's KV view a function of ITS OWN `kvl.len` (`off = len-win` when past the
1984 /// window), and the z-batched seqs kernels take one shared t_kv/rung, not per-row
1985 /// offsets. This is the same shape as `decode_batch_layers`' per-seq fallback arm,
1986 /// and it costs launches, not weight bandwidth (KV is per-session state either way).
1987 ///
1988 /// PER-LAYER GEOMETRY (the five mechanisms that make the generic body wrong here, all
1989 /// from `step35_geom`/cfg): n_head 64 full / 96 SWA (wq/wo/attn_gate widths per layer),
1990 /// partial rope (n_rot 64 full / 128 SWA), dual base (5e6/1e4) + `rope_freqs` factors
1991 /// on FULL layers only, SWA window 512 with per-SESSION view offsets, and the separate
1992 /// head-wise `attn_gate` (one pre-sigmoid scalar per (token, head), input = the
1993 /// post-attn_norm hidden, applied before wo).
1994 ///
1995 /// EXACTNESS (the isolation contract, decode-batch-gate gate2's bar): every kernel here
1996 /// is row-independent at m=B or per-session:
1997 /// * `rms_norm`/`add_rms_norm`/`quantize_q8_1`/`attn_head_gate`/activations: per-row
1998 /// programs, grid over rows — row bi's bytes are the 1-row call's bytes.
1999 /// * projections via `matmul_pre` at m=2..8: Q8_0/Q6_K-class rides the b2/b4/b8
2000 /// batched-mmvq tier (bit-identical per (token,row) to m=1 mmvq); IQ4_XS — this
2001 /// SKU's trunk class — has no mmvq/batched kernel, so BOTH m=1 decode and the m=B
2002 /// walk ride `qmatvec_iq4_XS_dp4a` (grid (out_f, m): each column IS the m=1 dp4a
2003 /// program). Same class at every width = the decode-parity law by construction.
2004 /// * `rope_neox2` takes per-row positions (tok = row / n_heads) — row bi rotates at
2005 /// ITS pos with the layer's (n_rot, base, ff), same bits as its solo call.
2006 /// * per-session append/fa_decode_kvmod: literally the eager arm's calls on that
2007 /// session's own cache and views.
2008 /// * MoE (`moe_ffn_il_zq8` at t=B): the router is per-column decode-exact at
2009 /// t < PRIME_MIN_T (m=1 program per column), sigmoid routing + expert dispatch are
2010 /// per-token — a session's experts are a function of its own row only.
2011 /// The known eager-vs-batched FP gap is why PP-N Step35 deliberately serves THIS walk at
2012 /// B=1 too: the scheduler can change width during a session, so one numeric class must
2013 /// cover every live width. `b2geo35` pins static widths and an explicit B=1 -> B>1
2014 /// transition under live defaults.
2015 ///
2016 /// STAGE-SCOPED FROM BIRTH: `[lo, hi)` + caller-supplied engine/pos_d, so
2017 /// `decode_step_batch_ppn` calls it per stage (per-stage engine, per-stage pos_d, the
2018 /// #87 entry fence and boundary slots unchanged) — the pp2-batch seam lesson.
2019 #[allow(clippy::too_many_arguments)]
2020 pub(crate) fn step35_decode_batch_layers(
2021 &self,
2022 e: &Engine,
2023 x: CudaSlice<f32>,
2024 caches: &mut [&mut Cache],
2025 pos_d: &CudaSlice<i32>,
2026 lo: usize,
2027 hi: usize,
2028 ph_last: &mut std::time::Instant,
2029 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2030 self.step35_decode_rows_layers(e, x, caches, pos_d, None, lo, hi, ph_last)
2031 }
2032
2033 /// Diagnostic generalization of the serving walk: `row_to_cache[r]` names the session
2034 /// whose KV row is consumed by hidden row `r`. Serving passes `None`, preserving the
2035 /// identity mapping and its launch sequence. The MoESD harness passes B groups of gamma
2036 /// consecutive rows so each session's verify columns append causally while projections and
2037 /// MoE dispatch see the full B*gamma target width.
2038 #[allow(clippy::too_many_arguments)]
2039 fn step35_decode_rows_layers(
2040 &self,
2041 e: &Engine,
2042 mut x: CudaSlice<f32>,
2043 caches: &mut [&mut Cache],
2044 pos_d: &CudaSlice<i32>,
2045 row_to_cache: Option<&[usize]>,
2046 lo: usize,
2047 hi: usize,
2048 ph_last: &mut std::time::Instant,
2049 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2050 let b_n = row_to_cache.map_or(caches.len(), |rows| rows.len());
2051 let cfg = &self.cfg;
2052 let n_embd = cfg.n_embd as usize;
2053 let eps = cfg.rms_eps;
2054 cfg.step35
2055 .as_ref()
2056 .ok_or("step35_decode_batch_layers requires step35 cfg")?;
2057 if b_n == 0 || x.len() != b_n * n_embd || pos_d.len() != b_n {
2058 return Err(format!(
2059 "step35 row mapping shape mismatch: rows={b_n} x={} pos={} n_embd={n_embd}",
2060 x.len(),
2061 pos_d.len(),
2062 )
2063 .into());
2064 }
2065 if row_to_cache.is_some_and(|rows| rows.iter().any(|&ci| ci >= caches.len())) {
2066 return Err("step35 row mapping names a missing cache".into());
2067 }
2068 let cache_index = |row: usize| row_to_cache.map_or(row, |rows| rows[row]);
2069 // b2geo35 gate evidence: one line, first B>1 walk only (grep-stable prefix).
2070 if b_n > 1 {
2071 static ONCE: std::sync::Once = std::sync::Once::new();
2072 ONCE.call_once(|| {
2073 eprintln!(
2074 "[step35-batch] first B>1 batched step35 walk: B={b_n} layers=[{lo},{hi})"
2075 );
2076 });
2077 }
2078
2079 for il in lo..hi {
2080 let layer = &self.layers[il];
2081 let Mixer::Full(fa) = &layer.mixer else {
2082 return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
2083 };
2084 let geometry = self.step35_geom(il);
2085 let hd = geometry.head_dim_k as usize;
2086 let nkv = geometry.n_head_kv as usize;
2087 let nh = geometry.n_head as usize;
2088 let rbase = geometry.rope_base;
2089 let scale = geometry.attention_scale();
2090 let swa = geometry.window.is_some();
2091 let win = geometry.window.unwrap_or(0) as usize;
2092 let n_rot = geometry.n_rot as usize;
2093 let q_dim = nh * hd;
2094 let kv_dim = nkv * hd;
2095
2096 // ---- attn_norm + q8_1 quantize, batched (B rows) ----
2097 let anorm = layer.attn_norm.float_data();
2098 let mut xn = e.uninit(b_n * n_embd)?;
2099 e.rms_norm(&x, anorm, &mut xn, n_embd, b_n, eps)?;
2100 let (hq, hdq) = e.quantize_q8_1(&xn, b_n, n_embd)?;
2101
2102 // ---- batched projections: q/k/v + the separate head-wise gate (one weight
2103 // stream for B rows; xn is the live f32 fallback for non-q8_1-fast classes) ----
2104 let q0 = e.matmul_pre(&fa.wq, &hq, &hdq, &xn, b_n)?;
2105 let k0 = e.matmul_pre(&fa.wk, &hq, &hdq, &xn, b_n)?;
2106 let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, &xn, b_n)?;
2107 let gw = fa
2108 .attn_gate
2109 .as_ref()
2110 .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
2111 // gate input = the post-attn_norm hidden (upstream `cur`) — same xn/q8 pair.
2112 let gt = e.matmul_pre(gw, &hq, &hdq, &xn, b_n)?;
2113
2114 // ---- q/k RMSNorm over head_dim rows + the per-layer PARTIAL rope ----
2115 let mut q = e.uninit(b_n * q_dim)?;
2116 e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, b_n * nh, eps)?;
2117 let mut k = e.uninit(b_n * kv_dim)?;
2118 e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, b_n * nkv, eps)?;
2119 let ff = if geometry.rope_factors {
2120 self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
2121 } else {
2122 None
2123 };
2124 e.rope_neox2(
2125 &mut q, &mut k, pos_d, hd, n_rot, nh, nkv, b_n, rbase, 1.0, ff,
2126 )?;
2127 ph_mark(e, 1, ph_last)?;
2128
2129 // ---- per-session: KV append + windowed/global fa_decode (each session's OWN
2130 // len drives its view offset — the iso-gap law, no cross-session term) ----
2131 let mut attn = e.uninit(b_n * q_dim)?;
2132 if b_n == 1 {
2133 // B=1 SPECIALIZED ENTRY (lane/cx-eagerpar): the general row loop below
2134 // materializes q_row and a_row because a B>1 FA call consumes/produces one
2135 // contiguous row at a time. At B=1, q and attn already ARE those whole rows.
2136 // Pass them directly to the same fa_decode_kvmod call: this removes two
2137 // arithmetic-free D2D copies (90 launches/token on Step3.7's 45 layers)
2138 // without changing any arithmetic kernel, shape, argument value, or order.
2139 // Keep the B>1 body verbatim below; b1fix's one-class/transition gates are
2140 // the promotion bar, not an FP-similarity tolerance.
2141 let kvl = caches[cache_index(0)].kv[il].as_mut().unwrap();
2142 let k_row = k.slice(0..kv_dim);
2143 let v_row = v0.slice(0..kv_dim);
2144 let next_len = kvl.len + 1;
2145 let (off, t_kv) = if swa && next_len > win {
2146 (next_len - win, win)
2147 } else {
2148 (0, next_len)
2149 };
2150 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
2151 e.append_kv_quantized_view(
2152 &k_row,
2153 &v_row,
2154 &mut kvl.k,
2155 &mut kvl.v,
2156 write_row,
2157 kvl.kv_dim_k,
2158 kvl.kv_dim_v,
2159 kvl.k_tok_bytes,
2160 kvl.v_tok_bytes,
2161 Engine::kv_fp8_on(),
2162 )?;
2163 kvl.len = next_len;
2164 ph_mark(e, 2, ph_last)?;
2165 let physical = kvl.physical_rows(off, off + t_kv)?;
2166 let k_view = e.view_u8_range(
2167 &kvl.k,
2168 physical.start * kvl.k_tok_bytes,
2169 physical.end * kvl.k_tok_bytes,
2170 );
2171 let v_view = e.view_u8_range(
2172 &kvl.v,
2173 physical.start * kvl.v_tok_bytes,
2174 physical.end * kvl.v_tok_bytes,
2175 );
2176 e.fa_decode_kvmod(
2177 &q,
2178 &k_view,
2179 &v_view,
2180 &mut attn,
2181 hd,
2182 nh,
2183 nkv,
2184 t_kv,
2185 scale,
2186 kvl.k_tok_bytes,
2187 kvl.v_tok_bytes,
2188 Engine::kv_fp8_on(),
2189 )?;
2190 ph_mark(e, 4, ph_last)?;
2191 } else {
2192 for bi in 0..b_n {
2193 let cache = &mut caches[cache_index(bi)];
2194 let kvl = cache.kv[il].as_mut().unwrap();
2195 let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
2196 let v_row = v0.slice(bi * kv_dim..(bi + 1) * kv_dim);
2197 let next_len = kvl.len + 1;
2198 let (off, t_kv) = if swa && next_len > win {
2199 (next_len - win, win)
2200 } else {
2201 (0, next_len)
2202 };
2203 let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
2204 e.append_kv_quantized_view(
2205 &k_row,
2206 &v_row,
2207 &mut kvl.k,
2208 &mut kvl.v,
2209 write_row,
2210 kvl.kv_dim_k,
2211 kvl.kv_dim_v,
2212 kvl.k_tok_bytes,
2213 kvl.v_tok_bytes,
2214 Engine::kv_fp8_on(),
2215 )?;
2216 kvl.len = next_len;
2217 ph_mark(e, 2, ph_last)?;
2218 // the eager arm's SWA view arithmetic, verbatim (step35_decode_attn):
2219 // token-aligned offset, keys carry absolute rope, mask is positional.
2220 let physical = kvl.physical_rows(off, off + t_kv)?;
2221 let k_view = e.view_u8_range(
2222 &kvl.k,
2223 physical.start * kvl.k_tok_bytes,
2224 physical.end * kvl.k_tok_bytes,
2225 );
2226 let v_view = e.view_u8_range(
2227 &kvl.v,
2228 physical.start * kvl.v_tok_bytes,
2229 physical.end * kvl.v_tok_bytes,
2230 );
2231 // The per-session cache view remains authoritative (including SWA's
2232 // physical-row rebase), while Q/O use their existing packed row views.
2233 // This preserves the exact FA program and removes only the two D2D copies.
2234 let q_row = q.slice(bi * q_dim..(bi + 1) * q_dim);
2235 let mut a_row = attn.slice_mut(bi * q_dim..(bi + 1) * q_dim);
2236 e.fa_decode_kvmod_view(
2237 &q_row,
2238 &k_view,
2239 &v_view,
2240 &mut a_row,
2241 hd,
2242 nh,
2243 nkv,
2244 t_kv,
2245 scale,
2246 kvl.k_tok_bytes,
2247 kvl.v_tok_bytes,
2248 Engine::kv_fp8_on(),
2249 )?;
2250 ph_mark(e, 4, ph_last)?;
2251 }
2252 }
2253
2254 // ---- head-wise gate (one sigmoid per (token, head), pre-wo) + o-proj at m=B ----
2255 let mut ag = e.uninit(b_n * q_dim)?;
2256 e.attn_head_gate(&attn, >, &mut ag, None, hd, nh, b_n)?;
2257 let mixed = e.matmul(&fa.wo, &ag, b_n)?;
2258 ph_mark(e, 5, ph_last)?;
2259
2260 // ---- residual add + post_attn_norm + FFN, batched ----
2261 let pnorm = layer.post_attn_norm.float_data();
2262 let mut x1 = e.uninit(b_n * n_embd)?;
2263 let mut z = e.uninit(b_n * n_embd)?;
2264 e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut z, n_embd, b_n, eps)?;
2265 let ffn_out = match &layer.ffn {
2266 crate::hybrid::Ffn::Dense {
2267 ffn_gate,
2268 ffn_up,
2269 ffn_down,
2270 } => {
2271 // A dense step35 FFN's clamp is the SHEXP array (upstream's one
2272 // build_ffn serves dense + shared expert, llama-graph.cpp:1751);
2273 // ffn_act_lim dispatches clamped/plain per layer. Layers 0-2 (the
2274 // leading dense) have no live limit on this artifact, but the route
2275 // is correct by construction, not by artifact.
2276 let n_ff = ffn_gate.out_features();
2277 let (zq, zd) = e.quantize_q8_1(&z, b_n, n_embd)?;
2278 let g = e.matmul_pre(ffn_gate, &zq, &zd, &z, b_n)?;
2279 let u = e.matmul_pre(ffn_up, &zq, &zd, &z, b_n)?;
2280 let mut act = e.uninit(b_n * n_ff)?;
2281 Self::ffn_act_lim(
2282 e,
2283 cfg,
2284 &g,
2285 &u,
2286 1.0,
2287 1.0,
2288 cfg.clamp_shexp_at(il as u32),
2289 &mut act,
2290 b_n * n_ff,
2291 )?;
2292 let (aq, ad) = e.quantize_q8_1(&act, b_n, n_ff)?;
2293 e.matmul_pre(ffn_down, &aq, &ad, &act, b_n)?
2294 }
2295 // t=B < PRIME_MIN_T: per-column decode-exact router + host sigmoid routing
2296 // + per-token expert dispatch — the same per-token program as eager t=1,
2297 // including the per-layer SwiGLU clamp (43/44) via the sequential path's
2298 // ffn_act_lim. The sigmoid-router deny on dev/pairs holds by predicate.
2299 crate::hybrid::Ffn::Moe(m) => {
2300 self.moe_ffn_il_zq8(e, m, &z, None, b_n, il as u16)?
2301 }
2302 };
2303 let mut x2 = e.uninit(b_n * n_embd)?;
2304 e.add(&x1, &ffn_out, &mut x2, b_n * n_embd)?;
2305 x = x2;
2306 ph_mark(e, 9, ph_last)?;
2307 }
2308 Ok(x)
2309 }
2310
2311 /// Kill-switch seam for the gemma4 dense-31B batched decode arm. DEFAULT ON since the
2312 /// 2026-08-16 owner flip ("if the performance are so strong in favor... we serve the
2313 /// correctness and best performance"): the arm's exactness battery is green at B=4/8,
2314 /// the served identity gate is byte-exact vs eager at c1/c4, and the served aggregate
2315 /// read 55→257 tok/s c16 on the NVFP4mix artifact at 450W (SERVED-AGGREGATE.md).
2316 /// `MEMRA_GEMMA4_BATCH=0` forces the eager per-session path (the rollback);
2317 /// `1` is the old opt-in spelling, still accepted. Any OTHER value REFUSES LOUD at
2318 /// first use — a mis-typed kill switch must not silently pick a serving path.
2319 pub fn gemma4_batch_on() -> bool {
2320 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2321 *ON.get_or_init(|| match std::env::var("MEMRA_GEMMA4_BATCH").as_deref() {
2322 Err(_) | Ok("1") => true,
2323 Ok("0") => false,
2324 Ok(v) => panic!(
2325 "MEMRA_GEMMA4_BATCH={v:?} is not a recognized value (want unset/1 = batched \
2326 decode, 0 = eager kill switch) — refusing to guess a serving path"
2327 ),
2328 })
2329 }
2330
2331 /// THE gemma4 dense-31B BATCHED DECODE ARM (lane/gemma-batched, 2026-08-16).
2332 ///
2333 /// gemma4 served eager-only — the c1→c8 aggregate was FLAT (~55 tok/s, per-stream
2334 /// collapse) because there was no batched arm, not because of quantization. This is it.
2335 ///
2336 /// SHAPE — batched where the weights are, per-session where the state is (the step35
2337 /// law, applied to gemma4's own geometry):
2338 /// * embed+scale, attn_norm+q8_1 quantize, wq/wk/wv projections, q/k RMSNorm +
2339 /// weightless-V norm + dual rope (fused `rms_norm_qkv_rope`), post_attn_norm, the
2340 /// layer-scale tail with its dense GEGLU FFN (`gemma4_layer_tail_add_nq`), output
2341 /// norm, softcapped head — ALL at m=B: one weight stream serves B rows (decode is
2342 /// weight-BW-bound; that is the entire aggregate win). Every one of these is the
2343 /// SAME batch-capable function the proven verify trunk (`gemma4_verify_trunk`) runs
2344 /// at width t, so this arm inherits the verify path's numerics wholesale.
2345 /// * KV append + fa_decode stay a PER-SESSION loop: each session appends its one new
2346 /// token to its own cache and attends its own [win_off .. len] view — the SWA
2347 /// window + global-vs-windowed geometry makes each session's t_kv independent, so
2348 /// there is no cross-session batched attention (identical to eager per session).
2349 ///
2350 /// EXACTNESS: v1 routes every session's attention through `fa_decode_kvmod` (the eager
2351 /// arm's unconditional fallback — same call `gemma4_decode_attn` makes with the rows_w
2352 /// fast arms off), so a B=1 run is the eager decode's own attention program and the
2353 /// batch is per-row independent by construction. The rows / rows_w per-session fast
2354 /// arms are a later perf increment gated behind their own seam.
2355 fn gemma4_decode_batch(
2356 &self,
2357 e: &Engine,
2358 tokens: &[u32],
2359 caches: &mut [&mut Cache],
2360 samp: &[Option<DevSamp>],
2361 masks: &[Option<(&CudaSlice<u32>, usize)>],
2362 lean: bool,
2363 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
2364 let b_n = tokens.len();
2365 if b_n == 0 || b_n != caches.len() {
2366 return Err(format!(
2367 "gemma4_decode_batch: tokens/caches mismatch (tokens={b_n}, caches={})",
2368 caches.len()
2369 )
2370 .into());
2371 }
2372 // Exactness tier boundary: the battery is green at B<=8 (per-row mmvq); m>8
2373 // crosses the dp4a-tail/GEMM numeric configs it never proved. The worker's chunk
2374 // policy caps gemma4 at 8; this is the per-request backstop (Err, never a panic —
2375 // the 2026-08-07 worker-FATAL law).
2376 if b_n > 8 {
2377 return Err(format!(
2378 "gemma4_decode_batch: B={b_n} > 8, past the proven exactness tier — \
2379 the scheduler must chunk gemma4 at <=8"
2380 )
2381 .into());
2382 }
2383 let n_embd = self.cfg.n_embd as usize;
2384 let eps = self.cfg.rms_eps;
2385 if b_n > 1 {
2386 static ONCE: std::sync::Once = std::sync::Once::new();
2387 ONCE.call_once(|| {
2388 eprintln!("[gemma4-batch] first B>1 batched gemma4 walk: B={b_n}");
2389 });
2390 }
2391 // per-session rope positions (each sequence at its own depth).
2392 let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
2393 let pos_d = e.htod_i32(&pos_v)?;
2394 let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
2395 e.scale_inplace(&mut x, (n_embd as f32).sqrt(), b_n * n_embd)?;
2396 // cross-layer carry: each tail emits the next layer's attn-normed q8_1 input.
2397 let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
2398 let n_layers = self.layers.len();
2399 for (il, layer) in self.layers.iter().enumerate() {
2400 let (hq, hdq) = match h_carry.take() {
2401 Some(p) => p,
2402 None => {
2403 e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, b_n, eps)?
2404 }
2405 };
2406 let Mixer::Full(fa) = &layer.mixer else {
2407 return Err(format!("gemma4 layer {il} not full-attn — corrupt config").into());
2408 };
2409 let o = self.gemma4_batch_attn(e, fa, il, &hq, &hdq, &pos_d, b_n, caches)?;
2410 let next_norm = if il + 1 < n_layers {
2411 Some(self.layers[il + 1].attn_norm.float_data())
2412 } else {
2413 None
2414 };
2415 // pn-fold front (lane/gemma-pnfold merge): the batched arm rides the SAME
2416 // tail front as the eager/verify trio, so batched == eager holds by
2417 // construction at either MEMRA_G4_PNFOLD value (seam-off falls through to
2418 // the unfused rms_norm + tail chain this arm shipped with).
2419 let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, b_n, next_norm)?;
2420 x = xn;
2421 h_carry = hn;
2422 }
2423 let mut hn = e.uninit(b_n * n_embd)?;
2424 e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
2425 let mut ld = e.matmul(&self.output, &hn, b_n)?;
2426 let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
2427 e.softcap(&mut ld, cap, b_n * self.output.out_features())?;
2428 self.gemma4_suppress(e, &mut ld, b_n)?; // non-monotonic — before any argmax/sample
2429 let mut ph_last = std::time::Instant::now();
2430 self.decode_batch_epilogue(e, caches, samp, masks, lean, ld, b_n, &mut ph_last)
2431 }
2432
2433 /// Per-session gemma4 attention for the batched arm: batched projections + fused
2434 /// q/k-norm + weightless-V-norm + dual rope over all B rows (per-row independent, the
2435 /// verify path's exact kernels), then a per-session KV append + `fa_decode_kvmod` over
2436 /// each session's own window/global view, then one batched wo matmul. Mirrors the eager
2437 /// `gemma4_decode_attn` fallback per row.
2438 #[allow(clippy::too_many_arguments)]
2439 fn gemma4_batch_attn(
2440 &self,
2441 e: &Engine,
2442 fa: &crate::hybrid::FullAttnLayer,
2443 il: usize,
2444 hq: &CudaSlice<i8>,
2445 hdq: &CudaSlice<f32>,
2446 pos_d: &CudaSlice<i32>,
2447 b_n: usize,
2448 caches: &mut [&mut Cache],
2449 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2450 let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
2451 let eps = self.cfg.rms_eps;
2452 let aux = self.gemma4_aux.as_ref().unwrap();
2453 let ones = aux.ones(e);
2454 let h0 = e.zeros(0)?;
2455 let h = &h0;
2456 // projections at m=B (the f32 fallback `h` is empty; matmul_pre uses the q8_1 pair).
2457 let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, b_n)?;
2458 let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, b_n)?;
2459 let v0 = if swa {
2460 e.matmul_pre(&fa.wv, hq, hdq, h, b_n)?
2461 } else {
2462 e.clone_dtod(&k0)? // globals: V := K clone (weightless V-norm, never roped)
2463 };
2464 let mut q = e.uninit(b_n * nh * hd)?;
2465 let mut k = e.uninit(b_n * nkv * hd)?;
2466 let mut v = e.uninit(b_n * nkv * hd)?;
2467 let ff = if swa {
2468 None
2469 } else {
2470 Some(
2471 aux.rope_freqs(e)
2472 .expect("gemma4 global rope needs rope_freqs.weight"),
2473 )
2474 };
2475 e.rms_norm_qkv_rope(
2476 &q0,
2477 &k0,
2478 &v0,
2479 fa.q_norm.float_data(),
2480 fa.k_norm.float_data(),
2481 ones,
2482 &mut q,
2483 &mut k,
2484 &mut v,
2485 hd,
2486 nh * b_n,
2487 nkv * b_n,
2488 pos_d,
2489 nh,
2490 nkv,
2491 base,
2492 1.0,
2493 ff,
2494 eps,
2495 )?;
2496 let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
2497 let q_dim = nh * hd;
2498 let kv_dim = nkv * hd;
2499 let mut attn = e.uninit(b_n * q_dim)?;
2500 for bi in 0..b_n {
2501 let kvl = caches[bi].kv[il].as_mut().unwrap();
2502 let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
2503 let v_row = v.slice(bi * kv_dim..(bi + 1) * kv_dim);
2504 // gemma4's KV is a linear buffer (no ring rebase — the SWA view below is a plain
2505 // token-offset), so append at kvl.len exactly as eager gemma4_decode_attn does.
2506 e.append_kv_quantized_view(
2507 &k_row,
2508 &v_row,
2509 &mut kvl.k,
2510 &mut kvl.v,
2511 kvl.len,
2512 kvl.kv_dim_k,
2513 kvl.kv_dim_v,
2514 kvl.k_tok_bytes,
2515 kvl.v_tok_bytes,
2516 (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
2517 )?;
2518 kvl.len += 1;
2519 // eager SWA view arithmetic (gemma4_decode_attn): token-aligned window offset;
2520 // keys carry absolute rope, the mask is purely positional.
2521 let (off_tok, t_kv) = if swa && kvl.len > win {
2522 (kvl.len - win, win)
2523 } else {
2524 (0, kvl.len)
2525 };
2526 let k_view = e.view_u8_range(
2527 &kvl.k,
2528 off_tok * kvl.k_tok_bytes,
2529 (off_tok + t_kv) * kvl.k_tok_bytes,
2530 );
2531 let v_view = e.view_u8_range(
2532 &kvl.v,
2533 off_tok * kvl.v_tok_bytes,
2534 (off_tok + t_kv) * kvl.v_tok_bytes,
2535 );
2536 let q_row = q.slice(bi * q_dim..(bi + 1) * q_dim);
2537 let mut a_row = attn.slice_mut(bi * q_dim..(bi + 1) * q_dim);
2538 e.fa_decode_kvmod_view(
2539 &q_row,
2540 &k_view,
2541 &v_view,
2542 &mut a_row,
2543 hd,
2544 nh,
2545 nkv,
2546 t_kv,
2547 scale,
2548 kvl.k_tok_bytes,
2549 kvl.v_tok_bytes,
2550 swa && crate::Engine::wkv_on(),
2551 )?;
2552 }
2553 Ok(e.matmul(&fa.wo, &attn, b_n)?)
2554 }
2555
2556 /// Standalone MoESD target forward. This entrypoint is not used by serving: it widens the
2557 /// existing Step-3.7 batched layer walk to B*gamma rows while preserving one causal KV chain
2558 /// per session. It returns device logits and performs no sampling or logits D2H, matching the
2559 /// target-model term T_T measured by the paper.
2560 pub fn moesd_target_forward(
2561 &self,
2562 e: &Engine,
2563 tokens: &[u32],
2564 batch: usize,
2565 gamma: usize,
2566 caches: &mut [&mut Cache],
2567 ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2568 if self.cfg.step35.is_none() {
2569 return Err("MoESD target forward currently requires Step-3.7/Step35 geometry".into());
2570 }
2571 if batch == 0 || gamma == 0 || caches.len() != batch || tokens.len() != batch * gamma {
2572 return Err(format!(
2573 "MoESD shape mismatch: B={batch} gamma={gamma} caches={} tokens={}",
2574 caches.len(),
2575 tokens.len(),
2576 )
2577 .into());
2578 }
2579 let rows = batch * gamma;
2580 if rows > 256 {
2581 return Err(format!("MoESD target width {rows} exceeds the frozen 32*8 matrix").into());
2582 }
2583 let n_embd = self.cfg.n_embd as usize;
2584 let eps = self.cfg.rms_eps;
2585 let payload = rows * n_embd;
2586 let row_to_cache: Vec<usize> = (0..batch)
2587 .flat_map(|session| (0..gamma).map(move |_| session))
2588 .collect();
2589 let positions: Vec<i32> = row_to_cache
2590 .iter()
2591 .enumerate()
2592 .map(|(row, &session)| (caches[session].pos + row % gamma) as i32)
2593 .collect();
2594 let mut ph_last = std::time::Instant::now();
2595
2596 let logits = if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2597 if fence.len() != 3 || crate::pp::pp2_streams_off() {
2598 return Err(
2599 "MoESD PP target forward requires the live two-stage stream split".into(),
2600 );
2601 }
2602 let rt = crate::pp::PpNRt::get(e)?;
2603 if rt.n_stages() != 2 {
2604 return Err(format!("MoESD expected two PP stages, got {}", rt.n_stages()).into());
2605 }
2606 let caller_stream = e.stream();
2607 rt.fence_stages_behind(&caller_stream)?;
2608 let slot = {
2609 let _st0 = rt.enter(0);
2610 let e0 = rt.engine(0, e);
2611 let pos_d = e0.htod_i32(&positions)?;
2612 let x = e0.htod(&self.embd.gather(n_embd, tokens))?;
2613 ph_mark(e0, 0, &mut ph_last)?;
2614 let x = self.step35_decode_rows_layers(
2615 e0,
2616 x,
2617 caches,
2618 &pos_d,
2619 Some(&row_to_cache),
2620 fence[0],
2621 fence[1],
2622 &mut ph_last,
2623 )?;
2624 rt.tx(0, &x, payload)?
2625 };
2626 let logits = {
2627 let _st1 = rt.enter(1);
2628 let e1 = rt.engine(1, e);
2629 let pos_d = e1.htod_i32(&positions)?;
2630 let x = rt.rx(0, slot, payload)?;
2631 let x = self.step35_decode_rows_layers(
2632 e1,
2633 x,
2634 caches,
2635 &pos_d,
2636 Some(&row_to_cache),
2637 fence[1],
2638 fence[2],
2639 &mut ph_last,
2640 )?;
2641 let mut hn = e1.uninit(payload)?;
2642 e1.rms_norm(
2643 &x,
2644 self.output_norm.float_data(),
2645 &mut hn,
2646 n_embd,
2647 rows,
2648 eps,
2649 )?;
2650 let logits = e1.matmul(&self.output, &hn, rows)?;
2651 rt.publish_to(1, &caller_stream)?;
2652 logits
2653 };
2654 logits
2655 } else {
2656 let pos_d = e.htod_i32(&positions)?;
2657 let x = e.htod(&self.embd.gather(n_embd, tokens))?;
2658 ph_mark(e, 0, &mut ph_last)?;
2659 let x = self.step35_decode_rows_layers(
2660 e,
2661 x,
2662 caches,
2663 &pos_d,
2664 Some(&row_to_cache),
2665 0,
2666 self.layers.len(),
2667 &mut ph_last,
2668 )?;
2669 let mut hn = e.uninit(payload)?;
2670 e.rms_norm(
2671 &x,
2672 self.output_norm.float_data(),
2673 &mut hn,
2674 n_embd,
2675 rows,
2676 eps,
2677 )?;
2678 e.matmul(&self.output, &hn, rows)?
2679 };
2680 for cache in caches.iter_mut() {
2681 cache.pos += gamma;
2682 }
2683 Ok(logits)
2684 }
2685
2686 /// The batched tick's TAIL, after the trunk: grammar masks -> device sampling -> lean
2687 /// logits park -> `pos` bump. Split out with the pp seam (`decode_batch_layers`) because
2688 /// under a stage split this runs on the LAST stage's engine and device — the lm_head, the
2689 /// masks, the sampler, and `cache.last_logits_dev` all live where the final residual
2690 /// lands, and the caller must be able to place them there without duplicating 90 lines of
2691 /// serving contract. `logits` is `[b_n, n_vocab]` already computed by the caller (the
2692 /// output_norm + lm_head pair stays at the call site so a stage split can fence around
2693 /// it); everything after it is here, verbatim.
2694 #[allow(clippy::too_many_arguments)]
2695 fn decode_batch_epilogue(
2696 &self,
2697 e: &Engine,
2698 caches: &mut [&mut Cache],
2699 samp: &[Option<DevSamp>],
2700 masks: &[Option<(&CudaSlice<u32>, usize)>],
2701 lean: bool,
2702 logits: CudaSlice<f32>,
2703 b_n: usize,
2704 ph_last: &mut std::time::Instant,
2705 ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
2706 // GRAMMAR MASKS (constrained decoding): preserve each masked row's PRISTINE logits
2707 // for its consumer (lean park into cache.last_logits_dev — the reuse-pool park stays
2708 // unmasked, the v1 contract — or the non-lean D2H), then ban in place BEFORE the
2709 // device sampler reads the row. All stream-ordered; masks=&[] takes no new branch.
2710 let n_vocab = self.output.out_features();
2711 let mut logits = logits;
2712 let mut pristine: Vec<Option<CudaSlice<f32>>> = Vec::new();
2713 if masks.iter().take(b_n).any(|m| m.is_some()) {
2714 pristine.resize_with(b_n, || None);
2715 for (bi, m) in masks.iter().take(b_n).enumerate() {
2716 let Some((mask, words)) = m else { continue };
2717 assert!(
2718 samp.get(bi).copied().flatten().is_some(),
2719 "grammar-masked row {bi} must request a device sample"
2720 );
2721 if lean {
2722 let cache = &mut caches[bi];
2723 if cache
2724 .last_logits_dev
2725 .as_ref()
2726 .map(|d| d.len() < n_vocab)
2727 .unwrap_or(true)
2728 {
2729 cache.last_logits_dev = Some(e.uninit(n_vocab)?);
2730 }
2731 let dst = cache.last_logits_dev.as_mut().unwrap();
2732 e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
2733 } else {
2734 let mut p = e.uninit(n_vocab)?;
2735 e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), &mut p)?;
2736 pristine[bi] = Some(p);
2737 }
2738 e.mask_logits_col(&mut logits, mask, bi, n_vocab, *words)?;
2739 }
2740 }
2741
2742 // Device-side sampling for requested rows (see the method doc). Enqueued before the
2743 // big logits D2H so the tiny [B] token readback rides the same sync.
2744 let mut next: Vec<Option<u32>> = vec![None; b_n];
2745 if samp.iter().take(b_n).any(|s| s.is_some()) {
2746 let mut toks = e.alloc_u32_zeroed(b_n)?;
2747 let mut perturb: Option<CudaSlice<f32>> = None;
2748 for (bi, s) in samp.iter().take(b_n).enumerate() {
2749 let Some((temp, seed, ctr, top_k, top_p, min_p)) = s else {
2750 continue;
2751 };
2752 let filtered = *temp > 0.0 && (*top_k > 0 || *top_p < 1.0 || *min_p > 0.0);
2753 if *temp <= 0.0 {
2754 e.argmax_token_device_col(&logits, bi, n_vocab, &mut toks, bi)?;
2755 } else if filtered {
2756 if perturb.is_none() {
2757 perturb = Some(e.zeros(n_vocab)?);
2758 }
2759 let pb = perturb.as_mut().unwrap();
2760 self.devsample_filtered_col(
2761 e, &logits, bi, n_vocab, *temp, *seed, *ctr, *top_k, *top_p, *min_p, pb,
2762 &mut toks, bi,
2763 )?;
2764 } else {
2765 if perturb.is_none() {
2766 perturb = Some(e.zeros(n_vocab)?);
2767 }
2768 let pb = perturb.as_mut().unwrap();
2769 e.gumbel_perturb_col(&logits, bi, pb, n_vocab, *seed, *ctr, *temp)?;
2770 e.argmax_token_device_col(pb, 0, n_vocab, &mut toks, bi)?;
2771 }
2772 }
2773 let host_toks = e.dtoh_u32(&toks)?;
2774 for (bi, s) in samp.iter().take(b_n).enumerate() {
2775 if s.is_some() {
2776 next[bi] = Some(host_toks[bi]);
2777 }
2778 }
2779 }
2780
2781 let lean_any = lean && samp.iter().take(b_n).any(|s| s.is_some());
2782 let rows: Vec<Vec<f32>> = if lean_any {
2783 // LEAN: park device-sampled rows on-device (per-cache buffer, dtod); D2H only
2784 // the rows that still need host logits. No sampled rows + no fallback rows =
2785 // the big D2H disappears (the [B] token readback above already synced).
2786 for (bi, s) in samp.iter().take(b_n).enumerate() {
2787 if s.is_none() {
2788 continue;
2789 }
2790 // grammar-masked rows already parked their PRISTINE copy above — the
2791 // in-place ban has since poisoned this row for the reuse-pool consumer.
2792 if masks.get(bi).copied().flatten().is_some() {
2793 continue;
2794 }
2795 let cache = &mut caches[bi];
2796 if cache
2797 .last_logits_dev
2798 .as_ref()
2799 .map(|d| d.len() < n_vocab)
2800 .unwrap_or(true)
2801 {
2802 cache.last_logits_dev = Some(e.uninit(n_vocab)?);
2803 }
2804 let dst = cache.last_logits_dev.as_mut().unwrap();
2805 e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
2806 }
2807 (0..b_n)
2808 .map(|bi| {
2809 if samp.get(bi).copied().flatten().is_some() {
2810 Ok(Vec::new())
2811 } else {
2812 e.dtoh_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab))
2813 }
2814 })
2815 .collect::<Result<_, _>>()?
2816 } else {
2817 let host = e.dtoh(&logits)?;
2818 (0..b_n)
2819 .map(|bi| {
2820 // grammar-masked non-lean rows return the PRISTINE copy (the in-place ban
2821 // must never leak into last_logits — reuse-pool/park semantics unchanged).
2822 if let Some(p) = pristine.get(bi).and_then(|p| p.as_ref()) {
2823 return e.dtoh(p);
2824 }
2825 Ok(host[bi * n_vocab..(bi + 1) * n_vocab].to_vec())
2826 })
2827 .collect::<Result<_, _>>()?
2828 };
2829 for c in caches.iter_mut() {
2830 c.pos += 1;
2831 }
2832 ph_mark(e, 11, ph_last)?;
2833 Ok((rows, next))
2834 }
2835}
2836
2837fn b1_fast_arch_eligible(arch: &Arch) -> bool {
2838 // The whole qwen35 family is excluded, not just MoE: spec verify for these archs runs
2839 // the generic batched numeric class (spec.rs qwen35_serving_class), so live B=1 serving
2840 // must stay in that same class. B1FAST's eager program would reopen the near-tie-flip
2841 // divergence the 2026-08-14 exactness fix closed (1 ULP at layer 2 -> 2.3e-1 head
2842 // maxdiff, amplified by the GDN recurrence).
2843 !matches!(arch, Arch::Qwen35 | Arch::Qwen35Moe)
2844}
2845
2846fn b1_fast_env_on(value: Option<&str>) -> bool {
2847 value == Some("1")
2848}
2849
2850#[cfg(test)]
2851mod tests {
2852 use super::{b1_fast_arch_eligible, b1_fast_env_on};
2853 use memra_gguf::config::Arch;
2854
2855 #[test]
2856 fn qwen35_family_stays_in_one_decode_numeric_class_across_widths() {
2857 assert!(!b1_fast_arch_eligible(&Arch::Qwen35Moe));
2858 assert!(!b1_fast_arch_eligible(&Arch::Qwen35));
2859 assert!(b1_fast_arch_eligible(&Arch::Qwen3Moe));
2860 }
2861
2862 #[test]
2863 fn b1_eager_program_requires_explicit_opt_in() {
2864 assert!(!b1_fast_env_on(None));
2865 assert!(!b1_fast_env_on(Some("0")));
2866 assert!(!b1_fast_env_on(Some("true")));
2867 assert!(b1_fast_env_on(Some("1")));
2868 }
2869}