[Feature] Mask Replay

docs/training/sampling_mask.md

@@ -0,0 +1,113 @@
+# Sampling Mask (Distribution Replay)
+
+When using top-k/top-p sampling for RL rollouts (e.g. GRPO), there is a
+systematic mismatch between the truncated distribution the sampler actually
+drew from and the full-vocabulary softmax used to compute log-probabilities
+during training. The **sampling mask** feature closes this gap by returning
+the exact set of token IDs that survived top-k/top-p/min-p filtering at each
+generation step, so the training side can normalize over the same support.
+
+## Background
+
+This feature implements the **Keep Sampling Mask** strategy described in the
+[DeepSeek-V3.2 technical report](https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/assets/paper.pdf)
+(Section 3.3). The key insight: top-k/top-p truncation during rollout sampling
+introduces a mismatch between the action spaces of `π_old` and `π_θ`, which
+violates the principles of importance sampling and destabilizes training. By
+preserving the truncation masks from `π_old` and applying them to `π_θ` during
+training, both policies share identical action subspaces. DeepSeek reports that
+combining top-p sampling with the Keep Sampling Mask strategy effectively
+preserves language consistency during RL training.
+
+## Quick start
+
+```bash
+vllm serve <model> \
+    --return-sampling-mask \
+    --logprobs-mode processed_logprobs
+```
+
+```python
+from vllm import LLM, SamplingParams
+
+llm = LLM(model, return_sampling_mask=True,
+           logprobs_mode="processed_logprobs")
+output = llm.generate(
+    "The capital of France is",
+    SamplingParams(temperature=1.0, top_k=50, top_p=0.95, logprobs=1),
+)
+mask = output[0].outputs[0].sampling_mask
+# mask.token_ids: [[187, 326, 512], [42, 88], ...]
+#   mask.token_ids[i] = token IDs in the sampling support for generated token i
+```
+
+The mask is also available via the `/inference/v1/generate` HTTP endpoint:
+
+```json
+{
+  "choices": [{
+    "token_ids": [187, 42, 303],
+    "sampling_mask": [[187, 326, 512], [42, 88], [303, 11, 22]],
+    "finish_reason": "stop"
+  }]
+}
+```
+
+## Requirements
+
+| Requirement | Reason |
+| --- | --- |
+| `--return-sampling-mask` | Engine-level opt-in (disables FlashInfer sampler) |
+| `--logprobs-mode processed_logprobs` | Returned logprobs are normalized over the nucleus, not full vocab |
+| `temperature > 0` | Greedy has no truncated distribution |
+| `top_k > 0` | Bounds mask size; pure top-p can produce vocab-sized masks |
+| Model Runner V2 | Required by the async D2H copy pipeline |
+
+The engine rejects unsupported combinations at startup or request time:
+
+- Speculative decoding
+- Diffusion models
+- Custom logits processors (engine-level `--logits-processors`)
+
+## How it works
+
+1. The sampler applies all logit processors (penalties, logit bias, bad words,
+   temperature, min-p) and then top-k/top-p filtering, which sets excluded
+   logits to `-inf`.
+2. After sampling, `torch.isfinite(processed_logits)` identifies the surviving
+   token IDs — this is the sampling mask.
+3. The mask is transferred GPU → CPU asynchronously alongside sampled tokens.
+4. On request completion, per-step masks are merged and converted to
+   `list[list[int]]` for the response.
+
+## RL training usage
+
+The training side needs two things for the importance ratio `π_θ/π_old`:
+
+**`π_old(a|s)` — old policy's nucleus-normalized logprob:**
+Already returned by vLLM when `--logprobs-mode processed_logprobs` is set.
+The `log_softmax` is computed over processed logits (where filtered tokens
+are `-inf`), so the denominator only includes the nucleus.
+
+**`π_θ(a|s)` — current policy's nucleus-normalized logprob:**
+Computed by the training framework using the mask:
+
+```python
+# mask_ids: list[int], the sampling support for this token
+# logits: the training model's raw logits for this position
+keep = torch.zeros(vocab_size, dtype=torch.bool)
+keep[mask_ids] = True
+masked_logits = logits.masked_fill(~keep, float("-inf"))
+log_prob = log_softmax(masked_logits)[sampled_token_id]
+```
+
+Both sides normalize over the same token set, so the importance ratio is
+consistent.
+
+## Limitations
+
+- **Engine-level flag:** `--return-sampling-mask` globally disables the
+  FlashInfer fused sampler. All requests pay the cost of the PyTorch sampling
+  path, even if they don't need the mask.
+- **No streaming support:** The mask is returned only in the final response,
+  not in intermediate streaming chunks.

