feat: top-p/top-k train sampling with sampling replay

[mikasenghaas]

Supersedes #2979 (same feature, re-cut from current main; see What changed vs #2979).

Adds top-p and top-k sampling support for train rollouts (both were hardcoded off). Truncated sampling renormalizes the rollout distribution over the surviving "kept set" of tokens; our rollout logprobs already reflect that (logprobs_mode = "processed_logprobs"), but the trainer normalizes over the full vocabulary — so every importance ratio is biased, and runs with truncated sampling collapse. This PR makes truncation safe by recording the kept set at sampling time and renormalizing trainer logprobs over the same set: DeepSeek V3.2's "Keep Sampling Mask" (arXiv:2512.02556 §3.1), also described in Cognition's SWE-1.7 post as "sampling distribution replay".

Usage

[orchestrator.train.sampling]
top_p = 0.95
top_k = 512   # optional — defaulted to 512 when truncation is on, since it bounds the kept sets

That's the whole config — there are no replay flags. Truncated train sampling (top_p < 1 and/or top_k) implies sampling replay end to end:

• Every truncating policy-sourced sampling config gets a top-k bound (top_k respected if set, else defaulted to 512; values above 512 are rejected — see below), so kept sets are never larger than the capture width. Truncation knobs must be the typed fields — smuggling them via extra_body is rejected. Frozen-source envs are exempt (external endpoints, no importance ratios).
• inference.enable_return_sampling_mask (bool, named after vLLM's in-flight native flag) turns on capture at a fixed width of 512; the orchestrator rejects train-sampling top_k > 512 so no kept set ever overflows — replay is exact at every position. The flag is auto-set and persisted into per-node configs; hand-setting is only for standalone-launched servers.
• The trainer is data-driven: it replays masks whenever a batch carries them, like any other per-token stream. The orchestrator enforces that truncating envs actually produce masks (fails fast if the server isn't capturing).
• Consumers of rollout logprobs that break under renormalization are rejected at config time: opd/opsd (reference logprobs are full-vocab prefill scores), and the gibberish/repetition filters (removed from the default lists, rejected if explicitly configured — their full-softmax thresholds misfire when singleton kept sets read as probability 1.0).

How it works

Inference (src/prime_rl/inference/vllm/kept_tokens.py, monkey patches over the stock vLLM 0.26 wheel):

• The kept set is read off the sampler's processed logprobs — the exact tensor the token was sampled from — so membership holds by construction. (slime and SGLang's original patches instead recompute the nucleus and must force-keep the sampled token against kernel boundary disagreements.)
• vLLM's inter-process output structs are fixed positional msgspec schemas, so the kept ids ride the existing logprobs channel as a -1-separated extension on the id tensor (ids only — nothing between sampler and API process pairs ids and logprob values column-wise), at a fixed device-side width (no host syncs). An API-process patch splits the extension back off before vLLM builds logprob dicts — chat/eval consumers see byte-identical logprobs — and /inference/v1/generate returns base64 {ids, counts} per choice, like routed_experts. Kept sets are decode-only, so PD-disaggregated serving needs no router changes.
• No env vars: the enable flag rides vLLM's additional_config as enable_return_sampling_mask, snapshotted at Sampler.__init__ where vLLM guarantees a config context (the fp32_lm_head mechanism). The API-process patches are data-driven off the separator id and install unconditionally — rows without extensions pass through untouched.
• Incompatible setups fail at startup instead of running silently biased: speculative decoding, logprobs_mode overrides, VLLM_USE_V2_MODEL_RUNNER=1 (prime-rl pins the V1 runner anyway).

Upstream path: vLLM is adding native support with the same semantics and constraints — vllm-project/vllm#49577 enable_return_sampling_mask, near-merge, earliest release ~0.28 (built for the V2 model runner). Once it ships in a release we pin, the two engine patches here reduce to the routed_experts-style API-layer glue (KeptTokensCapture + serializer). On released vLLM the only patch-free alternative today is requesting logprobs = top_k per token, which ships k ids+floats per position through vLLM's per-position logprob-dict machinery — orders of magnitude more transport and API-process work than this extension (~32 B/token measured).

Trainer (data-driven — replays masks whenever the batch carries them):

• Masked positions compute logprob = logits[label]/T - logsumexp(logits[kept]/T) in both the chunked fused LM head (backward restricted to kept ids) and the vanilla path. Positions without a mask (context tokens, non-policy samples) use full-vocab logprobs.
• Singleton kept sets (top token above the top-p threshold) give logprob 0 and exactly zero gradient — the entropy-preserving property from the SWE-1.7 post.
• Entropy stays full-vocab (it's a collapse diagnostic; matches slime).
• Gemma-family softcapped lm_heads don't implement kept-set renormalization and fail loudly on their head assert.

Transport: KeptTokens {ids, counts} (int32 bytes, CSR-style) on TrainingSample/MicroBatch, appended last to keep the positional wire layout stable; packed/truncated/padded alongside the other per-token streams; tensorized as [1, seq, max_kept] with -1 padding.

What changed vs #2979

• Re-cut as a single commit on current main, adapted to the vLLM 0.26 bump, the pass-through [inference.vllm] config, the multi-tenant removal, and the v0 env-compat drop — which also retires feat: top-p/top-k train sampling with sampling replay #2979's known gap (kept tokens were v1-only; v0 envs no longer exist).
• The engine-side extension now rides the logprob id tensor only; the -inf float filler rows were pure IPC overhead (halves the extension's engine→API traffic).
• The config knob is a bool named after vLLM's proposed flag (enable_return_sampling_mask) with a hardcoded capture width, transported via additional_config instead of env vars; the orchestrator rejects top_k > 512 instead of deriving a width.
• Documented the upstream alignment path (vLLM #49577, above).

Paired dep PRs (already merged and pinned)

• feat: parse kept_tokens sampling masks from /generate responses renderers#102 — parse kept_tokens from /generate responses (base64 splice fast path).
• feat: thread kept-set sampling masks from generate responses to Branch verifiers#1978 — KeptTokens dataclass: TurnTokens carrier → graph attribution → Branch.kept_tokens.

Both are ancestors of main's current submodule pins — this PR does not touch submodules.

Verification

Checks on this branch:

• uv run ruff check / ruff format --check, uv lock --check
• uv run pytest tests/unit/test_configs.py tests/unit/inference/ (120 passed)
• uv run pytest tests/unit/train/ tests/unit/orchestrator/ (165 passed; the one failure, test_qwen3_vl_e2e.py, fails identically on main — pre-existing, fix: token_id-formatted logprob tokens in the qwen3-vl fake engine #3161)
• Config auto-wiring dry-run: top_p 0.97 on reverse-text resolves to top_k = 512 (with warning) and inference.toml: kept_tokens = 512.
• CPU numeric check: selective_log_softmax_with_kept and the fused _SequenceChunkedLogProbEntropyFn (forward + backward) match a dense masked-renormalization reference (float32 error ≤ 5e-7; misaligned-mask fallback; singleton kept set → logprob 0, exactly zero grad).

End-to-end on reverse-text (Qwen3-0.6B-Reverse-Text-SFT, 20 steps, 1 trainer + 1 inference GPU), both runs from this branch:

• Baseline (no truncation) — regression check with the API-side patches installed unconditionally: reward 0.18 → 0.73, 0% rollout errors, 100% trainable, mismatch_kl 0.0007–0.0153. Traces carry top_p = 1.0; capture stays off (no capture ENABLED engine log), logprobs unaffected.
• top_p 0.97 (replay): reward 0.19 → 0.75, 0% rollout errors, mismatch_kl bounded 0.0006–0.0143, entropy healthy. Traces carry top_p = 0.97 / top_k = 512; the engine logs Kept-set sampling-mask capture ENABLED for this Sampler instance (cap=512) (via additional_config); and since the orchestrator raises on any truncating sample without masks, completing 20/20 steps means every trainable sample shipped its kept sets. W&B: reverse-text/reverse-text-{baseline,topp0.97}-pr3235.
• Rejection dry-run: top_k = 1024 with truncation fails config validation with the fixed-capture-width error.

Prior validation on #2979 (same logic; CPU tests were out-of-band, GPU runs on H200):

• Fused and vanilla logprob paths match a dense masked-renormalization reference, forward and backward — including singleton zero-grad, misaligned-mask fallback, temperature ≠ 1, and an exp-overflow regression case.
• hendrycks sanity (R1-Distill-Qwen-1.5B, 200 steps, batch 512 × 8k ctx): mismatch_kl 0.0003–0.0004 flat from step 1 to 200 — below an untruncated control's noise floor; train reward 0.49 → 0.68; AIME2024 eval 0.1875 → 0.2458; entropy flat; replay ≈5% MFU vs control.
• Measured mask load (2.5M-token steps): kept-set sizes mean ~8 / median 2 / p99 ~82, 45–48% singletons, top_k = 512 never binds, 100.00% mask coverage on sampled tokens, ~32 B/token on the wire.

🤖 Generated with Claude Code