Skip to main content

memra_engine/
pp.rs

1//! M2 pipeline-parallel N-stage runtime (generalizes the M1 2-stage seam).
2//!
3//! Door: `MEMRA_PP_STAGES=N` (default OFF — unset/0/1 = no behavior change anywhere).
4//! Stage map: N stages over the trunk layers with N-1 cuts. `MEMRA_PP_SPLITS=c1,..,cN-1`
5//! sets the cuts explicitly (strictly increasing, in (0, n_layers)); `MEMRA_PP_SPLIT=<i>`
6//! is the N=2 back-compat spelling; default = even split (cut s = s*n_layers/N).
7//! Placement: `MEMRA_PP_DEVICES=d0,..,dN-1` maps stage s to device ds (default: all on
8//! the primary engine's device).
9//!
10//! M1 history (increments 1-2, merged + hardened on the 8x box 2026-08-02): seam + gate
11//! single-device; then real transport — per-stage streams/events, device placement,
12//! peer-copy boundary (M0: cudaMemcpyPeerAsync beats NCCL 2.8x at PP activation sizes),
13//! per-context PDL module caches, default-mempool peer grants. All five r3 gates PASS
14//! bit-identical (receipts ~/receipts/m1-pp2/ on darklanes-bench).
15//!
16//! M2 increment 1 (this file): N-STAGE GENERALIZATION — `Pp2Rt` becomes `PpNRt`:
17//!   - `stages`: Vec of per-stage execution homes (device, context, stream, remote Engine);
18//!   - `boundaries`: N-1 boundary runtimes, each with TWO persistent double-buffered slots
19//!     (ev_tx/ev_rx per slot) and its own overlap step counter; transport is selected PER
20//!     BOUNDARY (dtod same-device / cudaMemcpyPeerAsync cross-device by default; opt-in
21//!     `MEMRA_PP_HOST_BOUNCE=1` uses pinned D2H + H2D instead);
22//!   - the default peer transport grants peer + default-mempool access between EVERY distinct
23//!     pair of devices in use. Host bounce skips serving-time grants; its boot diagnostics
24//!     transiently enable peer + pool access, then revoke the pool grants and disable peer access
25//!     before proceeding. Sharded weights plus stage-local auxiliary buffers ensure that no peer
26//!     read can bypass the bounced boundary.
27//!
28//! M2 increment 2 (weight sharding): the loader uploads each stage's layer range THROUGH
29//! that stage's engine (`layer_engine`), so weights land on the device that runs them —
30//! the bring-up peer-read placement dies. `output_norm` + lm head load through the LAST
31//! stage's engine; the embed table stays host-side with stage 0. Split-plane/f16 decode
32//! mirrors are built per layer through the owning stage's engine too (the rp4 mirrors ARE
33//! the decode weights on the q8 path — leaving them on dev0 would fake the kill).
34//! Rollback seam: `MEMRA_PP_SHARD=0` = M1 bring-up placement (all weights on primary,
35//! remote stages peer-read).
36//!
37//! M2 increment 3 (deferred readback — the pipelining seed): `PendingLogits` — the eager
38//! decode arm can END a step without the logits D2H (`decode_step_h_ppn_deferred`): the
39//! logits stay device-resident with a completion event; `wait()` drains them through a
40//! DEDICATED readback stream (waits the event, copies, syncs) so tokens t+1.. keep
41//! enqueuing on the stage streams while token t drains. Per-token math is fully
42//! event-ordered (same slots, same ev_tx/ev_rx chain) — scheduling changes, math does
43//! not; the pipelined replay arm of `ppn-gate` proves bit-identity per step.
44//!
45//! Ownership across a boundary (unchanged from M1):
46//!   - hidden state [n_embd] f32 is the ONLY tensor that crosses;
47//!   - KV/linear-attn cache entries are per-layer: stage s exclusively owns cache state
48//!     for its layer range (and, under MEMRA_PP_DEVICES, allocates it on its device);
49//!   - position/rope state is the scalar `cache.pos` snapshot taken once per step; every stage
50//!     uploads its own position buffer on its own stream (no cross-device position pointer);
51//!   - the embed table lives with stage 0, output_norm + lm head with the last stage.
52//!
53//! THE MULTI-STREAM LAW (why this is safe with cudarc event tracking disabled): all
54//! cross-stage bytes flow through the persistent boundary slots, ordered by ev_tx/ev_rx;
55//! per-stage scratch is allocated AND freed on that stage's stream (stream-ordered); the
56//! async mem pool runs with opportunistic reuse OFF + internal dependencies ON
57//! (memra-runtime), so a block freed on stream A and reused on stream B carries a
58//! driver-inserted dependency. Weights are load-time state no stage stream can precede,
59//! and the step's terminal logits readback (sync D2H, or PendingLogits' event-ordered
60//! readback stream) drains the last stage, whose TX-wait chain transitively drains all.
61//!
62//! Scope: plain eager decode only (generic arm N-stage; gemma4 arm 2-stage). NOT wired:
63//! batch/dc/graph/spec loops and the gemma4-E4B eager arm.
64//!
65//! CORRECTION (pp2-hardening 2026-08-06): this header used to add "(`warn_unwired_once`
66//! fires)" to that list, which was wrong. `warn_unwired_once` has exactly two call sites
67//! and BOTH are gemma4-specific (decode.rs, hybrid_forward.rs) — the batch/dc/graph/spec
68//! loops never warned. Worse, the batched loop did not merely run unsplit: it walked the
69//! whole trunk on the primary stream and, under a sharded cross-device placement,
70//! peer-read every remote stage's weights each step — 28x slower at B=1 with all three
71//! `decode-batch-gate` gates PASSING (peer reads are byte-exact, so only perf broke).
72//! `decode_step_batch` now FAILS CLOSED in that regime via `pp_sharded_cross_device()`
73//! (`MEMRA_PP_ALLOW_UNSPLIT_BATCH=1` = measurement override). "Unwired" for dc/graph/spec
74//! still means "runs unsplit, silently" — audit each before trusting it on a pair.
75
76use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
77use std::sync::{Arc, Mutex, OnceLock};
78
79use cudarc::driver::{CudaContext, CudaEvent, CudaSlice, CudaStream};
80
81use crate::Engine;
82
83/// Returns the stage fence iff the ppN door is open: `MEMRA_PP_STAGES=N` (N >= 2) with a
84/// valid cut list. The fence has N+1 entries: `[0, c1, .., cN-1, n_layers]`; stage s runs
85/// layers `[fence[s], fence[s+1])`. Reads the environment on every call (gates toggle the
86/// door in-process); the cost is a few getenv per decode step, eager-loop noise.
87pub fn pp_cuts(n_layers: usize) -> Option<Vec<usize>> {
88    let n_st: usize = match std::env::var("MEMRA_PP_STAGES") {
89        Ok(v) if v.is_empty() || v == "0" || v == "1" => return None,
90        Ok(v) => match v.parse::<usize>() {
91            Ok(n) => n,
92            Err(_) => {
93                warn_bad_once(&format!("MEMRA_PP_STAGES={v} unparseable; door stays OFF"));
94                return None;
95            }
96        },
97        Err(_) => return None,
98    };
99    if n_st < 2 || n_st > n_layers {
100        warn_bad_once(&format!(
101            "MEMRA_PP_STAGES={n_st} outside [2, n_layers={n_layers}]; door stays OFF"
102        ));
103        return None;
104    }
105    let mut fence = Vec::with_capacity(n_st + 1);
106    fence.push(0usize);
107    if let Ok(s) = std::env::var("MEMRA_PP_SPLITS") {
108        let parts: Result<Vec<usize>, _> =
109            s.split(',').map(|p| p.trim().parse::<usize>()).collect();
110        match parts {
111            Ok(cuts) if cuts.len() == n_st - 1 => fence.extend(cuts),
112            _ => {
113                warn_bad_once(&format!(
114                    "MEMRA_PP_SPLITS={s} invalid (want {} comma-separated cuts); door stays OFF",
115                    n_st - 1
116                ));
117                return None;
118            }
119        }
120    } else if let Ok(v) = std::env::var("MEMRA_PP_SPLIT") {
121        // N=2 back-compat spelling. With N>2 a single split is ambiguous — fail the door
122        // loudly rather than guess (a silent even-split would fake a gate config).
123        if n_st != 2 {
124            warn_bad_once(&format!(
125                "MEMRA_PP_SPLIT={v} set with MEMRA_PP_STAGES={n_st}; use MEMRA_PP_SPLITS \
126                 for N>2 — door stays OFF"
127            ));
128            return None;
129        }
130        match v.parse::<usize>() {
131            Ok(c) => fence.push(c),
132            Err(_) => {
133                warn_bad_once(&format!("MEMRA_PP_SPLIT={v} unparseable; door stays OFF"));
134                return None;
135            }
136        }
137    } else {
138        for s in 1..n_st {
139            fence.push(s * n_layers / n_st);
140        }
141    }
142    fence.push(n_layers);
143    for w in fence.windows(2) {
144        if w[0] >= w[1] {
145            warn_bad_once(&format!(
146                "pp stage fence {fence:?} not strictly increasing over [0, {n_layers}]; \
147                 door stays OFF"
148            ));
149            return None;
150        }
151    }
152    Some(fence)
153}
154
155/// N=2 back-compat view of the door (the gemma4 arm and `pp2-gate` are 2-stage): `Some(cut)`
156/// iff the door is open with EXACTLY two stages.
157pub fn pp2_split(n_layers: usize) -> Option<usize> {
158    pp_cuts(n_layers).filter(|f| f.len() == 3).map(|f| f[1])
159}
160
161/// The stage that owns layer `il` under `fence` (see `pp_cuts`).
162pub fn stage_of(fence: &[usize], il: usize) -> usize {
163    debug_assert!(fence.len() >= 2);
164    match fence[1..fence.len() - 1].binary_search(&il) {
165        // fence[1..][k] == il means il is the FIRST layer of stage k+1
166        Ok(k) => k + 1,
167        Err(k) => k,
168    }
169}
170
171/// MEMRA_PP_STREAMS=0: rollback to the increment-1 same-stream seam (boundary = two plain
172/// dtod copies on the ambient compute stream, no per-stage streams/events/devices).
173pub fn pp2_streams_off() -> bool {
174    matches!(std::env::var("MEMRA_PP_STREAMS").as_deref(), Ok("0"))
175}
176
177/// True iff the ppN door would put TWO OR MORE stage streams on ONE device (devices
178/// unset = all stages on the primary; or an explicit placement with a repeated device).
179/// The deferred-readback (pipelined) arm is REFUSED in this regime: the 2026-08-02 x20
180/// soak record — singledev pipelined 13/20 PASS default, 7 failures each diverging at a
181/// different step (timing-race signature); MEMRA_PDL=0 went 20/20 on one soak but a
182/// second same-config soak on the auto-gated build failed 2/20 (n2) and battery-4 failed
183/// n4 — so PDL narrows the window without closing it, and the true root cause (same
184/// Engine kernels concurrent on two streams of one device) is NOT fixed by any flag yet.
185/// Cross-device pipelined (one stage stream per device) is 23/23 clean post-fix. Refuse
186/// loudly rather than return silently-wrong logits. Env-only read (callable pre-runtime).
187pub fn pp_multi_stream_same_device() -> bool {
188    let stages_open = std::env::var("MEMRA_PP_STAGES")
189        .map(|v| v.parse::<usize>().map(|n| n >= 2).unwrap_or(false))
190        .unwrap_or(false);
191    let devices = std::env::var("MEMRA_PP_DEVICES")
192        .ok()
193        .filter(|v| !v.is_empty());
194    if (!stages_open && devices.is_none()) || pp2_streams_off() {
195        return false;
196    }
197    match devices {
198        None => true, // door open, no placement: every stage stream lands on the primary
199        Some(s) => {
200            let mut v: Vec<&str> = s.split(',').map(|p| p.trim()).collect();
201            let n = v.len();
202            v.sort_unstable();
203            v.dedup();
204            v.len() < n // repeated device = shared-device streams
205        }
206    }
207}
208
209/// True iff the ppN door is open AND the placement spans 2+ DISTINCT devices AND the
210/// per-stage sharded loader is on — i.e. some layers' weights live on a device other than
211/// the primary. Any path that walks the WHOLE trunk on one stream in this regime reads
212/// those weights over PCIe every step. Env-only read (callable pre-runtime).
213///
214/// Measured cost of doing that (pp2-hardening 2026-08-06, 2x RTX PRO 6000, PCIe Gen5 x16
215/// P2P, decode-batch-bench q9, N=5 interleaved, `research/pp2-hardening-20260806`):
216/// **B=1 7.4 vs 208.9 tok/s (28x), B=4 29.8 vs 491.3 (16.5x), B=8 47.4 vs 657.0 (13.9x)**.
217/// The same sweep with `MEMRA_PP_SHARD=0` (weights all home) returns 178.5/491.1/656.6 —
218/// identical to the single-device door-open arm — so the entire cliff is the peer read,
219/// not the door and not the placement plumbing. Exactness is NOT the issue: peer reads
220/// return identical bytes and every `decode-batch-gate` gate PASSED on this config, which
221/// is precisely why it needs a refusal rather than a gate.
222pub fn pp_sharded_cross_device() -> bool {
223    let stages_open = std::env::var("MEMRA_PP_STAGES")
224        .map(|v| v.parse::<usize>().map(|n| n >= 2).unwrap_or(false))
225        .unwrap_or(false);
226    // MEMRA_PP_STREAMS=0 (2026-08-06, pp2-batch): the same-stream rollback seam ALSO turns
227    // the sharded loader off — `layer_engine` returns the primary engine whenever
228    // `pp2_streams_off()`, and `new_cache` skips `Cache::new_ppn` on the same condition. So
229    // in that regime every weight and every cache is home on the primary and an unsplit walk
230    // peer-reads NOTHING. Without this term the guard refused that config too: a spurious
231    // refusal of a placement that is sound and full-speed. Found wiring the batched pp arm.
232    if !stages_open || pp_shard_off() || pp2_streams_off() {
233        return false;
234    }
235    match pp2_devices_env() {
236        None => false, // no placement: every stage is the primary device, nothing remote
237        Some(s) => {
238            let mut v: Vec<&str> = s.split(',').map(|p| p.trim()).collect();
239            v.sort_unstable();
240            v.dedup();
241            v.len() >= 2
242        }
243    }
244}
245
246/// The shared fail-closed guard for EVERY decode path that has no pp stage split.
247/// Returns `Err` iff `pp_sharded_cross_device()` — i.e. the caller would walk the whole
248/// trunk on one stream while some layers' weights live on another device, peer-reading
249/// them every step. `path` names the refusing function so the operator knows which loop
250/// they hit; `alt` names the working alternative for that loop.
251///
252/// One helper rather than four copies because the audit found FOUR paths with the same
253/// hole (`decode_step_batch`, `decode_step_dc`, the graph capture that wraps dc, and
254/// `decode_step_t*` verify), and a per-path copy is how one gets missed on the next
255/// addition. Override: `MEMRA_PP_ALLOW_UNSPLIT_BATCH=1` (one door for all of them —
256/// they are the same measurement question).
257pub fn refuse_unsplit_if_remote(path: &str, alt: &str) -> Result<(), Box<dyn std::error::Error>> {
258    if pp_host_bounce_active() {
259        return Err(format!(
260            "{path}: refused with MEMRA_PP_HOST_BOUNCE=1 on sharded cross-device PP — \
261             this unsplit path peer-reads remote weights, while host bounce covers only \
262             explicit stage-boundary transfers. Use {alt}; the \
263             MEMRA_PP_ALLOW_UNSPLIT_BATCH override is unavailable on a broken-peer host."
264        )
265        .into());
266    }
267    if pp_sharded_cross_device()
268        && std::env::var("MEMRA_PP_ALLOW_UNSPLIT_BATCH").as_deref() != Ok("1")
269    {
270        return Err(format!(
271            "{path}: refused with the ppN door open across 2+ devices — this path has no pp \
272             stage split, so it would walk ALL layers on one stream and peer-read every \
273             remote stage's weights each step (measured 28x slower at B=1, 13.9x at B=8 on \
274             a PRO 6000 pair over PCIe Gen5 x16 P2P; research/pp2-hardening-20260806). \
275             Exactness is unaffected — peer reads return identical bytes and the exactness \
276             gates PASS on this config — which is exactly why it must refuse instead of \
277             being caught by a gate. Fixes, in order: {alt}; or MEMRA_PP_SHARD=0 (all \
278             weights home on the primary — full speed, forfeits the capacity PP-2 exists \
279             for); or close the pp door. MEMRA_PP_ALLOW_UNSPLIT_BATCH=1 overrides for \
280             measurement."
281        )
282        .into());
283    }
284    Ok(())
285}
286
287/// MEMRA_BATCH_PP=0: rollback/A-B seam for the BATCHED stage split (pp2-batch 2026-08-06).
288/// Default ON — with the ppN door open the batched decode step takes its own stage split
289/// (`decode_step_batch_ppn`) exactly as the eager step does. Setting 0 sends the batched
290/// path back through the unsplit body, which under a sharded cross-device placement is
291/// then caught by `refuse_unsplit_if_remote` (the 28x peer-read regime) rather than run
292/// silently. Exists so the bit-identity gate can A/B split vs unsplit IN ONE PROCESS
293/// against the same loaded weights — read per step, never memoized, for that reason.
294pub fn batch_pp_on() -> bool {
295    std::env::var("MEMRA_BATCH_PP").as_deref() != Ok("0")
296}
297
298/// MEMRA_DUAL_PP three-state mode for the dual-active PP-2 batched decode path.
299/// Default ON (owner flip 2026-08-11) after the box1 PRO-pair re-gate: correctness
300/// bit-identity B=1..5, servestress no-thrash, 10-boot soak 929/929 golden matches with
301/// 0 slot collisions across 9123 pairs (research/dualpp2-20260811/RESULTS-regate.md), plus
302/// the dualpp1 c>=8 interleaved perf floor (+20.753% minimum,
303/// research/dualpp1-20260811/RESULTS.md).
304///
305/// The three states carry different failure semantics on purpose:
306/// - `Off` (`MEMRA_DUAL_PP=0`): the serial rollback seam. Overlap also follows OFF unless
307///   `MEMRA_PP_OVERLAP` is set explicitly, so one flag restores the exact pre-flip naked path.
308/// - `Forced` (`MEMRA_DUAL_PP=1`): the pre-flip explicit request. A placement that cannot
309///   run dual (single-slot boundary, host bounce, non-PP-2 fence) REFUSES with the binding
310///   quoted reason before any token or cache advance — the gate negative cells pin this.
311/// - `Auto` (unset): the flipped default. Dual runs where the re-gate validated it
312///   (PP-2 fence, double-slot, peer transport, B>=2) and silently degrades to the serial
313///   PP-N walker everywhere else — naked PP-3 serving and the MEMRA_PP_HOST_BOUNCE=1
314///   broken-peer escape hatch must keep decoding, not refuse.
315#[derive(Clone, Copy, PartialEq, Eq, Debug)]
316pub enum DualPpMode {
317    Off,
318    Forced,
319    Auto,
320}
321
322/// Pure resolution for MEMRA_DUAL_PP, split from the env read so the flip regression tests
323/// cannot race parallel test threads on process env.
324pub fn dual_pp_mode_resolve(v: Option<&str>) -> DualPpMode {
325    match v {
326        Some("0") => DualPpMode::Off,
327        Some("1") => DualPpMode::Forced,
328        _ => DualPpMode::Auto,
329    }
330}
331
332pub fn dual_pp_mode() -> DualPpMode {
333    dual_pp_mode_resolve(std::env::var("MEMRA_DUAL_PP").ok().as_deref())
334}
335
336/// True when the dual-active door is open (Forced or Auto). Read per step so the
337/// model-level gate can replay serial and waved arms against one loaded checkpoint.
338pub fn dual_pp_on() -> bool {
339    dual_pp_mode() != DualPpMode::Off
340}
341
342/// Engine-entry routing for the dual-active path, kept pure for the flip regression
343/// tests. `Forced` routes every B>=2 PP-2 call into `decode_step_batch_dual` even when
344/// the placement cannot run it, so the binding refusals stay reachable and loud.
345/// `Auto` routes only the exact re-gated regime and leaves everything else on the serial
346/// PP-N walker. `dual_pp_eligibility` remains behind this as defense in depth.
347pub fn dual_pp_route(
348    mode: DualPpMode,
349    batch: usize,
350    stages: usize,
351    double_slot: bool,
352    host_bounce: bool,
353) -> bool {
354    if batch < 2 {
355        return false;
356    }
357    match mode {
358        DualPpMode::Off => false,
359        DualPpMode::Forced => true,
360        DualPpMode::Auto => stages == 2 && double_slot && !host_bounce,
361    }
362}
363
364/// Binding-amendment refusal text. The negative gate quotes this exact line and requires the
365/// decode call to return before producing a token or advancing a cache.
366pub const DUAL_PP_SINGLE_SLOT_REFUSAL: &str = "decode_step_batch_dual: refused: PP boundary is single-slot; set MEMRA_PP_OVERLAP=1 so both alternating boundary slots are prepared before dual-active decode";
367pub const DUAL_PP_HOST_BOUNCE_REFUSAL: &str = "decode_step_batch_dual: refused: MEMRA_PP_HOST_BOUNCE=1 is unvalidated for dual-active decode; disable MEMRA_DUAL_PP or use peer transport";
368
369/// Pure schedule policy shared by the runtime and kernel-check manifest cells. A single row
370/// has no second wave and must stay on the serial PP-N walker.
371pub fn dual_pp_wave_mid(batch: usize) -> Option<usize> {
372    (batch >= 2).then_some((batch + 1) / 2)
373}
374
375/// Fail-closed eligibility check kept pure so the negative manifest cell cannot accidentally
376/// initialize CUDA state. Slot preparation itself remains `PpNRt::prepare_overlap_slots`.
377pub fn dual_pp_eligibility(
378    stages: usize,
379    double_slot: bool,
380    host_bounce: bool,
381) -> Result<(), &'static str> {
382    if stages != 2 {
383        return Err(
384            "decode_step_batch_dual: refused: dual-active decode requires exactly two PP stages",
385        );
386    }
387    if !double_slot {
388        return Err(DUAL_PP_SINGLE_SLOT_REFUSAL);
389    }
390    if host_bounce {
391        return Err(DUAL_PP_HOST_BOUNCE_REFUSAL);
392    }
393    Ok(())
394}
395
396/// Liveness is counted only while the two host-driven decode layer walkers are both active.
397/// Enqueue order is not proof for Step: its router readback synchronizes the issuing thread.
398static DUAL_PP_OVERLAPS: AtomicUsize = AtomicUsize::new(0);
399static DUAL_PP_ACTIVE_STAGES: AtomicUsize = AtomicUsize::new(0);
400static DUAL_PP_STAGE_NS: [AtomicU64; 4] = [
401    AtomicU64::new(0),
402    AtomicU64::new(0),
403    AtomicU64::new(0),
404    AtomicU64::new(0),
405];
406static DUAL_PP_STAGE_SAMPLES: [AtomicUsize; 4] = [
407    AtomicUsize::new(0),
408    AtomicUsize::new(0),
409    AtomicUsize::new(0),
410    AtomicUsize::new(0),
411];
412static DUAL_PP_TIMING_DROPPED: AtomicUsize = AtomicUsize::new(0);
413static DUAL_PP_SLOT_PAIRS: AtomicUsize = AtomicUsize::new(0);
414static DUAL_PP_SLOT_USES: [AtomicUsize; 2] = [AtomicUsize::new(0), AtomicUsize::new(0)];
415static DUAL_PP_SLOT_COLLISIONS: AtomicUsize = AtomicUsize::new(0);
416
417pub const DUAL_PP_STAGE_NAMES: [&str; 4] = [
418    "wave_a_stage0",
419    "wave_a_stage1",
420    "wave_b_stage0",
421    "wave_b_stage1",
422];
423
424pub fn dual_pp_overlaps() -> usize {
425    DUAL_PP_OVERLAPS.load(Ordering::Relaxed)
426}
427
428/// Record the two boundary slots selected for one dual-active wave pair. A same-slot pair is
429/// rejected by the caller before wave B can consume a residual; the collision counter makes that
430/// fail-closed path observable to the detached soak instead of relying only on log scanning.
431pub(crate) fn record_dual_pp_slot_pair(slot_a: usize, slot_b: usize) -> bool {
432    debug_assert!(slot_a < DUAL_PP_SLOT_USES.len());
433    debug_assert!(slot_b < DUAL_PP_SLOT_USES.len());
434    if slot_a == slot_b {
435        DUAL_PP_SLOT_COLLISIONS.fetch_add(1, Ordering::Relaxed);
436        return false;
437    }
438    DUAL_PP_SLOT_USES[slot_a].fetch_add(1, Ordering::Relaxed);
439    DUAL_PP_SLOT_USES[slot_b].fetch_add(1, Ordering::Relaxed);
440    DUAL_PP_SLOT_PAIRS.fetch_add(1, Ordering::Relaxed);
441    true
442}
443
444/// `(completed wave pairs, [slot 0 uses, slot 1 uses], rejected same-slot pairs)`.
445pub fn dual_pp_slot_snapshot() -> (usize, [usize; 2], usize) {
446    (
447        DUAL_PP_SLOT_PAIRS.load(Ordering::Relaxed),
448        std::array::from_fn(|i| DUAL_PP_SLOT_USES[i].load(Ordering::Relaxed)),
449        DUAL_PP_SLOT_COLLISIONS.load(Ordering::Relaxed),
450    )
451}
452
453/// CUDA-event timing is a diagnostic-only process door. The scored N=5 block runs without
454/// it; the companion box1 diagnostic process enables it and exports cumulative per-wave
455/// stage spans through `/metrics`.
456pub fn dual_pp_timing_on() -> bool {
457    static ON: OnceLock<bool> = OnceLock::new();
458    *ON.get_or_init(|| std::env::var("MEMRA_DUAL_PP_TIMING").as_deref() == Ok("1"))
459}
460
461pub(crate) fn record_dual_pp_stage_ms(stage: usize, ms: f32) {
462    assert!(
463        stage < DUAL_PP_STAGE_NS.len(),
464        "dual PP timing stage out of range"
465    );
466    let ns = (f64::from(ms) * 1_000_000.0).round() as u64;
467    DUAL_PP_STAGE_NS[stage].fetch_add(ns, Ordering::Relaxed);
468    DUAL_PP_STAGE_SAMPLES[stage].fetch_add(1, Ordering::Relaxed);
469}
470
471/// Timing is diagnostic only: a CUDA event that is not ready (or otherwise fails) must not
472/// change decode control flow. Count and warn once, then leave the scored-path result intact.
473pub(crate) fn record_dual_pp_timing_drop(context: &str, err: &dyn std::fmt::Display) {
474    let previous = DUAL_PP_TIMING_DROPPED.fetch_add(1, Ordering::Relaxed);
475    if previous == 0 {
476        eprintln!(
477            "[dual-pp] WARN: skipped diagnostic timing sample at {context}: {err}; decode continues"
478        );
479    }
480}
481
482pub(crate) fn record_dual_pp_stage_result<E: std::fmt::Display>(
483    stage: usize,
484    elapsed: Result<f32, E>,
485) {
486    match elapsed {
487        Ok(ms) => record_dual_pp_stage_ms(stage, ms),
488        Err(err) => record_dual_pp_timing_drop(DUAL_PP_STAGE_NAMES[stage], &err),
489    }
490}
491
492pub fn dual_pp_timing_dropped() -> usize {
493    DUAL_PP_TIMING_DROPPED.load(Ordering::Relaxed)
494}
495
496/// `(total_nanoseconds, samples)` for wave-A stage0/stage1 then wave-B stage0/stage1.
497pub fn dual_pp_timing_snapshot() -> ([u64; 4], [usize; 4]) {
498    (
499        std::array::from_fn(|i| DUAL_PP_STAGE_NS[i].load(Ordering::Relaxed)),
500        std::array::from_fn(|i| DUAL_PP_STAGE_SAMPLES[i].load(Ordering::Relaxed)),
501    )
502}
503
504pub(crate) struct DualPpStageGuard;
505
506pub(crate) fn enter_dual_pp_stage() -> DualPpStageGuard {
507    let active = DUAL_PP_ACTIVE_STAGES.fetch_add(1, Ordering::AcqRel);
508    if active > 0 {
509        DUAL_PP_OVERLAPS.fetch_add(1, Ordering::Relaxed);
510    }
511    DualPpStageGuard
512}
513
514impl Drop for DualPpStageGuard {
515    fn drop(&mut self) {
516        let active = DUAL_PP_ACTIVE_STAGES.fetch_sub(1, Ordering::AcqRel);
517        debug_assert!(active > 0, "dual PP active-stage counter underflow");
518    }
519}
520
521/// MEMRA_PRIME_PP=0: rollback/A-B seam for the PRIME (chunked prefill) stage split
522/// (lane/pp-leverb 2026-08-08). Default ON — with the ppN door open the chunked prime takes
523/// its own per-stage range walk exactly as the eager/batched/verify steps do. Setting 0 sends
524/// prime back through the unsplit whole-trunk walk. NOTE: unlike batch/dc/graph/spec, prime
525/// keeps NO `refuse_unsplit_if_remote` — its unsplit walk over a sharded placement is the
526/// measured 22% amortized peer-read tax (research/pp-prefill-20260807 anatomy: m=4096
527/// amortizes the weight reads), not the decode 28x cliff, and the unsplit walk IS the
528/// split-vs-unsplit gate's reference arm (`prime-split-gate`), so it must stay callable.
529/// Read per call, never memoized (the gate A/Bs both arms in one process).
530pub fn prime_pp_on() -> bool {
531    std::env::var("MEMRA_PRIME_PP").as_deref() != Ok("0")
532}
533
534/// MEMRA_PRIME_PIPE=0: rollback/A-B seam for the PP-2 PRIME CHUNK PIPELINE
535/// (lane/cx-pipeline-prime 2026-08-08). Default ON when the prime stage split is live;
536/// setting 0 keeps the serial per-chunk stage walk. Read per prime call so the exactness
537/// gate can replay both schedules against one loaded model.
538pub fn prime_pipe_on() -> bool {
539    std::env::var("MEMRA_PRIME_PIPE").as_deref() != Ok("0")
540}
541
542/// SPLIT-LIVENESS COUNTER for the prime stage split: bumped ONCE per prime chunk that
543/// actually executed the per-stage walk. The `prime-split-gate` requires this to ADVANCE
544/// during its split arm — bit-identity of two identical UNSPLIT walks is vacuous, so a gate
545/// that only compared bits would go green while the walker doesn't exist. With the counter,
546/// the gate is RED until the walker lands (the tickinv35 pattern: the gate exists and fails
547/// before the mechanism does). Relaxed ordering: single-threaded host issue, count-only.
548pub static PRIME_SPLIT_CHUNKS: AtomicUsize = AtomicUsize::new(0);
549
550/// Read the split-liveness counter (gate-side).
551pub fn prime_split_chunks() -> usize {
552    PRIME_SPLIT_CHUNKS.load(Ordering::Relaxed)
553}
554
555/// PIPELINE-LIVENESS COUNTER: bumped only when a second PP-2 prime stage enters its layer
556/// walker while the other stage's walker is still active. Step's per-layer router readback
557/// synchronizes the host, so enqueue order alone is not liveness: a single host thread can
558/// call stage 0(N+1) before the stage-1 epilogue and still serialize all trunk computation.
559pub static PRIME_PIPE_OVERLAPS: AtomicUsize = AtomicUsize::new(0);
560
561/// Read the prime-pipeline overlap counter (gate-side).
562pub fn prime_pipe_overlaps() -> usize {
563    PRIME_PIPE_OVERLAPS.load(Ordering::Relaxed)
564}
565
566static PRIME_PIPE_ACTIVE_STAGES: AtomicUsize = AtomicUsize::new(0);
567
568pub(crate) struct PrimePipeStageGuard;
569
570/// Mark one host-driven stage walker active. With PP-2, a transition 1 -> 2 proves the
571/// two device walkers overlap in wall time; exactly one transition is counted per pair.
572pub(crate) fn enter_prime_pipe_stage() -> PrimePipeStageGuard {
573    let active = PRIME_PIPE_ACTIVE_STAGES.fetch_add(1, Ordering::AcqRel);
574    if active > 0 {
575        PRIME_PIPE_OVERLAPS.fetch_add(1, Ordering::Relaxed);
576    }
577    PrimePipeStageGuard
578}
579
580impl Drop for PrimePipeStageGuard {
581    fn drop(&mut self) {
582        let active = PRIME_PIPE_ACTIVE_STAGES.fetch_sub(1, Ordering::AcqRel);
583        debug_assert!(active > 0, "prime pipeline active-stage counter underflow");
584    }
585}
586
587/// Step35 cross-request prime liveness counters (lane/cx-prime-batch, 2026-08-08).
588/// The exactness gate requires BOTH to advance: a successful step35 batch alone is not
589/// sufficient under PP-N if it walked the whole sharded trunk on one stream.
590pub static STEP35_PRIME_BATCHES: AtomicUsize = AtomicUsize::new(0);
591pub static STEP35_PRIME_BATCH_SPLITS: AtomicUsize = AtomicUsize::new(0);
592
593pub fn step35_prime_batches() -> usize {
594    STEP35_PRIME_BATCHES.load(Ordering::Relaxed)
595}
596
597pub fn step35_prime_batch_splits() -> usize {
598    STEP35_PRIME_BATCH_SPLITS.load(Ordering::Relaxed)
599}
600
601/// MEMRA_SPEC_PP=0: rollback/A-B seam for the SPEC VERIFY stage split (pp2-spec 2026-08-06).
602/// Default ON — with the ppN door open the verify forward (`decode_step_t_core_ppn`) takes its
603/// own stage split exactly as the eager and batched steps do. Setting 0 sends verify back through
604/// the unsplit trunk walk, which under a sharded cross-device placement is then caught by
605/// `refuse_unsplit_if_remote` (the 28x peer-read regime) rather than running silently. Exists so
606/// the bit-identity gate can A/B split vs unsplit IN ONE PROCESS against the same loaded weights
607/// — read per verify call, never memoized, for that reason.
608pub fn spec_pp_on() -> bool {
609    std::env::var("MEMRA_SPEC_PP").as_deref() != Ok("0")
610}
611
612/// MEMRA_PP_OVERLAP: alternate the double-buffered boundary slots per step (the
613/// pipelining seed). Scheduling structure only, never math. Read per step so gates can
614/// A/B in-process.
615///
616/// Unset follows the dual-PP mode (owner flip 2026-08-11): `Auto` resolves ON — the naked
617/// serve path is the box1 re-gate's dual arm (MEMRA_DUAL_PP=1 MEMRA_PP_OVERLAP=1,
618/// 929/929 golden, 0/9123 slot collisions). `Off` resolves OFF so MEMRA_DUAL_PP=0 alone
619/// restores the exact pre-flip serial naked path. `Forced` resolves OFF so the binding
620/// single-slot refusal of the explicit pre-flip request stays reachable — the
621/// decode-batch-gate negative cell pins and asserts precisely that combination.
622pub fn pp2_overlap() -> bool {
623    pp2_overlap_resolve(
624        std::env::var("MEMRA_PP_OVERLAP").ok().as_deref(),
625        dual_pp_mode(),
626    )
627}
628
629/// Pure resolution for MEMRA_PP_OVERLAP, split from the env read for the flip
630/// regression tests.
631pub fn pp2_overlap_resolve(v: Option<&str>, mode: DualPpMode) -> bool {
632    match v {
633        Some("1") => true,
634        Some(_) => false,
635        None => mode == DualPpMode::Auto,
636    }
637}
638
639/// Broken-peer escape hatch: stage-boundary activations travel through page-locked host
640/// memory instead of `cudaMemcpyPeerAsync`. Default OFF; captured when `PpNRt` is built.
641pub fn pp_host_bounce_on() -> bool {
642    matches!(std::env::var("MEMRA_PP_HOST_BOUNCE").as_deref(), Ok("1"))
643}
644
645/// True when host bounce is the live transport for a sharded cross-device placement.
646/// Callers use this to close paths that still peer-read non-boundary state.
647pub fn pp_host_bounce_active() -> bool {
648    (pp_host_bounce_on() || PEER_RUNTIME_HOST_BOUNCE.load(Ordering::Acquire))
649        && pp_sharded_cross_device()
650}
651
652/// M2 increment 2 rollback seam: MEMRA_PP_SHARD=0 = the M1 bring-up placement (all
653/// weights upload through the primary engine; remote stages peer-read). Default ON —
654/// under MEMRA_PP_DEVICES each stage's layer range uploads through its own engine.
655pub fn pp_shard_off() -> bool {
656    matches!(std::env::var("MEMRA_PP_SHARD").as_deref(), Ok("0"))
657}
658
659/// Raw `MEMRA_PP_DEVICES` (parsed/validated at PpNRt build — a bad string must fail the
660/// decode step loudly, never silently fall back to same-device and fake a gate PASS).
661fn pp2_devices_env() -> Option<String> {
662    std::env::var("MEMRA_PP_DEVICES")
663        .ok()
664        .filter(|v| !v.is_empty())
665}
666
667static WARNED_BAD: AtomicBool = AtomicBool::new(false);
668fn warn_bad_once(msg: &str) {
669    if !WARNED_BAD.swap(true, Ordering::Relaxed) {
670        eprintln!("[pp] {msg}");
671    }
672}
673
674static WARNED_UNWIRED: AtomicBool = AtomicBool::new(false);
675/// One-time notice when the door is set but the executing path has no pp arm
676/// (M2 wires the generic eager decode at any N and the gemma4 eager arm at N=2).
677pub fn warn_unwired_once(path: &str) {
678    let open = std::env::var("MEMRA_PP_STAGES")
679        .map(|v| !v.is_empty() && v != "0" && v != "1")
680        .unwrap_or(false);
681    if open && !WARNED_UNWIRED.swap(true, Ordering::Relaxed) {
682        eprintln!("[pp] MEMRA_PP_STAGES set but `{path}` has no pp arm at this N; running unsplit");
683    }
684}
685
686// ======================================================================================
687//  PpNRt: the M2 transport runtime (per-stage streams, per-boundary events + slots)
688// ======================================================================================
689
690/// One pipeline stage's execution home: device, context, launch stream, and (for a stage
691/// remote to the primary engine's device) a dedicated Engine in that device's primary
692/// context (CUmodules are per-context).
693pub struct StageRt {
694    pub dev: usize,
695    pub ctx: Arc<CudaContext>,
696    pub stream: Arc<CudaStream>,
697    /// `Some` only when `dev` differs from the primary engine's device.
698    engine: Option<Engine>,
699}
700
701/// One boundary slot: a persistent RX-side buffer + its TX/RX completion events.
702/// PERSISTENT because the buffer is written by the TX stage's stream and read by the RX
703/// stage's: a per-step alloc/free would enqueue the free on ONE stream while the other
704/// might still be reading (the cross-stream free hazard) — a never-freed slot cannot race.
705struct BoundarySlot {
706    buf: Mutex<Option<CudaSlice<f32>>>,
707    /// Recorded on the TX stage's stream after the TX copy; RX waits on it. Created in
708    /// the TX stage's context (cuEventRecord requires event ctx == stream ctx).
709    ev_tx: CudaEvent,
710    /// Recorded on the RX stage's stream after the RX copy; the NEXT TX into this slot
711    /// waits on it (write-after-read guard). Created in the RX stage's context. Waiting
712    /// on a never-recorded event is a defined no-op, so step 0 needs no special case.
713    ev_rx: CudaEvent,
714}
715
716/// Boundary b sits between stage b (TX) and stage b+1 (RX). Two slots, alternating per
717/// step under MEMRA_PP_OVERLAP=1 (each boundary counts its own steps — a decode step
718/// crosses every boundary exactly once, so the counters stay in lockstep).
719struct BoundaryRt {
720    slots: [BoundarySlot; 2],
721    step: AtomicUsize,
722    /// true iff stage b and stage b+1 live on different devices (peer transport).
723    cross: bool,
724}
725
726#[derive(Clone, Copy, Debug, PartialEq, Eq)]
727enum BoundaryTransport {
728    Local,
729    Peer,
730    HostBounce,
731}
732
733#[derive(Clone, Copy)]
734struct BoundaryPath {
735    boundary: usize,
736    src_stage: usize,
737    dst_stage: usize,
738    transport: BoundaryTransport,
739}
740
741fn boundary_transport(cross: bool, host_bounce: bool) -> BoundaryTransport {
742    match (cross, host_bounce) {
743        (false, _) => BoundaryTransport::Local,
744        (true, false) => BoundaryTransport::Peer,
745        (true, true) => BoundaryTransport::HostBounce,
746    }
747}
748
749const PEER_PROBE_FIXED_BYTES: usize = 16 * 1024;
750const PEER_PROBE_TOKEN_WIDTHS: [usize; 4] = [1, 8, 16, crate::cache::PRIME_CHUNK_MAX_TOKENS];
751
752/// Native cross-device boundary copies between low-frequency runtime integrity probes.
753/// Fixed rather than operator-tunable: this is a safety gate, not a performance experiment.
754pub const PEER_RUNTIME_PROBE_INTERVAL_COPIES: u64 = 8 * 1024;
755/// One complete runtime width rotation. The maximum-chunk rung runs once per cycle.
756pub const PEER_RUNTIME_PROBE_CYCLE_COPIES: u64 =
757    PEER_RUNTIME_PROBE_INTERVAL_COPIES * PEER_PROBE_TOKEN_WIDTHS.len() as u64;
758/// Consecutive runnable probe intervals that may be blocked by live speculative UVA state before
759/// integrity coverage becomes explicitly degraded. Four intervals are one full width rotation.
760pub const PEER_RUNTIME_PROBE_DEFERRAL_BOUND_INTERVALS: u64 = PEER_PROBE_TOKEN_WIDTHS.len() as u64;
761/// Maximum measured owner-thread wall cost that may remain on an interactive scheduler boundary.
762const PEER_RUNTIME_PROBE_BUDGET_NS: u64 = 5_000_000;
763
764pub const PEER_PROBE_REQUIRED_REFUSAL: &str = "PP bring-up refused: MEMRA_PEER_PROBE=0 cannot authorize native peer transport for a \
765     sharded cross-device placement while MEMRA_PP_HOST_BOUNCE!=1; leave MEMRA_PEER_PROBE \
766     enabled or set MEMRA_PP_HOST_BOUNCE=1";
767
768#[derive(Clone, Copy, Debug, PartialEq, Eq)]
769pub enum PeerProbeStartupPolicy {
770    Allowed,
771    BypassedWithHostBounce,
772}
773
774/// Pure startup policy so unit tests and kernel-check pin the entire refusal matrix without
775/// mutating process-global environment variables.
776pub fn peer_probe_startup_policy(
777    probe_on: bool,
778    sharded_cross_device: bool,
779    host_bounce: bool,
780) -> Result<PeerProbeStartupPolicy, &'static str> {
781    match (probe_on, sharded_cross_device, host_bounce) {
782        (false, true, false) => Err(PEER_PROBE_REQUIRED_REFUSAL),
783        (false, true, true) => Ok(PeerProbeStartupPolicy::BypassedWithHostBounce),
784        _ => Ok(PeerProbeStartupPolicy::Allowed),
785    }
786}
787
788static PEER_PROBE_BYPASSED: AtomicU64 = AtomicU64::new(0);
789static PEER_BOUNDARY_COPIES: AtomicU64 = AtomicU64::new(0);
790static PEER_RUNTIME_PROBES: AtomicU64 = AtomicU64::new(0);
791static PEER_RUNTIME_PROBE_FAILURES: AtomicU64 = AtomicU64::new(0);
792static PEER_RUNTIME_PROBE_DEFERRED: AtomicU64 = AtomicU64::new(0);
793static PEER_RUNTIME_PROBE_INTEGRITY_DEGRADED: AtomicBool = AtomicBool::new(false);
794static PEER_RUNTIME_PROBE_FAILED: AtomicBool = AtomicBool::new(false);
795static PEER_RUNTIME_HOST_BOUNCE: AtomicBool = AtomicBool::new(false);
796static PEER_RUNTIME_NEXT_PROBE_COPY: [AtomicU64; PEER_PROBE_TOKEN_WIDTHS.len()] = [
797    AtomicU64::new(PEER_RUNTIME_PROBE_INTERVAL_COPIES),
798    AtomicU64::new(2 * PEER_RUNTIME_PROBE_INTERVAL_COPIES),
799    AtomicU64::new(3 * PEER_RUNTIME_PROBE_INTERVAL_COPIES),
800    AtomicU64::new(4 * PEER_RUNTIME_PROBE_INTERVAL_COPIES),
801];
802static PEER_RUNTIME_PROBE_MAX_COST_NS: [AtomicU64; PEER_PROBE_TOKEN_WIDTHS.len()] = [
803    AtomicU64::new(0),
804    AtomicU64::new(0),
805    AtomicU64::new(0),
806    AtomicU64::new(0),
807];
808
809#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
810pub struct PeerProbeMetrics {
811    pub bypassed: u64,
812    pub boundary_copies: u64,
813    pub runtime_probes: u64,
814    pub runtime_failures: u64,
815    pub deferred_total: u64,
816    pub integrity_degraded: bool,
817    pub degraded_to_host_bounce: bool,
818}
819
820pub fn peer_probe_metrics() -> PeerProbeMetrics {
821    PeerProbeMetrics {
822        bypassed: PEER_PROBE_BYPASSED.load(Ordering::Relaxed),
823        boundary_copies: PEER_BOUNDARY_COPIES.load(Ordering::Relaxed),
824        runtime_probes: PEER_RUNTIME_PROBES.load(Ordering::Relaxed),
825        runtime_failures: PEER_RUNTIME_PROBE_FAILURES.load(Ordering::Relaxed),
826        deferred_total: PEER_RUNTIME_PROBE_DEFERRED.load(Ordering::Relaxed),
827        integrity_degraded: PEER_RUNTIME_PROBE_INTEGRITY_DEGRADED.load(Ordering::Acquire),
828        degraded_to_host_bounce: PEER_RUNTIME_HOST_BOUNCE.load(Ordering::Acquire),
829    }
830}
831
832#[derive(Clone, Copy, Debug, PartialEq, Eq)]
833pub enum RuntimePeerProbeStatus {
834    NotRun,
835    Deferred,
836    Passed,
837    DegradedToHostBounce,
838}
839
840impl RuntimePeerProbeStatus {
841    pub fn ran(self) -> bool {
842        matches!(self, Self::Passed | Self::DegradedToHostBounce)
843    }
844}
845
846fn publish_runtime_peer_probe_deferral(
847    deferred_total: &AtomicU64,
848    integrity_degraded: &AtomicBool,
849    intervals: u64,
850    bound_reached: bool,
851) {
852    deferred_total.fetch_add(intervals, Ordering::Relaxed);
853    if bound_reached {
854        integrity_degraded.store(true, Ordering::Release);
855    }
856}
857
858/// Publish newly observed copy-count intervals where a runnable peer probe was blocked by live
859/// speculative UVA state. The worker coalesces scheduler polls before calling this function.
860pub fn record_runtime_peer_probe_deferral(intervals: u64, bound_reached: bool) {
861    publish_runtime_peer_probe_deferral(
862        &PEER_RUNTIME_PROBE_DEFERRED,
863        &PEER_RUNTIME_PROBE_INTEGRITY_DEGRADED,
864        intervals,
865        bound_reached,
866    );
867}
868
869/// A completed native probe or validated transport failover restores an explicit integrity state.
870pub fn clear_runtime_peer_probe_integrity_degraded() {
871    PEER_RUNTIME_PROBE_INTEGRITY_DEGRADED.store(false, Ordering::Release);
872}
873
874fn runtime_peer_probe_idle_only(width_index: usize, measured_cost_ns: u64) -> bool {
875    width_index + 1 == PEER_PROBE_TOKEN_WIDTHS.len()
876        || measured_cost_ns > PEER_RUNTIME_PROBE_BUDGET_NS
877}
878
879/// Pick the oldest runnable per-width deadline. Idle-only overdue work is skipped rather than
880/// blocking later cheap deadlines, so the small integrity ladder keeps its copy-count cadence.
881fn runtime_peer_probe_candidate(
882    copies: u64,
883    next_probe_copy: [u64; PEER_PROBE_TOKEN_WIDTHS.len()],
884    measured_cost_ns: [u64; PEER_PROBE_TOKEN_WIDTHS.len()],
885    scheduler_idle: bool,
886) -> Option<(usize, usize)> {
887    let mut selected: Option<(usize, u64)> = None;
888    for width_index in 0..PEER_PROBE_TOKEN_WIDTHS.len() {
889        let due = next_probe_copy[width_index];
890        if copies < due
891            || (!scheduler_idle
892                && runtime_peer_probe_idle_only(width_index, measured_cost_ns[width_index]))
893        {
894            continue;
895        }
896        if selected.is_none_or(|(_, selected_due)| due < selected_due) {
897            selected = Some((width_index, due));
898        }
899    }
900    selected.map(|(width_index, _)| (width_index, PEER_PROBE_TOKEN_WIDTHS[width_index]))
901}
902
903/// Advance a late per-width deadline to the first future cycle. Missed idle opportunities
904/// collapse into one probe instead of producing an owner-thread catch-up burst.
905fn runtime_peer_probe_next_copy(due: u64, copies: u64) -> u64 {
906    let cycles = copies.saturating_sub(due) / PEER_RUNTIME_PROBE_CYCLE_COPIES + 1;
907    due.saturating_add(PEER_RUNTIME_PROBE_CYCLE_COPIES.saturating_mul(cycles))
908}
909
910/// Fail closed before arming the fallback, then publish host bounce only after its staging check
911/// succeeds. The two atomics are parameters so unit tests never mutate process-global state.
912fn latch_runtime_host_bounce<E>(
913    native_failed: &AtomicBool,
914    degraded_to_host_bounce: &AtomicBool,
915    arm_and_validate: impl FnOnce() -> Result<(), E>,
916) -> Result<(), E> {
917    native_failed.store(true, Ordering::Release);
918    arm_and_validate()?;
919    degraded_to_host_bounce.store(true, Ordering::Release);
920    Ok(())
921}
922
923fn peer_probe_on() -> bool {
924    std::env::var("MEMRA_PEER_PROBE").as_deref() != Ok("0")
925}
926
927#[derive(Clone, Copy, Debug, PartialEq, Eq)]
928enum PeerProbeDecision {
929    Clean,
930    ProceedWithHostBounce { mismatches: usize },
931}
932
933fn peer_probe_mismatch_count(expected: &[u8], readback: &[u8]) -> usize {
934    expected
935        .iter()
936        .zip(readback)
937        .filter(|(a, b)| a != b)
938        .count()
939        + expected.len().abs_diff(readback.len())
940}
941
942fn peer_probe_decision(
943    expected: &[u8],
944    readback: &[u8],
945    host_bounce: bool,
946) -> Result<PeerProbeDecision, String> {
947    let mismatches = peer_probe_mismatch_count(expected, readback);
948    if mismatches == 0 {
949        Ok(PeerProbeDecision::Clean)
950    } else if host_bounce {
951        Ok(PeerProbeDecision::ProceedWithHostBounce { mismatches })
952    } else {
953        Err(format!("{mismatches} mismatched byte(s)"))
954    }
955}
956
957fn peer_probe_pattern(bytes: usize, boundary: usize, src_dev: usize, dst_dev: usize) -> Vec<u8> {
958    let mut state = 0xD1B5_4A32_D192_ED03u64
959        ^ (bytes as u64).rotate_left(7)
960        ^ (boundary as u64).rotate_left(19)
961        ^ (src_dev as u64).rotate_left(31)
962        ^ (dst_dev as u64).rotate_left(43);
963    (0..bytes)
964        .map(|_| {
965            state ^= state << 13;
966            state ^= state >> 7;
967            state ^= state << 17;
968            state as u8
969        })
970        .collect()
971}
972
973fn peer_probe_bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
974    assert_eq!(bytes.len() % std::mem::size_of::<f32>(), 0);
975    bytes
976        .chunks_exact(std::mem::size_of::<f32>())
977        .map(|chunk| f32::from_bits(u32::from_ne_bytes(chunk.try_into().unwrap())))
978        .collect()
979}
980
981fn peer_probe_f32_to_bytes(values: &[f32]) -> Vec<u8> {
982    values
983        .iter()
984        .flat_map(|value| value.to_bits().to_ne_bytes())
985        .collect()
986}
987
988/// A legacy `cuMemAlloc` buffer used only by the boot probe. Unlike memra's normal
989/// stream-ordered allocations, it becomes peer-visible through `cuCtxEnablePeerAccess`
990/// without requiring the default-pool grants that deliberately happen after the probe.
991struct PeerProbeBuffer {
992    ctx: Arc<CudaContext>,
993    ptr: cudarc::driver::sys::CUdeviceptr,
994}
995
996impl PeerProbeBuffer {
997    fn new(ctx: &Arc<CudaContext>, bytes: usize) -> Result<Self, Box<dyn std::error::Error>> {
998        ctx.bind_to_thread()?;
999        let ptr = unsafe { cudarc::driver::result::malloc_sync(bytes)? };
1000        Ok(Self {
1001            ctx: ctx.clone(),
1002            ptr,
1003        })
1004    }
1005}
1006
1007impl Drop for PeerProbeBuffer {
1008    fn drop(&mut self) {
1009        if self.ctx.bind_to_thread().is_ok() {
1010            let _ = unsafe { cudarc::driver::result::free_sync(self.ptr) };
1011        }
1012    }
1013}
1014
1015fn peer_probe_copy(
1016    src: &StageRt,
1017    dst: &StageRt,
1018    expected: &[u8],
1019) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
1020    let bytes = expected.len();
1021    let src_buf = PeerProbeBuffer::new(&src.ctx, bytes)?;
1022    unsafe {
1023        cudarc::driver::result::memcpy_htod_sync(src_buf.ptr, expected)?;
1024    }
1025
1026    let dst_buf = PeerProbeBuffer::new(&dst.ctx, bytes)?;
1027    let poison: Vec<u8> = expected.iter().map(|b| !b).collect();
1028    unsafe {
1029        cudarc::driver::result::memcpy_htod_sync(dst_buf.ptr, &poison)?;
1030    }
1031
1032    src.ctx.bind_to_thread()?;
1033    unsafe {
1034        cudarc::driver::result::memcpy_peer_async(
1035            dst.ctx.cu_ctx(),
1036            dst_buf.ptr,
1037            src.ctx.cu_ctx(),
1038            src_buf.ptr,
1039            bytes,
1040            src.stream.cu_stream(),
1041        )?;
1042    }
1043    src.stream.synchronize()?;
1044
1045    dst.ctx.bind_to_thread()?;
1046    let mut readback = vec![0u8; bytes];
1047    unsafe {
1048        cudarc::driver::result::memcpy_dtoh_sync(&mut readback, dst_buf.ptr)?;
1049    }
1050    Ok(readback)
1051}
1052
1053fn run_peer_probe_pass(
1054    stages: &[StageRt],
1055    peer_capable: &[(usize, usize)],
1056    host_bounce: bool,
1057    label: &str,
1058    bytes: usize,
1059) -> Result<(), Box<dyn std::error::Error>> {
1060    if bytes == 0 {
1061        return Err(format!("PP peer byte-integrity probe {label} size is zero").into());
1062    }
1063    let started = std::time::Instant::now();
1064    let mut copies = 0usize;
1065    let mut skipped = 0usize;
1066    let mut total_mismatches = 0usize;
1067
1068    for boundary in 0..stages.len() - 1 {
1069        if stages[boundary].dev == stages[boundary + 1].dev {
1070            continue;
1071        }
1072        for (src_idx, dst_idx) in [(boundary, boundary + 1), (boundary + 1, boundary)] {
1073            let src = &stages[src_idx];
1074            let dst = &stages[dst_idx];
1075            if !peer_capable.contains(&(src.dev, dst.dev)) {
1076                if host_bounce {
1077                    skipped += 1;
1078                    eprintln!(
1079                        "[pp] peer byte-integrity probe SKIP: boundary={boundary} \
1080                         dev{}->dev{} label={label} bytes={bytes} (peer capability unavailable; \
1081                         MEMRA_PP_HOST_BOUNCE=1 remains fail-safe)",
1082                        src.dev, dst.dev,
1083                    );
1084                    continue;
1085                }
1086                return Err(format!(
1087                    "PP peer byte-integrity probe cannot run boundary={boundary} \
1088                     dev{}->dev{}: peer access was not enabled",
1089                    src.dev, dst.dev,
1090                )
1091                .into());
1092            }
1093
1094            let expected = peer_probe_pattern(bytes, boundary, src.dev, dst.dev);
1095            let readback = match peer_probe_copy(src, dst, &expected) {
1096                Ok(readback) => readback,
1097                Err(err) if host_bounce => {
1098                    skipped += 1;
1099                    eprintln!(
1100                        "[pp] peer byte-integrity probe ERROR: boundary={boundary} \
1101                         dev{}->dev{} label={label} bytes={bytes}: {err}; \
1102                         MEMRA_PP_HOST_BOUNCE=1, proceeding on the host-staged path",
1103                        src.dev, dst.dev,
1104                    );
1105                    continue;
1106                }
1107                Err(err) => {
1108                    return Err(format!(
1109                        "PP peer byte-integrity probe FAILED: boundary={boundary} \
1110                         dev{}->dev{} label={label} bytes={bytes}: {err}; refusing native P2P \
1111                         (set MEMRA_PP_HOST_BOUNCE=1 to use the host-staged path; \
1112                         MEMRA_PEER_PROBE=0 cannot authorize sharded native peer transport)",
1113                        src.dev, dst.dev,
1114                    )
1115                    .into());
1116                }
1117            };
1118            copies += 1;
1119            match peer_probe_decision(&expected, &readback, host_bounce) {
1120                Ok(PeerProbeDecision::Clean) => {}
1121                Ok(PeerProbeDecision::ProceedWithHostBounce { mismatches }) => {
1122                    total_mismatches += mismatches;
1123                    eprintln!(
1124                        "[pp] peer byte-integrity probe CORRUPTION: boundary={boundary} \
1125                         dev{}->dev{} label={label} bytes={bytes} mismatches={mismatches}; \
1126                         MEMRA_PP_HOST_BOUNCE=1, proceeding on the host-staged path",
1127                        src.dev, dst.dev,
1128                    );
1129                }
1130                Err(mismatch) => {
1131                    return Err(format!(
1132                        "PP peer byte-integrity probe FAILED: boundary={boundary} \
1133                         dev{}->dev{} label={label} bytes={bytes}: {mismatch}; refusing native \
1134                         P2P (set MEMRA_PP_HOST_BOUNCE=1 to use the host-staged path; \
1135                         MEMRA_PEER_PROBE=0 cannot authorize sharded native peer transport)",
1136                        src.dev, dst.dev,
1137                    )
1138                    .into());
1139                }
1140            }
1141        }
1142    }
1143
1144    let status = if total_mismatches > 0 {
1145        "BOUNCE"
1146    } else if skipped > 0 && copies > 0 {
1147        "PARTIAL"
1148    } else if skipped > 0 {
1149        "SKIP"
1150    } else {
1151        "PASS"
1152    };
1153    eprintln!(
1154        "[pp] peer byte-integrity probe {}: label={label} bytes={bytes} copies={copies} \
1155         skipped={skipped} mismatches={total_mismatches} elapsed_ms={:.3}",
1156        status,
1157        started.elapsed().as_secs_f64() * 1e3,
1158    );
1159    Ok(())
1160}
1161
1162fn host_bounce_capacity(n_embd: usize) -> Result<(usize, usize), String> {
1163    if n_embd == 0 {
1164        return Err("MEMRA_PP_HOST_BOUNCE needs non-zero model n_embd".into());
1165    }
1166    let elems = n_embd
1167        .checked_mul(crate::cache::PRIME_CHUNK_MAX_TOKENS)
1168        .ok_or_else(|| format!("host-bounce element count overflows for n_embd={n_embd}"))?;
1169    let bytes = elems
1170        .checked_mul(std::mem::size_of::<f32>())
1171        .ok_or_else(|| format!("host-bounce byte count overflows for n_embd={n_embd}"))?;
1172    Ok((elems, bytes))
1173}
1174
1175/// One bidirectional-DMA staging allocation. `CU_MEMHOSTALLOC_PORTABLE` matters here: the
1176/// D2H producer and H2D consumer are in distinct CUDA primary contexts. Cacheable memory is
1177/// intentional (rather than cudarc's write-combined pinned slice) because this allocation is
1178/// the destination of D2H as well as the source of H2D.
1179struct PinnedHostBounce {
1180    ptr: *mut f32,
1181    len: usize,
1182}
1183
1184unsafe impl Send for PinnedHostBounce {}
1185unsafe impl Sync for PinnedHostBounce {}
1186
1187impl PinnedHostBounce {
1188    fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
1189        let bytes = len
1190            .checked_mul(std::mem::size_of::<f32>())
1191            .ok_or("host-bounce pinned allocation size overflow")?;
1192        let ptr = unsafe {
1193            cudarc::driver::result::malloc_host(
1194                bytes,
1195                cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE,
1196            )?
1197        } as *mut f32;
1198        if ptr.is_null() {
1199            return Err("cuMemHostAlloc returned a null host-bounce pointer".into());
1200        }
1201        Ok(Self { ptr, len })
1202    }
1203
1204    fn prefix(&self, n: usize) -> &[f32] {
1205        assert!(
1206            n <= self.len,
1207            "host-bounce source {n} > capacity {}",
1208            self.len
1209        );
1210        unsafe { std::slice::from_raw_parts(self.ptr, n) }
1211    }
1212
1213    fn prefix_mut(&mut self, n: usize) -> &mut [f32] {
1214        assert!(
1215            n <= self.len,
1216            "host-bounce destination {n} > capacity {}",
1217            self.len
1218        );
1219        unsafe { std::slice::from_raw_parts_mut(self.ptr, n) }
1220    }
1221}
1222
1223impl Drop for PinnedHostBounce {
1224    fn drop(&mut self) {
1225        let _ = unsafe { cudarc::driver::result::free_host(self.ptr.cast()) };
1226    }
1227}
1228
1229struct HostBounceRt {
1230    n_embd: usize,
1231    capacity: usize,
1232    slots: Vec<Option<[Mutex<PinnedHostBounce>; 2]>>,
1233}
1234
1235impl HostBounceRt {
1236    fn new(n_embd: usize, boundaries: &[BoundaryRt]) -> Result<Self, Box<dyn std::error::Error>> {
1237        let (capacity, _) = host_bounce_capacity(n_embd)?;
1238        let mut slots = Vec::with_capacity(boundaries.len());
1239        for boundary in boundaries {
1240            slots.push(if boundary.cross {
1241                Some([
1242                    Mutex::new(PinnedHostBounce::new(capacity)?),
1243                    Mutex::new(PinnedHostBounce::new(capacity)?),
1244                ])
1245            } else {
1246                None
1247            });
1248        }
1249        Ok(Self {
1250            n_embd,
1251            capacity,
1252            slots,
1253        })
1254    }
1255
1256    fn slot(
1257        &self,
1258        boundary: usize,
1259        slot: usize,
1260    ) -> Result<&Mutex<PinnedHostBounce>, Box<dyn std::error::Error>> {
1261        self.slots
1262            .get(boundary)
1263            .and_then(Option::as_ref)
1264            .and_then(|slots| slots.get(slot))
1265            .ok_or_else(|| format!("host-bounce slot {boundary}:{slot} is not initialized").into())
1266    }
1267}
1268
1269pub struct PpNRt {
1270    stages: Vec<StageRt>,
1271    boundaries: Vec<BoundaryRt>,
1272    /// true iff ANY boundary crosses devices.
1273    cross_any: bool,
1274    /// Startup selection captured at runtime construction. A runtime probe failure may promote
1275    /// the process-wide one-way host-bounce latch without mutating this value.
1276    host_bounce: bool,
1277    /// Boot-time peer validation is default-on; `MEMRA_PEER_PROBE=0` is diagnostics-only.
1278    peer_probe: bool,
1279    /// Directed device pairs for which `cuDeviceCanAccessPeer` succeeded.
1280    peer_capable: Vec<(usize, usize)>,
1281    /// Sticky one-time model-width probe result. The value is the one-row geometry byte count.
1282    peer_probe_geometry: OnceLock<Result<usize, String>>,
1283    /// Lazily allocated after the authoritative model width is known at cache creation.
1284    bounce: OnceLock<Result<HostBounceRt, String>>,
1285    /// Dedicated readback stream in the LAST stage's context (deferred logits D2H —
1286    /// waiting there instead of on the compute stream keeps later tokens enqueuable).
1287    readback: Arc<CudaStream>,
1288}
1289
1290/// M1 name kept alive for external callers (`pp-transport-smoke`, receipts, docs).
1291pub type Pp2Rt = PpNRt;
1292
1293static RTN: OnceLock<Result<PpNRt, String>> = OnceLock::new();
1294
1295impl PpNRt {
1296    /// The process-wide transport runtime, built on first use against the primary engine.
1297    /// The stage count + device map freeze at first build (one config per process — gates
1298    /// run one placement per invocation). Build errors are sticky and loud.
1299    pub fn get(e: &Engine) -> Result<&'static PpNRt, Box<dyn std::error::Error>> {
1300        RTN.get_or_init(|| Self::build(e).map_err(|err| err.to_string()))
1301            .as_ref()
1302            .map_err(|s| -> Box<dyn std::error::Error> { s.clone().into() })
1303    }
1304
1305    fn build(e: &Engine) -> Result<PpNRt, Box<dyn std::error::Error>> {
1306        let primary_dev = e.ctx().ordinal();
1307        // Stage count: MEMRA_PP_DEVICES length wins when set (it IS the placement);
1308        // else MEMRA_PP_STAGES; else 2 (the M1 default — pp-transport-smoke runs doorless).
1309        let devices: Vec<usize> =
1310            match pp2_devices_env() {
1311                Some(s) => {
1312                    let parts: Result<Vec<usize>, _> =
1313                        s.split(',').map(|p| p.trim().parse::<usize>()).collect();
1314                    match parts {
1315                        Ok(v) if v.len() >= 2 => v,
1316                        _ => return Err(format!(
1317                            "MEMRA_PP_DEVICES={s} unparseable (want <d0>,..,<dN-1> e.g. 0,1,2,3)"
1318                        )
1319                        .into()),
1320                    }
1321                }
1322                None => {
1323                    let n_st = std::env::var("MEMRA_PP_STAGES")
1324                        .ok()
1325                        .and_then(|v| v.parse::<usize>().ok())
1326                        .filter(|&n| n >= 2)
1327                        .unwrap_or(2);
1328                    vec![primary_dev; n_st]
1329                }
1330            };
1331        if let Ok(v) = std::env::var("MEMRA_PP_STAGES") {
1332            if let Ok(n) = v.parse::<usize>() {
1333                if n >= 2 && n != devices.len() {
1334                    return Err(format!(
1335                        "MEMRA_PP_DEVICES lists {} devices but MEMRA_PP_STAGES={n} — \
1336                         refusing an ambiguous placement",
1337                        devices.len()
1338                    )
1339                    .into());
1340                }
1341            }
1342        }
1343        let n_st = devices.len();
1344        let cross_any = devices.iter().any(|&d| d != devices[0]);
1345        let host_bounce = pp_host_bounce_on();
1346        let peer_probe = peer_probe_on();
1347        let sharded_cross_device = cross_any && !pp_shard_off();
1348        if host_bounce && cross_any {
1349            if pp_shard_off() {
1350                return Err(
1351                    "MEMRA_PP_HOST_BOUNCE=1 refuses MEMRA_PP_SHARD=0: the boundary can bounce, \
1352                     but remote stages would still peer-read primary-device weights"
1353                        .into(),
1354                );
1355            }
1356            if devices.last().copied() != Some(primary_dev) {
1357                return Err(format!(
1358                    "MEMRA_PP_HOST_BOUNCE=1 requires the primary engine on the last/head stage \
1359                     (primary dev{primary_dev}, placement {devices:?}); otherwise returned \
1360                     logits/hidden state remain peer reads"
1361                )
1362                .into());
1363            }
1364        }
1365        let peer_probe_policy =
1366            peer_probe_startup_policy(peer_probe, sharded_cross_device, host_bounce)?;
1367        if peer_probe_policy == PeerProbeStartupPolicy::BypassedWithHostBounce {
1368            PEER_PROBE_BYPASSED.fetch_add(1, Ordering::Relaxed);
1369            eprintln!(
1370                "[pp] SECURITY RED: peer_probe_bypassed: MEMRA_PEER_PROBE=0 on a sharded \
1371                 cross-device placement; MEMRA_PP_HOST_BOUNCE=1 is the only enabled transport"
1372            );
1373        }
1374
1375        // Validate every placement ordinal in both transports. Native peer transport requires
1376        // access both ways. Host bounce remains usable without it, but records any capable pairs
1377        // so the byte probe can still diagnose a lying peer path before selecting the fallback.
1378        let mut used: Vec<usize> = devices.clone();
1379        used.push(primary_dev);
1380        used.sort_unstable();
1381        used.dedup();
1382        let mut peer_capable = Vec::new();
1383        if used.len() > 1 {
1384            let n = cudarc::driver::result::device::get_count()? as usize;
1385            for &d in &used {
1386                if d >= n {
1387                    return Err(format!(
1388                        "MEMRA_PP_DEVICES={devices:?} but only {n} CUDA device(s) present"
1389                    )
1390                    .into());
1391                }
1392            }
1393            if !host_bounce || peer_probe {
1394                for &a in &used {
1395                    for &b in &used {
1396                        if a == b {
1397                            continue;
1398                        }
1399                        let da = cudarc::driver::result::device::get(a as i32)?;
1400                        let db = cudarc::driver::result::device::get(b as i32)?;
1401                        let mut can: i32 = 0;
1402                        let capability = unsafe {
1403                            cudarc::driver::sys::cuDeviceCanAccessPeer(&mut can, da, db).result()
1404                        };
1405                        if let Err(err) = capability {
1406                            if host_bounce {
1407                                eprintln!(
1408                                    "[pp] peer byte-integrity probe capability query failed for \
1409                                     dev{a}->dev{b}: {err}; MEMRA_PP_HOST_BOUNCE=1 remains active"
1410                                );
1411                                continue;
1412                            }
1413                            return Err(err.into());
1414                        }
1415                        if can == 0 {
1416                            if !host_bounce {
1417                                return Err(format!(
1418                                    "device {a} cannot peer-access device {b} \
1419                                     (cuDeviceCanAccessPeer=0); ppN cross-device needs P2P — \
1420                                     refusing a silently-staged path"
1421                                )
1422                                .into());
1423                            }
1424                        } else {
1425                            peer_capable.push((a, b));
1426                        }
1427                    }
1428                }
1429            }
1430        }
1431
1432        // PER-STAGE ENGINE ISOLATION (2026-08-02 singledev pipelined find): Engine owns
1433        // lazily-grown SHARED scratch pools (fa_part_pool, fa_vf16_scratch, argmax
1434        // partials, ...) that are stable-pointer by design — safe on one stream, a data
1435        // race the moment two stage streams run concurrently through the SAME Engine
1436        // (deferred readback, >=2 tokens in flight: token t+1's stage-0 fa memsets the
1437        // partials while token t's stage-s fa still reads them — the nondeterministic
1438        // all-logits divergence; cross-device arms were immune because remote stages
1439        // already got their own Engine). Every stage s>0 gets its OWN Engine even on the
1440        // primary device: same CUcontext (primary retain), so the per-context CUmodule
1441        // cache makes it cheap; scratch pools are per-Engine, so stages never share.
1442        // Stage 0 keeps the primary engine (single-threaded host issue: the only
1443        // concurrent user of `e` during a pp walk is stage 0 itself).
1444        let mk_stage = |dev: usize, s: usize| -> Result<StageRt, Box<dyn std::error::Error>> {
1445            if dev == primary_dev && s == 0 {
1446                let ctx = e.ctx().clone();
1447                let stream = ctx.new_stream()?;
1448                Ok(StageRt {
1449                    dev,
1450                    ctx,
1451                    stream,
1452                    engine: None,
1453                })
1454            } else {
1455                let eng = Engine::new(dev)?;
1456                let ctx = eng.ctx().clone();
1457                let stream = ctx.new_stream()?;
1458                Ok(StageRt {
1459                    dev,
1460                    ctx,
1461                    stream,
1462                    engine: Some(eng),
1463                })
1464            }
1465        };
1466        let mut stages = Vec::with_capacity(n_st);
1467        for (s, &d) in devices.iter().enumerate() {
1468            stages.push(mk_stage(d, s)?);
1469        }
1470
1471        if cross_any
1472            && !peer_probe
1473            && peer_probe_policy != PeerProbeStartupPolicy::BypassedWithHostBounce
1474        {
1475            eprintln!(
1476                "[pp] WARNING: MEMRA_PEER_PROBE=0 skips the boot-time peer byte-integrity \
1477                 gate; diagnostics escape hatch active"
1478            );
1479        }
1480
1481        if used.len() > 1 {
1482            if !host_bounce {
1483                // A context per distinct device (first stage that lives there; the primary's
1484                // context for the primary device).
1485                let ctx_of = |d: usize| -> &Arc<CudaContext> {
1486                    if d == primary_dev {
1487                        e.ctx()
1488                    } else {
1489                        &stages.iter().find(|s| s.dev == d).unwrap().ctx
1490                    }
1491                };
1492                // Enable peer access BOTH ways for every distinct pair (idempotent;
1493                // ALREADY_ENABLED is success).
1494                for &a in &used {
1495                    for &b in &used {
1496                        if a == b {
1497                            continue;
1498                        }
1499                        ctx_of(a).bind_to_thread()?;
1500                        let rc = unsafe {
1501                            cudarc::driver::sys::cuCtxEnablePeerAccess(ctx_of(b).cu_ctx(), 0)
1502                        };
1503                        use cudarc::driver::sys::cudaError_enum as E;
1504                        if rc != E::CUDA_SUCCESS && rc != E::CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED
1505                        {
1506                            return Err(format!(
1507                                "cuCtxEnablePeerAccess(dev{a} -> dev{b}) failed: {rc:?}"
1508                            )
1509                            .into());
1510                        }
1511                    }
1512                }
1513                // The fixed-size byte gate runs immediately after peer enable and before pool
1514                // grants. Legacy allocations make it exercise the exact `cuMemcpyPeerAsync` API
1515                // without depending on the pool setup that follows.
1516                if peer_probe && cross_any {
1517                    let probe = run_peer_probe_pass(
1518                        &stages,
1519                        &peer_capable,
1520                        host_bounce,
1521                        "fixed-16KiB",
1522                        PEER_PROBE_FIXED_BYTES,
1523                    );
1524                    e.ctx().bind_to_thread()?;
1525                    probe?;
1526                }
1527                // MEM-POOL access grant (8x box 2026-08-02, M1 cross-device fix #2):
1528                // cuCtxEnablePeerAccess does NOT map STREAM-ORDERED POOL allocations, and every
1529                // engine buffer/weight goes through the device default pool (cuMemAllocAsync via
1530                // cudarc; memra-runtime configures that pool). A stage kernel dereferencing
1531                // another device's weights — or a boundary peer TX writing the RX slot — needs
1532                // cuMemPoolSetAccess on the OWNING device's default pool for the ACCESSING
1533                // device; without it the first remote dereference is CUDA_ERROR_ILLEGAL_ADDRESS
1534                // (reported at the next API call in the poisoned context). Grant all pairs.
1535                for &owner in &used {
1536                    for &accessor in &used {
1537                        if owner == accessor {
1538                            continue;
1539                        }
1540                        let dev = cudarc::driver::result::device::get(owner as i32)?;
1541                        let mut pool: cudarc::driver::sys::CUmemoryPool = std::ptr::null_mut();
1542                        unsafe {
1543                            cudarc::driver::sys::cuDeviceGetDefaultMemPool(&mut pool, dev)
1544                                .result()?;
1545                        }
1546                        let desc = cudarc::driver::sys::CUmemAccessDesc {
1547                        location: cudarc::driver::sys::CUmemLocation {
1548                            type_: cudarc::driver::sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
1549                            id: accessor as i32,
1550                        },
1551                        flags: cudarc::driver::sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_READWRITE,
1552                    };
1553                        let rc = unsafe { cudarc::driver::sys::cuMemPoolSetAccess(pool, &desc, 1) };
1554                        if rc != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
1555                            return Err(format!(
1556                            "cuMemPoolSetAccess(dev{owner} pool -> dev{accessor}) failed: {rc:?}"
1557                        )
1558                        .into());
1559                        }
1560                    }
1561                }
1562                // MEM-POOL access grant (8x box 2026-08-02, cross-device fix #2):
1563                // cuCtxEnablePeerAccess does NOT map STREAM-ORDERED POOL allocations, and every
1564                // engine buffer/weight goes through the device default pool (cuMemAllocAsync via
1565                // cudarc; memra-runtime configures that pool). A stage-1 kernel dereferencing
1566                // dev0 weights — or the stage-0 peer TX writing dev1's RX slot — needs
1567                // cuMemPoolSetAccess on the OWNING device's default pool for the ACCESSING
1568                // device; without it the first remote dereference is CUDA_ERROR_ILLEGAL_ADDRESS
1569                // (reported at the next API call in the poisoned context). Grant both ways.
1570                for (owner, accessor) in [
1571                    (stages[0].dev, stages[1].dev),
1572                    (stages[1].dev, stages[0].dev),
1573                ] {
1574                    let dev = cudarc::driver::result::device::get(owner as i32)?;
1575                    let mut pool: cudarc::driver::sys::CUmemoryPool = std::ptr::null_mut();
1576                    unsafe {
1577                        cudarc::driver::sys::cuDeviceGetDefaultMemPool(&mut pool, dev).result()?;
1578                    }
1579                    let desc = cudarc::driver::sys::CUmemAccessDesc {
1580                    location: cudarc::driver::sys::CUmemLocation {
1581                        type_: cudarc::driver::sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
1582                        id: accessor as i32,
1583                    },
1584                    flags: cudarc::driver::sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_READWRITE,
1585                };
1586                    let rc = unsafe { cudarc::driver::sys::cuMemPoolSetAccess(pool, &desc, 1) };
1587                    if rc != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
1588                        return Err(format!(
1589                            "cuMemPoolSetAccess(dev{owner} pool -> dev{accessor}) failed: {rc:?}"
1590                        )
1591                        .into());
1592                    }
1593                }
1594                // restore the primary context for the caller's subsequent work
1595                e.ctx().bind_to_thread()?;
1596                eprintln!(
1597                    "[pp] cross-device transport: {} (cudaMemcpyPeerAsync per cross boundary; \
1598                 peer + default-pool access granted all pairs over {used:?}; weight home: {})",
1599                    devices
1600                        .iter()
1601                        .enumerate()
1602                        .map(|(s, d)| format!("stage{s}=dev{d}"))
1603                        .collect::<Vec<_>>()
1604                        .join(" "),
1605                    if pp_shard_off() {
1606                        format!("dev{primary_dev} (MEMRA_PP_SHARD=0 bring-up placement)")
1607                    } else {
1608                        "per-stage (sharded loader)".to_string()
1609                    }
1610                );
1611            } else {
1612                e.ctx().bind_to_thread()?;
1613                eprintln!(
1614                    "[pp] cross-device transport: {} (HOST-STAGED pinned D2H -> H2D per cross \
1615                     boundary; MEMRA_PP_HOST_BOUNCE=1; peer-pool grants bypassed; \
1616                     diagnostic peer access is removed before host-staged serving; \
1617                     weight home: per-stage (sharded loader))",
1618                    devices
1619                        .iter()
1620                        .enumerate()
1621                        .map(|(s, d)| format!("stage{s}=dev{d}"))
1622                        .collect::<Vec<_>>()
1623                        .join(" "),
1624                );
1625            }
1626        }
1627
1628        let mk_slot =
1629            |tx: &StageRt, rx: &StageRt| -> Result<BoundarySlot, Box<dyn std::error::Error>> {
1630                Ok(BoundarySlot {
1631                    buf: Mutex::new(None),
1632                    ev_tx: tx.ctx.new_event(None)?,
1633                    ev_rx: rx.ctx.new_event(None)?,
1634                })
1635            };
1636        let mut boundaries = Vec::with_capacity(n_st - 1);
1637        for b in 0..n_st - 1 {
1638            let (tx, rx) = (&stages[b], &stages[b + 1]);
1639            boundaries.push(BoundaryRt {
1640                slots: [mk_slot(tx, rx)?, mk_slot(tx, rx)?],
1641                step: AtomicUsize::new(0),
1642                cross: tx.dev != rx.dev,
1643            });
1644        }
1645        let readback = stages[n_st - 1].ctx.new_stream()?;
1646        let rt = PpNRt {
1647            stages,
1648            boundaries,
1649            cross_any,
1650            host_bounce,
1651            peer_probe,
1652            peer_capable,
1653            peer_probe_geometry: OnceLock::new(),
1654            bounce: OnceLock::new(),
1655            readback,
1656        };
1657        if rt.peer_probe && rt.cross_any && rt.host_bounce {
1658            rt.run_host_bounce_legacy_probe(e)?;
1659        }
1660        Ok(rt)
1661    }
1662
1663    pub fn n_stages(&self) -> usize {
1664        self.stages.len()
1665    }
1666
1667    /// True iff any boundary crosses devices.
1668    pub fn cross_device(&self) -> bool {
1669        self.cross_any
1670    }
1671
1672    fn host_bounce_active(&self) -> bool {
1673        self.host_bounce || PEER_RUNTIME_HOST_BOUNCE.load(Ordering::Acquire)
1674    }
1675
1676    fn context_for_dev<'a>(
1677        &'a self,
1678        e: &'a Engine,
1679        dev: usize,
1680    ) -> Result<&'a Arc<CudaContext>, Box<dyn std::error::Error>> {
1681        if dev == e.ctx().ordinal() {
1682            return Ok(e.ctx());
1683        }
1684        self.stages
1685            .iter()
1686            .find(|stage| stage.dev == dev)
1687            .map(|stage| &stage.ctx)
1688            .ok_or_else(|| format!("PP peer probe has no CUDA context for dev{dev}").into())
1689    }
1690
1691    fn enable_probe_peer_access(
1692        &self,
1693        e: &Engine,
1694        pairs: &[(usize, usize)],
1695    ) -> Result<Vec<(usize, usize)>, Box<dyn std::error::Error>> {
1696        let mut enabled = Vec::new();
1697        for &(src_dev, dst_dev) in pairs {
1698            let enable = (|| -> Result<(), Box<dyn std::error::Error>> {
1699                let src_ctx = self.context_for_dev(e, src_dev)?;
1700                let dst_ctx = self.context_for_dev(e, dst_dev)?;
1701                src_ctx.bind_to_thread()?;
1702                let rc = unsafe { cudarc::driver::sys::cuCtxEnablePeerAccess(dst_ctx.cu_ctx(), 0) };
1703                use cudarc::driver::sys::cudaError_enum as E;
1704                if rc == E::CUDA_SUCCESS || rc == E::CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED {
1705                    Ok(())
1706                } else {
1707                    Err(format!("{rc:?}").into())
1708                }
1709            })();
1710            if let Err(err) = enable {
1711                eprintln!(
1712                    "[pp] peer byte-integrity probe could not enable \
1713                     dev{src_dev}->dev{dst_dev}: {err}; MEMRA_PP_HOST_BOUNCE=1 remains active"
1714                );
1715            } else {
1716                enabled.push((src_dev, dst_dev));
1717            }
1718        }
1719        Ok(enabled)
1720    }
1721
1722    fn disable_probe_peer_access(
1723        &self,
1724        e: &Engine,
1725        pairs: &[(usize, usize)],
1726    ) -> Result<(), Box<dyn std::error::Error>> {
1727        let mut failures = Vec::new();
1728        for &(src_dev, dst_dev) in pairs {
1729            let disable = (|| -> Result<(), Box<dyn std::error::Error>> {
1730                let src_ctx = self.context_for_dev(e, src_dev)?;
1731                let dst_ctx = self.context_for_dev(e, dst_dev)?;
1732                src_ctx.bind_to_thread()?;
1733                let rc = unsafe { cudarc::driver::sys::cuCtxDisablePeerAccess(dst_ctx.cu_ctx()) };
1734                use cudarc::driver::sys::cudaError_enum as E;
1735                if rc == E::CUDA_SUCCESS || rc == E::CUDA_ERROR_PEER_ACCESS_NOT_ENABLED {
1736                    Ok(())
1737                } else {
1738                    Err(format!("{rc:?}").into())
1739                }
1740            })();
1741            if let Err(err) = disable {
1742                failures.push(format!("dev{src_dev}->dev{dst_dev}: {err}"));
1743            }
1744        }
1745        e.ctx().bind_to_thread()?;
1746        if failures.is_empty() {
1747            eprintln!(
1748                "[pp] peer byte-integrity probe teardown: disabled {} diagnostic pair(s); \
1749                 host-bounce serving has no probe-enabled peer access",
1750                pairs.len(),
1751            );
1752            Ok(())
1753        } else {
1754            Err(format!(
1755                "PP peer probe could not disable diagnostic peer access ({}); \
1756                 refusing host-bounce serving",
1757                failures.join(", "),
1758            )
1759            .into())
1760        }
1761    }
1762
1763    fn grant_probe_pool_access(
1764        &self,
1765        e: &Engine,
1766        pairs: &[(usize, usize)],
1767    ) -> Result<Vec<(usize, usize)>, Box<dyn std::error::Error>> {
1768        let mut granted = Vec::new();
1769        for &(src_dev, dst_dev) in pairs {
1770            let grant = (|| -> Result<(), Box<dyn std::error::Error>> {
1771                self.context_for_dev(e, dst_dev)?.bind_to_thread()?;
1772                let dev = cudarc::driver::result::device::get(dst_dev as i32)?;
1773                let mut pool: cudarc::driver::sys::CUmemoryPool = std::ptr::null_mut();
1774                unsafe {
1775                    cudarc::driver::sys::cuDeviceGetDefaultMemPool(&mut pool, dev).result()?;
1776                }
1777                let desc = cudarc::driver::sys::CUmemAccessDesc {
1778                    location: cudarc::driver::sys::CUmemLocation {
1779                        type_: cudarc::driver::sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
1780                        id: src_dev as i32,
1781                    },
1782                    flags:
1783                        cudarc::driver::sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_READWRITE,
1784                };
1785                let rc = unsafe { cudarc::driver::sys::cuMemPoolSetAccess(pool, &desc, 1) };
1786                if rc == cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
1787                    Ok(())
1788                } else {
1789                    Err(format!("{rc:?}").into())
1790                }
1791            })();
1792            if let Err(err) = grant {
1793                eprintln!(
1794                    "[pp] production-slot probe could not grant dev{src_dev} access to \
1795                     dev{dst_dev}'s default pool: {err}; MEMRA_PP_HOST_BOUNCE=1 remains active"
1796                );
1797            } else {
1798                granted.push((src_dev, dst_dev));
1799            }
1800        }
1801        Ok(granted)
1802    }
1803
1804    fn revoke_probe_pool_access(
1805        &self,
1806        e: &Engine,
1807        pairs: &[(usize, usize)],
1808    ) -> Result<(), Box<dyn std::error::Error>> {
1809        let mut failures = Vec::new();
1810        for &(src_dev, dst_dev) in pairs {
1811            let revoke = (|| -> Result<(), Box<dyn std::error::Error>> {
1812                self.context_for_dev(e, dst_dev)?.bind_to_thread()?;
1813                let dev = cudarc::driver::result::device::get(dst_dev as i32)?;
1814                let mut pool: cudarc::driver::sys::CUmemoryPool = std::ptr::null_mut();
1815                unsafe {
1816                    cudarc::driver::sys::cuDeviceGetDefaultMemPool(&mut pool, dev).result()?;
1817                }
1818                let desc = cudarc::driver::sys::CUmemAccessDesc {
1819                    location: cudarc::driver::sys::CUmemLocation {
1820                        type_: cudarc::driver::sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
1821                        id: src_dev as i32,
1822                    },
1823                    flags: cudarc::driver::sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_NONE,
1824                };
1825                let rc = unsafe { cudarc::driver::sys::cuMemPoolSetAccess(pool, &desc, 1) };
1826                if rc == cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
1827                    Ok(())
1828                } else {
1829                    Err(format!("{rc:?}").into())
1830                }
1831            })();
1832            if let Err(err) = revoke {
1833                failures.push(format!("dev{src_dev}->dev{dst_dev}: {err}"));
1834            }
1835        }
1836        e.ctx().bind_to_thread()?;
1837        if failures.is_empty() {
1838            Ok(())
1839        } else {
1840            Err(format!(
1841                "PP peer probe could not revoke diagnostic pool access ({}); \
1842                 refusing host-bounce serving",
1843                failures.join(", "),
1844            )
1845            .into())
1846        }
1847    }
1848
1849    fn run_host_bounce_legacy_probe(&self, e: &Engine) -> Result<(), Box<dyn std::error::Error>> {
1850        let enabled = self.enable_probe_peer_access(e, &self.peer_capable)?;
1851        let probe = run_peer_probe_pass(
1852            &self.stages,
1853            &enabled,
1854            true,
1855            "fixed-16KiB-legacy-preflight",
1856            PEER_PROBE_FIXED_BYTES,
1857        );
1858        let disable = self.disable_probe_peer_access(e, &enabled);
1859        disable?;
1860        probe
1861    }
1862
1863    fn new_peer_probe_boundary(
1864        &self,
1865        src_stage: usize,
1866        dst_stage: usize,
1867    ) -> Result<BoundaryRt, Box<dyn std::error::Error>> {
1868        let tx = &self.stages[src_stage];
1869        let rx = &self.stages[dst_stage];
1870        let mk_slot = || -> Result<BoundarySlot, Box<dyn std::error::Error>> {
1871            Ok(BoundarySlot {
1872                buf: Mutex::new(None),
1873                ev_tx: tx.ctx.new_event(None)?,
1874                ev_rx: rx.ctx.new_event(None)?,
1875            })
1876        };
1877        Ok(BoundaryRt {
1878            slots: [mk_slot()?, mk_slot()?],
1879            step: AtomicUsize::new(0),
1880            cross: tx.dev != rx.dev,
1881        })
1882    }
1883
1884    fn production_probe_readback(
1885        &self,
1886        path: BoundaryPath,
1887        boundary: &BoundaryRt,
1888        expected: &[u8],
1889        n: usize,
1890        slot_idx: usize,
1891    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
1892        debug_assert_eq!(expected.len(), n * std::mem::size_of::<f32>());
1893        let host = peer_probe_bytes_to_f32(expected);
1894        let poison_bytes: Vec<u8> = expected.iter().map(|byte| !byte).collect();
1895        let poison = peer_probe_bytes_to_f32(&poison_bytes);
1896        let src = &self.stages[path.src_stage];
1897        let dst = &self.stages[path.dst_stage];
1898
1899        // Pre-poison the exact stream-ordered BoundarySlot allocation so a missing or partial
1900        // peer write cannot accidentally agree where the deterministic source contains zeroes.
1901        dst.ctx.bind_to_thread()?;
1902        let poison_buf = dst.stream.clone_htod(&poison)?;
1903        dst.stream.synchronize()?;
1904        let replaced = boundary.slots[slot_idx]
1905            .buf
1906            .lock()
1907            .unwrap()
1908            .replace(poison_buf);
1909        drop(replaced);
1910        dst.stream.synchronize()?;
1911
1912        src.ctx.bind_to_thread()?;
1913        let x = src.stream.clone_htod(&host)?;
1914        self.tx_slot_path(path, boundary, &x, n, slot_idx)?;
1915
1916        dst.ctx.bind_to_thread()?;
1917        let work = self.rx_slot_path(path, boundary, slot_idx, n)?;
1918        let back = dst.stream.clone_dtoh(&work)?;
1919        dst.stream.synchronize()?;
1920        Ok(peer_probe_f32_to_bytes(&back))
1921    }
1922
1923    fn clear_peer_probe_boundary(
1924        &self,
1925        boundary: &BoundaryRt,
1926        src_stage: usize,
1927        dst_stage: usize,
1928    ) -> Result<(), Box<dyn std::error::Error>> {
1929        self.stages[dst_stage].ctx.bind_to_thread()?;
1930        for slot in &boundary.slots {
1931            let buffer = slot.buf.lock().unwrap().take();
1932            drop(buffer);
1933        }
1934        self.stages[src_stage].stream.synchronize()?;
1935        self.stages[dst_stage].stream.synchronize()?;
1936        Ok(())
1937    }
1938
1939    fn run_production_peer_probe(
1940        &self,
1941        enabled_pairs: &[(usize, usize)],
1942        host_bounce: bool,
1943        n_embd: usize,
1944    ) -> Result<(), Box<dyn std::error::Error>> {
1945        let started = std::time::Instant::now();
1946        let mut copies = 0usize;
1947        let mut skipped = 0usize;
1948        let mut total_mismatches = 0usize;
1949        let mut largest_clean_payload = 0usize;
1950
1951        for boundary_idx in 0..self.stages.len() - 1 {
1952            if self.stages[boundary_idx].dev == self.stages[boundary_idx + 1].dev {
1953                continue;
1954            }
1955            for (src_stage, dst_stage) in [
1956                (boundary_idx, boundary_idx + 1),
1957                (boundary_idx + 1, boundary_idx),
1958            ] {
1959                let src_dev = self.stages[src_stage].dev;
1960                let dst_dev = self.stages[dst_stage].dev;
1961                if !enabled_pairs.contains(&(src_dev, dst_dev)) {
1962                    if host_bounce {
1963                        skipped += PEER_PROBE_TOKEN_WIDTHS.len();
1964                        eprintln!(
1965                            "[pp] production-slot peer probe SKIP: boundary={boundary_idx} \
1966                             dev{src_dev}->dev{dst_dev} widths_tokens={:?} \
1967                             (peer or pool access unavailable; MEMRA_PP_HOST_BOUNCE=1 remains \
1968                             fail-safe)",
1969                            PEER_PROBE_TOKEN_WIDTHS,
1970                        );
1971                        continue;
1972                    }
1973                    return Err(format!(
1974                        "PP production-slot peer probe cannot run boundary={boundary_idx} \
1975                         dev{src_dev}->dev{dst_dev}: peer/pool access is not enabled"
1976                    )
1977                    .into());
1978                }
1979
1980                let probe_boundary = self.new_peer_probe_boundary(src_stage, dst_stage)?;
1981                let path = BoundaryPath {
1982                    boundary: boundary_idx,
1983                    src_stage,
1984                    dst_stage,
1985                    transport: BoundaryTransport::Peer,
1986                };
1987                let mut direction_copies = 0usize;
1988                let mut direction_skipped = 0usize;
1989                let mut direction_mismatches = 0usize;
1990                let mut direction_largest_clean = 0usize;
1991                let mut failure = None;
1992
1993                for (width_idx, tokens) in PEER_PROBE_TOKEN_WIDTHS.into_iter().enumerate() {
1994                    let n = n_embd.checked_mul(tokens).ok_or_else(|| {
1995                        format!(
1996                            "PP production-slot probe element count overflows for \
1997                             n_embd={n_embd} tokens={tokens}"
1998                        )
1999                    })?;
2000                    let bytes = n.checked_mul(std::mem::size_of::<f32>()).ok_or_else(|| {
2001                        format!(
2002                            "PP production-slot probe byte count overflows for \
2003                             n_embd={n_embd} tokens={tokens}"
2004                        )
2005                    })?;
2006                    let expected = peer_probe_pattern(bytes, boundary_idx, src_dev, dst_dev);
2007                    let readback = match self.production_probe_readback(
2008                        path,
2009                        &probe_boundary,
2010                        &expected,
2011                        n,
2012                        width_idx % 2,
2013                    ) {
2014                        Ok(readback) => readback,
2015                        Err(err) if host_bounce => {
2016                            skipped += 1;
2017                            direction_skipped += 1;
2018                            eprintln!(
2019                                "[pp] production-slot peer probe ERROR: \
2020                                 boundary={boundary_idx} dev{src_dev}->dev{dst_dev} \
2021                                 tokens={tokens} bytes={bytes}: {err}; \
2022                                 MEMRA_PP_HOST_BOUNCE=1, proceeding on the host-staged path"
2023                            );
2024                            continue;
2025                        }
2026                        Err(err) => {
2027                            failure = Some(format!(
2028                                "PP production-slot peer probe FAILED: \
2029                                 boundary={boundary_idx} dev{src_dev}->dev{dst_dev} \
2030                                 tokens={tokens} bytes={bytes}: {err}; refusing native P2P \
2031                                 (set MEMRA_PP_HOST_BOUNCE=1 to use the host-staged path; \
2032                                 MEMRA_PEER_PROBE=0 cannot authorize sharded native peer \
2033                                 transport)"
2034                            ));
2035                            break;
2036                        }
2037                    };
2038                    copies += 1;
2039                    direction_copies += 1;
2040                    let mismatches = peer_probe_mismatch_count(&expected, &readback);
2041                    if mismatches == 0 {
2042                        largest_clean_payload = largest_clean_payload.max(bytes);
2043                        direction_largest_clean = direction_largest_clean.max(bytes);
2044                    } else if host_bounce {
2045                        total_mismatches += mismatches;
2046                        direction_mismatches += mismatches;
2047                        eprintln!(
2048                            "[pp] production-slot peer probe CORRUPTION: \
2049                             boundary={boundary_idx} dev{src_dev}->dev{dst_dev} tokens={tokens} \
2050                             bytes={bytes} mismatches={mismatches}; MEMRA_PP_HOST_BOUNCE=1, \
2051                             proceeding on the host-staged path"
2052                        );
2053                    } else {
2054                        failure = Some(format!(
2055                            "PP production-slot peer probe FAILED: boundary={boundary_idx} \
2056                             dev{src_dev}->dev{dst_dev} tokens={tokens} bytes={bytes}: \
2057                             {mismatches} mismatched byte(s); refusing native P2P \
2058                             (set MEMRA_PP_HOST_BOUNCE=1 to use the host-staged path; \
2059                             MEMRA_PEER_PROBE=0 cannot authorize sharded native peer transport)"
2060                        ));
2061                        break;
2062                    }
2063                }
2064
2065                self.clear_peer_probe_boundary(&probe_boundary, src_stage, dst_stage)?;
2066                if let Some(err) = failure {
2067                    return Err(err.into());
2068                }
2069                eprintln!(
2070                    "[pp] production-slot peer probe direction: boundary={boundary_idx} \
2071                     dev{src_dev}->dev{dst_dev} copies={direction_copies} \
2072                     skipped={direction_skipped} mismatches={direction_mismatches} \
2073                     largest_clean_payload_bytes={direction_largest_clean}"
2074                );
2075            }
2076        }
2077
2078        let status = if total_mismatches > 0 {
2079            "BOUNCE"
2080        } else if skipped > 0 && copies > 0 {
2081            "PARTIAL"
2082        } else if skipped > 0 {
2083            "SKIP"
2084        } else {
2085            "PASS"
2086        };
2087        eprintln!(
2088            "[pp] production-slot peer probe {status}: widths_tokens={:?} copies={copies} \
2089             skipped={skipped} mismatches={total_mismatches} \
2090             largest_clean_payload_bytes={largest_clean_payload} elapsed_ms={:.3}",
2091            PEER_PROBE_TOKEN_WIDTHS,
2092            started.elapsed().as_secs_f64() * 1e3,
2093        );
2094        Ok(())
2095    }
2096
2097    fn run_host_bounce_production_probe(
2098        &self,
2099        e: &Engine,
2100        n_embd: usize,
2101    ) -> Result<(), Box<dyn std::error::Error>> {
2102        let enabled = self.enable_probe_peer_access(e, &self.peer_capable)?;
2103        let granted = self.grant_probe_pool_access(e, &enabled)?;
2104        let probe = self.run_production_peer_probe(&granted, true, n_embd);
2105        // Teardown always runs, but the probe verdict wins: a CORRUPTION verdict (probe is
2106        // Err) must never be masked by a teardown failure. `revoke?; disable?; probe`
2107        // short-circuited teardown errors BEFORE probe was inspected, discarding the byte-
2108        // integrity signal on any teardown hiccup (hermes 9d6ae8d3). Surface teardown errors
2109        // only when the probe itself succeeded.
2110        let revoke = self.revoke_probe_pool_access(e, &granted);
2111        let disable = self.disable_probe_peer_access(e, &enabled);
2112        probe?;
2113        revoke?;
2114        disable?;
2115        Ok(())
2116    }
2117
2118    fn init_peer_probe_geometry(
2119        &self,
2120        e: &Engine,
2121        n_embd: usize,
2122    ) -> Result<(), Box<dyn std::error::Error>> {
2123        if !self.peer_probe || !self.cross_any {
2124            return Ok(());
2125        }
2126        let bytes = n_embd
2127            .checked_mul(std::mem::size_of::<f32>())
2128            .ok_or_else(|| format!("PP boundary-slot byte count overflows for n_embd={n_embd}"))?;
2129        let result = self.peer_probe_geometry.get_or_init(|| {
2130            let probe = if self.host_bounce_active() {
2131                self.run_host_bounce_production_probe(e, n_embd)
2132            } else {
2133                self.run_production_peer_probe(&self.peer_capable, false, n_embd)
2134            };
2135            let restore = e.ctx().bind_to_thread();
2136            match (probe, restore) {
2137                (Ok(()), Ok(())) => Ok(bytes),
2138                (Err(err), _) => Err(err.to_string()),
2139                (_, Err(err)) => Err(err.to_string()),
2140            }
2141        });
2142        let probed = result
2143            .as_ref()
2144            .map_err(|err| -> Box<dyn std::error::Error> { err.clone().into() })?;
2145        if *probed != bytes {
2146            return Err(format!(
2147                "peer probe initialized for boundary-slot bytes={probed} but model requests \
2148                 bytes={bytes}; one PP runtime supports one model geometry per process"
2149            )
2150            .into());
2151        }
2152        Ok(())
2153    }
2154
2155    fn init_host_bounce_staging(
2156        &self,
2157        e: &Engine,
2158        n_embd: usize,
2159    ) -> Result<(), Box<dyn std::error::Error>> {
2160        if !self.cross_any {
2161            return Ok(());
2162        }
2163        e.ctx().bind_to_thread()?;
2164        let result = self.bounce.get_or_init(|| {
2165            HostBounceRt::new(n_embd, &self.boundaries)
2166                .map(|rt| {
2167                    let bytes = rt.capacity * std::mem::size_of::<f32>();
2168                    eprintln!(
2169                        "[pp] host-bounce staging ready: n_embd={n_embd} max_tokens={} \
2170                         slot_bytes={bytes} slots_per_cross_boundary=2",
2171                        crate::cache::PRIME_CHUNK_MAX_TOKENS,
2172                    );
2173                    rt
2174                })
2175                .map_err(|err| err.to_string())
2176        });
2177        let bounce = result
2178            .as_ref()
2179            .map_err(|err| -> Box<dyn std::error::Error> { err.clone().into() })?;
2180        if bounce.n_embd != n_embd {
2181            return Err(format!(
2182                "host-bounce runtime initialized for n_embd={} but model requests n_embd={n_embd}; \
2183                 one PP runtime supports one model geometry per process",
2184                bounce.n_embd,
2185            )
2186            .into());
2187        }
2188        Ok(())
2189    }
2190
2191    /// Exercise the newly armed staging through the real D2H/event/H2D boundary path before the
2192    /// live transport latch can observe it. One row per cross boundary is enough to validate the
2193    /// pinned capacity, event ordering, contexts, and byte continuity without touching peer DMA.
2194    fn validate_host_bounce_staging(
2195        &self,
2196        e: &Engine,
2197        n_embd: usize,
2198    ) -> Result<(), Box<dyn std::error::Error>> {
2199        let bytes = n_embd
2200            .checked_mul(std::mem::size_of::<f32>())
2201            .ok_or_else(|| {
2202                format!("host-bounce validation byte count overflows for n_embd={n_embd}")
2203            })?;
2204        for boundary_idx in 0..self.stages.len() - 1 {
2205            if !self.boundaries[boundary_idx].cross {
2206                continue;
2207            }
2208            let src_stage = boundary_idx;
2209            let dst_stage = boundary_idx + 1;
2210            let probe_boundary = self.new_peer_probe_boundary(src_stage, dst_stage)?;
2211            let path = BoundaryPath {
2212                boundary: boundary_idx,
2213                src_stage,
2214                dst_stage,
2215                transport: BoundaryTransport::HostBounce,
2216            };
2217            let expected = peer_probe_pattern(
2218                bytes,
2219                boundary_idx,
2220                self.stages[src_stage].dev,
2221                self.stages[dst_stage].dev,
2222            );
2223            let readback =
2224                self.production_probe_readback(path, &probe_boundary, &expected, n_embd, 0);
2225            let clear = self.clear_peer_probe_boundary(&probe_boundary, src_stage, dst_stage);
2226            let readback = readback?;
2227            clear?;
2228            let mismatches = peer_probe_mismatch_count(&expected, &readback);
2229            if mismatches > 0 {
2230                return Err(format!(
2231                    "runtime host-bounce staging validation FAILED: boundary={boundary_idx} \
2232                     bytes={bytes} mismatches={mismatches}"
2233                )
2234                .into());
2235            }
2236        }
2237        e.ctx().bind_to_thread()?;
2238        eprintln!(
2239            "[pp] runtime host-bounce staging validation PASS: row_bytes={bytes} \
2240             cross_boundaries={}",
2241            self.boundaries
2242                .iter()
2243                .filter(|boundary| boundary.cross)
2244                .count(),
2245        );
2246        Ok(())
2247    }
2248
2249    fn arm_runtime_host_bounce(
2250        &self,
2251        e: &Engine,
2252        row_bytes: usize,
2253    ) -> Result<(), Box<dyn std::error::Error>> {
2254        if row_bytes == 0 || row_bytes % std::mem::size_of::<f32>() != 0 {
2255            return Err(format!(
2256                "runtime host-bounce cannot recover n_embd from row_bytes={row_bytes}"
2257            )
2258            .into());
2259        }
2260        let n_embd = row_bytes / std::mem::size_of::<f32>();
2261        self.init_host_bounce_staging(e, n_embd)?;
2262        self.validate_host_bounce_staging(e, n_embd)
2263    }
2264
2265    /// Finish boot-time transport setup from the authoritative model width. This runs the
2266    /// production `BoundarySlot` ladder at 1/8/16/`PRIME_CHUNK_MAX_TOKENS` `[n_embd] f32` rows
2267    /// once, then allocates host-bounce slots when selected. The loader calls it before uploading
2268    /// the first model weight; `new_cache` repeats the call as an idempotent guard before the first
2269    /// forward.
2270    pub fn init_boundary_transport(
2271        &self,
2272        e: &Engine,
2273        n_embd: usize,
2274    ) -> Result<(), Box<dyn std::error::Error>> {
2275        if PEER_RUNTIME_PROBE_FAILED.load(Ordering::Acquire)
2276            && !PEER_RUNTIME_HOST_BOUNCE.load(Ordering::Acquire)
2277        {
2278            return Err(
2279                "PP runtime peer byte-integrity probe previously failed; refusing native P2P \
2280                 reuse because runtime host-bounce staging could not be armed"
2281                    .into(),
2282            );
2283        }
2284        self.init_peer_probe_geometry(e, n_embd)?;
2285        if !self.host_bounce_active() || !self.cross_any {
2286            return Ok(());
2287        }
2288        self.init_host_bounce_staging(e, n_embd)
2289    }
2290
2291    /// Run one due peer re-probe at a scheduler boundary on the CUDA owner thread. Each width has
2292    /// an independent copy-count deadline: an idle-only rung can remain pending while later cheap
2293    /// rungs keep running. The probe synchronizes the stage streams it exercises; no background
2294    /// thread touches CUDA.
2295    fn service_runtime_peer_probe(
2296        &self,
2297        e: &Engine,
2298        scheduler_idle: bool,
2299        probe_allowed: bool,
2300    ) -> Result<RuntimePeerProbeStatus, Box<dyn std::error::Error>> {
2301        if !self.peer_probe || !self.cross_any || self.host_bounce_active() {
2302            return Ok(RuntimePeerProbeStatus::NotRun);
2303        }
2304        if PEER_RUNTIME_PROBE_FAILED.load(Ordering::Acquire) {
2305            return Err(
2306                "PP runtime peer byte-integrity probe previously failed; native P2P is latched off"
2307                    .into(),
2308            );
2309        }
2310        let row_bytes = match self.peer_probe_geometry.get() {
2311            Some(Ok(bytes)) => *bytes,
2312            _ => return Ok(RuntimePeerProbeStatus::NotRun),
2313        };
2314
2315        let copies = PEER_BOUNDARY_COPIES.load(Ordering::Relaxed);
2316        let (width_index, tokens) = loop {
2317            let next_probe_copy = std::array::from_fn(|width_index| {
2318                PEER_RUNTIME_NEXT_PROBE_COPY[width_index].load(Ordering::Relaxed)
2319            });
2320            let measured_cost_ns = std::array::from_fn(|width_index| {
2321                PEER_RUNTIME_PROBE_MAX_COST_NS[width_index].load(Ordering::Relaxed)
2322            });
2323            let Some(candidate) = runtime_peer_probe_candidate(
2324                copies,
2325                next_probe_copy,
2326                measured_cost_ns,
2327                scheduler_idle,
2328            ) else {
2329                return Ok(RuntimePeerProbeStatus::NotRun);
2330            };
2331            // A mismatch immediately revokes native peer access before validated host bounce is
2332            // published. Live speculative sessions still dereference token/position state through
2333            // UVA outside the bounced boundary, so the worker may defer a runnable cheap rung until
2334            // those sessions retire. Do not consume its deadline or completed-probe counter.
2335            if !probe_allowed {
2336                return Ok(RuntimePeerProbeStatus::Deferred);
2337            }
2338            let due = next_probe_copy[candidate.0];
2339            let next = runtime_peer_probe_next_copy(due, copies);
2340            if PEER_RUNTIME_NEXT_PROBE_COPY[candidate.0]
2341                .compare_exchange(due, next, Ordering::AcqRel, Ordering::Relaxed)
2342                .is_ok()
2343            {
2344                break candidate;
2345            }
2346        };
2347        let probe_index = PEER_RUNTIME_PROBES.fetch_add(1, Ordering::Relaxed);
2348        let probe_bytes = row_bytes.checked_mul(tokens);
2349        let scheduler_class = if scheduler_idle { "idle" } else { "busy" };
2350        let label = format!("runtime-{scheduler_class}-{tokens}tok");
2351        let started = std::time::Instant::now();
2352        let probe = match probe_bytes {
2353            Some(bytes) => {
2354                run_peer_probe_pass(&self.stages, &self.peer_capable, false, &label, bytes)
2355            }
2356            None => Err(format!(
2357                "PP runtime peer probe byte count overflows for row_bytes={row_bytes} \
2358                 tokens={tokens}"
2359            )
2360            .into()),
2361        };
2362        let restore = e.ctx().bind_to_thread();
2363        let elapsed_ns = started.elapsed().as_nanos().min(u64::MAX as u128) as u64;
2364        let previous_max =
2365            PEER_RUNTIME_PROBE_MAX_COST_NS[width_index].fetch_max(elapsed_ns, Ordering::Relaxed);
2366        let verdict = match (probe, restore) {
2367            (Ok(()), Ok(())) => Ok(()),
2368            (Err(err), _) => Err(err.to_string()),
2369            (_, Err(err)) => Err(err.to_string()),
2370        };
2371        if let Err(err) = verdict {
2372            PEER_RUNTIME_PROBE_FAILURES.fetch_add(1, Ordering::Relaxed);
2373            let arm = latch_runtime_host_bounce(
2374                &PEER_RUNTIME_PROBE_FAILED,
2375                &PEER_RUNTIME_HOST_BOUNCE,
2376                || {
2377                    self.arm_runtime_host_bounce(e, row_bytes)
2378                        .map_err(|arm_err| arm_err.to_string())
2379                },
2380            );
2381            if let Err(arm_err) = arm {
2382                let message = format!(
2383                    "PP runtime peer byte-integrity re-probe FAILED after \
2384                     boundary_copies={copies} rung={}/{} tokens={tokens}: {err}; native P2P is \
2385                     latched off and host-bounce staging could not be armed: {arm_err}",
2386                    width_index + 1,
2387                    PEER_PROBE_TOKEN_WIDTHS.len(),
2388                );
2389                eprintln!("[pp] SECURITY RED: {message}; worker must stop");
2390                return Err(message.into());
2391            }
2392            eprintln!(
2393                "[pp] SECURITY RED: PP runtime peer byte-integrity re-probe FAILED after \
2394                 boundary_copies={copies} rung={}/{} tokens={tokens}: {err}; native P2P is \
2395                 latched off and the live transport DEGRADED to validated host bounce for the \
2396                 remainder of this process",
2397                width_index + 1,
2398                PEER_PROBE_TOKEN_WIDTHS.len(),
2399            );
2400            return Ok(RuntimePeerProbeStatus::DegradedToHostBounce);
2401        }
2402        if width_index + 1 != PEER_PROBE_TOKEN_WIDTHS.len()
2403            && previous_max <= PEER_RUNTIME_PROBE_BUDGET_NS
2404            && elapsed_ns > PEER_RUNTIME_PROBE_BUDGET_NS
2405        {
2406            eprintln!(
2407                "[pp] runtime peer re-probe rung exceeded the {:.3}ms owner-thread budget: \
2408                 tokens={tokens} measured_ms={:.3}; future runs are idle-only",
2409                PEER_RUNTIME_PROBE_BUDGET_NS as f64 / 1e6,
2410                elapsed_ns as f64 / 1e6,
2411            );
2412        }
2413        eprintln!(
2414            "[pp] runtime peer byte-integrity re-probe PASS: \
2415             boundary_copies={copies} interval_copies={PEER_RUNTIME_PROBE_INTERVAL_COPIES} \
2416             rung={}/{} tokens={tokens} bytes={} probe_index={probe_index} elapsed_ms={:.3} \
2417             scheduler_idle={scheduler_idle}",
2418            width_index + 1,
2419            PEER_PROBE_TOKEN_WIDTHS.len(),
2420            probe_bytes.unwrap(),
2421            elapsed_ns as f64 / 1e6,
2422        );
2423        Ok(RuntimePeerProbeStatus::Passed)
2424    }
2425
2426    fn bounce_rt(&self) -> Result<&HostBounceRt, Box<dyn std::error::Error>> {
2427        self.bounce
2428            .get()
2429            .ok_or_else(|| -> Box<dyn std::error::Error> {
2430                "MEMRA_PP_HOST_BOUNCE=1 staging was not initialized from model geometry".into()
2431            })?
2432            .as_ref()
2433            .map_err(|err| -> Box<dyn std::error::Error> { err.clone().into() })
2434    }
2435
2436    /// The engine a stage's subgraph must run through: the primary engine when the stage
2437    /// lives on the primary device, else the stage's own (remote-context) engine.
2438    pub fn engine<'a>(&'a self, s: usize, primary: &'a Engine) -> &'a Engine {
2439        self.stages[s].engine.as_ref().unwrap_or(primary)
2440    }
2441
2442    /// Bind this OS thread to stage `s`'s CUDA context before issuing work there.
2443    pub fn bind_stage(&self, s: usize) -> Result<(), Box<dyn std::error::Error>> {
2444        self.stages[s].ctx.bind_to_thread()?;
2445        Ok(())
2446    }
2447
2448    /// Enter stage `s`: until the guard drops, every engine op on this thread launches on
2449    /// the stage's stream (memra_runtime ambient-stream override).
2450    pub fn enter(&self, s: usize) -> memra_runtime::StreamOverride {
2451        memra_runtime::push_stream_override(self.stages[s].stream.clone())
2452    }
2453
2454    /// Allocate/grow BOTH slots for a boundary before pipelined issue starts. `tx()` can
2455    /// grow a slot lazily, but first-use ordering requires synchronizing the RX stream
2456    /// after that allocation. If slot 1 first grows after stage 1 of chunk N has already
2457    /// been queued, that sync drains chunk N and erases the only overlap in a two-chunk
2458    /// prime. Prewarming both slots pays the same one-time sync before either stage starts.
2459    pub fn prepare_overlap_slots(
2460        &self,
2461        b: usize,
2462        n: usize,
2463    ) -> Result<(), Box<dyn std::error::Error>> {
2464        let bd = &self.boundaries[b];
2465        let s_rx = &self.stages[b + 1].stream;
2466        let mut grew = false;
2467        for sl in &bd.slots {
2468            let mut guard = sl.buf.lock().unwrap();
2469            if guard.as_ref().map(|bf| bf.len() < n).unwrap_or(true) {
2470                *guard = Some(s_rx.alloc_zeros::<f32>(n)?);
2471                grew = true;
2472            }
2473        }
2474        if grew {
2475            s_rx.synchronize()?;
2476        }
2477        Ok(())
2478    }
2479
2480    /// Boundary TX at boundary `b` (call within the stage-`b` scope; `x` = the
2481    /// materialized [n] residual): wait for the slot's previous RX (write-after-read
2482    /// guard), copy `x` into the slot's persistent buffer via the boundary's transport on
2483    /// stage-b's stream (the owning-stream/publication law), record ev_tx. Returns the
2484    /// slot index for the paired rx().
2485    ///
2486    /// `n` is the PAYLOAD ELEMENT COUNT, not a fixed model constant: the eager arm passes
2487    /// `n_embd` (one row), the batched arm passes `b_n * n_embd` (B stacked rows, the
2488    /// [B, n_embd] boundary). The slot buffer is GROW-ONLY and the transport moves exactly
2489    /// the first `n` elements — batched serving changes B every tick (chunk fill), and a
2490    /// realloc-on-every-size-change would host-sync the RX stream per width change (see the
2491    /// SLOT FIRST-USE ORDERING note below for why each allocation needs that sync). Growing
2492    /// to the high-water mark makes the syncs O(distinct widths) instead of O(width changes).
2493    pub fn tx(
2494        &self,
2495        b: usize,
2496        x: &CudaSlice<f32>,
2497        n: usize,
2498    ) -> Result<usize, Box<dyn std::error::Error>> {
2499        assert_eq!(x.len(), n, "pp tx: residual length mismatch");
2500        let bd = &self.boundaries[b];
2501        let slot_idx = if pp2_overlap() {
2502            bd.step.fetch_add(1, Ordering::Relaxed) % 2
2503        } else {
2504            0
2505        };
2506        self.tx_slot(b, x, n, slot_idx)
2507    }
2508
2509    /// Pipelined boundary TX: always alternate the shared double-buffer slots, independent
2510    /// of the decode-side `MEMRA_PP_OVERLAP` experiment flag. The boundary-local atomic
2511    /// keeps concurrent callers on one slot sequence rather than each restarting at A.
2512    pub fn tx_pipelined(
2513        &self,
2514        b: usize,
2515        x: &CudaSlice<f32>,
2516        n: usize,
2517    ) -> Result<usize, Box<dyn std::error::Error>> {
2518        assert_eq!(x.len(), n, "pp tx: residual length mismatch");
2519        let slot_idx = self.boundaries[b].step.fetch_add(1, Ordering::Relaxed) % 2;
2520        self.tx_slot(b, x, n, slot_idx)
2521    }
2522
2523    fn tx_slot(
2524        &self,
2525        b: usize,
2526        x: &CudaSlice<f32>,
2527        n: usize,
2528        slot_idx: usize,
2529    ) -> Result<usize, Box<dyn std::error::Error>> {
2530        let bd = &self.boundaries[b];
2531        let path = BoundaryPath {
2532            boundary: b,
2533            src_stage: b,
2534            dst_stage: b + 1,
2535            transport: boundary_transport(bd.cross, self.host_bounce_active()),
2536        };
2537        let copied_slot = self.tx_slot_path(path, bd, x, n, slot_idx)?;
2538        if path.transport == BoundaryTransport::Peer {
2539            PEER_BOUNDARY_COPIES.fetch_add(1, Ordering::Relaxed);
2540        }
2541        Ok(copied_slot)
2542    }
2543
2544    fn tx_slot_path(
2545        &self,
2546        path: BoundaryPath,
2547        bd: &BoundaryRt,
2548        x: &CudaSlice<f32>,
2549        n: usize,
2550        slot_idx: usize,
2551    ) -> Result<usize, Box<dyn std::error::Error>> {
2552        debug_assert!(slot_idx < 2);
2553        let sl = &bd.slots[slot_idx];
2554        let s_tx = &self.stages[path.src_stage].stream;
2555        s_tx.wait(&sl.ev_rx)?;
2556        let mut guard = sl.buf.lock().unwrap();
2557        if guard.as_ref().map(|bf| bf.len() < n).unwrap_or(true) {
2558            // allocated on the RX stage's stream: the buffer lives on the RX device.
2559            let s_rx = &self.stages[path.dst_stage].stream;
2560            *guard = Some(s_rx.alloc_zeros::<f32>(n)?);
2561            // SLOT FIRST-USE ORDERING (2026-08-02 pipelined-gate find): the lazy alloc's
2562            // pool-alloc + memset enqueue on the RX stream; the TX copy below issues on
2563            // the TX stream, and on a slot's FIRST use ev_rx has never been recorded —
2564            // nothing orders them. With >=2 tokens in flight the RX stream is still busy
2565            // with the previous token, the memset lands AFTER the TX copy, and the
2566            // boundary residual is zeroed (window=1 passed, window>=2 failed at the
2567            // slot-1 first-use step; -overlap arms passed because the synchronous serial
2568            // arm pre-warmed both slots). Host-sync the RX stream once per slot
2569            // allocation — at most 2*(N-1) one-time syncs per process, all during prime.
2570            s_rx.synchronize()?;
2571        }
2572        let buf = guard.as_mut().unwrap();
2573        match path.transport {
2574            BoundaryTransport::Local => s_tx.memcpy_dtod(x, buf)?,
2575            BoundaryTransport::HostBounce => {
2576                debug_assert_eq!(path.src_stage, path.boundary);
2577                debug_assert_eq!(path.dst_stage, path.boundary + 1);
2578                let bounce = self.bounce_rt()?;
2579                if n > bounce.capacity {
2580                    return Err(format!(
2581                        "pp host-bounce payload {n} exceeds geometry-sized capacity {} \
2582                         (n_embd={}, max prime tokens={})",
2583                        bounce.capacity,
2584                        bounce.n_embd,
2585                        crate::cache::PRIME_CHUNK_MAX_TOKENS,
2586                    )
2587                    .into());
2588                }
2589                let mut host = bounce.slot(path.boundary, slot_idx)?.lock().unwrap();
2590                // D2H is issued on the producing stage's stream. ev_tx below publishes the
2591                // completed host bytes to the receiving stream; the exact prefix avoids moving
2592                // a full 64 MiB slot for a one-row decode, and no peer pointer is formed here.
2593                s_tx.memcpy_dtoh(x, host.prefix_mut(n))?;
2594            }
2595            BoundaryTransport::Peer => {
2596                // cudaMemcpyPeerAsync (M0: 2.8x NCCL at PP activation sizes), issued on the
2597                // publishing TX stream with explicit src/dst contexts.
2598                use cudarc::driver::{DevicePtr, DevicePtrMut};
2599                let (sp, _g0) = x.device_ptr(s_tx);
2600                let (dp, _g1) = buf.device_ptr_mut(s_tx);
2601                self.stages[path.src_stage].ctx.bind_to_thread()?;
2602                unsafe {
2603                    cudarc::driver::result::memcpy_peer_async(
2604                        self.stages[path.dst_stage].ctx.cu_ctx(),
2605                        dp,
2606                        self.stages[path.src_stage].ctx.cu_ctx(),
2607                        sp,
2608                        n * std::mem::size_of::<f32>(),
2609                        s_tx.cu_stream(),
2610                    )?;
2611                }
2612            }
2613        }
2614        sl.ev_tx.record(s_tx)?;
2615        Ok(slot_idx)
2616    }
2617
2618    /// Boundary RX at boundary `b` (call within the stage-`b+1` scope): wait on the slot's
2619    /// ev_tx, copy the boundary buffer into a fresh working buffer (dtod on the RX stream —
2620    /// local on the RX device in both transports), record ev_rx. The returned buffer is
2621    /// RX-stage-owned: allocated, consumed, and eventually freed on that stage's stream.
2622    pub fn rx(
2623        &self,
2624        b: usize,
2625        slot_idx: usize,
2626        n: usize,
2627    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2628        let bd = &self.boundaries[b];
2629        let path = BoundaryPath {
2630            boundary: b,
2631            src_stage: b,
2632            dst_stage: b + 1,
2633            transport: boundary_transport(bd.cross, self.host_bounce_active()),
2634        };
2635        self.rx_slot_path(path, bd, slot_idx, n)
2636    }
2637
2638    fn rx_slot_path(
2639        &self,
2640        path: BoundaryPath,
2641        bd: &BoundaryRt,
2642        slot_idx: usize,
2643        n: usize,
2644    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
2645        let sl = &bd.slots[slot_idx];
2646        let s_rx = &self.stages[path.dst_stage].stream;
2647        s_rx.wait(&sl.ev_tx)?;
2648        let mut guard = sl.buf.lock().unwrap();
2649        let buf = guard.as_mut().expect("pp rx before tx");
2650        assert!(
2651            buf.len() >= n,
2652            "pp rx: slot holds {} < requested {n}",
2653            buf.len()
2654        );
2655        if path.transport == BoundaryTransport::HostBounce {
2656            debug_assert_eq!(path.src_stage, path.boundary);
2657            debug_assert_eq!(path.dst_stage, path.boundary + 1);
2658            let bounce = self.bounce_rt()?;
2659            let host = bounce.slot(path.boundary, slot_idx)?.lock().unwrap();
2660            let mut dst = buf.slice_mut(0..n);
2661            // The destination stream already waits ev_tx, so this H2D cannot observe the
2662            // staging slot before the source stream's D2H completes.
2663            s_rx.memcpy_htod(host.prefix(n), &mut dst)?;
2664        }
2665        // uninit working buffer (fully overwritten by the copy), allocated explicitly on
2666        // the stage stream so rx() is correct even outside an enter() scope.
2667        let mut work = unsafe { s_rx.alloc::<f32>(n)? };
2668        // Slice the slot to the payload: the buffer is grow-only (see tx), so at a narrower
2669        // width it is LONGER than `work` and cudarc's memcpy_dtod (dst.len() >= src.len())
2670        // would assert. The paired tx wrote exactly these first n elements.
2671        s_rx.memcpy_dtod(&buf.slice(0..n), &mut work)?;
2672        sl.ev_rx.record(s_rx)?;
2673        Ok(work)
2674    }
2675
2676    /// PUBLISH a DEVICE-RESIDENT result off the last stage to the caller's stream
2677    /// (lane/pp2-spec 2026-08-06).
2678    ///
2679    /// Every ppN body before this one returned HOST values — `decode_step_h_ppn` and
2680    /// `decode_step_batch_ppn` both `dtoh` inside the last-stage scope, and a dtoh on the
2681    /// producing stream is self-ordering. The verify trunk is the FIRST ppN body whose
2682    /// contract is device-resident output (`decode_step_t_h_emb_dev` exists precisely so the
2683    /// accept walk argmaxes on-device instead of moving T x n_vocab f32 per round), and
2684    /// device slices carry no stream affinity: the caller resumes on the PRIMARY stream and
2685    /// dereferences buffers whose producing kernels are still queued on the last stage's
2686    /// stream. Nothing orders them.
2687    ///
2688    /// Why this only ever failed on ONE device: with stages on separate devices the caller's
2689    /// first touch is a cross-device copy that the driver orders against the source context,
2690    /// and the readback path syncs. Two streams on the SAME device genuinely overlap, so the
2691    /// primary stream reads a buffer whose matmul has not run — nondeterministic garbage
2692    /// (measured: NaN, 3155.677, and 2.87e-5 where the reference had -2.0048926), and it
2693    /// poisons the NEXT arm in the same process because the corrupted KV persists. This is
2694    /// the same class as the SLOT FIRST-USE ORDERING find above, one level up: there the
2695    /// unordered pair was alloc-memset vs TX copy, here it is stage-N compute vs the
2696    /// caller's consumer.
2697    ///
2698    /// Fix = the boundary law applied to the exit: record an event on the producing stage
2699    /// stream, make the caller's stream wait on it. Event-wait, not a device sync, so the
2700    /// stage streams keep running for the deferred-readback arm. Call INSIDE the last-stage
2701    /// scope, after the last enqueue, with the caller's (pre-`enter`) stream.
2702    pub fn publish_to(
2703        &self,
2704        s: usize,
2705        dst: &Arc<CudaStream>,
2706    ) -> Result<(), Box<dyn std::error::Error>> {
2707        let st = &self.stages[s];
2708        // Same stream (STREAMS=0 rollback, or a caller already on the stage stream): the
2709        // stream orders itself; recording+waiting would be a no-op with a stray event.
2710        if Arc::ptr_eq(&st.stream, dst) {
2711            return Ok(());
2712        }
2713        let ev = st.ctx.new_event(None)?;
2714        ev.record(&st.stream)?;
2715        dst.wait(&ev)?;
2716        Ok(())
2717    }
2718
2719    /// REVERSE PUBLICATION (#87 root cause, lane/pp2spec-crash 2026-08-07): order every
2720    /// STAGE stream behind the CALLER's stream — the mirror of `publish_to`.
2721    ///
2722    /// `publish_to` orders caller READS behind stage COMPUTE. Nothing ordered the other
2723    /// direction: buffers ALLOCATED on a stage stream (the verify's returned logits/hidden,
2724    /// the VerifyCkpt stashes) are CONSUMED by kernels the caller enqueues on the PRIMARY
2725    /// stream, and when they drop, cudarc enqueues `free_async` on the ALLOCATING (stage)
2726    /// stream. With event tracking elided (the decode-path default) the drop carries no
2727    /// read-guard, so the pool can hand the block to the NEXT stage-stream allocation and
2728    /// its writes overwrite memory the queued primary-stream consumer has not read yet.
2729    /// Measured (research/pp2spec-crash-20260807): the spec round-seed read 13/4096 NaN =
2730    /// the uninitialized-bits signature (P(NaN|random u32) ~ 1/256), clean by host re-read
2731    /// time — a read-before-write race, fatal via the argmax-sentinel -> embed_gather MMU
2732    /// fault, and gated on c>=2 because a backed-up primary stream widens the window.
2733    ///
2734    /// Fix law: before a ppN body enqueues NEW stage-stream work (allocations that may
2735    /// reuse freed blocks), every stage stream waits the caller's stream at its current
2736    /// point. All primary consumers of the previous round's stage-allocated buffers are
2737    /// enqueued by then (single host thread), so reuse-writes land strictly after them.
2738    /// Call at ppN-body ENTRY with the pre-`enter` caller stream. Door-shut configs never
2739    /// build a PpNRt, so single-card behavior is untouched.
2740    pub fn fence_stages_behind(
2741        &self,
2742        src: &Arc<CudaStream>,
2743    ) -> Result<(), Box<dyn std::error::Error>> {
2744        let ev = src.context().new_event(None)?;
2745        ev.record(src)?;
2746        for st in &self.stages {
2747            if Arc::ptr_eq(&st.stream, src) {
2748                continue;
2749            }
2750            st.stream.wait(&ev)?;
2751        }
2752        Ok(())
2753    }
2754
2755    /// Deferred readback: record a fresh completion event on the LAST stage's stream
2756    /// (call after the step's logits matmul has been enqueued there).
2757    pub fn record_done(&self) -> Result<CudaEvent, Box<dyn std::error::Error>> {
2758        let last = &self.stages[self.stages.len() - 1];
2759        let ev = last.ctx.new_event(None)?;
2760        ev.record(&last.stream)?;
2761        Ok(ev)
2762    }
2763
2764    /// The dedicated readback stream (last stage's context).
2765    pub fn readback_stream(&self) -> &Arc<CudaStream> {
2766        &self.readback
2767    }
2768}
2769
2770/// Service a due runtime peer probe without constructing a PP runtime on door-shut placements.
2771/// Must be called by the CUDA owner thread at a scheduling boundary.
2772pub fn service_runtime_peer_probe(
2773    e: &Engine,
2774    scheduler_idle: bool,
2775    probe_allowed: bool,
2776) -> Result<RuntimePeerProbeStatus, Box<dyn std::error::Error>> {
2777    let Some(rt) = RTN.get() else {
2778        return Ok(RuntimePeerProbeStatus::NotRun);
2779    };
2780    let rt = rt
2781        .as_ref()
2782        .map_err(|err| -> Box<dyn std::error::Error> { err.clone().into() })?;
2783    rt.service_runtime_peer_probe(e, scheduler_idle, probe_allowed)
2784}
2785
2786/// M2 increment 3: a step's logits, still device-resident on the LAST stage. `wait()`
2787/// orders the readback stream behind the step's completion event, copies, and syncs —
2788/// tokens enqueued after this step keep running on the stage streams while the caller
2789/// drains token t. Dropping without waiting is safe (buffers free stream-ordered).
2790pub struct PendingLogits {
2791    logits: CudaSlice<f32>,
2792    ev: CudaEvent,
2793    rb: Arc<CudaStream>,
2794}
2795
2796impl PendingLogits {
2797    pub fn new(logits: CudaSlice<f32>, ev: CudaEvent, rb: Arc<CudaStream>) -> Self {
2798        PendingLogits { logits, ev, rb }
2799    }
2800
2801    /// Blocks until this step's logits are computed, returns them host-side. Only this
2802    /// step's work is waited on (event-ordered) — NOT later tokens already enqueued on
2803    /// the stage streams.
2804    pub fn wait(self) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
2805        self.rb.wait(&self.ev)?;
2806        let host = self.rb.clone_dtoh(&self.logits)?;
2807        self.rb.synchronize()?;
2808        // logits drop AFTER the sync: the D2H has fully completed, so the stream-ordered
2809        // free on the compute stream cannot race the copy.
2810        Ok(host)
2811    }
2812}
2813
2814/// Bring up the PP transport while model geometry is known but before model weights upload.
2815/// Door-shut and placement-free loads remain untouched.
2816pub fn init_model_transport(
2817    e: &Engine,
2818    cfg: &memra_gguf::config::ModelConfig,
2819    n_trunk: usize,
2820) -> Result<(), Box<dyn std::error::Error>> {
2821    if pp2_streams_off() || pp2_devices_env().is_none() || pp_cuts(n_trunk).is_none() {
2822        return Ok(());
2823    }
2824    PpNRt::get(e)?.init_boundary_transport(e, cfg.n_embd as usize)
2825}
2826
2827/// Stage-owned cache allocation door: when the ppN door is open AND `MEMRA_PP_DEVICES`
2828/// is set (placement plumbing), each layer's cache is allocated by its OWNING stage's
2829/// engine — on one device this is byte-for-byte today's allocation (gated); cross-device
2830/// it puts each stage's KV on that stage's HBM. Door shut or devices unset: plain
2831/// `Cache::new` (zero behavior change). Trailing MTP/NextN layers (beyond the trunk)
2832/// map to the LAST stage.
2833pub fn new_cache(
2834    e: &Engine,
2835    cfg: &memra_gguf::config::ModelConfig,
2836    max_ctx: usize,
2837) -> Result<crate::cache::Cache, Box<dyn std::error::Error>> {
2838    let n_trunk = (cfg.n_layer - cfg.nextn_predict_layers) as usize;
2839    if let Some(fence) = pp_cuts(n_trunk) {
2840        if pp2_devices_env().is_some() && !pp2_streams_off() {
2841            let rt = PpNRt::get(e)?;
2842            rt.init_boundary_transport(e, cfg.n_embd as usize)?;
2843            let n_st = fence.len() - 1;
2844            assert_eq!(
2845                rt.n_stages(),
2846                n_st,
2847                "PpNRt stage count {} != fence stages {n_st}",
2848                rt.n_stages()
2849            );
2850            // #87 REVERSE PUBLICATION at ADMISSION (lane/pp2spec-crash): this is the one
2851            // stage-stream allocation site OUTSIDE the ppN step bodies — a NEW session's
2852            // KV alloc_zeros enqueue on the STAGE streams, and their pool blocks can be
2853            // reuse of buffers freed from ANOTHER session's in-flight verify whose
2854            // primary-stream reads are still queued (the c=2 residual: exactly one trap
2855            // per admission collision, round 0, after the step-body fences landed).
2856            // Order the stage streams behind the caller before the memsets can clobber.
2857            // Anatomy: `PpNRt::fence_stages_behind`.
2858            rt.fence_stages_behind(&e.stream())?;
2859            let devs: Vec<&dyn memra_kv::KvDev> = (0..n_st)
2860                .map(|s| rt.engine(s, e) as &dyn memra_kv::KvDev)
2861                .collect();
2862            let cache = crate::cache::Cache::new_ppn(&devs, &fence, cfg, max_ctx)?;
2863            sync_stages_after_load(e, n_trunk)?;
2864            return Ok(cache);
2865        }
2866        if !pp2_streams_off() {
2867            // CACHE BIRTH BARRIER (2026-08-02 pipelined-arm residual race): with the door
2868            // open but no device placement, Cache::new's alloc_zeros memsets enqueue on
2869            // the PRIMARY worker stream while the first KV appends / recurrent-state
2870            // reads run on the per-stage streams — no event orders them, and under
2871            // deferred readback the stage streams are hot immediately (a memset tail
2872            // can zero an already-appended KV row; intermittent, ~1-in-3 gate FAIL).
2873            // One context-sync per cache creation kills the class.
2874            let cache = crate::cache::Cache::new(e, cfg, max_ctx)?;
2875            sync_stages_after_load(e, n_trunk)?;
2876            return Ok(cache);
2877        }
2878    }
2879    crate::cache::Cache::new(e, cfg, max_ctx)
2880}
2881
2882/// M2 increment 2 LOAD BARRIER: weight uploads and decode-mirror builds enqueue on the
2883/// loading engines' WORKER streams; the first consumer launches on a DIFFERENT stream
2884/// with no load->decode event — the door-off reference walk on the primary worker
2885/// stream (sharded load: remote builds still in flight), or a fresh per-stage stream.
2886/// The 2026-08-02 gate finds (n2-dev01 step-0 168k-logit graze; split5 ref=0.0 head —
2887/// a half-built rp4 mirror — poisoning step-0 KV and every later step): one
2888/// context-wide synchronize per stage at load end kills the class. No-op when the door
2889/// is shut at load (single-stream load+decode is ordered by the stream itself).
2890pub fn sync_stages_after_load(
2891    e: &Engine,
2892    n_trunk: usize,
2893) -> Result<(), Box<dyn std::error::Error>> {
2894    if pp2_streams_off() || pp_cuts(n_trunk).is_none() {
2895        return Ok(());
2896    }
2897    let rt = PpNRt::get(e)?;
2898    for s in 0..rt.n_stages() {
2899        rt.stages[s].ctx.bind_to_thread()?;
2900        unsafe {
2901            cudarc::driver::sys::cuCtxSynchronize().result()?;
2902        }
2903    }
2904    e.ctx().bind_to_thread()?;
2905    unsafe {
2906        cudarc::driver::sys::cuCtxSynchronize().result()?;
2907    }
2908    Ok(())
2909}
2910
2911/// M2 increment 2 (weight sharding): the engine that should UPLOAD layer `il`'s weights
2912/// (and build its decode mirrors) — the owning stage's engine when the door is open with
2913/// device placement and sharding not rolled back; else the primary. `il >= n_trunk`
2914/// (MTP/NextN blocks) maps to the last stage. The head (output_norm + lm head) belongs
2915/// to the last trunk layer's stage — call with `il = n_trunk - 1`.
2916pub fn layer_engine<'a>(
2917    e: &'a Engine,
2918    n_trunk: usize,
2919    il: usize,
2920) -> Result<&'a Engine, Box<dyn std::error::Error>> {
2921    if pp_shard_off() || pp2_devices_env().is_none() || pp2_streams_off() {
2922        return Ok(e);
2923    }
2924    let Some(fence) = pp_cuts(n_trunk) else {
2925        return Ok(e);
2926    };
2927    let rt = PpNRt::get(e)?;
2928    let s = stage_of(&fence, il.min(n_trunk - 1));
2929    Ok(rt.engine(s, e))
2930}
2931
2932/// Restore a cache checkpoint through each layer's owning engine.
2933///
2934/// `source = None` is an in-place rewind: the target already owns the append-only KV bytes and
2935/// only its lengths plus recurrent state move back to the snapshot. `Some(source)` restores into
2936/// a freshly allocated larger cache: checkpoint-valid KV rows are copied from the parked cache,
2937/// while recurrent state always comes from the checkpoint's owned device copies.
2938///
2939/// This cannot use `Cache::rollback(e, ...)` under cross-device PP: a single primary engine is
2940/// not the owner of every stage's cache buffers. The rare rewind/grow boundary synchronizes open
2941/// PP contexts before publishing the restored cache to the next request.
2942pub fn restore_cache_checkpoint(
2943    e: &Engine,
2944    cfg: &memra_gguf::config::ModelConfig,
2945    source: Option<&crate::cache::Cache>,
2946    target: &mut crate::cache::Cache,
2947    snap: &crate::cache::CacheSnapshot,
2948) -> Result<(), Box<dyn std::error::Error>> {
2949    let n = target.kv.len();
2950    if target.recur.len() != n
2951        || snap.kv_len.len() != n
2952        || snap.conv.len() != n
2953        || snap.ssm.len() != n
2954        || source.is_some_and(|s| s.kv.len() != n || s.recur.len() != n)
2955    {
2956        return Err("checkpoint cache layer-count mismatch".into());
2957    }
2958    if snap.pos > target.max_ctx {
2959        return Err(format!(
2960            "checkpoint pos {} exceeds target capacity {}",
2961            snap.pos, target.max_ctx,
2962        )
2963        .into());
2964    }
2965
2966    let n_trunk = (cfg.n_layer - cfg.nextn_predict_layers) as usize;
2967    for il in 0..n {
2968        let owner = layer_engine(e, n_trunk, il)?;
2969        let src_kv = source.map(|s| &s.kv[il]);
2970        match (src_kv, target.kv[il].as_mut(), snap.kv_len[il]) {
2971            (Some(Some(src)), Some(dst), Some(len)) => {
2972                if len > src.len || len > target.max_ctx {
2973                    return Err(format!(
2974                        "checkpoint layer {il} len {len} exceeds source {} or target {}",
2975                        src.len, target.max_ctx,
2976                    )
2977                    .into());
2978                }
2979                if src.kv_dim_k != dst.kv_dim_k
2980                    || src.kv_dim_v != dst.kv_dim_v
2981                    || src.k_tok_bytes != dst.k_tok_bytes
2982                    || src.v_tok_bytes != dst.v_tok_bytes
2983                {
2984                    return Err(format!("checkpoint KV layout mismatch at layer {il}").into());
2985                }
2986                let kb = len * src.k_tok_bytes;
2987                let vb = len * src.v_tok_bytes;
2988                if kb > 0 {
2989                    owner.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
2990                }
2991                if vb > 0 {
2992                    owner.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
2993                }
2994                dst.len = len;
2995                owner.set_i32_one(&mut dst.len_d, len as i32)?;
2996            }
2997            (None, Some(dst), Some(len)) => {
2998                if len > dst.len || len > target.max_ctx {
2999                    return Err(format!(
3000                        "checkpoint layer {il} len {len} exceeds live {} or target {}",
3001                        dst.len, target.max_ctx,
3002                    )
3003                    .into());
3004                }
3005                dst.len = len;
3006                owner.set_i32_one(&mut dst.len_d, len as i32)?;
3007            }
3008            (Some(None), None, None) | (None, None, None) => {}
3009            _ => return Err(format!("checkpoint KV kind mismatch at layer {il}").into()),
3010        }
3011
3012        match (target.recur[il].as_mut(), &snap.conv[il], &snap.ssm[il]) {
3013            (Some(dst), Some(conv), Some(ssm)) => {
3014                if conv.len() != dst.conv_state.len() || ssm.len() != dst.ssm_state.len() {
3015                    return Err(
3016                        format!("checkpoint recurrent layout mismatch at layer {il}").into(),
3017                    );
3018                }
3019                owner.copy_into(&mut dst.conv_state, 0, conv, conv.len())?;
3020                owner.copy_into(&mut dst.ssm_state, 0, ssm, ssm.len())?;
3021            }
3022            (None, None, None) => {}
3023            _ => {
3024                return Err(format!("checkpoint recurrent kind mismatch at layer {il}").into());
3025            }
3026        }
3027    }
3028    target.pos = snap.pos;
3029
3030    // Open PP uses per-stage streams/contexts; publish every restored plane before the caller
3031    // starts the next prime. Door-shut single-stream restores remain naturally ordered.
3032    sync_stages_after_load(e, n_trunk)?;
3033    if source.is_some() {
3034        // A grown cache replaces and drops the source immediately after this returns. Bound the
3035        // D2D copies first so an async-pool free cannot recycle a source plane prematurely.
3036        e.stream().synchronize()?;
3037    }
3038    Ok(())
3039}
3040
3041#[cfg(test)]
3042mod host_bounce_tests {
3043    use super::{
3044        BoundaryTransport, DUAL_PP_HOST_BOUNCE_REFUSAL, DUAL_PP_SINGLE_SLOT_REFUSAL,
3045        PEER_PROBE_FIXED_BYTES, PEER_PROBE_REQUIRED_REFUSAL, PEER_PROBE_TOKEN_WIDTHS,
3046        PEER_RUNTIME_PROBE_BUDGET_NS, PEER_RUNTIME_PROBE_CYCLE_COPIES,
3047        PEER_RUNTIME_PROBE_DEFERRAL_BOUND_INTERVALS, PEER_RUNTIME_PROBE_INTERVAL_COPIES,
3048        PeerProbeDecision, PeerProbeStartupPolicy, boundary_transport, dual_pp_eligibility,
3049        dual_pp_timing_dropped, dual_pp_timing_snapshot, dual_pp_wave_mid, host_bounce_capacity,
3050        latch_runtime_host_bounce, peer_probe_bytes_to_f32, peer_probe_decision,
3051        peer_probe_f32_to_bytes, peer_probe_mismatch_count, peer_probe_pattern,
3052        peer_probe_startup_policy, publish_runtime_peer_probe_deferral,
3053        record_dual_pp_stage_result, runtime_peer_probe_candidate, runtime_peer_probe_next_copy,
3054    };
3055
3056    // ---- 2026-08-11 default-flip safety regression (owner-ordered) ----------------------
3057    // All pure-resolution tests: no env mutation (parallel test threads share process env).
3058
3059    #[test]
3060    fn flip_default_is_dual_auto_with_explicit_off_and_forced_seams() {
3061        use super::{DualPpMode, dual_pp_mode_resolve};
3062        assert_eq!(dual_pp_mode_resolve(None), DualPpMode::Auto);
3063        assert_eq!(dual_pp_mode_resolve(Some("0")), DualPpMode::Off);
3064        assert_eq!(dual_pp_mode_resolve(Some("1")), DualPpMode::Forced);
3065        // Any other value is not a silent third state: treat as the default.
3066        assert_eq!(dual_pp_mode_resolve(Some("2")), DualPpMode::Auto);
3067        assert_eq!(dual_pp_mode_resolve(Some("")), DualPpMode::Auto);
3068    }
3069
3070    #[test]
3071    fn flip_overlap_follows_mode_and_one_flag_restores_preflip_serial() {
3072        use super::{DualPpMode, pp2_overlap_resolve};
3073        // Naked default = the re-gated dual arm: overlap ON.
3074        assert!(pp2_overlap_resolve(None, DualPpMode::Auto));
3075        // MEMRA_DUAL_PP=0 ALONE restores the exact pre-flip naked path (single-slot serial).
3076        assert!(!pp2_overlap_resolve(None, DualPpMode::Off));
3077        // The explicit pre-flip request keeps its binding single-slot refusal reachable.
3078        assert!(!pp2_overlap_resolve(None, DualPpMode::Forced));
3079        // Explicit values always win over the mode.
3080        for mode in [DualPpMode::Off, DualPpMode::Forced, DualPpMode::Auto] {
3081            assert!(pp2_overlap_resolve(Some("1"), mode));
3082            assert!(!pp2_overlap_resolve(Some("0"), mode));
3083        }
3084    }
3085
3086    #[test]
3087    fn flip_auto_routes_only_the_regated_regime_and_degrades_serially_elsewhere() {
3088        use super::{DualPpMode, dual_pp_route};
3089        // The exact box1 re-gate regime: PP-2, double-slot, peer transport, B>=2.
3090        assert!(dual_pp_route(DualPpMode::Auto, 2, 2, true, false));
3091        assert!(dual_pp_route(DualPpMode::Auto, 17, 2, true, false));
3092        // Outside it, Auto must DEGRADE (serial PP-N walker), never refuse:
3093        assert!(!dual_pp_route(DualPpMode::Auto, 1, 2, true, false)); // no second wave
3094        assert!(!dual_pp_route(DualPpMode::Auto, 2, 3, true, false)); // naked PP-3 keeps serving
3095        assert!(!dual_pp_route(DualPpMode::Auto, 2, 2, false, false)); // single-slot boundary
3096        assert!(!dual_pp_route(DualPpMode::Auto, 2, 2, true, true)); // host-bounce escape hatch
3097        // Forced routes every B>=2 call into the dual body so the binding refusals fire loud.
3098        assert!(dual_pp_route(DualPpMode::Forced, 2, 3, false, true));
3099        assert!(!dual_pp_route(DualPpMode::Forced, 1, 2, true, false));
3100        // Off is the rollback seam: never dual.
3101        assert!(!dual_pp_route(DualPpMode::Off, 8, 2, true, false));
3102    }
3103
3104    #[test]
3105    fn dual_pp_split_is_honest_at_one_and_ceil_first_afterward() {
3106        assert_eq!(dual_pp_wave_mid(1), None);
3107        assert_eq!(dual_pp_wave_mid(2), Some(1));
3108        assert_eq!(dual_pp_wave_mid(3), Some(2));
3109        assert_eq!(dual_pp_wave_mid(8), Some(4));
3110        assert_eq!(dual_pp_wave_mid(16), Some(8));
3111        assert_eq!(dual_pp_wave_mid(31), Some(16));
3112        assert_eq!(dual_pp_wave_mid(32), Some(16));
3113    }
3114
3115    #[test]
3116    fn dual_pp_refuses_single_slot_and_non_pp2_shapes() {
3117        assert_eq!(
3118            dual_pp_eligibility(2, false, false),
3119            Err(DUAL_PP_SINGLE_SLOT_REFUSAL)
3120        );
3121        assert!(dual_pp_eligibility(2, true, false).is_ok());
3122        assert!(dual_pp_eligibility(3, true, false).is_err());
3123    }
3124
3125    #[test]
3126    fn dual_pp_refuses_unvalidated_host_bounce_transport() {
3127        assert_eq!(
3128            dual_pp_eligibility(2, true, true),
3129            Err(DUAL_PP_HOST_BOUNCE_REFUSAL),
3130        );
3131    }
3132
3133    #[test]
3134    fn dual_pp_timing_error_is_counted_without_recording_a_sample() {
3135        let dropped_before = dual_pp_timing_dropped();
3136        let (_, samples_before) = dual_pp_timing_snapshot();
3137        record_dual_pp_stage_result(0, Err::<f32, _>("CUDA_ERROR_NOT_READY"));
3138        let (_, samples_after) = dual_pp_timing_snapshot();
3139        assert_eq!(samples_after[0], samples_before[0]);
3140        assert!(dual_pp_timing_dropped() >= dropped_before + 1);
3141    }
3142
3143    #[test]
3144    fn corrupted_peer_readback_fails_closed_unless_host_bounce_is_selected() {
3145        assert_eq!(
3146            PEER_PROBE_TOKEN_WIDTHS,
3147            [1, 8, 16, crate::cache::PRIME_CHUNK_MAX_TOKENS],
3148        );
3149        let largest_payload_bytes = PEER_PROBE_TOKEN_WIDTHS[3] * 4096 * std::mem::size_of::<f32>();
3150        assert_eq!(largest_payload_bytes, 64 * 1024 * 1024);
3151        assert!(largest_payload_bytes >= 1024 * 1024);
3152        let expected = peer_probe_pattern(PEER_PROBE_FIXED_BYTES, 2, 0, 1);
3153        assert_eq!(
3154            peer_probe_f32_to_bytes(&peer_probe_bytes_to_f32(&expected)),
3155            expected,
3156        );
3157        let mut corrupted = expected.clone();
3158        for offset in [0, 8_191, PEER_PROBE_FIXED_BYTES - 1] {
3159            corrupted[offset] ^= 0x5a;
3160        }
3161
3162        assert_eq!(peer_probe_mismatch_count(&expected, &corrupted), 3);
3163        assert_eq!(
3164            peer_probe_decision(&expected, &corrupted, false),
3165            Err("3 mismatched byte(s)".to_string()),
3166        );
3167        assert_eq!(
3168            peer_probe_decision(&expected, &corrupted, true),
3169            Ok(PeerProbeDecision::ProceedWithHostBounce { mismatches: 3 }),
3170        );
3171    }
3172
3173    #[test]
3174    fn probe_off_refusal_matrix_is_fail_closed_only_for_sharded_native_peer() {
3175        for probe_on in [false, true] {
3176            for sharded in [false, true] {
3177                for host_bounce in [false, true] {
3178                    let got = peer_probe_startup_policy(probe_on, sharded, host_bounce);
3179                    let expected = match (probe_on, sharded, host_bounce) {
3180                        (false, true, false) => Err(PEER_PROBE_REQUIRED_REFUSAL),
3181                        (false, true, true) => Ok(PeerProbeStartupPolicy::BypassedWithHostBounce),
3182                        _ => Ok(PeerProbeStartupPolicy::Allowed),
3183                    };
3184                    assert_eq!(
3185                        got, expected,
3186                        "probe_on={probe_on} sharded={sharded} host_bounce={host_bounce}",
3187                    );
3188                }
3189            }
3190        }
3191        assert!(PEER_PROBE_REQUIRED_REFUSAL.contains("MEMRA_PEER_PROBE=0"));
3192        assert!(PEER_PROBE_REQUIRED_REFUSAL.contains("MEMRA_PP_HOST_BOUNCE!=1"));
3193    }
3194
3195    #[test]
3196    fn runtime_reprobe_keeps_cheap_deadlines_live_while_expensive_work_waits_for_idle() {
3197        let every = PEER_RUNTIME_PROBE_INTERVAL_COPIES;
3198        assert_eq!(PEER_RUNTIME_PROBE_CYCLE_COPIES, 4 * every);
3199        let mut next = [every, 2 * every, 3 * every, 4 * every];
3200        let measured_ns = [1_000_000, 2_000_000, 3_000_000, 0];
3201
3202        assert_eq!(
3203            runtime_peer_probe_candidate(every - 1, next, measured_ns, false),
3204            None,
3205        );
3206        assert_eq!(
3207            runtime_peer_probe_candidate(every, next, measured_ns, false),
3208            Some((0, 1)),
3209        );
3210
3211        // Pretend the three cheap deadlines completed. The maximum rung is due but must not run
3212        // on the interactive boundary.
3213        next[..3].copy_from_slice(&[5 * every, 6 * every, 7 * every]);
3214        assert_eq!(
3215            runtime_peer_probe_candidate(4 * every, next, measured_ns, false),
3216            None,
3217        );
3218        // Once the next cheap deadline arrives, it remains runnable even though the older max
3219        // deadline is still pending.
3220        assert_eq!(
3221            runtime_peer_probe_candidate(5 * every, next, measured_ns, false),
3222            Some((0, 1)),
3223        );
3224        // An idle boundary drains the oldest pending rung first.
3225        assert_eq!(
3226            runtime_peer_probe_candidate(5 * every, next, measured_ns, true),
3227            Some((3, crate::cache::PRIME_CHUNK_MAX_TOKENS)),
3228        );
3229    }
3230
3231    #[test]
3232    fn runtime_reprobe_moves_any_measured_over_budget_rung_to_idle_only() {
3233        let every = PEER_RUNTIME_PROBE_INTERVAL_COPIES;
3234        let next = [u64::MAX, every, u64::MAX, u64::MAX];
3235        let mut measured_ns = [0; PEER_PROBE_TOKEN_WIDTHS.len()];
3236        measured_ns[1] = PEER_RUNTIME_PROBE_BUDGET_NS + 1;
3237        assert_eq!(
3238            runtime_peer_probe_candidate(every, next, measured_ns, false),
3239            None
3240        );
3241        assert_eq!(
3242            runtime_peer_probe_candidate(every, next, measured_ns, true),
3243            Some((1, 8)),
3244        );
3245    }
3246
3247    #[test]
3248    fn late_runtime_reprobe_advances_once_instead_of_bursting_catchup() {
3249        let every = PEER_RUNTIME_PROBE_INTERVAL_COPIES;
3250        let due = every;
3251        assert_eq!(runtime_peer_probe_next_copy(due, due), due + 4 * every);
3252        assert_eq!(runtime_peer_probe_next_copy(due, 20 * every), 21 * every);
3253    }
3254
3255    #[test]
3256    fn runtime_reprobe_deferral_metric_counts_intervals_and_publishes_bound_state() {
3257        use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
3258
3259        assert_eq!(
3260            PEER_RUNTIME_PROBE_DEFERRAL_BOUND_INTERVALS * PEER_RUNTIME_PROBE_INTERVAL_COPIES,
3261            PEER_RUNTIME_PROBE_CYCLE_COPIES,
3262        );
3263        let deferred = AtomicU64::new(0);
3264        let degraded = AtomicBool::new(false);
3265        publish_runtime_peer_probe_deferral(&deferred, &degraded, 1, false);
3266        assert_eq!(deferred.load(Ordering::Relaxed), 1);
3267        assert!(!degraded.load(Ordering::Acquire));
3268
3269        publish_runtime_peer_probe_deferral(
3270            &deferred,
3271            &degraded,
3272            PEER_RUNTIME_PROBE_DEFERRAL_BOUND_INTERVALS - 1,
3273            true,
3274        );
3275        assert_eq!(
3276            deferred.load(Ordering::Relaxed),
3277            PEER_RUNTIME_PROBE_DEFERRAL_BOUND_INTERVALS,
3278        );
3279        assert!(degraded.load(Ordering::Acquire));
3280    }
3281
3282    #[test]
3283    fn runtime_probe_failure_latches_native_before_publishing_validated_bounce() {
3284        use std::sync::atomic::{AtomicBool, Ordering};
3285
3286        let failed = AtomicBool::new(false);
3287        let degraded = AtomicBool::new(false);
3288        let armed = latch_runtime_host_bounce(&failed, &degraded, || Ok::<_, String>(()));
3289        assert!(armed.is_ok());
3290        assert!(failed.load(Ordering::Acquire));
3291        assert!(degraded.load(Ordering::Acquire));
3292
3293        let failed = AtomicBool::new(false);
3294        let degraded = AtomicBool::new(false);
3295        let refused = latch_runtime_host_bounce(&failed, &degraded, || {
3296            Err::<(), _>("injected staging mismatch".to_string())
3297        });
3298        assert_eq!(refused, Err("injected staging mismatch".to_string()));
3299        assert!(failed.load(Ordering::Acquire));
3300        assert!(!degraded.load(Ordering::Acquire));
3301    }
3302
3303    #[test]
3304    fn transport_selection_keeps_peer_default_and_bounces_only_cross_device() {
3305        assert_eq!(boundary_transport(false, false), BoundaryTransport::Local);
3306        assert_eq!(boundary_transport(false, true), BoundaryTransport::Local);
3307        assert_eq!(boundary_transport(true, false), BoundaryTransport::Peer);
3308        assert_eq!(
3309            boundary_transport(true, true),
3310            BoundaryTransport::HostBounce
3311        );
3312    }
3313
3314    #[test]
3315    fn step37_geometry_sizes_each_slot_from_the_prime_cap() {
3316        let (elems, bytes) = host_bounce_capacity(4096).expect("valid Step-3.7 geometry");
3317        assert_eq!(elems, 4096 * crate::cache::PRIME_CHUNK_MAX_TOKENS);
3318        assert_eq!(bytes, 64 * 1024 * 1024);
3319    }
3320
3321    #[test]
3322    fn host_bounce_capacity_rejects_invalid_or_overflowing_geometry() {
3323        assert!(host_bounce_capacity(0).is_err());
3324        assert!(host_bounce_capacity(usize::MAX).is_err());
3325    }
3326}