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::cache::Cache;
30use crate::hybrid::{HybridModel, Mixer};
31use crate::Engine;
32use cudarc::driver::CudaSlice;
33
34// ---- MEMRA_BATCH_PHASE=1 (diagnostics): sync-bounded per-phase accumulators for the batched
35// tick. Each boundary syncs the stream, so the TOTAL inflates (launch pipelining is destroyed);
36// the value is the RANKING/shares, not absolute ms. Read via `batch_phase_report()`.
37pub(crate) static BATCH_PHASE: std::sync::Mutex<[f64; 12]> = std::sync::Mutex::new([0.0; 12]);
38pub const BATCH_PHASE_NAMES: [&str; 12] = [
39    "setup(ptrs+embed H2D)",
40    "attn batched pre (norm/qkv/rope)",
41    "attn per-seq: kv append",
42    "attn per-seq: q/a dtod copies",
43    "attn per-seq: fa_decode",
44    "attn post (gate+o-proj)",
45    "gdn batched projections",
46    "gdn state ops (conv/prep/scan)",
47    "gdn out (gated norm+proj)",
48    "ffn (add/norm/gate/up/act/down)",
49    "lm_head (norm+matmul)",
50    "logits D2H + host split",
51];
52pub fn batch_phase_on() -> bool {
53    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
54    *ON.get_or_init(|| std::env::var("MEMRA_BATCH_PHASE").as_deref() == Ok("1"))
55}
56pub fn batch_phase_report() -> String {
57    let ph = BATCH_PHASE.lock().unwrap();
58    let tot: f64 = ph.iter().sum();
59    let mut rows: Vec<(usize, f64)> = ph.iter().copied().enumerate().collect();
60    rows.sort_by(|a, b| b.1.total_cmp(&a.1));
61    let mut s = format!("[batch-phase] total {:.1} ms (sync-bounded; shares rank, not walltime)\n", tot * 1e3);
62    for (i, v) in rows {
63        s += &format!("  {:>6.1} ms {:>5.1}%  {}\n", v * 1e3, v / tot * 100.0, BATCH_PHASE_NAMES[i]);
64    }
65    s
66}
67
68impl HybridModel {
69    /// Batched-decode width cap. 8 = the exactness-tier default (see the assert below);
70    /// MEMRA_DECODE_BATCH_CAP overrides for tier-probe measurement, clamped to 32.
71    pub fn decode_batch_cap() -> usize {
72        use std::sync::OnceLock;
73        static CAP: OnceLock<usize> = OnceLock::new();
74        *CAP.get_or_init(|| {
75            std::env::var("MEMRA_DECODE_BATCH_CAP").ok()
76                .and_then(|v| v.parse().ok())
77                .map(|c: usize| c.clamp(1, 32))
78                .unwrap_or(8)
79        })
80    }
81
82    /// EXACT-16 TIER admission (increment 3a, 2026-08-01, 5090 receipts
83    /// research/batched-tick-inc3-20260801): true iff EVERY matmul the batched decode step
84    /// runs has a per-(token,row) bit-exact kernel class at m=9..16 under the verify_exact
85    /// scope — i.e. the batched-mmvq b16 family (32-thread warp reduce, the exact m=1 mmvq
86    /// program per column) or the e4m3 grid.y=m mmvq catch-all. Q8_0 qualifies only with
87    /// the split-plane mirror (rp4, MEMRA_Q8RP): its b16 kernel exists only as the _rp twin.
88    /// Float matmuls (cuBLASLt, n-dependent reductions) and MoE FFNs disqualify the model.
89    /// Measured attribution for WHY the naked m=16 tier is not exact: the m>=16 arms
90    /// (MMQ int8-MMA `mul_mat_q` — MEMRA_PP_Q8MMQ default-on — and `qmatvec_gemm`, both
91    /// block-scale f32) and the m=9..15 dp4a tail (128-thread two-level reduce) all break
92    /// per-row bit-identity vs isolated decode (gate2 step-0 bit-diffs, maxdiff ~1.3-2.3e-1).
93    pub fn decode_batch_exact16_ok(&self) -> bool {
94        fn ok(w: &crate::model::GpuTensor) -> bool {
95            match w {
96                crate::model::GpuTensor::Quant { qtype, rp4, .. } =>
97                    *qtype == crate::QT_Q4_0 || *qtype == crate::QT_Q6_K
98                    || *qtype == crate::QT_F8_E4M3
99                    || (*qtype == crate::QT_Q8_0 && rp4.is_some()),
100                _ => false,
101            }
102        }
103        if self.cfg.m3.is_some() || self.is_gemma4_e4b() || self.cfg.gemma4.is_some() {
104            return false;
105        }
106        self.layers.iter().all(|l| {
107            let mix_ok = match &l.mixer {
108                Mixer::Full(fa) => [&fa.wq, &fa.wk, &fa.wv, &fa.wo].into_iter().all(ok),
109                Mixer::Linear(la) => [&la.wqkv, &la.wqkv_gate, &la.ssm_beta,
110                                      &la.ssm_alpha, &la.ssm_out].into_iter().all(ok),
111                // MLA rides its own increment-4 arm; never admitted to the exact-16 tier here.
112                Mixer::Mla(_) => false,
113            };
114            let ffn_ok = match &l.ffn {
115                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } =>
116                    [ffn_gate, ffn_up, ffn_down].into_iter().all(ok),
117                crate::hybrid::Ffn::Moe(_) => false,
118            };
119            mix_ok && ffn_ok
120        }) && ok(&self.output)
121    }
122
123    /// H3 rollback/A-B seam (serve-path phase 2): `MEMRA_SERVE_B1FAST=0` sends B=1 back
124    /// through the batched body (the pre-change tick, bit-for-bit). Default ON.
125    ///
126    /// EXACTNESS, stated precisely (measured on-box 2026-08-05, sm_120 q9 NVFP4-MTP):
127    /// the fast path is BIT-IDENTICAL TO `decode_step_h` — decode-batch-gate's STRICT
128    /// gate1 (`--mode strict`) PASSes with it ON and FAILs with it OFF at maxdiff
129    /// 1.591e-1. It is deliberately NOT bit-identical to the batched body: the two
130    /// carry the long-accepted decode-config FP-composition gap (same class gate1's
131    /// config mode tolerates), and this lever moves solo sessions onto the NAKED side
132    /// of it. That is the desired direction — a c=1 serve request now computes exactly
133    /// what `run-gen` computes for the same prompt. Token-stream receipts:
134    /// research/servepath-p2-20260805 (greedy 150 ids + seeded-sampled identical to the
135    /// run-gen oracle AND cross-arm, so the gap is sub-token here as designed).
136    ///
137    /// Read fresh (an `AtomicU8` memo, not a `OnceLock`): decode-batch-gate flips this
138    /// seam BETWEEN gates in-process — gate1 needs the fast path ON to prove bit-identity,
139    /// gate2 needs it pinned OFF to keep testing the batched body. A latch-once read would
140    /// bake whichever gate ran first, so the gate could never test both sides. The memo
141    /// caches the parse but `set_b1_fast` invalidates it.
142    pub fn b1_fast_on() -> bool {
143        // 0 = unknown/invalidated, 1 = off, 2 = on
144        match Self::b1_fast_memo().load(std::sync::atomic::Ordering::Relaxed) {
145            1 => false,
146            2 => true,
147            _ => {
148                let on = std::env::var("MEMRA_SERVE_B1FAST").as_deref() != Ok("0");
149                Self::b1_fast_memo()
150                    .store(if on { 2 } else { 1 }, std::sync::atomic::Ordering::Relaxed);
151                on
152            }
153        }
154    }
155
156    fn b1_fast_memo() -> &'static std::sync::atomic::AtomicU8 {
157        static MEMO: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
158        &MEMO
159    }
160
161    /// Test/gate seam: force the B=1 fast path on or off for the rest of the process,
162    /// overriding the env. Used by decode-batch-gate to pin gate2's reference arm.
163    pub fn set_b1_fast(on: bool) {
164        Self::b1_fast_memo()
165            .store(if on { 2 } else { 1 }, std::sync::atomic::Ordering::Relaxed);
166    }
167
168    /// H3 body: the m=1 FUSED trunk (`decode_layers_eager` — shared verbatim with
169    /// `decode_step_h`/the ppN stages) plus the batched path's own serving epilogue
170    /// (grammar mask, device sample, lean-logits park). See the call-site comment in
171    /// `decode_step_batch_sampled_lean_masked` for why this is bit-identical.
172    fn decode_step_b1_fast(
173        &self,
174        e: &Engine,
175        token: u32,
176        caches: &mut [&mut Cache],
177        samp: &[Option<(f32, u64, u32)>],
178        masks: &[Option<(&CudaSlice<u32>, usize)>],
179        lean: bool,
180    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
181        let n_embd = self.cfg.n_embd as usize;
182        let eps = self.cfg.rms_eps;
183        let pos = caches[0].pos;
184        let pos_d = e.htod_i32(&[pos as i32])?;
185        let x = e.htod(&self.embd.gather(n_embd, &[token]))?;
186        // the SHARED m=1 trunk: same function decode_step_h runs, so every m=1 fusion
187        // (cross-layer add+norm+q8_1, fused SwiGLU, lever 1's gate+up dual) fires here.
188        let x = self.decode_layers_eager(e, x, 0, self.layers.len(), &pos_d, pos, caches[0])?;
189        let mut hn = e.uninit(n_embd)?;
190        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, 1, eps)?;
191        let logits = e.matmul(&self.output, &hn, 1)?;
192
193        // ---- epilogue: byte-for-byte the batched path's, at b_n=1 ----
194        let n_vocab = self.output.out_features();
195        let mut logits = logits;
196        let mut pristine: Option<CudaSlice<f32>> = None;
197        if let Some((mask, words)) = masks.first().copied().flatten() {
198            assert!(samp.first().copied().flatten().is_some(),
199                    "grammar-masked row 0 must request a device sample");
200            if lean {
201                let cache = &mut caches[0];
202                if cache.last_logits_dev.as_ref().map(|d| d.len() < n_vocab).unwrap_or(true) {
203                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
204                }
205                let dst = cache.last_logits_dev.as_mut().unwrap();
206                e.dtod_copy_view(&logits.slice(0..n_vocab), dst)?;
207            } else {
208                let mut p = e.uninit(n_vocab)?;
209                e.dtod_copy_view(&logits.slice(0..n_vocab), &mut p)?;
210                pristine = Some(p);
211            }
212            e.mask_logits_col(&mut logits, mask, 0, n_vocab, words)?;
213        }
214
215        let mut next: Vec<Option<u32>> = vec![None; 1];
216        if let Some((temp, seed, ctr)) = samp.first().copied().flatten() {
217            let mut toks = e.alloc_u32_zeroed(1)?;
218            if temp <= 0.0 {
219                e.argmax_token_device_col(&logits, 0, n_vocab, &mut toks, 0)?;
220            } else {
221                let mut pb = e.zeros(n_vocab)?;
222                e.gumbel_perturb_col(&logits, 0, &mut pb, n_vocab, seed, ctr, temp)?;
223                e.argmax_token_device_col(&pb, 0, n_vocab, &mut toks, 0)?;
224            }
225            next[0] = Some(e.dtoh_u32(&toks)?[0]);
226        }
227
228        let sampled = samp.first().copied().flatten().is_some();
229        let rows: Vec<Vec<f32>> = if lean && sampled {
230            if masks.first().copied().flatten().is_none() {
231                let cache = &mut caches[0];
232                if cache.last_logits_dev.as_ref().map(|d| d.len() < n_vocab).unwrap_or(true) {
233                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
234                }
235                let dst = cache.last_logits_dev.as_mut().unwrap();
236                e.dtod_copy_view(&logits.slice(0..n_vocab), dst)?;
237            }
238            vec![Vec::new()]
239        } else if let Some(p) = pristine.as_ref() {
240            vec![e.dtoh(p)?]
241        } else {
242            vec![e.dtoh(&logits)?]
243        };
244        // decode_layers_eager does NOT advance cache.pos (decode_step_h advances it after
245        // the head); the batched path advances every cache at the tail — same here.
246        caches[0].pos += 1;
247        Ok((rows, next))
248    }
249
250    /// One batched greedy-decode step over B independent sequences.
251    /// `tokens[b]` is sequence b's input token; `caches[b]` its private cache (position,
252    /// quantized KV, GDN/conv state). Returns the B logits rows (host, [n_vocab] each).
253    /// Each cache's pos/len advance exactly as `decode_step_h` would.
254    pub fn decode_step_batch(
255        &self,
256        e: &Engine,
257        tokens: &[u32],
258        caches: &mut [&mut Cache],
259    ) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
260        let (rows, _) = self.decode_step_batch_sampled(e, tokens, caches, &[])?;
261        Ok(rows)
262    }
263
264    /// `decode_step_batch` + DEVICE-SIDE SAMPLING for eligible rows (the batched-tick lever,
265    /// 2026-08-01): the host sampler's temp-path is O(n_vocab) with a full-vocab exp per row
266    /// (measured 1.36 ms/row at the 9B's 248320 vocab = 10.9 ms/tick at B=8 — the single
267    /// largest component of the serving tick). Here each requested row samples ON DEVICE
268    /// between the lm_head matmul and the logits D2H:
269    ///   temp <= 0 (greedy): the 2-pass device argmax — bit-identical to host argmax
270    ///     (argmax-gate contract, same kernels as the dc serving path).
271    ///   temp > 0: gumbel_perturb(seed, ctr, temp) + the same argmax = ONE categorical draw
272    ///     from softmax(logits/temp) — the sampled-spec Philox machinery. Deterministic per
273    ///     (seed, ctr) and INDEPENDENT of batch composition (the isolation contract;
274    ///     decode-batch-gate gate3). NOTE: the draw stream differs from the host sampler's
275    ///     SplitMix64 (distribution-equal, seed-deterministic, NOT byte-equal to the old
276    ///     host draws) — greedy rows are unchanged bit-exact.
277    /// `samp[bi] = Some((temp, seed, ctr))` requests a device sample for row bi; the full
278    /// logits rows are still returned (worker keeps last_logits semantics + fallback rows).
279    pub fn decode_step_batch_sampled(
280        &self,
281        e: &Engine,
282        tokens: &[u32],
283        caches: &mut [&mut Cache],
284        samp: &[Option<(f32, u64, u32)>],
285    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
286        self.decode_step_batch_sampled_lean(e, tokens, caches, samp, false)
287    }
288
289    /// `decode_step_batch_sampled` + LEAN LOGITS (increment 2 component 3, 2026-08-01):
290    /// with `lean`, device-sampled rows SKIP the [n_vocab] logits D2H (9.4%/32.5% of the
291    /// pre-/post-inc2 tick profile) — their returned row is EMPTY. The audit-mapped
292    /// consumers: (a) the next tick's host sample — never fires, `device_next` carries the
293    /// token; (b) the graph-promotion argmax — reads only prefill logits (generated empty);
294    /// (c) the KV-reuse pool park at retire — the REAL consumer, served by a per-cache
295    /// device park: the row is dtod-copied into `cache.last_logits_dev` (device bandwidth)
296    /// and D2H'd ONCE at retire by the worker. Rows without a device sample keep a per-row
297    /// D2H. `lean=false` is bit-for-bit the previous method (gates + non-serving callers).
298    pub fn decode_step_batch_sampled_lean(
299        &self,
300        e: &Engine,
301        tokens: &[u32],
302        caches: &mut [&mut Cache],
303        samp: &[Option<(f32, u64, u32)>],
304        lean: bool,
305    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
306        self.decode_step_batch_sampled_lean_masked(e, tokens, caches, samp, &[], lean)
307    }
308
309    /// `decode_step_batch_sampled_lean` + GRAMMAR MASKS (constrained decoding, 2026-08-03):
310    /// `masks[bi] = Some((packed_bitset, words))` bans every unset-bit vocab id on row bi
311    /// (mask_logits_f32, -FLT_MAX) BETWEEN the lm_head matmul and the device sampler, so a
312    /// constrained row rides the SAME device-sample/lean-logits tick as everyone else — no
313    /// full-row D2H, no host O(n_vocab) sample. Contract: a masked row must also request a
314    /// device sample. The row's PRISTINE logits are preserved for their consumers before the
315    /// in-place ban: lean rows park the unmasked row into `cache.last_logits_dev` (the
316    /// retire-time reuse-pool park stays unmasked — continuations resume grammar-free, the
317    /// v1 host-path contract), non-lean rows D2H the unmasked row. `masks = &[]` is
318    /// bit-for-bit the unmasked method.
319    pub fn decode_step_batch_sampled_lean_masked(
320        &self,
321        e: &Engine,
322        tokens: &[u32],
323        caches: &mut [&mut Cache],
324        samp: &[Option<(f32, u64, u32)>],
325        masks: &[Option<(&CudaSlice<u32>, usize)>],
326        lean: bool,
327    ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn std::error::Error>> {
328        // NOTE (inc3 3c, 2026-08-01, KILLED ARM): a deferred-token-readback variant (all
329        // chunks of a tick writing device-sampled tokens into one shared buffer, ONE
330        // dtoh_u32 after the last chunk instead of one per chunk) measured FLAT at serve
331        // level on the 5090 (N=4 medians within +-0.7% at c=8/16/32 — 3 saved syncs
332        // against a ~100 ms weight-bound tick is ~0.1%, below resolution). Killed per the
333        // flags doctrine; receipts research/batched-tick-inc3-20260801 (serve-points.jsonl
334        // base vs defer arms) are the record. The per-chunk [B]-u32 readback below IS the
335        // tick's only steady-state D2H — one per chunk, none per seq.
336        let b_n = tokens.len();
337        assert!(b_n >= 1 && b_n == caches.len(), "tokens/caches length mismatch");
338        // ---- H3: B=1 FAST-PATH (serve-path phase 2, 2026-08-05) ----------------------------
339        // At b_n==1 every projection below calls `matmul_pre(.., b_n)` with m=1, which is
340        // ALREADY the m=1 mmvq dispatch — so the m=1 *kernel family* was never the gap. What
341        // this body does NOT have is the m=1 *fusion chain* that `decode_step_h` carries:
342        //   - the cross-layer add+norm+quantize fusion (`add_rms_norm_q8_1`: 3 launches -> 1),
343        //   - the fused SwiGLU epilogue (`silu_mul_scaled_q8_1`: folds ffn_down's quantize
344        //     into its producer) and, with it, `matmul_pre_dual_noscale`'s gate+up pair
345        //     fusion — i.e. phase-1 LEVER 1.
346        // Routing b_n==1 through `decode_layers_eager` (the SHARED trunk `decode_step_h` and
347        // the ppN stages already use, lifted verbatim — not a copy) makes every present and
348        // future m=1 lever fire on the serve path automatically, which is the durable half of
349        // this change. The epilogue (grammar mask -> device sample -> lean logits park) is
350        // kept EXACTLY as the batched path runs it, so the serving contract is untouched.
351        // BIT-IDENTITY: the trunk is the same function `decode_step_h` calls, and every
352        // fusion it enables is kernel-check-pinned bit-identical to its unfused sequence
353        // (add_rms_norm == add;rms_norm | _q8_1 == +quantize_q8_1 | dual_noscale == two
354        // matmul_pre_noscale). Gate: decode-batch-gate B=1 vs decode_step_h + serve stream
355        // identity. MEMRA_SERVE_B1FAST=0 is the rollback/A-B seam.
356        if b_n == 1
357            && Self::b1_fast_on()
358            && !self.is_gemma4_e4b()
359            && self.cfg.gemma4.is_none()
360            && self.cfg.m3.is_none()
361            && crate::pp::pp_cuts(self.layers.len()).is_none()
362            && !e.verify_exact_on()
363        {
364            return self.decode_step_b1_fast(e, tokens[0], caches, samp, masks, lean);
365        }
366        // MEMRA_DECODE_BATCH_CAP (experimental door, serving-lane tier probe 2026-08-01):
367        // default 8 keeps the v1 exactness policy — B=2..8 rides the verify-tier batched
368        // mmvq arms, per-row bit-identical to isolated m=1 decode. Values >8 are a
369        // MEASUREMENT DOOR ONLY: m=9..15 falls to the grid.y=m dp4a tail (m weight
370        // re-reads + a different reduce shape) and m>=16 crosses into the GEMM tier
371        // (block-scale f32 rounding) — BOTH break the "byte-identical to isolated"
372        // serving contract. Never default this above 8 without the batched-tier
373        // exactness policy landing.
374        let cap = Self::decode_batch_cap();
375        // EXACT-16 TIER (increment 3a): chunks of 9..=16 are admitted WITHOUT the env door
376        // when every matmul has a bit-exact b16-class kernel (see decode_batch_exact16_ok).
377        // The verify_exact scope below pins that dispatch for the whole step: it turns off
378        // the m>=16 GEMM arms (qmatvec_gemm + MMQ + fp8/f16/fp4 — all block-scale/foreign
379        // numeric configs) so every projection rides the batched-mmvq b16 tier, which is
380        // per-(token,row) bit-identical to isolated m=1 decode (gate2 bit-strength PASS at
381        // B=12/16, s32+s160, 5090 receipts research/batched-tick-inc3-20260801). Without
382        // the exact tier, B>cap stays refused; the env door (MEMRA_DECODE_BATCH_CAP) keeps
383        // its old meaning as the non-exact measurement probe.
384        let exact16 = b_n > 8 && b_n <= 16 && self.decode_batch_exact16_ok();
385        assert!(
386            b_n <= cap || exact16,
387            "decode_step_batch: B={b_n} > cap {cap} with no exact tier (Q8_0 m>8 needs the \
388             q8rp mirror's b16 class; m>16 crosses GEMM/dp4a numeric configs) — refused"
389        );
390        struct ExactScope<'a>(&'a Engine, bool);
391        impl Drop for ExactScope<'_> {
392            fn drop(&mut self) {
393                if self.1 {
394                    self.0.set_verify_exact(false);
395                }
396            }
397        }
398        let _exact_scope = ExactScope(e, exact16);
399        if exact16 {
400            e.set_verify_exact(true);
401        }
402        assert!(
403            !self.is_gemma4_e4b() && self.cfg.gemma4.is_none(),
404            "decode_step_batch v1 covers the hybrid non-gemma4 trunk only"
405        );
406        let cfg = &self.cfg;
407        let n_embd = cfg.n_embd as usize;
408        let eps = cfg.rms_eps;
409        let n_head = cfg.n_head as usize;
410        let n_head_kv = cfg.n_head_kv as usize;
411        let head_dim = cfg.head_dim_k as usize;
412        let scale = 1.0 / (head_dim as f32).sqrt();
413        let rope_dims = cfg.rope_dim_count as usize;
414
415        // Per-row rope positions (each sequence at its own depth).
416        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
417        let pos_d = e.htod_i32(&pos_v)?;
418
419        // Per-step STATE POINTER TABLE (one H2D): for every linear layer, [conv x B]
420        // [ssm_in x B][ssm_out x B] device addresses. The batched state kernels read their
421        // sequence's pointer from these arrays — states stay per-cache (no pooling refactor),
422        // yet conv/prep/scan collapse from 3xB launches per layer to 3. Rebuilt every step
423        // because the ssm ping-pong swaps pointers host-side after each scan.
424        // INCREMENT 2 (2026-08-01): the SAME table now also carries, for every FULL-attn
425        // layer, [k0,v0,k1,v1,...] cache base addresses — the z-batched seqs append and
426        // seqs fa_decode kernels read their sequence's cache through it (the MoE
427        // expert-table pattern), collapsing 2xB launches per attn layer to 2.
428        let mut lin_base: Vec<Option<usize>> = vec![None; self.layers.len()];
429        let mut attn_base: Vec<Option<usize>> = vec![None; self.layers.len()];
430        let mut ptrs: Vec<u64> = Vec::new();
431        {
432            use cudarc::driver::DevicePtr;
433            let s = &e.gpu.stream();
434            for (il, layer) in self.layers.iter().enumerate() {
435                match &layer.mixer {
436                    Mixer::Linear(_) => {
437                        lin_base[il] = Some(ptrs.len());
438                        for c in caches.iter() {
439                            let rl = c.recur[il].as_ref().unwrap();
440                            let (p, _g) = rl.conv_state.device_ptr(s);
441                            ptrs.push(p as u64);
442                        }
443                        for c in caches.iter() {
444                            let rl = c.recur[il].as_ref().unwrap();
445                            let (p, _g) = rl.ssm_state.device_ptr(s);
446                            ptrs.push(p as u64);
447                        }
448                        for c in caches.iter() {
449                            let rl = c.recur[il].as_ref().unwrap();
450                            let (p, _g) = rl.ssm_state_alt.device_ptr(s);
451                            ptrs.push(p as u64);
452                        }
453                    }
454                    Mixer::Full(_) => {
455                        attn_base[il] = Some(ptrs.len());
456                        for c in caches.iter() {
457                            let kvl = c.kv[il].as_ref().unwrap();
458                            let (pk, _g) = kvl.k.device_ptr(s);
459                            let (pv, _g2) = kvl.v.device_ptr(s);
460                            ptrs.push(pk as u64);
461                            ptrs.push(pv as u64);
462                        }
463                    }
464                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
465                }
466            }
467        }
468        let ptr_table = if ptrs.is_empty() { None } else { Some(e.htod_u64(&ptrs)?) };
469
470        // INCREMENT 2 arm picks (per STEP — t_kv is layer-invariant within a tick):
471        // - seqs APPEND: format-only condition (per-row program is t_kv-independent);
472        //   default flash module only (fp8-KV rides the per-seq g-module path).
473        // - seqs FA: every row must take the v4 eager arm at ITS OWN t_kv AND all rows
474        //   must share ONE fa_split_keys rung (the rows-twins' straddle law) — a rung
475        //   crossing inside the batch keeps the per-seq loop for that step, so each
476        //   sequence always executes the exact program its isolated run would.
477        // MEMRA_BATCH_APPEND=0 / MEMRA_BATCH_FA=0 are the rollback/A-B seams.
478        let t_kvs: Vec<usize> = caches.iter().map(|c| c.pos + 1).collect();
479        let t_kv_max = *t_kvs.iter().max().unwrap();
480        let seqs_append = {
481            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
482            *ON.get_or_init(|| std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0"))
483        } && !Engine::kv_fp8_on();
484        let sp0 = crate::fa_split_keys(t_kvs[0], cfg.n_head_kv as usize);
485        let seqs_fa = {
486            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
487            *ON.get_or_init(|| std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0"))
488        } && t_kvs.iter().all(|&t| crate::fa_seqs_eligible(t, head_dim))
489          && t_kvs.iter().all(|&t| crate::fa_split_keys(t, cfg.n_head_kv as usize) == sp0);
490
491        // Embed all B tokens -> x [B, n_embd] (host gather, one H2D).
492        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
493
494        // MEMRA_BATCH_PHASE=1: sync-bounded phase accumulation (diagnostics — see header note).
495        let ph_on = batch_phase_on();
496        let mut ph_last = std::time::Instant::now();
497        let ph_mark = |slot: usize,
498                       last: &mut std::time::Instant|
499         -> Result<(), Box<dyn std::error::Error>> {
500            if ph_on {
501                e.stream().synchronize()?;
502                let now = std::time::Instant::now();
503                BATCH_PHASE.lock().unwrap()[slot] += (now - *last).as_secs_f64();
504                *last = now;
505            }
506            Ok(())
507        };
508        ph_mark(0, &mut ph_last)?;
509
510        for (il, layer) in self.layers.iter().enumerate() {
511            // ---- attn_norm + q8_1 quantize, batched (B rows) ----
512            let anorm = layer.attn_norm.float_data();
513            let mut xn = e.uninit(b_n * n_embd)?;
514            e.rms_norm(&x, anorm, &mut xn, n_embd, b_n, eps)?;
515            let (hq, hd) = e.quantize_q8_1(&xn, b_n, n_embd)?;
516
517            // ---- mixer ----
518            let mixed: CudaSlice<f32> = match &layer.mixer {
519                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
520                Mixer::Full(fa) => {
521                    // Batched projections: one weight read serves all B rows.
522                    let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, b_n)?;
523                    let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, b_n)?;
524                    let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, b_n)?;
525
526                    let gated = cfg.attn_out_gate();
527                    let (mut q, gate) = if gated {
528                        let mut qs = e.uninit(b_n * n_head * head_dim)?;
529                        let mut gs = e.uninit(b_n * n_head * head_dim)?;
530                        e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, b_n)?;
531                        (qs, Some(gs))
532                    } else {
533                        (qf, None)
534                    };
535
536                    // QK-norm over B*n_head rows, rope with per-row positions.
537                    let mut qn = e.uninit(b_n * n_head * head_dim)?;
538                    e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, b_n * n_head, eps)?;
539                    q = qn;
540                    let mut kn = e.uninit(b_n * n_head_kv * head_dim)?;
541                    e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, b_n * n_head_kv, eps)?;
542                    k = kn;
543                    e.rope_neox(&mut q, &pos_d, head_dim, rope_dims, n_head, b_n,
544                                cfg.rope_freq_base, 1.0)?;
545                    e.rope_neox(&mut k, &pos_d, head_dim, rope_dims, n_head_kv, b_n,
546                                cfg.rope_freq_base, 1.0)?;
547                    ph_mark(1, &mut ph_last)?;
548
549                    // INCREMENT 2 (2026-08-01): the per-seq (append, attend) launch train
550                    // becomes two phases. Phase A appends all B rows (one z-batched launch,
551                    // or the per-seq loop on the seam/fp8 path); phase B attends all B
552                    // sequences (one blockIdx.z launch + one combine on the batched arm —
553                    // which also reads q / writes attn at row offsets, killing the per-seq
554                    // q/a dtod copies — or the per-seq loop when any row is outside the v4
555                    // arm / a split rung crosses inside the batch). Caches are disjoint per
556                    // sequence, so the phase split leaves every row's math untouched.
557                    let q_dim = n_head * head_dim;
558                    let kv_dim = n_head_kv * head_dim;
559                    let mut attn = e.uninit(b_n * q_dim)?;
560                    // ---- phase A: KV append (all B rows) ----
561                    if seqs_append {
562                        let (kdk, kdv, ktb, vtb) = {
563                            let kvl = caches[0].kv[il].as_ref().unwrap();
564                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
565                        };
566                        let base = attn_base[il].expect("full layer missing from pointer table");
567                        let table = ptr_table.as_ref().expect("pointer table missing");
568                        let kv_view = table.slice(base..base + 2 * b_n);
569                        e.append_kv_quantized_seqs(&k, &v, &kv_view, &pos_d, b_n,
570                                                   kdk, kdv, ktb, vtb)?;
571                        for cache in caches.iter_mut() {
572                            let kvl = cache.kv[il].as_mut().unwrap();
573                            debug_assert_eq!(kvl.len, cache.pos, "kv len / pos out of lockstep");
574                            kvl.len += 1;
575                        }
576                    } else {
577                        for (bi, cache) in caches.iter_mut().enumerate() {
578                            let kvl = cache.kv[il].as_mut().unwrap();
579                            let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
580                            let v_row = v.slice(bi * kv_dim..(bi + 1) * kv_dim);
581                            e.append_kv_quantized_view(
582                                &k_row, &v_row, &mut kvl.k, &mut kvl.v, kvl.len,
583                                kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
584                                Engine::kv_fp8_on(),
585                            )?;
586                            kvl.len += 1;
587                        }
588                    }
589                    ph_mark(2, &mut ph_last)?;
590                    // ---- phase B: attention (all B sequences) ----
591                    if seqs_fa {
592                        let (ktb, vtb) = {
593                            let kvl = caches[0].kv[il].as_ref().unwrap();
594                            (kvl.k_tok_bytes, kvl.v_tok_bytes)
595                        };
596                        let base = attn_base[il].expect("full layer missing from pointer table");
597                        let table = ptr_table.as_ref().expect("pointer table missing");
598                        let kv_view = table.slice(base..base + 2 * b_n);
599                        e.fa_decode_batch_seqs_v4(&q, &kv_view, &pos_d, &mut attn,
600                                                  head_dim, n_head, n_head_kv, b_n,
601                                                  t_kv_max, scale, sp0, ktb, vtb)?;
602                        ph_mark(4, &mut ph_last)?;
603                    } else {
604                        for (bi, cache) in caches.iter_mut().enumerate() {
605                            let kvl = cache.kv[il].as_mut().unwrap();
606                            let t_kv = kvl.len;
607                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
608                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
609                            // fa_decode wants a q slice starting at row bi: the fallback arm
610                            // scratch-copies the row (q8-class µs cost); the seqs arm above
611                            // reads/writes row offsets in place.
612                            let mut q_row = e.uninit(q_dim)?;
613                            e.dtod_copy_view(&q.slice(bi * q_dim..(bi + 1) * q_dim), &mut q_row)?;
614                            ph_mark(3, &mut ph_last)?;
615                            let mut a_row = e.uninit(q_dim)?;
616                            e.fa_decode_kvmod(
617                                &q_row, &k_view, &v_view, &mut a_row, head_dim, n_head, n_head_kv,
618                                t_kv, scale, kvl.k_tok_bytes, kvl.v_tok_bytes, Engine::kv_fp8_on(),
619                            )?;
620                            ph_mark(4, &mut ph_last)?;
621                            e.dtod_copy_into(&a_row, &mut attn, bi * q_dim)?;
622                            ph_mark(3, &mut ph_last)?;
623                        }
624                    }
625
626                    // Output gate (element-wise — batches whole) + o-proj at m=B.
627                    let attn_g = match &gate {
628                        Some(g) => {
629                            let n = b_n * q_dim;
630                            let mut gsig = e.uninit(n)?;
631                            e.sigmoid(g, &mut gsig, n)?;
632                            let mut ag = e.uninit(n)?;
633                            e.mul(&attn, &gsig, &mut ag, n)?;
634                            ag
635                        }
636                        None => attn,
637                    };
638                    let o = e.matmul(&fa.wo, &attn_g, b_n)?;
639                    ph_mark(5, &mut ph_last)?;
640                    o
641                }
642                Mixer::Linear(la) => {
643                    // v2 (the B-scaling fix): the GDN mixer's PROJECTIONS carry the layer's
644                    // weight mass — batch them at m=B so wqkv/gate/beta/alpha/ssm_out stream
645                    // ONCE per step instead of once per sequence. Only the recurrent state ops
646                    // (fused conv ring, gdn prep, gdn scan) stay per-seq — they are state-bound
647                    // micro-kernels, not weight readers. Composition unchanged vs v1 (matmul_pre
648                    // == fused2 per (tensor,row); _bN mmvq per-row == m=1): same numeric config.
649                    let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
650                    let d_state = ssm.state_size as usize;
651                    let num_k = ssm.group_count as usize;
652                    let num_v = ssm.time_step_rank as usize;
653                    let d_conv = ssm.conv_kernel as usize;
654                    let key_dim = d_state * num_k;
655                    let value_dim = d_state * num_v;
656                    let conv_dim = key_dim * 2 + value_dim;
657                    let gdn_scale = 1.0 / (d_state as f32).sqrt();
658
659                    // ---- batched projections (the weight win) ----
660                    let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, b_n)?;
661                    let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, b_n)?;
662                    let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, b_n)?;
663                    let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, b_n)?;
664                    ph_mark(6, &mut ph_last)?;
665
666                    // ---- batched recurrent state ops (3 launches for all B sequences) ----
667                    let base = lin_base[il].expect("linear layer missing from pointer table");
668                    let table = ptr_table.as_ref().expect("pointer table missing");
669                    let conv_view = table.slice(base..base + b_n);
670                    let in_view = table.slice(base + b_n..base + 2 * b_n);
671                    let out_view = table.slice(base + 2 * b_n..base + 3 * b_n);
672                    let mut conv_outs = e.uninit(b_n * conv_dim)?;
673                    e.ssm_conv1d_fused_decode_b(&qkv_mixed, &conv_view,
674                                                la.ssm_conv1d.float_data(), &mut conv_outs,
675                                                conv_dim, d_conv, b_n)?;
676                    let mut q_l2 = e.uninit(b_n * value_dim)?;
677                    let mut k_l2 = e.uninit(b_n * value_dim)?;
678                    let mut v_gd = e.uninit(b_n * value_dim)?;
679                    let mut beta_b = e.uninit(b_n * num_v)?;
680                    let mut g_log = e.uninit(b_n * num_v)?;
681                    e.gdn_prep_decode_b(&conv_outs, &beta_raw, &alpha,
682                                        la.ssm_dt.float_data(), la.ssm_a.float_data(),
683                                        &mut q_l2, &mut k_l2, &mut v_gd, &mut beta_b, &mut g_log,
684                                        d_state, num_v, num_k, key_dim, eps, conv_dim, b_n)?;
685                    let mut o_all = e.uninit(b_n * value_dim)?;
686                    e.gdn_scan_s128_batched(&q_l2, &k_l2, &v_gd, &g_log, &beta_b,
687                                            &in_view, &out_view, &mut o_all,
688                                            num_v, b_n, gdn_scale)?;
689                    // ping-pong: scan wrote each seq's alt buffer; swap host handles (the
690                    // NEXT step's table rebuild picks up the new canonical pointers).
691                    for cache in caches.iter_mut() {
692                        let rl = cache.recur[il].as_mut().unwrap();
693                        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
694                    }
695                    ph_mark(7, &mut ph_last)?;
696
697                    // ---- batched gated norm + out-projection ----
698                    let o = if e.uses_q8_1_fast(&la.ssm_out) {
699                        let (gq, gd) = e.gated_rmsnorm_q8_1(&o_all, la.ssm_norm.float_data(),
700                                                            &z, d_state, b_n * num_v, eps)?;
701                        let g0 = e.zeros(0)?;
702                        e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, b_n)?
703                    } else {
704                        let mut gn = e.uninit(b_n * value_dim)?;
705                        e.gated_rmsnorm(&o_all, la.ssm_norm.float_data(), &z, &mut gn,
706                                        d_state, b_n * num_v, eps)?;
707                        e.matmul(&la.ssm_out, &gn, b_n)?
708                    };
709                    ph_mark(8, &mut ph_last)?;
710                    o
711                }
712            };
713
714            // ---- residual add + post_attn_norm + FFN, batched ----
715            let pnorm = layer.post_attn_norm.float_data();
716            let mut x1 = e.uninit(b_n * n_embd)?;
717            let mut z = e.uninit(b_n * n_embd)?;
718            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut z, n_embd, b_n, eps)?;
719            let ffn_out = match &layer.ffn {
720                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
721                    // v1 covers the SiLU family; M3's swigluoai clamp rides a scaled epilogue
722                    // (m=1 fused tier) — batched M3 lands with the batched-fusion pass.
723                    assert!(self.cfg.m3.is_none(),
724                            "decode_step_batch v1: M3 swigluoai FFN not yet batched");
725                    let n_ff = ffn_gate.out_features();
726                    let (zq, zd) = e.quantize_q8_1(&z, b_n, n_embd)?;
727                    // REFUTED ARM (lane/q27-deepdive, 2026-08-05): fusing this gate+up pair
728                    // into `matmul_q8_fused2_t` (the fused2_b8 tier) measured FLAT-TO-NEGATIVE
729                    // at the serving tick — bench c=8 213.1/213.8, 213.9/214.4, 214.4/213.5
730                    // (sign flips) and serve c=8 paired mean −0.20% over 3 passes. Mechanism:
731                    // unlike m=1 (where the pair is 128 of 1015 launches in a 7.67%-gap tick),
732                    // the c=8 tick is 73.2% one weight-bound kernel class with launch cost
733                    // already hidden — halving 128 launches of ~28k buys nothing. The m=1 arm
734                    // in `matmul_pre_dual_noscale` (+0.94%) stays; this call site keeps the two
735                    // launches. Kernel + fused2_b8 wrapper retained: kernel-check gates it at
736                    // m=5/8 and matmul_q8_fused2_t serves the verify tier. Receipts:
737                    // research/q27-deepdive-20260805/ (lever3-bench-*, serve-points.jsonl).
738                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &z, b_n)?;
739                    let u = e.matmul_pre(ffn_up, &zq, &zd, &z, b_n)?;
740                    let mut act = e.uninit(b_n * n_ff)?;
741                    e.silu_mul(&g, &u, &mut act, b_n * n_ff)?;
742                    let (aq, ad) = e.quantize_q8_1(&act, b_n, n_ff)?;
743                    e.matmul_pre(ffn_down, &aq, &ad, &act, b_n)?
744                }
745                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &z, None, b_n, il as u16)?,
746            };
747            // next-layer input x = x1 + ffn_out (batched element-wise add)
748            let mut x2 = e.uninit(b_n * n_embd)?;
749            e.add(&x1, &ffn_out, &mut x2, b_n * n_embd)?;
750            x = x2;
751            ph_mark(9, &mut ph_last)?;
752        }
753
754        // ---- output norm + lm_head at m=B, one D2H ----
755        let mut hn = e.uninit(b_n * n_embd)?;
756        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
757        let logits = e.matmul(&self.output, &hn, b_n)?;
758        ph_mark(10, &mut ph_last)?;
759
760        // GRAMMAR MASKS (constrained decoding): preserve each masked row's PRISTINE logits
761        // for its consumer (lean park into cache.last_logits_dev — the reuse-pool park stays
762        // unmasked, the v1 contract — or the non-lean D2H), then ban in place BEFORE the
763        // device sampler reads the row. All stream-ordered; masks=&[] takes no new branch.
764        let n_vocab = self.output.out_features();
765        let mut logits = logits;
766        let mut pristine: Vec<Option<CudaSlice<f32>>> = Vec::new();
767        if masks.iter().take(b_n).any(|m| m.is_some()) {
768            pristine.resize_with(b_n, || None);
769            for (bi, m) in masks.iter().take(b_n).enumerate() {
770                let Some((mask, words)) = m else { continue };
771                assert!(samp.get(bi).copied().flatten().is_some(),
772                        "grammar-masked row {bi} must request a device sample");
773                if lean {
774                    let cache = &mut caches[bi];
775                    if cache.last_logits_dev.as_ref().map(|d| d.len() < n_vocab).unwrap_or(true) {
776                        cache.last_logits_dev = Some(e.uninit(n_vocab)?);
777                    }
778                    let dst = cache.last_logits_dev.as_mut().unwrap();
779                    e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
780                } else {
781                    let mut p = e.uninit(n_vocab)?;
782                    e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), &mut p)?;
783                    pristine[bi] = Some(p);
784                }
785                e.mask_logits_col(&mut logits, mask, bi, n_vocab, *words)?;
786            }
787        }
788
789        // Device-side sampling for requested rows (see the method doc). Enqueued before the
790        // big logits D2H so the tiny [B] token readback rides the same sync.
791        let mut next: Vec<Option<u32>> = vec![None; b_n];
792        if samp.iter().take(b_n).any(|s| s.is_some()) {
793            let mut toks = e.alloc_u32_zeroed(b_n)?;
794            let mut perturb: Option<CudaSlice<f32>> = None;
795            for (bi, s) in samp.iter().take(b_n).enumerate() {
796                let Some((temp, seed, ctr)) = s else { continue };
797                if *temp <= 0.0 {
798                    e.argmax_token_device_col(&logits, bi, n_vocab, &mut toks, bi)?;
799                } else {
800                    if perturb.is_none() {
801                        perturb = Some(e.zeros(n_vocab)?);
802                    }
803                    let pb = perturb.as_mut().unwrap();
804                    e.gumbel_perturb_col(&logits, bi, pb, n_vocab, *seed, *ctr, *temp)?;
805                    e.argmax_token_device_col(pb, 0, n_vocab, &mut toks, bi)?;
806                }
807            }
808            let host_toks = e.dtoh_u32(&toks)?;
809            for (bi, s) in samp.iter().take(b_n).enumerate() {
810                if s.is_some() {
811                    next[bi] = Some(host_toks[bi]);
812                }
813            }
814        }
815
816        let lean_any = lean && samp.iter().take(b_n).any(|s| s.is_some());
817        let rows: Vec<Vec<f32>> = if lean_any {
818            // LEAN: park device-sampled rows on-device (per-cache buffer, dtod); D2H only
819            // the rows that still need host logits. No sampled rows + no fallback rows =
820            // the big D2H disappears (the [B] token readback above already synced).
821            for (bi, s) in samp.iter().take(b_n).enumerate() {
822                if s.is_none() { continue; }
823                // grammar-masked rows already parked their PRISTINE copy above — the
824                // in-place ban has since poisoned this row for the reuse-pool consumer.
825                if masks.get(bi).copied().flatten().is_some() { continue; }
826                let cache = &mut caches[bi];
827                if cache.last_logits_dev.as_ref().map(|d| d.len() < n_vocab).unwrap_or(true) {
828                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
829                }
830                let dst = cache.last_logits_dev.as_mut().unwrap();
831                e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
832            }
833            (0..b_n)
834                .map(|bi| {
835                    if samp.get(bi).copied().flatten().is_some() {
836                        Ok(Vec::new())
837                    } else {
838                        e.dtoh_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab))
839                    }
840                })
841                .collect::<Result<_, _>>()?
842        } else {
843            let host = e.dtoh(&logits)?;
844            (0..b_n).map(|bi| {
845                // grammar-masked non-lean rows return the PRISTINE copy (the in-place ban
846                // must never leak into last_logits — reuse-pool/park semantics unchanged).
847                if let Some(p) = pristine.get(bi).and_then(|p| p.as_ref()) {
848                    return e.dtoh(p);
849                }
850                Ok(host[bi * n_vocab..(bi + 1) * n_vocab].to_vec())
851            }).collect::<Result<_, _>>()?
852        };
853        for c in caches.iter_mut() {
854            c.pos += 1;
855        }
856        ph_mark(11, &mut ph_last)?;
857        Ok((rows, next))
858    }
859}