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}