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.

CORRECTION (pp2-hardening 2026-08-06): this header used to add “(warn_unwired_once fires)” to that list, which was wrong. warn_unwired_once has exactly two call sites and BOTH are gemma4-specific (decode.rs, hybrid_forward.rs) — the batch/dc/graph/spec loops never warned. Worse, the batched loop did not merely run unsplit: it walked the whole trunk on the primary stream and, under a sharded cross-device placement, peer-read every remote stage’s weights each step — 28x slower at B=1 with all three decode-batch-gate gates PASSING (peer reads are byte-exact, so only perf broke). decode_step_batch now FAILS CLOSED in that regime via pp_sharded_cross_device() (MEMRA_PP_ALLOW_UNSPLIT_BATCH=1 = measurement override). “Unwired” for dc/graph/spec still means “runs unsplit, silently” — audit each before trusting it on a pair.

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).

Statics

PRIME_PIPE_OVERLAPS

PIPELINE-LIVENESS COUNTER: bumped only when a second PP-2 prime stage enters its layer walker while the other stage’s walker is still active. Step’s per-layer router readback synchronizes the host, so enqueue order alone is not liveness: a single host thread can call stage 0(N+1) before the stage-1 epilogue and still serialize all trunk computation.

PRIME_SPLIT_CHUNKS

SPLIT-LIVENESS COUNTER for the prime stage split: bumped ONCE per prime chunk that actually executed the per-stage walk. The prime-split-gate requires this to ADVANCE during its split arm — bit-identity of two identical UNSPLIT walks is vacuous, so a gate that only compared bits would go green while the walker doesn’t exist. With the counter, the gate is RED until the walker lands (the tickinv35 pattern: the gate exists and fails before the mechanism does). Relaxed ordering: single-threaded host issue, count-only.

STEP35_PRIME_BATCHES

Step35 cross-request prime liveness counters (lane/cx-prime-batch, 2026-08-08). The exactness gate requires BOTH to advance: a successful step35 batch alone is not sufficient under PP-N if it walked the whole sharded trunk on one stream.

STEP35_PRIME_BATCH_SPLITS

Functions

batch_pp_on

MEMRA_BATCH_PP=0: rollback/A-B seam for the BATCHED stage split (pp2-batch 2026-08-06). Default ON — with the ppN door open the batched decode step takes its own stage split (decode_step_batch_ppn) exactly as the eager step does. Setting 0 sends the batched path back through the unsplit body, which under a sharded cross-device placement is then caught by refuse_unsplit_if_remote (the 28x peer-read regime) rather than run silently. Exists so the bit-identity gate can A/B split vs unsplit IN ONE PROCESS against the same loaded weights — read per step, never memoized, for that reason.

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.

pp_sharded_cross_device

True iff the ppN door is open AND the placement spans 2+ DISTINCT devices AND the per-stage sharded loader is on — i.e. some layers’ weights live on a device other than the primary. Any path that walks the WHOLE trunk on one stream in this regime reads those weights over PCIe every step. Env-only read (callable pre-runtime).

prime_pipe_on

MEMRA_PRIME_PIPE=0: rollback/A-B seam for the PP-2 PRIME CHUNK PIPELINE (lane/cx-pipeline-prime 2026-08-08). Default ON when the prime stage split is live; setting 0 keeps the serial per-chunk stage walk. Read per prime call so the exactness gate can replay both schedules against one loaded model.

prime_pipe_overlaps

Read the prime-pipeline overlap counter (gate-side).

prime_pp_on

MEMRA_PRIME_PP=0: rollback/A-B seam for the PRIME (chunked prefill) stage split (lane/pp-leverb 2026-08-08). Default ON — with the ppN door open the chunked prime takes its own per-stage range walk exactly as the eager/batched/verify steps do. Setting 0 sends prime back through the unsplit whole-trunk walk. NOTE: unlike batch/dc/graph/spec, prime keeps NO refuse_unsplit_if_remote — its unsplit walk over a sharded placement is the measured 22% amortized peer-read tax (research/pp-prefill-20260807 anatomy: m=4096 amortizes the weight reads), not the decode 28x cliff, and the unsplit walk IS the split-vs-unsplit gate’s reference arm (prime-split-gate), so it must stay callable. Read per call, never memoized (the gate A/Bs both arms in one process).

prime_split_chunks

Read the split-liveness counter (gate-side).

refuse_unsplit_if_remote

The shared fail-closed guard for EVERY decode path that has no pp stage split. Returns Err iff pp_sharded_cross_device() — i.e. the caller would walk the whole trunk on one stream while some layers’ weights live on another device, peer-reading them every step. path names the refusing function so the operator knows which loop they hit; alt names the working alternative for that loop.

spec_pp_on

MEMRA_SPEC_PP=0: rollback/A-B seam for the SPEC VERIFY stage split (pp2-spec 2026-08-06). Default ON — with the ppN door open the verify forward (decode_step_t_core_ppn) takes its own stage split exactly as the eager and batched steps do. Setting 0 sends verify back through the unsplit trunk walk, which under a sharded cross-device placement is then caught by refuse_unsplit_if_remote (the 28x peer-read regime) rather than running silently. Exists so the bit-identity gate can A/B split vs unsplit IN ONE PROCESS against the same loaded weights — read per verify call, never memoized, for that reason.

stage_of

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

step35_prime_batch_splits

step35_prime_batches

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).