# 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 `PYTHONHASHSEED`s). 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.