[faresobeid]

One RL loss for ipo, icepop and the new reinforce, with score centering always on and an optional sampler KL, over one ragged (CSR) per-token support from #4232:

loss_t = -adv_tau * A_t * (w(a) * log p(a) - sum_v sg[u_v * w_v] * log p_v) + beta_t * KL(q || p)

sum_v u_v f(v) estimates E_q[f] with one rule per token:

• Truncated sampling (top-p and/or top-k, replayed): the replayed support S with u_v = q_v. Centering is exact.
• Untruncated sampling: K = score_centering_k (default 16) i.i.d. sampler draws with u_v = 1/K, captured on the GPU next to the sampled token.

What changes

• One RLLoss replaces IPOLoss / IcePopLoss. Types differ only in w: ipo 1(|p-q| <= eps) * min(p/q, max_importance_ratio), icepop 1(low <= p/q <= high) * p/q, reinforce 1.
• Score centering is always on for these types; custom is unaffected. Config: score_centering_k only (replaces score_centering / score_centering_topk). The top-K head + tail estimator and its top-id capture are gone: compact-logprob requests now need replay or prl_sampler_draws.
• KL over the same support (kl_coeff, default 0). kl_estimator = exact | k1; kl_staleness = none | inv_sqrt uses per-token policy age from the engines' weight-version labels (main's per-token staleness stream).
• RLConfig requests the support for every policy-sampled train source: replay masks with sampler probabilities when truncated, sampling.sampler_draws = K when not (kept across live config updates). Train sampling needs temperature > 0.
• Ragged support end to end. TokenSupport (int32 ids, host int64 offsets, replayed flag, sampler logprobs) replaces the padded [N, K] mask / top-k tensors in data.py, CP sharding, the model forwards and the loss. The fused LM head slices fixed token chunks from the CSR, with no padding, sort or device-to-host sync.
• Sparse LM-head backward. When every RL token's logit gradient lies on its label and support (replay, or reinforce without KL), the head skips the dense softmax backward and uses a deterministic CSR matmul.
• Sampler-evidence validation runs on the GPU on the CSR entries (main's perf(trainer): validate score-centering data on the GPU #4284 checks), checked once per step before the optimizer.
• Metrics: score_centering/support_mass (E_q[p/q]), score_centering/mask_tv, and with the KL kl, kl/beta, kl/age0.
• Behavior change: every RL run is now centered. Untruncated runs capture K draws per token.

Why

Each token needs the sampler distribution over its support. Padding every support to the micro batch's widest row wasted more than 99% of those tensors on chat with uncapped top-p (about 1.5 GiB at 64k tokens, K=2048). The per-type losses duplicated the same centering math.

Numbers

Head + RLLoss fwd+bwd, H200, H=2048, V≈152k, real chat kept counts at top_p 0.97 (ms / peak GiB):

	16k tokens	64k tokens
no support	148 / 2.08	602 / 2.5
padded, K=2048 (before)	170 / 3.06	664 / 8.4
ragged, uncapped	170 / 2.21	674 / 2.6
ragged, sparse backward (replay or K8 draws, no KL)	72 / 0.94	259 / 1.14

• Ragged vs padded: bitwise on non-replayed logprobs and entropy; replayed logprobs within 1e-6; gradients ≤ 1 bf16 ulp. Sparse vs dense backward: rounding level in fp32.

Validation

• CPU Slurm: pytest tests/unit -m "not gpu": 1024 passed. ruff check, ruff format --check, uv lock --check (uv 0.13.5): clean.
• GPU (H200): test_serving_tokens, test_loss, test_fused_lm_head (incl. sparse vs dense backward), test_batch with -m gpu: 50 passed.
• Unit tests cover dense-reference gradients for each weight on a replayed support, the KL (exact and k1) against a dense KL, the K-draw estimator against the exact center, and the config capture rules.
• End to end (Qwen3-1.7B reverse-text, ipo, 4+4 H200): untruncated sampling with K=16 draws and top_p 0.95 replay with exact centering both train without errors. A/B with 2 runs per arm, 50 steps: exact centering at top_p 0.95 and K=16 / K=8 draws (simulated from the replayed q) reach the same reward (0.26–0.27 over the last 10 steps), and so does untruncated sampling with K=16 / K=8. Decode inter-token latency is about 10% higher for untruncated sampling with draws than for top_p 0.95 replay (3.4–3.5 vs 3.05–3.17 ms), the same for K=8 and K=16.

Before merge

• Merge feat(rl): one paired GPU capture and one support gather for exact score centering (follow-up to #4176) #4232 first, then retarget this PR to main.

---

Note

Medium Risk
Large refactor of RL loss math, batch wire format, and LM-head autograd on the hot training path; truncated runs now center by default, so regressions would show up as wrong gradients or KL rather than silent UI bugs.

Overview
Unifies RL policy losses into one RLLoss for ipo, icepop, and new reinforce, all sharing the same per-token objective: importance weight w, optional score centering under the sampler, and optional KL(q‖p) with exact / k1 estimators and inv_sqrt staleness from token policy age.

Replaces padded mask/top-k tensors with TokenSupport (CSR ids, host offsets, replay flag, sampler logprobs) from batch packing through model forward, CP gather, validation, and loss. IcePop-only centering knobs (score_centering_topk, IcePop-only flag) move to shared fields: score_centering (auto/true/false), score_centering_estimator (sampled/topk), score_centering_k. auto now exactly centers replayed (truncated) tokens by default—a behavior change vs IcePop-only centering on main.

Inference/orchestrator wiring: RLConfig auto-requests logprobs and, for untruncated sampled centering/KL, sampler_draws → prl_sampler_draws; vLLM rejects draws without compact logprobs. Fused LM head computes support_logprobs over CSR rows and can take a sparse backward (Triton segment_reduce / CSR matmul) when RL gradients lie only on label + support—large memory/time win for replay or REINFORCE+sampled centering without KL.

Docs and tests are updated for REINFORCE, StableReinforce mapping, GPU-side CSR validation, and dense-vs-sparse gradient parity.

^{Reviewed by Cursor Bugbot for commit 4509953. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit fc8dc13. Configure here.}