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::collections::VecDeque;
34use std::sync::Arc;
35
36type DualPpCudaSpan = Option<(CudaEvent, CudaEvent)>;
37
38/// One disjoint request wave moving through the PP3/PP4 anti-diagonal schedule. The activation
39/// itself lives in `PpNRt`'s persistent boundary slot; this host state carries only ownership of
40/// the request caches and the slot selected by the preceding stage.
41struct PpDecodeWave<'slice, 'cache> {
42    row_lo: usize,
43    tokens: &'slice [u32],
44    caches: &'slice mut [&'cache mut Cache],
45    phase_last: std::time::Instant,
46    #[allow(clippy::type_complexity)]
47    // allow: one-shot composite type; naming it would hide the shape that matters at the call site
48    result: Option<(Vec<Vec<f32>>, Vec<Option<u32>>)>,
49    committed: bool,
50}
51
52impl Drop for PpDecodeWave<'_, '_> {
53    fn drop(&mut self) {
54        if !self.committed {
55            for cache in self.caches.iter_mut() {
56                cache.mark_tainted();
57            }
58        }
59    }
60}
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63struct PpWaveTransfer {
64    boundary: usize,
65    wave: usize,
66    slot: usize,
67}
68
69#[derive(Debug)]
70enum PpWaveMessage {
71    Transfer(PpWaveTransfer),
72    WorkerError {
73        boundary: usize,
74        wave: usize,
75        error: String,
76    },
77}
78
79struct PpWaveIncoming {
80    boundary: usize,
81    transfers: std::sync::mpsc::Receiver<PpWaveMessage>,
82    acknowledgements: std::sync::mpsc::Sender<PpWaveTransfer>,
83}
84
85impl PpWaveIncoming {
86    fn receive(&self, expected_wave: usize) -> Result<PpWaveTransfer, String> {
87        match self.transfers.recv() {
88            Ok(PpWaveMessage::Transfer(transfer)) => {
89                if transfer.boundary != self.boundary || transfer.wave != expected_wave {
90                    return Err(format!(
91                        "PP wave boundary {} expected wave {expected_wave}, got boundary {} wave {} slot {}",
92                        self.boundary, transfer.boundary, transfer.wave, transfer.slot,
93                    ));
94                }
95                if transfer.slot >= 2 {
96                    return Err(format!(
97                        "PP wave boundary {} wave {expected_wave} carried invalid slot {}",
98                        self.boundary, transfer.slot,
99                    ));
100                }
101                Ok(transfer)
102            }
103            Ok(PpWaveMessage::WorkerError {
104                boundary,
105                wave,
106                error,
107            }) => {
108                if boundary != self.boundary {
109                    return Err(format!(
110                        "PP wave boundary {} received worker failure for boundary {boundary} wave {wave}: {error}",
111                        self.boundary,
112                    ));
113                }
114                Err(format!(
115                    "PP wave boundary {boundary} upstream worker failed at wave {wave} while receiver expected wave {expected_wave}: {error}"
116                ))
117            }
118            Err(_) => Err(format!(
119                "PP wave boundary {} transfer channel closed before wave {expected_wave}",
120                self.boundary,
121            )),
122        }
123    }
124
125    fn acknowledge(&self, transfer: PpWaveTransfer) -> Result<(), String> {
126        if transfer.boundary != self.boundary {
127            return Err(format!(
128                "PP wave acknowledgement boundary mismatch: receiver {} transfer {}",
129                self.boundary, transfer.boundary,
130            ));
131        }
132        self.acknowledgements.send(transfer).map_err(|_| {
133            format!(
134                "PP wave boundary {} acknowledgement channel closed at wave {} slot {}",
135                self.boundary, transfer.wave, transfer.slot,
136            )
137        })
138    }
139}
140
141struct PpWaveOutgoing {
142    boundary: usize,
143    transfers: std::sync::mpsc::Sender<PpWaveMessage>,
144    acknowledgements: std::sync::mpsc::Receiver<PpWaveTransfer>,
145    slot_owner: [Option<PpWaveTransfer>; 2],
146    pending: VecDeque<PpWaveTransfer>,
147    next_slot: Option<usize>,
148    next_wave: usize,
149}
150
151impl PpWaveOutgoing {
152    fn new(
153        boundary: usize,
154        transfers: std::sync::mpsc::Sender<PpWaveMessage>,
155        acknowledgements: std::sync::mpsc::Receiver<PpWaveTransfer>,
156    ) -> Self {
157        Self {
158            boundary,
159            transfers,
160            acknowledgements,
161            slot_owner: [None, None],
162            pending: VecDeque::new(),
163            next_slot: None,
164            next_wave: 0,
165        }
166    }
167
168    fn receive_ack(&mut self, expected: PpWaveTransfer) -> Result<(), String> {
169        let actual = self.acknowledgements.recv().map_err(|_| {
170            format!(
171                "PP wave boundary {} acknowledgement channel closed waiting for wave {} slot {}",
172                self.boundary, expected.wave, expected.slot,
173            )
174        })?;
175        if actual != expected {
176            return Err(format!(
177                "PP wave boundary {} expected acknowledgement wave {} slot {}, got boundary {} wave {} slot {}",
178                self.boundary,
179                expected.wave,
180                expected.slot,
181                actual.boundary,
182                actual.wave,
183                actual.slot,
184            ));
185        }
186        let pending = self.pending.pop_front().ok_or_else(|| {
187            "PP wave acknowledgement arrived with no pending transfer".to_string()
188        })?;
189        if pending != expected {
190            return Err(format!(
191                "PP wave boundary {} acknowledgement order mismatch: pending wave {} slot {}, expected wave {} slot {}",
192                self.boundary, pending.wave, pending.slot, expected.wave, expected.slot,
193            ));
194        }
195        self.slot_owner[expected.slot] = None;
196        Ok(())
197    }
198
199    /// Return the slot `tx_pipelined` must select next. If that slot still belongs to an older
200    /// wave, wait for the exact downstream acknowledgement proving `rx` recorded `ev_rx` for it.
201    fn prepare(&mut self, wave: usize) -> Result<Option<usize>, String> {
202        if wave != self.next_wave {
203            return Err(format!(
204                "PP wave boundary {} producer order mismatch: expected wave {}, got {wave}",
205                self.boundary, self.next_wave,
206            ));
207        }
208        if let Some(slot) = self.next_slot
209            && let Some(owner) = self.slot_owner[slot]
210        {
211            self.receive_ack(owner)?;
212        }
213        Ok(self.next_slot)
214    }
215
216    fn publish(
217        &mut self,
218        wave: usize,
219        slot: usize,
220        expected_slot: Option<usize>,
221    ) -> Result<(), String> {
222        if wave != self.next_wave {
223            return Err(format!(
224                "PP wave boundary {} publish order mismatch: expected wave {}, got {wave}",
225                self.boundary, self.next_wave,
226            ));
227        }
228        if slot >= 2 {
229            return Err(format!(
230                "PP wave boundary {} wave {wave} selected invalid slot {slot}",
231                self.boundary,
232            ));
233        }
234        if let Some(expected) = expected_slot
235            && slot != expected
236        {
237            return Err(format!(
238                "PP wave boundary {} wave {wave} broke slot alternation: expected {expected}, got {slot}",
239                self.boundary,
240            ));
241        }
242        if let Some(owner) = self.slot_owner[slot] {
243            return Err(format!(
244                "PP wave boundary {} attempted to reuse slot {slot} for wave {wave} before acknowledgement of wave {}",
245                self.boundary, owner.wave,
246            ));
247        }
248        let transfer = PpWaveTransfer {
249            boundary: self.boundary,
250            wave,
251            slot,
252        };
253        self.slot_owner[slot] = Some(transfer);
254        self.pending.push_back(transfer);
255        self.next_slot = Some(slot ^ 1);
256        self.next_wave += 1;
257        self.transfers
258            .send(PpWaveMessage::Transfer(transfer))
259            .map_err(|_| {
260                format!(
261                    "PP wave boundary {} transfer channel closed publishing wave {wave} slot {slot}",
262                    self.boundary,
263                )
264            })
265    }
266
267    fn finish(&mut self) -> Result<(), String> {
268        while let Some(expected) = self.pending.front().copied() {
269            self.receive_ack(expected)?;
270        }
271        Ok(())
272    }
273
274    fn publish_worker_error(&self, error: &str) {
275        let _ = self.transfers.send(PpWaveMessage::WorkerError {
276            boundary: self.boundary,
277            wave: self.next_wave,
278            error: error.to_string(),
279        });
280    }
281}
282
283fn pp_wave_channels(
284    boundaries: usize,
285) -> (Vec<Option<PpWaveOutgoing>>, Vec<Option<PpWaveIncoming>>) {
286    let mut outgoing = Vec::with_capacity(boundaries);
287    let mut incoming = Vec::with_capacity(boundaries);
288    for boundary in 0..boundaries {
289        let (transfer_tx, transfer_rx) = std::sync::mpsc::channel();
290        let (ack_tx, ack_rx) = std::sync::mpsc::channel();
291        outgoing.push(Some(PpWaveOutgoing::new(boundary, transfer_tx, ack_rx)));
292        incoming.push(Some(PpWaveIncoming {
293            boundary,
294            transfers: transfer_rx,
295            acknowledgements: ack_tx,
296        }));
297    }
298    (outgoing, incoming)
299}
300
301fn dual_pp_timing_event(e: &Engine, context: &str) -> Option<CudaEvent> {
302    if !crate::pp::dual_pp_timing_on() {
303        return None;
304    }
305    match e
306        .stream()
307        .record_event(Some(cudarc::driver::sys::CUevent_flags::CU_EVENT_DEFAULT))
308    {
309        Ok(event) => Some(event),
310        Err(err) => {
311            crate::pp::record_dual_pp_timing_drop(context, &err);
312            None
313        }
314    }
315}
316
317/// Per-step, per-LAYER-RANGE invariants the batched trunk needs: the device state-pointer
318/// table for the range's layers, the arm picks, and the per-row `t_kv` snapshot. Built once
319/// per step per range by `HybridModel::batch_layer_ctx`, consumed by `decode_batch_layers`.
320///
321/// WHY IT IS RANGE-SCOPED AND NOT STEP-SCOPED (this is the whole point of the struct):
322/// `ptr_table` is a `CudaSlice<u64>` of DEVICE ADDRESSES, uploaded through `e` — so it lives
323/// on `e`'s device, and its entries are pointers into caches that live on the device that
324/// OWNS those layers. Under a pp stage split, stage s runs layers [fence[s], fence[s+1])
325/// whose cache state was allocated by stage s's engine (`pp::new_cache` -> `Cache::new_ppn`),
326/// so stage s must build its OWN table through its OWN engine. One step-wide table built on
327/// the primary would put every stage's kernel arguments in stage-0's HBM — a peer read per
328/// pointer fetch, which is the exact cliff `pp::refuse_unsplit_if_remote` exists to stop.
329/// `lo`/`hi` are recorded so the consumer can assert the ctx it was handed matches the range
330/// it was asked to run (the offsets in `lin_base`/`attn_base` are only valid for that range).
331pub(crate) struct BatchLayerCtx {
332    /// Offset into `ptr_table` of layer il's [conv x B][ssm_in x B][ssm_out x B] block
333    /// (linear-attn layers only). Indexed by ABSOLUTE layer id; `None` off-range.
334    lin_base: Vec<Option<usize>>,
335    /// Offset into `ptr_table` of layer il's [k0,v0,k1,v1,..] block (full-attn layers only).
336    /// Indexed by ABSOLUTE layer id; `None` off-range.
337    attn_base: Vec<Option<usize>>,
338    ptr_table: Option<CudaSlice<u64>>,
339    /// Per-row `pos + 1` — the t_kv each sequence attends at this step. Layer-invariant
340    /// within a step, so the arm picks below are decided once.
341    t_kvs: Vec<usize>,
342    t_kv_max: usize,
343    /// The single `fa_split_keys` rung every row shares (the rows-twins straddle law).
344    sp0: usize,
345    seqs_fa: bool,
346    lo: usize,
347    hi: usize,
348}
349
350// ---- MEMRA_BATCH_PHASE=1 (diagnostics): sync-bounded per-phase accumulators for the batched
351// tick. Each boundary syncs the stream, so the TOTAL inflates (launch pipelining is destroyed);
352// the value is the RANKING/shares, not absolute ms. Read via `batch_phase_report()`.
353pub(crate) static BATCH_PHASE: std::sync::Mutex<[f64; 12]> = std::sync::Mutex::new([0.0; 12]);
354/// Device-sample request for one batched row.
355/// `top_k=0` / `top_p>=1.0` / `min_p<=0.0` = that filter off. Greedy = temp<=0 (device
356/// argmax); pure temperature = seeded gumbel; any filter on = filter_stats floor + the
357/// filtered gumbel draw. `penalty` carries host-maintained sparse counts for the exact active
358/// history window; the epilogue applies them on device before filters and sampling.
359#[derive(Clone, Debug)]
360pub struct DevSamp {
361    pub temp: f32,
362    pub seed: u64,
363    pub ctr: u32,
364    pub top_k: i32,
365    pub top_p: f32,
366    pub min_p: f32,
367    pub penalty: Option<DevPenalty>,
368}
369
370#[derive(Clone, Debug)]
371pub struct DevPenalty {
372    repeat: f32,
373    freq: f32,
374    present: f32,
375    counts: Vec<(u32, u32)>,
376}
377
378/// A one-row decode whose device work has been enqueued but whose result has not crossed back
379/// to the host yet. The worker owns the CUDA context, so this is deliberately a poll-at-the-next
380/// scheduler-boundary handoff rather than a background CUDA thread. Keeping the completion event
381/// and output buffers alive prevents the async-pool from recycling them while the next step runs.
382pub struct PendingBatchStep {
383    logits: CudaSlice<f32>,
384    pristine: Vec<Option<CudaSlice<f32>>>,
385    tokens: Option<CudaSlice<u32>>,
386    sampled: Vec<bool>,
387    n_vocab: usize,
388    lean: bool,
389    done: CudaEvent,
390    readback: Arc<cudarc::driver::CudaStream>,
391}
392
393impl PendingBatchStep {
394    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
395    fn new(
396        logits: CudaSlice<f32>,
397        pristine: Vec<Option<CudaSlice<f32>>>,
398        tokens: Option<CudaSlice<u32>>,
399        sampled: Vec<bool>,
400        n_vocab: usize,
401        lean: bool,
402        done: CudaEvent,
403        readback: Arc<cudarc::driver::CudaStream>,
404    ) -> Self {
405        Self {
406            logits,
407            pristine,
408            tokens,
409            sampled,
410            n_vocab,
411            lean,
412            done,
413            readback,
414        }
415    }
416
417    /// Wait for this step only, then perform one ordered readback of its host-visible results.
418    /// The compute stream may already be carrying the following step when this runs.
419    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
420    pub fn wait(self) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
421        self.readback.wait(&self.done)?;
422        // Lean device-sampled rows already parked their pristine logits in the session cache;
423        // their only host-visible result is the sampled token id. Avoid recreating the large
424        // vocab-row D2H that this path was introduced to remove.
425        let need_logits = !self.lean || self.sampled.iter().any(|sampled| !sampled);
426        let host_logits = need_logits
427            .then(|| self.readback.clone_dtoh(&self.logits))
428            .transpose()?;
429        let host_pristine: Vec<Option<Vec<f32>>> = self
430            .pristine
431            .iter()
432            .map(|row| {
433                row.as_ref()
434                    .map(|row| self.readback.clone_dtoh(row))
435                    .transpose()
436            })
437            .collect::<Result<_, _>>()?;
438        let host_tokens = self
439            .tokens
440            .as_ref()
441            .map(|tokens| self.readback.clone_dtoh(tokens))
442            .transpose()?;
443        self.readback.synchronize()?;
444
445        let mut rows = Vec::with_capacity(self.sampled.len());
446        for (bi, sampled) in self.sampled.iter().copied().enumerate() {
447            if self.lean && sampled {
448                rows.push(Vec::new());
449            } else if let Some(row) = host_pristine[bi].as_ref() {
450                rows.push(row.clone());
451            } else {
452                let start = bi * self.n_vocab;
453                let logits = host_logits
454                    .as_ref()
455                    .ok_or("pending step did not retain host logits for an unsampled row")?;
456                rows.push(logits[start..start + self.n_vocab].to_vec());
457            }
458        }
459        let next = host_tokens.map_or_else(
460            || vec![None; self.sampled.len()],
461            |tokens| {
462                self.sampled
463                    .iter()
464                    .enumerate()
465                    .map(|(bi, sampled)| sampled.then_some(tokens[bi]))
466                    .collect()
467            },
468        );
469        Ok((rows, next))
470    }
471}
472
473impl DevPenalty {
474    /// Checked constructor for callers that do not already own a unique count map.
475    pub fn try_new(
476        repeat: f32,
477        freq: f32,
478        present: f32,
479        counts: Vec<(u32, u32)>,
480    ) -> Result<Self, &'static str> {
481        let mut seen = std::collections::HashSet::with_capacity(counts.len());
482        for &(id, count) in &counts {
483            if count == 0 {
484                return Err("device penalty counts must be positive");
485            }
486            if !seen.insert(id) {
487                return Err("device penalty token ids must be unique");
488            }
489        }
490        Ok(Self {
491            repeat,
492            freq,
493            present,
494            counts,
495        })
496    }
497
498    /// Zero-copy validation seam for a producer that already owns a unique count map.
499    ///
500    /// # Safety
501    ///
502    /// `counts` must contain each token id at most once and every count must be positive. The
503    /// batched kernel assigns one CUDA thread to each entry and performs a non-atomic
504    /// read/modify/write of that token's logit.
505    pub unsafe fn from_unique_counts_unchecked(
506        repeat: f32,
507        freq: f32,
508        present: f32,
509        counts: Vec<(u32, u32)>,
510    ) -> Self {
511        Self {
512            repeat,
513            freq,
514            present,
515            counts,
516        }
517    }
518}
519
520impl DevSamp {
521    pub fn new(temp: f32, seed: u64, ctr: u32, top_k: i32, top_p: f32, min_p: f32) -> Self {
522        Self {
523            temp,
524            seed,
525            ctr,
526            top_k,
527            top_p,
528            min_p,
529            penalty: None,
530        }
531    }
532
533    pub fn with_penalty(mut self, penalty: DevPenalty) -> Self {
534        self.penalty = Some(penalty);
535        self
536    }
537}
538
539pub const BATCH_PHASE_NAMES: [&str; 12] = [
540    "setup(ptrs+embed H2D)",
541    "attn batched pre (norm/qkv/rope)",
542    "attn per-seq: kv append",
543    "attn per-seq: q/a dtod copies",
544    "attn per-seq: fa_decode",
545    "attn post (gate+o-proj)",
546    "gdn batched projections",
547    "gdn state ops (conv/prep/scan)",
548    "gdn out (gated norm+proj)",
549    "ffn (add/norm/gate/up/act/down)",
550    "lm_head (norm+matmul)",
551    "logits D2H + host split",
552];
553pub fn batch_phase_on() -> bool {
554    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
555    *ON.get_or_init(|| std::env::var("MEMRA_BATCH_PHASE").as_deref() == Ok("1"))
556}
557/// Accumulate the elapsed time since `last` into phase slot `slot` and re-stamp `last`.
558/// No-op unless `MEMRA_BATCH_PHASE=1`. Syncs the ambient stream first, so under a pp stage
559/// scope this bounds the STAGE's stream, which is what the caller is timing.
560///
561/// A free fn rather than the closure it replaced: `decode_batch_layers` (the pp stage seam)
562/// runs the instrumented layer loop, so the marker has to be callable from both the seam
563/// and its caller's epilogue. `batch_phase_on()` is a `OnceLock` memo, so per-call cost is
564/// the same atomic load the hoisted `ph_on` local was.
565fn ph_mark(
566    e: &Engine,
567    slot: usize,
568    last: &mut std::time::Instant,
569) -> Result<(), Box<dyn std::error::Error>> {
570    if batch_phase_on() {
571        e.stream().synchronize()?;
572        let now = std::time::Instant::now();
573        BATCH_PHASE.lock().unwrap()[slot] += (now - *last).as_secs_f64();
574        *last = now;
575    }
576    Ok(())
577}
578
579pub fn batch_phase_report() -> String {
580    let ph = BATCH_PHASE.lock().unwrap();
581    let tot: f64 = ph.iter().sum();
582    let mut rows: Vec<(usize, f64)> = ph.iter().copied().enumerate().collect();
583    rows.sort_by(|a, b| b.1.total_cmp(&a.1));
584    let mut s = format!(
585        "[batch-phase] total {:.1} ms (sync-bounded; shares rank, not walltime)\n",
586        tot * 1e3
587    );
588    for (i, v) in rows {
589        s += &format!(
590            "  {:>6.1} ms {:>5.1}%  {}\n",
591            v * 1e3,
592            v / tot * 100.0,
593            BATCH_PHASE_NAMES[i]
594        );
595    }
596    s
597}
598
599impl HybridModel {
600    /// Batched-decode width cap. 8 = the exactness-tier default (see the assert below);
601    /// MEMRA_DECODE_BATCH_CAP overrides for tier-probe measurement, clamped to 32.
602    pub fn decode_batch_cap() -> usize {
603        use std::sync::OnceLock;
604        static CAP: OnceLock<usize> = OnceLock::new();
605        *CAP.get_or_init(|| {
606            std::env::var("MEMRA_DECODE_BATCH_CAP")
607                .ok()
608                .and_then(|v| v.parse().ok())
609                .map(|c: usize| c.clamp(1, 32))
610                .unwrap_or(8)
611        })
612    }
613
614    /// EXACT-16 TIER admission (increment 3a, 2026-08-01, 5090 receipts
615    /// research/batched-tick-inc3-20260801): true iff EVERY matmul the batched decode step
616    /// runs has a per-(token,row) bit-exact kernel class at m=9..16 under the verify_exact
617    /// scope — i.e. the batched-mmvq b16 family (32-thread warp reduce, the exact m=1 mmvq
618    /// program per column) or the e4m3 grid.y=m mmvq catch-all. Q8_0 qualifies only with
619    /// the split-plane mirror (rp4, MEMRA_Q8RP): its b16 kernel exists only as the _rp twin.
620    /// Float matmuls (cuBLASLt, n-dependent reductions) and MoE FFNs disqualify the model.
621    /// Measured attribution for WHY the naked m=16 tier is not exact: the m>=16 arms
622    /// (MMQ int8-MMA `mul_mat_q` — MEMRA_PP_Q8MMQ default-on — and `qmatvec_gemm`, both
623    /// block-scale f32) and the m=9..15 dp4a tail (128-thread two-level reduce) all break
624    /// per-row bit-identity vs isolated decode (gate2 step-0 bit-diffs, maxdiff ~1.3-2.3e-1).
625    pub fn decode_batch_exact16_ok(&self) -> bool {
626        fn ok(w: &crate::model::GpuTensor) -> bool {
627            match w {
628                crate::model::GpuTensor::Quant { qtype, .. } => {
629                    *qtype == crate::QT_Q4_0 || *qtype == crate::QT_Q6_K
630                    || *qtype == crate::QT_F8_E4M3
631                    // BLOCK-128 FP8-ST (lane/rp-on-st, 2026-08-06): admitted now that the class
632                    // has a b16 batched kernel (`qmatvec_e4m3_blk_mmvq_b16`), bit-identical per
633                    // (token,row) to its m=1 launch. Before that kernel existed this class fell to
634                    // the grid.y=m form at every width — still EXACT, so the tier's correctness
635                    // bar was met, but it re-read the weight m times, which is why admitting it
636                    // without the kernel would have been a throughput trap rather than a win.
637                    || *qtype == crate::QT_F8_E4M3_BLK
638                    // NVFP4 (lane/rp-on-st, 2026-08-06) — THE blocker this lane measured. The
639                    // mixed FP8-ST 27B is 193 NVFP4 dense-MLP tensors, and this predicate is an
640                    // ALL over every matmul, so NVFP4's missing b16 refused the whole checkpoint
641                    // (`B=16 > cap 8 with no exact tier ... refused`) even with both e4m3 classes
642                    // admitted. It now has base + _rp b16 twins off its existing batched template
643                    // (bit-identical per (token,row) to the m=1 mmvq: same nibble decode, dp4a
644                    // order, ue4m3 scale, warp reduce). This also opens the tier for pure-NVFP4
645                    // GGUF models, which is a behavior change on the primary format — hence the
646                    // full decode-batch config+strict battery on both.
647                    || *qtype == crate::QT_NVFP4
648                    // Q4_K (lane/rp-on-st): named by MEMRA_EXACT16_WHY as the 9B NVFP4 GGUF's
649                    // refusing class (`L0.wqkv qtype=1`) — mixed NVFP4 checkpoints keep Q4_K
650                    // attention. Now has base + _rp b16.
651                    || *qtype == crate::QT_Q4_K
652                    // Q5_K (lane/rp-on-st): the FOURTH class the diagnostic named on the same 9B
653                    // GGUF (`L0.wqkv_gate qtype=3`). A shipped mixed checkpoint spreads ~500
654                    // matmuls over four/five classes, and this predicate is an ALL — so chunk 16
655                    // was unreachable for every real artifact until every class had a b16.
656                    || *qtype == crate::QT_Q5_K
657                    // Q8_0 NO LONGER requires the mirror (rp4): it has a base b16 too, so the
658                    // tier is reachable at zero VRAM. Named by the diagnostic as the FP8-ST
659                    // refusal — `L0.ssm_beta qtype=0 rp4=false`, a 23.9 MiB residual class that
660                    // was gating chunk 16 for a 16.4 GiB checkpoint.
661                    || *qtype == crate::QT_Q8_0
662                }
663                _ => false,
664            }
665        }
666        // WHY-NOT DIAGNOSTIC (lane/rp-on-st, 2026-08-06): this predicate is a bare bool over
667        // ~500 tensors, so a refusal produced only `B=16 > cap 8 with no exact tier ... refused`
668        // with no way to tell WHICH class refused. That cost this lane two wrong hypotheses (the
669        // rp mirror, then e4m3-only) before the NVFP4 gap was found. MEMRA_EXACT16_WHY=1 names
670        // the first refusing tensor + its qtype. Diagnostic-only per flags doctrine; default off,
671        // zero cost when unread.
672        let why = std::env::var("MEMRA_EXACT16_WHY").is_ok();
673        macro_rules! chk {
674            ($t:expr, $label:expr) => {{
675                let r = ok($t);
676                if !r && why {
677                    // qtype = -1 means the tensor is NOT Quant at all (a float/BF16/F16
678                    // container), which the tier can never admit — a distinct diagnosis from
679                    // "quantized, but in a class with no b16 kernel".
680                    let (qt, rp4) = match $t {
681                        crate::model::GpuTensor::Quant { qtype, rp4, .. } => {
682                            (*qtype, rp4.is_some())
683                        }
684                        _ => (-1, false),
685                    };
686                    eprintln!("[exact16] REFUSED by {} qtype={qt} rp4={rp4}", $label);
687                }
688                r
689            }};
690        }
691        let operations = self.plan.trunk_operations();
692        if operations.contains(&memra_gguf::model_plan::OperationKind::SwiGluOaiActivation)
693            || self.is_gemma4_e4b()
694            || crate::plan_backend::decode_batch_program(&self.plan)
695                == crate::plan_backend::DecodeBatchProgram::Gemma
696        {
697            if why {
698                eprintln!("[exact16] REFUSED by architecture (m3/gemma4)");
699            }
700            return false;
701        }
702        self.layers.iter().enumerate().all(|(li, l)| {
703            let mix_ok = match &l.mixer {
704                Mixer::Full(fa) => {
705                    chk!(&fa.wq, format!("L{li}.wq"))
706                        && chk!(&fa.wk, format!("L{li}.wk"))
707                        && chk!(&fa.wv, format!("L{li}.wv"))
708                        && chk!(&fa.wo, format!("L{li}.wo"))
709                }
710                Mixer::Linear(la) => {
711                    chk!(&la.wqkv, format!("L{li}.wqkv"))
712                        && chk!(&la.wqkv_gate, format!("L{li}.wqkv_gate"))
713                        && chk!(&la.ssm_beta, format!("L{li}.ssm_beta"))
714                        && chk!(&la.ssm_alpha, format!("L{li}.ssm_alpha"))
715                        && chk!(&la.ssm_out, format!("L{li}.ssm_out"))
716                }
717                // MLA rides its own increment-4 arm; never admitted to the exact-16 tier here.
718                Mixer::Mla(_) => {
719                    if why {
720                        eprintln!("[exact16] REFUSED by L{li} MLA mixer");
721                    }
722                    false
723                }
724                // KDA rides its own eager arm; never admitted to the exact-16 tier here.
725                Mixer::Kda(_) => {
726                    if why {
727                        eprintln!("[exact16] REFUSED by L{li} KDA mixer");
728                    }
729                    false
730                }
731            };
732            let ffn_ok = match &l.ffn {
733                crate::hybrid::Ffn::Dense {
734                    ffn_gate,
735                    ffn_up,
736                    ffn_down,
737                } => {
738                    chk!(ffn_gate, format!("L{li}.ffn_gate"))
739                        && chk!(ffn_up, format!("L{li}.ffn_up"))
740                        && chk!(ffn_down, format!("L{li}.ffn_down"))
741                }
742                crate::hybrid::Ffn::Moe(m) => {
743                    // lane/orndecode-20260822: the categorical refusal here was the c16 wall on
744                    // MoE checkpoints — serve chunked c16 into two B<=8 waves (agg flat ~700 on
745                    // ornith15 while the frozen vLLM column reads ~1190). The MoE stage itself is
746                    // width-exact by construction at decode widths: the dev/pairs expert kernels
747                    // replay one per-(token,expert) program whose arithmetic never sees batch
748                    // width, the router (gemv f32 + sigmoid + topk) is row-wise, and the shexp
749                    // trio rides the per-column decode-exact arm at every verify width
750                    // (t in 2..PRIME_MIN_T), so no b16 qmatvec class is ever demanded of it.
751                    // "By construction" is NOT the qualification — the CSR-NVFP4
752                    // batch-composition defect (v0.99.0, research/samplat-20260821) shipped on
753                    // exactly that reasoning. STATUS (orndecode, 2026-08-22): byte gates are
754                    // GREEN on ornith15 (decode-batch-gate config gate2+gate3 PASS at B=12 and
755                    // B=16, bit-checked vs isolated) but the tier LOSES throughput today —
756                    // B=16 exact measured 220 agg vs 551 at B=8 same-window, because the
757                    // exact-verify scope drives the shexp trio (and friends) to per-column m=1
758                    // decode-exact launches. MEMRA_EXACT16_MOE=1 is therefore an OPT-IN
759                    // measurement door until the b16-class stage kernels land; serve must not
760                    // pick a tier that halves the aggregate it exists to raise.
761                    if std::env::var("MEMRA_EXACT16_MOE").as_deref() != Ok("1") {
762                        if why {
763                            eprintln!(
764                                "[exact16] REFUSED by L{li} MoE ffn (opt-in: MEMRA_EXACT16_MOE=1 \
765                                 — byte-safe but slower than two B<=8 waves today)"
766                            );
767                        }
768                        false
769                    } else {
770                        match (&m.gate_shexp, &m.up_shexp, &m.down_shexp) {
771                            (Some(g), Some(u), Some(d)) => {
772                                chk!(g, format!("L{li}.gate_shexp"))
773                                    && chk!(u, format!("L{li}.up_shexp"))
774                                    && chk!(d, format!("L{li}.down_shexp"))
775                            }
776                            _ => true,
777                        }
778                    }
779                }
780            };
781            mix_ok && ffn_ok
782        }) && chk!(&self.output, "output".to_string())
783    }
784
785    /// Opt-in/A-B seam for the eager B=1 fusion program. `MEMRA_SERVE_B1FAST=1` sends an
786    /// eligible solo tick through that program; unset/other values keep B=1 on the generic
787    /// batched body, the same numeric class used at B>=2.
788    ///
789    /// EXACTNESS, stated precisely (measured on-box 2026-08-05, sm_120 q9 NVFP4-MTP):
790    /// the fast path is BIT-IDENTICAL TO `decode_step_h` — decode-batch-gate's STRICT
791    /// gate1 (`--mode strict`) PASSes with it ON and FAILs with it OFF at maxdiff
792    /// 1.591e-1. It is deliberately NOT bit-identical to the batched body: the two
793    /// carry a decode-config FP-composition gap (same class gate1's config mode measures).
794    /// That gap became correctness-visible under live load: Step35, Q35-MoE, and finally
795    /// dense Q27 all produced load-history-dependent token streams, including early EOS,
796    /// when a request crossed between the two programs. The generic body is therefore the
797    /// correctness default; the eager program remains available only for fixed-solo A/Bs.
798    /// Historical token-stream/performance receipts:
799    /// research/servepath-p2-20260805 (greedy 150 ids + seeded-sampled identical to the
800    /// run-gen oracle AND cross-arm, so the gap is sub-token here as designed).
801    ///
802    /// Read fresh (an `AtomicU8` memo, not a `OnceLock`): decode-batch-gate flips this
803    /// seam BETWEEN gates in-process — gate1 needs the fast path ON to prove bit-identity,
804    /// gate2 needs it pinned OFF to keep testing the batched body. A latch-once read would
805    /// bake whichever gate ran first, so the gate could never test both sides. The memo
806    /// caches the parse but `set_b1_fast` invalidates it.
807    pub fn b1_fast_on() -> bool {
808        // 0 = unknown/invalidated, 1 = off, 2 = on
809        match Self::b1_fast_memo().load(std::sync::atomic::Ordering::Relaxed) {
810            1 => false,
811            2 => true,
812            _ => {
813                let value = std::env::var("MEMRA_SERVE_B1FAST").ok();
814                let on = b1_fast_env_on(value.as_deref());
815                Self::b1_fast_memo()
816                    .store(if on { 2 } else { 1 }, std::sync::atomic::Ordering::Relaxed);
817                on
818            }
819        }
820    }
821
822    fn b1_fast_memo() -> &'static std::sync::atomic::AtomicU8 {
823        static MEMO: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
824        &MEMO
825    }
826
827    /// Test/gate seam: force the B=1 fast path on or off for the rest of the process,
828    /// overriding the env. Used by decode-batch-gate to exercise the opt-in eager arm and
829    /// pin gate2's default reference arm.
830    pub fn set_b1_fast(on: bool) {
831        Self::b1_fast_memo().store(if on { 2 } else { 1 }, std::sync::atomic::Ordering::Relaxed);
832    }
833
834    /// Whether this architecture may switch a live serving row onto the eager B=1 fusion
835    /// class. Qwen35-MoE must stay on the batched trunk at every width: its eager and batched
836    /// hybrid/MoE walks are each deterministic, but crossing B=1 -> B>=2 changes greedy token
837    /// ids and can introduce an early EOS (Q35 sellgate, 2026-08-12).
838    pub fn b1_fast_plan_eligible(&self) -> bool {
839        b1_fast_plan_eligible(&self.plan)
840    }
841
842    /// H3 body: the m=1 FUSED trunk (`decode_layers_eager` — shared verbatim with
843    /// `decode_step_h`/the ppN stages) plus the batched path's own serving epilogue
844    /// (grammar mask, device sample, lean-logits park). See the call-site comment in
845    /// `decode_step_batch_sampled_lean_masked` for why this is bit-identical.
846    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
847    fn decode_step_b1_fast(
848        &self,
849        e: &Engine,
850        token: u32,
851        caches: &mut [&mut Cache],
852        samp: &[Option<DevSamp>],
853        masks: &[Option<(&CudaSlice<u32>, usize)>],
854        lean: bool,
855    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
856        let n_embd = self.cfg.n_embd as usize;
857        let eps = self.cfg.rms_eps;
858        let pos = caches[0].pos;
859        let pos_d = e.htod_i32(&[pos as i32])?;
860        let x = e.htod(&self.embd.try_gather(n_embd, &[token])?)?;
861        // the SHARED m=1 trunk: same function decode_step_h runs, so every m=1 fusion
862        // (cross-layer add+norm+q8_1, fused SwiGLU, lever 1's gate+up dual) fires here.
863        let x = self.decode_layers_eager(e, x, 0, self.layers.len(), &pos_d, pos, caches[0])?;
864        let mut hn = e.uninit(n_embd)?;
865        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
866        let logits = e.matmul(&self.output, &hn, 1)?;
867
868        // ---- epilogue: byte-for-byte the batched path's, at b_n=1 ----
869        let n_vocab = self.output.out_features();
870        let mut logits = logits;
871        let mut pristine: Option<CudaSlice<f32>> = None;
872        if let Some((mask, words)) = masks.first().copied().flatten() {
873            assert!(
874                samp.first().and_then(Option::as_ref).is_some(),
875                "grammar-masked row 0 must request a device sample"
876            );
877            if lean {
878                let cache = &mut caches[0];
879                if cache
880                    .last_logits_dev
881                    .as_ref()
882                    .map(|d| d.len() < n_vocab)
883                    .unwrap_or(true)
884                {
885                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
886                }
887                let dst = cache.last_logits_dev.as_mut().unwrap();
888                e.dtod_copy_view(&logits.slice(0..n_vocab), dst)?;
889            } else {
890                let mut p = e.uninit(n_vocab)?;
891                e.dtod_copy_view(&logits.slice(0..n_vocab), &mut p)?;
892                pristine = Some(p);
893            }
894            e.mask_logits_col(&mut logits, mask, 0, n_vocab, words)?;
895        }
896
897        let mut next: Vec<Option<u32>> = vec![None; 1];
898        if let Some(s) = samp.first().and_then(Option::as_ref) {
899            let mut toks = e.alloc_u32_zeroed(1)?;
900            // Filtered-greedy degenerates to plain argmax (the max always survives every
901            // truncation filter), so temp<=0 short-circuits regardless of filters.
902            let filtered = s.temp > 0.0 && (s.top_k > 0 || s.top_p < 1.0 || s.min_p > 0.0);
903            if s.temp <= 0.0 {
904                e.argmax_token_device_col(&logits, 0, n_vocab, &mut toks, 0)?;
905            } else if filtered {
906                let mut pb = e.zeros(n_vocab)?;
907                self.devsample_filtered_col(
908                    e, &logits, 0, n_vocab, s.temp, s.seed, s.ctr, s.top_k, s.top_p, s.min_p,
909                    &mut pb, &mut toks, 0,
910                )?;
911            } else {
912                let mut pb = e.zeros(n_vocab)?;
913                e.gumbel_perturb_col(&logits, 0, &mut pb, n_vocab, s.seed, s.ctr, s.temp)?;
914                e.argmax_token_device_col(&pb, 0, n_vocab, &mut toks, 0)?;
915            }
916            next[0] = Some(e.dtoh_u32(&toks)?[0]);
917        }
918
919        let sampled = samp.first().and_then(Option::as_ref).is_some();
920        let rows: Vec<Vec<f32>> = if lean && sampled {
921            if masks.first().copied().flatten().is_none() {
922                let cache = &mut caches[0];
923                if cache
924                    .last_logits_dev
925                    .as_ref()
926                    .map(|d| d.len() < n_vocab)
927                    .unwrap_or(true)
928                {
929                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
930                }
931                let dst = cache.last_logits_dev.as_mut().unwrap();
932                e.dtod_copy_view(&logits.slice(0..n_vocab), dst)?;
933            }
934            vec![Vec::new()]
935        } else if let Some(p) = pristine.as_ref() {
936            vec![e.dtoh(p)?]
937        } else {
938            vec![e.dtoh(&logits)?]
939        };
940        // decode_layers_eager does NOT advance cache.pos (decode_step_h advances it after
941        // the head); the batched path advances every cache at the tail — same here.
942        caches[0].pos += 1;
943        Ok((rows, next))
944    }
945
946    /// One filtered device draw for stacked-logits row `col`: `filter_stats` solves the
947    /// single unnormalized-prob floor that encodes top-k AND top-p AND min-p (block-internal
948    /// binary search, bit-stable), then the filtered gumbel perturb + argmax draws one token
949    /// from the truncated softmax into `toks[slot]`. All device-side — no stat D2H, no row
950    /// copy; the only host traffic stays the caller's one [B]-u32 token readback.
951    #[allow(clippy::too_many_arguments)]
952    fn devsample_filtered_col(
953        &self,
954        e: &Engine,
955        logits: &CudaSlice<f32>,
956        col: usize,
957        n_vocab: usize,
958        temp: f32,
959        seed: u64,
960        ctr: u32,
961        top_k: i32,
962        top_p: f32,
963        min_p: f32,
964        pb: &mut CudaSlice<f32>,
965        toks: &mut CudaSlice<u32>,
966        slot: usize,
967    ) -> Result<(), Box<dyn std::error::Error>> {
968        let rows = e.htod_i32(&[col as i32])?;
969        let mut th = e.zeros(1)?;
970        let mut z = e.zeros(1)?;
971        let mut mx = e.zeros(1)?;
972        e.filter_stats(
973            logits, n_vocab, &rows, &mut th, &mut z, &mut mx, n_vocab, 1, temp, top_k, top_p, min_p,
974        )?;
975        e.gumbel_perturb_filtered_col(logits, col, pb, n_vocab, seed, ctr, temp, &mx, &th, 0)?;
976        e.argmax_token_device_col(pb, 0, n_vocab, toks, slot)?;
977        Ok(())
978    }
979
980    /// One batched greedy-decode step over B independent sequences.
981    /// `tokens[b]` is sequence b's input token; `caches[b]` its private cache (position,
982    /// quantized KV, GDN/conv state). Returns the B logits rows (host, [n_vocab] each).
983    /// Each cache's pos/len advance exactly as `decode_step_h` would.
984    pub fn decode_step_batch(
985        &self,
986        e: &Engine,
987        tokens: &[u32],
988        caches: &mut [&mut Cache],
989    ) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
990        let (rows, _) = self.decode_step_batch_sampled(e, tokens, caches, &[])?;
991        Ok(rows)
992    }
993
994    /// `decode_step_batch` + DEVICE-SIDE SAMPLING for eligible rows (the batched-tick lever,
995    /// 2026-08-01): the host sampler's temp-path is O(n_vocab) with a full-vocab exp per row
996    /// (measured 1.36 ms/row at the 9B's 248320 vocab = 10.9 ms/tick at B=8 — the single
997    /// largest component of the serving tick). Here each requested row samples ON DEVICE
998    /// between the lm_head matmul and the logits D2H:
999    ///   temp <= 0 (greedy): the 2-pass device argmax — bit-identical to host argmax
1000    ///     (argmax-gate contract, same kernels as the dc serving path).
1001    ///   temp > 0: gumbel_perturb(seed, ctr, temp) + the same argmax = ONE categorical draw
1002    ///     from softmax(logits/temp) — the sampled-spec Philox machinery. Deterministic per
1003    ///     (seed, ctr) and INDEPENDENT of batch composition (the isolation contract;
1004    ///     decode-batch-gate gate3). NOTE: the draw stream differs from the host sampler's
1005    ///     SplitMix64 (distribution-equal, seed-deterministic, NOT byte-equal to the old
1006    ///     host draws) — greedy rows are unchanged bit-exact.
1007    /// `samp[bi] = Some(DevSamp { .. })` requests a device sample for row bi; the full
1008    /// logits rows are still returned (worker keeps last_logits semantics + fallback rows).
1009    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1010    pub fn decode_step_batch_sampled(
1011        &self,
1012        e: &Engine,
1013        tokens: &[u32],
1014        caches: &mut [&mut Cache],
1015        samp: &[Option<DevSamp>],
1016    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1017        self.decode_step_batch_sampled_lean(e, tokens, caches, samp, false)
1018    }
1019
1020    /// `decode_step_batch_sampled` + LEAN LOGITS (increment 2 component 3, 2026-08-01):
1021    /// with `lean`, device-sampled rows SKIP the [n_vocab] logits D2H (9.4%/32.5% of the
1022    /// pre-/post-inc2 tick profile) — their returned row is EMPTY. The audit-mapped
1023    /// consumers: (a) the next tick's host sample — never fires, `device_next` carries the
1024    /// token; (b) the graph-promotion argmax — reads only prefill logits (generated empty);
1025    /// (c) the KV-reuse pool park at retire — the REAL consumer, served by a per-cache
1026    /// device park: the row is dtod-copied into `cache.last_logits_dev` (device bandwidth)
1027    /// and D2H'd ONCE at retire by the worker. Rows without a device sample keep a per-row
1028    /// D2H. `lean=false` is bit-for-bit the previous method (gates + non-serving callers).
1029    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1030    pub fn decode_step_batch_sampled_lean(
1031        &self,
1032        e: &Engine,
1033        tokens: &[u32],
1034        caches: &mut [&mut Cache],
1035        samp: &[Option<DevSamp>],
1036        lean: bool,
1037    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1038        self.decode_step_batch_sampled_lean_masked(e, tokens, caches, samp, &[], lean)
1039    }
1040
1041    /// `decode_step_batch_sampled_lean` + GRAMMAR MASKS (constrained decoding, 2026-08-03):
1042    /// `masks[bi] = Some((packed_bitset, words))` bans every unset-bit vocab id on row bi
1043    /// (mask_logits_f32, -FLT_MAX) BETWEEN the lm_head matmul and the device sampler, so a
1044    /// constrained row rides the SAME device-sample/lean-logits tick as everyone else — no
1045    /// full-row D2H, no host O(n_vocab) sample. Contract: a masked row must also request a
1046    /// device sample. The row's PRISTINE logits are preserved for their consumers before the
1047    /// in-place ban: lean rows park the unmasked row into `cache.last_logits_dev` (the
1048    /// retire-time reuse-pool park stays unmasked — continuations resume grammar-free, the
1049    /// v1 host-path contract), non-lean rows D2H the unmasked row. `masks = &[]` is
1050    /// bit-for-bit the unmasked method.
1051    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1052    pub fn decode_step_batch_sampled_lean_masked(
1053        &self,
1054        e: &Engine,
1055        tokens: &[u32],
1056        caches: &mut [&mut Cache],
1057        samp: &[Option<DevSamp>],
1058        masks: &[Option<(&CudaSlice<u32>, usize)>],
1059        lean: bool,
1060    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1061        self.decode_step_batch_sampled_lean_masked_schedule(
1062            e, tokens, caches, samp, masks, lean, None, None,
1063        )
1064    }
1065
1066    /// Whether the generic, unsplit batched trunk can leave its result on the device for one
1067    /// scheduler boundary. The pending path is intentionally c=1-only today: PP stages, model
1068    /// specific batched programs, and the fixed-solo fusion arm each have different output
1069    /// ownership and keep their established synchronous readback contract.
1070    pub fn decode_step_overlap_eligible(&self) -> bool {
1071        !batch_phase_on()
1072            && crate::pp::pp_cuts(self.layers.len()).is_none()
1073            // mHC trunks are excluded: the pending (deferred-readback) epilogue is only
1074            // wired for the generic trunk body, and the hyper walk keeps the synchronous
1075            // readback contract (see the named refusal in `_pending`).
1076            && self.hyper.is_none()
1077            && !Self::b1_fast_on()
1078            && self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeBatch)
1079            && crate::plan_backend::decode_batch_program(&self.plan)
1080                == crate::plan_backend::DecodeBatchProgram::Generic
1081    }
1082
1083    /// Enqueue one generic B=1 decode and defer its D2H until [`PendingBatchStep::wait`]. This
1084    /// is the engine half of the overlap scheduler: the server can publish the token selected
1085    /// from step n before it waits for step n+1's logits.
1086    pub fn decode_step_batch_sampled_lean_masked_pending(
1087        &self,
1088        e: &Engine,
1089        tokens: &[u32],
1090        caches: &mut [&mut Cache],
1091        samp: &[Option<DevSamp>],
1092        masks: &[Option<(&CudaSlice<u32>, usize)>],
1093        lean: bool,
1094    ) -> Result<PendingBatchStep, Box<dyn std::error::Error>> {
1095        if self.hyper.is_some() {
1096            return Err(
1097                "decode_step_batch_sampled_lean_masked_pending: the overlap scheduler's \
1098                 deferred-readback step is not wired for the HyperConnections residual — \
1099                 the hyper batched walk keeps the synchronous epilogue (its `pending_out` \
1100                 plumbing through `decode_step_batch_hyper` does not exist yet). Serve mHC \
1101                 sessions through the synchronous batched chain or the eager per-session \
1102                 loop; `decode_step_overlap_eligible` already reports false for this trunk."
1103                    .into(),
1104            );
1105        }
1106        if tokens.len() != 1 || caches.len() != 1 {
1107            return Err("overlap scheduler requires a single decode row".into());
1108        }
1109        if !self.decode_step_overlap_eligible() {
1110            return Err(
1111                "overlap scheduler is unavailable for this model, topology, or diagnostic arm"
1112                    .into(),
1113            );
1114        }
1115        let mut pending = None;
1116        let _ = self.decode_step_batch_sampled_lean_masked_schedule(
1117            e,
1118            tokens,
1119            caches,
1120            samp,
1121            masks,
1122            lean,
1123            None,
1124            Some(&mut pending),
1125        )?;
1126        pending.ok_or_else(|| "overlap scheduler did not produce a pending step".into())
1127    }
1128
1129    /// Worker-scheduled twin of [`Self::decode_step_batch_sampled_lean_masked`]. The worker
1130    /// supplies the balanced dual-wave boundary it used when forming this tick. Direct engine
1131    /// callers keep the automatic midpoint above; the explicit seam makes scheduler chunking and
1132    /// engine execution one checked contract instead of two coincident width calculations.
1133    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1134    #[allow(clippy::too_many_arguments)] // allow: the parameter list mirrors the kernel/FFI/call contract; bundling into a struct is a refactor, not a lint fix
1135    pub fn decode_step_batch_sampled_lean_masked_scheduled(
1136        &self,
1137        e: &Engine,
1138        tokens: &[u32],
1139        caches: &mut [&mut Cache],
1140        samp: &[Option<DevSamp>],
1141        masks: &[Option<(&CudaSlice<u32>, usize)>],
1142        lean: bool,
1143        dual_wave_mid: usize,
1144    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1145        if self.hyper.is_some() {
1146            return Err(
1147                "decode_step_batch_sampled_lean_masked_scheduled: the dual-active PP-2 \
1148                 wave schedule has no HyperConnections trunk — `decode_step_batch_dual`'s \
1149                 two host walkers drive the generic/step35 layer bodies only, and no \
1150                 dual-wave twin of `hyper_batch_range_decode` exists. mHC chunks are \
1151                 serial ticks: the worker's chunk policy must not form dual waves for \
1152                 this topology (decode_step_batch_hyper owns the serial PP-N split)."
1153                    .into(),
1154            );
1155        }
1156        self.decode_step_batch_sampled_lean_masked_schedule(
1157            e,
1158            tokens,
1159            caches,
1160            samp,
1161            masks,
1162            lean,
1163            Some(dual_wave_mid),
1164            None,
1165        )
1166    }
1167
1168    #[allow(clippy::too_many_arguments)]
1169    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1170    fn decode_step_batch_sampled_lean_masked_schedule(
1171        &self,
1172        e: &Engine,
1173        tokens: &[u32],
1174        caches: &mut [&mut Cache],
1175        samp: &[Option<DevSamp>],
1176        masks: &[Option<(&CudaSlice<u32>, usize)>],
1177        lean: bool,
1178        scheduled_dual_mid: Option<usize>,
1179        pending_out: Option<&mut Option<PendingBatchStep>>,
1180    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1181        for cache in caches.iter() {
1182            cache.ensure_usable("decode_step_batch")?;
1183        }
1184        if crate::pp::pp_cuts(self.layers.len()).is_some()
1185            && !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline)
1186        {
1187            return Err("pipeline rewrite is not qualified for batched decode".into());
1188        }
1189        if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeBatch) {
1190            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::DecodeEager) {
1191                return Err("neither batch nor eager decode rewrite is qualified".into());
1192            }
1193            if masks.iter().any(Option::is_some) {
1194                return Err(
1195                    "unqualified batch rewrite cannot fall back with device grammar masks".into(),
1196                );
1197            }
1198            if tokens.len() != caches.len() {
1199                return Err("batch fallback token/cache shape mismatch".into());
1200            }
1201            static ONCE: std::sync::Once = std::sync::Once::new();
1202            ONCE.call_once(|| {
1203                eprintln!(
1204                    "[rewrite] decode-batch.v1 unqualified; using receipt-backed native eager rows"
1205                );
1206            });
1207            let mut rows = Vec::with_capacity(tokens.len());
1208            for (token, cache) in tokens.iter().copied().zip(caches.iter_mut()) {
1209                rows.push(self.decode_step_h(e, token, cache)?.0);
1210            }
1211            return Ok((rows, vec![None; tokens.len()]));
1212        }
1213        // NOTE (inc3 3c, 2026-08-01, KILLED ARM): a deferred-token-readback variant (all
1214        // chunks of a tick writing device-sampled tokens into one shared buffer, ONE
1215        // dtoh_u32 after the last chunk instead of one per chunk) measured FLAT at serve
1216        // level on the 5090 (N=4 medians within +-0.7% at c=8/16/32 — 3 saved syncs
1217        // against a ~100 ms weight-bound tick is ~0.1%, below resolution). Killed per the
1218        // flags doctrine; receipts research/batched-tick-inc3-20260801 (serve-points.jsonl
1219        // base vs defer arms) are the record. The per-chunk [B]-u32 readback below IS the
1220        // tick's only steady-state D2H — one per chunk, none per seq.
1221        let b_n = tokens.len();
1222        assert!(
1223            b_n >= 1 && b_n == caches.len(),
1224            "tokens/caches length mismatch"
1225        );
1226        // ---- mHC DOOR (lane/glm53-batched-decode, 2026-08-28): the HyperConnections trunk
1227        // takes its OWN batched walk. Every body below this point runs the serial residual —
1228        // on an hc model that is a DIFFERENT function computed fluently (the failure class
1229        // `refuse_hyper` exists for) — so the hyper route must come before the pp door, the
1230        // b1 fast path, and the width tiers, and it owns its own PP-N stage split inside.
1231        // The dual-wave and pending entries refused above with named reasons; this guard is
1232        // the defense-in-depth backstop for a future caller that reaches here with either.
1233        if self.hyper.is_some() {
1234            if scheduled_dual_mid.is_some() {
1235                return Err(
1236                    "decode_step_batch: a dual-wave schedule reached the hyper trunk — \
1237                     no dual-wave twin of hyper_batch_range_decode exists; mHC chunks are \
1238                     serial ticks"
1239                        .into(),
1240                );
1241            }
1242            if pending_out.is_some() {
1243                return Err(
1244                    "decode_step_batch: the pending (deferred-readback) epilogue reached \
1245                     the hyper trunk — the hyper walk keeps the synchronous readback \
1246                     contract"
1247                        .into(),
1248                );
1249            }
1250            return self.decode_step_batch_hyper(e, tokens, caches, samp, masks, lean);
1251        }
1252        let _pp_walk =
1253            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
1254                let rt = crate::pp::PpNRt::get(e)?;
1255                Some(rt.acquire_walk("decode_step_batch")?)
1256            } else {
1257                None
1258            };
1259        // ---- PP DOOR: THE BATCHED STAGE SPLIT (pp2-batch 2026-08-06) ----------------------
1260        // Until this increment this body had NO pp arm: it walked lo=0..n_layers on the
1261        // primary engine's stream, with no stage split, no boundary, and no `rt.enter()`. With
1262        // the door open and a sharded cross-device placement, every projection for the remote
1263        // stages' layers was read over PCIe, per step, silently — measured 7.4 vs 208.9 tok/s
1264        // at B=1 (28x), 47.4 vs 657.0 at B=8 (13.9x) on a PRO 6000 pair over Gen5 x16 P2P.
1265        // Nothing failed or warned, because peer reads return identical bytes and all three
1266        // `decode-batch-gate` gates PASS on that config — the failure mode was performance,
1267        // and a green exactness battery hid it. `pp2-hardening` made that regime FAIL CLOSED
1268        // (research/pp2-hardening-20260806); this lane makes it legitimately split, so the
1269        // refusal lifts for the batched path.
1270        //
1271        // `decode_step_batch_ppn` runs each stage's layer range through that stage's engine
1272        // and stream with a [B, n_embd] boundary transfer between them, i.e. every stage
1273        // touches only LOCAL weights and LOCAL cache state. The refusal below still guards
1274        // the residue: the door open with `MEMRA_PP_STREAMS=0` (the same-stream rollback,
1275        // which also disables the sharded loader, so nothing is remote — `pp_shard_off` and
1276        // `pp2_streams_off` both make `pp_sharded_cross_device()` false) or a placement whose
1277        // PpNRt fails to build. Keeping the call means a future path that reaches here in a
1278        // remote regime still refuses instead of regressing 28x.
1279        if let Some(fence) = crate::pp::pp_cuts(self.layers.len())
1280            && !crate::pp::pp2_streams_off()
1281            && crate::pp::batch_pp_on()
1282        {
1283            let n_stages = fence.len() - 1;
1284            let wave_on = crate::pp::pp_wave_on()
1285                .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
1286            if crate::pp::pp_wave_route_enabled(wave_on, crate::pp::pp2_overlap(), n_stages, b_n) {
1287                if scheduled_dual_mid.is_some() {
1288                    return Err(
1289                        "decode_step_batch: worker supplied a PP2 midpoint to a PP3/PP4 wavefront"
1290                            .into(),
1291                    );
1292                }
1293                return self
1294                    .decode_step_batch_wavefront(e, tokens, caches, samp, masks, lean, &fence);
1295            }
1296            // Auto (flipped default) routes dual only in the re-gated regime and
1297            // degrades serially elsewhere; Forced keeps every ineligible placement on
1298            // the refusing dual body so the binding negative cells stay reachable.
1299            let route_dual = crate::pp::dual_pp_route(
1300                crate::pp::dual_pp_mode(),
1301                b_n,
1302                fence.len() - 1,
1303                crate::pp::pp2_overlap(),
1304                crate::pp::pp_host_bounce_active(),
1305            );
1306            if route_dual {
1307                let mid = scheduled_dual_mid
1308                    .or_else(|| crate::pp::dual_pp_wave_mid(b_n))
1309                    .expect("dual PP B>=2 must have a wave midpoint");
1310                return self
1311                    .decode_step_batch_dual(e, tokens, caches, samp, masks, lean, &fence, mid);
1312            }
1313            return self.decode_step_batch_ppn(e, tokens, caches, samp, masks, lean, &fence);
1314        }
1315        if scheduled_dual_mid.is_some() {
1316            return Err(
1317                "decode_step_batch: worker supplied a dual-wave schedule but the PP-2 dual path is unavailable"
1318                    .into(),
1319            );
1320        }
1321        crate::pp::refuse_unsplit_if_remote(
1322            "decode_step_batch",
1323            "drop MEMRA_PP_STREAMS=0 / MEMRA_BATCH_PP=0 so the batched path takes its OWN \
1324             stage split (decode_step_batch_ppn), or serve single-stream over the eager pp \
1325             arm (decode_step_h), which is also split",
1326        )?;
1327        // ---- H3: B=1 FAST-PATH (serve-path phase 2, 2026-08-05) ----------------------------
1328        // At b_n==1 every projection below calls `matmul_pre(.., b_n)` with m=1, which is
1329        // ALREADY the m=1 mmvq dispatch — so the m=1 *kernel family* was never the gap. What
1330        // this body does NOT have is the m=1 *fusion chain* that `decode_step_h` carries:
1331        //   - the cross-layer add+norm+quantize fusion (`add_rms_norm_q8_1`: 3 launches -> 1),
1332        //   - the fused SwiGLU epilogue (`silu_mul_scaled_q8_1`: folds ffn_down's quantize
1333        //     into its producer) and, with it, `matmul_pre_dual_noscale`'s gate+up pair
1334        //     fusion — i.e. phase-1 LEVER 1.
1335        // Routing b_n==1 through `decode_layers_eager` (the SHARED trunk `decode_step_h` and
1336        // the ppN stages already use, lifted verbatim — not a copy) makes every present and
1337        // future m=1 lever fire on the opt-in path automatically. The epilogue (grammar mask ->
1338        // device sample -> lean logits park) stays exactly as the batched path runs it; the trunk's
1339        // different FP composition is why this path cannot be a load-changing default.
1340        // BIT-IDENTITY: the trunk is the same function `decode_step_h` calls, and every
1341        // fusion it enables is kernel-check-pinned bit-identical to its unfused sequence
1342        // (add_rms_norm == add;rms_norm | _q8_1 == +quantize_q8_1 | dual_noscale == two
1343        // matmul_pre_noscale). Gate: decode-batch-gate B=1 vs decode_step_h + serve stream
1344        // identity. MEMRA_SERVE_B1FAST=1 is the fixed-solo opt-in/A-B seam; the default
1345        // stays on this function's generic body so batch-width changes cannot change the
1346        // FP program mid-request.
1347        if b_n == 1
1348            && Self::b1_fast_on()
1349            && !samp.iter().flatten().any(|s| s.penalty.is_some())
1350            && self.b1_fast_plan_eligible()
1351            && !self.is_gemma4_e4b()
1352            && crate::plan_backend::decode_batch_program(&self.plan)
1353                == crate::plan_backend::DecodeBatchProgram::Generic
1354            && !self
1355                .plan
1356                .trunk_operations()
1357                .contains(&memra_gguf::model_plan::OperationKind::SwiGluOaiActivation)
1358            && crate::pp::pp_cuts(self.layers.len()).is_none()
1359            && !e.verify_exact_on()
1360        {
1361            return self.decode_step_b1_fast(e, tokens[0], caches, samp, masks, lean);
1362        }
1363        // MEMRA_DECODE_BATCH_CAP (experimental door, serving-lane tier probe 2026-08-01):
1364        // default 8 keeps the v1 exactness policy — B=2..8 rides the verify-tier batched
1365        // mmvq arms, per-row bit-identical to isolated m=1 decode. Values >8 are a
1366        // MEASUREMENT DOOR ONLY: m=9..15 falls to the grid.y=m dp4a tail (m weight
1367        // re-reads + a different reduce shape) and m>=16 crosses into the GEMM tier
1368        // (block-scale f32 rounding) — BOTH break the "byte-identical to isolated"
1369        // serving contract. Never default this above 8 without the batched-tier
1370        // exactness policy landing.
1371        let cap = Self::decode_batch_cap();
1372        // EXACT-16 TIER (increment 3a): chunks of 9..=16 are admitted WITHOUT the env door
1373        // when every matmul has a bit-exact b16-class kernel (see decode_batch_exact16_ok).
1374        // The verify_exact scope below pins that dispatch for the whole step: it turns off
1375        // the m>=16 GEMM arms (qmatvec_gemm + MMQ + fp8/f16/fp4 — all block-scale/foreign
1376        // numeric configs) so every projection rides the batched-mmvq b16 tier, which is
1377        // per-(token,row) bit-identical to isolated m=1 decode (gate2 bit-strength PASS at
1378        // B=12/16, s32+s160, 5090 receipts research/batched-tick-inc3-20260801). Without
1379        // the exact tier, B>cap stays refused; the env door (MEMRA_DECODE_BATCH_CAP) keeps
1380        // its old meaning as the non-exact measurement probe.
1381        let exact16 = b_n > 8 && b_n <= 16 && self.decode_batch_exact16_ok();
1382        assert!(
1383            b_n <= cap || exact16,
1384            "decode_step_batch: B={b_n} > cap {cap} with no exact tier — refused. Either \
1385             B>16 (there is NO exact kernel class above 16: m>16 crosses GEMM/dp4a numeric \
1386             configs; the serve scheduler chunks wider concurrency into <=16 groups instead), \
1387             or some matmul in this checkpoint has no bit-exact b16 kernel — run with \
1388             MEMRA_EXACT16_WHY=1 to see which tensor and qtype refuses"
1389        );
1390        struct ExactScope<'a>(&'a Engine, bool);
1391        impl Drop for ExactScope<'_> {
1392            fn drop(&mut self) {
1393                if self.1 {
1394                    self.0.set_verify_exact(false);
1395                }
1396            }
1397        }
1398        let _exact_scope = ExactScope(e, exact16);
1399        if exact16 {
1400            e.set_verify_exact(true);
1401        }
1402        // gemma4: NO batched arm at any B (per-layer SWA/global geometry, hd-512 MQA globals,
1403        // weightless V-norm, softcapped head — none of it in the generic body below). This was
1404        // an assert until 2026-08-07: one serve request panicked the worker, the respawn
1405        // re-panicked on the queued request, and the process FATALed
1406        // (research/gemma4-serve-20260807/raw/repro-panic-server-*.log). The worker now routes
1407        // gemma4 sessions to the per-session eager loop and never calls here; this Err is the
1408        // defense-in-depth backstop — a future path that reaches it refuses PER-REQUEST
1409        // instead of killing the process. The eager arm (gemma4_decode_step_h) is the
1410        // supported decode.
1411        let batch_program = crate::plan_backend::decode_batch_program(&self.plan);
1412        if self.is_gemma4_e4b() || batch_program == crate::plan_backend::DecodeBatchProgram::Gemma {
1413            // BATCHED ARM (lane/gemma-batched, 2026-08-16): the dense 31B gets its own
1414            // per-session batched walk (gemma4_decode_batch) — DEFAULT ON since the owner
1415            // flip (MEMRA_GEMMA4_BATCH=0 = the eager kill switch). Same shape law as
1416            // step35: projections/norms/rope/FFN/head run at m=B (one weight stream, B
1417            // rows — decode is weight-BW-bound), KV append + fa_decode stay a per-session
1418            // loop (each session's own len drives its SWA/global view). E4B keeps its
1419            // dedicated decode; it never enters here.
1420            if batch_program == crate::plan_backend::DecodeBatchProgram::Gemma
1421                && !self.is_gemma4_e4b()
1422                && Self::gemma4_batch_on()
1423            {
1424                return self.gemma4_decode_batch(e, tokens, caches, samp, masks, lean);
1425            }
1426            return Err(
1427                "decode_step_batch has no gemma4 arm for this model class (per-layer \
1428                        swa/global geometry, softcapped head; the dense-31B batched arm is \
1429                        default-on, MEMRA_GEMMA4_BATCH=0 forces eager) — serve gemma4 on the \
1430                        eager per-session path"
1431                    .into(),
1432            );
1433        }
1434        // step35 (lane/step35-batched-decode, 2026-08-08): its OWN batched walk. The generic
1435        // body below is the uniform Full arm — global n_head, 128-dim rope on every layer, no
1436        // SWA window, no head-wise gate — which on step35 produced HTTP-200 GARBAGE at c>1
1437        // (research/step-sku-20260807/raw/b2ab-pre-*.log), so step35 NEVER enters it at any B.
1438        // `step35_decode_batch_layers` carries the real geometry: per-layer n_head (64/96),
1439        // partial rope (64 full / 128 SWA, dual base, rope_freqs on FULL only), per-SESSION
1440        // SWA view offsets from each session's own kvl.len, the separate head-wise gate at
1441        // m=B, and the sigmoid-router MoE via the same moe_ffn_il_zq8 the eager path uses.
1442        // MEMRA_STEP35_BATCH=0 = the fail-closed rollback seam. The server caps chunks at
1443        // B=1; on PP-N the B=1 correctness default also refuses the eager numeric class, while
1444        // an unsplit deployment can still use its existing eager B=1 route.
1445        if batch_program == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe {
1446            if !Self::step35_batch_on() {
1447                return Err(
1448                    "step35 batched decode is disabled (MEMRA_STEP35_BATCH=0) — \
1449                            only a non-PP eager B=1 route remains available"
1450                        .into(),
1451                );
1452            }
1453            let n_embd = self.cfg.n_embd as usize;
1454            let eps = self.cfg.rms_eps;
1455            let mut ph_last = std::time::Instant::now();
1456            let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1457            let pos_d = e.htod_i32(&pos_v)?;
1458            let x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
1459            ph_mark(e, 0, &mut ph_last)?;
1460            let x = self.step35_decode_batch_layers(
1461                e,
1462                x,
1463                caches,
1464                &pos_v,
1465                &pos_d,
1466                0,
1467                self.layers.len(),
1468                &mut ph_last,
1469            )?;
1470            let mut hn = e.uninit(b_n * n_embd)?;
1471            e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
1472            let logits = e.matmul(&self.output, &hn, b_n)?;
1473            ph_mark(e, 10, &mut ph_last)?;
1474            return self.decode_batch_epilogue(
1475                e,
1476                caches,
1477                samp,
1478                masks,
1479                lean,
1480                logits,
1481                b_n,
1482                &mut ph_last,
1483                None,
1484            );
1485        }
1486        let n_embd = self.cfg.n_embd as usize;
1487        let eps = self.cfg.rms_eps;
1488
1489        // MEMRA_BATCH_PHASE=1: sync-bounded phase accumulation (diagnostics — see header note).
1490        // Initialized BEFORE the tick-input assembly below so slot 0 covers the HOST side of
1491        // setup (pos_v/ptr-table builds, embed gather) as well as the H2D sync — the audit-fix
1492        // lane's Q6 instrumentation gap (research/audit-fixes2-20260805): the old placement
1493        // started the clock after the assembly, so slot 0 under-reported setup.
1494        let mut ph_last = std::time::Instant::now();
1495
1496        // Per-row rope positions (each sequence at its own depth).
1497        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
1498        let pos_d = e.htod_i32(&pos_v)?;
1499
1500        // Per-step, whole-trunk layer context: state pointer table + arm picks. Under a pp
1501        // split this call is made once PER STAGE with that stage's engine and range instead
1502        // (see `batch_layer_ctx`'s doc for why the table cannot be shared across devices).
1503        let n_layers = self.layers.len();
1504        let ctx = self.batch_layer_ctx(e, caches, 0, n_layers)?;
1505
1506        // Embed all B tokens -> x [B, n_embd] (host gather, one H2D).
1507        let x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
1508        ph_mark(e, 0, &mut ph_last)?;
1509
1510        let x = self.decode_batch_layers(e, x, caches, &ctx, &pos_d, &mut ph_last)?;
1511
1512        // ---- output norm + lm_head at m=B, one D2H ----
1513        let mut hn = e.uninit(b_n * n_embd)?;
1514        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
1515        let logits = e.matmul(&self.output, &hn, b_n)?;
1516        ph_mark(e, 10, &mut ph_last)?;
1517
1518        self.decode_batch_epilogue(
1519            e,
1520            caches,
1521            samp,
1522            masks,
1523            lean,
1524            logits,
1525            b_n,
1526            &mut ph_last,
1527            pending_out,
1528        )
1529    }
1530
1531    /// PP3/PP4 WAVEFRONT DECODE: split one scheduler tick into up to one wave per stage and drive
1532    /// the `(wave, stage)` grid through one persistent host worker per non-head stage. The caller
1533    /// owns the head stage. Explicit boundary messages carry `(wave, slot)` forward and exact
1534    /// post-rx acknowledgements carry slot credit back; every simultaneous cell owns distinct
1535    /// request caches and a distinct stage Engine. The arithmetic inside every cell remains the
1536    /// existing stage-scoped batched program verbatim.
1537    #[allow(clippy::too_many_arguments)]
1538    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
1539    fn decode_step_batch_wavefront<'cache>(
1540        &self,
1541        e: &Engine,
1542        tokens: &[u32],
1543        caches: &mut [&'cache mut Cache],
1544        samp: &[Option<DevSamp>],
1545        masks: &[Option<(&CudaSlice<u32>, usize)>],
1546        lean: bool,
1547        fence: &[usize],
1548    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
1549        let batch = tokens.len();
1550        if batch < 2 || batch != caches.len() {
1551            return Err("PP wavefront requires matching token/cache batches with B>=2".into());
1552        }
1553        if !samp.is_empty() && samp.len() != batch {
1554            return Err("PP wavefront sampling metadata must be empty or match B".into());
1555        }
1556        if !masks.is_empty() && masks.len() != batch {
1557            return Err("PP wavefront grammar masks must be empty or match B".into());
1558        }
1559        let stages = fence.len().saturating_sub(1);
1560        let rt = crate::pp::PpNRt::get(e)?;
1561        crate::pp::pp_wave_eligibility(
1562            stages,
1563            crate::pp::pp2_overlap(),
1564            rt.host_bounce_active(),
1565            rt.repeated_stage_device(),
1566        )
1567        .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
1568        crate::pp::pp_wave_numeric_eligibility(
1569            self.cfg
1570                .hy3
1571                .as_ref()
1572                .is_some_and(|hy3| hy3.weight_only_nvfp4),
1573            Engine::bf16_mmv_on(),
1574        )
1575        .map_err(|reason| -> Box<dyn std::error::Error> { reason.into() })?;
1576        if self.is_gemma4_e4b()
1577            || crate::plan_backend::decode_batch_program(&self.plan)
1578                == crate::plan_backend::DecodeBatchProgram::Gemma
1579        {
1580            return Err(
1581                "PP wavefront has no Gemma batched arm; use the model's qualified eager path"
1582                    .into(),
1583            );
1584        }
1585
1586        let ranges = crate::pp::pp_wave_ranges(batch, stages);
1587        let max_wave = ranges.iter().map(|(lo, hi)| hi - lo).max().unwrap_or(0);
1588        let cap = Self::decode_batch_cap();
1589        let exact16 = max_wave > 8 && max_wave <= 16 && self.decode_batch_exact16_ok();
1590        if max_wave > cap && !exact16 {
1591            return Err(format!(
1592                "PP wavefront B={batch} produces a {max_wave}-row wave above cap {cap} with no exact tier"
1593            )
1594            .into());
1595        }
1596        let step35_batched = crate::plan_backend::decode_batch_program(&self.plan)
1597            == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe;
1598        if step35_batched && !Self::step35_batch_on() {
1599            return Err(
1600                "step35 batched decode is disabled; PP wavefront has no correct fallback trunk"
1601                    .into(),
1602            );
1603        }
1604
1605        if rt.n_stages() != stages {
1606            return Err(format!(
1607                "PpNRt stage count {} != PP wavefront stages {stages}",
1608                rt.n_stages()
1609            )
1610            .into());
1611        }
1612        let caller_stream = e.stream();
1613        let primary_context = crate::pp::PrimaryContextRestore::new(e);
1614        rt.fence_stages_behind(&caller_stream)?;
1615        let n_embd = self.cfg.n_embd as usize;
1616        let slot_capacity = max_wave.saturating_mul(n_embd);
1617        for boundary in 0..stages - 1 {
1618            rt.prepare_overlap_slots(boundary, slot_capacity)?;
1619        }
1620
1621        let _exact_scopes = if exact16 {
1622            let engines: Vec<&Engine> = (0..stages).map(|stage| rt.engine(stage, e)).collect();
1623            engines
1624                .into_iter()
1625                .map(|engine| engine.exact_scope(true))
1626                .collect::<Vec<_>>()
1627        } else {
1628            Vec::new()
1629        };
1630
1631        let mut cache_tail: &mut [&'cache mut Cache] = caches;
1632        let mut waves = Vec::with_capacity(ranges.len());
1633        for &(lo, hi) in &ranges {
1634            let width = hi - lo;
1635            let (wave_caches, tail) = cache_tail.split_at_mut(width);
1636            cache_tail = tail;
1637            waves.push(std::sync::Mutex::new(PpDecodeWave {
1638                row_lo: lo,
1639                tokens: &tokens[lo..hi],
1640                caches: wave_caches,
1641                phase_last: std::time::Instant::now(),
1642                result: None,
1643                committed: false,
1644            }));
1645        }
1646        debug_assert!(cache_tail.is_empty());
1647
1648        let (mut outgoing, mut incoming) = pp_wave_channels(stages - 1);
1649        let walk_result = std::thread::scope(|scope| -> Result<(), Box<dyn std::error::Error>> {
1650            let mut workers = Vec::with_capacity(stages - 1);
1651            for stage in 0..stages - 1 {
1652                let stage_incoming = if stage == 0 {
1653                    None
1654                } else {
1655                    Some(
1656                        incoming[stage - 1]
1657                            .take()
1658                            .expect("PP wave incoming endpoint already moved"),
1659                    )
1660                };
1661                let stage_outgoing = outgoing[stage]
1662                    .take()
1663                    .expect("PP wave outgoing endpoint already moved");
1664                let wave_states = &waves;
1665                workers.push(scope.spawn(move || {
1666                    self.decode_step_batch_wave_worker(
1667                        e,
1668                        rt,
1669                        wave_states,
1670                        stage,
1671                        stage_incoming,
1672                        stage_outgoing,
1673                        fence,
1674                        step35_batched,
1675                    )
1676                }));
1677            }
1678
1679            let head_incoming = incoming[stages - 2]
1680                .take()
1681                .expect("PP wave head incoming endpoint already moved");
1682            let head_result = self.decode_step_batch_wave_head(
1683                e,
1684                rt,
1685                &waves,
1686                head_incoming,
1687                fence,
1688                step35_batched,
1689                samp,
1690                masks,
1691                lean,
1692            );
1693
1694            let mut worker_errors = Vec::new();
1695            let mut worker_panic = None;
1696            for worker in workers {
1697                match worker.join() {
1698                    Ok(Ok(())) => {}
1699                    Ok(Err(error)) => {
1700                        worker_errors.push(error);
1701                    }
1702                    Err(payload) => {
1703                        if worker_panic.is_none() {
1704                            worker_panic = Some(payload);
1705                        }
1706                    }
1707                }
1708            }
1709            if let Some(payload) = worker_panic {
1710                std::panic::resume_unwind(payload);
1711            }
1712            if let Some(error) = worker_errors.iter().find(|error| {
1713                !error.contains("channel closed") && !error.contains("upstream worker failed")
1714            }) {
1715                return Err(error.clone().into());
1716            }
1717            match head_result {
1718                Err(error) => Err(error),
1719                Ok(()) => match worker_errors.into_iter().next() {
1720                    Some(error) => Err(error.into()),
1721                    None => Ok(()),
1722                },
1723            }
1724        });
1725        let publish_result = if walk_result.is_ok() {
1726            Some(rt.publish_to(stages - 1, &caller_stream))
1727        } else {
1728            None
1729        };
1730        let restore_result = primary_context.restore();
1731        walk_result?;
1732        if let Some(result) = publish_result {
1733            result?;
1734        }
1735        restore_result?;
1736        static LOGGED: std::sync::Once = std::sync::Once::new();
1737        LOGGED.call_once(|| {
1738            eprintln!(
1739                "[pp-wave] PP{stages} decode wavefront engaged: waves={} max_wave={} (experimental, MEMRA_PP_WAVE=1)",
1740                ranges.len(),
1741                max_wave,
1742            );
1743        });
1744
1745        let mut completed = Vec::with_capacity(waves.len());
1746        for wave in waves {
1747            let state = wave
1748                .into_inner()
1749                .map_err(|_| "PP wave state lock poisoned")?;
1750            if state.result.is_none() {
1751                return Err("PP wavefront completed without a head-stage result".into());
1752            }
1753            completed.push(state);
1754        }
1755        let mut rows = Vec::with_capacity(batch);
1756        let mut next = Vec::with_capacity(batch);
1757        for state in &mut completed {
1758            let (wave_rows, wave_next) = state.result.take().expect("validated PP wave result");
1759            rows.extend(wave_rows);
1760            next.extend(wave_next);
1761            state.committed = true;
1762        }
1763        crate::pp::record_pp_wave_tick();
1764        Ok((rows, next))
1765    }
1766
1767    #[allow(clippy::too_many_arguments)]
1768    fn decode_step_batch_wave_worker<'slice, 'cache>(
1769        &self,
1770        e: &Engine,
1771        rt: &crate::pp::PpNRt,
1772        waves: &[std::sync::Mutex<PpDecodeWave<'slice, 'cache>>],
1773        stage: usize,
1774        incoming: Option<PpWaveIncoming>,
1775        mut outgoing: PpWaveOutgoing,
1776        fence: &[usize],
1777        step35_batched: bool,
1778    ) -> Result<(), String> {
1779        let result = (|| -> Result<(), String> {
1780            if (stage == 0) != incoming.is_none() {
1781                return Err(format!(
1782                    "PP wave stage {stage} incoming endpoint shape is invalid"
1783                ));
1784            }
1785            for (wave_index, state) in waves.iter().enumerate() {
1786                let transfer = match incoming.as_ref() {
1787                    Some(incoming) => Some(incoming.receive(wave_index)?),
1788                    None => None,
1789                };
1790                let mut state = state
1791                    .lock()
1792                    .map_err(|_| "PP wave state lock poisoned".to_string())?;
1793                self.decode_step_batch_wave_stage(
1794                    e,
1795                    rt,
1796                    &mut state,
1797                    wave_index,
1798                    stage,
1799                    transfer,
1800                    incoming.as_ref(),
1801                    &mut outgoing,
1802                    fence,
1803                    step35_batched,
1804                )
1805                .map_err(|error| error.to_string())?;
1806            }
1807            outgoing.finish()
1808        })();
1809        if let Err(error) = &result {
1810            outgoing.publish_worker_error(error);
1811        }
1812        result
1813    }
1814
1815    #[allow(clippy::too_many_arguments)]
1816    fn decode_step_batch_wave_head<'slice, 'cache>(
1817        &self,
1818        e: &Engine,
1819        rt: &crate::pp::PpNRt,
1820        waves: &[std::sync::Mutex<PpDecodeWave<'slice, 'cache>>],
1821        incoming: PpWaveIncoming,
1822        fence: &[usize],
1823        step35_batched: bool,
1824        samp: &[Option<DevSamp>],
1825        masks: &[Option<(&CudaSlice<u32>, usize)>],
1826        lean: bool,
1827    ) -> Result<(), Box<dyn std::error::Error>> {
1828        for (wave_index, state) in waves.iter().enumerate() {
1829            let transfer = incoming
1830                .receive(wave_index)
1831                .map_err(|error| -> Box<dyn std::error::Error> { error.into() })?;
1832            let mut state = state.lock().map_err(|_| "PP wave state lock poisoned")?;
1833            self.decode_step_batch_wave_final(
1834                e,
1835                rt,
1836                &mut state,
1837                transfer,
1838                &incoming,
1839                fence,
1840                step35_batched,
1841                samp,
1842                masks,
1843                lean,
1844            )?;
1845        }
1846        Ok(())
1847    }
1848
1849    #[allow(clippy::too_many_arguments)]
1850    fn decode_step_batch_wave_stage(
1851        &self,
1852        e: &Engine,
1853        rt: &crate::pp::PpNRt,
1854        wave: &mut PpDecodeWave<'_, '_>,
1855        wave_index: usize,
1856        stage: usize,
1857        transfer: Option<PpWaveTransfer>,
1858        incoming: Option<&PpWaveIncoming>,
1859        outgoing: &mut PpWaveOutgoing,
1860        fence: &[usize],
1861        step35_batched: bool,
1862    ) -> Result<(), Box<dyn std::error::Error>> {
1863        debug_assert!(stage + 1 < fence.len() - 1);
1864        rt.bind_stage(stage)?;
1865        let _stage = rt.enter(stage);
1866        let engine = rt.engine(stage, e);
1867        wave.phase_last = std::time::Instant::now();
1868        let width = wave.tokens.len();
1869        let n_embd = self.cfg.n_embd as usize;
1870        let payload = width * n_embd;
1871        let positions: Vec<i32> = wave.caches.iter().map(|cache| cache.pos as i32).collect();
1872        let positions_d = engine.htod_i32(&positions)?;
1873        let x = if stage == 0 {
1874            if transfer.is_some() || incoming.is_some() {
1875                return Err("PP wave stage 0 received an incoming transfer".into());
1876            }
1877            let x = engine.htod(&self.embd.try_gather(n_embd, wave.tokens)?)?;
1878            ph_mark(engine, 0, &mut wave.phase_last)?;
1879            x
1880        } else {
1881            let transfer = transfer.ok_or("PP wavefront stage has no incoming transfer")?;
1882            let incoming = incoming.ok_or("PP wavefront stage has no incoming endpoint")?;
1883            let x = rt.rx(stage - 1, transfer.slot, payload)?;
1884            incoming
1885                .acknowledge(transfer)
1886                .map_err(|error| -> Box<dyn std::error::Error> { error.into() })?;
1887            x
1888        };
1889        let expected_slot = outgoing
1890            .prepare(wave_index)
1891            .map_err(|error| -> Box<dyn std::error::Error> { error.into() })?;
1892        let _active = crate::pp::enter_pp_wave_cell();
1893        let x = if step35_batched {
1894            self.step35_decode_batch_layers(
1895                engine,
1896                x,
1897                wave.caches,
1898                &positions,
1899                &positions_d,
1900                fence[stage],
1901                fence[stage + 1],
1902                &mut wave.phase_last,
1903            )?
1904        } else {
1905            let ctx = self.batch_layer_ctx(engine, wave.caches, fence[stage], fence[stage + 1])?;
1906            self.decode_batch_layers(
1907                engine,
1908                x,
1909                wave.caches,
1910                &ctx,
1911                &positions_d,
1912                &mut wave.phase_last,
1913            )?
1914        };
1915        let slot = rt.tx_pipelined(stage, &x, payload)?;
1916        outgoing
1917            .publish(wave_index, slot, expected_slot)
1918            .map_err(|error| -> Box<dyn std::error::Error> { error.into() })?;
1919        Ok(())
1920    }
1921
1922    #[allow(clippy::too_many_arguments)]
1923    fn decode_step_batch_wave_final(
1924        &self,
1925        e: &Engine,
1926        rt: &crate::pp::PpNRt,
1927        wave: &mut PpDecodeWave<'_, '_>,
1928        transfer: PpWaveTransfer,
1929        incoming: &PpWaveIncoming,
1930        fence: &[usize],
1931        step35_batched: bool,
1932        samp: &[Option<DevSamp>],
1933        masks: &[Option<(&CudaSlice<u32>, usize)>],
1934        lean: bool,
1935    ) -> Result<(), Box<dyn std::error::Error>> {
1936        let stage = fence.len() - 2;
1937        rt.bind_stage(stage)?;
1938        let _stage = rt.enter(stage);
1939        let engine = rt.engine(stage, e);
1940        wave.phase_last = std::time::Instant::now();
1941        let width = wave.tokens.len();
1942        let n_embd = self.cfg.n_embd as usize;
1943        let payload = width * n_embd;
1944        let positions: Vec<i32> = wave.caches.iter().map(|cache| cache.pos as i32).collect();
1945        let positions_d = engine.htod_i32(&positions)?;
1946        let x = rt.rx(stage - 1, transfer.slot, payload)?;
1947        incoming
1948            .acknowledge(transfer)
1949            .map_err(|error| -> Box<dyn std::error::Error> { error.into() })?;
1950        let _active = crate::pp::enter_pp_wave_cell();
1951        let x = if step35_batched {
1952            self.step35_decode_batch_layers(
1953                engine,
1954                x,
1955                wave.caches,
1956                &positions,
1957                &positions_d,
1958                fence[stage],
1959                fence[stage + 1],
1960                &mut wave.phase_last,
1961            )?
1962        } else {
1963            let ctx = self.batch_layer_ctx(engine, wave.caches, fence[stage], fence[stage + 1])?;
1964            self.decode_batch_layers(
1965                engine,
1966                x,
1967                wave.caches,
1968                &ctx,
1969                &positions_d,
1970                &mut wave.phase_last,
1971            )?
1972        };
1973        let mut normalized = engine.uninit(payload)?;
1974        engine.rms_norm(
1975            &x,
1976            self.output_norm.float_data(),
1977            &mut normalized,
1978            n_embd,
1979            width,
1980            self.cfg.rms_eps,
1981        )?;
1982        let logits = engine.matmul(&self.output, &normalized, width)?;
1983        ph_mark(engine, 10, &mut wave.phase_last)?;
1984        let hi = wave.row_lo + width;
1985        let wave_samp = if samp.is_empty() {
1986            &[][..]
1987        } else {
1988            &samp[wave.row_lo..hi]
1989        };
1990        let wave_masks = if masks.is_empty() {
1991            &[][..]
1992        } else {
1993            &masks[wave.row_lo..hi]
1994        };
1995        wave.result = Some(self.decode_batch_epilogue(
1996            engine,
1997            wave.caches,
1998            wave_samp,
1999            wave_masks,
2000            lean,
2001            logits,
2002            width,
2003            &mut wave.phase_last,
2004            None,
2005        )?);
2006        Ok(())
2007    }
2008
2009    /// DUAL-ACTIVE PP-2 DECODE (increment 0): split one batch into wave A/B and drive
2010    /// stage 0(B) from a scoped host walker while this thread drives stage 1(A). Step's
2011    /// per-layer router readback synchronizes the host, so two CUDA streams issued by one
2012    /// host thread would remain serial; this mirrors the proven prime PP-2 host schedule.
2013    ///
2014    /// This arm is the naked PP-2 default since the 2026-08-11 owner flip (`MEMRA_DUAL_PP`
2015    /// unset = Auto; `0` is the serial rollback seam). It is fail-closed unless the
2016    /// double-slot door is open, prewarms both slots, and uses `tx_pipelined` exclusively.
2017    #[allow(clippy::too_many_arguments)]
2018    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2019    fn decode_step_batch_dual(
2020        &self,
2021        e: &Engine,
2022        tokens: &[u32],
2023        caches: &mut [&mut Cache],
2024        samp: &[Option<DevSamp>],
2025        masks: &[Option<(&CudaSlice<u32>, usize)>],
2026        lean: bool,
2027        fence: &[usize],
2028        mid: usize,
2029    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
2030        let b_n = tokens.len();
2031        assert!(
2032            b_n >= 1 && b_n == caches.len(),
2033            "tokens/caches length mismatch"
2034        );
2035        let Some(expected_mid) = crate::pp::dual_pp_wave_mid(b_n) else {
2036            return self.decode_step_batch_ppn(e, tokens, caches, samp, masks, lean, fence);
2037        };
2038        if mid != expected_mid {
2039            return Err(format!(
2040                "decode_step_batch_dual: worker midpoint {mid} is not the balanced midpoint {expected_mid} for B={b_n}"
2041            ).into());
2042        }
2043        if self.is_gemma4_e4b()
2044            || crate::plan_backend::decode_batch_program(&self.plan)
2045                == crate::plan_backend::DecodeBatchProgram::Gemma
2046        {
2047            return Err(
2048                "decode_step_batch_dual has no gemma4 arm — serve gemma4 on the eager \
2049                        per-session path"
2050                    .into(),
2051            );
2052        }
2053        assert!(
2054            samp.is_empty() || samp.len() == b_n,
2055            "decode_step_batch_dual: samp must be empty or have one entry per row"
2056        );
2057        assert!(
2058            masks.is_empty() || masks.len() == b_n,
2059            "decode_step_batch_dual: masks must be empty or have one entry per row"
2060        );
2061
2062        let cap = Self::decode_batch_cap();
2063        let max_wave = mid.max(b_n - mid);
2064        let exact16 = max_wave > 8 && max_wave <= 16 && self.decode_batch_exact16_ok();
2065        if max_wave > cap && !exact16 {
2066            return Err(format!(
2067                "decode_step_batch_dual: B={b_n} waves {mid}+{} exceed per-wave cap {cap} with no exact tier — refused",
2068                b_n - mid,
2069            ).into());
2070        }
2071        let n_st = fence.len() - 1;
2072        crate::pp::dual_pp_eligibility(
2073            n_st,
2074            crate::pp::pp2_overlap(),
2075            crate::pp::pp_host_bounce_active(),
2076        )
2077        .map_err(|msg| -> Box<dyn std::error::Error> { msg.into() })?;
2078        let rt = crate::pp::PpNRt::get(e)?;
2079        assert_eq!(
2080            rt.n_stages(),
2081            n_st,
2082            "PpNRt stage count {} != fence stages {n_st}",
2083            rt.n_stages()
2084        );
2085        let caller_stream = e.stream();
2086        rt.fence_stages_behind(&caller_stream)?;
2087
2088        let n_embd = self.cfg.n_embd as usize;
2089        let wave_cap = mid.max(b_n - mid) * n_embd;
2090        rt.prepare_overlap_slots(0, wave_cap)?;
2091
2092        // EXACT-16 is a property of either scheduled wave, not the combined live width. Keep
2093        // the scope live across both host walkers and set it on both stage-owned Engines.
2094        let _exact_scopes = if exact16 {
2095            let engines: Vec<&Engine> = (0..n_st).map(|s| rt.engine(s, e)).collect();
2096            engines
2097                .into_iter()
2098                .map(|engine| engine.exact_scope(true))
2099                .collect::<Vec<_>>()
2100        } else {
2101            Vec::new()
2102        };
2103
2104        let step35_batched = crate::plan_backend::decode_batch_program(&self.plan)
2105            == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe;
2106        if step35_batched && !Self::step35_batch_on() {
2107            return Err(
2108                "step35 batched decode is disabled (MEMRA_STEP35_BATCH=0) — \
2109                        dual-active PP-2 decode has no correct fallback trunk"
2110                    .into(),
2111            );
2112        }
2113
2114        let (tokens_a, tokens_b) = tokens.split_at(mid);
2115        let (caches_a, caches_b) = caches.split_at_mut(mid);
2116        let (samp_a, samp_b) = if samp.is_empty() {
2117            (&[][..], &[][..])
2118        } else {
2119            samp.split_at(mid)
2120        };
2121        let (masks_a, masks_b) = if masks.is_empty() {
2122            (&[][..], &[][..])
2123        } else {
2124            masks.split_at(mid)
2125        };
2126
2127        let (slot_a, ph_a, span_a0) = self.decode_step_batch_dual_stage0(
2128            e,
2129            rt,
2130            tokens_a,
2131            caches_a,
2132            fence,
2133            step35_batched,
2134            false,
2135        )?;
2136
2137        static LOGGED: std::sync::Once = std::sync::Once::new();
2138        LOGGED.call_once(|| {
2139            eprintln!("[dual-pp] dual-active PP-2 decode engaged (naked default since 2026-08-11; two waves)");
2140        });
2141
2142        let (out_a, out_b, span_b0, span_b1) = std::thread::scope(
2143            |scope| -> Result<_, Box<dyn std::error::Error>> {
2144                let stage0_b = scope.spawn(move || {
2145                    let staged = self
2146                        .decode_step_batch_dual_stage0(
2147                            e,
2148                            rt,
2149                            tokens_b,
2150                            caches_b,
2151                            fence,
2152                            step35_batched,
2153                            true,
2154                        )
2155                        .map_err(|err| err.to_string())?;
2156                    Ok::<_, String>((staged, caches_b))
2157                });
2158
2159                let out_a = self.decode_step_batch_dual_stage1(
2160                    e,
2161                    rt,
2162                    slot_a,
2163                    caches_a,
2164                    samp_a,
2165                    masks_a,
2166                    lean,
2167                    fence,
2168                    step35_batched,
2169                    ph_a,
2170                    true,
2171                )?;
2172                let ((slot_b, ph_b, span_b0), caches_b) = stage0_b
2173                    .join()
2174                    .map_err(|_| "dual PP stage-0 wave-B host walker panicked")?
2175                    .map_err(|err| -> Box<dyn std::error::Error> { err.into() })?;
2176                if !crate::pp::record_dual_pp_slot_pair(slot_a, slot_b) {
2177                    return Err(format!(
2178                        "decode_step_batch_dual: refused: wave A and B both selected boundary slot {slot_a}"
2179                    ).into());
2180                }
2181                let (out_b, span_b1) = self.decode_step_batch_dual_stage1(
2182                    e,
2183                    rt,
2184                    slot_b,
2185                    caches_b,
2186                    samp_b,
2187                    masks_b,
2188                    lean,
2189                    fence,
2190                    step35_batched,
2191                    ph_b,
2192                    false,
2193                )?;
2194                Ok((out_a, out_b, span_b0, span_b1))
2195            },
2196        )?;
2197
2198        // Wave B is the final producer. One event publishes all last-stage work back to the
2199        // caller after both epilogues, preserving the ordinary PP-N exit law.
2200        rt.publish_to(1, &caller_stream)?;
2201        let (out_a, span_a1) = out_a;
2202        for (stage, span) in [span_a0, span_a1, span_b0, span_b1].into_iter().enumerate() {
2203            if let Some((start, end)) = span {
2204                crate::pp::record_dual_pp_stage_result(stage, start.elapsed_ms(&end));
2205            }
2206        }
2207        let (mut rows, mut next) = out_a;
2208        rows.extend(out_b.0);
2209        next.extend(out_b.1);
2210        Ok((rows, next))
2211    }
2212
2213    #[allow(clippy::too_many_arguments)]
2214    fn decode_step_batch_dual_stage0(
2215        &self,
2216        e: &Engine,
2217        rt: &crate::pp::PpNRt,
2218        tokens: &[u32],
2219        caches: &mut [&mut Cache],
2220        fence: &[usize],
2221        step35_batched: bool,
2222        track_overlap: bool,
2223    ) -> Result<(usize, std::time::Instant, DualPpCudaSpan), Box<dyn std::error::Error>> {
2224        let b_n = tokens.len();
2225        let n_embd = self.cfg.n_embd as usize;
2226        let mut ph_last = std::time::Instant::now();
2227        rt.bind_stage(0)?;
2228        let _st0 = rt.enter(0);
2229        let e0 = rt.engine(0, e);
2230        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
2231        let pos_d = e0.htod_i32(&pos_v)?;
2232        let x = e0.htod(&self.embd.try_gather(n_embd, tokens)?)?;
2233        ph_mark(e0, 0, &mut ph_last)?;
2234        let timing_start = dual_pp_timing_event(e0, "stage0 start event");
2235        let x = {
2236            let _overlap = track_overlap.then(crate::pp::enter_dual_pp_stage);
2237            if step35_batched {
2238                self.step35_decode_batch_layers(
2239                    e0,
2240                    x,
2241                    caches,
2242                    &pos_v,
2243                    &pos_d,
2244                    fence[0],
2245                    fence[1],
2246                    &mut ph_last,
2247                )?
2248            } else {
2249                let ctx = self.batch_layer_ctx(e0, caches, fence[0], fence[1])?;
2250                self.decode_batch_layers(e0, x, caches, &ctx, &pos_d, &mut ph_last)?
2251            }
2252        };
2253        let timing = timing_start.zip(dual_pp_timing_event(e0, "stage0 end event"));
2254        let slot = rt.tx_pipelined(0, &x, b_n * n_embd)?;
2255        Ok((slot, ph_last, timing))
2256    }
2257
2258    #[allow(clippy::too_many_arguments)]
2259    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2260    fn decode_step_batch_dual_stage1(
2261        &self,
2262        e: &Engine,
2263        rt: &crate::pp::PpNRt,
2264        slot: usize,
2265        caches: &mut [&mut Cache],
2266        samp: &[Option<DevSamp>],
2267        masks: &[Option<(&CudaSlice<u32>, usize)>],
2268        lean: bool,
2269        fence: &[usize],
2270        step35_batched: bool,
2271        mut ph_last: std::time::Instant,
2272        track_overlap: bool,
2273    ) -> Result<((Vec<Vec<f32>>, Vec<Option<u32>>), DualPpCudaSpan), Box<dyn std::error::Error>>
2274    {
2275        let b_n = caches.len();
2276        let n_embd = self.cfg.n_embd as usize;
2277        let eps = self.cfg.rms_eps;
2278        rt.bind_stage(1)?;
2279        let _st1 = rt.enter(1);
2280        let e1 = rt.engine(1, e);
2281        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
2282        let pos_d = e1.htod_i32(&pos_v)?;
2283        let x = rt.rx(0, slot, b_n * n_embd)?;
2284        let timing_start = dual_pp_timing_event(e1, "stage1 start event");
2285        let x = {
2286            let _overlap = track_overlap.then(crate::pp::enter_dual_pp_stage);
2287            if step35_batched {
2288                self.step35_decode_batch_layers(
2289                    e1,
2290                    x,
2291                    caches,
2292                    &pos_v,
2293                    &pos_d,
2294                    fence[1],
2295                    fence[2],
2296                    &mut ph_last,
2297                )?
2298            } else {
2299                let ctx = self.batch_layer_ctx(e1, caches, fence[1], fence[2])?;
2300                self.decode_batch_layers(e1, x, caches, &ctx, &pos_d, &mut ph_last)?
2301            }
2302        };
2303        let timing = timing_start.zip(dual_pp_timing_event(e1, "stage1 end event"));
2304        let mut hn = e1.uninit(b_n * n_embd)?;
2305        e1.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
2306        let logits = e1.matmul(&self.output, &hn, b_n)?;
2307        ph_mark(e1, 10, &mut ph_last)?;
2308        Ok((
2309            self.decode_batch_epilogue(
2310                e1,
2311                caches,
2312                samp,
2313                masks,
2314                lean,
2315                logits,
2316                b_n,
2317                &mut ph_last,
2318                None,
2319            )?,
2320            timing,
2321        ))
2322    }
2323
2324    /// THE BATCHED PP-N STEP (pp2-batch increment 2, 2026-08-06): the batched tick split
2325    /// across `fence.len()-1` stages, each stage running ONLY its own layer range through
2326    /// ITS OWN engine and stream, with a `[B, n_embd]` boundary activation between them.
2327    /// The batched twin of `decode_step_h_ppn`, and the #1 item on the PP-2 serving bill —
2328    /// without it a >VRAM SKU (Step-3.7-Flash: 105 GB, fits only across two cards) serves
2329    /// SINGLE-STREAM only, because the batched path was the one loop with no stage split.
2330    ///
2331    /// STRUCTURE (mirrors the eager arm exactly, so the two stay comparable):
2332    ///   stage 0        `rt.enter(0)` -> per-stage pos_d + embed -> range -> `rt.tx`
2333    ///   middle stages  `rt.rx` -> per-stage pos_d -> range -> `rt.tx`
2334    ///   last stage     `rt.rx` -> per-stage pos_d -> range -> output_norm + lm_head ->
2335    ///                  the batched serving epilogue (masks, device sample, lean park)
2336    ///
2337    /// FOUR THINGS ARE PER-STAGE, and each is per-stage for a measured reason:
2338    ///
2339    /// 1. THE ENGINE (`rt.engine(s, e)`). Not just for the remote device: `Engine` owns
2340    ///    lazily-grown stable-pointer scratch pools (`fa_part_pool`, `fa_vf16_scratch`,
2341    ///    `argmax_partials`) that are single-stream-safe BY DESIGN. Two stage streams
2342    ///    through one Engine is the shared-scratch race the pp2 lane hit (2026-08-02
2343    ///    nondeterministic all-logits divergence, 35% flake). `PpNRt::build` already gives
2344    ///    every stage s>0 its own Engine even on the primary device, so honouring
2345    ///    `rt.engine(s, e)` here is what scopes the pools per stage — the batched path
2346    ///    allocates MORE of that scratch than the eager one (fa at m=B), so this is the
2347    ///    load-bearing half of the trap's mitigation, not an inherited nicety.
2348    ///
2349    /// 2. THE POINTER TABLE (`batch_layer_ctx(es, caches, lo, hi)`). See [`BatchLayerCtx`]:
2350    ///    it holds DEVICE ADDRESSES of that range's cache state, uploaded through that
2351    ///    stage's engine. One step-wide table on the primary would put every stage's kernel
2352    ///    arguments in stage-0's HBM — a peer read per pointer fetch, the exact cliff this
2353    ///    whole lane exists to remove.
2354    ///
2355    /// 3. `pos_d` (the M2 pipelining law, learned on the eager arm): each stage uploads its
2356    ///    own copy of the step's per-row positions on ITS stream, so the buffer is
2357    ///    allocated, consumed and freed on one stream. A shared stage-0 `pos_d` freed at fn
2358    ///    return breaks under deferred readback — the free enqueues on stream 0 while later
2359    ///    stages still dereference it.
2360    ///
2361    /// 4. THE HEAD + EPILOGUE run on the LAST stage: `output_norm`/`output` were uploaded
2362    ///    through the last stage's engine by the sharded loader (`hybrid.rs`: `e_head =
2363    ///    layer_engine(e, n_trunk, n_trunk-1)`), and `cache.last_logits_dev` must be
2364    ///    allocated where the logits are.
2365    ///
2366    /// EXACTNESS: PP-N adds ZERO deviation. Each stage runs the SAME kernels on the SAME
2367    /// bytes in the same order — the split only moves where the residual is materialized,
2368    /// and the boundary is a straight f32 copy (dtod same-device / `cudaMemcpyPeerAsync`
2369    /// cross-device, no conversion). So batched PP-N must be BIT-IDENTICAL to single-device
2370    /// batched at the same B, in both placement orders. Gate: `decode-batch-gate --mode
2371    /// pp` (logit-dump, both orders) — the batched analogue of the eager arm's 48 steps x
2372    /// 248,320 f32 logits with zero differing bits.
2373    ///
2374    /// The B=1 fast path is NOT taken here (its condition already excludes an open door):
2375    /// it routes through `decode_layers_eager` whole-trunk on one engine, which is exactly
2376    /// the unsplit walk. B=1 under the door rides this function's B=1 case instead — the
2377    /// same trade the eager arm's own ppn step makes, and the reason the pp2 lane measured
2378    /// B=1 door-open at 0.854x (the lost fusion chain), not a cliff.
2379    #[allow(clippy::too_many_arguments)]
2380    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2381    fn decode_step_batch_ppn(
2382        &self,
2383        e: &Engine,
2384        tokens: &[u32],
2385        caches: &mut [&mut Cache],
2386        samp: &[Option<DevSamp>],
2387        masks: &[Option<(&CudaSlice<u32>, usize)>],
2388        lean: bool,
2389        fence: &[usize],
2390    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
2391        let b_n = tokens.len();
2392        assert!(
2393            b_n >= 1 && b_n == caches.len(),
2394            "tokens/caches length mismatch"
2395        );
2396        // gemma4: same no-arm refusal as the unsplit body (see decode_step_batch), Err not
2397        // assert — a request must never kill the worker process.
2398        if self.is_gemma4_e4b()
2399            || crate::plan_backend::decode_batch_program(&self.plan)
2400                == crate::plan_backend::DecodeBatchProgram::Gemma
2401        {
2402            return Err(
2403                "decode_step_batch_ppn has no gemma4 arm — serve gemma4 on the eager \
2404                        per-session path"
2405                    .into(),
2406            );
2407        }
2408        // Same width policy as the unsplit body — the stage split changes WHERE kernels run,
2409        // never WHICH tier admits the width. Duplicated deliberately rather than hoisted:
2410        // the exact-16 scope must wrap the whole multi-stage walk (`set_verify_exact` is
2411        // per-Engine state read at dispatch on every stage), so it has to be established
2412        // here, and a shared helper returning a guard would have to own `e` plus the flag.
2413        let cap = Self::decode_batch_cap();
2414        let exact16 = b_n > 8 && b_n <= 16 && self.decode_batch_exact16_ok();
2415        assert!(
2416            b_n <= cap || exact16,
2417            "decode_step_batch_ppn: B={b_n} > cap {cap} with no exact tier — refused"
2418        );
2419        let rt = crate::pp::PpNRt::get(e)?;
2420        let n_st = fence.len() - 1;
2421        assert_eq!(
2422            rt.n_stages(),
2423            n_st,
2424            "PpNRt stage count {} != fence stages {n_st}",
2425            rt.n_stages()
2426        );
2427        // #87 REVERSE PUBLICATION (lane/pp2spec-crash): order every stage stream behind
2428        // the caller before this body's first stage allocation can reuse a pool block
2429        // whose queued primary-stream consumer has not read it yet. Anatomy:
2430        // `PpNRt::fence_stages_behind`. (This body dtoh+syncs its own logits, but its
2431        // PP-mode callers interleave with the spec verify's device-resident outputs in
2432        // the same worker, so the entry fence is the uniform law, not an optimization.)
2433        rt.fence_stages_behind(&e.stream())?;
2434        let n_embd = self.cfg.n_embd as usize;
2435        let eps = self.cfg.rms_eps;
2436        let payload = b_n * n_embd;
2437
2438        // EXACT-16 SCOPE, PER STAGE ENGINE: `verify_exact` is per-Engine state (an AtomicBool
2439        // on the Engine the dispatch reads), and each stage runs through a DIFFERENT Engine —
2440        // so setting it on the primary alone would leave stages 1..N-1 dispatching the m>=16
2441        // GEMM/MMQ arms while stage 0 used the exact b16 tier. That is a silent per-stage
2442        // numeric split (the failure this tier exists to prevent), so the flag is set on
2443        // every stage engine and cleared on all of them at scope exit.
2444        let _exact_scopes = if exact16 {
2445            let engines: Vec<&Engine> = (0..n_st).map(|s| rt.engine(s, e)).collect();
2446            engines
2447                .into_iter()
2448                .map(|engine| engine.exact_scope(true))
2449                .collect::<Vec<_>>()
2450        } else {
2451            Vec::new()
2452        };
2453
2454        let mut ph_last = std::time::Instant::now();
2455
2456        // B=1 PER-STAGE FAST PATH (measured 2026-08-06, PRO 6000 pair). The unsplit body's
2457        // b1_fast guard includes `pp_cuts().is_none()`, so opening the pp door dropped every
2458        // solo session off the m=1 FUSION chain (cross-layer add+norm+q8_1, fused SwiGLU,
2459        // lever 1's gate+up dual) and onto the batched m=1 walk. Cost, arm A vs arm C at B=1:
2460        // 208.5 vs 177.3 tok/s = -15.0% — and NOT a split cost, since arm B (stages=2 on ONE
2461        // card) pays the same 177, and the prior lane's `MEMRA_PP_SHARD=0` batched-body B=1
2462        // was 178.5. It was the fusion chain going missing, on the config the Step SKU serves
2463        // solo requests from.
2464        //
2465        // `decode_layers_eager(lo, hi)` is ALREADY range-scoped and is exactly what the eager
2466        // ppn arm (`decode_step_h_ppn`) calls per stage, so B=1 rides the same per-stage
2467        // structure: same engines, same streams, same [1, n_embd] boundary slots, same
2468        // stage-owned caches. Only the trunk kernels differ, and they differ identically to
2469        // how they differ off-door. Exactness is therefore the SAME accepted decode-config FP
2470        // class the unsplit b1_fast lever already carries (strict gate1 PASSes with it on,
2471        // FAILs with it off at maxdiff 1.591e-1) — which is why the pp gate pins
2472        // `set_b1_fast(false)`: with it on, the B=1 reference and the split arm would
2473        // legitimately sit on opposite sides of that gap and the bit-identity arm would
2474        // report a fake stage-split failure.
2475        //
2476        // Step3.5/Step3.7 are an exception (lane/cx-b1fix, 2026-08-10): their B>1 route is
2477        // `step35_decode_batch_layers`, and the live scheduler may move a session from B=1
2478        // to B>1. The eager/fused class and that batched class produce different greedy bytes,
2479        // so selecting the eager arm at B=1 made output depend on load history. Keep one
2480        // numeric class for this model family: Step35 always takes its stage-scoped batched
2481        // trunk at every width. The live transition gate in step35-b2-geometry-gate pins it.
2482        // Qwen35-MoE is the second exception (lane/cx-q35bug, 2026-08-12): on the Q35
2483        // sellgate workload the eager-B1 -> batched-B2 transition changed emitted token ids and
2484        // selected EOS at tokens 15/17/25. Keep that family on this generic batched trunk at B=1
2485        // too; dense Qwen35 retains the measured eager fast path.
2486        let b1_stage_fast = b_n == 1
2487            && Self::b1_fast_on()
2488            && self.b1_fast_plan_eligible()
2489            && !self.is_gemma4_e4b()
2490            && crate::plan_backend::decode_batch_program(&self.plan)
2491                == crate::plan_backend::DecodeBatchProgram::Generic
2492            && !self
2493                .plan
2494                .trunk_operations()
2495                .contains(&memra_gguf::model_plan::OperationKind::SwiGluOaiActivation)
2496            && !e.verify_exact_on();
2497        // step35 (lane/step35-batched-decode, 2026-08-08): B>1 rides its OWN stage-scoped
2498        // batched walk (`step35_decode_batch_layers`) — the generic `decode_batch_layers`
2499        // remains OFF-LIMITS for this arch at every B (its uniform geometry produced the
2500        // b2ab HTTP-200 garbage: research/step-sku-20260807/raw/b2ab-pre-*.log). Since
2501        // lane/cx-b1fix, B=1 also takes this walk: a Step35 PP-N session must not change
2502        // numeric class when live decode width changes. The refusal below guards the
2503        // rollback residue; under PP-N, disabling the only correct trunk makes Step35
2504        // requests fail closed instead of falling back to the eager class.
2505        let step35_batched = crate::plan_backend::decode_batch_program(&self.plan)
2506            == crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe;
2507        if step35_batched && !Self::step35_batch_on() {
2508            return Err(
2509                "step35 batched decode is disabled (MEMRA_STEP35_BATCH=0) — \
2510                        PP-N Step35 decode is unavailable because eager B=1 is a different \
2511                        numeric class"
2512                    .into(),
2513            );
2514        }
2515        // Hoisted: `caches[0].pos` as a value argument alongside `caches[0]` as `&mut` in one
2516        // call is a borrow conflict; `pos` is Copy and the epilogue is what advances it.
2517        let pos0 = if b1_stage_fast { caches[0].pos } else { 0 };
2518
2519        // ---- STAGE 0: embed (the table lives with stage 0) + layers [0, fence[1]) + TX ----
2520        let mut slot = {
2521            let _st0 = rt.enter(0);
2522            let e0 = rt.engine(0, e);
2523            let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
2524            let pos_d = e0.htod_i32(&pos_v)?;
2525            let x = e0.htod(&self.embd.try_gather(n_embd, tokens)?)?;
2526            ph_mark(e0, 0, &mut ph_last)?;
2527            let x = if b1_stage_fast {
2528                self.decode_layers_eager(e0, x, fence[0], fence[1], &pos_d, pos0, caches[0])?
2529            } else if step35_batched {
2530                self.step35_decode_batch_layers(
2531                    e0,
2532                    x,
2533                    caches,
2534                    &pos_v,
2535                    &pos_d,
2536                    fence[0],
2537                    fence[1],
2538                    &mut ph_last,
2539                )?
2540            } else {
2541                let ctx = self.batch_layer_ctx(e0, caches, fence[0], fence[1])?;
2542                self.decode_batch_layers(e0, x, caches, &ctx, &pos_d, &mut ph_last)?
2543            };
2544            rt.tx(0, &x, payload)?
2545            // x + pos_d + ctx.ptr_table drop here: freed stream-ordered on stage-0's stream.
2546        };
2547
2548        // ---- MIDDLE STAGES: RX boundary s-1 -> range -> TX boundary s ----
2549        for s in 1..n_st - 1 {
2550            let _st = rt.enter(s);
2551            let es = rt.engine(s, e);
2552            let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
2553            let pos_d = es.htod_i32(&pos_v)?;
2554            let x = rt.rx(s - 1, slot, payload)?;
2555            let x = if b1_stage_fast {
2556                self.decode_layers_eager(es, x, fence[s], fence[s + 1], &pos_d, pos0, caches[0])?
2557            } else if step35_batched {
2558                self.step35_decode_batch_layers(
2559                    es,
2560                    x,
2561                    caches,
2562                    &pos_v,
2563                    &pos_d,
2564                    fence[s],
2565                    fence[s + 1],
2566                    &mut ph_last,
2567                )?
2568            } else {
2569                let ctx = self.batch_layer_ctx(es, caches, fence[s], fence[s + 1])?;
2570                self.decode_batch_layers(es, x, caches, &ctx, &pos_d, &mut ph_last)?
2571            };
2572            slot = rt.tx(s, &x, payload)?;
2573        }
2574
2575        // ---- LAST STAGE: RX + final range + head + the batched serving epilogue ----
2576        let _stl = rt.enter(n_st - 1);
2577        let el = rt.engine(n_st - 1, e);
2578        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
2579        let pos_d = el.htod_i32(&pos_v)?;
2580        let x = rt.rx(n_st - 2, slot, payload)?;
2581        let x = if b1_stage_fast {
2582            self.decode_layers_eager(el, x, fence[n_st - 1], fence[n_st], &pos_d, pos0, caches[0])?
2583        } else if step35_batched {
2584            self.step35_decode_batch_layers(
2585                el,
2586                x,
2587                caches,
2588                &pos_v,
2589                &pos_d,
2590                fence[n_st - 1],
2591                fence[n_st],
2592                &mut ph_last,
2593            )?
2594        } else {
2595            let ctx = self.batch_layer_ctx(el, caches, fence[n_st - 1], fence[n_st])?;
2596            self.decode_batch_layers(el, x, caches, &ctx, &pos_d, &mut ph_last)?
2597        };
2598
2599        let mut hn = el.uninit(payload)?;
2600        el.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
2601        let logits = el.matmul(&self.output, &hn, b_n)?;
2602        ph_mark(el, 10, &mut ph_last)?;
2603
2604        self.decode_batch_epilogue(
2605            el,
2606            caches,
2607            samp,
2608            masks,
2609            lean,
2610            logits,
2611            b_n,
2612            &mut ph_last,
2613            None,
2614        )
2615    }
2616
2617    /// The mHC (HyperConnections) batched decode arm (lane/glm53-batched-decode,
2618    /// 2026-08-28). DEFAULT ON since 2026-08-31, on the hbatch-battery box receipts
2619    /// (research/glm53-flash-bringup-20260827/hbatch-battery-20260831/): interleaved x3
2620    /// ladder on the 3-card serving shape — ON wins every rung c>=2 (aggregate 1.095x at
2621    /// c=2 up to 1.214x at c=12, plateau ~1.20x from c=8), B=1 cost -0.30%, TTFT under
2622    /// load ON <= OFF at every rung, 36/36 concurrent tapes byte-identical to solo (incl.
2623    /// ON-solo == OFF-solo), admission clean to c=20, loop-law 0/448. `MEMRA_HYPER_BATCH=0`
2624    /// is the rollback seam (eager per-session decode). Any OTHER value REFUSES LOUD at
2625    /// first use — a mis-typed serving switch must not silently pick a path.
2626    pub fn hyper_batch_on() -> bool {
2627        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2628        *ON.get_or_init(|| match std::env::var("MEMRA_HYPER_BATCH").as_deref() {
2629            Ok("0") => false,
2630            Err(_) | Ok("1") => true,
2631            Ok(v) => panic!(
2632                "MEMRA_HYPER_BATCH={v:?} is not a recognized value (want unset/0 = eager \
2633                 per-session decode, 1 = batched mHC decode chunks) — refusing to guess a \
2634                 serving path"
2635            ),
2636        })
2637    }
2638
2639    /// The mHC batched-decode width cap, DERIVED rather than inherited (owner challenge,
2640    /// 2026-08-28: "why 8?"). The audit of every term that grows with B found exactly ONE
2641    /// numeric-class knee, and it is not memory, not the boundary payload, and not the
2642    /// per-session mixer loop:
2643    ///
2644    ///   * per-session mixers (KDA + MLA/kpool): step latency grows ~linearly in B on that
2645    ///     segment — a throughput term, no correctness wall at any width; per-session state
2646    ///     is ~104 MB MLA latent at 8k ctx + trivial KDA, so memory does not bind either.
2647    ///   * hc glue: block-per-token kernels (grid.y chunked at 65535 — B=64 is 256 rows on
2648    ///     `hc_post`), and the hc-mix GEMM runs per-row m=1 by construction (`pre_exact`).
2649    ///   * lm_head via `matmul_decode_exact`: per-row exact at every m (float per-token
2650    ///     m=1; quant b-tier to 16, grid.y=m mmvq above — re-reads, not rounding).
2651    ///   * MoE router (`router_gemv`): m-invariant at every t under defaults; sigmoid
2652    ///     top-k and routed-expert execution are per-token programs at any t.
2653    ///   * MoE SHARED EXPERT — THE BINDER: `hybrid_forward.rs` shexp trio,
2654    ///     `verify_t = t > 1 && t < PRIME_MIN_T`. At t >= 16 gate/up/down cross from
2655    ///     `matmul_decode_exact` onto the plain prefill matmul (cuBLASLt n-dependent for
2656    ///     float; the m>16 MMQ/GEMM block-scale class for quant) and per-row bit-identity
2657    ///     vs the isolated t=1 chain breaks — measured, not argued: the gate's knee probe
2658    ///     at B=16 mismatches from the first tick (`31-KNEE-b16-forced.log`), B=15 is green.
2659    ///
2660    /// So the exact tier is `1..=PRIME_MIN_T-1` = 15. Widening to 32/64 needs a
2661    /// decode-exact shexp arm for t >= 16 — which must NOT be flipped inside the shared
2662    /// `!prefill` branch, because step35's MoESD target forward (t up to 256) rides the
2663    /// same branch and its banked spec receipts pin the current bytes. Named follow-up.
2664    /// `MEMRA_DECODE_BATCH_CAP` narrows only, never widens past the knee.
2665    pub fn hyper_batch_cap() -> usize {
2666        let knee = crate::hybrid_forward::PRIME_MIN_T - 1;
2667        std::env::var("MEMRA_DECODE_BATCH_CAP")
2668            .ok()
2669            .and_then(|v| v.parse::<usize>().ok())
2670            .map(|c| c.clamp(1, knee))
2671            .unwrap_or(knee)
2672    }
2673
2674    /// THE mHC BATCHED DECODE STEP (lane/glm53-batched-decode, 2026-08-28): B sessions,
2675    /// one walk over the `[B, streams, n_embd]` stream state. This is the production
2676    /// blocker this lane lifts — with every batched entry refusing the hc residual,
2677    /// GLM-5.3-Flash served SINGLE-STREAM ONLY at any `MEMRA_MAX_SESSIONS`.
2678    ///
2679    /// The trunk is `hyper_batch_range_decode` (hybrid_forward.rs — see its doc for the
2680    /// batched/per-session/decode-exact shape law); the exit is `hyper::collapse` +
2681    /// output_norm + a DECODE-EXACT lm_head at m=B (each row's head program is the m=1
2682    /// program its solo step runs — `matmul_decode_exact`); the tail is the SAME
2683    /// `decode_batch_epilogue` every other batched arm serves (masks, device sampling,
2684    /// lean park, pos bump), so the serving contract is shared rather than duplicated.
2685    ///
2686    /// CONCURRENCY SHAPES: sessions may sit at DIFFERENT positions with different KDA
2687    /// recurrent states and different kpool index planes — each row carries its own
2688    /// single-position buffer and its own cache, which is what the gate's staggered-depth
2689    /// arm pins. One token per session per tick (pure decode); there is no mixed
2690    /// prefill/decode shape at this entry by construction. Width: B <= 8, the per-row
2691    /// exactness tier — there is NO exact16 tier here (`decode_batch_exact16_ok` refuses
2692    /// the Mla/Kda mixers), and the MoE router's fixed per-row program is only the decode
2693    /// arm below PRIME_MIN_T; wider concurrency is the scheduler's job to chunk.
2694    ///
2695    /// EXACTNESS BAR AND GATE: row b of a B-row tick is BIT-IDENTICAL (full logits, every
2696    /// step) to session b decoding alone through `decode_step_hyper`, including B=1 — one
2697    /// numeric class at every live width, the step35/Q35 class-crossing law. Gate:
2698    /// `glm5-hyper-batch-gate`, red-armed with a swapped-row and a wrong-cache-slot
2699    /// mutation (cross-session contamination is the silent-corruption failure mode).
2700    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2701    pub(crate) fn decode_step_batch_hyper(
2702        &self,
2703        e: &Engine,
2704        tokens: &[u32],
2705        caches: &mut [&mut Cache],
2706        samp: &[Option<DevSamp>],
2707        masks: &[Option<(&CudaSlice<u32>, usize)>],
2708        lean: bool,
2709    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
2710        let topology = *self
2711            .hyper
2712            .as_ref()
2713            .ok_or("decode_step_batch_hyper on a model with no HyperConnections topology")?;
2714        if !Self::hyper_batch_on() {
2715            return Err(
2716                "mHC batched decode is disabled (MEMRA_HYPER_BATCH unset/0, the fail-closed \
2717                 default until serving-box receipts land) — serve hyper-connection sessions \
2718                 on the eager per-session path, or set MEMRA_HYPER_BATCH=1"
2719                    .into(),
2720            );
2721        }
2722        // THE DOOR IS NOT ON THIS WALK, AND A SILENT NO-OP IS THE SAME FAILURE CLASS AS A
2723        // VACUOUS GATE. Take 13's serving A/B (2026-09-03) passed `MEMRA_GLM5_DECODE_GRAPH=1`
2724        // through the serve script and got SIX boots with zero `[glm5-decode-graph]` lines of any
2725        // kind, refusals included, while the same binary's gate harness armed the door fine
2726        // (`replays=564`). The missing conjunct is the walk: `MEMRA_GLM5_DECODE_GRAPH` is wired
2727        // into `hybrid_forward::hyper_range_decode`, the per-session SERIAL hc walk that
2728        // `decode_step_hyper` / `decode_step_hyper_ppn` run, and serving with
2729        // `MEMRA_HYPER_BATCH=1` routes every session through THIS batched walk instead, including
2730        // B=1. So the A/B measured the door's absence and read flat, which is the correct number
2731        // for the wrong question.
2732        //
2733        // Extending capture to the batched walk itself is still a separate lane (its trunk is
2734        // `decode_batch_layers`, with its own per-row geometry). Until then the honest behaviour
2735        // is to say so, once, rather than let a serving log's silence be read as a refusal.
2736        //
2737        // `MEMRA_HYPER_BATCH_SOLO=1` CHANGES WHAT IS HONEST HERE, so this line is now keyed on
2738        // it. At B=1 that door delegates to `hyper_range_decode`, which is exactly the walk the
2739        // graph door is wired into, so the door DOES engage and printing "NOT ON THIS PATH"
2740        // would contradict the `[glm5-decode-graph] engaged` lines a few entries below it in the
2741        // same log. Measured on the 2x B200 pair 2026-09-03: with the solo door armed, serving
2742        // printed `engaged dev=0 stage=[0, 24) runs=6 captured_layers=18` and `engaged dev=1
2743        // stage=[24, 45) runs=6 captured_layers=16` — the first time this door has engaged in
2744        // serving. A stale warning next to a working door is the same failure class the warning
2745        // was written to prevent.
2746        if crate::glm5_decode_graph_on() {
2747            static SAID: std::sync::Once = std::sync::Once::new();
2748            SAID.call_once(|| {
2749                if crate::hyper_batch_solo_on() {
2750                    eprintln!(
2751                        "[glm5-decode-graph] reachable via MEMRA_HYPER_BATCH_SOLO=1: this session \
2752                         enters the BATCHED hc walk, which delegates to the serial walk \
2753                         (hyper_range_decode) at B=1, and that is the walk the door is wired \
2754                         into. Expect [glm5-decode-graph] engaged/eager lines. At B>1 the batched \
2755                         trunk runs and the door is still not on that path."
2756                    );
2757                } else {
2758                    eprintln!(
2759                        "[glm5-decode-graph] NOT ON THIS PATH: MEMRA_GLM5_DECODE_GRAPH is on (the \
2760                         default since 2026-09-04; =0 disarms it) but this session decodes \
2761                         through the BATCHED hc walk (MEMRA_HYPER_BATCH=1), and \
2762                         the door is wired into the per-session serial walk (hyper_range_decode) \
2763                         only. The door will not engage, and will not refuse either, for as long \
2764                         as the batched walk is in use. Set MEMRA_HYPER_BATCH_SOLO=1 to reach it \
2765                         at B=1, unset MEMRA_HYPER_BATCH to price the serial walk, or read this \
2766                         line as the reason a serving log carries no [glm5-decode-graph] lines."
2767                    );
2768                }
2769            });
2770        }
2771        let b_n = tokens.len();
2772        if b_n == 0 || b_n != caches.len() {
2773            return Err("decode_step_batch_hyper: tokens/caches length mismatch".into());
2774        }
2775        // Width: Err, never assert — a request must not kill the worker (the gemma4
2776        // process-FATAL lesson). The cap is DERIVED, not inherited — see `hyper_batch_cap`.
2777        let cap = Self::hyper_batch_cap();
2778        if b_n > cap {
2779            return Err(format!(
2780                "decode_step_batch_hyper: B={b_n} > cap {cap} — at t >= PRIME_MIN_T (16) \
2781                 the MoE shared-expert trio crosses from matmul_decode_exact onto the \
2782                 prefill matmul class (cuBLASLt n-dependent / m>16 MMQ-GEMM), so per-row \
2783                 bit-identity vs isolated decode breaks at exactly B=16 (gate knee probe \
2784                 31-KNEE-b16-forced). Every other term is width-safe; widening needs a \
2785                 decode-exact shexp arm for t>=16 (named follow-up — the shared !prefill \
2786                 branch also carries step35 MoESD bytes and must not be flipped). Chunk \
2787                 wider concurrency into <={cap} groups"
2788            )
2789            .into());
2790        }
2791        let n_embd = self.cfg.n_embd as usize;
2792        let eps = self.cfg.rms_eps;
2793
2794        // M2 ppN door — the batched hc walk owns its own stage split, exactly as the
2795        // serial hc walks do (forward_hyper's note). Loud refusal on an unqualified
2796        // pipeline rewrite, never a single-engine walk over stage-sharded weights.
2797        if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
2798            if !self.rewrite_allowed(memra_gguf::execution_manifest::RewriteSurface::Pipeline) {
2799                return Err("pipeline rewrite is not qualified for this ModelPlan".into());
2800            }
2801            return self.decode_step_batch_hyper_ppn(
2802                e, tokens, caches, samp, masks, lean, &topology, &fence,
2803            );
2804        }
2805
2806        let mut ph_last = std::time::Instant::now();
2807        let pos_rows = Self::hyper_batch_pos_rows(e, caches)?;
2808        let embedded = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
2809        let mut x = crate::hyper::expand(e, &topology, &embedded, b_n, n_embd)?;
2810        ph_mark(e, 0, &mut ph_last)?;
2811        x = self.hyper_batch_range_decode(
2812            e,
2813            &topology,
2814            x,
2815            0,
2816            self.layers.len(),
2817            &pos_rows,
2818            caches,
2819        )?;
2820        let logits = self.hyper_batch_head_logits(e, &topology, &x, b_n, n_embd, eps)?;
2821        ph_mark(e, 10, &mut ph_last)?;
2822        self.decode_batch_epilogue(
2823            e,
2824            caches,
2825            samp,
2826            masks,
2827            lean,
2828            logits,
2829            b_n,
2830            &mut ph_last,
2831            None,
2832        )
2833    }
2834
2835    /// Per-row single-position device buffers, uploaded through THIS engine (under a pp
2836    /// split, the stage's engine — the per-stage pos_d law). The mixers take a t=1 `pos_d`
2837    /// exactly as their solo step does, so each session's row is a one-element buffer, not
2838    /// a shared [B] table.
2839    fn hyper_batch_pos_rows(
2840        e: &Engine,
2841        caches: &[&mut Cache],
2842    ) -> Result<Vec<CudaSlice<i32>>, Box<dyn std::error::Error>> {
2843        caches.iter().map(|c| e.htod_i32(&[c.pos as i32])).collect()
2844    }
2845
2846    /// The batched hc trunk exit: mean/gated collapse + output_norm + DECODE-EXACT lm_head.
2847    /// `matmul_decode_exact` at m=B runs each row through the m=1 head program the serial
2848    /// `hyper_decode_tail` runs (float: per-token m=1 cuBLASLt; quant: the per-(token,row)
2849    /// bit-exact batched mmvq tier), so the head cannot be the arm that breaks per-row
2850    /// identity. Returns device logits `[B, n_vocab]` for the shared epilogue.
2851    fn hyper_batch_head_logits(
2852        &self,
2853        e: &Engine,
2854        topology: &crate::hyper::HyperTopology,
2855        x: &CudaSlice<f32>,
2856        b_n: usize,
2857        n_embd: usize,
2858        eps: f32,
2859    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2860        let collapsed =
2861            crate::hyper::collapse(e, topology, self.hyper_head.as_ref(), x, b_n, n_embd)?;
2862        let mut hn = e.uninit(b_n * n_embd)?;
2863        e.rms_norm(
2864            &collapsed,
2865            self.output_norm.float_data(),
2866            &mut hn,
2867            n_embd,
2868            b_n,
2869            eps,
2870        )?;
2871        e.matmul_decode_exact(&self.output, &hn, b_n)
2872    }
2873
2874    /// ppN twin of `decode_step_batch_hyper`: the batched hc tick as N stage subgraphs,
2875    /// mirroring `decode_step_hyper_ppn` (per-stage engine, per-stage pos uploads, a
2876    /// `[B, streams, n_embd]` boundary payload) and `decode_step_batch_ppn` (the #87 entry
2877    /// fence, head + epilogue on the LAST stage's engine, where the loader put the head and
2878    /// where `cache.last_logits_dev` must live). No exact16 scope (no exact16 tier here)
2879    /// and no B=1 fast path (one numeric class at every width).
2880    #[allow(clippy::too_many_arguments)]
2881    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
2882    fn decode_step_batch_hyper_ppn(
2883        &self,
2884        e: &Engine,
2885        tokens: &[u32],
2886        caches: &mut [&mut Cache],
2887        samp: &[Option<DevSamp>],
2888        masks: &[Option<(&CudaSlice<u32>, usize)>],
2889        lean: bool,
2890        topology: &crate::hyper::HyperTopology,
2891        fence: &[usize],
2892    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
2893        let b_n = tokens.len();
2894        let n_embd = self.cfg.n_embd as usize;
2895        let eps = self.cfg.rms_eps;
2896        let payload = b_n * topology.streams * n_embd;
2897        let mut ph_last = std::time::Instant::now();
2898
2899        // The same-stream seam (MEMRA_PP_STREAMS=0 also disables the sharded loader, so
2900        // nothing is remote): one engine, boundary copies between ranges — the shape the
2901        // serial hc ppn walk uses for this knob.
2902        if crate::pp::pp2_streams_off() {
2903            let pos_rows = Self::hyper_batch_pos_rows(e, caches)?;
2904            let embedded = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
2905            let mut x = crate::hyper::expand(e, topology, &embedded, b_n, n_embd)?;
2906            ph_mark(e, 0, &mut ph_last)?;
2907            x = self
2908                .hyper_batch_range_decode(e, topology, x, fence[0], fence[1], &pos_rows, caches)?;
2909            for s in 1..fence.len() - 1 {
2910                let boundary_tx = e.clone_dtod(&x)?;
2911                let boundary_rx = e.clone_dtod(&boundary_tx)?;
2912                x = self.hyper_batch_range_decode(
2913                    e,
2914                    topology,
2915                    boundary_rx,
2916                    fence[s],
2917                    fence[s + 1],
2918                    &pos_rows,
2919                    caches,
2920                )?;
2921            }
2922            let logits = self.hyper_batch_head_logits(e, topology, &x, b_n, n_embd, eps)?;
2923            ph_mark(e, 10, &mut ph_last)?;
2924            return self.decode_batch_epilogue(
2925                e,
2926                caches,
2927                samp,
2928                masks,
2929                lean,
2930                logits,
2931                b_n,
2932                &mut ph_last,
2933                None,
2934            );
2935        }
2936
2937        let rt = crate::pp::PpNRt::get(e)?;
2938        let n_st = fence.len() - 1;
2939        assert_eq!(
2940            rt.n_stages(),
2941            n_st,
2942            "PpNRt stage count {} != fence stages {n_st}",
2943            rt.n_stages()
2944        );
2945        // #87 reverse publication (see decode_step_batch_ppn): order every stage stream
2946        // behind the caller before this body's first stage allocation.
2947        rt.fence_stages_behind(&e.stream())?;
2948
2949        // ---- STAGE 0: embed + expand (no weights) + layers [0, fence[1]) + TX ----
2950        let mut slot = {
2951            let _st0 = rt.enter(0);
2952            let e0 = rt.engine(0, e);
2953            let pos_rows = Self::hyper_batch_pos_rows(e0, caches)?;
2954            let embedded = e0.htod(&self.embd.try_gather(n_embd, tokens)?)?;
2955            let x = crate::hyper::expand(e0, topology, &embedded, b_n, n_embd)?;
2956            ph_mark(e0, 0, &mut ph_last)?;
2957            let x = self
2958                .hyper_batch_range_decode(e0, topology, x, fence[0], fence[1], &pos_rows, caches)?;
2959            rt.tx(0, &x, payload)?
2960        };
2961
2962        // ---- MIDDLE STAGES: RX -> range -> TX ----
2963        for s in 1..n_st - 1 {
2964            let _st = rt.enter(s);
2965            let es = rt.engine(s, e);
2966            let pos_rows = Self::hyper_batch_pos_rows(es, caches)?;
2967            let x = rt.rx(s - 1, slot, payload)?;
2968            let x = self.hyper_batch_range_decode(
2969                es,
2970                topology,
2971                x,
2972                fence[s],
2973                fence[s + 1],
2974                &pos_rows,
2975                caches,
2976            )?;
2977            slot = rt.tx(s, &x, payload)?;
2978        }
2979
2980        // ---- LAST STAGE: RX + final range + collapse/head + the shared epilogue ----
2981        let _stl = rt.enter(n_st - 1);
2982        let el = rt.engine(n_st - 1, e);
2983        let pos_rows = Self::hyper_batch_pos_rows(el, caches)?;
2984        let x = rt.rx(n_st - 2, slot, payload)?;
2985        let x = self.hyper_batch_range_decode(
2986            el,
2987            topology,
2988            x,
2989            fence[n_st - 1],
2990            fence[n_st],
2991            &pos_rows,
2992            caches,
2993        )?;
2994        let logits = self.hyper_batch_head_logits(el, topology, &x, b_n, n_embd, eps)?;
2995        ph_mark(el, 10, &mut ph_last)?;
2996        self.decode_batch_epilogue(
2997            el,
2998            caches,
2999            samp,
3000            masks,
3001            lean,
3002            logits,
3003            b_n,
3004            &mut ph_last,
3005            None,
3006        )
3007    }
3008
3009    /// Build the per-step layer context for layers `[lo, hi)`: the device state-pointer
3010    /// table plus the step's arm picks. See [`BatchLayerCtx`] for why this is RANGE-scoped
3011    /// (the table holds device addresses and must be uploaded through the engine whose
3012    /// device runs those layers).
3013    ///
3014    /// Table layout is unchanged from the whole-trunk version — `lin_base`/`attn_base` are
3015    /// still indexed by ABSOLUTE layer id, so `decode_batch_layers`' body indexes them
3016    /// exactly as the old inline loop did. Only layers in `[lo, hi)` contribute entries; the
3017    /// rest stay `None`, which is a loud `expect` if a range ever reads outside its own.
3018    pub(crate) fn batch_layer_ctx(
3019        &self,
3020        e: &Engine,
3021        caches: &[&mut Cache],
3022        lo: usize,
3023        hi: usize,
3024    ) -> Result<BatchLayerCtx, Box<dyn std::error::Error>> {
3025        let cfg = &self.cfg;
3026        let head_dim = cfg.head_dim_k as usize;
3027        // Per-step STATE POINTER TABLE (one H2D): for every linear layer, [conv x B]
3028        // [ssm_in x B][ssm_out x B] device addresses. The batched state kernels read their
3029        // sequence's pointer from these arrays — states stay per-cache (no pooling refactor),
3030        // yet conv/prep/scan collapse from 3xB launches per layer to 3. Rebuilt every step
3031        // because the ssm ping-pong swaps pointers host-side after each scan.
3032        // INCREMENT 2 (2026-08-01): the SAME table now also carries, for every FULL-attn
3033        // layer, [k0,v0,k1,v1,...] cache base addresses — the z-batched seqs append and
3034        // seqs fa_decode kernels read their sequence's cache through it (the MoE
3035        // expert-table pattern), collapsing 2xB launches per attn layer to 2.
3036        let mut lin_base: Vec<Option<usize>> = vec![None; self.layers.len()];
3037        let mut attn_base: Vec<Option<usize>> = vec![None; self.layers.len()];
3038        let mut ptrs: Vec<u64> = Vec::new();
3039        {
3040            use cudarc::driver::DevicePtr;
3041            let s = &e.gpu.stream();
3042            for il in lo..hi {
3043                match &self.layers[il].mixer {
3044                    Mixer::Linear(_) => {
3045                        lin_base[il] = Some(ptrs.len());
3046                        for c in caches.iter() {
3047                            let rl = c.recur[il].as_ref().unwrap();
3048                            let (p, _g) = rl.conv_state.device_ptr(s);
3049                            ptrs.push(p);
3050                        }
3051                        for c in caches.iter() {
3052                            let rl = c.recur[il].as_ref().unwrap();
3053                            let (p, _g) = rl.ssm_state.device_ptr(s);
3054                            ptrs.push(p);
3055                        }
3056                        for c in caches.iter() {
3057                            let rl = c.recur[il].as_ref().unwrap();
3058                            let (p, _g) = rl.ssm_state_alt.device_ptr(s);
3059                            ptrs.push(p);
3060                        }
3061                    }
3062                    Mixer::Full(_) => {
3063                        attn_base[il] = Some(ptrs.len());
3064                        for c in caches.iter() {
3065                            let kvl = c.kv[il].as_ref().unwrap();
3066                            let (pk, _g) = kvl.k.device_ptr(s);
3067                            let (pv, _g2) = kvl.v.device_ptr(s);
3068                            ptrs.push(pk);
3069                            ptrs.push(pv);
3070                        }
3071                    }
3072                    Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("batched PP decode"),
3073                    Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("batched decode layer"),
3074                }
3075            }
3076        }
3077        let ptr_table = if ptrs.is_empty() {
3078            None
3079        } else {
3080            Some(e.htod_u64(&ptrs)?)
3081        };
3082
3083        // INCREMENT 2 arm picks (per STEP — t_kv is layer-invariant within a tick):
3084        // - seqs APPEND: format-only condition (per-row program is t_kv-independent);
3085        //   default flash module only (fp8-KV rides the per-seq g-module path).
3086        // - seqs FA: every row must take the v4 eager arm at ITS OWN t_kv AND all rows
3087        //   must share ONE fa_split_keys rung (the rows-twins' straddle law) — a rung
3088        //   crossing inside the batch keeps the per-seq loop for that step, so each
3089        //   sequence always executes the exact program its isolated run would.
3090        //
3091        // The picks are t_kv-driven, and t_kv is layer-INVARIANT within a step, so every
3092        // stage of a pp split independently computes the SAME arms from the same `caches`
3093        // — a stage cannot silently take a different program than its unsplit self.
3094        let t_kvs: Vec<usize> = caches.iter().map(|c| c.pos + 1).collect();
3095        let t_kv_max = *t_kvs.iter().max().unwrap();
3096        let sp0 = crate::fa_split_keys(t_kvs[0], cfg.n_head_kv as usize);
3097        let seqs_fa = t_kvs.iter().all(|&t| crate::fa_seqs_eligible(t, head_dim))
3098            && t_kvs
3099                .iter()
3100                .all(|&t| crate::fa_split_keys(t, cfg.n_head_kv as usize) == sp0);
3101
3102        Ok(BatchLayerCtx {
3103            lin_base,
3104            attn_base,
3105            ptr_table,
3106            t_kvs,
3107            t_kv_max,
3108            sp0,
3109            seqs_fa,
3110            lo,
3111            hi,
3112        })
3113    }
3114
3115    /// THE PP SEAM (pp2-batch increment 1, 2026-08-06): run the batched trunk over layers
3116    /// `[ctx.lo, ctx.hi)`, entering with a materialized `[B, n_embd]` residual and exiting
3117    /// with the range's final residual materialized. The batched twin of
3118    /// `decode_layers_eager` — the eager arm has had this seam since M1-PP2 and every ppN
3119    /// stage calls it; the batched body had no equivalent, which is why every later PP-2
3120    /// increment (and spec-over-PP2, whose verify is a batched T=K+1 forward) waited on this
3121    /// extraction (`research/pp2-hardening-20260806/PROGRESS.md` bill item 1).
3122    ///
3123    /// SINGLE-DEVICE SEMANTICS ARE UNCHANGED BY CONSTRUCTION: the body is the old
3124    /// `for (il, layer) in self.layers.iter().enumerate()` loop moved verbatim, with `for il
3125    /// in ctx.lo..ctx.hi` as the header and the per-step invariants (`ptr_table`, arm picks,
3126    /// `t_kv`) read from `ctx` instead of enclosing locals. At `lo=0, hi=n_layers` — every
3127    /// call today — the launch sequence is identical, so the exactness contract in this
3128    /// module's header carries over untouched rather than needing a re-proof.
3129    ///
3130    /// UNLIKE the eager seam, this one is NOT yet stage-callable: `caches` is `&mut [&mut
3131    /// Cache]` mutated in place (KV `len` bumps, ssm ping-pong swaps), and `pos_d`/`x` come
3132    /// from the caller's device. Wiring a stage split means per-stage `pos_d` + a boundary
3133    /// `[B, n_embd]` transfer around this call, which is the NEXT increment. The seam exists
3134    /// so that increment is a call-site change, not a 250-line surgery.
3135    #[allow(clippy::too_many_arguments)]
3136    pub(crate) fn decode_batch_layers(
3137        &self,
3138        e: &Engine,
3139        mut x: CudaSlice<f32>,
3140        caches: &mut [&mut Cache],
3141        ctx: &BatchLayerCtx,
3142        pos_d: &CudaSlice<i32>,
3143        ph_last: &mut std::time::Instant,
3144    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3145        let b_n = caches.len();
3146        let cfg = &self.cfg;
3147        let n_embd = cfg.n_embd as usize;
3148        let eps = cfg.rms_eps;
3149        let (lin_base, attn_base) = (&ctx.lin_base, &ctx.attn_base);
3150        let ptr_table = &ctx.ptr_table;
3151        let (seqs_fa, sp0, t_kv_max) = (ctx.seqs_fa, ctx.sp0, ctx.t_kv_max);
3152        debug_assert_eq!(
3153            ctx.t_kvs.len(),
3154            b_n,
3155            "ctx built for a different batch width"
3156        );
3157
3158        for il in ctx.lo..ctx.hi {
3159            let layer = &self.layers[il];
3160            // ---- attn_norm + q8_1 quantize, batched (B rows) ----
3161            let anorm = layer.attn_norm.float_data();
3162            let mut xn = e.uninit(b_n * n_embd)?;
3163            e.rms_norm(&x, anorm, &mut xn, n_embd, b_n, eps)?;
3164            let (hq, hd) = e.quantize_q8_1(&xn, b_n, n_embd)?;
3165
3166            // ---- mixer ----
3167            let mixed: CudaSlice<f32> = match &layer.mixer {
3168                Mixer::Mla(_) => crate::hybrid::mla_path_unimplemented("batched PP decode"),
3169                Mixer::Kda(_) => crate::hybrid::kda_path_unimplemented("batched decode"),
3170                Mixer::Full(fa) => {
3171                    let geometry = cfg.full_attention_geometry_at(il as u32);
3172                    let n_head = geometry.n_head as usize;
3173                    let n_head_kv = geometry.n_head_kv as usize;
3174                    let head_dim = geometry.head_dim_k as usize;
3175                    let rope_dims = geometry.n_rot as usize;
3176                    let rope_base = geometry.rope_base;
3177                    let scale = geometry.attention_scale();
3178                    // Batched projections: one weight read serves all B rows. At B=1 the
3179                    // QKV triple fuses into ONE launch (rig-native decode increment 1 —
3180                    // bit-identical per (tensor,row), RIG-NATIVE-DECODE.md); B>1 and
3181                    // non-NVFP4 trunks keep the three singles.
3182                    let (qf, mut k, v) =
3183                        match e.matmul_nvfp4_fused3(&fa.wq, &fa.wk, &fa.wv, &hq, &hd, b_n)? {
3184                            Some(t) => t,
3185                            None => (
3186                                e.matmul_pre(&fa.wq, &hq, &hd, &xn, b_n)?,
3187                                e.matmul_pre(&fa.wk, &hq, &hd, &xn, b_n)?,
3188                                e.matmul_pre(&fa.wv, &hq, &hd, &xn, b_n)?,
3189                            ),
3190                        };
3191
3192                    let gated =
3193                        geometry.attention_gate == memra_gguf::config::AttentionGateKind::FusedQ;
3194                    let (mut q, gate) = if gated {
3195                        let mut qs = e.uninit(b_n * n_head * head_dim)?;
3196                        let mut gs = e.uninit(b_n * n_head * head_dim)?;
3197                        e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, b_n)?;
3198                        (qs, Some(gs))
3199                    } else {
3200                        (qf, None)
3201                    };
3202
3203                    // QK-norm over B*n_head rows, rope with per-row positions.
3204                    let mut qn = e.uninit(b_n * n_head * head_dim)?;
3205                    e.rms_norm(
3206                        &q,
3207                        fa.q_norm.float_data(),
3208                        &mut qn,
3209                        head_dim,
3210                        b_n * n_head,
3211                        eps,
3212                    )?;
3213                    q = qn;
3214                    let mut kn = e.uninit(b_n * n_head_kv * head_dim)?;
3215                    e.rms_norm(
3216                        &k,
3217                        fa.k_norm.float_data(),
3218                        &mut kn,
3219                        head_dim,
3220                        b_n * n_head_kv,
3221                        eps,
3222                    )?;
3223                    k = kn;
3224                    e.rope_neox(
3225                        &mut q, pos_d, head_dim, rope_dims, n_head, b_n, rope_base, 1.0,
3226                    )?;
3227                    e.rope_neox(
3228                        &mut k, pos_d, head_dim, rope_dims, n_head_kv, b_n, rope_base, 1.0,
3229                    )?;
3230                    ph_mark(e, 1, ph_last)?;
3231
3232                    // INCREMENT 2 (2026-08-01): the per-seq (append, attend) launch train
3233                    // becomes two phases. Phase A appends all B rows (one z-batched launch,
3234                    // or the per-seq loop on the seam/fp8 path); phase B attends all B
3235                    // sequences (one blockIdx.z launch + one combine on the batched arm —
3236                    // which also reads q / writes attn at row offsets, killing the per-seq
3237                    // q/a dtod copies — or the per-seq loop when any row is outside the v4
3238                    // arm / a split rung crosses inside the batch). Caches are disjoint per
3239                    // sequence, so the phase split leaves every row's math untouched.
3240                    let q_dim = n_head * head_dim;
3241                    let mut attn = e.uninit(b_n * q_dim)?;
3242                    // ---- phase A: KV append (all B rows) ----
3243                    let (kdk, kdv, ktb, vtb) = {
3244                        let kvl = caches[0].kv[il].as_ref().unwrap();
3245                        (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
3246                    };
3247                    let base = attn_base[il].expect("full layer missing from pointer table");
3248                    let table = ptr_table.as_ref().expect("pointer table missing");
3249                    let kv_view = table.slice(base..base + 2 * b_n);
3250                    e.append_kv_quantized_seqs(&k, &v, &kv_view, pos_d, b_n, kdk, kdv, ktb, vtb)?;
3251                    for cache in caches.iter_mut() {
3252                        let kvl = cache.kv[il].as_mut().unwrap();
3253                        debug_assert_eq!(kvl.len, cache.pos, "kv len / pos out of lockstep");
3254                        kvl.len += 1;
3255                    }
3256                    ph_mark(e, 2, ph_last)?;
3257                    // ---- phase B: attention (all B sequences) ----
3258                    if seqs_fa {
3259                        let (ktb, vtb) = {
3260                            let kvl = caches[0].kv[il].as_ref().unwrap();
3261                            (kvl.k_tok_bytes, kvl.v_tok_bytes)
3262                        };
3263                        let base = attn_base[il].expect("full layer missing from pointer table");
3264                        let table = ptr_table.as_ref().expect("pointer table missing");
3265                        let kv_view = table.slice(base..base + 2 * b_n);
3266                        e.fa_decode_batch_seqs_v4(
3267                            &q, &kv_view, pos_d, &mut attn, head_dim, n_head, n_head_kv, b_n,
3268                            t_kv_max, scale, sp0, ktb, vtb,
3269                        )?;
3270                        ph_mark(e, 4, ph_last)?;
3271                    } else {
3272                        for (bi, cache) in caches.iter_mut().enumerate() {
3273                            let kvl = cache.kv[il].as_mut().unwrap();
3274                            let t_kv = kvl.len;
3275                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
3276                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
3277                            // The fallback keeps one FA launch per distinct KV view, but Q and
3278                            // attention already live in packed row-major buffers. Pass those row
3279                            // views directly; only the arithmetic-free materialization copies go.
3280                            let q_row = q.slice(bi * q_dim..(bi + 1) * q_dim);
3281                            let mut a_row = attn.slice_mut(bi * q_dim..(bi + 1) * q_dim);
3282                            e.fa_decode_kvmod_view(
3283                                &q_row,
3284                                &k_view,
3285                                &v_view,
3286                                &mut a_row,
3287                                head_dim,
3288                                n_head,
3289                                n_head_kv,
3290                                t_kv,
3291                                scale,
3292                                kvl.k_tok_bytes,
3293                                kvl.v_tok_bytes,
3294                                false,
3295                            )?;
3296                            ph_mark(e, 4, ph_last)?;
3297                        }
3298                    }
3299
3300                    // Output gate (element-wise — batches whole) + o-proj at m=B.
3301                    let attn_g = match &gate {
3302                        Some(g) => {
3303                            let n = b_n * q_dim;
3304                            let mut gsig = e.uninit(n)?;
3305                            e.sigmoid(g, &mut gsig, n)?;
3306                            let mut ag = e.uninit(n)?;
3307                            e.mul(&attn, &gsig, &mut ag, n)?;
3308                            ag
3309                        }
3310                        None => attn,
3311                    };
3312                    let o = e.matmul(&fa.wo, &attn_g, b_n)?;
3313                    ph_mark(e, 5, ph_last)?;
3314                    o
3315                }
3316                Mixer::Linear(la) => {
3317                    // v2 (the B-scaling fix): the GDN mixer's PROJECTIONS carry the layer's
3318                    // weight mass — batch them at m=B so wqkv/gate/beta/alpha/ssm_out stream
3319                    // ONCE per step instead of once per sequence. Only the recurrent state ops
3320                    // (fused conv ring, gdn prep, gdn scan) stay per-seq — they are state-bound
3321                    // micro-kernels, not weight readers. Composition unchanged vs v1 (matmul_pre
3322                    // == fused2 per (tensor,row); _bN mmvq per-row == m=1): same numeric config.
3323                    let geometry = la.geometry;
3324                    let d_state = geometry.key_head_dim as usize;
3325                    let num_k = geometry.key_heads as usize;
3326                    let num_v = geometry.value_heads as usize;
3327                    let d_conv = geometry.conv_kernel as usize;
3328                    let key_dim = d_state * num_k;
3329                    let value_dim = geometry.value_head_dim as usize * num_v;
3330                    let conv_dim = key_dim * 2 + value_dim;
3331                    let gdn_scale = 1.0 / (d_state as f32).sqrt();
3332
3333                    // ---- batched projections (the weight win) ----
3334                    // At B=1 the mixer quartet fuses into ONE launch (rig-native decode
3335                    // increment 2 — bit-identical per (tensor,row), RIG-NATIVE-DECODE.md);
3336                    // B>1 and non-NVFP4 trunks keep the four singles.
3337                    let (qkv_mixed, z, beta_raw, alpha) = match e.matmul_nvfp4_fused4(
3338                        &la.wqkv,
3339                        &la.wqkv_gate,
3340                        &la.ssm_beta,
3341                        &la.ssm_alpha,
3342                        &hq,
3343                        &hd,
3344                        b_n,
3345                    )? {
3346                        Some(t) => t,
3347                        None => (
3348                            e.matmul_pre(&la.wqkv, &hq, &hd, &xn, b_n)?,
3349                            e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, b_n)?,
3350                            e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, b_n)?,
3351                            e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, b_n)?,
3352                        ),
3353                    };
3354                    ph_mark(e, 6, ph_last)?;
3355
3356                    // ---- batched recurrent state ops (3 launches for all B sequences) ----
3357                    let base = lin_base[il].expect("linear layer missing from pointer table");
3358                    let table = ptr_table.as_ref().expect("pointer table missing");
3359                    let conv_view = table.slice(base..base + b_n);
3360                    let in_view = table.slice(base + b_n..base + 2 * b_n);
3361                    let out_view = table.slice(base + 2 * b_n..base + 3 * b_n);
3362                    let mut conv_outs = e.uninit(b_n * conv_dim)?;
3363                    e.ssm_conv1d_fused_decode_b(
3364                        &qkv_mixed,
3365                        &conv_view,
3366                        la.ssm_conv1d.float_data(),
3367                        &mut conv_outs,
3368                        conv_dim,
3369                        d_conv,
3370                        b_n,
3371                    )?;
3372                    let mut q_l2 = e.uninit(b_n * value_dim)?;
3373                    let mut k_l2 = e.uninit(b_n * value_dim)?;
3374                    let mut v_gd = e.uninit(b_n * value_dim)?;
3375                    let mut beta_b = e.uninit(b_n * num_v)?;
3376                    let mut g_log = e.uninit(b_n * num_v)?;
3377                    e.gdn_prep_decode_b(
3378                        &conv_outs,
3379                        &beta_raw,
3380                        &alpha,
3381                        la.ssm_dt.float_data(),
3382                        la.ssm_a.float_data(),
3383                        &mut q_l2,
3384                        &mut k_l2,
3385                        &mut v_gd,
3386                        &mut beta_b,
3387                        &mut g_log,
3388                        d_state,
3389                        num_v,
3390                        num_k,
3391                        key_dim,
3392                        eps,
3393                        conv_dim,
3394                        b_n,
3395                    )?;
3396                    let mut o_all = e.uninit(b_n * value_dim)?;
3397                    e.gdn_scan_s128_batched(
3398                        &q_l2, &k_l2, &v_gd, &g_log, &beta_b, &in_view, &out_view, &mut o_all,
3399                        num_v, b_n, gdn_scale,
3400                    )?;
3401                    // ping-pong: scan wrote each seq's alt buffer; swap host handles (the
3402                    // NEXT step's table rebuild picks up the new canonical pointers).
3403                    for cache in caches.iter_mut() {
3404                        let rl = cache.recur[il].as_mut().unwrap();
3405                        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
3406                    }
3407                    ph_mark(e, 7, ph_last)?;
3408
3409                    // ---- batched gated norm + out-projection ----
3410                    let o = if e.uses_q8_1_fast(&la.ssm_out) {
3411                        let (gq, gd) = e.gated_rmsnorm_q8_1(
3412                            &o_all,
3413                            la.ssm_norm.float_data(),
3414                            &z,
3415                            d_state,
3416                            b_n * num_v,
3417                            eps,
3418                        )?;
3419                        let g0 = e.zeros(0)?;
3420                        e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, b_n)?
3421                    } else {
3422                        let mut gn = e.uninit(b_n * value_dim)?;
3423                        e.gated_rmsnorm(
3424                            &o_all,
3425                            la.ssm_norm.float_data(),
3426                            &z,
3427                            &mut gn,
3428                            d_state,
3429                            b_n * num_v,
3430                            eps,
3431                        )?;
3432                        e.matmul(&la.ssm_out, &gn, b_n)?
3433                    };
3434                    ph_mark(e, 8, ph_last)?;
3435                    o
3436                }
3437            };
3438
3439            // ---- residual add + post_attn_norm + FFN, batched ----
3440            let pnorm = layer.post_attn_norm.float_data();
3441            let mut x1 = e.uninit(b_n * n_embd)?;
3442            let mut z = e.uninit(b_n * n_embd)?;
3443            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut z, n_embd, b_n, eps)?;
3444            let ffn_out = match &layer.ffn {
3445                crate::hybrid::Ffn::Dense {
3446                    ffn_gate,
3447                    ffn_up,
3448                    ffn_down,
3449                } => {
3450                    // v1 covers the SiLU family; M3's swigluoai clamp rides a scaled epilogue
3451                    // (m=1 fused tier) — batched M3 lands with the batched-fusion pass.
3452                    assert!(
3453                        !self
3454                            .plan
3455                            .trunk_operations()
3456                            .contains(&memra_gguf::model_plan::OperationKind::SwiGluOaiActivation,),
3457                        "decode_step_batch v1: M3 swigluoai FFN not yet batched"
3458                    );
3459                    let n_ff = ffn_gate.out_features();
3460                    let (zq, zd) = e.quantize_q8_1(&z, b_n, n_embd)?;
3461                    // REFUTED ARM (lane/q27-deepdive, 2026-08-05): fusing this gate+up pair
3462                    // into `matmul_q8_fused2_t` (the fused2_b8 tier) measured FLAT-TO-NEGATIVE
3463                    // at the serving tick — bench c=8 213.1/213.8, 213.9/214.4, 214.4/213.5
3464                    // (sign flips) and serve c=8 paired mean −0.20% over 3 passes. Mechanism:
3465                    // unlike m=1 (where the pair is 128 of 1015 launches in a 7.67%-gap tick),
3466                    // the c=8 tick is 73.2% one weight-bound kernel class with launch cost
3467                    // already hidden — halving 128 launches of ~28k buys nothing. The m=1 arm
3468                    // in `matmul_pre_dual_noscale` (+0.94%) stays; this call site keeps the two
3469                    // launches. Kernel + fused2_b8 wrapper retained: kernel-check gates it at
3470                    // m=5/8 and matmul_q8_fused2_t serves the verify tier. Receipts:
3471                    // research/q27-deepdive-20260805/ (lever3-bench-*, serve-points.jsonl).
3472                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &z, b_n)?;
3473                    let u = e.matmul_pre(ffn_up, &zq, &zd, &z, b_n)?;
3474                    let mut act = e.uninit(b_n * n_ff)?;
3475                    e.silu_mul(&g, &u, &mut act, b_n * n_ff)?;
3476                    let (aq, ad) = e.quantize_q8_1(&act, b_n, n_ff)?;
3477                    e.matmul_pre(ffn_down, &aq, &ad, &act, b_n)?
3478                }
3479                crate::hybrid::Ffn::Moe(m) => {
3480                    // b_n==1: feed the zq8 seam (orndecode B2, see decode.rs twin). Wider
3481                    // ticks keep None — the dev arm quantizes per-token views there and the
3482                    // shexp pair rides the batched matmul, so there is nothing to share.
3483                    if b_n == 1 {
3484                        let zq8 = e.quantize_q8_1(&z, 1, n_embd)?;
3485                        self.moe_ffn_il_zq8(e, m, &z, Some(&zq8), b_n, il as u16)?
3486                    } else {
3487                        self.moe_ffn_il_zq8(e, m, &z, None, b_n, il as u16)?
3488                    }
3489                }
3490            };
3491            // next-layer input x = x1 + ffn_out (batched element-wise add)
3492            let mut x2 = e.uninit(b_n * n_embd)?;
3493            e.add(&x1, &ffn_out, &mut x2, b_n * n_embd)?;
3494            x = x2;
3495            ph_mark(e, 9, ph_last)?;
3496        }
3497        Ok(x)
3498    }
3499
3500    /// Rollback seam for the step35 batched decode arm (lane/step35-batched-decode,
3501    /// 2026-08-08). Default ON; `MEMRA_STEP35_BATCH=0` caps serving at B=1 and makes the
3502    /// batched bodies return Err. Since lane/cx-b1fix, PP-N also refuses the eager B=1
3503    /// numeric class, so the seam disables PP-N Step35 decode rather than serving unstable
3504    /// bytes. Also the b2geo35 gate's CANARY seam — the live assertions must fail under it.
3505    pub fn step35_batch_on() -> bool {
3506        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3507        *ON.get_or_init(|| std::env::var("MEMRA_STEP35_BATCH").as_deref() != Ok("0"))
3508    }
3509
3510    /// THE step35 BATCHED LAYER WALK (lane/step35-batched-decode, 2026-08-08): B sequences
3511    /// share one pass over layers `[lo, hi)` with the REAL step35 geometry — the arm that
3512    /// kills the B=1 pin (34 tok/s aggregate FLAT across c=1..8, round-robin serialized;
3513    /// research/step-sku-20260807 §4) without re-opening the b2ab garbage hole (the generic
3514    /// `decode_batch_layers` ran uniform n_head/full-width rope/no window/no gate over
3515    /// step35 weights and returned HTTP-200 garbage at c>1).
3516    ///
3517    /// SHAPE — batched where the weights are, per-session where the state is:
3518    ///   * attn_norm + quantize + wq/wk/wv/attn_gate projections + q/k norms + rope + head
3519    ///     gate + wo + residual/post-norm + FFN all run at m=B: ONE weight stream serves B
3520    ///     rows (decode is weight-BW-bound; this is the entire win).
3521    ///   * KV append + fa_decode stay a per-session loop — the SWA window makes each
3522    ///     session's KV view a function of ITS OWN `kvl.len` (`off = len-win` when past the
3523    ///     window), and the z-batched seqs kernels take one shared t_kv/rung, not per-row
3524    ///     offsets. This is the same shape as `decode_batch_layers`' per-seq fallback arm,
3525    ///     and it costs launches, not weight bandwidth (KV is per-session state either way).
3526    ///
3527    /// PER-LAYER GEOMETRY (the five mechanisms that make the generic body wrong here, all
3528    /// from `step35_geom`/cfg): n_head 64 full / 96 SWA (wq/wo/attn_gate widths per layer),
3529    /// partial rope (n_rot 64 full / 128 SWA), dual base (5e6/1e4) + `rope_freqs` factors
3530    /// on FULL layers only, SWA window 512 with per-SESSION view offsets, and the separate
3531    /// head-wise `attn_gate` (one pre-sigmoid scalar per (token, head), input = the
3532    /// post-attn_norm hidden, applied before wo).
3533    ///
3534    /// EXACTNESS (the isolation contract, decode-batch-gate gate2's bar): every kernel here
3535    /// is row-independent at m=B or per-session:
3536    ///   * `rms_norm`/`add_rms_norm`/`quantize_q8_1`/`attn_head_gate`/activations: per-row
3537    ///     programs, grid over rows — row bi's bytes are the 1-row call's bytes.
3538    ///   * projections via `matmul_pre` at m=2..8: Q8_0/Q6_K-class rides the b2/b4/b8
3539    ///     batched-mmvq tier (bit-identical per (token,row) to m=1 mmvq); IQ4_XS — this
3540    ///     SKU's trunk class — has no mmvq/batched kernel, so BOTH m=1 decode and the m=B
3541    ///     walk ride `qmatvec_iq4_XS_dp4a` (grid (out_f, m): each column IS the m=1 dp4a
3542    ///     program). Same class at every width = the decode-parity law by construction.
3543    ///   * `rope_neox2` takes per-row positions (tok = row / n_heads) — row bi rotates at
3544    ///     ITS pos with the layer's (n_rot, base, ff), same bits as its solo call.
3545    ///   * per-session append/fa_decode_kvmod: literally the eager arm's calls on that
3546    ///     session's own cache and views.
3547    ///   * MoE (`moe_ffn_il_zq8` at t=B): the router is per-column decode-exact at
3548    ///     t < PRIME_MIN_T (m=1 program per column), sigmoid routing + expert dispatch are
3549    ///     per-token — a session's experts are a function of its own row only.
3550    ///     The known eager-vs-batched FP gap is why PP-N Step35 deliberately serves THIS walk at
3551    ///     B=1 too: the scheduler can change width during a session, so one numeric class must
3552    ///     cover every live width. `b2geo35` pins static widths and an explicit B=1 -> B>1
3553    ///     transition under live defaults.
3554    ///
3555    /// STAGE-SCOPED FROM BIRTH: `[lo, hi)` + caller-supplied engine/pos_d, so
3556    /// `decode_step_batch_ppn` calls it per stage (per-stage engine, per-stage pos_d, the
3557    /// #87 entry fence and boundary slots unchanged) — the pp2-batch seam lesson.
3558    #[allow(clippy::too_many_arguments)]
3559    pub(crate) fn step35_decode_batch_layers(
3560        &self,
3561        e: &Engine,
3562        x: CudaSlice<f32>,
3563        caches: &mut [&mut Cache],
3564        positions: &[i32],
3565        pos_d: &CudaSlice<i32>,
3566        lo: usize,
3567        hi: usize,
3568        ph_last: &mut std::time::Instant,
3569    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3570        self.step35_decode_rows_layers(e, x, caches, positions, pos_d, None, lo, hi, ph_last)
3571    }
3572
3573    /// Diagnostic generalization of the serving walk: `row_to_cache[r]` names the session
3574    /// whose KV row is consumed by hidden row `r`. Serving passes `None`, preserving the
3575    /// identity mapping and its launch sequence. The MoESD harness passes B groups of gamma
3576    /// consecutive rows so each session's verify columns append causally while projections and
3577    /// MoE dispatch see the full B*gamma target width.
3578    #[allow(clippy::too_many_arguments)]
3579    fn step35_decode_rows_layers(
3580        &self,
3581        e: &Engine,
3582        mut x: CudaSlice<f32>,
3583        caches: &mut [&mut Cache],
3584        positions: &[i32],
3585        pos_d: &CudaSlice<i32>,
3586        row_to_cache: Option<&[usize]>,
3587        lo: usize,
3588        hi: usize,
3589        ph_last: &mut std::time::Instant,
3590    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
3591        let b_n = row_to_cache.map_or(caches.len(), |rows| rows.len());
3592        let cfg = &self.cfg;
3593        let n_embd = cfg.n_embd as usize;
3594        let eps = cfg.rms_eps;
3595        if !self.uses_sliding_gated_moe_program() {
3596            return Err(
3597                "sliding-gated-MoE batch rewrite requires its canonical operation class".into(),
3598            );
3599        }
3600        if b_n == 0 || x.len() != b_n * n_embd || positions.len() != b_n || pos_d.len() != b_n {
3601            return Err(format!(
3602                "step35 row mapping shape mismatch: rows={b_n} x={} host_pos={} device_pos={} \
3603                 n_embd={n_embd}",
3604                x.len(),
3605                positions.len(),
3606                pos_d.len(),
3607            )
3608            .into());
3609        }
3610        if row_to_cache.is_some_and(|rows| rows.iter().any(|&ci| ci >= caches.len())) {
3611            return Err("step35 row mapping names a missing cache".into());
3612        }
3613        let cache_index = |row: usize| row_to_cache.map_or(row, |rows| rows[row]);
3614        let has_rank_local_tp = self.layers[lo..hi].iter().any(|layer| {
3615            matches!(
3616                &layer.mixer,
3617                Mixer::Full(fa)
3618                    if fa
3619                        .step_tp_qkv
3620                        .as_ref()
3621                        .is_some_and(|tp| tp.attention.is_some())
3622            )
3623        });
3624        // MEMRA_STEP_TP_BATCH=1: the t-row batched step-TP walk — per layer, ONE t-grid
3625        // attn norm + ONE weight-amortized QKV over all rows, per-row attention on its
3626        // OWN session cache (the unmodified t=1 program via the col-select door), the
3627        // o_proj deferred and joined once per layer, one t-grid residual norm, one
3628        // t-row routed-expert sweep with a single combine per rank, and the exact t=1
3629        // shexp per row. Every kernel is the per-row-exact twin from the verify walk's
3630        // pedigree, so each session's greedy output is bit-equal to the layer-major-b1
3631        // replay below. Rows chunk at the tcol width (8).
3632        static TPB: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3633        let tp_batch =
3634            *TPB.get_or_init(|| std::env::var("MEMRA_STEP_TP_BATCH").as_deref() == Ok("1"));
3635        if b_n > 1
3636            && b_n <= 8
3637            && has_rank_local_tp
3638            && tp_batch
3639            && crate::tp::step_tp_qkv_fused_enabled().unwrap_or(false)
3640            && self.layers[lo..hi].iter().all(|layer| {
3641                matches!(
3642                    &layer.mixer,
3643                    Mixer::Full(fa)
3644                        if fa.step_tp_qkv.as_ref().is_some_and(|tp| {
3645                            tp.attention.is_some() && tp.runtime.native_p2p()
3646                        })
3647                )
3648            })
3649        {
3650            static ONCE: std::sync::Once = std::sync::Once::new();
3651            ONCE.call_once(|| {
3652                eprintln!(
3653                    "[step-tp-batch-trow] rows={b_n} execution=t-row-batched \
3654                     attention=per-session-rank-local kv_cache=per-session-distributed \
3655                     exactness=per-row-b1-twins performance_claim=false"
3656                );
3657            });
3658            let mut row_positions = Vec::with_capacity(b_n);
3659            for &position in positions {
3660                row_positions.push(e.htod_i32(&[position])?);
3661            }
3662            let mut x_t = x;
3663            let mut h_row = e.uninit(n_embd)?;
3664            let mut mixed_row = e.uninit(n_embd)?;
3665            let t = b_n;
3666            let mut pos_staged = false;
3667            for il in lo..hi {
3668                let layer = &self.layers[il];
3669                let mut h_t = e.uninit(t * n_embd)?;
3670                e.rms_norm(&x_t, layer.attn_norm.float_data(), &mut h_t, n_embd, t, eps)?;
3671                if !self.step35_verify_qkv_precompute(e, il, &h_t, t)? {
3672                    return Err(format!(
3673                        "step-tp-batch layer {il} lost tcol eligibility mid-walk \
3674                         (weights/doors changed under a live batch)"
3675                    )
3676                    .into());
3677                }
3678                // Per-session t-row fa: when every row's session clears the dcw doors,
3679                // the per-row pass stashes q+gate (append still lands per session) and
3680                // ONE table-kernel launch per rank attends all rows.
3681                let fa_rows =
3682                    self.step35_batch_fa_rows_precheck(caches, cache_index, positions, il)?;
3683                let mut next = e.uninit(t * n_embd)?;
3684                let mut deferred: Vec<usize> = Vec::new();
3685                let mut fa_deferred: Vec<usize> = Vec::new();
3686                // FULL t-row attention pass (rope/append + fa + combine + o_proj join in
3687                // 3 launches/rank): skips the per-row loop entirely. The device counters
3688                // advance in-kernel; mirror the HOST cache bookkeeping exactly as the
3689                // per-row tail would (staged/committed txn + local len + lazy mirror).
3690                static RR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3691                let rope_rows_on =
3692                    *RR.get_or_init(|| std::env::var("MEMRA_ROPE_ROWS").as_deref() != Ok("0"));
3693                static RRL: std::sync::OnceLock<Option<(usize, usize)>> =
3694                    std::sync::OnceLock::new();
3695                let rr_layer = *RRL.get_or_init(|| {
3696                    let v = std::env::var("MEMRA_ROPE_ROWS_LAYER").ok()?;
3697                    if let Some((a, b)) = v.split_once('-') {
3698                        Some((a.parse().ok()?, b.parse().ok()?))
3699                    } else {
3700                        let x: usize = v.parse().ok()?;
3701                        Some((x, x))
3702                    }
3703                });
3704                let rr_this = rr_layer.is_none_or(|(a, b)| il >= a && il <= b);
3705                let full_mixed = if fa_rows && rope_rows_on && rr_this {
3706                    self.step35_batch_rope_fa_pass(
3707                        e,
3708                        il,
3709                        caches,
3710                        cache_index,
3711                        positions,
3712                        t,
3713                        !pos_staged,
3714                    )?
3715                } else {
3716                    None
3717                };
3718                if let Some(mixed_t) = &full_mixed {
3719                    pos_staged = true;
3720                    #[allow(clippy::needless_range_loop)]
3721                    // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
3722                    for r in 0..t {
3723                        let ci = cache_index(r);
3724                        let cache = &mut *caches[ci];
3725                        let tp_kv = cache.tp_kv[il]
3726                            .as_mut()
3727                            .expect("precheck verified the distributed cache");
3728                        let transaction = tp_kv.begin_transaction()?;
3729                        let Mixer::Full(fa) = &self.layers[il].mixer else {
3730                            return Err("step-tp-batch expects full attention".into());
3731                        };
3732                        let tp = fa
3733                            .step_tp_qkv
3734                            .as_ref()
3735                            .ok_or("step-tp-batch lost its TP state")?;
3736                        let empty: [CudaSlice<f32>; 0] = [];
3737                        tp.runtime.append_tp_kv_transaction_inner(
3738                            tp_kv,
3739                            transaction,
3740                            &empty,
3741                            &empty,
3742                            1,
3743                            true,
3744                        )?;
3745                        tp.runtime
3746                            .commit_tp_kv_transaction_external(tp_kv, transaction, 1)?;
3747                        if let Some(local) = cache.kv[il].as_mut() {
3748                            local.len = positions[r] as usize + 1;
3749                            if !crate::tp::len_mirror_lazy_on() {
3750                                let _main = e.gpu.enter_main()?;
3751                                e.set_i32_one(&mut local.len_d, local.len as i32)?;
3752                            }
3753                        }
3754                    }
3755                    let o_out = mixed_t.len() / t;
3756                    {
3757                        for r in 0..t {
3758                            e.dtod_copy_view(
3759                                &mixed_t.slice(r * o_out..(r + 1) * o_out),
3760                                &mut mixed_row,
3761                            )?;
3762                            let mut x_row = e.uninit(n_embd)?;
3763                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
3764                            let (x1, ffn_out) = self
3765                                .residual_norm_ffn(e, layer, &x_row, &mixed_row, n_embd, il, eps)?;
3766                            let mut x2 = e.uninit(n_embd)?;
3767                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
3768                            e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
3769                        }
3770                    }
3771                    x_t = next;
3772                    continue;
3773                }
3774                for r in 0..t {
3775                    e.dtod_copy_view(&h_t.slice(r * n_embd..(r + 1) * n_embd), &mut h_row)?;
3776                    crate::tp::set_verify_tcol(Some(r));
3777                    if fa_rows {
3778                        crate::tp::set_spec_fa2_defer(Some(r));
3779                    } else {
3780                        crate::tp::set_tcol_oproj_defer(Some(r));
3781                    }
3782                    let mixed = match &layer.mixer {
3783                        Mixer::Full(fa) => {
3784                            let ci = cache_index(r);
3785                            self.full_attn_decode(
3786                                e,
3787                                fa,
3788                                &h_row,
3789                                &row_positions[r],
3790                                positions[r] as usize,
3791                                &mut *caches[ci],
3792                                il,
3793                            )
3794                        }
3795                        _ => Err("step-tp-batch expects full attention".into()),
3796                    };
3797                    crate::tp::set_verify_tcol(None);
3798                    crate::tp::set_spec_fa2_defer(None);
3799                    crate::tp::set_tcol_oproj_defer(None);
3800                    let mixed = mixed?;
3801                    if fa_rows && crate::tp::take_spec_fa2_stashed() {
3802                        fa_deferred.push(r);
3803                    } else if crate::tp::take_tcol_oproj_stashed() {
3804                        deferred.push(r);
3805                    } else {
3806                        // Ineligible column (sub-floor ctx / rebase): finish this row
3807                        // with the ordinary per-row body.
3808                        let mut x_row = e.uninit(n_embd)?;
3809                        e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
3810                        let (x1, ffn_out) =
3811                            self.residual_norm_ffn(e, layer, &x_row, &mixed, n_embd, il, eps)?;
3812                        let mut x2 = e.uninit(n_embd)?;
3813                        e.add(&x1, &ffn_out, &mut x2, n_embd)?;
3814                        e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
3815                    }
3816                }
3817                if !fa_deferred.is_empty() && fa_deferred.len() != t {
3818                    return Err("step-tp-batch fa rows stashed a strict subset of rows".into());
3819                }
3820                if fa_deferred.len() == t {
3821                    deferred = fa_deferred;
3822                }
3823                if !deferred.is_empty() {
3824                    let mixed_t = if fa_rows && deferred.len() == t {
3825                        self.step35_batch_fa_rows_join(e, il, caches, cache_index, positions, t)?
3826                    } else {
3827                        self.step35_verify_oproj_tcol(e, il, t)?
3828                    };
3829                    let o_out = mixed_t.len() / t;
3830                    {
3831                        for &r in &deferred {
3832                            e.dtod_copy_view(
3833                                &mixed_t.slice(r * o_out..(r + 1) * o_out),
3834                                &mut mixed_row,
3835                            )?;
3836                            let mut x_row = e.uninit(n_embd)?;
3837                            e.dtod_copy_view(&x_t.slice(r * n_embd..(r + 1) * n_embd), &mut x_row)?;
3838                            let (x1, ffn_out) = self
3839                                .residual_norm_ffn(e, layer, &x_row, &mixed_row, n_embd, il, eps)?;
3840                            let mut x2 = e.uninit(n_embd)?;
3841                            e.add(&x1, &ffn_out, &mut x2, n_embd)?;
3842                            e.dtod_copy_into(&x2, &mut next, r * n_embd)?;
3843                        }
3844                    }
3845                }
3846                x_t = next;
3847            }
3848            return Ok(x_t);
3849        }
3850        if b_n > 1 && has_rank_local_tp {
3851            static ONCE: std::sync::Once = std::sync::Once::new();
3852            ONCE.call_once(|| {
3853                eprintln!(
3854                    "[step-tp-batch-exact] rows={b_n} execution=layer-major-b1 \
3855                     attention=rank-local kv_cache=per-session-distributed \
3856                     transport=native-p2p exactness=b1-full-layer-program \
3857                     performance_claim=false"
3858                );
3859            });
3860            // Preserve the isolated B=1 numerical program for every live session. The scheduler
3861            // may change width after any token; allowing norms, residuals, experts, or the head
3862            // to select a B-dependent kernel changes greedy output even when attention itself is
3863            // rowwise. Replay one layer across all rows before advancing so the same TP/EP
3864            // weights remain hot, while every row still executes the qualified B=1 program.
3865            let mut row_states = Vec::with_capacity(b_n);
3866            let mut row_positions = Vec::with_capacity(b_n);
3867            #[allow(clippy::needless_range_loop)]
3868            // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
3869            for row in 0..b_n {
3870                let mut h_row = e.uninit(n_embd)?;
3871                e.copy_view_into(
3872                    &mut h_row,
3873                    0,
3874                    &x.slice(row * n_embd..(row + 1) * n_embd),
3875                    n_embd,
3876                )?;
3877                row_states.push(h_row);
3878                row_positions.push(e.htod_i32(&[positions[row]])?);
3879            }
3880            for il in lo..hi {
3881                let mut next_states = Vec::with_capacity(b_n);
3882                for (row, h_row) in row_states.into_iter().enumerate() {
3883                    let position = [positions[row]];
3884                    let cache = cache_index(row);
3885                    let mut one = [&mut *caches[cache]];
3886                    next_states.push(self.step35_decode_rows_layers(
3887                        e,
3888                        h_row,
3889                        &mut one,
3890                        &position,
3891                        &row_positions[row],
3892                        None,
3893                        il,
3894                        il + 1,
3895                        ph_last,
3896                    )?);
3897                }
3898                row_states = next_states;
3899            }
3900            let mut outputs = e.uninit(b_n * n_embd)?;
3901            for (row, output) in row_states.iter().enumerate() {
3902                e.copy_into(&mut outputs, row * n_embd, output, n_embd)?;
3903            }
3904            return Ok(outputs);
3905        }
3906        let rank_local_positions = if has_rank_local_tp {
3907            let mut device_positions = Vec::with_capacity(b_n);
3908            for &position in positions {
3909                device_positions.push(e.htod_i32(&[position])?);
3910            }
3911            Some(device_positions)
3912        } else {
3913            None
3914        };
3915        // b2geo35 gate evidence: one line, first B>1 walk only (grep-stable prefix).
3916        if b_n > 1 {
3917            static ONCE: std::sync::Once = std::sync::Once::new();
3918            ONCE.call_once(|| {
3919                eprintln!(
3920                    "[step35-batch] first B>1 batched step35 walk: B={b_n} layers=[{lo},{hi})"
3921                );
3922            });
3923        }
3924
3925        for il in lo..hi {
3926            let layer = &self.layers[il];
3927            let Mixer::Full(fa) = &layer.mixer else {
3928                return Err(format!("step35 layer {il} is not full-attn — corrupt config").into());
3929            };
3930            let geometry = self.step35_geom(il);
3931            let hd = geometry.head_dim_k as usize;
3932            let nkv = geometry.n_head_kv as usize;
3933            let nh = geometry.n_head as usize;
3934            let rbase = geometry.rope_base;
3935            let scale = geometry.attention_scale();
3936            let swa = geometry.window.is_some();
3937            let win = geometry.window.unwrap_or(0) as usize;
3938            let n_rot = geometry.n_rot as usize;
3939            let q_dim = nh * hd;
3940            let kv_dim = nkv * hd;
3941
3942            // ---- attn_norm + q8_1 quantize, batched (B rows) ----
3943            let anorm = layer.attn_norm.float_data();
3944            let mut xn = e.uninit(b_n * n_embd)?;
3945            e.rms_norm(&x, anorm, &mut xn, n_embd, b_n, eps)?;
3946            let rank_local_tp = fa
3947                .step_tp_qkv
3948                .as_ref()
3949                .is_some_and(|tp| tp.attention.is_some());
3950            let mixed = if rank_local_tp {
3951                // The B>1 path returns through the full-row oracle above. This branch is therefore
3952                // the qualified B=1 rank-local TP attention program.
3953                let row_positions = rank_local_positions
3954                    .as_ref()
3955                    .expect("rank-local TP positions were prepared");
3956                let mut outputs = e.uninit(b_n * n_embd)?;
3957                #[allow(clippy::needless_range_loop)]
3958                // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
3959                for row in 0..b_n {
3960                    let mut h_row = e.uninit(n_embd)?;
3961                    e.copy_view_into(
3962                        &mut h_row,
3963                        0,
3964                        &xn.slice(row * n_embd..(row + 1) * n_embd),
3965                        n_embd,
3966                    )?;
3967                    let cache = cache_index(row);
3968                    let output = self.step35_decode_attn(
3969                        e,
3970                        fa,
3971                        il,
3972                        &h_row,
3973                        None,
3974                        &row_positions[row],
3975                        caches[cache],
3976                    )?;
3977                    e.copy_into(&mut outputs, row * n_embd, &output, n_embd)?;
3978                }
3979                outputs
3980            } else {
3981                let (hq, hdq) = e.quantize_q8_1(&xn, b_n, n_embd)?;
3982
3983                // ---- batched projections: q/k/v + the separate head-wise gate (one weight
3984                // stream for B rows; xn is the live f32 fallback for non-q8_1-fast classes) ----
3985                let q0 = e.matmul_pre(&fa.wq, &hq, &hdq, &xn, b_n)?;
3986                let k0 = e.matmul_pre(&fa.wk, &hq, &hdq, &xn, b_n)?;
3987                let v0 = e.matmul_pre(&fa.wv, &hq, &hdq, &xn, b_n)?;
3988                let gw = fa
3989                    .attn_gate
3990                    .as_ref()
3991                    .ok_or("step35 layer is missing attn_gate.weight (head-wise attention gate)")?;
3992                // gate input = the post-attn_norm hidden (upstream `cur`) — same xn/q8 pair.
3993                let gt = e.matmul_pre(gw, &hq, &hdq, &xn, b_n)?;
3994
3995                // ---- q/k RMSNorm over head_dim rows + the per-layer PARTIAL rope ----
3996                let mut q = e.uninit(b_n * q_dim)?;
3997                e.rms_norm(&q0, fa.q_norm.float_data(), &mut q, hd, b_n * nh, eps)?;
3998                let mut k = e.uninit(b_n * kv_dim)?;
3999                e.rms_norm(&k0, fa.k_norm.float_data(), &mut k, hd, b_n * nkv, eps)?;
4000                let ff = if geometry.rope_factors {
4001                    self.step35_aux.as_ref().and_then(|a| a.rope_freqs(e))
4002                } else {
4003                    None
4004                };
4005                e.rope_neox2(
4006                    &mut q, &mut k, pos_d, hd, n_rot, nh, nkv, b_n, rbase, 1.0, ff,
4007                )?;
4008                ph_mark(e, 1, ph_last)?;
4009
4010                // ---- per-session: KV append + windowed/global fa_decode (each session's OWN
4011                // len drives its view offset — the iso-gap law, no cross-session term) ----
4012                let mut attn = e.uninit(b_n * q_dim)?;
4013                if b_n == 1 {
4014                    // B=1 SPECIALIZED ENTRY (lane/cx-eagerpar): the general row loop below
4015                    // materializes q_row and a_row because a B>1 FA call consumes/produces one
4016                    // contiguous row at a time. At B=1, q and attn already ARE those whole rows.
4017                    // Pass them directly to the same fa_decode_kvmod call: this removes two
4018                    // arithmetic-free D2D copies (90 launches/token on Step3.7's 45 layers)
4019                    // without changing any arithmetic kernel, shape, argument value, or order.
4020                    // Keep the B>1 body verbatim below; b1fix's one-class/transition gates are
4021                    // the promotion bar, not an FP-similarity tolerance.
4022                    let kvl = caches[cache_index(0)].kv[il].as_mut().unwrap();
4023                    let k_row = k.slice(0..kv_dim);
4024                    let v_row = v0.slice(0..kv_dim);
4025                    let next_len = kvl.len + 1;
4026                    let (off, t_kv) = if swa && next_len > win {
4027                        (next_len - win, win)
4028                    } else {
4029                        (0, next_len)
4030                    };
4031                    let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
4032                    e.append_kv_quantized_view(
4033                        &k_row,
4034                        &v_row,
4035                        &mut kvl.k,
4036                        &mut kvl.v,
4037                        write_row,
4038                        kvl.kv_dim_k,
4039                        kvl.kv_dim_v,
4040                        kvl.k_tok_bytes,
4041                        kvl.v_tok_bytes,
4042                        false,
4043                    )?;
4044                    kvl.len = next_len;
4045                    ph_mark(e, 2, ph_last)?;
4046                    let physical = kvl.physical_rows(off, off + t_kv)?;
4047                    let k_view = e.view_u8_range(
4048                        &kvl.k,
4049                        physical.start * kvl.k_tok_bytes,
4050                        physical.end * kvl.k_tok_bytes,
4051                    );
4052                    let v_view = e.view_u8_range(
4053                        &kvl.v,
4054                        physical.start * kvl.v_tok_bytes,
4055                        physical.end * kvl.v_tok_bytes,
4056                    );
4057                    e.fa_decode_kvmod(
4058                        &q,
4059                        &k_view,
4060                        &v_view,
4061                        &mut attn,
4062                        hd,
4063                        nh,
4064                        nkv,
4065                        t_kv,
4066                        scale,
4067                        kvl.k_tok_bytes,
4068                        kvl.v_tok_bytes,
4069                        false,
4070                    )?;
4071                    ph_mark(e, 4, ph_last)?;
4072                } else {
4073                    for bi in 0..b_n {
4074                        let cache = &mut caches[cache_index(bi)];
4075                        let kvl = cache.kv[il].as_mut().unwrap();
4076                        let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
4077                        let v_row = v0.slice(bi * kv_dim..(bi + 1) * kv_dim);
4078                        let next_len = kvl.len + 1;
4079                        let (off, t_kv) = if swa && next_len > win {
4080                            (next_len - win, win)
4081                        } else {
4082                            (0, next_len)
4083                        };
4084                        let write_row = e.prepare_kv_append(kvl, off & !31usize, 1)?;
4085                        e.append_kv_quantized_view(
4086                            &k_row,
4087                            &v_row,
4088                            &mut kvl.k,
4089                            &mut kvl.v,
4090                            write_row,
4091                            kvl.kv_dim_k,
4092                            kvl.kv_dim_v,
4093                            kvl.k_tok_bytes,
4094                            kvl.v_tok_bytes,
4095                            false,
4096                        )?;
4097                        kvl.len = next_len;
4098                        ph_mark(e, 2, ph_last)?;
4099                        // the eager arm's SWA view arithmetic, verbatim (step35_decode_attn):
4100                        // token-aligned offset, keys carry absolute rope, mask is positional.
4101                        let physical = kvl.physical_rows(off, off + t_kv)?;
4102                        let k_view = e.view_u8_range(
4103                            &kvl.k,
4104                            physical.start * kvl.k_tok_bytes,
4105                            physical.end * kvl.k_tok_bytes,
4106                        );
4107                        let v_view = e.view_u8_range(
4108                            &kvl.v,
4109                            physical.start * kvl.v_tok_bytes,
4110                            physical.end * kvl.v_tok_bytes,
4111                        );
4112                        // The per-session cache view remains authoritative (including SWA's
4113                        // physical-row rebase), while Q/O use their existing packed row views.
4114                        // This preserves the exact FA program and removes only the two D2D copies.
4115                        let q_row = q.slice(bi * q_dim..(bi + 1) * q_dim);
4116                        let mut a_row = attn.slice_mut(bi * q_dim..(bi + 1) * q_dim);
4117                        e.fa_decode_kvmod_view(
4118                            &q_row,
4119                            &k_view,
4120                            &v_view,
4121                            &mut a_row,
4122                            hd,
4123                            nh,
4124                            nkv,
4125                            t_kv,
4126                            scale,
4127                            kvl.k_tok_bytes,
4128                            kvl.v_tok_bytes,
4129                            false,
4130                        )?;
4131                        ph_mark(e, 4, ph_last)?;
4132                    }
4133                }
4134
4135                // ---- head-wise gate (one sigmoid per (token, head), pre-wo) + o-proj at m=B ----
4136                let mut ag = e.uninit(b_n * q_dim)?;
4137                e.attn_head_gate(&attn, &gt, &mut ag, None, hd, nh, b_n)?;
4138                e.matmul(&fa.wo, &ag, b_n)?
4139            };
4140            ph_mark(e, 5, ph_last)?;
4141
4142            // ---- residual add + post_attn_norm + FFN, batched ----
4143            let pnorm = layer.post_attn_norm.float_data();
4144            let mut x1 = e.uninit(b_n * n_embd)?;
4145            let mut z = e.uninit(b_n * n_embd)?;
4146            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut z, n_embd, b_n, eps)?;
4147            let ffn_out = match &layer.ffn {
4148                crate::hybrid::Ffn::Dense {
4149                    ffn_gate,
4150                    ffn_up,
4151                    ffn_down,
4152                } => {
4153                    // A dense step35 FFN's clamp is the SHEXP array (upstream's one
4154                    // build_ffn serves dense + shared expert, llama-graph.cpp:1751);
4155                    // ffn_act_lim dispatches clamped/plain per layer. Layers 0-2 (the
4156                    // leading dense) have no live limit on this artifact, but the route
4157                    // is correct by construction, not by artifact.
4158                    let n_ff = ffn_gate.out_features();
4159                    let (zq, zd) = e.quantize_q8_1(&z, b_n, n_embd)?;
4160                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &z, b_n)?;
4161                    let u = e.matmul_pre(ffn_up, &zq, &zd, &z, b_n)?;
4162                    let mut act = e.uninit(b_n * n_ff)?;
4163                    Self::ffn_act_lim(
4164                        e,
4165                        cfg,
4166                        &g,
4167                        &u,
4168                        1.0,
4169                        1.0,
4170                        cfg.clamp_shexp_at(il as u32),
4171                        &mut act,
4172                        b_n * n_ff,
4173                    )?;
4174                    let (aq, ad) = e.quantize_q8_1(&act, b_n, n_ff)?;
4175                    e.matmul_pre(ffn_down, &aq, &ad, &act, b_n)?
4176                }
4177                // t=B < PRIME_MIN_T: per-column decode-exact router + host sigmoid routing
4178                // + per-token expert dispatch — the same per-token program as eager t=1,
4179                // including the per-layer SwiGLU clamp (43/44) via the sequential path's
4180                // ffn_act_lim. The sigmoid-router deny on dev/pairs holds by predicate.
4181                crate::hybrid::Ffn::Moe(m) => {
4182                    // b_n==1: feed the zq8 seam (orndecode B2, see decode.rs twin). Wider
4183                    // ticks keep None — the dev arm quantizes per-token views there and the
4184                    // shexp pair rides the batched matmul, so there is nothing to share.
4185                    if b_n == 1 {
4186                        let zq8 = e.quantize_q8_1(&z, 1, n_embd)?;
4187                        self.moe_ffn_il_zq8(e, m, &z, Some(&zq8), b_n, il as u16)?
4188                    } else {
4189                        self.moe_ffn_il_zq8(e, m, &z, None, b_n, il as u16)?
4190                    }
4191                }
4192            };
4193            let mut x2 = e.uninit(b_n * n_embd)?;
4194            e.add(&x1, &ffn_out, &mut x2, b_n * n_embd)?;
4195            x = x2;
4196            ph_mark(e, 9, ph_last)?;
4197        }
4198        Ok(x)
4199    }
4200
4201    /// Kill-switch seam for the gemma4 dense-31B batched decode arm. DEFAULT ON since the
4202    /// 2026-08-16 owner flip ("if the performance are so strong in favor... we serve the
4203    /// correctness and best performance"): the arm's exactness battery is green at B=4/8,
4204    /// the served identity gate is byte-exact vs eager at c1/c4, and the served aggregate
4205    /// read 55→257 tok/s c16 on the NVFP4mix artifact at 450W (SERVED-AGGREGATE.md).
4206    /// `MEMRA_GEMMA4_BATCH=0` forces the eager per-session path (the rollback);
4207    /// `1` is the old opt-in spelling, still accepted. Any OTHER value REFUSES LOUD at
4208    /// first use — a mis-typed kill switch must not silently pick a serving path.
4209    pub fn gemma4_batch_on() -> bool {
4210        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4211        *ON.get_or_init(|| match std::env::var("MEMRA_GEMMA4_BATCH").as_deref() {
4212            Err(_) | Ok("1") => true,
4213            Ok("0") => false,
4214            Ok(v) => panic!(
4215                "MEMRA_GEMMA4_BATCH={v:?} is not a recognized value (want unset/1 = batched \
4216                 decode, 0 = eager kill switch) — refusing to guess a serving path"
4217            ),
4218        })
4219    }
4220
4221    /// THE gemma4 dense-31B BATCHED DECODE ARM (lane/gemma-batched, 2026-08-16).
4222    ///
4223    /// gemma4 served eager-only — the c1→c8 aggregate was FLAT (~55 tok/s, per-stream
4224    /// collapse) because there was no batched arm, not because of quantization. This is it.
4225    ///
4226    /// SHAPE — batched where the weights are, per-session where the state is (the step35
4227    /// law, applied to gemma4's own geometry):
4228    ///   * embed+scale, attn_norm+q8_1 quantize, wq/wk/wv projections, q/k RMSNorm +
4229    ///     weightless-V norm + dual rope (fused `rms_norm_qkv_rope`), post_attn_norm, the
4230    ///     layer-scale tail with its dense GEGLU FFN (`gemma4_layer_tail_add_nq`), output
4231    ///     norm, softcapped head — ALL at m=B: one weight stream serves B rows (decode is
4232    ///     weight-BW-bound; that is the entire aggregate win). Every one of these is the
4233    ///     SAME batch-capable function the proven verify trunk (`gemma4_verify_trunk`) runs
4234    ///     at width t, so this arm inherits the verify path's numerics wholesale.
4235    ///   * KV append + fa_decode stay a PER-SESSION loop: each session appends its one new
4236    ///     token to its own cache and attends its own [win_off .. len] view — the SWA
4237    ///     window + global-vs-windowed geometry makes each session's t_kv independent, so
4238    ///     there is no cross-session batched attention (identical to eager per session).
4239    ///
4240    /// EXACTNESS: v1 routes every session's attention through `fa_decode_kvmod` (the eager
4241    /// arm's unconditional fallback — same call `gemma4_decode_attn` makes with the rows_w
4242    /// fast arms off), so a B=1 run is the eager decode's own attention program and the
4243    /// batch is per-row independent by construction. The rows / rows_w per-session fast
4244    /// arms are a later perf increment gated behind their own seam.
4245    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4246    fn gemma4_decode_batch(
4247        &self,
4248        e: &Engine,
4249        tokens: &[u32],
4250        caches: &mut [&mut Cache],
4251        samp: &[Option<DevSamp>],
4252        masks: &[Option<(&CudaSlice<u32>, usize)>],
4253        lean: bool,
4254    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
4255        let b_n = tokens.len();
4256        if b_n == 0 || b_n != caches.len() {
4257            return Err(format!(
4258                "gemma4_decode_batch: tokens/caches mismatch (tokens={b_n}, caches={})",
4259                caches.len()
4260            )
4261            .into());
4262        }
4263        // Exactness tier boundary: the battery is green at B<=8 (per-row mmvq); m>8
4264        // crosses the dp4a-tail/GEMM numeric configs it never proved. The worker's chunk
4265        // policy caps gemma4 at 8; this is the per-request backstop (Err, never a panic —
4266        // the 2026-08-07 worker-FATAL law).
4267        if b_n > 8 {
4268            return Err(format!(
4269                "gemma4_decode_batch: B={b_n} > 8, past the proven exactness tier — \
4270                 the scheduler must chunk gemma4 at <=8"
4271            )
4272            .into());
4273        }
4274        let n_embd = self.cfg.n_embd as usize;
4275        let eps = self.cfg.rms_eps;
4276        if b_n > 1 {
4277            static ONCE: std::sync::Once = std::sync::Once::new();
4278            ONCE.call_once(|| {
4279                eprintln!("[gemma4-batch] first B>1 batched gemma4 walk: B={b_n}");
4280            });
4281        }
4282        // per-session rope positions (each sequence at its own depth).
4283        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
4284        let pos_d = e.htod_i32(&pos_v)?;
4285        let mut x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
4286        e.scale_inplace(&mut x, (n_embd as f32).sqrt(), b_n * n_embd)?;
4287        // cross-layer carry: each tail emits the next layer's attn-normed q8_1 input.
4288        let mut h_carry: Option<(CudaSlice<i8>, CudaSlice<f32>)> = None;
4289        let n_layers = self.layers.len();
4290        for (il, layer) in self.layers.iter().enumerate() {
4291            let (hq, hdq) = match h_carry.take() {
4292                Some(p) => p,
4293                None => {
4294                    e.rms_norm_q8_1(&x, self.layers[0].attn_norm.float_data(), n_embd, b_n, eps)?
4295                }
4296            };
4297            let Mixer::Full(fa) = &layer.mixer else {
4298                return Err(format!("gemma4 layer {il} not full-attn — corrupt config").into());
4299            };
4300            // STAGE-A ORACLE ARM (MEMRA_FAST=0) ONLY. `matmul_pre`'s raw-f32 escape needs the f32
4301            // attn-normed activation, and this trunk never materializes one — `rms_norm_q8_1`
4302            // above returns just the (i8, f32-scales) pair, which is exactly why the projections
4303            // used to be handed `e.zeros(0)` and read out of bounds.
4304            //
4305            // `rms_norm_decode` is the right producer and not merely a convenient one: it is
4306            // documented BIT-IDENTICAL to `rms_norm_q8_1`'s sum-of-squares reduction (same
4307            // blockDim=1024, same shfl tree), which is the property the spec verify path already
4308            // depends on. So the f32 recomputed here is precisely the tensor `rms_norm_q8_1`
4309            // quantized — the oracle compares against the same activation the fast path saw,
4310            // differing only in the weight-side arithmetic it is meant to be checking.
4311            //
4312            // Cost on the daily path: ONE branch on a OnceLock bool. Nothing is allocated and no
4313            // kernel is launched unless MEMRA_FAST=0.
4314            let h_raw = if Engine::stage_a_raw_needed() {
4315                let mut hf = e.uninit(b_n * n_embd)?;
4316                e.rms_norm_decode(&x, layer.attn_norm.float_data(), &mut hf, n_embd, b_n, eps)?;
4317                Some(hf)
4318            } else {
4319                None
4320            };
4321            let o =
4322                self.gemma4_batch_attn(e, fa, il, &hq, &hdq, h_raw.as_ref(), &pos_d, b_n, caches)?;
4323            let next_norm = if il + 1 < n_layers {
4324                Some(self.layers[il + 1].attn_norm.float_data())
4325            } else {
4326                None
4327            };
4328            // pn-fold front (lane/gemma-pnfold merge): the batched arm rides the SAME
4329            // tail front as the eager/verify trio, so batched == eager holds by
4330            // construction at either MEMRA_G4_PNFOLD value (seam-off falls through to
4331            // the unfused rms_norm + tail chain this arm shipped with).
4332            let (xn, hn) = self.gemma4_layer_tail_add_nq_pn(e, layer, &o, &x, b_n, next_norm)?;
4333            x = xn;
4334            h_carry = hn;
4335        }
4336        let mut hn = e.uninit(b_n * n_embd)?;
4337        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
4338        let mut ld = e.matmul(&self.output, &hn, b_n)?;
4339        let cap = self.cfg.gemma4.as_ref().unwrap().final_logit_softcapping;
4340        e.softcap(&mut ld, cap, b_n * self.output.out_features())?;
4341        self.gemma4_suppress(e, &mut ld, b_n)?; // non-monotonic — before any argmax/sample
4342        let mut ph_last = std::time::Instant::now();
4343        self.decode_batch_epilogue(e, caches, samp, masks, lean, ld, b_n, &mut ph_last, None)
4344    }
4345
4346    /// Per-session gemma4 attention for the batched arm: batched projections + fused
4347    /// q/k-norm + weightless-V-norm + dual rope over all B rows (per-row independent, the
4348    /// verify path's exact kernels), then a per-session KV append + `fa_decode_kvmod` over
4349    /// each session's own window/global view, then one batched wo matmul. Mirrors the eager
4350    /// `gemma4_decode_attn` fallback per row.
4351    #[allow(clippy::too_many_arguments)]
4352    fn gemma4_batch_attn(
4353        &self,
4354        e: &Engine,
4355        fa: &crate::hybrid::FullAttnLayer,
4356        il: usize,
4357        hq: &CudaSlice<i8>,
4358        hdq: &CudaSlice<f32>,
4359        h_raw: Option<&CudaSlice<f32>>,
4360        pos_d: &CudaSlice<i32>,
4361        b_n: usize,
4362        caches: &mut [&mut Cache],
4363    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4364        let (hd, nkv, nh, base, scale, swa) = self.gemma4_geom(il);
4365        let eps = self.cfg.rms_eps;
4366        let aux = self.gemma4_aux.as_ref().unwrap();
4367        let ones = aux.ones(e);
4368        // `h_raw` is Some ONLY under MEMRA_FAST=0, where matmul_pre takes its raw-f32 escape and
4369        // therefore needs a real activation; on the daily path it is None and the empty slice keeps
4370        // the old behaviour exactly (matmul_pre reads the q8_1 pair and never touches this buffer).
4371        let h0 = e.zeros(0)?;
4372        let h = h_raw.unwrap_or(&h0);
4373        // projections at m=B (on the fast path the f32 fallback `h` is empty and matmul_pre uses
4374        // the q8_1 pair; under the Stage-A oracle `h` carries the real f32 attn-normed rows).
4375        let q0 = e.matmul_pre(&fa.wq, hq, hdq, h, b_n)?;
4376        let k0 = e.matmul_pre(&fa.wk, hq, hdq, h, b_n)?;
4377        let v0 = if swa {
4378            e.matmul_pre(&fa.wv, hq, hdq, h, b_n)?
4379        } else {
4380            e.clone_dtod(&k0)? // globals: V := K clone (weightless V-norm, never roped)
4381        };
4382        let mut q = e.uninit(b_n * nh * hd)?;
4383        let mut k = e.uninit(b_n * nkv * hd)?;
4384        let mut v = e.uninit(b_n * nkv * hd)?;
4385        let ff = if swa {
4386            None
4387        } else {
4388            Some(
4389                aux.rope_freqs(e)
4390                    .expect("gemma4 global rope needs rope_freqs.weight"),
4391            )
4392        };
4393        e.rms_norm_qkv_rope(
4394            &q0,
4395            &k0,
4396            &v0,
4397            fa.q_norm.float_data(),
4398            fa.k_norm.float_data(),
4399            ones,
4400            &mut q,
4401            &mut k,
4402            &mut v,
4403            hd,
4404            self.gemma4_rope_dims(il),
4405            nh * b_n,
4406            nkv * b_n,
4407            pos_d,
4408            nh,
4409            nkv,
4410            base,
4411            1.0,
4412            ff,
4413            eps,
4414        )?;
4415        let win = self.cfg.gemma4.as_ref().unwrap().sliding_window as usize;
4416        let q_dim = nh * hd;
4417        let kv_dim = nkv * hd;
4418        let mut attn = e.uninit(b_n * q_dim)?;
4419        #[allow(clippy::needless_range_loop)]
4420        // allow: the explicit index loop keeps the offset arithmetic visible and aligned with the device-side indexing
4421        for bi in 0..b_n {
4422            let kvl = caches[bi].kv[il].as_mut().unwrap();
4423            let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
4424            let v_row = v.slice(bi * kv_dim..(bi + 1) * kv_dim);
4425            // gemma4's KV is a linear buffer (no ring rebase — the SWA view below is a plain
4426            // token-offset), so append at kvl.len exactly as eager gemma4_decode_attn does.
4427            e.append_kv_quantized_view(
4428                &k_row,
4429                &v_row,
4430                &mut kvl.k,
4431                &mut kvl.v,
4432                kvl.len,
4433                kvl.kv_dim_k,
4434                kvl.kv_dim_v,
4435                kvl.k_tok_bytes,
4436                kvl.v_tok_bytes,
4437                (!swa && crate::Engine::gkv_on()) || (swa && crate::Engine::wkv_on()),
4438            )?;
4439            kvl.len += 1;
4440            // eager SWA view arithmetic (gemma4_decode_attn): token-aligned window offset;
4441            // keys carry absolute rope, the mask is purely positional.
4442            let (off_tok, t_kv) = if swa && kvl.len > win {
4443                (kvl.len - win, win)
4444            } else {
4445                (0, kvl.len)
4446            };
4447            let k_view = e.view_u8_range(
4448                &kvl.k,
4449                off_tok * kvl.k_tok_bytes,
4450                (off_tok + t_kv) * kvl.k_tok_bytes,
4451            );
4452            let v_view = e.view_u8_range(
4453                &kvl.v,
4454                off_tok * kvl.v_tok_bytes,
4455                (off_tok + t_kv) * kvl.v_tok_bytes,
4456            );
4457            let q_row = q.slice(bi * q_dim..(bi + 1) * q_dim);
4458            let mut a_row = attn.slice_mut(bi * q_dim..(bi + 1) * q_dim);
4459            e.fa_decode_kvmod_view(
4460                &q_row,
4461                &k_view,
4462                &v_view,
4463                &mut a_row,
4464                hd,
4465                nh,
4466                nkv,
4467                t_kv,
4468                scale,
4469                kvl.k_tok_bytes,
4470                kvl.v_tok_bytes,
4471                swa && crate::Engine::wkv_on(),
4472            )?;
4473        }
4474        e.matmul(&fa.wo, &attn, b_n)
4475    }
4476
4477    /// Standalone MoESD target forward. This entrypoint is not used by serving: it widens the
4478    /// existing Step-3.7 batched layer walk to B*gamma rows while preserving one causal KV chain
4479    /// per session. It returns device logits and performs no sampling or logits D2H, matching the
4480    /// target-model term T_T measured by the paper.
4481    pub fn moesd_target_forward(
4482        &self,
4483        e: &Engine,
4484        tokens: &[u32],
4485        batch: usize,
4486        gamma: usize,
4487        caches: &mut [&mut Cache],
4488    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
4489        if self.hyper.is_some() {
4490            return Err(
4491                "moesd_target_forward: the MoESD speculative target walk has no \
4492                 HyperConnections trunk — it drives `step35_decode_rows_layers`, a serial \
4493                 residual rows-walk, and no [B*gamma, streams, n_embd] hyper rows-walk with \
4494                 causal per-session verify appends exists. mHC speculative verify is a \
4495                 separate lane, not this entry point."
4496                    .into(),
4497            );
4498        }
4499        for cache in caches.iter() {
4500            cache.ensure_usable("moesd_target_forward")?;
4501        }
4502        if crate::plan_backend::decode_batch_program(&self.plan)
4503            != crate::plan_backend::DecodeBatchProgram::SlidingGatedMoe
4504        {
4505            return Err("MoESD target forward currently requires Step-3.7/Step35 geometry".into());
4506        }
4507        if batch == 0 || gamma == 0 || caches.len() != batch || tokens.len() != batch * gamma {
4508            return Err(format!(
4509                "MoESD shape mismatch: B={batch} gamma={gamma} caches={} tokens={}",
4510                caches.len(),
4511                tokens.len(),
4512            )
4513            .into());
4514        }
4515        let rows = batch * gamma;
4516        if rows > 256 {
4517            return Err(format!("MoESD target width {rows} exceeds the frozen 32*8 matrix").into());
4518        }
4519        let _pp_walk =
4520            if crate::pp::pp_cuts(self.layers.len()).is_some() && !crate::pp::pp2_streams_off() {
4521                let rt = crate::pp::PpNRt::get(e)?;
4522                Some(rt.acquire_walk("moesd_target_forward")?)
4523            } else {
4524                None
4525            };
4526        let n_embd = self.cfg.n_embd as usize;
4527        let eps = self.cfg.rms_eps;
4528        let payload = rows * n_embd;
4529        let row_to_cache: Vec<usize> = (0..batch)
4530            .flat_map(|session| (0..gamma).map(move |_| session))
4531            .collect();
4532        let positions: Vec<i32> = row_to_cache
4533            .iter()
4534            .enumerate()
4535            .map(|(row, &session)| (caches[session].pos + row % gamma) as i32)
4536            .collect();
4537        let mut ph_last = std::time::Instant::now();
4538
4539        let logits = if let Some(fence) = crate::pp::pp_cuts(self.layers.len()) {
4540            if fence.len() != 3 || crate::pp::pp2_streams_off() {
4541                return Err(
4542                    "MoESD PP target forward requires the live two-stage stream split".into(),
4543                );
4544            }
4545            let rt = crate::pp::PpNRt::get(e)?;
4546            if rt.n_stages() != 2 {
4547                return Err(format!("MoESD expected two PP stages, got {}", rt.n_stages()).into());
4548            }
4549            let caller_stream = e.stream();
4550            rt.fence_stages_behind(&caller_stream)?;
4551            let slot = {
4552                let _st0 = rt.enter(0);
4553                let e0 = rt.engine(0, e);
4554                let pos_d = e0.htod_i32(&positions)?;
4555                let x = e0.htod(&self.embd.try_gather(n_embd, tokens)?)?;
4556                ph_mark(e0, 0, &mut ph_last)?;
4557                let x = self.step35_decode_rows_layers(
4558                    e0,
4559                    x,
4560                    caches,
4561                    &positions,
4562                    &pos_d,
4563                    Some(&row_to_cache),
4564                    fence[0],
4565                    fence[1],
4566                    &mut ph_last,
4567                )?;
4568                rt.tx(0, &x, payload)?
4569            };
4570
4571            {
4572                let _st1 = rt.enter(1);
4573                let e1 = rt.engine(1, e);
4574                let pos_d = e1.htod_i32(&positions)?;
4575                let x = rt.rx(0, slot, payload)?;
4576                let x = self.step35_decode_rows_layers(
4577                    e1,
4578                    x,
4579                    caches,
4580                    &positions,
4581                    &pos_d,
4582                    Some(&row_to_cache),
4583                    fence[1],
4584                    fence[2],
4585                    &mut ph_last,
4586                )?;
4587                let mut hn = e1.uninit(payload)?;
4588                e1.rms_norm(
4589                    &x,
4590                    self.output_norm.float_data(),
4591                    &mut hn,
4592                    n_embd,
4593                    rows,
4594                    eps,
4595                )?;
4596                let logits = e1.matmul(&self.output, &hn, rows)?;
4597                rt.publish_to(1, &caller_stream)?;
4598                logits
4599            }
4600        } else {
4601            let pos_d = e.htod_i32(&positions)?;
4602            let x = e.htod(&self.embd.try_gather(n_embd, tokens)?)?;
4603            ph_mark(e, 0, &mut ph_last)?;
4604            let x = self.step35_decode_rows_layers(
4605                e,
4606                x,
4607                caches,
4608                &positions,
4609                &pos_d,
4610                Some(&row_to_cache),
4611                0,
4612                self.layers.len(),
4613                &mut ph_last,
4614            )?;
4615            let mut hn = e.uninit(payload)?;
4616            e.rms_norm(
4617                &x,
4618                self.output_norm.float_data(),
4619                &mut hn,
4620                n_embd,
4621                rows,
4622                eps,
4623            )?;
4624            e.matmul(&self.output, &hn, rows)?
4625        };
4626        for cache in caches.iter_mut() {
4627            cache.pos += gamma;
4628        }
4629        Ok(logits)
4630    }
4631
4632    /// The batched tick's TAIL, after the trunk: grammar masks -> device sampling -> lean
4633    /// logits park -> `pos` bump. Split out with the pp seam (`decode_batch_layers`) because
4634    /// under a stage split this runs on the LAST stage's engine and device — the lm_head, the
4635    /// masks, the sampler, and `cache.last_logits_dev` all live where the final residual
4636    /// lands, and the caller must be able to place them there without duplicating 90 lines of
4637    /// serving contract. `logits` is `[b_n, n_vocab]` already computed by the caller (the
4638    /// output_norm + lm_head pair stays at the call site so a stage split can fence around
4639    /// it); everything after it is here, verbatim.
4640    #[allow(clippy::too_many_arguments)]
4641    #[allow(clippy::type_complexity)] // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4642    fn decode_batch_epilogue(
4643        &self,
4644        e: &Engine,
4645        caches: &mut [&mut Cache],
4646        samp: &[Option<DevSamp>],
4647        masks: &[Option<(&CudaSlice<u32>, usize)>],
4648        lean: bool,
4649        logits: CudaSlice<f32>,
4650        b_n: usize,
4651        ph_last: &mut std::time::Instant,
4652        pending_out: Option<&mut Option<PendingBatchStep>>,
4653    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
4654        // Grammar masks and penalties both mutate the sampling copy. Preserve each affected
4655        // row's PRISTINE logits first: continuation/reuse consumers must never inherit a mask
4656        // or get penalized twice after restore.
4657        let n_vocab = self.output.out_features();
4658        let mut logits = logits;
4659        let mut pristine: Vec<Option<CudaSlice<f32>>> = Vec::new();
4660        let row_mutates = |bi: usize| {
4661            masks.get(bi).is_some_and(Option::is_some)
4662                || samp
4663                    .get(bi)
4664                    .and_then(Option::as_ref)
4665                    .is_some_and(|s| s.penalty.is_some())
4666        };
4667        if (0..b_n).any(row_mutates) {
4668            pristine.resize_with(b_n, || None);
4669            for bi in 0..b_n {
4670                if !row_mutates(bi) {
4671                    continue;
4672                }
4673                if lean {
4674                    let cache = &mut caches[bi];
4675                    if cache
4676                        .last_logits_dev
4677                        .as_ref()
4678                        .map(|d| d.len() < n_vocab)
4679                        .unwrap_or(true)
4680                    {
4681                        cache.last_logits_dev = Some(e.uninit(n_vocab)?);
4682                    }
4683                    let dst = cache.last_logits_dev.as_mut().unwrap();
4684                    e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
4685                } else {
4686                    let mut p = e.uninit(n_vocab)?;
4687                    e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), &mut p)?;
4688                    pristine[bi] = Some(p);
4689                }
4690            }
4691        }
4692
4693        // Penalties precede grammar and probability filters, matching the host sampler chain.
4694        // Flatten only unique sparse counts for affected rows; heterogeneous requests keep
4695        // independent windows and coefficients in one launch.
4696        let penalized: Vec<(usize, &DevPenalty)> = samp
4697            .iter()
4698            .take(b_n)
4699            .enumerate()
4700            .filter_map(|(bi, s)| s.as_ref()?.penalty.as_ref().map(|p| (bi, p)))
4701            .filter(|(_, p)| !p.counts.is_empty())
4702            .collect();
4703        if !penalized.is_empty() {
4704            static ONCE: std::sync::Once = std::sync::Once::new();
4705            ONCE.call_once(|| {
4706                let unique: usize = penalized.iter().map(|(_, p)| p.counts.len()).sum();
4707                eprintln!(
4708                    "[device-penalty] sparse sampled rows={} unique-counts={} \
4709                     execution=one-ragged-launch raw-logits=preserved",
4710                    penalized.len(),
4711                    unique,
4712                );
4713            });
4714            let mut ids = Vec::new();
4715            let mut counts = Vec::new();
4716            let mut offsets = Vec::with_capacity(penalized.len() + 1);
4717            let mut rows = Vec::with_capacity(penalized.len());
4718            let mut reps = Vec::with_capacity(penalized.len());
4719            let mut freqs = Vec::with_capacity(penalized.len());
4720            let mut presents = Vec::with_capacity(penalized.len());
4721            offsets.push(0i32);
4722            for (bi, p) in penalized {
4723                rows.push(bi as i32);
4724                reps.push(p.repeat);
4725                freqs.push(p.freq);
4726                presents.push(p.present);
4727                for &(id, count) in &p.counts {
4728                    ids.push(id);
4729                    counts.push(count);
4730                }
4731                offsets.push(ids.len() as i32);
4732            }
4733            // SAFETY: rows come from `enumerate()` over this batch; DevPenalty's opaque count
4734            // set guarantees unique ids; and offsets are appended from the flattened vectors.
4735            unsafe {
4736                e.penalize_logits_sparse_rows_unchecked(
4737                    &mut logits,
4738                    &ids,
4739                    &counts,
4740                    &offsets,
4741                    &rows,
4742                    &reps,
4743                    &freqs,
4744                    &presents,
4745                    n_vocab,
4746                )?;
4747            }
4748        }
4749
4750        // GRAMMAR MASKS (constrained decoding): ban in place AFTER penalties and before the
4751        // device sampler. Penalized constrained rows remain on the host until their combined
4752        // composition gate exists, but keep the ordering correct as defense in depth.
4753        for (bi, m) in masks.iter().take(b_n).enumerate() {
4754            if let Some((mask, words)) = m {
4755                assert!(
4756                    samp.get(bi).and_then(Option::as_ref).is_some(),
4757                    "grammar-masked row {bi} must request a device sample"
4758                );
4759                e.mask_logits_col(&mut logits, mask, bi, n_vocab, *words)?;
4760            }
4761        }
4762
4763        // Device-side sampling for requested rows (see the method doc). Enqueued before the
4764        // big logits D2H so the tiny [B] token readback rides the same sync.
4765        let pending = pending_out.is_some();
4766        let mut next: Vec<Option<u32>> = vec![None; b_n];
4767        let mut device_tokens: Option<CudaSlice<u32>> = None;
4768        if samp.iter().take(b_n).any(|s| s.is_some()) {
4769            let mut toks = e.alloc_u32_zeroed(b_n)?;
4770            let mut perturb: Option<CudaSlice<f32>> = None;
4771            // FILTERED rows batch their filter_stats (lane/moebatch-q35moe): the per-row
4772            // devsample_filtered_col shape paid 1 HtoD + 3 tiny allocs + a 1-block launch PER
4773            // ROW PER TICK, serializing B single-SM kernels on the stream — measured as the
4774            // whole filtered-vs-temp-only serve gap at c8 (487 vs 700+ agg tok/s). Group rows
4775            // by (temp, top_k, top_p, min_p) — filter_stats takes scalar knobs — and solve
4776            // each group's thresholds in ONE grid=F launch over shared stat buffers, then
4777            // per-row perturb+argmax read their stat slot. Same kernels, same expressions,
4778            // same per-row (seed, ctr) draw — only the launch/alloc shape changes.
4779            let filt: Vec<(usize, &DevSamp)> = samp
4780                .iter()
4781                .take(b_n)
4782                .enumerate()
4783                .filter_map(|(bi, s)| s.as_ref().map(|s| (bi, s)))
4784                .filter(|(_, s)| s.temp > 0.0 && (s.top_k > 0 || s.top_p < 1.0 || s.min_p > 0.0))
4785                .collect();
4786            // Per-group stat buffers (one filter_stats launch per distinct knob tuple —
4787            // usually exactly one group per tick). Z is computed for output-shape parity
4788            // with the per-row form; the draw itself reads th/max only.
4789            let mut group_stats: Vec<(CudaSlice<f32>, CudaSlice<f32>)> = Vec::new();
4790            let mut row_stat: Vec<Option<(usize, usize)>> = vec![None; b_n];
4791            if !filt.is_empty() {
4792                #[allow(clippy::type_complexity)]
4793                // allow: one-shot composite type; naming it would hide the shape that matters at the call site
4794                let mut groups: Vec<((f32, i32, f32, f32), Vec<usize>)> = Vec::new();
4795                for &(bi, s) in &filt {
4796                    let key = (s.temp, s.top_k, s.top_p, s.min_p);
4797                    match groups.iter_mut().find(|(k, _)| *k == key) {
4798                        Some((_, rows)) => rows.push(bi),
4799                        None => groups.push((key, vec![bi])),
4800                    }
4801                }
4802                for ((temp, top_k, top_p, min_p), rows) in &groups {
4803                    let rows_i32: Vec<i32> = rows.iter().map(|&bi| bi as i32).collect();
4804                    let rows_d = e.htod_i32(&rows_i32)?;
4805                    let mut th = e.zeros(rows.len())?;
4806                    let mut z = e.zeros(rows.len())?;
4807                    let mut mx = e.zeros(rows.len())?;
4808                    e.filter_stats(
4809                        &logits,
4810                        n_vocab,
4811                        &rows_d,
4812                        &mut th,
4813                        &mut z,
4814                        &mut mx,
4815                        n_vocab,
4816                        rows.len(),
4817                        *temp,
4818                        *top_k,
4819                        *top_p,
4820                        *min_p,
4821                    )?;
4822                    let g = group_stats.len();
4823                    for (i, &bi) in rows.iter().enumerate() {
4824                        row_stat[bi] = Some((g, i));
4825                    }
4826                    group_stats.push((th, mx));
4827                }
4828            }
4829            for (bi, s) in samp.iter().take(b_n).enumerate() {
4830                let Some(s) = s else {
4831                    continue;
4832                };
4833                let filtered = s.temp > 0.0 && (s.top_k > 0 || s.top_p < 1.0 || s.min_p > 0.0);
4834                if s.temp <= 0.0 {
4835                    e.argmax_token_device_col(&logits, bi, n_vocab, &mut toks, bi)?;
4836                } else if filtered {
4837                    if perturb.is_none() {
4838                        perturb = Some(e.zeros(n_vocab)?);
4839                    }
4840                    let pb = perturb.as_mut().unwrap();
4841                    let (g, i) = row_stat[bi].expect("filtered row missing batched stats");
4842                    let (th, mx) = &group_stats[g];
4843                    e.gumbel_perturb_filtered_col(
4844                        &logits, bi, pb, n_vocab, s.seed, s.ctr, s.temp, mx, th, i,
4845                    )?;
4846                    e.argmax_token_device_col(pb, 0, n_vocab, &mut toks, bi)?;
4847                } else {
4848                    if perturb.is_none() {
4849                        perturb = Some(e.zeros(n_vocab)?);
4850                    }
4851                    let pb = perturb.as_mut().unwrap();
4852                    e.gumbel_perturb_col(&logits, bi, pb, n_vocab, s.seed, s.ctr, s.temp)?;
4853                    e.argmax_token_device_col(pb, 0, n_vocab, &mut toks, bi)?;
4854                }
4855            }
4856            if !pending {
4857                let host_toks = e.dtoh_u32(&toks)?;
4858                for (bi, s) in samp.iter().take(b_n).enumerate() {
4859                    if s.is_some() {
4860                        next[bi] = Some(host_toks[bi]);
4861                    }
4862                }
4863            }
4864            device_tokens = Some(toks);
4865        }
4866
4867        if let Some(slot) = pending_out {
4868            for c in caches.iter_mut() {
4869                c.pos += 1;
4870            }
4871            ph_mark(e, 11, ph_last)?;
4872            let done = e.stream().record_event(None)?;
4873            *slot = Some(PendingBatchStep::new(
4874                logits,
4875                pristine,
4876                device_tokens,
4877                samp.iter().take(b_n).map(Option::is_some).collect(),
4878                n_vocab,
4879                lean,
4880                done,
4881                e.copy_stream.clone(),
4882            ));
4883            return Ok((Vec::new(), vec![None; b_n]));
4884        }
4885
4886        let lean_any = lean && samp.iter().take(b_n).any(|s| s.is_some());
4887        let rows: Vec<Vec<f32>> = if lean_any {
4888            // LEAN: park device-sampled rows on-device (per-cache buffer, dtod); D2H only
4889            // the rows that still need host logits. No sampled rows + no fallback rows =
4890            // the big D2H disappears (the [B] token readback above already synced).
4891            for (bi, s) in samp.iter().take(b_n).enumerate() {
4892                if s.is_none() {
4893                    continue;
4894                }
4895                // Mutated rows already parked their PRISTINE copy above — neither a grammar
4896                // ban nor a penalty may poison the reuse-pool consumer.
4897                if masks.get(bi).copied().flatten().is_some()
4898                    || s.as_ref().is_some_and(|s| s.penalty.is_some())
4899                {
4900                    continue;
4901                }
4902                let cache = &mut caches[bi];
4903                if cache
4904                    .last_logits_dev
4905                    .as_ref()
4906                    .map(|d| d.len() < n_vocab)
4907                    .unwrap_or(true)
4908                {
4909                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
4910                }
4911                let dst = cache.last_logits_dev.as_mut().unwrap();
4912                e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
4913            }
4914            (0..b_n)
4915                .map(|bi| {
4916                    if samp.get(bi).and_then(Option::as_ref).is_some() {
4917                        Ok(Vec::new())
4918                    } else {
4919                        e.dtoh_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab))
4920                    }
4921                })
4922                .collect::<Result<_, _>>()?
4923        } else {
4924            let host = e.dtoh(&logits)?;
4925            (0..b_n)
4926                .map(|bi| {
4927                    // grammar-masked non-lean rows return the PRISTINE copy (the in-place ban
4928                    // must never leak into last_logits — reuse-pool/park semantics unchanged).
4929                    if let Some(p) = pristine.get(bi).and_then(|p| p.as_ref()) {
4930                        return e.dtoh(p);
4931                    }
4932                    Ok(host[bi * n_vocab..(bi + 1) * n_vocab].to_vec())
4933                })
4934                .collect::<Result<_, _>>()?
4935        };
4936        for c in caches.iter_mut() {
4937            c.pos += 1;
4938        }
4939        ph_mark(e, 11, ph_last)?;
4940        Ok((rows, next))
4941    }
4942}
4943
4944fn b1_fast_plan_eligible(plan: &memra_gguf::model_plan::ModelPlan) -> bool {
4945    // Every GDN plan is excluded: spec verify for this recurrent operation runs
4946    // the generic batched numeric class (spec.rs batched_serving_numeric_class), so live B=1 serving
4947    // must stay in that same class. B1FAST's eager program would reopen the near-tie-flip
4948    // divergence the 2026-08-14 exactness fix closed (1 ULP at layer 2 -> 2.3e-1 head
4949    // maxdiff, amplified by the GDN recurrence).
4950    !plan
4951        .trunk_operations()
4952        .contains(&memra_gguf::model_plan::OperationKind::GatedDeltaNet)
4953}
4954
4955fn b1_fast_env_on(value: Option<&str>) -> bool {
4956    value == Some("1")
4957}
4958
4959#[cfg(test)]
4960mod tests {
4961    use super::{
4962        PpWaveIncoming, PpWaveOutgoing, b1_fast_env_on, b1_fast_plan_eligible, pp_wave_channels,
4963    };
4964    use memra_gguf::config::{HfConfig, ModelConfig};
4965
4966    fn protocol_pair(boundary: usize) -> (PpWaveOutgoing, PpWaveIncoming) {
4967        let (mut outgoing, mut incoming) = pp_wave_channels(boundary + 1);
4968        (
4969            outgoing[boundary].take().unwrap(),
4970            incoming[boundary].take().unwrap(),
4971        )
4972    }
4973
4974    #[test]
4975    fn pp_wave_credit_requires_exact_ack_before_slot_reuse() {
4976        let (mut outgoing, incoming) = protocol_pair(0);
4977
4978        let expected0 = outgoing.prepare(0).unwrap();
4979        assert_eq!(expected0, None);
4980        outgoing.publish(0, 1, expected0).unwrap();
4981        let transfer0 = incoming.receive(0).unwrap();
4982
4983        let expected1 = outgoing.prepare(1).unwrap();
4984        assert_eq!(expected1, Some(0));
4985        outgoing.publish(1, 0, expected1).unwrap();
4986        let transfer1 = incoming.receive(1).unwrap();
4987
4988        // Wave 2 wants slot 1 again. Credit arrives only through the exact wave-0/slot-1
4989        // acknowledgement that a real consumer sends after rt.rx records ev_rx.
4990        incoming.acknowledge(transfer0).unwrap();
4991        let expected2 = outgoing.prepare(2).unwrap();
4992        assert_eq!(expected2, Some(1));
4993        outgoing.publish(2, 1, expected2).unwrap();
4994        let transfer2 = incoming.receive(2).unwrap();
4995
4996        incoming.acknowledge(transfer1).unwrap();
4997        incoming.acknowledge(transfer2).unwrap();
4998        outgoing.finish().unwrap();
4999        assert!(outgoing.pending.is_empty());
5000        assert_eq!(outgoing.slot_owner, [None, None]);
5001    }
5002
5003    #[test]
5004    fn pp_wave_protocol_rejects_order_and_propagates_worker_error() {
5005        let (mut outgoing, incoming) = protocol_pair(0);
5006        let expected = outgoing.prepare(0).unwrap();
5007        outgoing.publish(0, 0, expected).unwrap();
5008        let order_error = incoming.receive(1).unwrap_err();
5009        assert!(order_error.contains("expected wave 1"), "{order_error}");
5010
5011        let (outgoing, incoming) = protocol_pair(1);
5012        outgoing.publish_worker_error("injected stage failure");
5013        let worker_error = incoming.receive(0).unwrap_err();
5014        assert!(
5015            worker_error.contains("injected stage failure"),
5016            "{worker_error}"
5017        );
5018        assert!(worker_error.contains("boundary 1"), "{worker_error}");
5019        assert!(worker_error.contains("wave 0"), "{worker_error}");
5020    }
5021
5022    #[test]
5023    fn pp_wave_credit_rejects_wrong_ack_and_slot_generation() {
5024        let (mut outgoing, incoming) = protocol_pair(0);
5025        let expected0 = outgoing.prepare(0).unwrap();
5026        outgoing.publish(0, 0, expected0).unwrap();
5027        let _transfer0 = incoming.receive(0).unwrap();
5028        let expected1 = outgoing.prepare(1).unwrap();
5029        assert_eq!(expected1, Some(1));
5030        let wrong_slot = outgoing.publish(1, 0, expected1).unwrap_err();
5031        assert!(
5032            wrong_slot.contains("broke slot alternation"),
5033            "{wrong_slot}"
5034        );
5035
5036        // Rebuild after the rejected TX and inject an acknowledgement for wave 1 before wave 0.
5037        let (mut outgoing, incoming) = protocol_pair(0);
5038        let expected0 = outgoing.prepare(0).unwrap();
5039        outgoing.publish(0, 0, expected0).unwrap();
5040        let transfer0 = incoming.receive(0).unwrap();
5041        let expected1 = outgoing.prepare(1).unwrap();
5042        outgoing.publish(1, 1, expected1).unwrap();
5043        let transfer1 = incoming.receive(1).unwrap();
5044        incoming.acknowledgements.send(transfer1).unwrap();
5045        let wrong_ack = outgoing.prepare(2).unwrap_err();
5046        assert!(wrong_ack.contains("expected acknowledgement wave 0 slot 0"));
5047        assert!(wrong_ack.contains("got boundary 0 wave 1 slot 1"));
5048
5049        // Keep the compiler honest that the expected transfer really was the earlier one.
5050        assert_eq!(transfer0.wave, 0);
5051    }
5052
5053    #[test]
5054    fn pp_wave_protocol_reports_forward_and_ack_channel_closure() {
5055        let (outgoing, incoming) = protocol_pair(0);
5056        drop(outgoing);
5057        let forward_closed = incoming.receive(0).unwrap_err();
5058        assert!(forward_closed.contains("transfer channel closed"));
5059
5060        let (mut outgoing, incoming) = protocol_pair(0);
5061        let expected0 = outgoing.prepare(0).unwrap();
5062        outgoing.publish(0, 0, expected0).unwrap();
5063        let _ = incoming.receive(0).unwrap();
5064        let expected1 = outgoing.prepare(1).unwrap();
5065        outgoing.publish(1, 1, expected1).unwrap();
5066        let _ = incoming.receive(1).unwrap();
5067        drop(incoming);
5068        let ack_closed = outgoing.prepare(2).unwrap_err();
5069        assert!(ack_closed.contains("acknowledgement channel closed"));
5070
5071        let (mut outgoing, incoming) = protocol_pair(0);
5072        drop(incoming);
5073        let expected = outgoing.prepare(0).unwrap();
5074        let publish_closed = outgoing.publish(0, 0, expected).unwrap_err();
5075        assert!(publish_closed.contains("transfer channel closed"));
5076    }
5077
5078    #[test]
5079    fn gdn_plans_stay_in_one_decode_numeric_class_across_widths() {
5080        let compile = |json| {
5081            memra_gguf::model_plan::ModelPlan::compile(&ModelConfig::from_hf(&HfConfig::parse(
5082                json,
5083            )))
5084            .unwrap()
5085        };
5086        let gdn = compile(
5087            r#"{"model_type":"qwen3_5","num_hidden_layers":2,"hidden_size":64,
5088            "num_attention_heads":2,"num_key_value_heads":1,"head_dim":32,
5089            "intermediate_size":128,"vocab_size":16,"max_position_embeddings":128,
5090            "full_attention_interval":2,"linear_conv_kernel_dim":3,
5091            "linear_key_head_dim":32,"linear_value_head_dim":32,
5092            "linear_num_key_heads":1,"linear_num_value_heads":2}"#,
5093        );
5094        let full = compile(
5095            r#"{"model_type":"qwen3","num_hidden_layers":1,"hidden_size":64,
5096            "num_attention_heads":2,"num_key_value_heads":1,"head_dim":32,
5097            "intermediate_size":128,"vocab_size":16,"max_position_embeddings":128}"#,
5098        );
5099        assert!(!b1_fast_plan_eligible(&gdn));
5100        assert!(b1_fast_plan_eligible(&full));
5101    }
5102
5103    #[test]
5104    fn b1_eager_program_requires_explicit_opt_in() {
5105        assert!(!b1_fast_env_on(None));
5106        assert!(!b1_fast_env_on(Some("0")));
5107        assert!(!b1_fast_env_on(Some("true")));
5108        assert!(b1_fast_env_on(Some("1")));
5109    }
5110}