feat: top-p/top-k train sampling with kept-set mask replay

[faresobeid]

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, top_k, or min_p > 0) implies replay end to end:

• Every truncating policy-sourced sampling config gets a top-k bound (top_k respected if set, else defaulted to 512, which rarely binds at typical top_p), so kept sets are never larger than k. 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.kept_tokens (None = capture off, int = capture width) is derived as the largest configured top_k, so no kept set ever overflows — replay is exact at every position. It's only set by hand for standalone-launched inference 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.24 wheel — same pattern as router replay; no upstream support exists):

• 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 instead recompute the nucleus and must force-keep the sampled token against kernel boundary disagreements.)
• vLLM's inter-process output structs are fixed msgspec schemas, so the kept ids ride the existing logprobs channel as a -1-separated row extension, 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.
• Incompatible setups fail at startup instead of running silently biased: speculative decoding, logprobs_mode overrides, VLLM_USE_V2_MODEL_RUNNER=1.

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.

Paired dep branches

Submodule pointers in this PR pin:

• PrimeIntellect-ai/renderers feat/kept-tokens-parse — parse kept_tokens from /generate responses (base64 splice fast path).
• PrimeIntellect-ai/verifiers feat/kept-tokens-branch — TurnTokens → graph attribution → Branch.kept_tokens.

PRs for both still need to be opened.

Verification

All CPU-level, against the real installed vLLM 0.24:

• 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.
• The patched Sampler.forward produces kept sets matching an independent implementation of vLLM's top-p semantics; the returned sampled-token logprob equals the renormalized reference; greedy/untruncated steps stay position-aligned.
• Serialize → parse → branch spreading → transport → pack/truncate/pad → tensorize round-trips, with mixed masked/maskless samples.
• Config policy: top-k defaulting (explicit, per-env, inherited, min_p-triggered), derived inference.kept_tokens (including explicit-value override), extra_body truncation rejection, opd/filter guards, frozen-source exemption, standalone-server warning.

End-to-end GPU validation (H200, top_p 0.97 / top_k 512):

• reverse-text (Qwen3-0.6B, 20 steps): reward 0.10 → 0.78, mismatch_kl bounded 0.0005–0.0094, entropy rising, 0% errors.
• hendrycks sanity (R1-Distill-Qwen-1.5B, 200 steps, batch 512 × 8k ctx, W&B hendrycks-sanity/topp0.97-topk512-replay): mismatch_kl 0.0003–0.0004 at every step, flat from step 1 to 200 — below an untruncated control run's noise floor (0.00042–0.00046), since singleton kept sets contribute exactly zero mismatch. Train reward 0.49 → 0.68; AIME2024 eval 0.1875 → 0.2458; entropy flat (no collapse); DPPO trust-region masking 0.0000 in both runs; replay costs ≈5% MFU vs control.
• Measured mask load (2.5M-token steps): kept-set sizes mean ~8 / median 2 / p99 ~82, 45–48% singletons (matching Cognition's observation), top_k=512 never binds, 100.00% mask coverage on sampled tokens, ~32 B/token on the wire (~40% of the trainer payload, trivial in absolute bandwidth).

Known gap from the ablation: legacy v0 envs are not wired for kept tokens — truncated sampling on a v0 env fails loudly at the first group (the orchestrator's missing-mask guard). Needs the v0 ResponseTokens path wired for parity (routed_experts supports v0) or documented v1-only support.

🤖 Generated with Claude Code