use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use cudarc::driver::{CudaContext, CudaEvent, CudaSlice, CudaStream};
use crate::Engine;
pub fn pp_cuts(n_layers: usize) -> Option<Vec<usize>> {
let n_st: usize = match std::env::var("MEMRA_PP_STAGES") {
Ok(v) if v.is_empty() || v == "0" || v == "1" => return None,
Ok(v) => match v.parse::<usize>() {
Ok(n) => n,
Err(_) => {
warn_bad_once(&format!("MEMRA_PP_STAGES={v} unparseable; door stays OFF"));
return None;
}
},
Err(_) => return None,
};
if n_st < 2 || n_st > n_layers {
warn_bad_once(&format!(
"MEMRA_PP_STAGES={n_st} outside [2, n_layers={n_layers}]; door stays OFF"
));
return None;
}
let mut fence = Vec::with_capacity(n_st + 1);
fence.push(0usize);
if let Ok(s) = std::env::var("MEMRA_PP_SPLITS") {
let parts: Result<Vec<usize>, _> =
s.split(',').map(|p| p.trim().parse::<usize>()).collect();
match parts {
Ok(cuts) if cuts.len() == n_st - 1 => fence.extend(cuts),
_ => {
warn_bad_once(&format!(
"MEMRA_PP_SPLITS={s} invalid (want {} comma-separated cuts); door stays OFF",
n_st - 1
));
return None;
}
}
} else if let Ok(v) = std::env::var("MEMRA_PP_SPLIT") {
if n_st != 2 {
warn_bad_once(&format!(
"MEMRA_PP_SPLIT={v} set with MEMRA_PP_STAGES={n_st}; use MEMRA_PP_SPLITS \
for N>2 — door stays OFF"
));
return None;
}
match v.parse::<usize>() {
Ok(c) => fence.push(c),
Err(_) => {
warn_bad_once(&format!("MEMRA_PP_SPLIT={v} unparseable; door stays OFF"));
return None;
}
}
} else {
for s in 1..n_st {
fence.push(s * n_layers / n_st);
}
}
fence.push(n_layers);
for w in fence.windows(2) {
if w[0] >= w[1] {
warn_bad_once(&format!(
"pp stage fence {fence:?} not strictly increasing over [0, {n_layers}]; \
door stays OFF"
));
return None;
}
}
Some(fence)
}
pub fn pp2_split(n_layers: usize) -> Option<usize> {
pp_cuts(n_layers).filter(|f| f.len() == 3).map(|f| f[1])
}
pub fn stage_of(fence: &[usize], il: usize) -> usize {
debug_assert!(fence.len() >= 2);
match fence[1..fence.len() - 1].binary_search(&il) {
Ok(k) => k + 1,
Err(k) => k,
}
}
pub fn pp2_streams_off() -> bool {
matches!(std::env::var("MEMRA_PP_STREAMS").as_deref(), Ok("0"))
}
pub fn pp_multi_stream_same_device() -> bool {
let stages_open = std::env::var("MEMRA_PP_STAGES")
.map(|v| v.parse::<usize>().map(|n| n >= 2).unwrap_or(false))
.unwrap_or(false);
let devices = std::env::var("MEMRA_PP_DEVICES").ok().filter(|v| !v.is_empty());
if (!stages_open && devices.is_none()) || pp2_streams_off() {
return false;
}
match devices {
None => true, Some(s) => {
let mut v: Vec<&str> = s.split(',').map(|p| p.trim()).collect();
let n = v.len();
v.sort_unstable();
v.dedup();
v.len() < n }
}
}
pub fn pp_sharded_cross_device() -> bool {
let stages_open = std::env::var("MEMRA_PP_STAGES")
.map(|v| v.parse::<usize>().map(|n| n >= 2).unwrap_or(false))
.unwrap_or(false);
if !stages_open || pp_shard_off() || pp2_streams_off() {
return false;
}
match pp2_devices_env() {
None => false, Some(s) => {
let mut v: Vec<&str> = s.split(',').map(|p| p.trim()).collect();
v.sort_unstable();
v.dedup();
v.len() >= 2
}
}
}
pub fn refuse_unsplit_if_remote(path: &str, alt: &str) -> Result<(), Box<dyn std::error::Error>> {
if pp_host_bounce_active() {
return Err(format!(
"{path}: refused with MEMRA_PP_HOST_BOUNCE=1 on sharded cross-device PP — \
this unsplit path peer-reads remote weights, while host bounce covers only \
explicit stage-boundary transfers. Use {alt}; the \
MEMRA_PP_ALLOW_UNSPLIT_BATCH override is unavailable on a broken-peer host."
)
.into());
}
if pp_sharded_cross_device()
&& std::env::var("MEMRA_PP_ALLOW_UNSPLIT_BATCH").as_deref() != Ok("1")
{
return Err(format!(
"{path}: refused with the ppN door open across 2+ devices — this path has no pp \
stage split, so it would walk ALL layers on one stream and peer-read every \
remote stage's weights each step (measured 28x slower at B=1, 13.9x at B=8 on \
a PRO 6000 pair over PCIe Gen5 x16 P2P; research/pp2-hardening-20260806). \
Exactness is unaffected — peer reads return identical bytes and the exactness \
gates PASS on this config — which is exactly why it must refuse instead of \
being caught by a gate. Fixes, in order: {alt}; or MEMRA_PP_SHARD=0 (all \
weights home on the primary — full speed, forfeits the capacity PP-2 exists \
for); or close the pp door. MEMRA_PP_ALLOW_UNSPLIT_BATCH=1 overrides for \
measurement."
)
.into());
}
Ok(())
}
pub fn batch_pp_on() -> bool {
std::env::var("MEMRA_BATCH_PP").as_deref() != Ok("0")
}
pub fn prime_pp_on() -> bool {
std::env::var("MEMRA_PRIME_PP").as_deref() != Ok("0")
}
pub fn prime_pipe_on() -> bool {
std::env::var("MEMRA_PRIME_PIPE").as_deref() != Ok("0")
}
pub static PRIME_SPLIT_CHUNKS: AtomicUsize = AtomicUsize::new(0);
pub fn prime_split_chunks() -> usize {
PRIME_SPLIT_CHUNKS.load(Ordering::Relaxed)
}
pub static PRIME_PIPE_OVERLAPS: AtomicUsize = AtomicUsize::new(0);
pub fn prime_pipe_overlaps() -> usize {
PRIME_PIPE_OVERLAPS.load(Ordering::Relaxed)
}
static PRIME_PIPE_ACTIVE_STAGES: AtomicUsize = AtomicUsize::new(0);
pub(crate) struct PrimePipeStageGuard;
pub(crate) fn enter_prime_pipe_stage() -> PrimePipeStageGuard {
let active = PRIME_PIPE_ACTIVE_STAGES.fetch_add(1, Ordering::AcqRel);
if active > 0 {
PRIME_PIPE_OVERLAPS.fetch_add(1, Ordering::Relaxed);
}
PrimePipeStageGuard
}
impl Drop for PrimePipeStageGuard {
fn drop(&mut self) {
let active = PRIME_PIPE_ACTIVE_STAGES.fetch_sub(1, Ordering::AcqRel);
debug_assert!(active > 0, "prime pipeline active-stage counter underflow");
}
}
pub static STEP35_PRIME_BATCHES: AtomicUsize = AtomicUsize::new(0);
pub static STEP35_PRIME_BATCH_SPLITS: AtomicUsize = AtomicUsize::new(0);
pub fn step35_prime_batches() -> usize {
STEP35_PRIME_BATCHES.load(Ordering::Relaxed)
}
pub fn step35_prime_batch_splits() -> usize {
STEP35_PRIME_BATCH_SPLITS.load(Ordering::Relaxed)
}
pub fn spec_pp_on() -> bool {
std::env::var("MEMRA_SPEC_PP").as_deref() != Ok("0")
}
pub fn pp2_overlap() -> bool {
matches!(std::env::var("MEMRA_PP_OVERLAP").as_deref(), Ok("1"))
}
pub fn pp_host_bounce_on() -> bool {
matches!(std::env::var("MEMRA_PP_HOST_BOUNCE").as_deref(), Ok("1"))
}
pub fn pp_host_bounce_active() -> bool {
pp_host_bounce_on() && pp_sharded_cross_device()
}
pub fn pp_shard_off() -> bool {
matches!(std::env::var("MEMRA_PP_SHARD").as_deref(), Ok("0"))
}
fn pp2_devices_env() -> Option<String> {
std::env::var("MEMRA_PP_DEVICES").ok().filter(|v| !v.is_empty())
}
static WARNED_BAD: AtomicBool = AtomicBool::new(false);
fn warn_bad_once(msg: &str) {
if !WARNED_BAD.swap(true, Ordering::Relaxed) {
eprintln!("[pp] {msg}");
}
}
static WARNED_UNWIRED: AtomicBool = AtomicBool::new(false);
pub fn warn_unwired_once(path: &str) {
let open = std::env::var("MEMRA_PP_STAGES")
.map(|v| !v.is_empty() && v != "0" && v != "1")
.unwrap_or(false);
if open && !WARNED_UNWIRED.swap(true, Ordering::Relaxed) {
eprintln!(
"[pp] MEMRA_PP_STAGES set but `{path}` has no pp arm at this N; running unsplit"
);
}
}
pub struct StageRt {
pub dev: usize,
pub ctx: Arc<CudaContext>,
pub stream: Arc<CudaStream>,
engine: Option<Engine>,
}
struct BoundarySlot {
buf: Mutex<Option<CudaSlice<f32>>>,
ev_tx: CudaEvent,
ev_rx: CudaEvent,
}
struct BoundaryRt {
slots: [BoundarySlot; 2],
step: AtomicUsize,
cross: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum BoundaryTransport {
Local,
Peer,
HostBounce,
}
fn boundary_transport(cross: bool, host_bounce: bool) -> BoundaryTransport {
match (cross, host_bounce) {
(false, _) => BoundaryTransport::Local,
(true, false) => BoundaryTransport::Peer,
(true, true) => BoundaryTransport::HostBounce,
}
}
fn host_bounce_capacity(n_embd: usize) -> Result<(usize, usize), String> {
if n_embd == 0 {
return Err("MEMRA_PP_HOST_BOUNCE needs non-zero model n_embd".into());
}
let elems = n_embd
.checked_mul(crate::cache::PRIME_CHUNK_MAX_TOKENS)
.ok_or_else(|| format!("host-bounce element count overflows for n_embd={n_embd}"))?;
let bytes = elems
.checked_mul(std::mem::size_of::<f32>())
.ok_or_else(|| format!("host-bounce byte count overflows for n_embd={n_embd}"))?;
Ok((elems, bytes))
}
struct PinnedHostBounce {
ptr: *mut f32,
len: usize,
}
unsafe impl Send for PinnedHostBounce {}
unsafe impl Sync for PinnedHostBounce {}
impl PinnedHostBounce {
fn new(len: usize) -> Result<Self, Box<dyn std::error::Error>> {
let bytes = len
.checked_mul(std::mem::size_of::<f32>())
.ok_or("host-bounce pinned allocation size overflow")?;
let ptr = unsafe {
cudarc::driver::result::malloc_host(
bytes,
cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE,
)?
} as *mut f32;
if ptr.is_null() {
return Err("cuMemHostAlloc returned a null host-bounce pointer".into());
}
Ok(Self { ptr, len })
}
fn prefix(&self, n: usize) -> &[f32] {
assert!(n <= self.len, "host-bounce source {n} > capacity {}", self.len);
unsafe { std::slice::from_raw_parts(self.ptr, n) }
}
fn prefix_mut(&mut self, n: usize) -> &mut [f32] {
assert!(n <= self.len, "host-bounce destination {n} > capacity {}", self.len);
unsafe { std::slice::from_raw_parts_mut(self.ptr, n) }
}
}
impl Drop for PinnedHostBounce {
fn drop(&mut self) {
let _ = unsafe { cudarc::driver::result::free_host(self.ptr.cast()) };
}
}
struct HostBounceRt {
n_embd: usize,
capacity: usize,
slots: Vec<Option<[Mutex<PinnedHostBounce>; 2]>>,
}
impl HostBounceRt {
fn new(n_embd: usize, boundaries: &[BoundaryRt]) -> Result<Self, Box<dyn std::error::Error>> {
let (capacity, _) = host_bounce_capacity(n_embd)?;
let mut slots = Vec::with_capacity(boundaries.len());
for boundary in boundaries {
slots.push(if boundary.cross {
Some([
Mutex::new(PinnedHostBounce::new(capacity)?),
Mutex::new(PinnedHostBounce::new(capacity)?),
])
} else {
None
});
}
Ok(Self { n_embd, capacity, slots })
}
fn slot(
&self,
boundary: usize,
slot: usize,
) -> Result<&Mutex<PinnedHostBounce>, Box<dyn std::error::Error>> {
self.slots
.get(boundary)
.and_then(Option::as_ref)
.and_then(|slots| slots.get(slot))
.ok_or_else(|| format!("host-bounce slot {boundary}:{slot} is not initialized").into())
}
}
pub struct PpNRt {
stages: Vec<StageRt>,
boundaries: Vec<BoundaryRt>,
cross_any: bool,
host_bounce: bool,
bounce: OnceLock<Result<HostBounceRt, String>>,
readback: Arc<CudaStream>,
}
pub type Pp2Rt = PpNRt;
static RTN: OnceLock<Result<PpNRt, String>> = OnceLock::new();
impl PpNRt {
pub fn get(e: &Engine) -> Result<&'static PpNRt, Box<dyn std::error::Error>> {
RTN.get_or_init(|| Self::build(e).map_err(|err| err.to_string()))
.as_ref()
.map_err(|s| -> Box<dyn std::error::Error> { s.clone().into() })
}
fn build(e: &Engine) -> Result<PpNRt, Box<dyn std::error::Error>> {
let primary_dev = e.ctx().ordinal();
let devices: Vec<usize> = match pp2_devices_env() {
Some(s) => {
let parts: Result<Vec<usize>, _> =
s.split(',').map(|p| p.trim().parse::<usize>()).collect();
match parts {
Ok(v) if v.len() >= 2 => v,
_ => {
return Err(format!(
"MEMRA_PP_DEVICES={s} unparseable (want <d0>,..,<dN-1> e.g. 0,1,2,3)"
)
.into())
}
}
}
None => {
let n_st = std::env::var("MEMRA_PP_STAGES")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|&n| n >= 2)
.unwrap_or(2);
vec![primary_dev; n_st]
}
};
if let Ok(v) = std::env::var("MEMRA_PP_STAGES") {
if let Ok(n) = v.parse::<usize>() {
if n >= 2 && n != devices.len() {
return Err(format!(
"MEMRA_PP_DEVICES lists {} devices but MEMRA_PP_STAGES={n} — \
refusing an ambiguous placement",
devices.len()
)
.into());
}
}
}
let n_st = devices.len();
let cross_any = devices.iter().any(|&d| d != devices[0]);
let host_bounce = pp_host_bounce_on();
if host_bounce && cross_any {
if pp_shard_off() {
return Err(
"MEMRA_PP_HOST_BOUNCE=1 refuses MEMRA_PP_SHARD=0: the boundary can bounce, \
but remote stages would still peer-read primary-device weights"
.into(),
);
}
if devices.last().copied() != Some(primary_dev) {
return Err(format!(
"MEMRA_PP_HOST_BOUNCE=1 requires the primary engine on the last/head stage \
(primary dev{primary_dev}, placement {devices:?}); otherwise returned \
logits/hidden state remain peer reads"
)
.into());
}
}
let mut used: Vec<usize> = devices.clone();
used.push(primary_dev);
used.sort_unstable();
used.dedup();
if used.len() > 1 {
let n = cudarc::driver::result::device::get_count()? as usize;
for &d in &used {
if d >= n {
return Err(format!(
"MEMRA_PP_DEVICES={devices:?} but only {n} CUDA device(s) present"
)
.into());
}
}
if !host_bounce {
for &a in &used {
for &b in &used {
if a == b {
continue;
}
let da = cudarc::driver::result::device::get(a as i32)?;
let db = cudarc::driver::result::device::get(b as i32)?;
let mut can: i32 = 0;
unsafe {
cudarc::driver::sys::cuDeviceCanAccessPeer(&mut can, da, db).result()?;
}
if can == 0 {
return Err(format!(
"device {a} cannot peer-access device {b} \
(cuDeviceCanAccessPeer=0); ppN cross-device needs P2P — \
refusing a silently-staged path"
)
.into());
}
}
}
}
}
let mk_stage = |dev: usize, s: usize| -> Result<StageRt, Box<dyn std::error::Error>> {
if dev == primary_dev && s == 0 {
let ctx = e.ctx().clone();
let stream = ctx.new_stream()?;
Ok(StageRt { dev, ctx, stream, engine: None })
} else {
let eng = Engine::new(dev)?;
let ctx = eng.ctx().clone();
let stream = ctx.new_stream()?;
Ok(StageRt { dev, ctx, stream, engine: Some(eng) })
}
};
let mut stages = Vec::with_capacity(n_st);
for (s, &d) in devices.iter().enumerate() {
stages.push(mk_stage(d, s)?);
}
if used.len() > 1 {
if !host_bounce {
let ctx_of = |d: usize| -> &Arc<CudaContext> {
if d == primary_dev {
e.ctx()
} else {
&stages.iter().find(|s| s.dev == d).unwrap().ctx
}
};
for &a in &used {
for &b in &used {
if a == b {
continue;
}
ctx_of(a).bind_to_thread()?;
let rc = unsafe {
cudarc::driver::sys::cuCtxEnablePeerAccess(ctx_of(b).cu_ctx(), 0)
};
use cudarc::driver::sys::cudaError_enum as E;
if rc != E::CUDA_SUCCESS && rc != E::CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED {
return Err(format!(
"cuCtxEnablePeerAccess(dev{a} -> dev{b}) failed: {rc:?}"
)
.into());
}
}
}
for &owner in &used {
for &accessor in &used {
if owner == accessor {
continue;
}
let dev = cudarc::driver::result::device::get(owner as i32)?;
let mut pool: cudarc::driver::sys::CUmemoryPool = std::ptr::null_mut();
unsafe {
cudarc::driver::sys::cuDeviceGetDefaultMemPool(&mut pool, dev).result()?;
}
let desc = cudarc::driver::sys::CUmemAccessDesc {
location: cudarc::driver::sys::CUmemLocation {
type_: cudarc::driver::sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
id: accessor as i32,
},
flags: cudarc::driver::sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_READWRITE,
};
let rc = unsafe { cudarc::driver::sys::cuMemPoolSetAccess(pool, &desc, 1) };
if rc != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
return Err(format!(
"cuMemPoolSetAccess(dev{owner} pool -> dev{accessor}) failed: {rc:?}"
)
.into());
}
}
}
for (owner, accessor) in [(stages[0].dev, stages[1].dev), (stages[1].dev, stages[0].dev)] {
let dev = cudarc::driver::result::device::get(owner as i32)?;
let mut pool: cudarc::driver::sys::CUmemoryPool = std::ptr::null_mut();
unsafe {
cudarc::driver::sys::cuDeviceGetDefaultMemPool(&mut pool, dev).result()?;
}
let desc = cudarc::driver::sys::CUmemAccessDesc {
location: cudarc::driver::sys::CUmemLocation {
type_: cudarc::driver::sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
id: accessor as i32,
},
flags: cudarc::driver::sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_READWRITE,
};
let rc = unsafe { cudarc::driver::sys::cuMemPoolSetAccess(pool, &desc, 1) };
if rc != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
return Err(format!(
"cuMemPoolSetAccess(dev{owner} pool -> dev{accessor}) failed: {rc:?}"
)
.into());
}
}
e.ctx().bind_to_thread()?;
eprintln!(
"[pp] cross-device transport: {} (cudaMemcpyPeerAsync per cross boundary; \
peer + default-pool access granted all pairs over {used:?}; weight home: {})",
devices
.iter()
.enumerate()
.map(|(s, d)| format!("stage{s}=dev{d}"))
.collect::<Vec<_>>()
.join(" "),
if pp_shard_off() {
format!("dev{primary_dev} (MEMRA_PP_SHARD=0 bring-up placement)")
} else {
"per-stage (sharded loader)".to_string()
}
);
} else {
e.ctx().bind_to_thread()?;
eprintln!(
"[pp] cross-device transport: {} (HOST-STAGED pinned D2H -> H2D per cross \
boundary; MEMRA_PP_HOST_BOUNCE=1; peer access and peer-pool grants \
bypassed; weight home: per-stage (sharded loader))",
devices
.iter()
.enumerate()
.map(|(s, d)| format!("stage{s}=dev{d}"))
.collect::<Vec<_>>()
.join(" "),
);
}
}
let mk_slot = |tx: &StageRt, rx: &StageRt| -> Result<BoundarySlot, Box<dyn std::error::Error>> {
Ok(BoundarySlot {
buf: Mutex::new(None),
ev_tx: tx.ctx.new_event(None)?,
ev_rx: rx.ctx.new_event(None)?,
})
};
let mut boundaries = Vec::with_capacity(n_st - 1);
for b in 0..n_st - 1 {
let (tx, rx) = (&stages[b], &stages[b + 1]);
boundaries.push(BoundaryRt {
slots: [mk_slot(tx, rx)?, mk_slot(tx, rx)?],
step: AtomicUsize::new(0),
cross: tx.dev != rx.dev,
});
}
let readback = stages[n_st - 1].ctx.new_stream()?;
Ok(PpNRt {
stages,
boundaries,
cross_any,
host_bounce,
bounce: OnceLock::new(),
readback,
})
}
pub fn n_stages(&self) -> usize {
self.stages.len()
}
pub fn cross_device(&self) -> bool {
self.cross_any
}
pub fn init_host_bounce(
&self,
e: &Engine,
n_embd: usize,
) -> Result<(), Box<dyn std::error::Error>> {
if !self.host_bounce || !self.cross_any {
return Ok(());
}
e.ctx().bind_to_thread()?;
let result = self.bounce.get_or_init(|| {
HostBounceRt::new(n_embd, &self.boundaries)
.map(|rt| {
let bytes = rt.capacity * std::mem::size_of::<f32>();
eprintln!(
"[pp] host-bounce staging ready: n_embd={n_embd} max_tokens={} \
slot_bytes={bytes} slots_per_cross_boundary=2",
crate::cache::PRIME_CHUNK_MAX_TOKENS,
);
rt
})
.map_err(|err| err.to_string())
});
let bounce = result
.as_ref()
.map_err(|err| -> Box<dyn std::error::Error> { err.clone().into() })?;
if bounce.n_embd != n_embd {
return Err(format!(
"host-bounce runtime initialized for n_embd={} but model requests n_embd={n_embd}; \
one PP runtime supports one model geometry per process",
bounce.n_embd,
)
.into());
}
Ok(())
}
fn bounce_rt(&self) -> Result<&HostBounceRt, Box<dyn std::error::Error>> {
self.bounce
.get()
.ok_or_else(|| -> Box<dyn std::error::Error> {
"MEMRA_PP_HOST_BOUNCE=1 staging was not initialized from model geometry".into()
})?
.as_ref()
.map_err(|err| -> Box<dyn std::error::Error> { err.clone().into() })
}
pub fn engine<'a>(&'a self, s: usize, primary: &'a Engine) -> &'a Engine {
self.stages[s].engine.as_ref().unwrap_or(primary)
}
pub fn bind_stage(&self, s: usize) -> Result<(), Box<dyn std::error::Error>> {
self.stages[s].ctx.bind_to_thread()?;
Ok(())
}
pub fn enter(&self, s: usize) -> memra_runtime::StreamOverride {
memra_runtime::push_stream_override(self.stages[s].stream.clone())
}
pub fn prepare_overlap_slots(&self, b: usize, n: usize)
-> Result<(), Box<dyn std::error::Error>> {
let bd = &self.boundaries[b];
let s_rx = &self.stages[b + 1].stream;
let mut grew = false;
for sl in &bd.slots {
let mut guard = sl.buf.lock().unwrap();
if guard.as_ref().map(|bf| bf.len() < n).unwrap_or(true) {
*guard = Some(s_rx.alloc_zeros::<f32>(n)?);
grew = true;
}
}
if grew {
s_rx.synchronize()?;
}
Ok(())
}
pub fn tx(&self, b: usize, x: &CudaSlice<f32>, n: usize)
-> Result<usize, Box<dyn std::error::Error>> {
assert_eq!(x.len(), n, "pp tx: residual length mismatch");
let bd = &self.boundaries[b];
let slot_idx = if pp2_overlap() {
bd.step.fetch_add(1, Ordering::Relaxed) % 2
} else {
0
};
self.tx_slot(b, x, n, slot_idx)
}
pub fn tx_pipelined(&self, b: usize, x: &CudaSlice<f32>, n: usize)
-> Result<usize, Box<dyn std::error::Error>> {
assert_eq!(x.len(), n, "pp tx: residual length mismatch");
let slot_idx = self.boundaries[b].step.fetch_add(1, Ordering::Relaxed) % 2;
self.tx_slot(b, x, n, slot_idx)
}
fn tx_slot(&self, b: usize, x: &CudaSlice<f32>, n: usize, slot_idx: usize)
-> Result<usize, Box<dyn std::error::Error>> {
debug_assert!(slot_idx < 2);
let bd = &self.boundaries[b];
let sl = &bd.slots[slot_idx];
let s_tx = &self.stages[b].stream;
s_tx.wait(&sl.ev_rx)?;
let mut guard = sl.buf.lock().unwrap();
if guard.as_ref().map(|bf| bf.len() < n).unwrap_or(true) {
let s_rx = &self.stages[b + 1].stream;
*guard = Some(s_rx.alloc_zeros::<f32>(n)?);
s_rx.synchronize()?;
}
let buf = guard.as_mut().unwrap();
match boundary_transport(bd.cross, self.host_bounce) {
BoundaryTransport::Local => s_tx.memcpy_dtod(x, buf)?,
BoundaryTransport::HostBounce => {
let bounce = self.bounce_rt()?;
if n > bounce.capacity {
return Err(format!(
"pp host-bounce payload {n} exceeds geometry-sized capacity {} \
(n_embd={}, max prime tokens={})",
bounce.capacity,
bounce.n_embd,
crate::cache::PRIME_CHUNK_MAX_TOKENS,
)
.into());
}
let mut host = bounce.slot(b, slot_idx)?.lock().unwrap();
s_tx.memcpy_dtoh(x, host.prefix_mut(n))?;
}
BoundaryTransport::Peer => {
use cudarc::driver::{DevicePtr, DevicePtrMut};
let (sp, _g0) = x.device_ptr(s_tx);
let (dp, _g1) = buf.device_ptr_mut(s_tx);
self.stages[b].ctx.bind_to_thread()?;
unsafe {
cudarc::driver::result::memcpy_peer_async(
self.stages[b + 1].ctx.cu_ctx(),
dp,
self.stages[b].ctx.cu_ctx(),
sp,
n * std::mem::size_of::<f32>(),
s_tx.cu_stream(),
)?;
}
}
}
sl.ev_tx.record(s_tx)?;
Ok(slot_idx)
}
pub fn rx(&self, b: usize, slot_idx: usize, n: usize)
-> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
let sl = &self.boundaries[b].slots[slot_idx];
let s_rx = &self.stages[b + 1].stream;
s_rx.wait(&sl.ev_tx)?;
let mut guard = sl.buf.lock().unwrap();
let buf = guard.as_mut().expect("pp rx before tx");
assert!(buf.len() >= n, "pp rx: slot holds {} < requested {n}", buf.len());
if boundary_transport(self.boundaries[b].cross, self.host_bounce)
== BoundaryTransport::HostBounce
{
let bounce = self.bounce_rt()?;
let host = bounce.slot(b, slot_idx)?.lock().unwrap();
let mut dst = buf.slice_mut(0..n);
s_rx.memcpy_htod(host.prefix(n), &mut dst)?;
}
let mut work = unsafe { s_rx.alloc::<f32>(n)? };
s_rx.memcpy_dtod(&buf.slice(0..n), &mut work)?;
sl.ev_rx.record(s_rx)?;
Ok(work)
}
pub fn publish_to(&self, s: usize, dst: &Arc<CudaStream>)
-> Result<(), Box<dyn std::error::Error>> {
let st = &self.stages[s];
if Arc::ptr_eq(&st.stream, dst) {
return Ok(());
}
let ev = st.ctx.new_event(None)?;
ev.record(&st.stream)?;
dst.wait(&ev)?;
Ok(())
}
pub fn fence_stages_behind(&self, src: &Arc<CudaStream>)
-> Result<(), Box<dyn std::error::Error>> {
let ev = src.context().new_event(None)?;
ev.record(src)?;
for st in &self.stages {
if Arc::ptr_eq(&st.stream, src) {
continue;
}
st.stream.wait(&ev)?;
}
Ok(())
}
pub fn record_done(&self) -> Result<CudaEvent, Box<dyn std::error::Error>> {
let last = &self.stages[self.stages.len() - 1];
let ev = last.ctx.new_event(None)?;
ev.record(&last.stream)?;
Ok(ev)
}
pub fn readback_stream(&self) -> &Arc<CudaStream> {
&self.readback
}
}
pub struct PendingLogits {
logits: CudaSlice<f32>,
ev: CudaEvent,
rb: Arc<CudaStream>,
}
impl PendingLogits {
pub fn new(logits: CudaSlice<f32>, ev: CudaEvent, rb: Arc<CudaStream>) -> Self {
PendingLogits { logits, ev, rb }
}
pub fn wait(self) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
self.rb.wait(&self.ev)?;
let host = self.rb.clone_dtoh(&self.logits)?;
self.rb.synchronize()?;
Ok(host)
}
}
pub fn new_cache(e: &Engine, cfg: &memra_gguf::config::ModelConfig, max_ctx: usize)
-> Result<crate::cache::Cache, Box<dyn std::error::Error>> {
let n_trunk = (cfg.n_layer - cfg.nextn_predict_layers) as usize;
if let Some(fence) = pp_cuts(n_trunk) {
if pp2_devices_env().is_some() && !pp2_streams_off() {
let rt = PpNRt::get(e)?;
rt.init_host_bounce(e, cfg.n_embd as usize)?;
let n_st = fence.len() - 1;
assert_eq!(
rt.n_stages(), n_st,
"PpNRt stage count {} != fence stages {n_st}", rt.n_stages()
);
rt.fence_stages_behind(&e.stream())?;
let devs: Vec<&dyn memra_kv::KvDev> =
(0..n_st).map(|s| rt.engine(s, e) as &dyn memra_kv::KvDev).collect();
let cache = crate::cache::Cache::new_ppn(&devs, &fence, cfg, max_ctx)?;
sync_stages_after_load(e, n_trunk)?;
return Ok(cache);
}
if !pp2_streams_off() {
let cache = crate::cache::Cache::new(e, cfg, max_ctx)?;
sync_stages_after_load(e, n_trunk)?;
return Ok(cache);
}
}
crate::cache::Cache::new(e, cfg, max_ctx)
}
pub fn sync_stages_after_load(e: &Engine, n_trunk: usize)
-> Result<(), Box<dyn std::error::Error>> {
if pp2_streams_off() || pp_cuts(n_trunk).is_none() {
return Ok(());
}
let rt = PpNRt::get(e)?;
for s in 0..rt.n_stages() {
rt.stages[s].ctx.bind_to_thread()?;
unsafe {
cudarc::driver::sys::cuCtxSynchronize().result()?;
}
}
e.ctx().bind_to_thread()?;
unsafe {
cudarc::driver::sys::cuCtxSynchronize().result()?;
}
Ok(())
}
pub fn layer_engine<'a>(e: &'a Engine, n_trunk: usize, il: usize)
-> Result<&'a Engine, Box<dyn std::error::Error>> {
if pp_shard_off() || pp2_devices_env().is_none() || pp2_streams_off() {
return Ok(e);
}
let Some(fence) = pp_cuts(n_trunk) else { return Ok(e) };
let rt = PpNRt::get(e)?;
let s = stage_of(&fence, il.min(n_trunk - 1));
Ok(rt.engine(s, e))
}
pub fn restore_cache_checkpoint(
e: &Engine,
cfg: &memra_gguf::config::ModelConfig,
source: Option<&crate::cache::Cache>,
target: &mut crate::cache::Cache,
snap: &crate::cache::CacheSnapshot,
) -> Result<(), Box<dyn std::error::Error>> {
let n = target.kv.len();
if target.recur.len() != n
|| snap.kv_len.len() != n
|| snap.conv.len() != n
|| snap.ssm.len() != n
|| source.is_some_and(|s| s.kv.len() != n || s.recur.len() != n)
{
return Err("checkpoint cache layer-count mismatch".into());
}
if snap.pos > target.max_ctx {
return Err(format!(
"checkpoint pos {} exceeds target capacity {}",
snap.pos, target.max_ctx,
)
.into());
}
let n_trunk = (cfg.n_layer - cfg.nextn_predict_layers) as usize;
for il in 0..n {
let owner = layer_engine(e, n_trunk, il)?;
let src_kv = source.map(|s| &s.kv[il]);
match (src_kv, target.kv[il].as_mut(), snap.kv_len[il]) {
(Some(Some(src)), Some(dst), Some(len)) => {
if len > src.len || len > target.max_ctx {
return Err(format!(
"checkpoint layer {il} len {len} exceeds source {} or target {}",
src.len, target.max_ctx,
)
.into());
}
if src.kv_dim_k != dst.kv_dim_k
|| src.kv_dim_v != dst.kv_dim_v
|| src.k_tok_bytes != dst.k_tok_bytes
|| src.v_tok_bytes != dst.v_tok_bytes
{
return Err(format!("checkpoint KV layout mismatch at layer {il}").into());
}
let kb = len * src.k_tok_bytes;
let vb = len * src.v_tok_bytes;
if kb > 0 {
owner.copy_u8_into(&mut dst.k, 0, &src.k, kb)?;
}
if vb > 0 {
owner.copy_u8_into(&mut dst.v, 0, &src.v, vb)?;
}
dst.len = len;
owner.set_i32_one(&mut dst.len_d, len as i32)?;
}
(None, Some(dst), Some(len)) => {
if len > dst.len || len > target.max_ctx {
return Err(format!(
"checkpoint layer {il} len {len} exceeds live {} or target {}",
dst.len, target.max_ctx,
)
.into());
}
dst.len = len;
owner.set_i32_one(&mut dst.len_d, len as i32)?;
}
(Some(None), None, None) | (None, None, None) => {}
_ => return Err(format!("checkpoint KV kind mismatch at layer {il}").into()),
}
match (
target.recur[il].as_mut(),
&snap.conv[il],
&snap.ssm[il],
) {
(Some(dst), Some(conv), Some(ssm)) => {
if conv.len() != dst.conv_state.len() || ssm.len() != dst.ssm_state.len() {
return Err(
format!("checkpoint recurrent layout mismatch at layer {il}").into(),
);
}
owner.copy_into(&mut dst.conv_state, 0, conv, conv.len())?;
owner.copy_into(&mut dst.ssm_state, 0, ssm, ssm.len())?;
}
(None, None, None) => {}
_ => {
return Err(
format!("checkpoint recurrent kind mismatch at layer {il}").into(),
);
}
}
}
target.pos = snap.pos;
sync_stages_after_load(e, n_trunk)?;
if source.is_some() {
e.stream().synchronize()?;
}
Ok(())
}
#[cfg(test)]
mod host_bounce_tests {
use super::{boundary_transport, host_bounce_capacity, BoundaryTransport};
#[test]
fn transport_selection_keeps_peer_default_and_bounces_only_cross_device() {
assert_eq!(boundary_transport(false, false), BoundaryTransport::Local);
assert_eq!(boundary_transport(false, true), BoundaryTransport::Local);
assert_eq!(boundary_transport(true, false), BoundaryTransport::Peer);
assert_eq!(
boundary_transport(true, true),
BoundaryTransport::HostBounce
);
}
#[test]
fn step37_geometry_sizes_each_slot_from_the_prime_cap() {
let (elems, bytes) = host_bounce_capacity(4096).expect("valid Step-3.7 geometry");
assert_eq!(elems, 4096 * crate::cache::PRIME_CHUNK_MAX_TOKENS);
assert_eq!(bytes, 64 * 1024 * 1024);
}
#[test]
fn host_bounce_capacity_rejects_invalid_or_overflowing_geometry() {
assert!(host_bounce_capacity(0).is_err());
assert!(host_bounce_capacity(usize::MAX).is_err());
}
}