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