Module pp 

M2 pipeline-parallel N-stage runtime (generalizes the M1 2-stage seam).

Door: MEMRA_PP_STAGES=N (default OFF — unset/0/1 = no behavior change anywhere). Stage map: N stages over the trunk layers with N-1 cuts. MEMRA_PP_SPLITS=c1,..,cN-1 sets the cuts explicitly (strictly increasing, in (0, n_layers)); MEMRA_PP_SPLIT=<i> is the N=2 back-compat spelling; default = even split (cut s = s*n_layers/N). Placement: MEMRA_PP_DEVICES=d0,..,dN-1 maps stage s to device ds (default: all on the primary engine’s device).

M1 history (increments 1-2, merged + hardened on the 8x box 2026-08-02): seam + gate single-device; then real transport — per-stage streams/events, device placement, peer-copy boundary (M0: cudaMemcpyPeerAsync beats NCCL 2.8x at PP activation sizes), per-context PDL module caches, default-mempool peer grants. All five r3 gates PASS bit-identical (receipts ~/receipts/m1-pp2/ on darklanes-bench).

M2 increment 1 (this file): N-STAGE GENERALIZATION — Pp2Rt becomes PpNRt:

• stages: Vec of per-stage execution homes (device, context, stream, remote Engine);
• boundaries: N-1 boundary runtimes, each with TWO persistent double-buffered slots (ev_tx/ev_rx per slot) and its own overlap step counter; transport is selected PER BOUNDARY (dtod same-device / cudaMemcpyPeerAsync cross-device);
• peer + default-mempool access is granted between EVERY distinct pair of devices in use (stage devices + the primary): stage kernels may dereference the primary’s weights (bring-up placement) and stage-0’s pos_d, and each boundary peer-copies.

M2 increment 2 (weight sharding): the loader uploads each stage’s layer range THROUGH that stage’s engine (layer_engine), so weights land on the device that runs them — the bring-up peer-read placement dies. output_norm + lm head load through the LAST stage’s engine; the embed table stays host-side with stage 0. Split-plane/f16 decode mirrors are built per layer through the owning stage’s engine too (the rp4 mirrors ARE the decode weights on the q8 path — leaving them on dev0 would fake the kill). Rollback seam: MEMRA_PP_SHARD=0 = M1 bring-up placement (all weights on primary, remote stages peer-read).

M2 increment 3 (deferred readback — the pipelining seed): PendingLogits — the eager decode arm can END a step without the logits D2H (decode_step_h_ppn_deferred): the logits stay device-resident with a completion event; wait() drains them through a DEDICATED readback stream (waits the event, copies, syncs) so tokens t+1.. keep enqueuing on the stage streams while token t drains. Per-token math is fully event-ordered (same slots, same ev_tx/ev_rx chain) — scheduling changes, math does not; the pipelined replay arm of ppn-gate proves bit-identity per step.

Ownership across a boundary (unchanged from M1):

• hidden state [n_embd] f32 is the ONLY tensor that crosses;
• KV/linear-attn cache entries are per-layer: stage s exclusively owns cache state for its layer range (and, under MEMRA_PP_DEVICES, allocates it on its device);
• position/rope state is the scalar cache.pos snapshot taken once per step, uploaded on stage-0’s stream BEFORE the first TX event — every later stage’s wait chain transitively orders it (stage s waits boundary s-1’s ev_tx, which was recorded after stage s-1’s work, which waited boundary s-2’s ev_tx, … back to stage 0);
• the embed table lives with stage 0, output_norm + lm head with the last stage.

THE MULTI-STREAM LAW (why this is safe with cudarc event tracking disabled): all cross-stage bytes flow through the persistent boundary slots, ordered by ev_tx/ev_rx; per-stage scratch is allocated AND freed on that stage’s stream (stream-ordered); the async mem pool runs with opportunistic reuse OFF + internal dependencies ON (memra-runtime), so a block freed on stream A and reused on stream B carries a driver-inserted dependency. Weights are load-time state no stage stream can precede, and the step’s terminal logits readback (sync D2H, or PendingLogits’ event-ordered readback stream) drains the last stage, whose TX-wait chain transitively drains all.

Scope: plain eager decode only (generic arm N-stage; gemma4 arm 2-stage). NOT wired: batch/dc/graph/spec loops and the gemma4-E4B eager arm (warn_unwired_once fires).

Structs

