Expand description
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 by default; opt-inMEMRA_PP_HOST_BOUNCE=1uses pinned D2H + H2D instead);- the default peer transport grants peer + default-mempool access between EVERY distinct pair of devices in use. Host bounce skips serving-time grants; its boot diagnostics transiently enable peer + pool access, then revoke the pool grants and disable peer access before proceeding. Sharded weights plus stage-local auxiliary buffers ensure that no peer read can bypass the bounced boundary.
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.possnapshot taken once per step; every stage uploads its own position buffer on its own stream (no cross-device position pointer); - 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§
- Peer
Probe Metrics - Pending
Logits - 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).
Enums§
- Dual
PpMode - MEMRA_DUAL_PP three-state mode for the dual-active PP-2 batched decode path. Default ON (owner flip 2026-08-11) after the box1 PRO-pair re-gate: correctness bit-identity B=1..5, servestress no-thrash, 10-boot soak 929/929 golden matches with 0 slot collisions across 9123 pairs (research/dualpp2-20260811/RESULTS-regate.md), plus the dualpp1 c>=8 interleaved perf floor (+20.753% minimum, research/dualpp1-20260811/RESULTS.md).
- Peer
Probe Startup Policy
Constants§
- DUAL_
PP_ HOST_ BOUNCE_ REFUSAL - DUAL_
PP_ SINGLE_ SLOT_ REFUSAL - Binding-amendment refusal text. The negative gate quotes this exact line and requires the decode call to return before producing a token or advancing a cache.
- DUAL_
PP_ STAGE_ NAMES - PEER_
PROBE_ REQUIRED_ REFUSAL - PEER_
RUNTIME_ PROBE_ CYCLE_ COPIES - One complete runtime width rotation. The maximum-chunk rung runs once per cycle.
- PEER_
RUNTIME_ PROBE_ INTERVAL_ COPIES - Native cross-device boundary copies between low-frequency runtime integrity probes. Fixed rather than operator-tunable: this is a safety gate, not a performance experiment.
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-gaterequires 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 byrefuse_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. - dual_
pp_ eligibility - Fail-closed eligibility check kept pure so the negative manifest cell cannot accidentally
initialize CUDA state. Slot preparation itself remains
PpNRt::prepare_overlap_slots. - dual_
pp_ mode - dual_
pp_ mode_ resolve - Pure resolution for MEMRA_DUAL_PP, split from the env read so the flip regression tests cannot race parallel test threads on process env.
- dual_
pp_ on - True when the dual-active door is open (Forced or Auto). Read per step so the model-level gate can replay serial and waved arms against one loaded checkpoint.
- dual_
pp_ overlaps - dual_
pp_ route - Engine-entry routing for the dual-active path, kept pure for the flip regression
tests.
Forcedroutes every B>=2 PP-2 call intodecode_step_batch_dualeven when the placement cannot run it, so the binding refusals stay reachable and loud.Autoroutes only the exact re-gated regime and leaves everything else on the serial PP-N walker.dual_pp_eligibilityremains behind this as defense in depth. - dual_
pp_ slot_ snapshot (completed wave pairs, [slot 0 uses, slot 1 uses], rejected same-slot pairs).- dual_
pp_ timing_ dropped - dual_
pp_ timing_ on - CUDA-event timing is a diagnostic-only process door. The scored N=5 block runs without
it; the companion box1 diagnostic process enables it and exports cumulative per-wave
stage spans through
/metrics. - dual_
pp_ timing_ snapshot (total_nanoseconds, samples)for wave-A stage0/stage1 then wave-B stage0/stage1.- dual_
pp_ wave_ mid - Pure schedule policy shared by the runtime and kernel-check manifest cells. A single row has no second wave and must stay on the serial PP-N walker.
- init_
model_ transport - Bring up the PP transport while model geometry is known but before model weights upload. Door-shut and placement-free loads remain untouched.
- 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 withil = n_trunk - 1. - new_
cache - Stage-owned cache allocation door: when the ppN door is open AND
MEMRA_PP_DEVICESis 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: plainCache::new(zero behavior change). Trailing MTP/NextN layers (beyond the trunk) map to the LAST stage. - peer_
probe_ metrics - peer_
probe_ startup_ policy - Pure startup policy so unit tests and kernel-check pin the entire refusal matrix without mutating process-global environment variables.
- pp2_
overlap - MEMRA_PP_OVERLAP: alternate the double-buffered boundary slots per step (the pipelining seed). Scheduling structure only, never math. Read per step so gates can A/B in-process.
- pp2_
overlap_ resolve - Pure resolution for MEMRA_PP_OVERLAP, split from the env read for the flip regression tests.
- pp2_
split - N=2 back-compat view of the door (the gemma4 arm and
pp2-gateare 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_
host_ bounce_ active - True when host bounce is the live transport for a sharded cross-device placement. Callers use this to close paths that still peer-read non-boundary state.
- pp_
host_ bounce_ on - Broken-peer escape hatch: stage-boundary activations travel through page-locked host
memory instead of
cudaMemcpyPeerAsync. Default OFF; captured whenPpNRtis built. - 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
Erriffpp_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.pathnames the refusing function so the operator knows which loop they hit;altnames the working alternative for that loop. - restore_
cache_ checkpoint - Restore a cache checkpoint through each layer’s owning engine.
- service_
runtime_ peer_ probe - Service a due runtime peer probe without constructing a PP runtime on door-shut placements. Must be called by the CUDA owner thread at a scheduling boundary.
- 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 byrefuse_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
ilunderfence(seepp_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).