rust/src/engine-core-client/src/protocol/output.rs

@@ -123,6 +123,8 @@ pub struct EngineCoreOutput {
     /// frontend grows one.
     #[serde(default)]
     pub mm_cache_miss_hashes: Option<Vec<String>>,
+    #[serde(default)]
+    pub new_sampling_mask: Option<OpaqueValue>,
 }
 
 impl EngineCoreOutput {
@@ -440,6 +442,7 @@ mod tests {
                             routed_experts: None,
                             num_nans_in_logits: 0,
                             mm_cache_miss_hashes: None,
+                            new_sampling_mask: None,
                         },
                     ],
                     scheduler_stats: None,

rust/src/engine-core-client/src/tests/client.rs

@@ -2498,6 +2498,8 @@ fn python_msgpack_fixtures_match_rust_encoding() {
     let defaults_request_hex = lines.next().expect("missing defaults request fixture line");
     let multimodal_request_hex = lines.next().expect("missing multimodal request fixture line");
     let outputs_hex = lines.next().expect("missing outputs fixture line");
+    let sampling_mask_outputs_hex =
+        lines.next().expect("missing sampling mask outputs fixture line");
     let inline_logprobs_frames = lines.next().expect("missing inline logprobs fixture line");
     let multipart_logprobs_frames = lines.next().expect("missing multipart logprobs fixture line");
     let inline_prompt_frames = lines.next().expect("missing inline prompt logprobs fixture line");
@@ -2508,6 +2510,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
     let request_bytes = hex::decode(request_hex).unwrap();
     let multimodal_request_bytes = hex::decode(multimodal_request_hex).unwrap();
     let outputs_bytes = hex::decode(outputs_hex).unwrap();
+    let sampling_mask_outputs_bytes = hex::decode(sampling_mask_outputs_hex).unwrap();
 
     let decoded_request: EngineCoreRequest = rmp_serde::from_slice(&request_bytes).unwrap();
     let expected_request = sample_request();
@@ -2571,6 +2574,16 @@ fn python_msgpack_fixtures_match_rust_encoding() {
         decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap());
     assert_eq!(python_mm_features, rust_mm_features);
 
+    let decoded_sampling_mask_outputs: EngineCoreOutputs =
+        rmp_serde::from_slice(&sampling_mask_outputs_bytes).unwrap();
+    let sampling_mask_output =
+        &decoded_sampling_mask_outputs.as_request_batch().unwrap().outputs[0];
+    assert!(sampling_mask_output.mm_cache_miss_hashes.is_none());
+    assert!(matches!(
+        sampling_mask_output.new_sampling_mask.as_ref(),
+        Some(Value::Array(fields)) if fields.len() == 3
+    ));
+
     let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
     expect_test::expect![[r#"
         RequestBatch(
@@ -2598,6 +2611,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
                         routed_experts: None,
                         num_nans_in_logits: 0,
                         mm_cache_miss_hashes: None,
+                        new_sampling_mask: None,
                     },
                 ],
                 scheduler_stats: None,

rust/src/engine-core-client/src/tests/python_compat.py

@@ -98,6 +98,7 @@ class EngineCoreOutput(
     routed_experts: object | None = None
     num_nans_in_logits: int = 0
     mm_cache_miss_hashes: list[str] | None = None
+    new_sampling_mask: object | None = None
 
 
 class EngineCoreOutputs(
@@ -202,6 +203,29 @@ class EngineCoreOutputs(
     finished_requests={"req-1"},
 )
 
+sampling_mask_wire = [
+    [
+        "<i4",
+        [5],
+        msgpack.ExtType(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
+    ],
+    [
+        "<i8",
+        [2],
+        msgpack.ExtType(3, np.array([0, 5], dtype=np.int64).tobytes()),
+    ],
+    None,
+]
+outputs_with_sampling_mask = EngineCoreOutputs(
+    outputs=[
+        EngineCoreOutput(
+            request_id="req-mask",
+            new_token_ids=[16],
+            new_sampling_mask=sampling_mask_wire,
+        )
+    ]
+)
+
 
 def encode_ndarray(
     array: np.ndarray,
@@ -422,6 +446,7 @@ class EngineCoreReadyResponse:
 print(msgspec.msgpack.encode(defaults_request).hex())
 print(msgpack.packb(multimodal_request_wire, use_bin_type=True).hex())
 print(msgspec.msgpack.encode(outputs).hex())
+print(msgspec.msgpack.encode(outputs_with_sampling_mask).hex())
 print(" ".join(frame.hex() for frame in encode_output_frames(inline_logprobs)))
 print(
     " ".join(

tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py

@@ -8,6 +8,7 @@
 import pytest_asyncio
 from transformers import AutoTokenizer
 
+import vllm.envs as envs
 from tests.utils import RemoteOpenAIServer
 from vllm.config import ModelConfig
 from vllm.config.utils import getattr_iter
@@ -107,6 +108,49 @@ async def test_generate_endpoint(client):
     resp.raise_for_status()
     data = resp.json()
     assert "choices" in data
+    assert data["choices"][0].get("sampling_mask") is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    envs.VLLM_USE_RUST_FRONTEND,
+    reason="sampling mask output is not supported by the Rust frontend",
+)
+@pytest.mark.parametrize(
+    "server",
+    [["--return-sampling-mask", "--logprobs-mode", "processed_logprobs"]],
+    indirect=True,
+)
+async def test_generate_sampling_mask(client):
+    top_k = 5
+    payload = {
+        "model": MODEL_NAME,
+        "token_ids": [1, 2, 3],
+        "sampling_params": {
+            "max_tokens": 5,
+            "temperature": 0.8,
+            "top_k": top_k,
+            "top_p": 0.9,
+            "ignore_eos": True,
+            "seed": 0,
+        },
+        "stream": False,
+    }
+    resp = await client.post(GEN_ENDPOINT, json=payload)
+    resp.raise_for_status()
+    choice = resp.json()["choices"][0]
+
+    token_ids = choice["token_ids"]
+    sampling_mask = choice["sampling_mask"]
+    assert sampling_mask is not None
+    assert len(token_ids) == len(sampling_mask)
+
+    vocab_size = get_vocab_size(MODEL_NAME)
+    for token_id, support in zip(token_ids, sampling_mask):
+        assert support
+        assert len(support) == len(set(support))
+        assert all(0 <= support_token_id < vocab_size for support_token_id in support)
+        assert token_id in support
 
 
 @pytest.mark.asyncio

tests/v1/core/test_async_scheduler.py

@@ -291,6 +291,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.vllm_config = Mock()
     scheduler.vllm_config.model_config.enable_return_routed_experts = False
     scheduler.enable_return_routed_experts = False
+    scheduler.return_sampling_mask = False
     scheduler.recompute_kv_load_failures = False
     scheduler.defer_block_free = False
     scheduler.make_stats = Mock(return_value=None)

tests/v1/core/test_scheduler.py

@@ -3281,6 +3281,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.vllm_config = Mock()
     scheduler.vllm_config.model_config.enable_return_routed_experts = False
     scheduler.enable_return_routed_experts = False
+    scheduler.return_sampling_mask = False
     scheduler.recompute_kv_load_failures = False
     scheduler.defer_block_free = False
     scheduler.make_stats = Mock(return_value=None)

tests/v1/test_outputs.py

@@ -2,9 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from unittest import TestCase
 
+import numpy as np
 import torch
 
 from vllm.v1.outputs import LogprobsLists, LogprobsTensors
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
 
 
 def test_logprobs_tensors_cat():
@@ -30,6 +33,76 @@ def test_logprobs_tensors_cat():
     assert LogprobsTensors.cat([first]) is first
 
 
+def test_sampling_mask_tensors_tolist():
+    tensors = SamplingMaskTensors(
+        packed_mask=torch.tensor(
+            [[0b00101000], [0b00000000], [0b10000000]],
+            dtype=torch.uint8,
+        ),
+        counts=torch.tensor([2, 0, 1], dtype=torch.int32),
+        vocab_size=8,
+    )
+
+    result = tensors.tolists(np.array([1, 0, 1]))
+
+    assert result.token_ids.tolist() == [3, 5, 7]
+    assert result.offsets.tolist() == [0, 2, 3]
+    assert result.cu_num_generated_tokens == [0, 1, 1, 2]
+
+
+def test_sampling_mask_lists_to_nested_list():
+    from vllm.v1.outputs import SamplingMaskLists
+
+    mask = SamplingMaskLists(
+        token_ids=np.array([10, 11, 12, 20, 21]),
+        offsets=np.array([0, 3, 5]),
+    )
+
+    nested = mask.to_nested_list()
+
+    assert nested == [[10, 11, 12], [20, 21]]
+
+
+def test_sampling_mask_tensors_from_logits():
+    tensors = SamplingMaskTensors.from_logits(
+        logits=torch.tensor(
+            [
+                [1.0, float("-inf"), 2.0],
+                [3.0, 4.0, float("-inf")],
+                [float("-inf"), 5.0, 6.0],
+            ],
+            device="cuda",
+        ),
+        num_sampled_tokens=torch.tensor([1, 0, 1], device="cuda"),
+    )
+
+    result = tensors.tolists(np.array([1, 0, 1]))
+
+    assert result.token_ids.tolist() == [0, 2, 1, 2]
+    assert result.offsets.tolist() == [0, 2, 4]
+    assert result.cu_num_generated_tokens == [0, 1, 1, 2]
+
+
+def test_sampling_mask_matches_processed_top_k_top_p_support():
+    processed_logits = apply_top_k_top_p(
+        logits=torch.tensor([[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device="cuda"),
+        k=torch.tensor([3], device="cuda"),
+        p=torch.tensor([0.9], device="cuda"),
+    )
+    expected_token_ids = (
+        torch.isfinite(processed_logits[0]).nonzero().flatten().tolist()
+    )
+    assert 0 < len(expected_token_ids) < processed_logits.shape[1]
+
+    tensors = SamplingMaskTensors.from_logits(
+        processed_logits,
+        num_sampled_tokens=torch.tensor([1], device="cuda"),
+    )
+    result = tensors.tolists(np.array([1]))
+
+    assert result.to_nested_list() == [expected_token_ids]
+
+
 class TestLogprobsLists(TestCase):
     def setUp(self):
         self.logprobsLists = LogprobsLists(

vllm/config/model.py

@@ -241,6 +241,8 @@ class ModelConfig:
     flexibility."""
     enable_return_routed_experts: bool = False
     """Whether to return routed experts."""
+    return_sampling_mask: bool = False
+    """Whether to return the post-processing token support for each sample."""
     max_logprobs: int = Field(default=20, ge=-1)
     """Maximum number of log probabilities to return when `logprobs` is
     specified in `SamplingParams`. The default value comes the default for the
@@ -422,6 +424,7 @@ def compute_hash(self) -> str:
             "tokenizer_revision",
             "spec_target_max_model_len",
             "enforce_eager",
+            "return_sampling_mask",
             "logprobs_mode",
             "use_fp64_gumbel",
             "disable_cascade_attn",