PendingLogits

M2 increment 3: a step’s logits, still device-resident on the LAST stage. wait() orders the readback stream behind the step’s completion event, copies, and syncs — tokens enqueued after this step keep running on the stage streams while the caller drains token t. Dropping without waiting is safe (buffers free stream-ordered).

PpNRt

StageRt

One pipeline stage’s execution home: device, context, launch stream, and (for a stage remote to the primary engine’s device) a dedicated Engine in that device’s primary context (CUmodules are per-context).

Functions

layer_engine

M2 increment 2 (weight sharding): the engine that should UPLOAD layer il’s weights (and build its decode mirrors) — the owning stage’s engine when the door is open with device placement and sharding not rolled back; else the primary. il >= n_trunk (MTP/NextN blocks) maps to the last stage. The head (output_norm + lm head) belongs to the last trunk layer’s stage — call with il = n_trunk - 1.

new_cache

Stage-owned cache allocation door: when the ppN door is open AND MEMRA_PP_DEVICES is set (placement plumbing), each layer’s cache is allocated by its OWNING stage’s engine — on one device this is byte-for-byte today’s allocation (gated); cross-device it puts each stage’s KV on that stage’s HBM. Door shut or devices unset: plain Cache::new (zero behavior change). Trailing MTP/NextN layers (beyond the trunk) map to the LAST stage.

pp2_overlap

MEMRA_PP_OVERLAP=1: alternate the double-buffered boundary slots per step (the pipelining seed). Default OFF — scheduling structure only, never math. Read per step so gates can A/B in-process.

pp2_split

N=2 back-compat view of the door (the gemma4 arm and pp2-gate are 2-stage): Some(cut) iff the door is open with EXACTLY two stages.

pp2_streams_off

MEMRA_PP_STREAMS=0: rollback to the increment-1 same-stream seam (boundary = two plain dtod copies on the ambient compute stream, no per-stage streams/events/devices).

pp_cuts

Returns the stage fence iff the ppN door is open: MEMRA_PP_STAGES=N (N >= 2) with a valid cut list. The fence has N+1 entries: [0, c1, .., cN-1, n_layers]; stage s runs layers [fence[s], fence[s+1]). Reads the environment on every call (gates toggle the door in-process); the cost is a few getenv per decode step, eager-loop noise.

pp_multi_stream_same_device

True iff the ppN door would put TWO OR MORE stage streams on ONE device (devices unset = all stages on the primary; or an explicit placement with a repeated device). The deferred-readback (pipelined) arm is REFUSED in this regime: the 2026-08-02 x20 soak record — singledev pipelined 13/20 PASS default, 7 failures each diverging at a different step (timing-race signature); MEMRA_PDL=0 went 20/20 on one soak but a second same-config soak on the auto-gated build failed 2/20 (n2) and battery-4 failed n4 — so PDL narrows the window without closing it, and the true root cause (same Engine kernels concurrent on two streams of one device) is NOT fixed by any flag yet. Cross-device pipelined (one stage stream per device) is 23/23 clean post-fix. Refuse loudly rather than return silently-wrong logits. Env-only read (callable pre-runtime).

pp_shard_off

M2 increment 2 rollback seam: MEMRA_PP_SHARD=0 = the M1 bring-up placement (all weights upload through the primary engine; remote stages peer-read). Default ON — under MEMRA_PP_DEVICES each stage’s layer range uploads through its own engine.

stage_of

The stage that owns layer il under fence (see pp_cuts).

sync_stages_after_load

M2 increment 2 LOAD BARRIER: weight uploads and decode-mirror builds enqueue on the loading engines’ WORKER streams; the first consumer launches on a DIFFERENT stream with no load->decode event — the door-off reference walk on the primary worker stream (sharded load: remote builds still in flight), or a fresh per-stage stream. The 2026-08-02 gate finds (n2-dev01 step-0 168k-logit graze; split5 ref=0.0 head — a half-built rp4 mirror — poisoning step-0 KV and every later step): one context-wide synchronize per stage at load end kills the class. No-op when the door is shut at load (single-stream load+decode is ordered by the stream itself).

warn_unwired_once

One-time notice when the door is set but the executing path has no pp arm (M2 wires the generic eager decode at any N and the gemma4 eager arm at N=2).

Type Aliases

Pp2Rt

M1 name kept alive for external callers (pp-transport-smoke, receipts, docs).