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}