[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> \
+    --enable-return-sampling-mask \
+    --logprobs-mode processed_logprobs
+```
+
+```python
+from vllm import LLM, SamplingParams
+
+llm = LLM(model, enable_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 |
+| --- | --- |
+| `--enable-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:** `--enable-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/chat/tests/chat.rs

@@ -58,6 +58,7 @@ fn request_output(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 
@@ -84,6 +85,7 @@ fn request_output_with_logprobs(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 

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

@@ -112,6 +112,8 @@ pub struct EngineCoreOutput {
     /// Number of NaNs seen in logits. Values above zero indicate corruption.
     #[serde(default)]
     pub num_nans_in_logits: u32,
+    #[serde(default)]
+    pub new_sampling_mask: Option<OpaqueValue>,
 }
 
 impl EngineCoreOutput {
@@ -384,6 +386,7 @@ mod tests {
                 prefill_stats: None,
                 routed_experts: None,
                 num_nans_in_logits: 0,
+                new_sampling_mask: None,
             }],
             finished_requests: Some(BTreeSet::from(["req-1".to_string()])),
             ..Default::default()
@@ -437,6 +440,7 @@ mod tests {
                             prefill_stats: None,
                             routed_experts: None,
                             num_nans_in_logits: 0,
+                            new_sampling_mask: None,
                         },
                     ],
                     scheduler_stats: None,

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

@@ -234,6 +234,7 @@ fn request_output(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 
@@ -2606,6 +2607,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
                         prefill_stats: None,
                         routed_experts: None,
                         num_nans_in_logits: 0,
+                        new_sampling_mask: None,
                     },
                 ],
                 scheduler_stats: None,

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

@@ -96,6 +96,7 @@ class EngineCoreOutput(
     prefill_stats: object | None = None
     routed_experts: object | None = None
     num_nans_in_logits: int = 0
+    new_sampling_mask: object | None = None
 
 
 class EngineCoreOutputs(

rust/src/llm/tests/generate.rs

@@ -59,6 +59,7 @@ fn request_output_with_events(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 
@@ -84,6 +85,7 @@ fn request_output_with_logprobs(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 
@@ -111,6 +113,7 @@ fn request_output_with_logprobs_and_kv(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 

rust/src/server/src/grpc/tests.rs

@@ -125,6 +125,7 @@ fn request_output(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 

rust/src/server/src/routes/http_client_tests.rs

@@ -108,6 +108,7 @@ fn request_output(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 

rust/src/server/src/routes/tests.rs

@@ -84,6 +84,7 @@ fn request_output_with_stop_reason(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 
@@ -110,6 +111,7 @@ fn request_output_with_logprobs(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 
@@ -138,6 +140,7 @@ fn request_output_with_logprobs_and_kv(
         prefill_stats: None,
         routed_experts: None,
         num_nans_in_logits: 0,
+        new_sampling_mask: None,
     }
 }
 

tests/test_outputs.py

@@ -3,7 +3,8 @@
 
 import pytest
 
-from vllm.outputs import RequestOutput
+from vllm.outputs import RequestOutput, SamplingMask
+from vllm.sampling_mask import SamplingMask as SamplingMaskFromModule
 
 pytestmark = pytest.mark.cpu_test
 
@@ -19,3 +20,7 @@ def test_request_output_forward_compatible():
         example_arg_added_in_new_version="some_value",
     )
     assert output is not None
+
+
+def test_sampling_mask_import():
+    assert SamplingMask is SamplingMaskFromModule

tests/v1/core/test_async_scheduler.py

@@ -290,6 +290,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.enable_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

@@ -3250,6 +3250,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.enable_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,11 @@
 # 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.worker.gpu.sample.output import SamplingMaskTensors
 
 
 def test_logprobs_tensors_cat():
@@ -30,6 +32,48 @@ def test_logprobs_tensors_cat():
     assert LogprobsTensors.cat([first]) is first
 
 
+def test_sampling_mask_tensors_tolist():
+    tensors = SamplingMaskTensors(
+        token_ids=torch.tensor([3, 5, 7]),
+        counts=torch.tensor([2, 1]),
+    )
+
+    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],
+            ]
+        ),
+        num_sampled_tokens=torch.tensor([1, 0, 1]),
+    )
+
+    assert tensors.token_ids.tolist() == [0, 2, 1, 2]
+    assert tensors.counts.tolist() == [2, 2]
+
+
 class TestLogprobsLists(TestCase):
     def setUp(self):
         self.logprobsLists = LogprobsLists(

vllm/config/model.py

@@ -235,6 +235,8 @@ class ModelConfig:
     flexibility."""
     enable_return_routed_experts: bool = False
     """Whether to return routed experts."""
+    enable_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
@@ -414,6 +416,7 @@ def compute_hash(self) -> str:
             "tokenizer_revision",
             "spec_target_max_model_len",
             "enforce_eager",
+            "enable_return_sampling_mask",
             "logprobs_mode",
             "use_fp64_gumbel",
             "disable_cascade_attn",

vllm/config/vllm.py

@@ -969,6 +969,31 @@ def _verify_kv_transfer_compat(self) -> None:
             "expandable_segments is automatically disabled)."
         )
 
+    def _verify_sampling_replay_config(self) -> None:
+        model_config = self.model_config
+        if model_config is None or not model_config.enable_return_sampling_mask:
+            return
+        if not self.use_v2_model_runner:
+            raise ValueError("sampling distribution replay requires Model Runner V2")
+        if self.speculative_config is not None:
+            raise ValueError(
+                "sampling distribution replay does not support speculative decoding"
+            )
+        if model_config.is_diffusion:
+            raise ValueError(
+                "sampling distribution replay does not support diffusion models"
+            )
+        if model_config.logits_processors:
+            raise ValueError(
+                "sampling distribution replay does not support custom logits processors"
+            )
+        if model_config.logprobs_mode != "processed_logprobs":
+            raise ValueError(
+                "sampling distribution replay requires "
+                "logprobs_mode='processed_logprobs' so that returned logprobs "
+                "are normalized over the same nucleus as the sampling mask"
+            )
+
     def __post_init__(self):
         """Verify configs are valid & consistent with each other."""
 
@@ -1010,6 +1035,8 @@ def __post_init__(self):
                     "connectors (PD disaggregation, KV cache offload)."
                 )
 
+        self._verify_sampling_replay_config()
+
         if self.lora_config is not None:
             self.lora_config.verify_with_model_config(self.model_config)