[faresobeid]

Every score-centered or replayed request now captures the sampler's probabilities on the GPU, in vLLM's sampling step, with no host sync. vLLM returns no logprobs for these requests. Net +584 / −770 on current main.

What changes

• Sync-free paired GPU capture. Replayed (prl_replay) requests get each kept id with its processed logprob. Untruncated requests get the sampled token plus the top ids, or with extra_args.prl_sampler_draws = K, K i.i.d. draws from q (torch.multinomial on the global generator, so the draws are independent of the sampled token). All host copies are pinned and non-blocking, so vLLM's async scheduling keeps overlapping.
• Top-p-only replay with ragged CSR capture. Two Triton passes over vLLM's masked logits count each row's kept ids and pack (logprob << 32) | id into one flat device buffer. No top-k cap, no full-vocab bitmask, no injected top_k. The buffer size is inference.sampling_support_capacity (default 16M entries, 128 MB). A step whose supports do not fit fails those requests with an explicit error; nothing is truncated. The server logs the max per-step total every minute.
• One support gather in the fused LM head. On replayed rows the support ids are the mask ids, so one gather gives the replay normalizer, the support logprobs and the backward. LossInputs takes trainer_support_logprobs, sampler_support_logprobs, support_valid and sampler_tail_mass.
• Sampler draws for the sampled estimator (used by refactor(rl): one RL loss (ipo/icepop/reinforce) with shared exact score centering and sampler KL #4233), returned as compact logprobs with "samples": true.
• Removed: top_p_before_top_k / sampling_order.py (the p→k reorder; exact centering doesn't need it, since mask and logprobs come from the same processed logits), and the vLLM logprob-return path for compact requests (compact_logprobs.py, logprobs_mode / max_logprobs setup). TRAIN_TOP_K_DEFAULT injection and the top-k bound are gone too.
• Config: sampling-mask capture is enabled whenever train sampling is truncated or records sampler logprobs (score centering sets them).

Why

The old path computed logprobs in vLLM with a raised max_logprobs, synced GPU→CPU in capture (which stalled async scheduling), and padded supports to a fixed width, so top-p alone needed an injected top_k cap.

Numbers

Qwen3-1.7B (1×H200, 512 concurrent) and GLM-4.5-Air (TP 4, 256 concurrent), T=1, top_p 0.95, 1024 forced tokens, tok/s. replay = mask capture; sc = replay + sampler logprobs.

	Qwen top_k 128	Qwen top-p only	GLM top_k 128	GLM top-p only
no mask	23164	21699	9535	9543
replay	23165	20827	9337	8771
sc	24537	16761	9593	8980

• Capture kernel: 0.87 ms for 512 rows × 151936 with 9.3M kept entries (vLLM's top-p masking itself: 3.1 ms).
• Largest per-step total in a top-p-only Qwen run: 762k entries (4.5% of the default capacity). No overflows.
• Qwen top-p-only sc is lower because the API server serializes full-vocab rows at forced lengths (mean 174 kept per token); on GLM it is −6%.
• Fused LM head fwd+bwd (16k tokens, H=2048, K=128): replay + centering +1.6% over replay.
• Paired capture matches vLLM's processed_logprobs to ≤ 8.3e-7 on identical seeded sequences.

Validation

• CPU Slurm, on current main: pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/inference tests/unit/train -m "not gpu": 660 passed. ruff check, ruff format --check, uv lock --check: clean.
• GPU (H200): test_serving_tokens (capture equals the masked set, logprobs equal log_softmax over it, sampler rows merge, overflow fails), test_loss, test_fused_lm_head, test_nemotron_h_kl: 53 passed.
• End to end (Qwen3-0.6B reverse-text, icepop + exact score centering, top_p 0.95 without top_k, 10 steps): reward 0.173 → 0.752, the same as with top_k 2048; every token centered over its full support.

Before merge

• Merge feat(client): keep compact sampler-draw rows with their repeats renderers#202 (keeps sampler-draw rows with their repeats), then re-pin deps/renderers to its merge commit.

#4233 is stacked on this branch.

---

Note

High Risk
Changes the inference–trainer contract for truncated RL (logprobs, masks, capacity failures) and score-centered IcePop; misconfiguration or overflow now fails requests instead of silently bounding supports.

Overview
Sampling replay and IcePop score centering now get the sampler’s exact kept set and processed logprobs from GPU capture in vLLM’s sampling step, instead of raising max_logprobs and shipping logprobs through vLLM’s normal path. Replay and score-centered requests set logprobs = None on the server; supports are packed as CSR rows (logprob << 32) | id via Triton over the already masked processed logits, so top-p alone works without an injected top_k cap or fixed-width padding.

Removed the top_p_before_top_k engine patch (sampling_order.py), standalone compact_logprobs serialization for generic compact requests, orchestrator TRAIN_TOP_K_DEFAULT / 512 bound and mixed top-k validation, and RL auto-setup that forced processed_logprobs + max_logprobs. Added inference.sampling_support_capacity (default 16M entries); steps that exceed capacity fail explicitly rather than truncating.

The fused LM head uses one support gather for replay normalization, support logprobs, and backward; IcePop centering inputs are renamed to *_support_* with sampler_tail_mass instead of a “complete support” flag. Mask capture auto-enables when train sampling is truncated or records sampler logprobs (score centering). Optional prl_sampler_rows / prl_sampler_draws returns GPU-packed compact logprobs for untruncated score centering (including i.i.d. draws).

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

[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 ed83b67. Configure here.}