@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 50ba4bc6b2

[Feature] Mask Replay (#49577)

Signed-off-by: vx120 <893600387@qq.com> Signed-off-by: vx120 <57470515+vx120@users.noreply.github.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>

vx120 committed Aug 13, 2026 at 16:44 UTC 50ba4bc6b2cacd70c4711a1904dd7ad9740a578b
24 files changed +556 -5
docs/training/sampling_mask.md new
+113
@@ -0,0 +1,113 @@
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.
rust/src/engine-core-client/src/protocol/output.rs
+3
@@ -123,6 +123,8 @@ pub struct EngineCoreOutput {
123 /// frontend grows one.
124 #[serde(default)]
125 pub mm_cache_miss_hashes: Option<Vec<String>>,
126 + #[serde(default)]
127 + pub new_sampling_mask: Option<OpaqueValue>,
128 }
129
130 impl EngineCoreOutput {
@@ -440,6 +442,7 @@ mod tests {
442 routed_experts: None,
443 num_nans_in_logits: 0,
444 mm_cache_miss_hashes: None,
445 + new_sampling_mask: None,
446 },
447 ],
448 scheduler_stats: None,
rust/src/engine-core-client/src/tests/client.rs
+14
@@ -2498,6 +2498,8 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2498 let defaults_request_hex = lines.next().expect("missing defaults request fixture line");
2499 let multimodal_request_hex = lines.next().expect("missing multimodal request fixture line");
2500 let outputs_hex = lines.next().expect("missing outputs fixture line");
2501 + let sampling_mask_outputs_hex =
2502 + lines.next().expect("missing sampling mask outputs fixture line");
2503 let inline_logprobs_frames = lines.next().expect("missing inline logprobs fixture line");
2504 let multipart_logprobs_frames = lines.next().expect("missing multipart logprobs fixture line");
2505 let inline_prompt_frames = lines.next().expect("missing inline prompt logprobs fixture line");
@@ -2508,6 +2510,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2510 let request_bytes = hex::decode(request_hex).unwrap();
2511 let multimodal_request_bytes = hex::decode(multimodal_request_hex).unwrap();
2512 let outputs_bytes = hex::decode(outputs_hex).unwrap();
2513 + let sampling_mask_outputs_bytes = hex::decode(sampling_mask_outputs_hex).unwrap();
2514
2515 let decoded_request: EngineCoreRequest = rmp_serde::from_slice(&request_bytes).unwrap();
2516 let expected_request = sample_request();
@@ -2571,6 +2574,16 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2574 decode_value(&rmp_serde::to_vec_named(&expected_multimodal_request.mm_features).unwrap());
2575 assert_eq!(python_mm_features, rust_mm_features);
2576
2577 + let decoded_sampling_mask_outputs: EngineCoreOutputs =
2578 + rmp_serde::from_slice(&sampling_mask_outputs_bytes).unwrap();
2579 + let sampling_mask_output =
2580 + &decoded_sampling_mask_outputs.as_request_batch().unwrap().outputs[0];
2581 + assert!(sampling_mask_output.mm_cache_miss_hashes.is_none());
2582 + assert!(matches!(
2583 + sampling_mask_output.new_sampling_mask.as_ref(),
2584 + Some(Value::Array(fields)) if fields.len() == 3
2585 + ));
2586 +
2587 let decoded_outputs: EngineCoreOutputs = rmp_serde::from_slice(&outputs_bytes).unwrap();
2588 expect_test::expect![[r#"
2589 RequestBatch(
@@ -2598,6 +2611,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2611 routed_experts: None,
2612 num_nans_in_logits: 0,
2613 mm_cache_miss_hashes: None,
2614 + new_sampling_mask: None,
2615 },
2616 ],
2617 scheduler_stats: None,
rust/src/engine-core-client/src/tests/python_compat.py
+25
@@ -98,6 +98,7 @@ class EngineCoreOutput(
98 routed_experts: object | None = None
99 num_nans_in_logits: int = 0
100 mm_cache_miss_hashes: list[str] | None = None
101 + new_sampling_mask: object | None = None
102
103
104 class EngineCoreOutputs(
@@ -202,6 +203,29 @@ outputs = EngineCoreOutputs(
203 finished_requests={"req-1"},
204 )
205
206 +sampling_mask_wire = [
207 + [
208 + "<i4",
209 + [5],
210 + msgpack.ExtType(3, np.array([2, 12, 16, 17, 18], dtype=np.int32).tobytes()),
211 + ],
212 + [
213 + "<i8",
214 + [2],
215 + msgpack.ExtType(3, np.array([0, 5], dtype=np.int64).tobytes()),
216 + ],
217 + None,
218 +]
219 +outputs_with_sampling_mask = EngineCoreOutputs(
220 + outputs=[
221 + EngineCoreOutput(
222 + request_id="req-mask",
223 + new_token_ids=[16],
224 + new_sampling_mask=sampling_mask_wire,
225 + )
226 + ]
227 +)
228 +
229
230 def encode_ndarray(
231 array: np.ndarray,
@@ -422,6 +446,7 @@ print(msgspec.msgpack.encode(request).hex())
446 print(msgspec.msgpack.encode(defaults_request).hex())
447 print(msgpack.packb(multimodal_request_wire, use_bin_type=True).hex())
448 print(msgspec.msgpack.encode(outputs).hex())
449 +print(msgspec.msgpack.encode(outputs_with_sampling_mask).hex())
450 print(" ".join(frame.hex() for frame in encode_output_frames(inline_logprobs)))
451 print(
452 " ".join(
tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+44
@@ -8,6 +8,7 @@ import pytest
8 import pytest_asyncio
9 from transformers import AutoTokenizer
10
11 +import vllm.envs as envs
12 from tests.utils import RemoteOpenAIServer
13 from vllm.config import ModelConfig
14 from vllm.config.utils import getattr_iter
@@ -107,6 +108,49 @@ async def test_generate_endpoint(client):
108 resp.raise_for_status()
109 data = resp.json()
110 assert "choices" in data
111 + assert data["choices"][0].get("sampling_mask") is None
112 +
113 +
114 +@pytest.mark.asyncio
115 +@pytest.mark.skipif(
116 + envs.VLLM_USE_RUST_FRONTEND,
117 + reason="sampling mask output is not supported by the Rust frontend",
118 +)
119 +@pytest.mark.parametrize(
120 + "server",
121 + [["--return-sampling-mask", "--logprobs-mode", "processed_logprobs"]],
122 + indirect=True,
123 +)
124 +async def test_generate_sampling_mask(client):
125 + top_k = 5
126 + payload = {
127 + "model": MODEL_NAME,
128 + "token_ids": [1, 2, 3],
129 + "sampling_params": {
130 + "max_tokens": 5,
131 + "temperature": 0.8,
132 + "top_k": top_k,
133 + "top_p": 0.9,
134 + "ignore_eos": True,
135 + "seed": 0,
136 + },
137 + "stream": False,
138 + }
139 + resp = await client.post(GEN_ENDPOINT, json=payload)
140 + resp.raise_for_status()
141 + choice = resp.json()["choices"][0]
142 +
143 + token_ids = choice["token_ids"]
144 + sampling_mask = choice["sampling_mask"]
145 + assert sampling_mask is not None
146 + assert len(token_ids) == len(sampling_mask)
147 +
148 + vocab_size = get_vocab_size(MODEL_NAME)
149 + for token_id, support in zip(token_ids, sampling_mask):
150 + assert support
151 + assert len(support) == len(set(support))
152 + assert all(0 <= support_token_id < vocab_size for support_token_id in support)
153 + assert token_id in support
154
155
156 @pytest.mark.asyncio
tests/v1/core/test_async_scheduler.py
+1
@@ -291,6 +291,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
291 scheduler.vllm_config = Mock()
292 scheduler.vllm_config.model_config.enable_return_routed_experts = False
293 scheduler.enable_return_routed_experts = False
294 + scheduler.return_sampling_mask = False
295 scheduler.recompute_kv_load_failures = False
296 scheduler.defer_block_free = False
297 scheduler.make_stats = Mock(return_value=None)
tests/v1/core/test_scheduler.py
+1
@@ -3281,6 +3281,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
3281 scheduler.vllm_config = Mock()
3282 scheduler.vllm_config.model_config.enable_return_routed_experts = False
3283 scheduler.enable_return_routed_experts = False
3284 + scheduler.return_sampling_mask = False
3285 scheduler.recompute_kv_load_failures = False
3286 scheduler.defer_block_free = False
3287 scheduler.make_stats = Mock(return_value=None)
tests/v1/test_outputs.py
+73
@@ -2,9 +2,12 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from unittest import TestCase
4
5 +import numpy as np
6 import torch
7
8 from vllm.v1.outputs import LogprobsLists, LogprobsTensors
9 +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
10 +from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors
11
12
13 def test_logprobs_tensors_cat():
@@ -30,6 +33,76 @@ def test_logprobs_tensors_cat():
33 assert LogprobsTensors.cat([first]) is first
34
35
36 +def test_sampling_mask_tensors_tolist():
37 + tensors = SamplingMaskTensors(
38 + packed_mask=torch.tensor(
39 + [[0b00101000], [0b00000000], [0b10000000]],
40 + dtype=torch.uint8,
41 + ),
42 + counts=torch.tensor([2, 0, 1], dtype=torch.int32),
43 + vocab_size=8,
44 + )
45 +
46 + result = tensors.tolists(np.array([1, 0, 1]))
47 +
48 + assert result.token_ids.tolist() == [3, 5, 7]
49 + assert result.offsets.tolist() == [0, 2, 3]
50 + assert result.cu_num_generated_tokens == [0, 1, 1, 2]
51 +
52 +
53 +def test_sampling_mask_lists_to_nested_list():
54 + from vllm.v1.outputs import SamplingMaskLists
55 +
56 + mask = SamplingMaskLists(
57 + token_ids=np.array([10, 11, 12, 20, 21]),
58 + offsets=np.array([0, 3, 5]),
59 + )
60 +
61 + nested = mask.to_nested_list()
62 +
63 + assert nested == [[10, 11, 12], [20, 21]]
64 +
65 +
66 +def test_sampling_mask_tensors_from_logits():
67 + tensors = SamplingMaskTensors.from_logits(
68 + logits=torch.tensor(
69 + [
70 + [1.0, float("-inf"), 2.0],
71 + [3.0, 4.0, float("-inf")],
72 + [float("-inf"), 5.0, 6.0],
73 + ],
74 + device="cuda",
75 + ),
76 + num_sampled_tokens=torch.tensor([1, 0, 1], device="cuda"),
77 + )
78 +
79 + result = tensors.tolists(np.array([1, 0, 1]))
80 +
81 + assert result.token_ids.tolist() == [0, 2, 1, 2]
82 + assert result.offsets.tolist() == [0, 2, 4]
83 + assert result.cu_num_generated_tokens == [0, 1, 1, 2]
84 +
85 +
86 +def test_sampling_mask_matches_processed_top_k_top_p_support():
87 + processed_logits = apply_top_k_top_p(
88 + logits=torch.tensor([[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device="cuda"),
89 + k=torch.tensor([3], device="cuda"),
90 + p=torch.tensor([0.9], device="cuda"),
91 + )
92 + expected_token_ids = (
93 + torch.isfinite(processed_logits[0]).nonzero().flatten().tolist()
94 + )
95 + assert 0 < len(expected_token_ids) < processed_logits.shape[1]
96 +
97 + tensors = SamplingMaskTensors.from_logits(
98 + processed_logits,
99 + num_sampled_tokens=torch.tensor([1], device="cuda"),
100 + )
101 + result = tensors.tolists(np.array([1]))
102 +
103 + assert result.to_nested_list() == [expected_token_ids]
104 +
105 +
106 class TestLogprobsLists(TestCase):
107 def setUp(self):
108 self.logprobsLists = LogprobsLists(
vllm/config/model.py
+3
@@ -241,6 +241,8 @@ class ModelConfig:
241 flexibility."""
242 enable_return_routed_experts: bool = False
243 """Whether to return routed experts."""
244 + return_sampling_mask: bool = False
245 + """Whether to return the post-processing token support for each sample."""
246 max_logprobs: int = Field(default=20, ge=-1)
247 """Maximum number of log probabilities to return when `logprobs` is
248 specified in `SamplingParams`. The default value comes the default for the
@@ -422,6 +424,7 @@ class ModelConfig:
424 "tokenizer_revision",
425 "spec_target_max_model_len",
426 "enforce_eager",
427 + "return_sampling_mask",
428 "logprobs_mode",
429 "use_fp64_gumbel",
430 "disable_cascade_attn",
vllm/config/vllm.py
+27
@@ -1039,6 +1039,31 @@ class VllmConfig:
1039 "expandable_segments is automatically disabled)."
1040 )
1041
1042 + def _verify_sampling_replay_config(self) -> None:
1043 + model_config = self.model_config
1044 + if model_config is None or not model_config.return_sampling_mask:
1045 + return
1046 + if not self.use_v2_model_runner:
1047 + raise ValueError("sampling distribution replay requires Model Runner V2")
1048 + if self.speculative_config is not None:
1049 + raise ValueError(
1050 + "sampling distribution replay does not support speculative decoding"
1051 + )
1052 + if model_config.is_diffusion:
1053 + raise ValueError(
1054 + "sampling distribution replay does not support diffusion models"
1055 + )
1056 + if model_config.logits_processors:
1057 + raise ValueError(
1058 + "sampling distribution replay does not support custom logits processors"
1059 + )
1060 + if model_config.logprobs_mode != "processed_logprobs":
1061 + raise ValueError(
1062 + "sampling distribution replay requires "
1063 + "logprobs_mode='processed_logprobs' so that returned logprobs "
1064 + "are normalized over the same nucleus as the sampling mask"
1065 + )
1066 +
1067 def __post_init__(self):
1068 """Verify configs are valid & consistent with each other."""
1069
@@ -1080,6 +1105,8 @@ class VllmConfig:
1105 "connectors (PD disaggregation, KV cache offload)."
1106 )
1107
1108 + self._verify_sampling_replay_config()
1109 +
1110 if self.lora_config is not None:
1111 self.lora_config.verify_with_model_config(self.model_config)
1112
vllm/engine/arg_utils.py
+6
@@ -426,6 +426,7 @@ class EngineArgs:
426
427 model: str = ModelConfig.model
428 enable_return_routed_experts: bool = ModelConfig.enable_return_routed_experts
429 + return_sampling_mask: bool = ModelConfig.return_sampling_mask
430 model_weights: str = ModelConfig.model_weights
431 served_model_name: str | list[str] | None = ModelConfig.served_model_name
432 tokenizer: str | None = ModelConfig.tokenizer
@@ -874,6 +875,10 @@ class EngineArgs:
875 "--enable-return-routed-experts",
876 **model_kwargs["enable_return_routed_experts"],
877 )
878 + model_group.add_argument(
879 + "--return-sampling-mask",
880 + **model_kwargs["return_sampling_mask"],
881 + )
882 model_group.add_argument("--max-logprobs", **model_kwargs["max_logprobs"])
883 model_group.add_argument("--logprobs-mode", **model_kwargs["logprobs_mode"])
884 model_group.add_argument("--use-fp64-gumbel", **model_kwargs["use_fp64_gumbel"])
@@ -1748,6 +1753,7 @@ class EngineArgs:
1753 allow_deprecated_quantization=self.allow_deprecated_quantization,
1754 enforce_eager=self.enforce_eager,
1755 enable_return_routed_experts=self.enable_return_routed_experts,
1756 + return_sampling_mask=self.return_sampling_mask,
1757 max_logprobs=self.max_logprobs,
1758 logprobs_mode=self.logprobs_mode,
1759 use_fp64_gumbel=self.use_fp64_gumbel,
vllm/entrypoints/llm.py
+4
@@ -163,6 +163,8 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
163 dictionary or an AttentionConfig instance. If a dictionary, it will
164 be converted to an AttentionConfig. Allows specifying the attention
165 backend and other attention-related settings.
166 + return_sampling_mask: Return each sampled token's post-processing
167 + support set. Requires Model Runner V2 and processed log probabilities.
168 spec_method: Top-level alias for `speculative_config["method"]`.
169 spec_model: Top-level alias for `speculative_config["model"]`.
170 spec_tokens: Top-level alias for
@@ -201,6 +203,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
203 offload_params: set[str] | None = None,
204 enforce_eager: bool = False,
205 enable_return_routed_experts: bool = False,
206 + return_sampling_mask: bool = False,
207 disable_custom_all_reduce: bool = False,
208 hf_token: bool | str | None = None,
209 hf_overrides: HfOverrides | None = None,
@@ -317,6 +320,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
320 offload_params=offload_params or set(),
321 enforce_eager=enforce_eager,
322 enable_return_routed_experts=enable_return_routed_experts,
323 + return_sampling_mask=return_sampling_mask,
324 disable_custom_all_reduce=disable_custom_all_reduce,
325 hf_token=hf_token,
326 hf_overrides=hf_overrides,
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+2
@@ -210,6 +210,7 @@ class GenerateResponseChoice(BaseModel):
210 # ``None`` if (a) the request was aborted before any forward pass,
211 # or (b) ``enable_return_routed_experts`` is off server-side.
212 routed_experts: str | None = None
213 + sampling_mask: list[list[int]] | None = None
214
215 @field_validator("token_ids")
216 @classmethod
@@ -225,6 +226,7 @@ class GenerateResponseStreamChoice(BaseModel):
226 finish_reason: str | None = None
227 token_ids: list[int] | None = None
228 routed_experts: str | None = None
229 + sampling_mask: list[list[int]] | None = None
230
231
232 class GenerateStreamResponse(BaseModel):
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+5
@@ -315,12 +315,17 @@ class ServingTokens(GenerateBaseServing):
315 np.save(buf, output.routed_experts)
316 routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
317
318 + sampling_mask = None
319 + if output.sampling_mask is not None:
320 + sampling_mask = output.sampling_mask.token_ids
321 +
322 choice_data = GenerateResponseChoice(
323 index=output.index,
324 logprobs=logprobs,
325 finish_reason=output.finish_reason if output.finish_reason else "stop",
326 token_ids=as_list(output.token_ids),
327 routed_experts=routed_experts_b64,
328 + sampling_mask=sampling_mask,
329 )
330
331 choices.append(choice_data)
vllm/outputs.py
+15
@@ -18,6 +18,17 @@ from vllm.v1.metrics.stats import RequestStateStats
18 logger = init_logger(__name__)
19
20
21 +@dataclass
22 +class SamplingMask:
23 + """Per-token sampling support sets aligned with completion token IDs.
24 +
25 + Each inner list contains the vocabulary token IDs that survived
26 + top-k / top-p / min-p filtering for the corresponding generated token.
27 + """
28 +
29 + token_ids: list[list[int]]
30 +
31 +
32 @dataclass
33 class CompletionOutput:
34 """The output data of one completion output of a request.
@@ -30,6 +41,8 @@ class CompletionOutput:
41 output text.
42 logprobs: The log probabilities of the top probability words at each
43 position if the logprobs are requested.
44 + sampling_mask: The post-processing token support set for each generated
45 + token, if requested.
46 finish_reason: The reason why the sequence is finished.
47 stop_reason: The stop string or token id that caused the completion
48 to stop, None if the completion finished for some other reason
@@ -46,6 +59,7 @@ class CompletionOutput:
59 finish_reason: str | None = None
60 stop_reason: int | str | None = None
61 lora_request: LoRARequest | None = None
62 + sampling_mask: SamplingMask | None = None
63
64 def finished(self) -> bool:
65 return self.finish_reason is not None
@@ -56,6 +70,7 @@ class CompletionOutput:
70 f"text={self.text!r}, "
71 f"token_ids={self.token_ids}, "
72 f"routed_experts={self.routed_experts}, "
73 + f"sampling_mask={self.sampling_mask}, "
74 f"cumulative_logprob={self.cumulative_logprob}, "
75 f"logprobs={self.logprobs}, "
76 f"finish_reason={self.finish_reason}, "
vllm/v1/core/sched/scheduler.py
+10
@@ -326,6 +326,7 @@ class Scheduler(SchedulerInterface):
326 self.enable_return_routed_experts = (
327 vllm_config.model_config.enable_return_routed_experts
328 )
329 + self.return_sampling_mask = vllm_config.model_config.return_sampling_mask
330
331 if self.enable_return_routed_experts:
332 assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (
@@ -1809,6 +1810,7 @@ class Scheduler(SchedulerInterface):
1810
1811 stopped = False
1812 new_logprobs = None
1813 + new_sampling_mask = None
1814 new_token_ids = generated_token_ids
1815 pooler_output = pooler_outputs[req_index] if pooler_outputs else None
1816 kv_transfer_params = None
@@ -1938,6 +1940,13 @@ class Scheduler(SchedulerInterface):
1940 ):
1941 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
1942
1943 + if self.return_sampling_mask:
1944 + sampling_masks = model_runner_output.sampling_masks
1945 + if new_token_ids and sampling_masks is not None:
1946 + new_sampling_mask = sampling_masks.slice_request(
1947 + req_index, len(new_token_ids)
1948 + )
1949 +
1950 if num_nans_in_logits is not None and req_id in num_nans_in_logits:
1951 request.num_nans_in_logits = num_nans_in_logits[req_id]
1952
@@ -1951,6 +1960,7 @@ class Scheduler(SchedulerInterface):
1960 new_token_ids=new_token_ids,
1961 finish_reason=finish_reason,
1962 new_logprobs=new_logprobs,
1963 + new_sampling_mask=new_sampling_mask,
1964 new_prompt_logprobs_tensors=prompt_logprobs_tensors,
1965 pooling_output=pooler_output,
1966 stop_reason=request.stop_reason,
vllm/v1/engine/__init__.py
+3 -1
@@ -17,7 +17,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec
17 from vllm.pooling_params import PoolingParams
18 from vllm.sampling_params import SamplingParams
19 from vllm.v1.metrics.stats import PrefillStats, SchedulerStats
20 -from vllm.v1.outputs import LogprobsLists, LogprobsTensors
20 +from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplingMaskLists
21 from vllm.v1.serial_utils import UtilityResult
22
23 # Type for pause_generation mode parameter.
@@ -217,6 +217,8 @@ class EngineCoreOutput(
217 # so `array_like` positional serialization stays backward compatible.
218 mm_cache_miss_hashes: list[str] | None = None
219
220 + new_sampling_mask: SamplingMaskLists | None = None
221 +
222 @property
223 def finished(self) -> bool:
224 return self.finish_reason is not None
vllm/v1/engine/input_processor.py
+11
@@ -108,6 +108,17 @@ class InputProcessor:
108 self.tokenizer,
109 )
110
111 + if self.model_config.return_sampling_mask:
112 + if params.temperature <= 0:
113 + raise ValueError(
114 + "sampling distribution replay requires temperature > 0"
115 + )
116 + if params.top_k <= 0:
117 + raise ValueError(
118 + "sampling distribution replay requires top_k > 0 to "
119 + "bound sampling mask size, reduce transfer overhead, "
120 + "and avoid potential OOMs"
121 + )
122 if params.thinking_token_budget is not None and (
123 self.vllm_config.reasoning_config is None
124 or not self.vllm_config.reasoning_config.enabled
vllm/v1/engine/output_processor.py
+13
@@ -17,6 +17,7 @@ from vllm.outputs import (
17 PoolingOutput,
18 PoolingRequestOutput,
19 RequestOutput,
20 + SamplingMask,
21 )
22 from vllm.sampling_params import RequestOutputKind
23 from vllm.tokenizers import TokenizerLike
@@ -37,6 +38,7 @@ from vllm.v1.metrics.stats import (
38 RequestStateStats,
39 SchedulerStats,
40 )
41 +from vllm.v1.outputs import SamplingMaskLists
42
43 # shared empty CPU tensor used as a placeholder pooling output
44 EMPTY_CPU_TENSOR = torch.empty(0, device="cpu")
@@ -178,6 +180,7 @@ class RequestState:
180
181 # Routed experts accumulation (prompt + sample chunks)
182 self.routed_experts_chunks: list[np.ndarray] = []
183 + self.sampling_mask_chunks: list[SamplingMaskLists] = []
184
185 # Stream Interval
186 self.stream_interval = stream_interval
@@ -406,6 +409,11 @@ class RequestState:
409 if delta and logprobs:
410 logprobs = logprobs[-len(token_ids) :]
411
412 + sampling_mask = None
413 + if finished and self.sampling_mask_chunks:
414 + merged = SamplingMaskLists.merge(self.sampling_mask_chunks)
415 + sampling_mask = SamplingMask(merged.to_nested_list())
416 +
417 # Concatenate routed experts on finish
418 routed_experts = None
419 if finished and self.routed_experts_chunks:
@@ -416,6 +424,7 @@ class RequestState:
424 text=text,
425 token_ids=token_ids,
426 routed_experts=routed_experts,
427 + sampling_mask=sampling_mask,
428 logprobs=logprobs,
429 cumulative_logprob=self.logprobs_processor.cumulative_logprob,
430 finish_reason=str(finish_reason) if finished else None,
@@ -652,6 +661,10 @@ class OutputProcessor:
661 if pooling_output is None:
662 assert req_state.detokenizer is not None
663 assert req_state.logprobs_processor is not None
664 + if engine_core_output.new_sampling_mask is not None:
665 + req_state.sampling_mask_chunks.append(
666 + engine_core_output.new_sampling_mask
667 + )
668 # 2) Detokenize the token ids into text and perform stop checks.
669 stop_string = req_state.detokenizer.update(
670 new_token_ids, finish_reason == FinishReason.STOP
vllm/v1/outputs.py
+42
@@ -50,6 +50,45 @@ class LogprobsLists(NamedTuple):
50 )
51
52
53 +class SamplingMaskLists(NamedTuple):
54 + # [num_kept_tokens]
55 + token_ids: np.ndarray
56 + # [num_generated_tokens + 1]
57 + offsets: np.ndarray
58 + # [num_reqs + 1]
59 + cu_num_generated_tokens: list[int] | None = None
60 +
61 + def slice_request(self, req_idx: int, num_positions: int) -> "SamplingMaskLists":
62 + if self.cu_num_generated_tokens is None:
63 + start_idx = req_idx
64 + else:
65 + start_idx = self.cu_num_generated_tokens[req_idx]
66 + end_idx = start_idx + num_positions
67 + flat_start = self.offsets[start_idx]
68 + flat_end = self.offsets[end_idx]
69 + return SamplingMaskLists(
70 + self.token_ids[flat_start:flat_end],
71 + self.offsets[start_idx : end_idx + 1] - flat_start,
72 + None,
73 + )
74 +
75 + def to_nested_list(self) -> list[list[int]]:
76 + """Convert CSR representation to ``list[list[int]]``."""
77 + return [
78 + self.token_ids[int(self.offsets[i]) : int(self.offsets[i + 1])].tolist()
79 + for i in range(len(self.offsets) - 1)
80 + ]
81 +
82 + @staticmethod
83 + def merge(chunks: Sequence["SamplingMaskLists"]) -> "SamplingMaskLists":
84 + token_ids = np.concatenate([chunk.token_ids for chunk in chunks])
85 + counts = np.concatenate([np.diff(chunk.offsets) for chunk in chunks])
86 + offsets = np.empty(len(counts) + 1, dtype=np.int64)
87 + offsets[0] = 0
88 + np.cumsum(counts, dtype=np.int64, out=offsets[1:])
89 + return SamplingMaskLists(token_ids, offsets)
90 +
91 +
92 class LogprobsTensors(NamedTuple):
93 # [num_reqs x num_generated_tokens, max_num_logprobs + 1]
94 logprob_token_ids: torch.Tensor
@@ -307,6 +346,9 @@ class ModelRunnerOutput:
346 # ``None`` when ``enable_return_routed_experts`` is off.
347 routed_experts: RoutedExpertsLists | None = None
348
349 + # ``None`` when ``return_sampling_mask`` is off.
350 + sampling_masks: SamplingMaskLists | None = None
351 +
352 @staticmethod
353 def with_kv_conn_output_only(
354 kv_connector_output: KVConnectorOutput | None,
vllm/v1/worker/gpu/async_utils.py
+11 -1
@@ -17,7 +17,7 @@ from vllm.v1.outputs import (
17 PoolerOutput,
18 RoutedExpertsTensors,
19 )
20 -from vllm.v1.worker.gpu.sample.output import SamplerOutput
20 +from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
21 from vllm.v1.worker.utils import raise_if_nan_logits
22
23 if TYPE_CHECKING:
@@ -147,6 +147,11 @@ class AsyncOutput(AsyncModelRunnerOutput):
147 if sampler_output.num_nans is not None:
148 self.num_nans = async_copy_to_np(sampler_output.num_nans)
149 self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
150 + self.sampling_mask_tensors: SamplingMaskTensors | None = None
151 + if sampler_output.sampling_mask_tensors is not None:
152 + self.sampling_mask_tensors = (
153 + sampler_output.sampling_mask_tensors.to_cpu_nonblocking()
154 + )
155 self.routed_experts_cpu: RoutedExpertsTensors | None = None
156 if routed_experts is not None:
157 self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
@@ -172,6 +177,11 @@ class AsyncOutput(AsyncModelRunnerOutput):
177 del token_ids[num_tokens:]
178 self.model_runner_output.sampled_token_ids = sampled_token_ids
179
180 + if self.sampling_mask_tensors is not None:
181 + self.model_runner_output.sampling_masks = (
182 + self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np)
183 + )
184 +
185 if self.num_nans is not None:
186 self.model_runner_output.num_nans_in_logits = dict(
187 zip(self.model_runner_output.req_ids, self.num_nans.tolist())
vllm/v1/worker/gpu/model_runner.py
+1
@@ -398,6 +398,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
398 num_speculative_tokens=self.decode_query_len,
399 use_fp64_gumbel=self.model_config.use_fp64_gumbel,
400 reasoning_config=self.vllm_config.reasoning_config,
401 + return_sampling_mask=self.model_config.return_sampling_mask,
402 )
403 custom = self.model_state.custom_sampler(self.sampler)
404
vllm/v1/worker/gpu/sample/output.py
+116 -1
@@ -1,10 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 from dataclasses import dataclass
6 +from typing import NamedTuple
7
8 +import numpy as np
9 import torch
10
7 -from vllm.v1.outputs import LogprobsTensors
11 +from vllm.triton_utils import tl, triton
12 +from vllm.v1.outputs import LogprobsTensors, SamplingMaskLists
13
14
15 @dataclass
@@ -14,3 +19,113 @@ class SamplerOutput:
19 num_nans: torch.Tensor | None
20 num_sampled: torch.Tensor | None
21 num_rejected: torch.Tensor | None = None
22 + sampling_mask_tensors: SamplingMaskTensors | None = None
23 +
24 +
25 +@triton.jit
26 +def _pack_sampling_mask_kernel(
27 + logits_ptr,
28 + logits_row_stride,
29 + logits_col_stride,
30 + num_sampled_tokens_ptr,
31 + packed_mask_ptr,
32 + packed_mask_row_stride,
33 + counts_ptr,
34 + vocab_size,
35 + BLOCK_SIZE: tl.constexpr,
36 +):
37 + req_idx = tl.program_id(0)
38 + is_active = tl.load(num_sampled_tokens_ptr + req_idx) > 0
39 + count = tl.zeros((), dtype=tl.int32)
40 +
41 + for start_idx in range(0, vocab_size, BLOCK_SIZE):
42 + offsets = start_idx + tl.arange(0, BLOCK_SIZE)
43 + valid = offsets < vocab_size
44 + logits = tl.load(
45 + logits_ptr + req_idx * logits_row_stride + offsets * logits_col_stride,
46 + mask=valid,
47 + other=-float("inf"),
48 + )
49 + keep = (logits > -float("inf")) & (logits < float("inf")) & is_active
50 + count += tl.sum(keep).to(tl.int32)
51 +
52 + keep = tl.reshape(keep.to(tl.int32), (BLOCK_SIZE // 8, 8))
53 + bit_shifts = tl.arange(0, 8)[None, :]
54 + packed = tl.sum(keep << bit_shifts, axis=1).to(tl.uint8)
55 + byte_offsets = start_idx // 8 + tl.arange(0, BLOCK_SIZE // 8)
56 + tl.store(
57 + packed_mask_ptr + req_idx * packed_mask_row_stride + byte_offsets,
58 + packed,
59 + mask=byte_offsets < tl.cdiv(vocab_size, 8),
60 + )
61 +
62 + tl.store(counts_ptr + req_idx, count)
63 +
64 +
65 +class SamplingMaskTensors(NamedTuple):
66 + """Bit-packed device-side sampling mask data pending async D2H."""
67 +
68 + # [num_requests, ceil(vocab_size / 8)]
69 + packed_mask: torch.Tensor
70 + # [num_requests]
71 + counts: torch.Tensor
72 + vocab_size: int
73 +
74 + @classmethod
75 + def from_logits(
76 + cls,
77 + logits: torch.Tensor,
78 + num_sampled_tokens: torch.Tensor,
79 + ) -> SamplingMaskTensors:
80 + """Pack the finite-logit support for requests that sampled tokens."""
81 + num_reqs, vocab_size = logits.shape
82 + packed_width = (vocab_size + 7) // 8
83 +
84 + packed_mask = torch.empty(
85 + (num_reqs, packed_width), dtype=torch.uint8, device=logits.device
86 + )
87 + counts = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
88 + _pack_sampling_mask_kernel[(num_reqs,)](
89 + logits,
90 + logits.stride(0),
91 + logits.stride(1),
92 + num_sampled_tokens,
93 + packed_mask,
94 + packed_mask.stride(0),
95 + counts,
96 + vocab_size,
97 + BLOCK_SIZE=8192,
98 + )
99 +
100 + return cls(packed_mask, counts, vocab_size)
101 +
102 + def to_cpu_nonblocking(self) -> SamplingMaskTensors:
103 + if self.packed_mask.device.type == "cpu":
104 + return self
105 + return SamplingMaskTensors(
106 + self.packed_mask.to("cpu", non_blocking=True),
107 + self.counts.to("cpu", non_blocking=True),
108 + self.vocab_size,
109 + )
110 +
111 + def tolists(self, num_sampled_tokens: np.ndarray) -> SamplingMaskLists:
112 + """Convert the packed masks to the scheduler's CSR representation."""
113 + sampled_rows = np.flatnonzero(num_sampled_tokens)
114 + counts = self.counts.cpu().numpy()[sampled_rows]
115 + offsets = np.empty(len(counts) + 1, dtype=np.int64)
116 + offsets[0] = 0
117 + np.cumsum(counts, dtype=np.int64, out=offsets[1:])
118 + unpacked = np.unpackbits(
119 + self.packed_mask.cpu().numpy()[sampled_rows],
120 + axis=1,
121 + count=self.vocab_size,
122 + bitorder="little",
123 + )
124 + token_ids = np.nonzero(unpacked)[1].astype(np.int32, copy=False)
125 + return SamplingMaskLists(
126 + token_ids=token_ids,
127 + offsets=offsets,
128 + cu_num_generated_tokens=np.cumsum(
129 + np.concatenate(([0], num_sampled_tokens))
130 + ).tolist(),
131 + )
vllm/v1/worker/gpu/sample/sampler.py
+13 -2
@@ -22,7 +22,7 @@ from vllm.v1.worker.gpu.sample.logprob import (
22 LogprobTokenIdsState,
23 compute_topk_scores,
24 )
25 -from vllm.v1.worker.gpu.sample.output import SamplerOutput
25 +from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
26 from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
27 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
28 from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
@@ -40,6 +40,7 @@ class Sampler:
40 num_speculative_tokens: int = 1,
41 use_fp64_gumbel: bool = False,
42 reasoning_config: ReasoningConfig | None = None,
43 + return_sampling_mask: bool = False,
44 ):
45 self.logprobs_mode = logprobs_mode
46 self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
@@ -53,7 +54,10 @@ class Sampler:
54 self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
55 self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
56 self.num_speculative_tokens = num_speculative_tokens
56 - self.use_flashinfer = flashinfer_sampler_supported()
57 + self.return_sampling_mask = return_sampling_mask
58 + self.use_flashinfer = (
59 + not return_sampling_mask and flashinfer_sampler_supported()
60 + )
61
62 def add_request(
63 self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
@@ -137,6 +141,12 @@ class Sampler:
141 self.req_states.prefill_len.gpu,
142 )
143
144 + sampling_mask_tensors = None
145 + if self.return_sampling_mask:
146 + sampling_mask_tensors = SamplingMaskTensors.from_logits(
147 + processed_logits, num_sampled
148 + )
149 +
150 # These are GPU tensors.
151 sampler_output = SamplerOutput(
152 # The sampled tokens are expanded to 2D tensor with shape
@@ -147,6 +157,7 @@ class Sampler:
157 num_nans=num_nans,
158 num_sampled=num_sampled,
159 num_rejected=num_rejected,
160 + sampling_mask_tensors=sampling_mask_tensors,
161 )
162 return sampler_output
163