|
| 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 | +## Requirements |
| 57 | + |
| 58 | +| Requirement | Reason | |
| 59 | +| --- | --- | |
| 60 | +| `--return-sampling-mask` | Engine-level opt-in (disables FlashInfer sampler) | |
| 61 | +| `--logprobs-mode processed_logprobs` | Returned logprobs are normalized over the nucleus, not full vocab | |
| 62 | +| `temperature > 0` | Greedy has no truncated distribution | |
| 63 | +| `top_k > 0` | Bounds mask size; pure top-p can produce vocab-sized masks | |
| 64 | +| Model Runner V2 | Required by the async D2H copy pipeline | |
| 65 | + |
| 66 | +The engine rejects unsupported combinations at startup or request time: |
| 67 | + |
| 68 | +- Speculative decoding |
| 69 | +- Diffusion models |
| 70 | +- Custom logits processors (engine-level `--logits-processors`) |
| 71 | + |
| 72 | +## How it works |
| 73 | + |
| 74 | +1. The sampler applies all logit processors (penalties, logit bias, bad words, |
| 75 | + temperature, min-p) and then top-k/top-p filtering, which sets excluded |
| 76 | + logits to `-inf`. |
| 77 | +2. After sampling, `torch.isfinite(processed_logits)` identifies the surviving |
| 78 | + token IDs — this is the sampling mask. |
| 79 | +3. The mask is transferred GPU → CPU asynchronously alongside sampled tokens. |
| 80 | +4. On request completion, per-step masks are merged and converted to |
| 81 | + `list[list[int]]` for the response. |
| 82 | + |
| 83 | +## RL training usage |
| 84 | + |
| 85 | +The training side needs two things for the importance ratio `π_θ/π_old`: |
| 86 | + |
| 87 | +**`π_old(a|s)` — old policy's nucleus-normalized logprob:** |
| 88 | +Already returned by vLLM when `--logprobs-mode processed_logprobs` is set. |
| 89 | +The `log_softmax` is computed over processed logits (where filtered tokens |
| 90 | +are `-inf`), so the denominator only includes the nucleus. |
| 91 | + |
| 92 | +**`π_θ(a|s)` — current policy's nucleus-normalized logprob:** |
| 93 | +Computed by the training framework using the mask: |
| 94 | + |
| 95 | +```python |
| 96 | +# mask_ids: list[int], the sampling support for this token |
| 97 | +# logits: the training model's raw logits for this position |
| 98 | +keep = torch.zeros(vocab_size, dtype=torch.bool) |
| 99 | +keep[mask_ids] = True |
| 100 | +masked_logits = logits.masked_fill(~keep, float("-inf")) |
| 101 | +log_prob = log_softmax(masked_logits)[sampled_token_id] |
| 102 | +``` |
| 103 | + |
| 104 | +Both sides normalize over the same token set, so the importance ratio is |
| 105 | +consistent. |
| 106 | + |
| 107 | +## Limitations |
| 108 | + |
| 109 | +- **Engine-level flag:** `--return-sampling-mask` globally disables the |
| 110 | + FlashInfer fused sampler. All requests pay the cost of the PyTorch sampling |
| 111 | + path, even if they don't need the mask. |
| 112 | +- **No streaming support:** The mask is returned only in the final response, |
| 113 | + not in intermediate streaming chunks. |
0 commit comments