A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 07:15 UTC
[Model Runner V2] Use packed mask for prompt bin counts (#29756)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 30, 2025 at 14:15 UTC
ec38a7368df7a14331e5f99f5de899df13f3b954
3 files changed
+35
-25
vllm/v1/worker/gpu/sample/metadata.py
+4
-4
@@ -26,7 +26,7 @@ class SamplingMetadata:
26
27
# For penalties
28
idx_mapping: torch.Tensor
29
- prompt_bin_counts: torch.Tensor
29
+ prompt_bin_mask: torch.Tensor
30
output_bin_counts: torch.Tensor
31
32
@classmethod
@@ -57,7 +57,7 @@ class SamplingMetadata:
57
# NOTE(woosuk): These are placeholder tensors to avoid None checks in the
58
# penalties kernel. We use 2 instead of 1 as vocab_size to avoid Triton
59
# specialization and re-compilation at runtime.
60
- prompt_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
60
+ prompt_bin_mask = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
61
output_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
62
63
return cls(
@@ -71,7 +71,7 @@ class SamplingMetadata:
71
pos=pos,
72
max_num_logprobs=max_num_logprobs,
73
idx_mapping=idx_mapping,
74
- prompt_bin_counts=prompt_bin_counts,
74
+ prompt_bin_mask=prompt_bin_mask,
75
output_bin_counts=output_bin_counts,
76
)
77
@@ -174,6 +174,6 @@ def expand_sampling_metadata(
174
max_num_logprobs=sampling_metadata.max_num_logprobs,
175
# TODO(woosuk): Support penalties with spec decoding.
176
idx_mapping=sampling_metadata.idx_mapping,
177
- prompt_bin_counts=sampling_metadata.prompt_bin_counts,
177
+ prompt_bin_mask=sampling_metadata.prompt_bin_mask,
178
output_bin_counts=sampling_metadata.output_bin_counts,
179
)
vllm/v1/worker/gpu/sample/penalties.py
+21
-15
@@ -15,8 +15,8 @@ def _penalties_and_temperature_kernel(
15
presence_penalty_ptr,
16
temperature_ptr,
17
idx_mapping_ptr,
18
- prompt_bin_counts_ptr,
19
- prompt_bin_counts_stride,
18
+ prompt_bin_mask_ptr,
19
+ prompt_bin_mask_stride,
20
output_bin_counts_ptr,
21
output_bin_counts_stride,
22
vocab_size,
@@ -54,13 +54,16 @@ def _penalties_and_temperature_kernel(
54
55
# Apply repetition penalties.
56
if use_rep_penalty:
57
- prompt_bin_counts = tl.load(
58
- prompt_bin_counts_ptr
59
- + req_state_idx * prompt_bin_counts_stride
60
- + block,
61
- mask=mask,
57
+ packed_block = block_idx * BLOCK_SIZE // 32 + tl.arange(0, BLOCK_SIZE // 32)
58
+ packed_mask = tl.load(
59
+ prompt_bin_mask_ptr
60
+ + req_state_idx * prompt_bin_mask_stride
61
+ + packed_block,
62
+ mask=packed_block < tl.cdiv(vocab_size, 32),
63
)
63
- prompt_bin_mask = prompt_bin_counts > 0
64
+ prompt_bin_mask = (packed_mask[:, None] >> (tl.arange(0, 32)[None, :])) & 1
65
+ prompt_bin_mask = prompt_bin_mask.reshape(BLOCK_SIZE)
66
+
67
# If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
68
scale = tl.where(prompt_bin_mask | output_bin_mask, rep_penalty, 1.0)
69
# If logits are positive, divide by penalty, otherwise multiply by penalty.
@@ -93,8 +96,8 @@ def apply_penalties_and_temperature(
96
sampling_metadata.presence_penalty,
97
sampling_metadata.temperature,
98
sampling_metadata.idx_mapping,
96
- sampling_metadata.prompt_bin_counts,
97
- sampling_metadata.prompt_bin_counts.stride(0),
99
+ sampling_metadata.prompt_bin_mask,
100
+ sampling_metadata.prompt_bin_mask.stride(0),
101
sampling_metadata.output_bin_counts,
102
sampling_metadata.output_bin_counts.stride(0),
103
vocab_size,
@@ -107,7 +110,7 @@ def _bincount_kernel(
110
prefill_token_ids_ptr,
111
prefill_len,
112
prompt_len,
110
- prompt_bin_counts_ptr,
113
+ prompt_bin_mask_ptr,
114
output_bin_counts_ptr,
115
BLOCK_SIZE: tl.constexpr,
116
):
@@ -119,7 +122,10 @@ def _bincount_kernel(
122
if block_idx * BLOCK_SIZE < prompt_len:
123
mask = block < prompt_len
124
prefill_tokens = tl.load(prefill_token_ids_ptr + block, mask=mask)
122
- tl.atomic_add(prompt_bin_counts_ptr + prefill_tokens, 1, mask=mask)
125
+ idx = prefill_tokens // 32
126
+ bit_idx = prefill_tokens % 32
127
+ bit = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_idx
128
+ tl.atomic_or(prompt_bin_mask_ptr + idx, bit, mask=mask)
129
if (block_idx + 1) * BLOCK_SIZE >= prompt_len:
130
mask = block < prefill_len
131
mask &= block >= prompt_len
@@ -131,10 +137,10 @@ def bincount(
137
prefill_token_ids: torch.Tensor,
138
prefill_len: int,
139
prompt_len: int,
134
- prompt_bin_counts: torch.Tensor,
140
+ prompt_bin_mask: torch.Tensor,
141
output_bin_counts: torch.Tensor,
142
) -> None:
137
- prompt_bin_counts.zero_()
143
+ prompt_bin_mask.zero_()
144
output_bin_counts.zero_()
145
BLOCK_SIZE = 1024
146
num_blocks = triton.cdiv(prefill_len, BLOCK_SIZE)
@@ -142,7 +148,7 @@ def bincount(
148
prefill_token_ids,
149
prefill_len,
150
prompt_len,
145
- prompt_bin_counts,
151
+ prompt_bin_mask,
152
output_bin_counts,
153
BLOCK_SIZE=BLOCK_SIZE,
154
)
vllm/v1/worker/gpu/states.py
+10
-6
@@ -7,6 +7,7 @@ import torch
7
8
from vllm.lora.request import LoRARequest
9
from vllm.sampling_params import SamplingParams
10
+from vllm.utils.math_utils import cdiv
11
from vllm.utils.platform_utils import is_uva_available
12
from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
13
from vllm.v1.outputs import LogprobsTensors
@@ -97,11 +98,14 @@ class RequestState:
98
self.needs_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
99
100
# Statistics for penalties.
100
- # TODO(woosuk): These tensors are rarely used but can be extremely large.
101
- # Optimize the memory usage.
102
- self.prompt_bin_counts = torch.zeros(
103
- self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
101
+ self.prompt_bin_mask = torch.zeros(
102
+ self.max_num_reqs,
103
+ cdiv(self.vocab_size, 32),
104
+ dtype=torch.int32,
105
+ device=self.device,
106
)
107
+ # TODO(woosuk): This tensor is rarely used but can be extremely large.
108
+ # Optimize the memory usage.
109
self.output_bin_counts = torch.zeros(
110
self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
111
)
@@ -167,7 +171,7 @@ class RequestState:
171
self.prefill_token_ids.gpu[req_idx],
172
prefill_len,
173
prompt_len,
170
- self.prompt_bin_counts[req_idx],
174
+ self.prompt_bin_mask[req_idx],
175
self.output_bin_counts[req_idx],
176
)
177
@@ -239,7 +243,7 @@ class RequestState:
243
pos=pos,
244
max_num_logprobs=max_num_logprobs,
245
idx_mapping=idx_mapping,
242
- prompt_bin_counts=self.prompt_bin_counts,
246
+ prompt_bin_mask=self.prompt_bin_mask,
247
output_bin_counts=self.output_bin_counts,
248
)
249