1use 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
83pub 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 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
155pub fn pp2_split(n_layers: usize) -> Option<usize> {
158 pp_cuts(n_layers).filter(|f| f.len() == 3).map(|f| f[1])
159}
160
161pub 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 Ok(k) => k + 1,
167 Err(k) => k,
168 }
169}
170
171pub fn pp2_streams_off() -> bool {
174 matches!(std::env::var("MEMRA_PP_STREAMS").as_deref(), Ok("0"))
175}
176
177pub 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, 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 }
206 }
207}
208
209pub 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 if !stages_open || pp_shard_off() || pp2_streams_off() {
233 return false;
234 }
235 match pp2_devices_env() {
236 None => false, 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
246pub 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
287pub fn batch_pp_on() -> bool {
295 std::env::var("MEMRA_BATCH_PP").as_deref() != Ok("0")
296}
297
298#[derive(Clone, Copy, PartialEq, Eq, Debug)]
316pub enum DualPpMode {
317 Off,
318 Forced,
319 Auto,
320}
321
322pub 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
336pub fn dual_pp_on() -> bool {
339 dual_pp_mode() != DualPpMode::Off
340}
341
342pub 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
364pub 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
369pub fn dual_pp_wave_mid(batch: usize) -> Option<usize> {
372 (batch >= 2).then_some((batch + 1) / 2)
373}
374
375pub 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
396static 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
428pub(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
444pub 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
453pub 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
471pub(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
496pub 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
521pub fn prime_pp_on() -> bool {
531 std::env::var("MEMRA_PRIME_PP").as_deref() != Ok("0")
532}
533
534pub fn prime_pipe_on() -> bool {
539 std::env::var("MEMRA_PRIME_PIPE").as_deref() != Ok("0")
540}
541
542pub static PRIME_SPLIT_CHUNKS: AtomicUsize = AtomicUsize::new(0);
549
550pub fn prime_split_chunks() -> usize {
552 PRIME_SPLIT_CHUNKS.load(Ordering::Relaxed)
553}
554
555pub static PRIME_PIPE_OVERLAPS: AtomicUsize = AtomicUsize::new(0);
560
561pub 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
570pub(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
587pub 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
601pub fn spec_pp_on() -> bool {
609 std::env::var("MEMRA_SPEC_PP").as_deref() != Ok("0")
610}
611
612pub 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
629pub 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
639pub fn pp_host_bounce_on() -> bool {
642 matches!(std::env::var("MEMRA_PP_HOST_BOUNCE").as_deref(), Ok("1"))
643}
644
645pub 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
652pub fn pp_shard_off() -> bool {
656 matches!(std::env::var("MEMRA_PP_SHARD").as_deref(), Ok("0"))
657}
658
659fn 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);
675pub 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
686pub struct StageRt {
694 pub dev: usize,
695 pub ctx: Arc<CudaContext>,
696 pub stream: Arc<CudaStream>,
697 engine: Option<Engine>,
699}
700
701struct BoundarySlot {
706 buf: Mutex<Option<CudaSlice<f32>>>,
707 ev_tx: CudaEvent,
710 ev_rx: CudaEvent,
714}
715
716struct BoundaryRt {
720 slots: [BoundarySlot; 2],
721 step: AtomicUsize,
722 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
752pub const PEER_RUNTIME_PROBE_INTERVAL_COPIES: u64 = 8 * 1024;
755pub const PEER_RUNTIME_PROBE_CYCLE_COPIES: u64 =
757 PEER_RUNTIME_PROBE_INTERVAL_COPIES * PEER_PROBE_TOKEN_WIDTHS.len() as u64;
758pub const PEER_RUNTIME_PROBE_DEFERRAL_BOUND_INTERVALS: u64 = PEER_PROBE_TOKEN_WIDTHS.len() as u64;
761const 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
774pub 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
858pub 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
869pub 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
879fn 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
903fn 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
910fn 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
988struct 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
1175struct 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 cross_any: bool,
1274 host_bounce: bool,
1277 peer_probe: bool,
1279 peer_capable: Vec<(usize, usize)>,
1281 peer_probe_geometry: OnceLock<Result<usize, String>>,
1283 bounce: OnceLock<Result<HostBounceRt, String>>,
1285 readback: Arc<CudaStream>,
1288}
1289
1290pub type Pp2Rt = PpNRt;
1292
1293static RTN: OnceLock<Result<PpNRt, String>> = OnceLock::new();
1294
1295impl PpNRt {
1296 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 pub fn enter(&self, s: usize) -> memra_runtime::StreamOverride {
2451 memra_runtime::push_stream_override(self.stages[s].stream.clone())
2452 }
2453
2454 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 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 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 let s_rx = &self.stages[path.dst_stage].stream;
2560 *guard = Some(s_rx.alloc_zeros::<f32>(n)?);
2561 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 s_tx.memcpy_dtoh(x, host.prefix_mut(n))?;
2594 }
2595 BoundaryTransport::Peer => {
2596 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 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 s_rx.memcpy_htod(host.prefix(n), &mut dst)?;
2664 }
2665 let mut work = unsafe { s_rx.alloc::<f32>(n)? };
2668 s_rx.memcpy_dtod(&buf.slice(0..n), &mut work)?;
2672 sl.ev_rx.record(s_rx)?;
2673 Ok(work)
2674 }
2675
2676 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 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 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 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 pub fn readback_stream(&self) -> &Arc<CudaStream> {
2766 &self.readback
2767 }
2768}
2769
2770pub 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
2786pub 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 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 Ok(host)
2811 }
2812}
2813
2814pub 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
2827pub 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 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 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
2882pub 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
2911pub 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
2932pub 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 sync_stages_after_load(e, n_trunk)?;
3033 if source.is_some() {
3034 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 #[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 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 assert!(pp2_overlap_resolve(None, DualPpMode::Auto));
3075 assert!(!pp2_overlap_resolve(None, DualPpMode::Off));
3077 assert!(!pp2_overlap_resolve(None, DualPpMode::Forced));
3079 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 assert!(dual_pp_route(DualPpMode::Auto, 2, 2, true, false));
3091 assert!(dual_pp_route(DualPpMode::Auto, 17, 2, true, false));
3092 assert!(!dual_pp_route(DualPpMode::Auto, 1, 2, true, false)); assert!(!dual_pp_route(DualPpMode::Auto, 2, 3, true, false)); assert!(!dual_pp_route(DualPpMode::Auto, 2, 2, false, false)); assert!(!dual_pp_route(DualPpMode::Auto, 2, 2, true, true)); assert!(dual_pp_route(DualPpMode::Forced, 2, 3, false, true));
3099 assert!(!dual_pp_route(DualPpMode::Forced, 1, 2, true, false));
3100 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 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 assert_eq!(
3221 runtime_peer_probe_candidate(5 * every, next, measured_ns, false),
3222 Some((0, 1)),
3223 );
3224 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, °raded, 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 °raded,
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, °raded, || 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, °raded, || {
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}