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        // MEMRA_BATCH_PHASE=1: sync-bounded phase accumulation (diagnostics — see header note).
416        // Initialized BEFORE the tick-input assembly below so slot 0 covers the HOST side of
417        // setup (pos_v/ptr-table builds, embed gather) as well as the H2D sync — the audit-fix
418        // lane's Q6 instrumentation gap (research/audit-fixes2-20260805): the old placement
419        // started the clock after the assembly, so slot 0 under-reported setup.
420        let ph_on = batch_phase_on();
421        let mut ph_last = std::time::Instant::now();
422        let ph_mark = |slot: usize,
423                       last: &mut std::time::Instant|
424         -> Result<(), Box<dyn std::error::Error>> {
425            if ph_on {
426                e.stream().synchronize()?;
427                let now = std::time::Instant::now();
428                BATCH_PHASE.lock().unwrap()[slot] += (now - *last).as_secs_f64();
429                *last = now;
430            }
431            Ok(())
432        };
433
434        // Per-row rope positions (each sequence at its own depth).
435        let pos_v: Vec<i32> = caches.iter().map(|c| c.pos as i32).collect();
436        let pos_d = e.htod_i32(&pos_v)?;
437
438        // Per-step STATE POINTER TABLE (one H2D): for every linear layer, [conv x B]
439        // [ssm_in x B][ssm_out x B] device addresses. The batched state kernels read their
440        // sequence's pointer from these arrays — states stay per-cache (no pooling refactor),
441        // yet conv/prep/scan collapse from 3xB launches per layer to 3. Rebuilt every step
442        // because the ssm ping-pong swaps pointers host-side after each scan.
443        // INCREMENT 2 (2026-08-01): the SAME table now also carries, for every FULL-attn
444        // layer, [k0,v0,k1,v1,...] cache base addresses — the z-batched seqs append and
445        // seqs fa_decode kernels read their sequence's cache through it (the MoE
446        // expert-table pattern), collapsing 2xB launches per attn layer to 2.
447        let mut lin_base: Vec<Option<usize>> = vec![None; self.layers.len()];
448        let mut attn_base: Vec<Option<usize>> = vec![None; self.layers.len()];
449        let mut ptrs: Vec<u64> = Vec::new();
450        {
451            use cudarc::driver::DevicePtr;
452            let s = &e.gpu.stream();
453            for (il, layer) in self.layers.iter().enumerate() {
454                match &layer.mixer {
455                    Mixer::Linear(_) => {
456                        lin_base[il] = Some(ptrs.len());
457                        for c in caches.iter() {
458                            let rl = c.recur[il].as_ref().unwrap();
459                            let (p, _g) = rl.conv_state.device_ptr(s);
460                            ptrs.push(p as u64);
461                        }
462                        for c in caches.iter() {
463                            let rl = c.recur[il].as_ref().unwrap();
464                            let (p, _g) = rl.ssm_state.device_ptr(s);
465                            ptrs.push(p as u64);
466                        }
467                        for c in caches.iter() {
468                            let rl = c.recur[il].as_ref().unwrap();
469                            let (p, _g) = rl.ssm_state_alt.device_ptr(s);
470                            ptrs.push(p as u64);
471                        }
472                    }
473                    Mixer::Full(_) => {
474                        attn_base[il] = Some(ptrs.len());
475                        for c in caches.iter() {
476                            let kvl = c.kv[il].as_ref().unwrap();
477                            let (pk, _g) = kvl.k.device_ptr(s);
478                            let (pv, _g2) = kvl.v.device_ptr(s);
479                            ptrs.push(pk as u64);
480                            ptrs.push(pv as u64);
481                        }
482                    }
483                    Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
484                }
485            }
486        }
487        let ptr_table = if ptrs.is_empty() { None } else { Some(e.htod_u64(&ptrs)?) };
488
489        // INCREMENT 2 arm picks (per STEP — t_kv is layer-invariant within a tick):
490        // - seqs APPEND: format-only condition (per-row program is t_kv-independent);
491        //   default flash module only (fp8-KV rides the per-seq g-module path).
492        // - seqs FA: every row must take the v4 eager arm at ITS OWN t_kv AND all rows
493        //   must share ONE fa_split_keys rung (the rows-twins' straddle law) — a rung
494        //   crossing inside the batch keeps the per-seq loop for that step, so each
495        //   sequence always executes the exact program its isolated run would.
496        // MEMRA_BATCH_APPEND=0 / MEMRA_BATCH_FA=0 are the rollback/A-B seams.
497        let t_kvs: Vec<usize> = caches.iter().map(|c| c.pos + 1).collect();
498        let t_kv_max = *t_kvs.iter().max().unwrap();
499        let seqs_append = {
500            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
501            *ON.get_or_init(|| std::env::var("MEMRA_BATCH_APPEND").as_deref() != Ok("0"))
502        } && !Engine::kv_fp8_on();
503        let sp0 = crate::fa_split_keys(t_kvs[0], cfg.n_head_kv as usize);
504        let seqs_fa = {
505            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
506            *ON.get_or_init(|| std::env::var("MEMRA_BATCH_FA").as_deref() != Ok("0"))
507        } && t_kvs.iter().all(|&t| crate::fa_seqs_eligible(t, head_dim))
508          && t_kvs.iter().all(|&t| crate::fa_split_keys(t, cfg.n_head_kv as usize) == sp0);
509
510        // Embed all B tokens -> x [B, n_embd] (host gather, one H2D).
511        let mut x = e.htod(&self.embd.gather(n_embd, tokens))?;
512        ph_mark(0, &mut ph_last)?;
513
514        for (il, layer) in self.layers.iter().enumerate() {
515            // ---- attn_norm + q8_1 quantize, batched (B rows) ----
516            let anorm = layer.attn_norm.float_data();
517            let mut xn = e.uninit(b_n * n_embd)?;
518            e.rms_norm(&x, anorm, &mut xn, n_embd, b_n, eps)?;
519            let (hq, hd) = e.quantize_q8_1(&xn, b_n, n_embd)?;
520
521            // ---- mixer ----
522            let mixed: CudaSlice<f32> = match &layer.mixer {
523                Mixer::Mla(_) => crate::hybrid::mla_forward_unimplemented(),
524                Mixer::Full(fa) => {
525                    // Batched projections: one weight read serves all B rows.
526                    let qf = e.matmul_pre(&fa.wq, &hq, &hd, &xn, b_n)?;
527                    let mut k = e.matmul_pre(&fa.wk, &hq, &hd, &xn, b_n)?;
528                    let v = e.matmul_pre(&fa.wv, &hq, &hd, &xn, b_n)?;
529
530                    let gated = cfg.attn_out_gate();
531                    let (mut q, gate) = if gated {
532                        let mut qs = e.uninit(b_n * n_head * head_dim)?;
533                        let mut gs = e.uninit(b_n * n_head * head_dim)?;
534                        e.q_gate_split(&qf, &mut qs, &mut gs, head_dim, n_head, b_n)?;
535                        (qs, Some(gs))
536                    } else {
537                        (qf, None)
538                    };
539
540                    // QK-norm over B*n_head rows, rope with per-row positions.
541                    let mut qn = e.uninit(b_n * n_head * head_dim)?;
542                    e.rms_norm(&q, fa.q_norm.float_data(), &mut qn, head_dim, b_n * n_head, eps)?;
543                    q = qn;
544                    let mut kn = e.uninit(b_n * n_head_kv * head_dim)?;
545                    e.rms_norm(&k, fa.k_norm.float_data(), &mut kn, head_dim, b_n * n_head_kv, eps)?;
546                    k = kn;
547                    e.rope_neox(&mut q, &pos_d, head_dim, rope_dims, n_head, b_n,
548                                cfg.rope_freq_base, 1.0)?;
549                    e.rope_neox(&mut k, &pos_d, head_dim, rope_dims, n_head_kv, b_n,
550                                cfg.rope_freq_base, 1.0)?;
551                    ph_mark(1, &mut ph_last)?;
552
553                    // INCREMENT 2 (2026-08-01): the per-seq (append, attend) launch train
554                    // becomes two phases. Phase A appends all B rows (one z-batched launch,
555                    // or the per-seq loop on the seam/fp8 path); phase B attends all B
556                    // sequences (one blockIdx.z launch + one combine on the batched arm —
557                    // which also reads q / writes attn at row offsets, killing the per-seq
558                    // q/a dtod copies — or the per-seq loop when any row is outside the v4
559                    // arm / a split rung crosses inside the batch). Caches are disjoint per
560                    // sequence, so the phase split leaves every row's math untouched.
561                    let q_dim = n_head * head_dim;
562                    let kv_dim = n_head_kv * head_dim;
563                    let mut attn = e.uninit(b_n * q_dim)?;
564                    // ---- phase A: KV append (all B rows) ----
565                    if seqs_append {
566                        let (kdk, kdv, ktb, vtb) = {
567                            let kvl = caches[0].kv[il].as_ref().unwrap();
568                            (kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes)
569                        };
570                        let base = attn_base[il].expect("full layer missing from pointer table");
571                        let table = ptr_table.as_ref().expect("pointer table missing");
572                        let kv_view = table.slice(base..base + 2 * b_n);
573                        e.append_kv_quantized_seqs(&k, &v, &kv_view, &pos_d, b_n,
574                                                   kdk, kdv, ktb, vtb)?;
575                        for cache in caches.iter_mut() {
576                            let kvl = cache.kv[il].as_mut().unwrap();
577                            debug_assert_eq!(kvl.len, cache.pos, "kv len / pos out of lockstep");
578                            kvl.len += 1;
579                        }
580                    } else {
581                        for (bi, cache) in caches.iter_mut().enumerate() {
582                            let kvl = cache.kv[il].as_mut().unwrap();
583                            let k_row = k.slice(bi * kv_dim..(bi + 1) * kv_dim);
584                            let v_row = v.slice(bi * kv_dim..(bi + 1) * kv_dim);
585                            e.append_kv_quantized_view(
586                                &k_row, &v_row, &mut kvl.k, &mut kvl.v, kvl.len,
587                                kvl.kv_dim_k, kvl.kv_dim_v, kvl.k_tok_bytes, kvl.v_tok_bytes,
588                                Engine::kv_fp8_on(),
589                            )?;
590                            kvl.len += 1;
591                        }
592                    }
593                    ph_mark(2, &mut ph_last)?;
594                    // ---- phase B: attention (all B sequences) ----
595                    if seqs_fa {
596                        let (ktb, vtb) = {
597                            let kvl = caches[0].kv[il].as_ref().unwrap();
598                            (kvl.k_tok_bytes, kvl.v_tok_bytes)
599                        };
600                        let base = attn_base[il].expect("full layer missing from pointer table");
601                        let table = ptr_table.as_ref().expect("pointer table missing");
602                        let kv_view = table.slice(base..base + 2 * b_n);
603                        e.fa_decode_batch_seqs_v4(&q, &kv_view, &pos_d, &mut attn,
604                                                  head_dim, n_head, n_head_kv, b_n,
605                                                  t_kv_max, scale, sp0, ktb, vtb)?;
606                        ph_mark(4, &mut ph_last)?;
607                    } else {
608                        for (bi, cache) in caches.iter_mut().enumerate() {
609                            let kvl = cache.kv[il].as_mut().unwrap();
610                            let t_kv = kvl.len;
611                            let k_view = e.view_u8(&kvl.k, t_kv * kvl.k_tok_bytes);
612                            let v_view = e.view_u8(&kvl.v, t_kv * kvl.v_tok_bytes);
613                            // fa_decode wants a q slice starting at row bi: the fallback arm
614                            // scratch-copies the row (q8-class µs cost); the seqs arm above
615                            // reads/writes row offsets in place.
616                            let mut q_row = e.uninit(q_dim)?;
617                            e.dtod_copy_view(&q.slice(bi * q_dim..(bi + 1) * q_dim), &mut q_row)?;
618                            ph_mark(3, &mut ph_last)?;
619                            let mut a_row = e.uninit(q_dim)?;
620                            e.fa_decode_kvmod(
621                                &q_row, &k_view, &v_view, &mut a_row, head_dim, n_head, n_head_kv,
622                                t_kv, scale, kvl.k_tok_bytes, kvl.v_tok_bytes, Engine::kv_fp8_on(),
623                            )?;
624                            ph_mark(4, &mut ph_last)?;
625                            e.dtod_copy_into(&a_row, &mut attn, bi * q_dim)?;
626                            ph_mark(3, &mut ph_last)?;
627                        }
628                    }
629
630                    // Output gate (element-wise — batches whole) + o-proj at m=B.
631                    let attn_g = match &gate {
632                        Some(g) => {
633                            let n = b_n * q_dim;
634                            let mut gsig = e.uninit(n)?;
635                            e.sigmoid(g, &mut gsig, n)?;
636                            let mut ag = e.uninit(n)?;
637                            e.mul(&attn, &gsig, &mut ag, n)?;
638                            ag
639                        }
640                        None => attn,
641                    };
642                    let o = e.matmul(&fa.wo, &attn_g, b_n)?;
643                    ph_mark(5, &mut ph_last)?;
644                    o
645                }
646                Mixer::Linear(la) => {
647                    // v2 (the B-scaling fix): the GDN mixer's PROJECTIONS carry the layer's
648                    // weight mass — batch them at m=B so wqkv/gate/beta/alpha/ssm_out stream
649                    // ONCE per step instead of once per sequence. Only the recurrent state ops
650                    // (fused conv ring, gdn prep, gdn scan) stay per-seq — they are state-bound
651                    // micro-kernels, not weight readers. Composition unchanged vs v1 (matmul_pre
652                    // == fused2 per (tensor,row); _bN mmvq per-row == m=1): same numeric config.
653                    let ssm = cfg.ssm.as_ref().expect("linear mixer requires ssm cfg");
654                    let d_state = ssm.state_size as usize;
655                    let num_k = ssm.group_count as usize;
656                    let num_v = ssm.time_step_rank as usize;
657                    let d_conv = ssm.conv_kernel as usize;
658                    let key_dim = d_state * num_k;
659                    let value_dim = d_state * num_v;
660                    let conv_dim = key_dim * 2 + value_dim;
661                    let gdn_scale = 1.0 / (d_state as f32).sqrt();
662
663                    // ---- batched projections (the weight win) ----
664                    let qkv_mixed = e.matmul_pre(&la.wqkv, &hq, &hd, &xn, b_n)?;
665                    let z = e.matmul_pre(&la.wqkv_gate, &hq, &hd, &xn, b_n)?;
666                    let beta_raw = e.matmul_pre(&la.ssm_beta, &hq, &hd, &xn, b_n)?;
667                    let alpha = e.matmul_pre(&la.ssm_alpha, &hq, &hd, &xn, b_n)?;
668                    ph_mark(6, &mut ph_last)?;
669
670                    // ---- batched recurrent state ops (3 launches for all B sequences) ----
671                    let base = lin_base[il].expect("linear layer missing from pointer table");
672                    let table = ptr_table.as_ref().expect("pointer table missing");
673                    let conv_view = table.slice(base..base + b_n);
674                    let in_view = table.slice(base + b_n..base + 2 * b_n);
675                    let out_view = table.slice(base + 2 * b_n..base + 3 * b_n);
676                    let mut conv_outs = e.uninit(b_n * conv_dim)?;
677                    e.ssm_conv1d_fused_decode_b(&qkv_mixed, &conv_view,
678                                                la.ssm_conv1d.float_data(), &mut conv_outs,
679                                                conv_dim, d_conv, b_n)?;
680                    let mut q_l2 = e.uninit(b_n * value_dim)?;
681                    let mut k_l2 = e.uninit(b_n * value_dim)?;
682                    let mut v_gd = e.uninit(b_n * value_dim)?;
683                    let mut beta_b = e.uninit(b_n * num_v)?;
684                    let mut g_log = e.uninit(b_n * num_v)?;
685                    e.gdn_prep_decode_b(&conv_outs, &beta_raw, &alpha,
686                                        la.ssm_dt.float_data(), la.ssm_a.float_data(),
687                                        &mut q_l2, &mut k_l2, &mut v_gd, &mut beta_b, &mut g_log,
688                                        d_state, num_v, num_k, key_dim, eps, conv_dim, b_n)?;
689                    let mut o_all = e.uninit(b_n * value_dim)?;
690                    e.gdn_scan_s128_batched(&q_l2, &k_l2, &v_gd, &g_log, &beta_b,
691                                            &in_view, &out_view, &mut o_all,
692                                            num_v, b_n, gdn_scale)?;
693                    // ping-pong: scan wrote each seq's alt buffer; swap host handles (the
694                    // NEXT step's table rebuild picks up the new canonical pointers).
695                    for cache in caches.iter_mut() {
696                        let rl = cache.recur[il].as_mut().unwrap();
697                        std::mem::swap(&mut rl.ssm_state, &mut rl.ssm_state_alt);
698                    }
699                    ph_mark(7, &mut ph_last)?;
700
701                    // ---- batched gated norm + out-projection ----
702                    let o = if e.uses_q8_1_fast(&la.ssm_out) {
703                        let (gq, gd) = e.gated_rmsnorm_q8_1(&o_all, la.ssm_norm.float_data(),
704                                                            &z, d_state, b_n * num_v, eps)?;
705                        let g0 = e.zeros(0)?;
706                        e.matmul_pre(&la.ssm_out, &gq, &gd, &g0, b_n)?
707                    } else {
708                        let mut gn = e.uninit(b_n * value_dim)?;
709                        e.gated_rmsnorm(&o_all, la.ssm_norm.float_data(), &z, &mut gn,
710                                        d_state, b_n * num_v, eps)?;
711                        e.matmul(&la.ssm_out, &gn, b_n)?
712                    };
713                    ph_mark(8, &mut ph_last)?;
714                    o
715                }
716            };
717
718            // ---- residual add + post_attn_norm + FFN, batched ----
719            let pnorm = layer.post_attn_norm.float_data();
720            let mut x1 = e.uninit(b_n * n_embd)?;
721            let mut z = e.uninit(b_n * n_embd)?;
722            e.add_rms_norm(&x, &mixed, pnorm, &mut x1, &mut z, n_embd, b_n, eps)?;
723            let ffn_out = match &layer.ffn {
724                crate::hybrid::Ffn::Dense { ffn_gate, ffn_up, ffn_down } => {
725                    // v1 covers the SiLU family; M3's swigluoai clamp rides a scaled epilogue
726                    // (m=1 fused tier) — batched M3 lands with the batched-fusion pass.
727                    assert!(self.cfg.m3.is_none(),
728                            "decode_step_batch v1: M3 swigluoai FFN not yet batched");
729                    let n_ff = ffn_gate.out_features();
730                    let (zq, zd) = e.quantize_q8_1(&z, b_n, n_embd)?;
731                    // REFUTED ARM (lane/q27-deepdive, 2026-08-05): fusing this gate+up pair
732                    // into `matmul_q8_fused2_t` (the fused2_b8 tier) measured FLAT-TO-NEGATIVE
733                    // at the serving tick — bench c=8 213.1/213.8, 213.9/214.4, 214.4/213.5
734                    // (sign flips) and serve c=8 paired mean −0.20% over 3 passes. Mechanism:
735                    // unlike m=1 (where the pair is 128 of 1015 launches in a 7.67%-gap tick),
736                    // the c=8 tick is 73.2% one weight-bound kernel class with launch cost
737                    // already hidden — halving 128 launches of ~28k buys nothing. The m=1 arm
738                    // in `matmul_pre_dual_noscale` (+0.94%) stays; this call site keeps the two
739                    // launches. Kernel + fused2_b8 wrapper retained: kernel-check gates it at
740                    // m=5/8 and matmul_q8_fused2_t serves the verify tier. Receipts:
741                    // research/q27-deepdive-20260805/ (lever3-bench-*, serve-points.jsonl).
742                    let g = e.matmul_pre(ffn_gate, &zq, &zd, &z, b_n)?;
743                    let u = e.matmul_pre(ffn_up, &zq, &zd, &z, b_n)?;
744                    let mut act = e.uninit(b_n * n_ff)?;
745                    e.silu_mul(&g, &u, &mut act, b_n * n_ff)?;
746                    let (aq, ad) = e.quantize_q8_1(&act, b_n, n_ff)?;
747                    e.matmul_pre(ffn_down, &aq, &ad, &act, b_n)?
748                }
749                crate::hybrid::Ffn::Moe(m) => self.moe_ffn_il_zq8(e, m, &z, None, b_n, il as u16)?,
750            };
751            // next-layer input x = x1 + ffn_out (batched element-wise add)
752            let mut x2 = e.uninit(b_n * n_embd)?;
753            e.add(&x1, &ffn_out, &mut x2, b_n * n_embd)?;
754            x = x2;
755            ph_mark(9, &mut ph_last)?;
756        }
757
758        // ---- output norm + lm_head at m=B, one D2H ----
759        let mut hn = e.uninit(b_n * n_embd)?;
760        e.rms_norm(&x, self.output_norm.float_data(), &mut hn, n_embd, b_n, eps)?;
761        let logits = e.matmul(&self.output, &hn, b_n)?;
762        ph_mark(10, &mut ph_last)?;
763
764        // GRAMMAR MASKS (constrained decoding): preserve each masked row's PRISTINE logits
765        // for its consumer (lean park into cache.last_logits_dev — the reuse-pool park stays
766        // unmasked, the v1 contract — or the non-lean D2H), then ban in place BEFORE the
767        // device sampler reads the row. All stream-ordered; masks=&[] takes no new branch.
768        let n_vocab = self.output.out_features();
769        let mut logits = logits;
770        let mut pristine: Vec<Option<CudaSlice<f32>>> = Vec::new();
771        if masks.iter().take(b_n).any(|m| m.is_some()) {
772            pristine.resize_with(b_n, || None);
773            for (bi, m) in masks.iter().take(b_n).enumerate() {
774                let Some((mask, words)) = m else { continue };
775                assert!(samp.get(bi).copied().flatten().is_some(),
776                        "grammar-masked row {bi} must request a device sample");
777                if lean {
778                    let cache = &mut caches[bi];
779                    if cache.last_logits_dev.as_ref().map(|d| d.len() < n_vocab).unwrap_or(true) {
780                        cache.last_logits_dev = Some(e.uninit(n_vocab)?);
781                    }
782                    let dst = cache.last_logits_dev.as_mut().unwrap();
783                    e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
784                } else {
785                    let mut p = e.uninit(n_vocab)?;
786                    e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), &mut p)?;
787                    pristine[bi] = Some(p);
788                }
789                e.mask_logits_col(&mut logits, mask, bi, n_vocab, *words)?;
790            }
791        }
792
793        // Device-side sampling for requested rows (see the method doc). Enqueued before the
794        // big logits D2H so the tiny [B] token readback rides the same sync.
795        let mut next: Vec<Option<u32>> = vec![None; b_n];
796        if samp.iter().take(b_n).any(|s| s.is_some()) {
797            let mut toks = e.alloc_u32_zeroed(b_n)?;
798            let mut perturb: Option<CudaSlice<f32>> = None;
799            for (bi, s) in samp.iter().take(b_n).enumerate() {
800                let Some((temp, seed, ctr)) = s else { continue };
801                if *temp <= 0.0 {
802                    e.argmax_token_device_col(&logits, bi, n_vocab, &mut toks, bi)?;
803                } else {
804                    if perturb.is_none() {
805                        perturb = Some(e.zeros(n_vocab)?);
806                    }
807                    let pb = perturb.as_mut().unwrap();
808                    e.gumbel_perturb_col(&logits, bi, pb, n_vocab, *seed, *ctr, *temp)?;
809                    e.argmax_token_device_col(pb, 0, n_vocab, &mut toks, bi)?;
810                }
811            }
812            let host_toks = e.dtoh_u32(&toks)?;
813            for (bi, s) in samp.iter().take(b_n).enumerate() {
814                if s.is_some() {
815                    next[bi] = Some(host_toks[bi]);
816                }
817            }
818        }
819
820        let lean_any = lean && samp.iter().take(b_n).any(|s| s.is_some());
821        let rows: Vec<Vec<f32>> = if lean_any {
822            // LEAN: park device-sampled rows on-device (per-cache buffer, dtod); D2H only
823            // the rows that still need host logits. No sampled rows + no fallback rows =
824            // the big D2H disappears (the [B] token readback above already synced).
825            for (bi, s) in samp.iter().take(b_n).enumerate() {
826                if s.is_none() { continue; }
827                // grammar-masked rows already parked their PRISTINE copy above — the
828                // in-place ban has since poisoned this row for the reuse-pool consumer.
829                if masks.get(bi).copied().flatten().is_some() { continue; }
830                let cache = &mut caches[bi];
831                if cache.last_logits_dev.as_ref().map(|d| d.len() < n_vocab).unwrap_or(true) {
832                    cache.last_logits_dev = Some(e.uninit(n_vocab)?);
833                }
834                let dst = cache.last_logits_dev.as_mut().unwrap();
835                e.dtod_copy_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab), dst)?;
836            }
837            (0..b_n)
838                .map(|bi| {
839                    if samp.get(bi).copied().flatten().is_some() {
840                        Ok(Vec::new())
841                    } else {
842                        e.dtoh_view(&logits.slice(bi * n_vocab..(bi + 1) * n_vocab))
843                    }
844                })
845                .collect::<Result<_, _>>()?
846        } else {
847            let host = e.dtoh(&logits)?;
848            (0..b_n).map(|bi| {
849                // grammar-masked non-lean rows return the PRISTINE copy (the in-place ban
850                // must never leak into last_logits — reuse-pool/park semantics unchanged).
851                if let Some(p) = pristine.get(bi).and_then(|p| p.as_ref()) {
852                    return e.dtoh(p);
853                }
854                Ok(host[bi * n_vocab..(bi + 1) * n_vocab].to_vec())
855            }).collect::<Result<_, _>>()?
856        };
857        for c in caches.iter_mut() {
858            c.pos += 1;
859        }
860        ph_mark(11, &mut ph_last)?;
861        Ok((rows, next))
862    }
863}