[Perf][Model Runner V2] Compact sampling masks on GPU instead of unpacking the full-vocab bitmask on CPU (#54901)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: vx120 <893600387@qq.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: vx120 <893600387@qq.com>
aoshen02 committed
Sep 4, 2026 at 10:03 UTC
560ef78bfe734ea894f8a866e50406574021b9f5
10 files changed
+187
-107
tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+7
-1
@@ -137,6 +137,7 @@ async def test_generate_sampling_mask(client):
137
"top_p": 0.9,
138
"ignore_eos": True,
139
"seed": 0,
140
+ "logprobs": top_k,
141
},
142
"stream": False,
143
}
@@ -150,11 +151,16 @@ async def test_generate_sampling_mask(client):
151
assert len(token_ids) == len(sampling_mask)
152
153
vocab_size = get_vocab_size(MODEL_NAME)
153
- for token_id, support in zip(token_ids, sampling_mask):
154
+ logprobs_content = choice["logprobs"]["content"]
155
+ for token_id, support, entry in zip(token_ids, sampling_mask, logprobs_content):
156
assert support
157
assert len(support) == len(set(support))
158
assert all(0 <= support_token_id < vocab_size for support_token_id in support)
159
assert token_id in support
160
+ # processed_logprobs: exactly the support carries finite probability.
161
+ for top in entry["top_logprobs"]:
162
+ in_support = int(top["token"].removeprefix("token_id:")) in support
163
+ assert in_support == (top["logprob"] > -9999.0)
164
165
166
@pytest.mark.asyncio
tests/test_config.py
+16
@@ -753,6 +753,22 @@ def test_v1_model_runner_rejects_v2_only_features():
753
VllmConfig._validate_v1_model_runner(config)
754
755
756
+def test_batch_sharded_sampling_rejects_return_sampling_mask():
757
+ """The batch-sharded gather drops sampling masks, so the combination must
758
+ fail loudly instead of returning ``sampling_mask=None``."""
759
+ config = SimpleNamespace(
760
+ parallel_config=SimpleNamespace(
761
+ enable_batch_sharded_sampling=True, tensor_parallel_size=2
762
+ ),
763
+ scheduler_config=SimpleNamespace(max_num_seqs=8),
764
+ model_config=SimpleNamespace(max_logprobs=20, return_sampling_mask=True),
765
+ speculative_config=None,
766
+ )
767
+
768
+ with pytest.raises(ValueError, match="sampling masks"):
769
+ VllmConfig._validate_batch_sharded_sampling(config)
770
+
771
+
772
@pytest.mark.skip_global_cleanup
773
def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping(
774
monkeypatch,
tests/v1/core/test_scheduler.py
+47
@@ -4,6 +4,7 @@ import dataclasses
4
from concurrent.futures import Future
5
from unittest.mock import Mock
6
7
+import numpy as np
8
import pytest
9
import torch
10
@@ -43,6 +44,7 @@ from vllm.v1.outputs import (
44
ECConnectorOutput,
45
KVConnectorOutput,
46
ModelRunnerOutput,
47
+ SamplingMaskLists,
48
make_empty_encoder_model_runner_output,
49
)
50
from vllm.v1.request import Request, RequestStatus
@@ -667,6 +669,51 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
669
assert output2.num_scheduled_tokens[requests[2].request_id] == 800 - 224 - 224
670
671
672
+def test_update_from_output_routes_sampling_masks_by_request():
673
+ """Each request receives the sampler row at its own batch index."""
674
+ scheduler = create_scheduler()
675
+ scheduler.return_sampling_mask = True
676
+ requests = create_requests(num_requests=3, max_tokens=10)
677
+ for req in requests:
678
+ req.num_computed_tokens = req.num_tokens
679
+ scheduler.requests[req.request_id] = req
680
+ scheduler.running.append(req)
681
+ req.status = RequestStatus.RUNNING
682
+
683
+ scheduler_output = SchedulerOutput(
684
+ scheduled_new_reqs=[],
685
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
686
+ num_scheduled_tokens={req.request_id: 1 for req in requests},
687
+ total_num_scheduled_tokens=3,
688
+ scheduled_encoder_inputs={},
689
+ scheduled_spec_decode_tokens={},
690
+ num_common_prefix_blocks=[],
691
+ finished_req_ids=set(),
692
+ free_encoder_mm_hashes=[],
693
+ )
694
+ model_output = ModelRunnerOutput(
695
+ req_ids=[req.request_id for req in requests],
696
+ req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
697
+ sampled_token_ids=[[1], [3], [4]],
698
+ logprobs=None,
699
+ prompt_logprobs_dict={},
700
+ pooler_output=[],
701
+ sampling_masks=SamplingMaskLists(
702
+ np.array([1, 2, 3, 4, 5, 6], dtype=np.int32), np.array([0, 2, 3, 6])
703
+ ),
704
+ )
705
+
706
+ outputs = scheduler.update_from_output(scheduler_output, model_output)[0].outputs
707
+
708
+ assert [out.request_id for out in outputs] == [req.request_id for req in requests]
709
+ assert [out.new_sampling_mask.token_ids.tolist() for out in outputs] == [
710
+ [1, 2],
711
+ [3],
712
+ [4, 5, 6],
713
+ ]
714
+ assert all(out.new_sampling_mask.offsets is None for out in outputs)
715
+
716
+
717
def test_stop_via_update_from_output():
718
"""Test stopping behavior through update_from_output"""
719
scheduler = create_scheduler(num_speculative_tokens=1)
tests/v1/test_outputs.py
+31
-37
@@ -3,6 +3,7 @@
3
from unittest import TestCase
4
5
import numpy as np
6
+import pytest
7
import torch
8
9
from vllm.platforms import current_platform
@@ -61,23 +62,6 @@ def test_logprobs_tensors_tolists_with_tensor_boundaries():
62
assert sliced.sampled_token_ranks.tolist() == [3]
63
64
64
-def test_sampling_mask_tensors_tolist():
65
- tensors = SamplingMaskTensors(
66
- packed_mask=torch.tensor(
67
- [[0b00101000], [0b00000000], [0b10000000]],
68
- dtype=torch.uint8,
69
- ),
70
- counts=torch.tensor([2, 0, 1], dtype=torch.int32),
71
- vocab_size=8,
72
- )
73
-
74
- result = tensors.tolists(np.array([1, 0, 1]))
75
-
76
- assert result.token_ids.tolist() == [3, 5, 7]
77
- assert result.offsets.tolist() == [0, 2, 3]
78
- assert result.cu_num_generated_tokens == [0, 1, 1, 2]
79
-
80
-
65
def test_sampling_mask_lists_to_nested_list():
66
from vllm.v1.outputs import SamplingMaskLists
67
@@ -89,46 +73,56 @@ def test_sampling_mask_lists_to_nested_list():
73
nested = mask.to_nested_list()
74
75
assert nested == [[10, 11, 12], [20, 21]]
76
+ assert SamplingMaskLists(np.array([7, 9])).to_nested_list() == [[7, 9]]
77
+
78
+
79
+@pytest.mark.parametrize("max_num_kept", [512, 20_001])
80
+def test_sampling_mask_tensors_match_finite_support(max_num_kept):
81
+ """Whatever its size (empty, one, around the compact width, beyond the
82
+ cap, the whole vocab), a sampled row's mask is exactly the ascending set
83
+ of its finite logits; rows that sampled nothing are empty."""
84
+ from vllm.v1.worker.gpu.sample.output import MAX_COMPACT_SUPPORT
85
+
86
+ vocab_size = 20_001
87
+ sizes = [0, 1, 7, 511, 512, 513, MAX_COMPACT_SUPPORT + 5, vocab_size, 40, 3]
88
+ gen = torch.Generator().manual_seed(0)
89
+ logits = torch.full((len(sizes), vocab_size), float("-inf"))
90
+ expected = []
91
+ for row, size in enumerate(sizes):
92
+ kept = torch.randperm(vocab_size, generator=gen)[:size].sort().values
93
+ logits[row, kept] = torch.randn(size, generator=gen)
94
+ expected.append(kept.tolist())
95
+ num_sampled_tokens = torch.ones(len(sizes), dtype=torch.int32)
96
+ num_sampled_tokens[[1, 7]] = 0
97
+ expected[1] = expected[7] = []
98
93
-
94
-def test_sampling_mask_tensors_from_logits():
99
tensors = SamplingMaskTensors.from_logits(
96
- logits=torch.tensor(
97
- [
98
- [1.0, float("-inf"), 2.0],
99
- [3.0, 4.0, float("-inf")],
100
- [float("-inf"), 5.0, 6.0],
101
- ],
102
- device=DEVICE_TYPE,
103
- ),
104
- num_sampled_tokens=torch.tensor([1, 0, 1], device=DEVICE_TYPE),
100
+ logits.to(DEVICE_TYPE), num_sampled_tokens.to(DEVICE_TYPE), max_num_kept
101
)
102
107
- result = tensors.tolists(np.array([1, 0, 1]))
108
-
109
- assert result.token_ids.tolist() == [0, 2, 1, 2]
110
- assert result.offsets.tolist() == [0, 2, 4]
111
- assert result.cu_num_generated_tokens == [0, 1, 1, 2]
103
+ assert tensors.token_ids.shape[1] == min(max_num_kept, MAX_COMPACT_SUPPORT)
104
+ assert tensors.to_cpu_nonblocking().tolists().to_nested_list() == expected
105
106
114
-def test_sampling_mask_matches_processed_top_k_top_p_support():
107
+def test_sampling_mask_preserves_top_k_boundary_ties():
108
processed_logits = apply_top_k_top_p(
109
logits=torch.tensor(
110
[[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE
111
),
112
k=torch.tensor([3], device=DEVICE_TYPE),
120
- p=torch.tensor([0.9], device=DEVICE_TYPE),
113
+ p=None,
114
)
115
expected_token_ids = (
116
torch.isfinite(processed_logits[0]).nonzero().flatten().tolist()
117
)
125
- assert 0 < len(expected_token_ids) < processed_logits.shape[1]
118
+ assert len(expected_token_ids) > 3
119
120
tensors = SamplingMaskTensors.from_logits(
121
processed_logits,
122
num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),
123
+ max_num_kept=3,
124
)
131
- result = tensors.tolists(np.array([1]))
125
+ result = tensors.tolists()
126
127
assert result.to_nested_list() == [expected_token_ids]
128
vllm/config/vllm.py
+7
@@ -2709,6 +2709,13 @@ class VllmConfig:
2709
# fixed-width logprobs gather cannot reasonably size for.
2710
blockers.append("max_logprobs is -1, allowing vocab-size logprob requests")
2711
2712
+ if self.model_config is not None and self.model_config.return_sampling_mask:
2713
+ # gather_sampler_output() drops SamplingMaskTensors: masks come back None.
2714
+ blockers.append(
2715
+ "return_sampling_mask is set and the batch-sharded gather does "
2716
+ "not forward sampling masks"
2717
+ )
2718
+
2719
if (
2720
self.speculative_config is not None
2721
and self.speculative_config.enable_adaptive_verification
vllm/v1/engine/output_processor.py
+3
-2
@@ -419,8 +419,9 @@ class RequestState:
419
420
sampling_mask = None
421
if finished and self.sampling_mask_chunks:
422
- merged = SamplingMaskLists.merge(self.sampling_mask_chunks)
423
- sampling_mask = SamplingMask(merged.to_nested_list())
422
+ sampling_mask = SamplingMask(
423
+ [chunk.token_ids.tolist() for chunk in self.sampling_mask_chunks]
424
+ )
425
426
# Concatenate routed experts on finish
427
routed_experts = None
vllm/v1/outputs.py
+12
-27
@@ -54,42 +54,27 @@ class LogprobsLists(NamedTuple):
54
55
56
class SamplingMaskLists(NamedTuple):
57
+ """CSR sampling masks; a step slice holds one position (``offsets=None``)."""
58
+
59
# [num_kept_tokens]
60
token_ids: np.ndarray
59
- # [num_generated_tokens + 1]
60
- offsets: np.ndarray
61
- # [num_reqs + 1]
61
+ # [num_positions + 1], or None for a single position
62
+ offsets: np.ndarray | None = None
63
+ # Unused with one position per request; kept for the wire layout.
64
cu_num_generated_tokens: list[int] | None = None
65
66
def slice_request(self, req_idx: int, num_positions: int) -> "SamplingMaskLists":
65
- if self.cu_num_generated_tokens is None:
66
- start_idx = req_idx
67
- else:
68
- start_idx = self.cu_num_generated_tokens[req_idx]
69
- end_idx = start_idx + num_positions
70
- flat_start = self.offsets[start_idx]
71
- flat_end = self.offsets[end_idx]
67
+ assert num_positions == 1 and self.offsets is not None
68
return SamplingMaskLists(
73
- self.token_ids[flat_start:flat_end],
74
- self.offsets[start_idx : end_idx + 1] - flat_start,
75
- None,
69
+ self.token_ids[self.offsets[req_idx] : self.offsets[req_idx + 1]]
70
)
71
72
def to_nested_list(self) -> list[list[int]]:
79
- """Convert CSR representation to ``list[list[int]]``."""
80
- return [
81
- self.token_ids[int(self.offsets[i]) : int(self.offsets[i + 1])].tolist()
82
- for i in range(len(self.offsets) - 1)
83
- ]
84
-
85
- @staticmethod
86
- def merge(chunks: Sequence["SamplingMaskLists"]) -> "SamplingMaskLists":
87
- token_ids = np.concatenate([chunk.token_ids for chunk in chunks])
88
- counts = np.concatenate([np.diff(chunk.offsets) for chunk in chunks])
89
- offsets = np.empty(len(counts) + 1, dtype=np.int64)
90
- offsets[0] = 0
91
- np.cumsum(counts, dtype=np.int64, out=offsets[1:])
92
- return SamplingMaskLists(token_ids, offsets)
73
+ token_ids = self.token_ids.tolist()
74
+ if self.offsets is None:
75
+ return [token_ids]
76
+ offsets = self.offsets.tolist()
77
+ return [token_ids[offsets[i] : offsets[i + 1]] for i in range(len(offsets) - 1)]
78
79
80
class LogprobsTensors(NamedTuple):
vllm/v1/worker/gpu/async_utils.py
+1
-1
@@ -179,7 +179,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
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)
182
+ self.sampling_mask_tensors.tolists()
183
)
184
185
if self.num_nans is not None:
vllm/v1/worker/gpu/sample/output.py
+60
-38
@@ -23,48 +23,63 @@ class SamplerOutput:
23
24
25
@triton.jit
26
-def _pack_sampling_mask_kernel(
26
+def _compact_sampling_mask_kernel(
27
logits_ptr,
28
logits_row_stride,
29
logits_col_stride,
30
num_sampled_tokens_ptr,
31
+ token_ids_ptr,
32
+ token_ids_row_stride,
33
packed_mask_ptr,
34
packed_mask_row_stride,
35
counts_ptr,
36
vocab_size,
37
+ max_num_kept,
38
BLOCK_SIZE: tl.constexpr,
39
):
40
+ """Per row: first ``max_num_kept`` finite-logit ids, the full count, the bitmask."""
41
req_idx = tl.program_id(0)
42
is_active = tl.load(num_sampled_tokens_ptr + req_idx) > 0
43
count = tl.zeros((), dtype=tl.int32)
44
45
for start_idx in range(0, vocab_size, BLOCK_SIZE):
46
offsets = start_idx + tl.arange(0, BLOCK_SIZE)
43
- valid = offsets < vocab_size
47
logits = tl.load(
48
logits_ptr + req_idx * logits_row_stride + offsets * logits_col_stride,
46
- mask=valid,
49
+ mask=offsets < vocab_size,
50
other=-float("inf"),
51
)
52
keep = (logits > -float("inf")) & (logits < float("inf")) & is_active
50
- count += tl.sum(keep).to(tl.int32)
53
+ keep_i32 = keep.to(tl.int32)
54
+ pos = count + tl.cumsum(keep_i32, axis=0) - keep_i32
55
+ tl.store(
56
+ token_ids_ptr + req_idx * token_ids_row_stride + pos,
57
+ offsets.to(tl.int32),
58
+ mask=keep & (pos < max_num_kept),
59
+ )
60
+ count += tl.sum(keep_i32)
61
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)
62
+ bits = tl.reshape(keep_i32, (BLOCK_SIZE // 8, 8)) << tl.arange(0, 8)[None, :]
63
byte_offsets = start_idx // 8 + tl.arange(0, BLOCK_SIZE // 8)
64
tl.store(
65
packed_mask_ptr + req_idx * packed_mask_row_stride + byte_offsets,
58
- packed,
66
+ tl.sum(bits, axis=1).to(tl.uint8),
67
mask=byte_offsets < tl.cdiv(vocab_size, 8),
68
)
69
70
tl.store(counts_ptr + req_idx, count)
71
72
73
+# Bounds the [num_reqs, width] int32 buffer; wider rows use the bitmask instead.
74
+MAX_COMPACT_SUPPORT = 2048
75
+
76
+
77
class SamplingMaskTensors(NamedTuple):
66
- """Bit-packed device-side sampling mask data pending async D2H."""
78
+ """Device-side masks pending async D2H: compact ids, plus the bitmask as
79
+ the exact fallback for rows wider than ``max_num_kept``."""
80
81
+ # [num_requests, max_num_kept]
82
+ token_ids: torch.Tensor
83
# [num_requests, ceil(vocab_size / 8)]
84
packed_mask: torch.Tensor
85
# [num_requests]
@@ -76,56 +91,63 @@ class SamplingMaskTensors(NamedTuple):
91
cls,
92
logits: torch.Tensor,
93
num_sampled_tokens: torch.Tensor,
94
+ max_num_kept: int,
95
) -> SamplingMaskTensors:
80
- """Pack the finite-logit support for requests that sampled tokens."""
96
+ """Capture the finite-logit support of every row with a sampled token."""
97
num_reqs, vocab_size = logits.shape
82
- packed_width = (vocab_size + 7) // 8
98
+ max_num_kept = min(max_num_kept, vocab_size, MAX_COMPACT_SUPPORT)
99
+ device = logits.device
100
101
+ token_ids = torch.empty(
102
+ (num_reqs, max_num_kept), dtype=torch.int32, device=device
103
+ )
104
packed_mask = torch.empty(
85
- (num_reqs, packed_width), dtype=torch.uint8, device=logits.device
105
+ (num_reqs, (vocab_size + 7) // 8), dtype=torch.uint8, device=device
106
)
87
- counts = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
88
- _pack_sampling_mask_kernel[(num_reqs,)](
107
+ counts = torch.empty(num_reqs, dtype=torch.int32, device=device)
108
+ _compact_sampling_mask_kernel[(num_reqs,)](
109
logits,
110
logits.stride(0),
111
logits.stride(1),
112
num_sampled_tokens,
113
+ token_ids,
114
+ token_ids.stride(0),
115
packed_mask,
116
packed_mask.stride(0),
117
counts,
118
vocab_size,
119
+ max_num_kept,
120
BLOCK_SIZE=8192,
121
)
99
-
100
- return cls(packed_mask, counts, vocab_size)
122
+ return cls(token_ids, packed_mask, counts, vocab_size)
123
124
def to_cpu_nonblocking(self) -> SamplingMaskTensors:
103
- if self.packed_mask.device.type == "cpu":
125
+ if self.token_ids.device.type == "cpu":
126
return self
127
return SamplingMaskTensors(
128
+ self.token_ids.to("cpu", non_blocking=True),
129
self.packed_mask.to("cpu", non_blocking=True),
130
self.counts.to("cpu", non_blocking=True),
131
self.vocab_size,
132
)
133
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
- )
134
+ def tolists(self) -> SamplingMaskLists:
135
+ """CSR over all requests; rows without a sampled token are empty."""
136
+ counts = self.counts.cpu().numpy()
137
+ token_ids = self.token_ids.cpu().numpy()
138
+ packed_mask = self.packed_mask.cpu().numpy()
139
+ width = token_ids.shape[1]
140
+
141
+ def support(row: int) -> np.ndarray:
142
+ if counts[row] <= width:
143
+ return token_ids[row, : counts[row]]
144
+ # Wider than the compact row (ties or a huge top_k): use the bitmask.
145
+ bits = np.unpackbits(
146
+ packed_mask[row], count=self.vocab_size, bitorder="little"
147
+ )
148
+ return np.flatnonzero(bits).astype(np.int32, copy=False)
149
+
150
+ supports = [support(row) for row in range(len(counts))]
151
+ offsets = np.zeros(len(supports) + 1, dtype=np.int64)
152
+ np.cumsum([len(s) for s in supports], out=offsets[1:])
153
+ return SamplingMaskLists(np.concatenate(supports), offsets)
vllm/v1/worker/gpu/sample/sampler.py
+3
-1
@@ -186,8 +186,10 @@ class Sampler:
186
187
sampling_mask_tensors = None
188
if self.return_sampling_mask:
189
+ # Size by the validated top_k batch max; wider supports use the bitmask.
190
+ max_num_kept = int(np.max(self.sampling_states.top_k.np[idx_mapping_np]))
191
sampling_mask_tensors = SamplingMaskTensors.from_logits(
190
- processed_logits, num_sampled
192
+ processed_logits, num_sampled, max_num_kept
193
)
194
195
# These are GPU tensors.