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