[faresobeid]

Stacked on #4232 (base feat/sc-exact-support). Allows sampling replay (and exact score centering) with top-p alone: no injected top_k, no width cap, no full-vocab bitmask.

Why

• vLLM's mask capture writes [rows, min(batch max top_k, 2048)] ids plus a full [rows, V/8] bitmask on every step. The bitmask is always copied to the host.
• feat(rl): one paired GPU capture and one support gather for exact score centering (follow-up to #4176) #4232's paired capture gathers only width columns, so a top-p-only row wider than the cap fails as incomplete evidence. That is why the config injects top_k = 128 and rejects top_k > 512.
• Uncapped top-p supports are usually small: mean 4-20 kept per token on chat at top_p 0.95-0.97. The tails are long, though, and every large batch contains some. At forced lengths they reach the whole vocabulary.

What

Capture (paired_capture.py, request_replay.py)

• The patched sampler stashes vLLM's from_logits arguments and builds the supports after sampling. It uses two Triton passes over the processed logits, which vLLM's sampler has already masked to the nucleus:
  1. For each replayed request: the count of finite logits and the log-normalizer over them.
  2. A prefix sum gives each row's offset. Each row's kept ids are then compacted, packed with their logprobs as (logprob bits << 32) | id, into one flat device buffer.
• Sampler-rows requests (score centering without replay) are written into the same buffer.
• There is no sync. The buffer has a fixed capacity, inference.sampling_support_capacity, default 16M entries (128 MB).
• If a step's supports do not fit, those requests get count -1 and fail with an explicit error. They are never truncated. The server logs the maximum per-step total every minute (vllm.prime_rl.sampling_support).
• The host copy has two phases. The counts ride vLLM's async copy. Once vLLM consumes the step, exactly the step's total entries are copied on a private stream, so the copy is sized by the real total, not rows × max width.
• The scheduler, output processor, serving and renderer formats are unchanged. They receive the same int64 CSR rows as before.

Validation

• Replay requests need temperature > 0. vLLM's top_k > 0 requirement for mask capture is dropped for them.

Config

• Removed TRAIN_TOP_K_DEFAULT injection, TRAIN_TOP_K_BOUND, and the "cannot mix top_k > 0 and -1" check. That check existed only because the capture width came from the batch max top_k.
• Score centering no longer requires a bounded top_k. For top-p-only sampling it sets logprobs >= 1 to opt into compact evidence.

vLLM's top-p is not exactly sort-based top-p

The sampled token is always drawn from exactly the captured set, because capture reads the logits the sampler masked. So replay stays consistent.

vLLM 0.32's Triton top-p (apply_top_k_top_p_triton) does not always keep the sort-based nucleus. Compared with an fp64 stable-sort reference, 64-128 rows × 151936:

rows	batch	p	rows differing	kept-count diff	min kept mass
N(0,4²) logits	64	0.99	40	-1..+6	0.99
N(0,4²)	128	0.95	118	-28..+28	0.949
N(0,4²)	128	0.99	128	-1057..+1661	0.987
rounded to 0.5 (ties)	64	0.95	64	-3085..+1825	0.936
peaked (σ=12, rounded to 1)	64	0.99	14	0..+5	0.990
peaked	128	0.95/0.99	0	0	—

On peaked, LLM-like rows it is nearly always exact. On flat rows the kept mass can fall below p, and rows with many tied logits differ widely. This PR leaves vLLM's top-p as is.

Numbers (H200, chat prompts, T=1, 1024 forced tokens per request, median of 3)

tok/s. replay = mask capture only; sc = replay plus compact sampler logprobs.

Qwen3-1.7B, 1 GPU, 512 concurrent

top_p	server	top_k 128	top_k 2048	top-p only
0.95	no mask	23164	21511	21699
0.95	#4232 replay	21694	20052	(rejected)
0.95	this PR, replay	23165	20654	20827
0.95	#4232 sc	24101	21950	(rejected)
0.95	this PR, sc	24537	21826	16761
0.97	no mask	23415	21421	23403
0.97	#4232 replay	22741	21128	(rejected)
0.97	this PR, replay	23252	20826	20217
0.97	#4232 sc	24164	22043	(rejected)
0.97	this PR, sc	24521	21763	15884

Kept per token with top-p only, Qwen3-1.7B: mean 174-213, p99.9 42k-69k, max 151934 (the whole vocab). This is the forced-length regime: once the answer is done, the model's distribution goes flat. With top_k 2048: mean 7-11.

GLM-4.5-Air, TP 4, 256 concurrent

top_p	server	top_k 128	top_k 2048	top-p only
0.95	no mask	9535	9114	9543
0.95	this PR, replay	9337	8887	8771
0.95	this PR, sc	9593	9080	8980
0.97	no mask	9564	9099	9510
0.97	this PR, replay	9396	8904	9244
0.97	this PR, sc	9558	9067	9375

GLM top-p only: kept mean 13-24, p99.9 250-520, max 151549 (the whole vocab).

• With top_k caps, capture is as fast as or faster than feat(rl): one paired GPU capture and one support gather for exact score centering (follow-up to #4176) #4232's fixed-width plus bitmask capture.
• Top-p-only replay vs no mask:
  • Qwen3-1.7B: -4% (p0.95) and -9..-14% (p0.97).
  • GLM-4.5-Air: -8% (p0.95) and -3% (p0.97).
  • Against top_k 2048 replay it is within ±4%.
• The bigger top-p-only sc drop on Qwen (16-18k vs 21k tok/s) comes from the API server serializing the compact logprobs of full-vocab rows (mean 174 entries per token). It tracks the forced-length tail. The GLM tail is milder, and there the drop is 1-6%.
• No capture overflows occurred at the default capacity. In a top-p-only Qwen3-1.7B rerun (512 × 1024), the largest per-step total was 761,765 entries (4.5% of capacity).

Capture kernel alone, 512 rows × 151936 with a wide synthetic distribution (9.3M kept): 0.87 ms. vLLM's fixed-width capture at K=2048 takes 0.22 ms (it can't hold those rows); top-p masking takes 3.1 ms.

End-to-end

Qwen3-0.6B reverse-text, 10 steps, icepop plus exact score centering, top_p 0.95 with no top_k. It was run together with the ragged trainer (stacked draft on #4233).

	top-p only (0.95)	top_p 0.95 + top_k 2048
reward step 1 → 10	0.173 → 0.752	0.177 → 0.752
mismatch_kl mean	2.24e-3	2.23e-3
is_masked mean	9e-5	1.6e-4
score_centering uncentered / tail_mass	0 / 0	0 / 0
score_centering mask_tv mean	0.0121	0.0122
support size mean (last step)	23.0	15.8
median step time (steps 2-10)	11.7 s	12.1 s

Every token was replayed and centered exactly over its full recorded support (head mass 1, tail 0).

Tests

• pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/inference -m "not gpu" on a CPU node: 413 passed. Ruff is clean.
• GPU: tests/unit/inference/test_serving_tokens.py, which includes a new GPU test, passes. The new test checks that the capture equals the masked set, that logprobs equal log_softmax over the kept set, that sampler rows are merged, and that overflow yields -1.

Note: #4233 rewrote the score-centering validator. Merging #4232 into #4233 after this PR will conflict in rl.py. The resolution is the same one-line change (k = (sampling.top_k or 1) for truncated sampling).