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