Skip to main content

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, &gt, &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}