[Feature] Mask Replay

[vx120]

Summary

This PR adds experimental support for sampling distribution replay.

A sampling mask represents the vocabulary support retained after top-k/top-p filtering. It is not an attention mask and does not affect causal attention or KV-cache behavior.

When enabled, vLLM returns the sampling support for each generated token in a CSR-style representation:

SamplingMask(
    token_ids=[...],
    offsets=[0, ..., len(token_ids)],
)

The feature is opt-in and does not change default generation behavior.

Motivation

With top-p sampling, rollout probabilities are normalized over a truncated vocabulary, while training commonly recomputes probabilities over the full vocabulary. This creates a systematic mismatch in importance ratios and KL estimates.

Sampling replay returns the rollout-time support so training can recompute log-probabilities over the same distribution.

Experiment

We compare GRPO training with sampling replay enabled and disabled under top-p sampling.

For each generated token:

• logp_diff = training_logprob - rollout_logprob
• importance_ratio = exp(logp_diff)

Ideally, the rollout and training distributions match, so logp_diff is close to 0 and importance_ratio is close to 1.

Mean Importance Ratio

With replay enabled, the mean importance ratio remains tightly centered around 1. Without replay, it is consistently below 1 and shows larger fluctuations.

This indicates that replay restores the rollout-time top-p normalization during training, while the non-replay baseline compares the truncated rollout distribution with a full-vocabulary training distribution.

Mean Log-Probability Difference

With replay enabled, the mean log-probability difference remains close to 0. Without replay, it has a persistent negative bias.

This is expected: a token probability normalized over the full vocabulary is generally lower than the probability normalized over the retained top-p support.

Standard Deviation of the Log-Probability Difference

Replay lowers the standard deviation of the log-probability difference. This means it improves not only the average alignment, but also token-level consistency between rollout and training probabilities.

Lower variance produces more stable importance weights and reduces the impact of probability-ratio outliers.

Approximate KL

Replay produces a lower and more stable approximate KL estimate, with fewer large spikes. This indicates that the training-time distribution stays closer to the rollout-time distribution throughout optimization.

Conclusion

Across all four metrics, sampling replay improves rollout/training distribution consistency:

• importance_ratio_mean stays close to 1;
• logp_diff_mean stays close to 0;
• logp_diff_std is lower;
• approx_kl is lower and more stable.

These results demonstrate the intended algorithmic effect of sampling replay: improved probability alignment and lower-variance importance weights.

This experiment does not, by itself, establish an improvement in downstream reward or task accuracy.

Limitations

The current implementation requires:

• Model Runner V2;
• logprobs_mode="processed_logprobs";
• SamplingParams.logprobs=1;
• temperature > 0.

Speculative decoding, custom samplers, and custom logits processors are not currently supported.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vx120.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[aoshen02]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5c9f82aeb9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve all tied tokens in sparse top-k masks

When multiple logits tie at the top-k boundary, apply_top_k_top_p_pytorch retains every tied token because it masks only values strictly below the kth value, so the actual sampling support can contain more than top_k tokens. Truncating the sparse representation with torch.topk(..., max_top_k) silently omits some supported tokens; if an omitted token is sampled, _build_sampling_mask_lists raises an error, and otherwise replay receives an incomplete support set and computes incorrect normalization. Build the sparse row from every finite processed logit, or otherwise preserve all boundary ties.

Useful? React with 👍 / 👎.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83627 for commit c9378745404c.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83630 for commit 2eae885da039.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83650 for commit b53b2a9cf3bc.

[njhill]

Thanks @vx120 @aoshen02

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83656 for commit ab01a119c053.

[aoshen02]

/ci retry

[github-actions]

✅ Queued 4 failed job(s) for retry in Buildkite CI #83656.