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