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 8, 05:00 UTC
[Model Runner V2] Minor refactoring for penalties (#34662)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 16, 2026 at 21:43 UTC
d00df624f313a6a5a7a6245b71448b068b080cd7
4 files changed
+93
-74
vllm/v1/worker/gpu/model_runner.py
+2
-8
@@ -155,9 +155,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
155
max_num_reqs=self.max_num_reqs,
156
vocab_size=self.vocab_size,
157
device=self.device,
158
- all_token_ids=self.req_states.all_token_ids.gpu,
159
- prompt_len=self.req_states.prompt_len.gpu,
160
- total_len=self.req_states.total_len.gpu,
158
+ req_states=self.req_states,
159
logprobs_mode=self.model_config.logprobs_mode,
160
num_speculative_tokens=self.num_speculative_steps + 1,
161
)
@@ -528,11 +526,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
526
527
if scheduler_output.scheduled_new_reqs:
528
self.req_states.apply_staged_writes()
531
- self.sampler.apply_staged_writes(
532
- self.req_states.all_token_ids.gpu,
533
- self.req_states.prefill_len.np,
534
- self.req_states.prompt_len.np,
535
- )
529
+ self.sampler.apply_staged_writes()
530
if self.uses_mrope:
531
self.mrope_states.apply_staged_writes()
532
vllm/v1/worker/gpu/sample/bad_words.py
+8
-15
@@ -6,24 +6,17 @@ import torch
6
from vllm.sampling_params import SamplingParams
7
from vllm.triton_utils import tl, triton
8
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
9
+from vllm.v1.worker.gpu.states import RequestState
10
11
MAX_BAD_WORDS_TOTAL_TOKENS = 1024 # Max total tokens for all bad words per request
12
MAX_NUM_BAD_WORDS = 128 # Max number of bad words per request
13
14
15
class BadWordsState:
15
- def __init__(
16
- self,
17
- all_token_ids: torch.Tensor,
18
- prompt_len: torch.Tensor,
19
- total_len: torch.Tensor,
20
- ):
21
- self.all_token_ids = all_token_ids
22
- self.prompt_len = prompt_len
23
- self.total_len = total_len
24
-
25
- self.max_num_reqs = prompt_len.shape[0]
26
- self.device = prompt_len.device
16
+ def __init__(self, req_states: RequestState):
17
+ self.req_states = req_states
18
+ self.max_num_reqs = req_states.max_num_reqs
19
+ self.device = req_states.device
20
21
# flattened bad word tokens: [max_num_reqs, MAX_BAD_WORDS_TOTAL_TOKENS]
22
self.bad_word_token_ids = StagedWriteTensor(
@@ -95,9 +88,9 @@ class BadWordsState:
88
self.bad_word_token_ids.gpu,
89
self.bad_word_offsets.gpu,
90
self.num_bad_words.gpu,
98
- self.all_token_ids,
99
- self.prompt_len,
100
- self.total_len,
91
+ self.req_states.all_token_ids.gpu,
92
+ self.req_states.prompt_len.gpu,
93
+ self.req_states.total_len.gpu,
94
input_ids,
95
expanded_local_pos,
96
max_num_bad_words,
vllm/v1/worker/gpu/sample/penalties.py
+77
-37
@@ -6,14 +6,18 @@ import torch
6
from vllm.sampling_params import SamplingParams
7
from vllm.triton_utils import tl, triton
8
from vllm.utils.math_utils import cdiv
9
+from vllm.utils.torch_utils import async_tensor_h2d
10
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
11
+from vllm.v1.worker.gpu.states import RequestState
12
13
14
class PenaltiesState:
13
- def __init__(self, max_num_reqs: int, vocab_size: int, device: torch.device):
14
- self.max_num_reqs = max_num_reqs
15
- self.vocab_size = vocab_size
16
- self.device = device
15
+ def __init__(self, req_states: RequestState):
16
+ self.req_states = req_states
17
+
18
+ max_num_reqs = req_states.max_num_reqs
19
+ self.vocab_size = req_states.vocab_size
20
+ self.device = req_states.device
21
22
self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
23
self.frequency_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
@@ -26,7 +30,7 @@ class PenaltiesState:
30
31
# Statistics for penalties.
32
self.prompt_bin_mask = torch.zeros(
29
- self.max_num_reqs,
33
+ max_num_reqs,
34
cdiv(self.vocab_size, 32),
35
dtype=torch.int32,
36
device=self.device,
@@ -34,10 +38,10 @@ class PenaltiesState:
38
# TODO(woosuk): This tensor is rarely used but can be very large, taking up
39
# GBs of GPU memory. Optimize the memory usage.
40
self.output_bin_counts = torch.zeros(
37
- self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
41
+ max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
42
)
43
40
- self._penalties_reqs: list[int] = []
44
+ self._new_penalties_reqs: list[int] = []
45
46
def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
47
self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
@@ -47,24 +51,29 @@ class PenaltiesState:
51
do_penalty = use_penalty(sampling_params)
52
self.use_penalty[req_idx] = do_penalty
53
if do_penalty:
50
- self._penalties_reqs.append(req_idx)
54
+ self._new_penalties_reqs.append(req_idx)
55
+
56
+ def apply_staged_writes(self) -> None:
57
+ if self._new_penalties_reqs:
58
+ idx_mapping = async_tensor_h2d(
59
+ self._new_penalties_reqs,
60
+ dtype=torch.int32,
61
+ target_device=self.device,
62
+ pin_memory=True,
63
+ )
64
52
- def apply_staged_writes(
53
- self,
54
- all_token_ids: torch.Tensor,
55
- prefill_lens: np.ndarray,
56
- prompt_lens: np.ndarray,
57
- ) -> None:
58
- # TODO(woosuk): Optimize this.
59
- for req_idx in self._penalties_reqs:
65
+ prefill_lens = self.req_states.prefill_len.np[self._new_penalties_reqs]
66
+ max_prefill_len = int(prefill_lens.max())
67
bincount(
61
- all_token_ids[req_idx],
62
- int(prefill_lens[req_idx]),
63
- int(prompt_lens[req_idx]),
64
- self.prompt_bin_mask[req_idx],
65
- self.output_bin_counts[req_idx],
68
+ idx_mapping,
69
+ self.req_states.all_token_ids.gpu,
70
+ self.req_states.prompt_len.gpu,
71
+ self.req_states.prefill_len.gpu,
72
+ self.prompt_bin_mask,
73
+ self.output_bin_counts,
74
+ max_prefill_len,
75
)
67
- self._penalties_reqs.clear()
76
+ self._new_penalties_reqs.clear()
77
78
self.repetition_penalty.copy_to_uva()
79
self.frequency_penalty.copy_to_uva()
@@ -214,51 +223,82 @@ def apply_penalties(
223
)
224
225
217
-@triton.jit(do_not_specialize=["prefill_len", "prompt_len"])
226
+@triton.jit
227
def _bincount_kernel(
228
+ idx_mapping_ptr,
229
all_token_ids_ptr,
220
- prefill_len,
221
- prompt_len,
230
+ all_token_ids_stride,
231
+ prompt_len_ptr,
232
+ prefill_len_ptr,
233
prompt_bin_mask_ptr,
234
+ prompt_bin_mask_stride,
235
output_bin_counts_ptr,
236
+ output_bin_counts_stride,
237
BLOCK_SIZE: tl.constexpr,
238
):
226
- block_idx = tl.program_id(0)
239
+ batch_idx = tl.program_id(0)
240
+ block_idx = tl.program_id(1)
241
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
242
+
243
+ prefill_len = tl.load(prefill_len_ptr + req_state_idx)
244
if block_idx * BLOCK_SIZE >= prefill_len:
245
return
246
247
+ prompt_len = tl.load(prompt_len_ptr + req_state_idx)
248
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
249
if block_idx * BLOCK_SIZE < prompt_len:
250
mask = block < prompt_len
233
- prompt_tokens = tl.load(all_token_ids_ptr + block, mask=mask)
251
+ prompt_tokens = tl.load(
252
+ all_token_ids_ptr + req_state_idx * all_token_ids_stride + block, mask=mask
253
+ )
254
idx = prompt_tokens // 32
255
bit_idx = prompt_tokens % 32
256
bit = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_idx
237
- tl.atomic_or(prompt_bin_mask_ptr + idx, bit, mask=mask)
257
+ tl.atomic_or(
258
+ prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + idx,
259
+ bit,
260
+ mask=mask,
261
+ )
262
+
263
if (block_idx + 1) * BLOCK_SIZE >= prompt_len:
264
mask = block < prefill_len
265
mask &= block >= prompt_len
241
- output_tokens = tl.load(all_token_ids_ptr + block, mask=mask)
242
- tl.atomic_add(output_bin_counts_ptr + output_tokens, 1, mask=mask)
266
+ output_tokens = tl.load(
267
+ all_token_ids_ptr + req_state_idx * all_token_ids_stride + block, mask=mask
268
+ )
269
+ tl.atomic_add(
270
+ output_bin_counts_ptr
271
+ + req_state_idx * output_bin_counts_stride
272
+ + output_tokens,
273
+ 1,
274
+ mask=mask,
275
+ )
276
277
278
def bincount(
279
+ idx_mapping: torch.Tensor,
280
all_token_ids: torch.Tensor,
247
- prefill_len: int,
248
- prompt_len: int,
281
+ prompt_len: torch.Tensor,
282
+ prefill_len: torch.Tensor,
283
prompt_bin_mask: torch.Tensor,
284
output_bin_counts: torch.Tensor,
285
+ max_prefill_len: int,
286
) -> None:
252
- prompt_bin_mask.zero_()
253
- output_bin_counts.zero_()
287
+ prompt_bin_mask[idx_mapping] = 0
288
+ output_bin_counts[idx_mapping] = 0
289
+ num_reqs = idx_mapping.shape[0]
290
BLOCK_SIZE = 1024
255
- num_blocks = triton.cdiv(prefill_len, BLOCK_SIZE)
256
- _bincount_kernel[(num_blocks,)](
291
+ num_blocks = triton.cdiv(max_prefill_len, BLOCK_SIZE)
292
+ _bincount_kernel[(num_reqs, num_blocks)](
293
+ idx_mapping,
294
all_token_ids,
258
- prefill_len,
295
+ all_token_ids.stride(0),
296
prompt_len,
297
+ prefill_len,
298
prompt_bin_mask,
299
+ prompt_bin_mask.stride(0),
300
output_bin_counts,
301
+ output_bin_counts.stride(0),
302
BLOCK_SIZE=BLOCK_SIZE,
303
)
304
vllm/v1/worker/gpu/sample/sampler.py
+6
-14
@@ -15,6 +15,7 @@ from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
15
from vllm.v1.worker.gpu.sample.output import SamplerOutput
16
from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
17
from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
18
+from vllm.v1.worker.gpu.states import RequestState
19
20
21
class Sampler:
@@ -23,9 +24,7 @@ class Sampler:
24
max_num_reqs: int,
25
vocab_size: int,
26
device: torch.device,
26
- all_token_ids: torch.Tensor,
27
- prompt_len: torch.Tensor,
28
- total_len: torch.Tensor,
27
+ req_states: RequestState,
28
logprobs_mode: LogprobsMode = "raw_logprobs",
29
num_speculative_tokens: int = 1,
30
):
@@ -35,9 +34,9 @@ class Sampler:
34
self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
35
36
self.sampling_states = SamplingStates(max_num_reqs, vocab_size)
38
- self.penalties_state = PenaltiesState(max_num_reqs, vocab_size, device)
37
+ self.penalties_state = PenaltiesState(req_states)
38
self.logit_bias_state = LogitBiasState(max_num_reqs, device)
40
- self.bad_words_state = BadWordsState(all_token_ids, prompt_len, total_len)
39
+ self.bad_words_state = BadWordsState(req_states)
40
self.num_speculative_tokens = num_speculative_tokens
41
42
def add_request(
@@ -48,16 +47,9 @@ class Sampler:
47
self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)
48
self.bad_words_state.add_request(req_idx, sampling_params)
49
51
- def apply_staged_writes(
52
- self,
53
- all_token_ids: torch.Tensor,
54
- prefill_lens: np.ndarray,
55
- prompt_lens: np.ndarray,
56
- ) -> None:
50
+ def apply_staged_writes(self) -> None:
51
self.sampling_states.apply_staged_writes()
58
- self.penalties_state.apply_staged_writes(
59
- all_token_ids, prefill_lens, prompt_lens
60
- )
52
+ self.penalties_state.apply_staged_writes()
53
self.logit_bias_state.apply_staged_writes()
54
self.bad_words_state.apply_staged_writes()
55