iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][Model Runner V2] Compact sampling masks on GPU instead of unpacking the full-vocab bitmask on CPU (#54901)

Browse Source 
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>
This commit is contained in:
aoshen02
2026-09-03 19:03:36 -07:00
committed by GitHub
co-authored by Claude Fable 5.1 vx120
parent d6bce42983
commit 560ef78bfe
10 changed files with 188 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
+7 -1
View File
@@ -137,6 +137,7 @@ async def test_generate_sampling_mask(client):
"top_p": 0.9,
"ignore_eos": True,
"seed": 0,
"logprobs": top_k,
},
"stream": False,
}
@@ -150,11 +151,16 @@ async def test_generate_sampling_mask(client):
assert len(token_ids) == len(sampling_mask)
vocab_size = get_vocab_size(MODEL_NAME)
for token_id, support in zip(token_ids, sampling_mask):
logprobs_content = choice["logprobs"]["content"]
for token_id, support, entry in zip(token_ids, sampling_mask, logprobs_content):
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
# processed_logprobs: exactly the support carries finite probability.
for top in entry["top_logprobs"]:
in_support = int(top["token"].removeprefix("token_id:")) in support
assert in_support == (top["logprob"] > -9999.0)
@pytest.mark.asyncio
tests/test_config.py
+16
View File
@@ -753,6 +753,22 @@ def test_v1_model_runner_rejects_v2_only_features():
VllmConfig._validate_v1_model_runner(config)
def test_batch_sharded_sampling_rejects_return_sampling_mask():
"""The batch-sharded gather drops sampling masks, so the combination must
fail loudly instead of returning ``sampling_mask=None``."""
config = SimpleNamespace(
parallel_config=SimpleNamespace(
enable_batch_sharded_sampling=True, tensor_parallel_size=2
),
scheduler_config=SimpleNamespace(max_num_seqs=8),
model_config=SimpleNamespace(max_logprobs=20, return_sampling_mask=True),
speculative_config=None,
)
with pytest.raises(ValueError, match="sampling masks"):
VllmConfig._validate_batch_sharded_sampling(config)
@pytest.mark.skip_global_cleanup
def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping(
monkeypatch,
tests/v1/core/test_scheduler.py
+47
View File
@@ -4,6 +4,7 @@ import dataclasses
from concurrent.futures import Future
from unittest.mock import Mock
import numpy as np
import pytest
import torch
@@ -43,6 +44,7 @@ from vllm.v1.outputs import (
ECConnectorOutput,
KVConnectorOutput,
ModelRunnerOutput,
SamplingMaskLists,
make_empty_encoder_model_runner_output,
)
from vllm.v1.request import Request, RequestStatus
@@ -667,6 +669,51 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
assert output2.num_scheduled_tokens[requests[2].request_id] == 800 - 224 - 224
def test_update_from_output_routes_sampling_masks_by_request():
"""Each request receives the sampler row at its own batch index."""
scheduler = create_scheduler()
scheduler.return_sampling_mask = True
requests = create_requests(num_requests=3, max_tokens=10)
for req in requests:
req.num_computed_tokens = req.num_tokens
scheduler.requests[req.request_id] = req
scheduler.running.append(req)
req.status = RequestStatus.RUNNING
scheduler_output = SchedulerOutput(
scheduled_new_reqs=[],
scheduled_cached_reqs=CachedRequestData.make_empty(),
num_scheduled_tokens={req.request_id: 1 for req in requests},
total_num_scheduled_tokens=3,
scheduled_encoder_inputs={},
scheduled_spec_decode_tokens={},
num_common_prefix_blocks=[],
finished_req_ids=set(),
free_encoder_mm_hashes=[],
)
model_output = ModelRunnerOutput(
req_ids=[req.request_id for req in requests],
req_id_to_index={req.request_id: i for i, req in enumerate(requests)},
sampled_token_ids=[[1], [3], [4]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[],
sampling_masks=SamplingMaskLists(
np.array([1, 2, 3, 4, 5, 6], dtype=np.int32), np.array([0, 2, 3, 6])
),
)
outputs = scheduler.update_from_output(scheduler_output, model_output)[0].outputs
assert [out.request_id for out in outputs] == [req.request_id for req in requests]
assert [out.new_sampling_mask.token_ids.tolist() for out in outputs] == [
[1, 2],
[3],
[4, 5, 6],
]
assert all(out.new_sampling_mask.offsets is None for out in outputs)
def test_stop_via_update_from_output():
"""Test stopping behavior through update_from_output"""
scheduler = create_scheduler(num_speculative_tokens=1)
tests/v1/test_outputs.py
+30 -36
View File
@@ -3,6 +3,7 @@
from unittest import TestCase
import numpy as np
import pytest
import torch
from vllm.platforms import current_platform
@@ -61,23 +62,6 @@ def test_logprobs_tensors_tolists_with_tensor_boundaries():
assert sliced.sampled_token_ranks.tolist() == [3]
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
@@ -89,46 +73,56 @@ def test_sampling_mask_lists_to_nested_list():
nested = mask.to_nested_list()
assert nested == [[10, 11, 12], [20, 21]]
assert SamplingMaskLists(np.array([7, 9])).to_nested_list() == [[7, 9]]
def test_sampling_mask_tensors_from_logits():
@pytest.mark.parametrize("max_num_kept", [512, 20_001])
def test_sampling_mask_tensors_match_finite_support(max_num_kept):
"""Whatever its size (empty, one, around the compact width, beyond the
cap, the whole vocab), a sampled row's mask is exactly the ascending set
of its finite logits; rows that sampled nothing are empty."""
from vllm.v1.worker.gpu.sample.output import MAX_COMPACT_SUPPORT
vocab_size = 20_001
sizes = [0, 1, 7, 511, 512, 513, MAX_COMPACT_SUPPORT + 5, vocab_size, 40, 3]
gen = torch.Generator().manual_seed(0)
logits = torch.full((len(sizes), vocab_size), float("-inf"))
expected = []
for row, size in enumerate(sizes):
kept = torch.randperm(vocab_size, generator=gen)[:size].sort().values
logits[row, kept] = torch.randn(size, generator=gen)
expected.append(kept.tolist())
num_sampled_tokens = torch.ones(len(sizes), dtype=torch.int32)
num_sampled_tokens[[1, 7]] = 0
expected[1] = expected[7] = []
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=DEVICE_TYPE,
),
num_sampled_tokens=torch.tensor([1, 0, 1], device=DEVICE_TYPE),
logits.to(DEVICE_TYPE), num_sampled_tokens.to(DEVICE_TYPE), max_num_kept
)
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]
assert tensors.token_ids.shape[1] == min(max_num_kept, MAX_COMPACT_SUPPORT)
assert tensors.to_cpu_nonblocking().tolists().to_nested_list() == expected
def test_sampling_mask_matches_processed_top_k_top_p_support():
def test_sampling_mask_preserves_top_k_boundary_ties():
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=DEVICE_TYPE
),
k=torch.tensor([3], device=DEVICE_TYPE),
p=torch.tensor([0.9], device=DEVICE_TYPE),
p=None,
)
expected_token_ids = (
torch.isfinite(processed_logits[0]).nonzero().flatten().tolist()
)
assert 0 < len(expected_token_ids) < processed_logits.shape[1]
assert len(expected_token_ids) > 3
tensors = SamplingMaskTensors.from_logits(
processed_logits,
num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE),
max_num_kept=3,
)
result = tensors.tolists(np.array([1]))
result = tensors.tolists()
assert result.to_nested_list() == [expected_token_ids]
vllm/config/vllm.py
+7
View File
@@ -2709,6 +2709,13 @@ class VllmConfig:
# fixed-width logprobs gather cannot reasonably size for.
blockers.append("max_logprobs is -1, allowing vocab-size logprob requests")
if self.model_config is not None and self.model_config.return_sampling_mask:
# gather_sampler_output() drops SamplingMaskTensors: masks come back None.
blockers.append(
"return_sampling_mask is set and the batch-sharded gather does "
"not forward sampling masks"
)
if (
self.speculative_config is not None
and self.speculative_config.enable_adaptive_verification
vllm/v1/engine/output_processor.py
+3 -2
View File
@@ -419,8 +419,9 @@ class RequestState:
sampling_mask = None
if finished and self.sampling_mask_chunks:
merged = SamplingMaskLists.merge(self.sampling_mask_chunks)
sampling_mask = SamplingMask(merged.to_nested_list())
sampling_mask = SamplingMask(
[chunk.token_ids.tolist() for chunk in self.sampling_mask_chunks]
)
# Concatenate routed experts on finish
routed_experts = None
vllm/v1/outputs.py
+12 -27
View File
@@ -54,42 +54,27 @@ class LogprobsLists(NamedTuple):
class SamplingMaskLists(NamedTuple):
"""CSR sampling masks; a step slice holds one position (``offsets=None``)."""
# [num_kept_tokens]
token_ids: np.ndarray
# [num_generated_tokens + 1]
offsets: np.ndarray
# [num_reqs + 1]
# [num_positions + 1], or None for a single position
offsets: np.ndarray | None = None
# Unused with one position per request; kept for the wire layout.
cu_num_generated_tokens: list[int] | None = None
def slice_request(self, req_idx: int, num_positions: int) -> "SamplingMaskLists":
if self.cu_num_generated_tokens is None:
start_idx = req_idx
else:
start_idx = self.cu_num_generated_tokens[req_idx]
end_idx = start_idx + num_positions
flat_start = self.offsets[start_idx]
flat_end = self.offsets[end_idx]
assert num_positions == 1 and self.offsets is not None
return SamplingMaskLists(
self.token_ids[flat_start:flat_end],
self.offsets[start_idx : end_idx + 1] - flat_start,
None,
self.token_ids[self.offsets[req_idx] : self.offsets[req_idx + 1]]
)
def to_nested_list(self) -> list[list[int]]:
"""Convert CSR representation to ``list[list[int]]``."""
return [
self.token_ids[int(self.offsets[i]) : int(self.offsets[i + 1])].tolist()
for i in range(len(self.offsets) - 1)
]
@staticmethod
def merge(chunks: Sequence["SamplingMaskLists"]) -> "SamplingMaskLists":
token_ids = np.concatenate([chunk.token_ids for chunk in chunks])
counts = np.concatenate([np.diff(chunk.offsets) for chunk in chunks])
offsets = np.empty(len(counts) + 1, dtype=np.int64)
offsets[0] = 0
np.cumsum(counts, dtype=np.int64, out=offsets[1:])
return SamplingMaskLists(token_ids, offsets)
token_ids = self.token_ids.tolist()
if self.offsets is None:
return [token_ids]
offsets = self.offsets.tolist()
return [token_ids[offsets[i] : offsets[i + 1]] for i in range(len(offsets) - 1)]
class LogprobsTensors(NamedTuple):
vllm/v1/worker/gpu/async_utils.py
+1 -1
View File
@@ -179,7 +179,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
if self.sampling_mask_tensors is not None:
self.model_runner_output.sampling_masks = (
self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np)
self.sampling_mask_tensors.tolists()
)
if self.num_nans is not None:
vllm/v1/worker/gpu/sample/output.py
+62 -40
View File
@@ -23,48 +23,63 @@ class SamplerOutput:
@triton.jit
def _pack_sampling_mask_kernel(
def _compact_sampling_mask_kernel(
logits_ptr,
logits_row_stride,
logits_col_stride,
num_sampled_tokens_ptr,
token_ids_ptr,
token_ids_row_stride,
packed_mask_ptr,
packed_mask_row_stride,
counts_ptr,
vocab_size,
max_num_kept,
BLOCK_SIZE: tl.constexpr,
):
"""Per row: first ``max_num_kept`` finite-logit ids, the full count, the bitmask."""
req_idx = tl.program_id(0)
is_active = tl.load(num_sampled_tokens_ptr + req_idx) > 0
count = tl.zeros((), dtype=tl.int32)
for start_idx in range(0, vocab_size, BLOCK_SIZE):
offsets = start_idx + tl.arange(0, BLOCK_SIZE)
valid = offsets < vocab_size
logits = tl.load(
logits_ptr + req_idx * logits_row_stride + offsets * logits_col_stride,
mask=valid,
mask=offsets < vocab_size,
other=-float("inf"),
)
keep = (logits > -float("inf")) & (logits < float("inf")) & is_active
count += tl.sum(keep).to(tl.int32)
keep_i32 = keep.to(tl.int32)
pos = count + tl.cumsum(keep_i32, axis=0) - keep_i32
tl.store(
token_ids_ptr + req_idx * token_ids_row_stride + pos,
offsets.to(tl.int32),
mask=keep & (pos < max_num_kept),
)
count += tl.sum(keep_i32)
keep = tl.reshape(keep.to(tl.int32), (BLOCK_SIZE // 8, 8))
bit_shifts = tl.arange(0, 8)[None, :]
packed = tl.sum(keep << bit_shifts, axis=1).to(tl.uint8)
bits = tl.reshape(keep_i32, (BLOCK_SIZE // 8, 8)) << tl.arange(0, 8)[None, :]
byte_offsets = start_idx // 8 + tl.arange(0, BLOCK_SIZE // 8)
tl.store(
packed_mask_ptr + req_idx * packed_mask_row_stride + byte_offsets,
packed,
tl.sum(bits, axis=1).to(tl.uint8),
mask=byte_offsets < tl.cdiv(vocab_size, 8),
)
tl.store(counts_ptr + req_idx, count)
class SamplingMaskTensors(NamedTuple):
"""Bit-packed device-side sampling mask data pending async D2H."""
# Bounds the [num_reqs, width] int32 buffer; wider rows use the bitmask instead.
MAX_COMPACT_SUPPORT = 2048
class SamplingMaskTensors(NamedTuple):
"""Device-side masks pending async D2H: compact ids, plus the bitmask as
the exact fallback for rows wider than ``max_num_kept``."""
# [num_requests, max_num_kept]
token_ids: torch.Tensor
# [num_requests, ceil(vocab_size / 8)]
packed_mask: torch.Tensor
# [num_requests]
@@ -76,56 +91,63 @@ class SamplingMaskTensors(NamedTuple):
cls,
logits: torch.Tensor,
num_sampled_tokens: torch.Tensor,
max_num_kept: int,
) -> SamplingMaskTensors:
"""Pack the finite-logit support for requests that sampled tokens."""
"""Capture the finite-logit support of every row with a sampled token."""
num_reqs, vocab_size = logits.shape
packed_width = (vocab_size + 7) // 8
max_num_kept = min(max_num_kept, vocab_size, MAX_COMPACT_SUPPORT)
device = logits.device
packed_mask = torch.empty(
(num_reqs, packed_width), dtype=torch.uint8, device=logits.device
token_ids = torch.empty(
(num_reqs, max_num_kept), dtype=torch.int32, device=device
)
counts = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
_pack_sampling_mask_kernel[(num_reqs,)](
packed_mask = torch.empty(
(num_reqs, (vocab_size + 7) // 8), dtype=torch.uint8, device=device
)
counts = torch.empty(num_reqs, dtype=torch.int32, device=device)
_compact_sampling_mask_kernel[(num_reqs,)](
logits,
logits.stride(0),
logits.stride(1),
num_sampled_tokens,
token_ids,
token_ids.stride(0),
packed_mask,
packed_mask.stride(0),
counts,
vocab_size,
max_num_kept,
BLOCK_SIZE=8192,
)
return cls(packed_mask, counts, vocab_size)
return cls(token_ids, packed_mask, counts, vocab_size)
def to_cpu_nonblocking(self) -> SamplingMaskTensors:
if self.packed_mask.device.type == "cpu":
if self.token_ids.device.type == "cpu":
return self
return SamplingMaskTensors(
self.token_ids.to("cpu", non_blocking=True),
self.packed_mask.to("cpu", non_blocking=True),
self.counts.to("cpu", non_blocking=True),
self.vocab_size,
)
def tolists(self, num_sampled_tokens: np.ndarray) -> SamplingMaskLists:
"""Convert the packed masks to the scheduler's CSR representation."""
sampled_rows = np.flatnonzero(num_sampled_tokens)
counts = self.counts.cpu().numpy()[sampled_rows]
offsets = np.empty(len(counts) + 1, dtype=np.int64)
offsets[0] = 0
np.cumsum(counts, dtype=np.int64, out=offsets[1:])
unpacked = np.unpackbits(
self.packed_mask.cpu().numpy()[sampled_rows],
axis=1,
count=self.vocab_size,
bitorder="little",
)
token_ids = np.nonzero(unpacked)[1].astype(np.int32, copy=False)
return SamplingMaskLists(
token_ids=token_ids,
offsets=offsets,
cu_num_generated_tokens=np.cumsum(
np.concatenate(([0], num_sampled_tokens))
).tolist(),
)
def tolists(self) -> SamplingMaskLists:
"""CSR over all requests; rows without a sampled token are empty."""
counts = self.counts.cpu().numpy()
token_ids = self.token_ids.cpu().numpy()
packed_mask = self.packed_mask.cpu().numpy()
width = token_ids.shape[1]
def support(row: int) -> np.ndarray:
if counts[row] <= width:
return token_ids[row, : counts[row]]
# Wider than the compact row (ties or a huge top_k): use the bitmask.
bits = np.unpackbits(
packed_mask[row], count=self.vocab_size, bitorder="little"
)
return np.flatnonzero(bits).astype(np.int32, copy=False)
supports = [support(row) for row in range(len(counts))]
offsets = np.zeros(len(supports) + 1, dtype=np.int64)
np.cumsum([len(s) for s in supports], out=offsets[1:])
return SamplingMaskLists(np.concatenate(supports), offsets)
vllm/v1/worker/gpu/sample/sampler.py
+3 -1
View File
@@ -186,8 +186,10 @@ class Sampler:
sampling_mask_tensors = None
if self.return_sampling_mask:
# Size by the validated top_k batch max; wider supports use the bitmask.
max_num_kept = int(np.max(self.sampling_states.top_k.np[idx_mapping_np]))
sampling_mask_tensors = SamplingMaskTensors.from_logits(
processed_logits, num_sampled
processed_logits, num_sampled, max_num_kept
)
# These are GPU tensors.