| 1 | # Sampling Mask (Distribution Replay) |
| 2 | |
| 3 | When using top-k/top-p sampling for RL rollouts (e.g. GRPO), there is a |
| 4 | systematic mismatch between the truncated distribution the sampler actually |
| 5 | drew from and the full-vocabulary softmax used to compute log-probabilities |
| 6 | during training. The **sampling mask** feature closes this gap by returning |
| 7 | the exact set of token IDs that survived top-k/top-p/min-p filtering at each |
| 8 | generation step, so the training side can normalize over the same support. |
| 9 | |
| 10 | ## Background |
| 11 | |
| 12 | This feature implements the **Keep Sampling Mask** strategy described in the |
| 13 | [DeepSeek-V3.2 technical report](https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/assets/paper.pdf) |
| 14 | (Section 3.3). The key insight: top-k/top-p truncation during rollout sampling |
| 15 | introduces a mismatch between the action spaces of `π_old` and `π_θ`, which |
| 16 | violates the principles of importance sampling and destabilizes training. By |
| 17 | preserving the truncation masks from `π_old` and applying them to `π_θ` during |
| 18 | training, both policies share identical action subspaces. DeepSeek reports that |
| 19 | combining top-p sampling with the Keep Sampling Mask strategy effectively |
| 20 | preserves language consistency during RL training. |
| 21 | |
| 22 | ## Quick start |
| 23 | |
| 24 | ```bash |
| 25 | vllm serve <model> \ |
| 26 | --return-sampling-mask \ |
| 27 | --logprobs-mode processed_logprobs |
| 28 | ``` |
| 29 | |
| 30 | ```python |
| 31 | from vllm import LLM, SamplingParams |
| 32 | |
| 33 | llm = LLM(model, return_sampling_mask=True, |
| 34 | logprobs_mode="processed_logprobs") |
| 35 | output = llm.generate( |
| 36 | "The capital of France is", |
| 37 | SamplingParams(temperature=1.0, top_k=50, top_p=0.95, logprobs=1), |
| 38 | ) |
| 39 | mask = output[0].outputs[0].sampling_mask |
| 40 | # mask.token_ids: [[187, 326, 512], [42, 88], ...] |
| 41 | # mask.token_ids[i] = token IDs in the sampling support for generated token i |
| 42 | ``` |
| 43 | |
| 44 | The mask is also available via the `/inference/v1/generate` HTTP endpoint: |
| 45 | |
| 46 | ```json |
| 47 | { |
| 48 | "choices": [{ |
| 49 | "token_ids": [187, 42, 303], |
| 50 | "sampling_mask": [[187, 326, 512], [42, 88], [303, 11, 22]], |
| 51 | "finish_reason": "stop" |
| 52 | }] |
| 53 | } |
| 54 | ``` |
| 55 | |
| 56 | With `stream=true`, each chunk carries the masks of the tokens it emits. |
| 57 | |
| 58 | ## Requirements |
| 59 | |
| 60 | | Requirement | Reason | |
| 61 | | --- | --- | |
| 62 | | `--return-sampling-mask` | Engine-level opt-in (disables FlashInfer sampler) | |
| 63 | | `--logprobs-mode processed_logprobs` | Returned logprobs are normalized over the nucleus, not full vocab | |
| 64 | | `temperature > 0` | Greedy has no truncated distribution | |
| 65 | | `top_k > 0` | Bounds mask size; pure top-p can produce vocab-sized masks | |
| 66 | | Model Runner V2 | Required by the async D2H copy pipeline | |
| 67 | |
| 68 | Speculative decoding is supported; each emitted token gets the support of the |
| 69 | target distribution at its position. |
| 70 | |
| 71 | The engine rejects unsupported combinations at startup or request time: |
| 72 | |
| 73 | - Speculative decoding with adaptive verification |
| 74 | - Diffusion models |
| 75 | - Custom logits processors (engine-level `--logits-processors`) |
| 76 | |
| 77 | ## How it works |
| 78 | |
| 79 | 1. The sampler applies all logit processors (penalties, logit bias, bad words, |
| 80 | temperature, min-p) and then top-k/top-p filtering, which sets excluded |
| 81 | logits to `-inf`. |
| 82 | 2. After sampling, `torch.isfinite(processed_logits)` identifies the surviving |
| 83 | token IDs — this is the sampling mask. |
| 84 | 3. The mask is transferred GPU → CPU asynchronously alongside sampled tokens. |
| 85 | 4. On request completion, per-step masks are merged and converted to |
| 86 | `list[list[int]]` for the response. |
| 87 | |
| 88 | ## RL training usage |
| 89 | |
| 90 | The training side needs two things for the importance ratio `π_θ/π_old`: |
| 91 | |
| 92 | **`π_old(a|s)` — old policy's nucleus-normalized logprob:** |
| 93 | Already returned by vLLM when `--logprobs-mode processed_logprobs` is set. |
| 94 | The `log_softmax` is computed over processed logits (where filtered tokens |
| 95 | are `-inf`), so the denominator only includes the nucleus. |
| 96 | |
| 97 | **`π_θ(a|s)` — current policy's nucleus-normalized logprob:** |
| 98 | Computed by the training framework using the mask: |
| 99 | |
| 100 | ```python |
| 101 | # mask_ids: list[int], the sampling support for this token |
| 102 | # logits: the training model's raw logits for this position |
| 103 | keep = torch.zeros(vocab_size, dtype=torch.bool) |
| 104 | keep[mask_ids] = True |
| 105 | masked_logits = logits.masked_fill(~keep, float("-inf")) |
| 106 | log_prob = log_softmax(masked_logits)[sampled_token_id] |
| 107 | ``` |
| 108 | |
| 109 | Both sides normalize over the same token set, so the importance ratio is |
| 110 | consistent. |
| 111 | |
| 112 | ## Score centering with top-k logprobs |
| 113 | |
| 114 | [Score centering](https://arxiv.org/abs/2609.20807) needs the sampler's top-k |
| 115 | token IDs and logprobs plus the sampled token's logprob; the paper uses |
| 116 | `k=128` and does not need the exact sampling support. vLLM already returns |
| 117 | this through the ordinary `logprobs` parameter: |
| 118 | |
| 119 | ```bash |
| 120 | vllm serve <model> --logprobs-mode processed_logprobs --max-logprobs 128 |
| 121 | ``` |
| 122 | |
| 123 | ```json |
| 124 | {"sampling_params": {"temperature": 1.0, "top_p": 0.9, "logprobs": 128}} |
| 125 | ``` |
| 126 | |
| 127 | Each position then carries the sampled token's logprob and the top-128 |
| 128 | `(token_id, logprob)` pairs, normalized over the post-truncation support: |
| 129 | |
| 130 | ```json |
| 131 | {"token": "token_id:11", "logprob": -0.757, |
| 132 | "top_logprobs": [{"token": "token_id:11", "logprob": -0.757}, |
| 133 | {"token": "token_id:13", "logprob": -0.632}, |
| 134 | {"token": "token_id:0", "logprob": -9999.0}, ...]} |
| 135 | ``` |
| 136 | |
| 137 | - `--max-logprobs` defaults to 20; raise it to the requested `k`. |
| 138 | - Tokens outside the support are returned as `-9999.0` with arbitrary IDs; |
| 139 | drop them. The finite entries are the sampling mask when it has at most `k` |
| 140 | tokens. |
| 141 | - The sampled token is always the first entry and may repeat inside the top-k. |
| 142 | - `/inference/v1/generate` encodes IDs as `token_id:<id>`; the OpenAI |
| 143 | endpoints need `return_tokens_as_token_ids`. |
| 144 | |
| 145 | ## Limitations |
| 146 | |
| 147 | - **Engine-level flag:** `--return-sampling-mask` globally disables the |
| 148 | FlashInfer fused sampler. All requests pay the cost of the PyTorch sampling |
| 149 | path, even if they don't need the mask. |