code/docs/P3_ENGINE.md

P3: native paged-KV continuous-batching sampler

Status: implemented, CPU-verified, opt-in. GPU validation is scripted but NOT run (the P2 GPU chain owns the GPUs). Files: align/rollout/paged.py, align/rollout/engine.py, lazy exports in align/rollout/__init__.py, tests/test_paged_engine.py (see the Review section: 27 test_mutant_* functions, each with an unmutated control that passes), tests/test_paged_engine_bench.py (plus the measurement CLI). Runtime deps unchanged (torch + stdlib; tests/test_dependency_firewall.py green).

Why

align.inference.engine.Engine (what NativeRolloutBackend drives) gives each slot a contiguous max_seq_len strip of KV for the engine's life. Prover rollouts have a heavy-tailed length distribution (most samples end at a few thousand tokens, some run to the 16k/40k cap), so concurrency is budget / cap whatever is actually in flight. Paging hands out block_size-token blocks on demand; waste is bounded by one partial block per sequence, and a sequence that outgrows its share can be preempted and recomputed instead of the whole run failing.

Design

paged.py

• BlockAllocator: LIFO free list + refcounts. free on refcount 0 raises (a double free would put one block on the list twice and give it to two sequences). check_consistent() proves free + referenced == total, no duplicates.
• PagedKVCache: per-layer [num_blocks, block_size, Hkv, D] K and V; block tables; copy-free append (write scatters into (block, offset) in place). ensure_capacity is all-or-nothing. slots() refuses a write into a block with refcount > 1 (a shared block is read-only for everyone, owner included). gather() returns contiguous [1, Hkv, L, D] (same bytes, same layout whatever else is in the batch). block_table_tensor() is the int32 [B, max_blocks] argument a kernel takes.
• TorchPagedAttention(mode): pure torch, the attention seam is the PagedAttention protocol (write new K/V, return [T, H, D]). 'row' = one SDPA per query row over exactly its own keys; 'segment' = one SDPA per prefill chunk with a bool mask + all decode rows padded into one batched SDPA.
• TritonPagedAttention: interface only (raises). The docstring pins the kernel contract (q, k/v cache blocks, block tables, cu_seqlens_q, ctx_lens, slot mapping) and the invariance requirement (fixed KV split per context length, no atomics).
• kv_bytes_per_token, plan_capacity: the memory accounting below.

engine.py

