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