• Model-agnostic WITHOUT editing any model file: _ModelView drives the model's own submodules layer by layer (attn_norm, attn.{q,k,v,o}_proj, q_norm/k_norm, mlp_norm, mlp, norm, lm_head, rope) and replaces only attention. Qwen3 per-head QK-norm vs OLMoE full-width QK-norm (+ clip_qkv) is told apart by the norm weight width; a non-canonical kv-head map or an unknown module layout is refused at build. Mistral, Qwen3 and OLMoE (MoE mlp returning a tuple) all run.
• Scheduler: every sequence is tokens + num_computed; a step schedules n_new = min(remaining, chunk_size, budget) per sequence. Prefill, chunked prefill, decode and post-preemption recompute are one path; a sequence samples when its scheduled tokens reach the end of tokens. Decodes are scheduled first; prefill chunks and decode tokens share a flat ragged step.
• Admit/evict: FIFO admission if the whole current token list fits; on block exhaustion the youngest unscheduled running sequence is preempted by recompute (blocks freed, num_computed = 0, tokens and sampler state kept). No host swap exists on this box, so recompute is the only eviction.
• Per-sequence SamplingParams: max_tokens, temperature, top_p, top_k (applied before top-p), seed, stop_token_ids, text stop (needs a tokenizer), plus engine eos_token_id and max_seq_len. One torch.Generator per sequence, advanced only by its own draws. Unseeded requests hash the request id (not admission order). The returned logprob is log_softmax(logits/T)[tok] untruncated, the same basis as native.LOGPROB_BASIS, so the rho contract is unchanged.
• Prefix sharing (opt-in): a waiting sequence adopts the leading FULL blocks of a running sequence (GRPO's G samples prefill the prompt once), defers while the leader is still computing them, and always recomputes its last prompt token.
• numerics: 'invariant' (default; every row-wise module call is issued per row, attention per query row; bit-reproducible, slow) or 'fast' (rows batched, decode attention padded; equal to invariant to rounding).

Opt-in: from align.rollout import PagedEngine, PagedEngineConfig, SamplingParams, PagedResult resolve lazily; importing align.rollout does not import engine/paged (tested in a subprocess). native.py, client.py, align/rl/* are untouched. It is not yet wired as a RolloutBackend (see open items).

Evidence (all measured, CPU, torch 2.11, 4-8 threads)

Run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_paged_engine.py tests/test_paged_engine_bench.py -m "not gpu and not oracle and not slow" -> 149 passed, 1 skipped (the dense-only reference check for OLMoE) after the review pass (69 before it). Before: the modules do not exist (ImportError); the oracles below also fail on every mutant.

1. Greedy equals independent references. Mistral, Qwen3 (QK-norm) and OLMoE tiny models in fp64, three prompts (3, 9, 17 tokens), 10 new tokens, both numerics, chunk 5: tokens identical to a naive full-recompute loop. For Mistral/Qwen3 the reference is a transformer written in the test from raw weights in fp64 (logprobs agree to < 1e-5; the models' RoPE rotates in fp32 by design, so < 1e-15 is not attainable against a true-fp64 reference). OLMoE is checked against the model's own full forward. A separate test pins the reference to the model forward (< 1e-5). Also identical (tokens, logprobs < 1e-4) to align.inference.engine.Engine and to NativeRolloutBackend at temperature 0 (fp32).
2. Batch invariance. fp64 Qwen3, invariant: a sequence alone vs among 1/3/7 others of random lengths 2-30 admitted at random steps 0-6, random chunk (3/5/64), budget (8/16/64) and max_num_seqs (2/4/16), 4 seeds x 3 compositions: tokens, logprobs and every step's full logit vector are bitwise equal (torch.equal). Same under chunk size 1/2/64, with prefix sharing on (>= 8 prompt tokens actually shared), and with a 17-block pool that forces preemption by recompute. fast vs invariant: dense models, logits within 1e-9 (fp64) and identical greedy tokens. OLMoE is NOT that tight: its router/MoE path has fp32 islands, so even an fp64 model drifts ~1.2e-7 between fast and invariant (tokens identical in 400 randomized configs, tolerance 1e-5); a near-tie in expert routing is therefore a place fast can flip.
3. Seeded sampling (temperature, top-p, top-k, combined) is bit-identical alone vs 2/5/8 companions, with different chunking and a 24-block pool; different seeds give different streams. row_distribution equals an independent fp64 reference (top-k then minimal-nucleus top-p) to 1e-6 with identical support size.
4. Allocator never leaks. Property run of 10,000 randomized events (admit with optional prefix share / grow / finish) on PagedKVCache, books checked after EVERY event against refcounts recomputed from the live block tables (>100 OOM events and

   100 shared admits occurred); plus an engine fuzz: 120 requests, 14-block pool, prefix sharing on, preemptions > 0, consistent after every step, pool fully free at the end.

5. Mutants killed (monkeypatched; the oracle must fail): gather drops last key; slot mapping ignores block offset; K heads flipped (GQA map); prefill causal mask off by one (fast); RoPE positions restart each step; allocator free leaks; free double-appends; incref removed; double-free guard removed; preemption forgets generated tokens; preemption forgets to free blocks; preemption picks the oldest victim (more recompute); one shared generator; row batching in fp32 (breaks bitwise invariance); top-p keeps extra token; top-k off by one; logprob from truncated distribution; prefix sharing of a partial block; prefix sharing without the last-token recompute. Each has an unmutated control that passes.

CPU throughput scaling (generated tokens/s, best of 3, 16-token prompts, 24 new tokens, 3.41M-param Qwen3-shaped fp32, 8 threads, shared busy host)

N seqs	fast tok/s	speedup	invariant tok/s	speedup
1	210	1.00x	174	1.00x
2	412	1.96x	207	1.19x
4	581	2.77x	237	1.36x
8	867	4.13x	272	1.56x
16	1033	4.92x	280	1.61x
32	1154	5.49x	n/a	n/a

Batching scales sub-linearly (Python-level per-sequence attention and per-row sampling dominate a tiny model); invariant buys bit-reproducibility at 4-5x lower throughput at N=16. CPU numbers say nothing about GPU absolute speed.

bf16 expectation (24 trials, 7 companions, random-init 4L/256H model, so near-flat logits)

• invariant: 0/24 greedy and 0/24 sampled sequences differ alone vs batched; max |dlogit| = 0 (bitwise).
• fast: 11/24 (45.8%) greedy sequences differ (first-divergence-limited max |dlogit| 0.047, mean-of-max 0.035, i.e. a few bf16 ulps of the logit scale); 0/24 sampled differ. This model is a pessimistic tie-heavy case for greedy (2048 near-equal logits); a trained model's peaked distributions will flip less, but that is not measured here. The 0/24 sampled figure is also an artefact of flat distributions (one uniform draw lands in the same bin), not evidence of safety. Rule: use invariant whenever the logprobs feed the rho correction or a replayable run; fast is for throughput evals.

Memory accounting (Qwen3-8B / Goedel-Prover-V2-8B, bf16 KV)

• KV bytes per token = 36 layers x 2 x 8 kv-heads x 128 x 2 B = 147,456 B (144 KiB) (kv_bytes_per_token; test pins it). Mistral-7B: 131,072 B.
• The vLLM probe's 455,888 tokens is a TP=4 number. .tmp/p2/P2-B/smoke.log: "Available KV cache memory: 15.65 GiB" per rank, KV heads sharded 4 ways -> 36,864 B/token/rank -> 15.65 GiB / 36,864 = 455,8xx (test reproduces it within 0.05%), and "Maximum concurrency for 40,960 tokens: 11.13x" (our plan gives 11.12x). The P2 baseline manifest records 456,528 tokens at util 0.9 and 876 tok/s end-to-end on 4 GPUs (9.6M tokens, 114 of 1952 length-capped).
• One A10G, native engine, no TP (device 22.06 GiB as vLLM sees it, util 0.9 -> 19.85 GiB; weights 15.26 GiB; 1.0 GiB activations + 0.4 GiB context assumed): KV budget 3.19 GiB = 23,248 tokens. At a 16k cap: 1 concurrent sequence; at a 40k cap: 0 (a single 40,960-token sequence cannot fit). At util 0.95: 31,280 tokens, still 1 / 0. At a 4k cap: 5. A dense per-slot strip allocates the same bytes, so paging does not conjure memory here; it makes the 23k tokens usable by many short sequences instead of one reserved strip, and removes the failure when one long sample appears.
• Planner for sharded layouts (tp_degree divides KV heads): TP=2 gives 164,912 tokens (10 x 16k, 4 x 40k); TP=4 gives 440,944 (26 x 16k, 10 x 40k) with 0.5 GiB activations.
• Consequence (decision input): the native engine as written is single-device. For the 8B prover it cannot replace the TP=4 vLLM pool (876 tok/s end-to-end) - it can serve Goedel-8B at ~4k caps on one GPU, or serve models <= ~4B at long caps. Matching vLLM needs TP (shard KV heads + projection columns) or an fp8 KV cache; neither is built here (open items).

GPU validation plan (run only after the P2 chain, PID 733530, has exited; GPUs 0-3 show 22 GB each now)

Preconditions (/apps/CLAUDE.md): free -g; uptime; awk '$3=="nvme2n1"{print $12}' /proc/diskstats; ps -eo stat | grep -c '^D'; nvidia-smi all clear and pgrep -f chain_after_baseline empty. Both steps below take the heavy-job lock (--wait-lock), so they serialise.

cd /apps/development/rlhf/align
export HF_HOME=/apps/development/rlhf/hf_cache HF_HUB_DISABLE_XET=1 PYTHONPATH=$PWD TMPDIR=$PWD/.tmp/g2/p3 \
       PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
M=/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a
mkdir -p .tmp/g2/p3/gpu

# 1) vLLM, ONE GPU, same memory the native engine gets, 16 problems x 4 samples, 4k cap (also writes the prompt_ids reused in step 2)
CUDA_VISIBLE_DEVICES=0 .venv-rollout/bin/python extras/vllm_rollout/lean_gen.py \
  --problems .tmp/p2/problems/minif2f.jsonl --limit-problems 16 --out .tmp/g2/p3/gpu/vllm_tp1.jsonl \
  --n 4 --cap-tokens 4096 --round 0 --tp 1 --max-model-len 8192 --gpu-mem-util 0.90 --max-num-seqs 64 --seed 0 --wait-lock --label p3-vllm-tp1

# 2) native paged engine, same GPU, same prompt ids, same T=1.0 / top_p=0.95 / 4096 cap; fills free memory minus 3 GiB with KV (see Review: the padded-decode transient is unmeasured)
CUDA_VISIBLE_DEVICES=0 .venv/bin/python tests/test_paged_engine_bench.py --gpu --model-dir $M \
  --prompts-jsonl .tmp/g2/p3/gpu/vllm_tp1.jsonl --n-prompts 16 --g 4 --max-tokens 4096 \
  --max-num-seqs 64 --budget 2048 --chunk 512 --numerics fast --wait-lock --out .tmp/g2/p3/gpu/native_fast.json

# 3) cost of bit-invariance (slow: fewer prompts)
CUDA_VISIBLE_DEVICES=0 .venv/bin/python tests/test_paged_engine_bench.py --gpu --model-dir $M \
  --prompts-jsonl .tmp/g2/p3/gpu/vllm_tp1.jsonl --n-prompts 2 --g 4 --max-tokens 1024 \
  --numerics invariant --wait-lock --out .tmp/g2/p3/gpu/native_invariant.json

Compare tokens_per_s_end_to_end, len_p50/len_max/finish_reasons and peak_cuda_mib in the native JSON with tokens_per_s_end_to_end, finish_length, gpu_mem_used_peak_mib and kv_cache.tokens in .tmp/g2/p3/gpu/vllm_tp1.jsonl.manifest.json. Sampling streams differ by construction (per-sequence generators vs vLLM's), so compare throughput and length distributions, not tokens. Greedy fidelity on GPU (the check that matters for the rho contract): run both on 8 prompts at temperature 0 and compare first-divergence positions; expect bf16 near-tie flips, report the median first-divergence index, not "identical". Acceptance to proceed with TP work: native tokens/s >= 0.5x vLLM TP=1 at this shape (CUDA-graph-free, per-sequence gather attention; if it is far below, the Triton kernel is the next step, not TP).

Open items

1. Not a RolloutBackend yet. A PagedRolloutBackend (same submit/collect, prompt_token_ids/logprob extras, G-sample expansion with prefix_sharing=True) is a small adapter on run_compare's request building; deferred because client.py/ native.py are outside this item's list.
2. Single-device only (no TP, no fp8 KV): see the consequence above.
3. Pure-torch attention is the bottleneck on GPU (a Python loop over segments, K/V gathered per sequence per layer per step; per-row float() syncs in sampling). The Triton kernel interface is written, the kernel is not; it must keep the per-row reduction independent of batch composition or invariant becomes fast.
4. Multi-turn masks / tool-turn continuation (ROADMAP P3) are not implemented: the engine supports add_request of an extended token list (prefix sharing makes re-prefilling the transcript cheap) but has no turn-boundary loss mask output.
5. Swap-out preemption and prefix caching of FINISHED sequences (hash-keyed block reuse across requests) are not built; sharing is only between concurrently running sequences.
6. invariant is ~4-5x slower than fast on CPU and has not been timed on GPU; a batch-invariant GEMM (fixed tile/split-K per row) would remove the per-row calls.
7. OlmoeAttention/MoE paths are exercised on tiny fp64 models only; expert-parallel modules (moe_ep) are not supported.
8. Greedy fidelity of fast mode in bf16 on a real checkpoint is unmeasured (plan above).

Review (adversarial pass, CPU only)

Found by reading the engine against its own claims, then probing; every item below has a failing-before / passing-after test and at least one mutant that the new oracle kills.

• Pool guard off by one (add_request): the last sampled token is never fed back, so a request ending with N tokens holds N-1 in the cache. The guard sized the pool on N and refused a request that fits exactly (prompt 3 + 6 new, block 4: 8 cached tokens = 2 blocks, 2-block pool rejected as "needs 3"). Now sized on N-1; swept over prompt 1-9 x max_tokens {1,2,3,4,5,8,9} at block 4 with a pool of exactly the needed blocks (peak blocks == needed, pool fully free after).
• Out-of-vocabulary prompt id wedged the engine: the IndexError fired inside step(), the poisoned request stayed resident, and every later step() re-raised (on CUDA it is a device-side assert). Refused at add_request.
• Greedy NaN logits were silent: argmax of a NaN row returns the NaN's index and log_softmax is NaN, so a bf16 overflow became a plausible token with a NaN logprob in the rollout (sampled rows already raised inside multinomial). step() now raises FloatingPointError before any sequence advances (retryable; -inf stays legal). The check also synchronises, so forward_seconds now includes GPU time instead of kernel-launch time only (steps with no sampled row still do not sync).
• kv_dtype crashed for any dtype other than the model's (SDPA dtype mismatch on the first step). K/V are now dequantised on read. bf16 KV under an fp32 model: logits within 6e-3, identical greedy tokens, half the bytes. float8_e4m3fn runs but is UNSCALED and visibly lossy (tiny-model logits off by 1.3, tokens differ), so "fp8 KV" in the decision above needs per-tensor scales, not just a dtype switch.
• Bench seeds were not reproducible: hash((pid, j)) is salted per interpreter; two launches of the same GPU command drew different sampled streams. Now a sha256-derived seed (tested across two PYTHONHASHSEEDs). The bench also passed only config.json's single eos id; vLLM also stops on generation_config's list (Goedel: 151645 and 151643), so the native engine could run past <|endoftext|> and skew the length comparison. stop_ids_from_generation_config now supplies the set.
• Mutation gaps closed: stop_token_ids was untested (the older stop test's or passed with the feature broken); the "truncated logprob" mutant never ran engine code (now swaps the engine's sample_row and checks logprob == log_softmax(logits/T) through the whole engine); a batched lm_head stayed bit-identical at hidden=128 on CPU so that mutant survived (now a per-site mutant on a 512-wide model for mlp, qkv, o_proj and lm_head, each asserted to be reached and to change bits).

Not fixed, measured or reasoned:

• Recompute preemption bill: with no admission watermark a new request takes the free blocks, the next step's residents need them, and the youngest is preempted again. Measured on 20-40 block pools (block 4), 1 long decoder + 6 arrivals of 24 prompt / 40 new tokens: 7-10 preemptions, 667-818 forward tokens against 445 minimal (1.5-1.8x). Correctness is unaffected (fuzzed bitwise); a block watermark on admission is the lever.
• Unmeasured GPU transient: fast pads decode K/V to the longest sequence: 64 seqs x 8 kv-heads x 4800 x 128 x 2 B x (K and V) = 1.17 GiB per layer call, and if the SDPA backend chosen for a masked GQA call is math, K/V are expanded to 32 heads (4.7 GiB). plan_capacity's 1.0 GiB activation allowance does not cover either; hence --reserve-gib 3.0 above. Read peak_cuda_mib.
• PagedEngineConfig.max_seq_len may exceed the model's max_position_embeddings without any warning (RoPE just extrapolates).

Tensor parallelism (inference, opt-in)

Status: implemented, CPU-verified on gloo at TP=2 and TP=4, opt-in. GPU validation is scripted and NOT run (the P2 chain owns the GPUs). Files (all new; no existing file was edited, so every default is byte-identical): align/model/tp.py (the sharded model), align/rollout/tp_engine.py (TPPagedEngine), tests/test_tp.py, tests/test_tp_load.py, tests/test_tp_engine.py, tests/test_tp_bench.py (capacity/latency model + the GPU runner), tests/test_tp_runner.py (the runner end to end on CPU/gloo). Not exported from align.rollout's lazy table (that file is not part of this item): from align.rollout.tp_engine import TPPagedEngine. Runtime deps unchanged (torch + stdlib, safetensors lazily; tests/test_dependency_firewall.py green).

Why (the bottleneck, in numbers)

One A10G holds the 8B prover's weights (15.26 GiB) and has ~3.2-3.6 GiB left: one 16k sequence, zero 40k sequences. TP=4 shrinks both terms of the KV budget: weights /4 (3.81 GiB per rank, computed from the config's parameter count; vLLM measured 3.83 GiB loaded) and KV bytes per token /4 (36,864 B vs 147,456 B), so the pool grows ~17x in tokens, not 4x.

Design

Megatron layout; the only collectives are on the residual stream: 2 all-reduces per layer + 1 for the embedding + 1 all-gather for the logits.

piece	sharding	comm
embed_tokens -> VocabParallelEmbedding	rows (ceil(V/tp) each, tail padded with zeros); masked lookup	all-reduce(sum); exact in any order (x + 0)
q/k/v_proj	plain nn.Linear, output = this rank's heads (q: H/tp, kv: Hkv/tp)	none
attention	the unchanged model code on the local head counts; Qwen3's per-head q/k-norm weight (width head_dim) is replicated	none
o_proj -> RowParallelLinear	input features = this rank's heads	all-reduce(sum)
gate/up_proj / down_proj	intermediate/tp columns / rows (RowParallelLinear)	all-reduce(sum) after down_proj
lm_head -> VocabParallelLMHead	rows, tail padded	all-gather of logits, cut back to V

• No second forward. The TP model IS Qwen3ForCausalLM / NmoMistralForCausalLM built with a LOCAL config (heads, kv heads, intermediate divided by tp) and five module types swapped. The dense forward, _head_norm/_swiglu seams, RoPE, RMSNorm and the GQA map are the original code, and PagedEngine drives the shard through its own submodules unchanged. Consequence: the paged KV cache is sharded by kv head for free (_ModelView reads the local n_kv_heads; Qwen3-8B at tp=4: 8 -> 2) and per-rank KV bytes are exactly 1/tp (tested).
• GQA. Blocked map q_head // group stays local and canonical iff tp divides n_kv_heads. tp > n_kv_heads would REPLICATE kv heads (vLLM's choice); that layout is refused (TPError naming every violated divisibility: kv heads, heads, intermediate) rather than guessed. Qwen3-8B supports tp in {1, 2, 4, 8}.
• QK-norm. Per head, local, replicated weight. OLMoE's full-width norm cannot be sharded this way; it is not buildable here (no rule -> TPError).
• Logits: all-gather, not a parallel softmax. Sampling needs the full row (top-k, top-p, the multinomial draw, the untruncated log-softmax the rho correction is defined on) and every rank must hold identical logits for the replicated sampler below. Cost: (tp-1)/tp x V x 2 B = 223 KiB per SAMPLED row at tp=4 bf16. vocab_parallel_logsumexp / _token_logprob / _argmax are the two-scalars-per-row alternative for greedy decoding and teacher-forced scoring (verified against the full softmax, ties to the lowest index like torch.argmax); the engine does not route through them.
• Vocab padding. A vocabulary that does not divide tp is padded per rank and cut after the gather (a padding column's logit 0.0 could otherwise win the argmax of an all-negative row); a tp larger than the vocabulary leaves a rank with zero real rows (tested at V=5, tp=4).
• Scheduler/sampler: replicated, rank 0 authoritative for INPUTS and verifier of the replication (chosen over "rank 0 broadcasts decisions"). Every rank runs the identical deterministic scheduler; generate takes the request stream from rank 0 and broadcasts it once. Why not broadcast decisions: a ragged object (segments, block tables, sampled-row list) pickled and host-synced every step, to ship what every rank can recompute from state it holds. Why not sample on rank 0 only: a token broadcast between every forward and the next, plus the full row on rank 0 anyway. Redundant sampling is free because the logits are bit-identical on every rank by construction (an all-gather hands every rank the same bytes) and the per-sequence generators are seeded identically (_stable_seed is sha256, not hash()). The one real risk of replication is SILENT DIVERGENCE (one differing token or schedule -> mismatched collectives -> an NCCL hang or garbage), closed by two 8-byte all-gathers per step: _schedule agrees a fingerprint of the step (ids, num_computed, n_new, the fed token ids) BEFORE any model collective, and has_unfinished agrees the loop predicate plus a hash of the previous step's events. Symmetric (all-gather), so a divergence raises TPDesyncError on ALL ranks. verify=False drops them. Measured cost on gloo (3-layer toy, loaded host): +0.7 ms/step at tp=2 (9.9 vs 9.2 ms, +8%), +2.9 ms/step at tp=4 (18.5 vs 15.6 ms, +18%); at 36 layers the two checks are 2 of 76 collectives per step. GPU cost unmeasured: the runner has --no-verify for the A/B.
• Weight loading, bounded host RAM. load_tp_model(snapshot, tp, dtype, device) reads per tensor ONLY this rank's slice via safe_open().get_slice(name)[a:b], makes it compact and owned (_own: safetensors hands back a view whose storage spans the WHOLE source tensor), moves it, casts on the target. Total and loud: the checkpoint's tensors must EQUAL the architecture's (extra/missing/wrong-shape/absent shard raise), names come from the single-device loaders' own tables. tp_from_model shards an in-memory model (copies, never views) for tests and weight refresh.
• Inference only: TPGroup.all_reduce raises on a tensor that requires grad. reduce_dtype=torch.float32 accumulates the all-reduce in fp32 (2x volume).

Evidence (all measured, CPU/gloo, torch 2.11; run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_tp.py tests/test_tp_load.py tests/test_tp_engine.py tests/test_tp_bench.py tests/test_tp_runner.py -m "not gpu and not oracle and not slow")

Before: align.model.tp and align.rollout.tp_engine do not exist (ImportError); a vocab-sharded model also fails PagedEngine.add_request (the engine reads the vocabulary from embed.weight.shape[0], i.e. V/tp) - one of the mutants below. Tests in the 5 files: 134 (131 default + 1 slow RSS measurement + 2 oracle HuggingFace-parity cases).

1. fp64 logits equal the unsharded model. Tiny Qwen3 (per-head QK-norm) and Mistral, heads 8 : kv 4, odd vocab 99 (padded shards), padded attention mask. Max |dlogit| over prefill: TP=2 5.6e-16 (qwen3) / 4.4e-16 (mistral), TP=4 6.7e-16 / 8.9e-16, at logit scale ~2.8; threshold 1e-12.
2. Engine = TP=1 engine. TPPagedEngine (requests given on rank 0 only) vs PagedEngine on the unsharded model, 5 variants (invariant, fast, 16-block pool with 5 preemptions, prefix sharing with 36 shared tokens, chunk 3/budget 8), greedy + seeded T/top-k/top-p + unseeded streams + a stop-token request: tokens, finish reasons, step counts, preemption counts, shared-prefix counts identical; logprob diff 0; every sampled logits row (fp64, captured at sample_row) within 8.9e-16 at TP=2 and TP=4 for both families. Decode logits vs a DENSE full-recompute forward of the unsharded model: 1.1e-15.
3. Batch invariance survives TP (invariant): a seeded sequence alone vs among 5 companions admitted at other steps, 3 chunk/pool configs (default, chunk 3 + budget 16, 40-block pool with max_num_seqs=3): tokens, logprobs and every step's logits bitwise equal, TP=2 and TP=4.
4. KV is sharded: per-rank KV bytes = unsharded / tp exactly (393,216 -> 196,608 at tp=2, 98,304 at tp=4); kv_bytes_per_token(36, 8, 128, bf16, tp_degree=4) = 36,864 B.
5. Vocab-parallel ops (embedding, lm_head gather, logsumexp, token logprob, argmax with ties) equal full-row fp64 references to 1e-12 at V = 99, 100 and 5.
6. Loading: shards loaded from single-file and 3-shard index snapshots, Qwen3 and Mistral, tp 2 and 4, are bit-identical to tp_from_model; a fp32->bf16 cast per slice equals casting the full model. Bytes returned by safetensors per rank (recorded at the safetensors layer, independent of tp.py) equal the rank's parameter bytes exactly; whole-tensor reads happen only for the replicated norms. Peak anonymous RSS materialising bf16 weights from a 490 MiB fp32 checkpoint (-m slow): unsharded loader 236 MiB, one tp=4 rank 66 MiB (file-backed mmap peak 178 / 183 MiB, reclaimable).
7. Refusals: non-divisible kv heads / heads / intermediate (all violations named), tp > kv heads, unknown parameter, extra / missing / wrong-shape tensor, index naming an absent shard, unsupported model_type, out-of-vocabulary id (IndexError, where a naive sharded embedding yields zeros).
8. HuggingFace parity on CPU (oracle, .venv-rollout, transformers 5.19): the written tiny Qwen3 snapshot, 6 prompts x 40 greedy tokens fp32, no eos: TP=2 and TP=4 tokens identical to transformers' in 6/6 prompts. This is the CPU twin of the Goedel-8B GPU parity run below.
9. bf16 drift (tiny Qwen3, TP=4, vs unsharded): max |dlogit| 0.039 on logit scale 2.88 (1.4%, one bf16 ulp at that scale), identical with and without reduce_dtype=float32 - on this toy the all-reduce rounding is not the dominant term, which says nothing about 36 layers of a real model (GPU run).
10. Defaults untouched: no existing file edited; test_paged_engine.py, test_paged_engine_bench.py, test_qwen3_parity.py, test_model_parity.py, test_public_api.py, test_dependency_firewall.py: 224 passed, 1 skipped (the existing OLMoE dense-only skip). Full CPU suite once at the end (-m "not gpu and not oracle and not slow", shared host, other items' edits in the tree): 2711 passed, 14 skipped, 1 failed - test_moe_ep.py::test_ep_layer_forward_backward_equal_single_process_bitwise[4-4] died with EADDRINUSE (a free-port race with a concurrent job on this shared host) and passes alone (3 passed).

Mutants (monkeypatched INSIDE the spawned ranks; each oracle must fail, each has an unmutated control that passes):

feature	mutant	how it dies
layers	no row-parallel all-reduce	logits wrong by 3.9
layers	k_proj slice taken from the mirrored rank	wrong by 2.6 (GQA reads the wrong heads)
layers	up_proj slice from the mirrored rank	wrong by 3.4
vocab	embedding: every id "mine" (clamped lookup summed tp times)	wrong by 3.7
vocab	logits gathered in reversed rank order	wrong by 4.5
vocab	padding columns not cut	gathered width != V (RuntimeError)
vocab	logsumexp with the LOCAL max	lse off by 3.5
vocab	embedding range check removed	OOV id embeds silently (caught)
loader	get_tensor + narrow instead of get_slice	bytes read > 2x the shard's
loader	no copy: view of the full tensor (source AND _own)	storage oracle fails (either guard alone is still safe)
loader	_own removed on the safetensors path	CPU params pin the file mapping (caught)
loader	tensor-set check removed	extra tensor loads silently
engine	constant schedule fingerprint	schedule mismatch is not a TPDesyncError (gloo abort / hang-until-timeout)
engine	consensus all-gather removed	a wrong token on rank 1 goes through: both ranks return normally
engine	_stable_seed = pid	TPDesyncError (unseeded streams differ per rank)
engine	full-vocab shim removed	add_request refuses valid ids >= V/tp

Compute: Goedel-Prover-V2-8B on 4 x A10G

Run .venv/bin/python tests/test_tp_bench.py --plan. Computed from measured inputs (config parameter count; the vLLM probe's memory split in .tmp/p2/P2-B/smoke.log; NCCL bandwidth in docs/MEASUREMENTS.md) except the rows marked ASSUMED.

• Per-rank weights, bf16: tp=1 15.26 GiB, tp=2 7.63, tp=4 3.81 GiB (2.048 B params; norms replicated, +0.3 M).
• The planner reproduces vLLM's TP=4 pool: weights 3.83 + activation 0.36 + non-torch 0.02 (CUDA graphs are charged after the pool) at util 0.9 of 22.06 GiB -> 456,112 tokens vs vLLM's 455,888 (+0.05%), 11.14x at 40,960 (vLLM 11.13x).
• Native engine, bf16 KV, block 16; activation 0.5 GiB (padded-decode transient/tp + logits) and 0.5 GiB NCCL/context are ASSUMED, no CUDA graphs:

tp	util	weights GiB	KV GiB	KV tokens	vs vLLM 455,888	seqs @16k	seqs @40k
1	0.90	15.26	3.60	26,192	0.06x	1	0
2	0.90	7.63	11.23	163,472	0.36x	9	3
4	0.90	3.81	15.04	438,048	0.96x	26	10
4	0.95	3.81	16.14	470,176	1.03x	28	11

(vLLM TP=4: 27 x 16k, 11 x 40k. The tp=1 row uses the same assumptions as the others; the earlier single-device section used 1.0 + 0.4 GiB, hence 23,248 tokens there.)

• Link volume per rank, tp=4, bf16: 73 all-reduces per step (2 x 36 layers + embedding), each rows x 4096 x 2 B and each moving 2(w-1)/w of that; 864 KiB of layer all-reduce + 12 KiB embedding per row; 223 KiB of logits gather per SAMPLED row.
• Decode latency model (lower bound: serial comm, no kernel-launch or Python overhead). Link: measured 4.0 GB/s all-gather algbw -> 3.0 GB/s bus bandwidth (docs/MEASUREMENTS.md: NCCL over host SHM, 4 ranks, "it is the shared-memory transport's bandwidth, not latency"); HBM 600 GB/s (datasheet, ASSUMED); small-message collective latency: NO measurement exists on this box - rows show 20 / 50 / 100 us (ASSUMED). --mode allreduce below measures it.

batch	ctx	latency	weights ms	KV ms	all-reduce ms	gather ms	step ms	tok/s (decode-only)
1	4096	20 / 50 / 100 us	6.83	0.25	1.76 / 3.95 / 7.60	0.10-0.18	8.9 / 11.2 / 14.9	112 / 90 / 67
16	4096	20 / 50 / 100 us	6.83	4.03	6.2 / 8.4 / 12.1	1.3	18.3 / 20.6 / 24.3	873 / 778 / 660
64	4096	20 / 50 / 100 us	6.83	16.1	20.6 / 22.8 / 26.4	4.9	48.4 / 50.6 / 54.3	1322 / 1264 / 1178
64	8192	20 / 50 / 100 us	6.83	32.2	20.6 / 22.8 / 26.4	4.9	64.5 / 66.7 / 70.4	992 / 959 / 909

Reading: batch 1 is latency-bound (73 collectives); batch 64 is bandwidth-bound on the 3 GB/s link (all-reduce 21-26 ms of a 48-54 ms step, ~40%) and then KV reads. The decode-only ceiling of ~1,200-1,300 tok/s at batch 64 is NOT comparable to vLLM's measured 876 tok/s end-to-end (that includes prefill, the length tail and Lean-length samples); the pure-torch engine will land well below its own ceiling (Python attention loop, no CUDA graphs).

• numerics='invariant' under TP is a latency problem: it issues the o_proj/down_proj/mlp calls per ROW, so a 64-row step makes 64 x 73 = 4,672 small all-reduces (>= 93 ms at even 20 us each). Use fast for throughput and invariant only where bit-reproducible logprobs are needed (docs above).

GPU validation plan (run only after the P2 chain, PID 733530, has exited; every step takes the heavy-job lock, so they serialise)

Preconditions (/apps/CLAUDE.md): free -g; uptime; awk '$3=="nvme2n1"{print $12}' /proc/diskstats; ps -eo stat | grep -c '^D'; nvidia-smi clear, pgrep -f chain_after_baseline empty. Each job below is rank-0-locked (--wait-lock queues behind another holder).

cd /apps/development/rlhf/align
export HF_HOME=/apps/development/rlhf/hf_cache HF_HUB_DISABLE_XET=1 PYTHONPATH=$PWD TMPDIR=$PWD/.tmp/g3/tp \
       PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
M=/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a
G=.tmp/g3/tp/gpu; mkdir -p $G

# 0) the missing measurement: NCCL all-reduce latency (8 KiB = one decode row ... 64 MiB) and the logits gather, 4 ranks
torchrun --standalone --nproc_per_node=4 tests/test_tp_bench.py --mode allreduce --wait-lock --out $G/allreduce.json

# 1) vLLM TP=4 baseline = the throughput comparator AND the source of prompt_ids: 32 miniF2F problems x n=4, cap 16384
.venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems .tmp/p2/problems/minif2f.jsonl --limit-problems 32 \
  --out $G/vllm_tp4_cap16k.jsonl --n 4 --cap-tokens 16384 --round 0 --tp 4 --max-model-len 40960 --gpu-mem-util 0.90 \
  --max-num-seqs 128 --seed 0 --wait-lock --label tp-vllm-cap16k

# 2) HuggingFace greedy reference, ONE GPU (fits: 16.4 GB), 8 prompts x 64 tokens, with the top1-top2 logit margin per step
CUDA_VISIBLE_DEVICES=0 .venv-rollout/bin/python tests/test_tp_bench.py --mode hf_ref --model-dir $M \
  --prompts-jsonl $G/vllm_tp4_cap16k.jsonl --n-prompts 8 --new-tokens 64 --dtype bfloat16 --wait-lock --out $G/hf_ref.json

# 3) TP=4 parity vs HF, 64 greedy tokens (invariant numerics = the RL-safe mode), no eos stop (the reference never stops)
torchrun --standalone --nproc_per_node=4 tests/test_tp_bench.py --mode parity --model-dir $M --hf-ref $G/hf_ref.json \
  --dtype bfloat16 --numerics invariant --no-eos --tie-margin 0.1 --max-num-seqs 8 --budget 512 --chunk 256 --wait-lock --out $G/parity.json

# 4) throughput: same prompt ids / T=1.0 / top_p=0.95 / cap 16384 / generation_config stop ids; KV pool = min free over ranks - 3 GiB
torchrun --standalone --nproc_per_node=4 tests/test_tp_bench.py --mode throughput --model-dir $M \
  --prompts-jsonl $G/vllm_tp4_cap16k.jsonl --n-prompts 32 --g 4 --cap-tokens 16384 --max-seq-len 40960 \
  --max-num-seqs 64 --budget 2048 --chunk 512 --numerics fast --reserve-gib 3.0 --wait-lock --out $G/native_tp4_fast.json
#    and the consensus A/B:   ... same command with --no-verify --out $G/native_tp4_fast_noverify.json
# 5) cost of bit-invariance under TP (slow: few prompts)
torchrun --standalone --nproc_per_node=4 tests/test_tp_bench.py --mode throughput --model-dir $M \
  --prompts-jsonl $G/vllm_tp4_cap16k.jsonl --n-prompts 2 --g 4 --cap-tokens 1024 --numerics invariant --wait-lock --out $G/native_tp4_invariant.json

Acceptance. Parity: n_bad == 0 (exit code 0): every divergence from HF's greedy stream sits where HF's own top1-top2 gap is below 0.1 logits (a bf16 near-tie), and report the median first_divergence (not "identical"). Throughput: compare tokens_per_s_end_to_end, len_p50/len_max, finish_reasons, kv_tokens and peak_cuda_mib_max_over_ranks with the vLLM manifest ($G/vllm_tp4_cap16k.jsonl.manifest.json: tokens_per_s_end_to_end, finish_length, kv_cache.tokens); the bar carried over from the single-device section is native >= 0.5x vLLM at this shape - if it is far below, the next step is the Triton paged-attention kernel (and CUDA graphs), not more TP. Sampled streams differ by construction, so compare distributions, not tokens. Replace the ASSUMED latency column above with allreduce.json (all_reduce[0].us is the 8 KiB decode-row latency) before quoting any decode estimate.

Things only a GPU run can show: NCCL all_gather_into_tensor / broadcast_object_list(group_src=) / the int64 consensus all-gather on CUDA tensors (gloo passes; NCCL was not run), whether NCCL's all-reduce is bitwise identical across ranks at bf16 (the design needs only the gathered LOGITS identical, which an all-gather guarantees, but identical residuals make invariant cleaner), the real step time vs the model above, and the reduce_dtype=float32 effect over 36 layers.

Open items

1. Not a RolloutBackend; TPPagedEngine is not in align.rollout's lazy exports (that module is outside this item). Same adapter work as the single-device engine.
2. GPU/NCCL path unrun (see above). The runner's every mode is exercised on CPU/gloo (tests/test_tp_runner.py), so a typo should not survive to the GPU.
3. invariant + TP issues per-row collectives (4,672 all-reduces for a 64-row step). Batching the all-reduce would keep bit-invariance only if NCCL's per-element reduction order is independent of message size/algorithm (it picks the algorithm by size); pin NCCL_ALGO/NCCL_PROTO and test, or keep invariant for small RL batches.
4. Greedy fast path unrouted: vocab_parallel_argmax/_token_logprob would cut the logits gather (223 KiB -> 16 B per row) for greedy requests; the engine still gathers the full row. Worth it only if the measured gather time (model: ~5 ms at 64 rows) shows up.
5. tp > n_kv_heads (kv-head replication) and MoE/OLMoE TP are not built. 8B at tp=8 would work (1 kv head per rank) but this box has 4 GPUs.
6. Per-rank KV is bf16; fp8 KV (with per-tensor scales, see the single-device section) would double the pool again.
7. The 2 consensus all-gathers cost 8-18% of a toy gloo step; GPU cost to be measured with --no-verify.
8. Ranks must be fed identical requests in identical order; generate enforces it (rank 0's list is broadcast), direct add_request callers must do it themselves (a miss is caught at the next schedule check, not silently absorbed).

Review findings from building and testing this (all fixed, each pinned by a test)

• gloo's all_gather_into_tensor wants a flat rank-major output buffer, not [world, ...] (first parity run died on a size mismatch).
• safetensors slices are views whose storage spans the whole source tensor; kept as parameters on the CPU path they pin the full file mapping (found by the storage oracle, fixed by _own); a dim-0 narrow of an in-memory tensor has the same defect, so StateDictSource copies.
• PagedEngine.add_request validates ids against embed.weight.shape[0], which is V/tp for a sharded embedding: TPPagedEngine presents a full-vocabulary shim.
• PagedEngine(return_logits=True) stores float32 rows, which caps any logit parity check at ~6e-8; the fp64 oracles capture the row at sample_row instead.
• A test fixture's model.to(torch.bfloat16) is in place and silently turned a shared fp32 model into bf16 for every later test (caught by a byte-count bound).
• Mutants that desynchronise gloo can abort a rank in C++ (Is there a distributed collective mismatch in your code?) or hang to the timeout: the harness treats that as a kill for mutants only, never for controls.