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 10, 05:15 UTC
[Model Runner V2] Refactor Sampler (#32245)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 13, 2026 at 17:58 UTC
90c0836902daf1766205f1644322231c274c3d7e
7 files changed
+289
-269
vllm/v1/worker/gpu/model_runner.py
+36
-25
@@ -49,7 +49,6 @@ from vllm.v1.worker.gpu.input_batch import (
49
)
50
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
51
from vllm.v1.worker.gpu.sample.logprob import compute_prompt_logprobs
52
-from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
52
from vllm.v1.worker.gpu.sample.output import SamplerOutput
53
from vllm.v1.worker.gpu.sample.sampler import Sampler
54
from vllm.v1.worker.gpu.spec_decode import init_speculator
@@ -139,7 +138,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
138
dtype=self.dtype,
139
device=self.device,
140
)
142
- self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
141
+ self.sampler = Sampler(
142
+ max_num_reqs=self.max_num_reqs,
143
+ vocab_size=self.vocab_size,
144
+ device=self.device,
145
+ logprobs_mode=self.model_config.logprobs_mode,
146
+ )
147
148
# CUDA graphs.
149
self.cudagraph_manager = CudaGraphManager(
@@ -310,12 +314,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
314
hidden_states: torch.Tensor,
315
) -> None:
316
num_reqs = hidden_states.shape[0]
313
- sampling_metadata = SamplingMetadata.make_dummy(
314
- num_reqs=num_reqs,
315
- device=self.device,
316
- )
317
logits = self.model.compute_logits(hidden_states)
318
- self.sampler(logits, sampling_metadata)
318
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device)
319
+ idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
320
+ pos = torch.zeros(num_reqs, dtype=torch.int64, device=self.device)
321
+ # NOTE(woosuk): During the initial memory profiling, the sampler may skip
322
+ # top_k, top_p, and logprobs, using less GPU memory than what is possible
323
+ # during actual execution.
324
+ self.sampler(logits, idx_mapping, idx_mapping_np, pos)
325
326
@torch.inference_mode()
327
def profile_run(self) -> None:
@@ -401,9 +407,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
407
assert new_req_data.prefill_token_ids is not None
408
assert new_req_data.sampling_params is not None
409
req_id = new_req_data.req_id
410
+ prompt_len = len(new_req_data.prompt_token_ids)
411
self.req_states.add_request(
412
req_id=req_id,
406
- prompt_len=len(new_req_data.prompt_token_ids),
413
+ prompt_len=prompt_len,
414
prefill_token_ids=new_req_data.prefill_token_ids,
415
num_computed_tokens=new_req_data.num_computed_tokens,
416
sampling_params=new_req_data.sampling_params,
@@ -423,6 +430,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
430
self.block_tables.append_block_ids(
431
req_index, new_req_data.block_ids, overwrite=True
432
)
433
+ self.sampler.add_request(
434
+ req_index, prompt_len, new_req_data.sampling_params
435
+ )
436
437
# Add new blocks for the existing requests.
438
cached_reqs = scheduler_output.scheduled_cached_reqs
@@ -436,6 +446,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
446
447
self.req_states.apply_staged_writes()
448
self.block_tables.apply_staged_writes()
449
+ self.sampler.apply_staged_writes(
450
+ self.req_states.prefill_token_ids.gpu,
451
+ self.req_states.prefill_len.np,
452
+ self.req_states.prompt_len,
453
+ )
454
if self.uses_mrope:
455
self.mrope_states.apply_staged_writes()
456
@@ -612,10 +627,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
627
self,
628
hidden_states: torch.Tensor,
629
input_batch: InputBatch,
615
- sampling_metadata: SamplingMetadata,
630
grammar_output: GrammarOutput | None,
631
) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
632
sample_hidden_states = hidden_states[input_batch.logits_indices]
633
+ sample_pos = input_batch.positions[input_batch.logits_indices]
634
logits = self.model.compute_logits(sample_hidden_states)
635
if grammar_output is not None:
636
# Apply grammar bitmask to the logits in-place.
@@ -627,7 +642,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
642
)
643
644
# Sample tokens and compute logprobs (if needed).
630
- sampler_output = self.sampler(logits, sampling_metadata)
645
+ sampler_output = self.sampler(
646
+ logits,
647
+ input_batch.expanded_idx_mapping,
648
+ input_batch.idx_mapping_np,
649
+ sample_pos,
650
+ )
651
652
if input_batch.num_draft_tokens == 0:
653
# No draft tokens (common case).
@@ -766,7 +786,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
786
input_batch.idx_mapping,
787
self.req_states.num_computed_tokens.gpu,
788
self.req_states.last_sampled_tokens,
769
- self.req_states.output_bin_counts,
789
+ self.sampler.penalties_state.output_bin_counts,
790
sampled_tokens,
791
num_sampled,
792
num_rejected,
@@ -786,7 +806,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
806
def propose_draft(
807
self,
808
input_batch: InputBatch,
789
- sampling_metadata: SamplingMetadata,
809
last_hidden_states: torch.Tensor,
810
aux_hidden_states: list[torch.Tensor] | None,
811
num_sampled: torch.Tensor,
@@ -801,13 +820,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
820
]
821
draft_tokens = self.speculator.propose(
822
input_batch,
804
- sampling_metadata,
823
last_hidden_states,
824
aux_hidden_states,
825
num_sampled,
826
num_rejected,
827
last_sampled_tokens,
828
next_prefill_tokens,
829
+ self.sampler.sampling_states.temperature.gpu,
830
+ self.sampler.sampling_states.seeds.gpu,
831
)
832
return draft_tokens
833
@@ -893,12 +913,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
913
scheduler_output,
914
num_tokens_after_padding,
915
)
896
-
897
- pos = input_batch.positions[input_batch.logits_indices]
898
- sampling_metadata = self.req_states.make_sampling_metadata(
899
- input_batch.expanded_idx_mapping, input_batch.idx_mapping_np, pos
900
- )
901
-
916
if self.lora_config:
917
# Activate LoRA adapters.
918
lora_inputs = self.req_states.make_lora_inputs(
@@ -917,7 +931,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
931
device=self.device,
932
)
933
self.prepare_dummy_attn_metadata(input_batch)
920
- sampling_metadata = None
934
935
# Run model.
936
if cudagraph_mode == CUDAGraphMode.FULL:
@@ -946,7 +959,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
959
positions=positions,
960
)
961
949
- self.execute_model_state = hidden_states, input_batch, sampling_metadata
962
+ self.execute_model_state = hidden_states, input_batch
963
return None
964
965
@torch.inference_mode()
@@ -955,12 +968,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
968
grammar_output: GrammarOutput | None,
969
) -> AsyncOutput | ModelRunnerOutput:
970
assert self.execute_model_state is not None
958
- hidden_states, input_batch, sampling_metadata = self.execute_model_state
971
+ hidden_states, input_batch = self.execute_model_state
972
self.execute_model_state = None # type: ignore
960
- assert sampling_metadata is not None
973
974
sampler_output, num_sampled, num_rejected = self.sample(
963
- hidden_states, input_batch, sampling_metadata, grammar_output
975
+ hidden_states, input_batch, grammar_output
976
)
977
prompt_logprobs_dict = self.compute_prompt_logprobs(hidden_states, input_batch)
978
@@ -992,7 +1004,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1004
if self.do_spec_decode:
1005
draft_tokens = self.propose_draft(
1006
input_batch,
995
- sampling_metadata,
1007
hidden_states,
1008
None, # aux_hidden_states
1009
num_sampled,
vllm/v1/worker/gpu/sample/metadata.py
deleted
-79
@@ -1,79 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
4
-
5
-import torch
6
-
7
-
8
-@dataclass
9
-class SamplingMetadata:
10
- idx_mapping: torch.Tensor
11
-
12
- temperature: torch.Tensor
13
-
14
- top_p: torch.Tensor | None
15
- top_k: torch.Tensor | None
16
- min_p: torch.Tensor | None
17
-
18
- # For penalties
19
- repetition_penalty: torch.Tensor
20
- frequency_penalty: torch.Tensor
21
- presence_penalty: torch.Tensor
22
- prompt_bin_mask: torch.Tensor
23
- output_bin_counts: torch.Tensor
24
-
25
- seeds: torch.Tensor
26
- pos: torch.Tensor
27
-
28
- # None means no logprobs, 0 means sampled token logprobs only
29
- max_num_logprobs: int | None
30
-
31
- @classmethod
32
- def make_dummy(
33
- cls,
34
- num_reqs: int,
35
- device: torch.device,
36
- ) -> "SamplingMetadata":
37
- assert num_reqs > 0
38
- idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
39
-
40
- temperature = torch.zeros(num_reqs, dtype=torch.float32, device=device)
41
- temperature[0] = 0.5
42
- # TODO(woosuk): Use top-p and top-k for dummy sampler.
43
- # Currently, they are disabled because of memory usage.
44
- # top_p = torch.full((num_reqs,), 0.95, dtype=torch.float32, device=device)
45
- # top_k = torch.full((num_reqs,), 20, dtype=torch.int32, device=device)
46
- top_p = None
47
- top_k = None
48
- min_p = torch.zeros(num_reqs, dtype=torch.float32, device=device)
49
- # NOTE(woosuk): We must set penalties to their default values to make sure
50
- # the penalties kernel does not touch the placeholder bin_counts tensors.
51
- repetition_penalty = torch.ones(num_reqs, dtype=torch.float32, device=device)
52
- frequency_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
53
- presence_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
54
-
55
- # NOTE(woosuk): These are placeholder tensors to avoid None checks in the
56
- # penalties kernel. We use 2 instead of 1 as vocab_size to avoid Triton
57
- # specialization and re-compilation at runtime.
58
- prompt_bin_mask = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
59
- output_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
60
-
61
- seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
62
- pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
63
- max_num_logprobs = 20
64
-
65
- return cls(
66
- idx_mapping=idx_mapping,
67
- temperature=temperature,
68
- top_p=top_p,
69
- top_k=top_k,
70
- min_p=min_p,
71
- repetition_penalty=repetition_penalty,
72
- frequency_penalty=frequency_penalty,
73
- presence_penalty=presence_penalty,
74
- prompt_bin_mask=prompt_bin_mask,
75
- output_bin_counts=output_bin_counts,
76
- seeds=seeds,
77
- pos=pos,
78
- max_num_logprobs=max_num_logprobs,
79
- )
vllm/v1/worker/gpu/sample/penalties.py
+103
-11
@@ -1,9 +1,87 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
import torch
5
6
+from vllm.sampling_params import SamplingParams
7
from vllm.triton_utils import tl, triton
6
-from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
8
+from vllm.utils.math_utils import cdiv
9
+from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
10
+
11
+
12
+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
17
+
18
+ self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
19
+ self.frequency_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
20
+ self.presence_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
21
+
22
+ # Initialize repetition penalty manually because 0 is an invalid value for it.
23
+ self.repetition_penalty.np.fill(1.0)
24
+ self.repetition_penalty.copy_to_uva()
25
+
26
+ # Statistics for penalties.
27
+ self.prompt_bin_mask = torch.zeros(
28
+ self.max_num_reqs,
29
+ cdiv(self.vocab_size, 32),
30
+ dtype=torch.int32,
31
+ device=self.device,
32
+ )
33
+ # TODO(woosuk): This tensor is rarely used but can be very large, taking up
34
+ # GBs of GPU memory. Optimize the memory usage.
35
+ self.output_bin_counts = torch.zeros(
36
+ self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
37
+ )
38
+
39
+ self._penalties_reqs: list[int] = []
40
+
41
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
42
+ self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
43
+ self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty
44
+ self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
45
+ if use_penalty(sampling_params):
46
+ self._penalties_reqs.append(req_idx)
47
+
48
+ def apply_staged_writes(
49
+ self,
50
+ prefill_token_ids: torch.Tensor,
51
+ prefill_lens: np.ndarray,
52
+ prompt_lens: np.ndarray,
53
+ ) -> None:
54
+ # TODO(woosuk): Optimize this.
55
+ for req_idx in self._penalties_reqs:
56
+ bincount(
57
+ prefill_token_ids[req_idx],
58
+ int(prefill_lens[req_idx]),
59
+ int(prompt_lens[req_idx]),
60
+ self.prompt_bin_mask[req_idx],
61
+ self.output_bin_counts[req_idx],
62
+ )
63
+ self._penalties_reqs.clear()
64
+
65
+ self.repetition_penalty.copy_to_uva()
66
+ self.frequency_penalty.copy_to_uva()
67
+ self.presence_penalty.copy_to_uva()
68
+
69
+ def apply_penalties_and_temperature(
70
+ self,
71
+ logits: torch.Tensor,
72
+ idx_mapping: torch.Tensor,
73
+ temperature: torch.Tensor,
74
+ ) -> None:
75
+ apply_penalties_and_temperature(
76
+ logits,
77
+ idx_mapping,
78
+ temperature,
79
+ self.repetition_penalty.gpu,
80
+ self.frequency_penalty.gpu,
81
+ self.presence_penalty.gpu,
82
+ self.prompt_bin_mask,
83
+ self.output_bin_counts,
84
+ )
85
86
87
@triton.jit
@@ -84,7 +162,13 @@ def _penalties_and_temperature_kernel(
162
163
def apply_penalties_and_temperature(
164
logits: torch.Tensor,
87
- sampling_metadata: SamplingMetadata,
165
+ idx_mapping: torch.Tensor,
166
+ temperature: torch.Tensor,
167
+ repetition_penalty: torch.Tensor,
168
+ frequency_penalty: torch.Tensor,
169
+ presence_penalty: torch.Tensor,
170
+ prompt_bin_mask: torch.Tensor,
171
+ output_bin_counts: torch.Tensor,
172
) -> None:
173
num_reqs, vocab_size = logits.shape
174
BLOCK_SIZE = 8192
@@ -92,15 +176,15 @@ def apply_penalties_and_temperature(
176
_penalties_and_temperature_kernel[(num_reqs, num_blocks)](
177
logits,
178
logits.stride(0),
95
- sampling_metadata.idx_mapping,
96
- sampling_metadata.repetition_penalty,
97
- sampling_metadata.frequency_penalty,
98
- sampling_metadata.presence_penalty,
99
- sampling_metadata.temperature,
100
- sampling_metadata.prompt_bin_mask,
101
- sampling_metadata.prompt_bin_mask.stride(0),
102
- sampling_metadata.output_bin_counts,
103
- sampling_metadata.output_bin_counts.stride(0),
179
+ idx_mapping,
180
+ repetition_penalty,
181
+ frequency_penalty,
182
+ presence_penalty,
183
+ temperature,
184
+ prompt_bin_mask,
185
+ prompt_bin_mask.stride(0),
186
+ output_bin_counts,
187
+ output_bin_counts.stride(0),
188
vocab_size,
189
BLOCK_SIZE=BLOCK_SIZE,
190
)
@@ -153,3 +237,11 @@ def bincount(
237
output_bin_counts,
238
BLOCK_SIZE=BLOCK_SIZE,
239
)
240
+
241
+
242
+def use_penalty(sampling_params: SamplingParams) -> bool:
243
+ return (
244
+ sampling_params.repetition_penalty != 1.0
245
+ or sampling_params.frequency_penalty != 0.0
246
+ or sampling_params.presence_penalty != 0.0
247
+ )
vllm/v1/worker/gpu/sample/sampler.py
+70
-22
@@ -1,23 +1,29 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import numpy as np
5
import torch
6
7
import vllm.envs as envs
8
from vllm.config.model import LogprobsMode
9
+from vllm.sampling_params import SamplingParams
10
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
12
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
13
+from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
14
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
12
-from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
15
from vllm.v1.worker.gpu.sample.min_p import apply_min_p
16
from vllm.v1.worker.gpu.sample.output import SamplerOutput
15
-from vllm.v1.worker.gpu.sample.penalties import apply_penalties_and_temperature
17
+from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
18
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
19
20
21
class Sampler:
22
def __init__(
23
self,
24
+ max_num_reqs: int,
25
+ vocab_size: int,
26
+ device: torch.device,
27
logprobs_mode: LogprobsMode = "raw_logprobs",
28
):
29
if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
@@ -25,26 +31,54 @@ class Sampler:
31
self.logprobs_mode = logprobs_mode
32
self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
33
34
+ self.sampling_states = SamplingStates(max_num_reqs, vocab_size)
35
+ self.penalties_state = PenaltiesState(max_num_reqs, vocab_size, device)
36
+ self.logit_bias_state = LogitBiasState(max_num_reqs, device)
37
+
38
+ def add_request(
39
+ self,
40
+ req_idx: int,
41
+ prompt_len: int,
42
+ sampling_params: SamplingParams,
43
+ ) -> None:
44
+ self.sampling_states.add_request(req_idx, sampling_params)
45
+ self.penalties_state.add_request(req_idx, sampling_params)
46
+ self.logit_bias_state.add_request(req_idx, prompt_len, sampling_params)
47
+
48
+ def apply_staged_writes(
49
+ self,
50
+ prefill_token_ids: torch.Tensor,
51
+ prefill_lens: np.ndarray,
52
+ prompt_lens: np.ndarray,
53
+ ) -> None:
54
+ self.sampling_states.apply_staged_writes()
55
+ self.penalties_state.apply_staged_writes(
56
+ prefill_token_ids, prefill_lens, prompt_lens
57
+ )
58
+ self.logit_bias_state.apply_staged_writes()
59
+
60
def __call__(
61
self,
62
logits: torch.Tensor,
31
- sampling_metadata: SamplingMetadata,
63
+ idx_mapping: torch.Tensor,
64
+ idx_mapping_np: np.ndarray,
65
+ pos: torch.Tensor,
66
) -> SamplerOutput:
67
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
68
# that num_nans is computed before applying penalties and temperature.
69
num_nans = get_num_nans(logits) if self.compute_nans else None
36
- sampled, processed_logits = self.sample(logits, sampling_metadata)
37
- if sampling_metadata.max_num_logprobs is not None:
70
+ sampled, processed_logits = self.sample(
71
+ logits, idx_mapping, idx_mapping_np, pos
72
+ )
73
+
74
+ max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
75
+ if max_num_logprobs != NO_LOGPROBS:
76
logits = (
77
processed_logits
78
if self.logprobs_mode == "processed_logprobs"
79
else logits
80
)
43
- logprobs_tensors = compute_topk_logprobs(
44
- logits,
45
- sampling_metadata.max_num_logprobs,
46
- sampled,
47
- )
81
+ logprobs_tensors = compute_topk_logprobs(logits, max_num_logprobs, sampled)
82
else:
83
logprobs_tensors = None
84
@@ -62,27 +96,41 @@ class Sampler:
96
def sample(
97
self,
98
logits: torch.Tensor,
65
- sampling_metadata: SamplingMetadata,
99
+ idx_mapping: torch.Tensor,
100
+ idx_mapping_np: np.ndarray,
101
+ pos: torch.Tensor,
102
) -> tuple[torch.Tensor, torch.Tensor]:
103
# Copy logits to a new FP32 tensor.
104
logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
105
106
+ # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
107
+ self.logit_bias_state.apply_logit_bias(logits, idx_mapping, pos)
108
+
109
# Apply penalties and temperature in place.
71
- apply_penalties_and_temperature(logits, sampling_metadata)
72
- # Apply min_p in place.
73
- if sampling_metadata.min_p is not None:
74
- apply_min_p(logits, sampling_metadata.idx_mapping, sampling_metadata.min_p)
75
- # Apply top_k and/or top_p. This might return a new tensor.
76
- logits = apply_top_k_top_p(
77
- logits, sampling_metadata.top_k, sampling_metadata.top_p
110
+ self.penalties_state.apply_penalties_and_temperature(
111
+ logits, idx_mapping, self.sampling_states.temperature.gpu
112
)
113
114
+ # Apply min_p in place if any request has a non-zero min_p.
115
+ do_min_p = self.sampling_states.do_min_p(idx_mapping_np)
116
+ if do_min_p:
117
+ apply_min_p(logits, idx_mapping, self.sampling_states.min_p.gpu)
118
+
119
+ # Apply top_k and/or top_p. This might return a new tensor.
120
+ do_top_k = self.sampling_states.do_top_k(idx_mapping_np)
121
+ top_k = self.sampling_states.top_k.gpu[idx_mapping] if do_top_k else None
122
+ do_top_p = self.sampling_states.do_top_p(idx_mapping_np)
123
+ top_p = self.sampling_states.top_p.gpu[idx_mapping] if do_top_p else None
124
+ if do_top_k or do_top_p:
125
+ logits = apply_top_k_top_p(logits, top_k, top_p)
126
+
127
+ # Sample the next token.
128
sampled = gumbel_sample(
129
logits,
82
- sampling_metadata.idx_mapping,
83
- sampling_metadata.temperature,
84
- sampling_metadata.seeds,
85
- sampling_metadata.pos,
130
+ idx_mapping,
131
+ self.sampling_states.temperature.gpu,
132
+ self.sampling_states.seeds.gpu,
133
+ pos,
134
apply_temperature=False,
135
)
136
return sampled, logits
vllm/v1/worker/gpu/sample/states.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+import torch
5
+
6
+from vllm.sampling_params import SamplingParams
7
+from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
8
+
9
+NO_LOGPROBS = -1
10
+_NP_INT64_MIN = np.iinfo(np.int64).min
11
+_NP_INT64_MAX = np.iinfo(np.int64).max
12
+
13
+
14
+class SamplingStates:
15
+ def __init__(self, max_num_reqs: int, vocab_size: int):
16
+ self.max_num_reqs = max_num_reqs
17
+ self.vocab_size = vocab_size
18
+
19
+ self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
20
+ self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
21
+ self.top_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
22
+ self.min_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
23
+ self.seeds = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
24
+
25
+ # Initialize top_k and top_p manually because 0 is an invalid value for them.
26
+ self.top_k.np.fill(self.vocab_size)
27
+ self.top_k.copy_to_uva()
28
+ self.top_p.np.fill(1.0)
29
+ self.top_p.copy_to_uva()
30
+
31
+ self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
32
+ # -1 means no logprobs are requested.
33
+ self.num_logprobs.fill(NO_LOGPROBS)
34
+
35
+ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
36
+ self.temperature.np[req_idx] = sampling_params.temperature
37
+ self.top_p.np[req_idx] = sampling_params.top_p
38
+ if 0 < sampling_params.top_k < self.vocab_size:
39
+ top_k = sampling_params.top_k
40
+ else:
41
+ top_k = self.vocab_size
42
+ self.top_k.np[req_idx] = top_k
43
+ self.min_p.np[req_idx] = sampling_params.min_p
44
+
45
+ if sampling_params.seed is not None:
46
+ seed = sampling_params.seed
47
+ else:
48
+ seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)
49
+ self.seeds.np[req_idx] = seed
50
+
51
+ if sampling_params.logprobs is not None:
52
+ num_logprobs = sampling_params.logprobs
53
+ else:
54
+ num_logprobs = NO_LOGPROBS
55
+ self.num_logprobs[req_idx] = num_logprobs
56
+
57
+ def apply_staged_writes(self) -> None:
58
+ self.temperature.copy_to_uva()
59
+ self.top_p.copy_to_uva()
60
+ self.top_k.copy_to_uva()
61
+ self.min_p.copy_to_uva()
62
+ self.seeds.copy_to_uva()
63
+
64
+ def do_min_p(self, idx_mapping_np: np.ndarray) -> bool:
65
+ return np.any(self.min_p.np[idx_mapping_np] != 0.0)
66
+
67
+ def do_top_k(self, idx_mapping_np: np.ndarray) -> bool:
68
+ return np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
69
+
70
+ def do_top_p(self, idx_mapping_np: np.ndarray) -> bool:
71
+ return np.any(self.top_p.np[idx_mapping_np] != 1.0)
72
+
73
+ def max_num_logprobs(self, idx_mapping_np: np.ndarray) -> int:
74
+ return int(np.max(self.num_logprobs[idx_mapping_np]))
vllm/v1/worker/gpu/spec_decode/eagle.py
+6
-4
@@ -17,7 +17,6 @@ from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
17
from vllm.v1.worker.gpu.block_table import BlockTables
18
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
19
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
20
-from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
20
from vllm.v1.worker.gpu.spec_decode.eagle_cudagraph import EagleCudaGraphManager
21
22
logger = init_logger(__name__)
@@ -188,7 +187,6 @@ class EagleSpeculator:
187
def propose(
188
self,
189
input_batch: InputBatch,
191
- sampling_metadata: SamplingMetadata,
190
# [num_tokens, hidden_size]
191
last_hidden_states: torch.Tensor,
192
# num_layers x [num_tokens, hidden_size]
@@ -201,6 +199,10 @@ class EagleSpeculator:
199
last_sampled: torch.Tensor,
200
# [num_reqs]
201
next_prefill_tokens: torch.Tensor,
202
+ # [max_num_reqs]
203
+ temperature: torch.Tensor,
204
+ # [max_num_reqs]
205
+ seeds: torch.Tensor,
206
) -> torch.Tensor:
207
# NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
208
# number of rejected tokens, we maintain the size of eagle's input_ids and
@@ -246,8 +248,8 @@ class EagleSpeculator:
248
# affect the output distribution after rejection sampling.
249
idx_mapping = self.idx_mapping[:num_reqs]
250
idx_mapping.copy_(input_batch.idx_mapping)
249
- self.temperature.copy_(sampling_metadata.temperature)
250
- self.seeds.copy_(sampling_metadata.seeds)
251
+ self.temperature.copy_(temperature)
252
+ self.seeds.copy_(seeds)
253
# Gather the values and copy them to the pre-allocated buffers.
254
pos = self.input_buffers.positions[:num_reqs]
255
torch.gather(input_batch.positions, 0, last_token_indices, out=pos)
vllm/v1/worker/gpu/states.py
-128
@@ -7,14 +7,9 @@ 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
10
from vllm.v1.outputs import LogprobsTensors
11
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
13
-from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
14
-from vllm.v1.worker.gpu.sample.penalties import bincount
12
16
-_NP_INT64_MIN = np.iinfo(np.int64).min
17
-_NP_INT64_MAX = np.iinfo(np.int64).max
13
NO_LORA_ID = 0
14
15
@@ -81,38 +76,8 @@ class RequestState:
76
self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
77
self.lora_ids.fill(NO_LORA_ID)
78
84
- # Sampling parameters.
85
- self.temperature = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
86
- self.top_p = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
87
- self.top_k = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
88
- self.min_p = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
89
- self.repetition_penalty = UvaBackedTensor(
90
- self.max_num_reqs, dtype=torch.float32
91
- )
92
- self.frequency_penalty = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
93
- self.presence_penalty = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
94
- self.seeds = UvaBackedTensor(self.max_num_reqs, dtype=torch.int64)
95
-
96
- self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
97
- # -1 means no logprobs are requested.
98
- self.num_logprobs.fill(-1)
79
self.needs_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
80
101
- # Statistics for penalties.
102
- self.prompt_bin_mask = torch.zeros(
103
- self.max_num_reqs,
104
- cdiv(self.vocab_size, 32),
105
- dtype=torch.int32,
106
- device=self.device,
107
- )
108
- # TODO(woosuk): This tensor is rarely used but can be extremely large.
109
- # Optimize the memory usage.
110
- self.output_bin_counts = torch.zeros(
111
- self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
112
- )
113
-
114
- self._penalties_reqs: list[int] = []
115
-
81
@property
82
def num_reqs(self) -> int:
83
return len(self.req_id_to_index)
@@ -147,33 +112,6 @@ class RequestState:
112
else:
113
self.lora_ids[req_idx] = NO_LORA_ID
114
150
- self.temperature.np[req_idx] = sampling_params.temperature
151
- self.top_p.np[req_idx] = sampling_params.top_p
152
- if 0 < sampling_params.top_k < self.vocab_size:
153
- top_k = sampling_params.top_k
154
- else:
155
- top_k = self.vocab_size
156
- self.top_k.np[req_idx] = top_k
157
- self.min_p.np[req_idx] = sampling_params.min_p
158
- self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
159
- self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty
160
- self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
161
-
162
- if use_penalty(sampling_params):
163
- self._penalties_reqs.append(req_idx)
164
-
165
- if sampling_params.seed is not None:
166
- seed = sampling_params.seed
167
- else:
168
- seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)
169
- self.seeds.np[req_idx] = seed
170
-
171
- if sampling_params.logprobs is not None:
172
- num_logprobs = sampling_params.logprobs
173
- else:
174
- num_logprobs = -1
175
- self.num_logprobs[req_idx] = num_logprobs
176
-
115
# For now, only support prompt logprobs for the prompt tokens.
116
needs_prompt_logprobs = sampling_params.prompt_logprobs is not None
117
self.needs_prompt_logprobs[req_idx] = needs_prompt_logprobs
@@ -183,17 +121,6 @@ class RequestState:
121
self.prefill_token_ids.apply_write()
122
self.num_computed_tokens.apply_write()
123
186
- # TODO(woosuk): Optimize this.
187
- for req_idx in self._penalties_reqs:
188
- bincount(
189
- self.prefill_token_ids.gpu[req_idx],
190
- int(self.prefill_len.np[req_idx]),
191
- int(self.prompt_len[req_idx]),
192
- self.prompt_bin_mask[req_idx],
193
- self.output_bin_counts[req_idx],
194
- )
195
- self._penalties_reqs.clear()
196
-
124
def remove_request(self, req_id: str) -> None:
125
self.extra_data.pop(req_id, None)
126
req_idx = self.req_id_to_index.pop(req_id, None)
@@ -203,53 +130,6 @@ class RequestState:
130
self.index_to_req_id.pop(req_idx, None)
131
self.free_indices.append(req_idx)
132
206
- def make_sampling_metadata(
207
- self,
208
- idx_mapping: torch.Tensor,
209
- idx_mapping_np: np.ndarray,
210
- pos: torch.Tensor,
211
- ) -> SamplingMetadata:
212
- temperature = self.temperature.copy_to_uva()
213
-
214
- top_p = self.top_p.np[idx_mapping_np]
215
- no_top_p = np.all(top_p == 1.0)
216
- top_p = self.top_p.copy_to_uva()[idx_mapping] if not no_top_p else None
217
-
218
- top_k = self.top_k.np[idx_mapping_np]
219
- no_top_k = np.all(top_k == self.vocab_size)
220
- top_k = self.top_k.copy_to_uva()[idx_mapping] if not no_top_k else None
221
-
222
- min_p = self.min_p.np[idx_mapping_np]
223
- no_min_p = np.all(min_p == 0.0)
224
- min_p = self.min_p.copy_to_uva() if not no_min_p else None
225
-
226
- rep_penalty = self.repetition_penalty.copy_to_uva()
227
- freq_penalty = self.frequency_penalty.copy_to_uva()
228
- pres_penalty = self.presence_penalty.copy_to_uva()
229
-
230
- seeds = self.seeds.copy_to_uva()
231
-
232
- num_logprobs = self.num_logprobs[idx_mapping_np]
233
- max_num_logprobs: int | None = int(np.max(num_logprobs))
234
- if max_num_logprobs == -1:
235
- max_num_logprobs = None
236
-
237
- return SamplingMetadata(
238
- idx_mapping=idx_mapping,
239
- temperature=temperature,
240
- top_p=top_p,
241
- top_k=top_k,
242
- min_p=min_p,
243
- repetition_penalty=rep_penalty,
244
- frequency_penalty=freq_penalty,
245
- presence_penalty=pres_penalty,
246
- prompt_bin_mask=self.prompt_bin_mask,
247
- output_bin_counts=self.output_bin_counts,
248
- seeds=seeds,
249
- pos=pos,
250
- max_num_logprobs=max_num_logprobs,
251
- )
252
-
133
def make_lora_inputs(
134
self,
135
req_ids: list[str],
@@ -272,11 +152,3 @@ class RequestState:
152
class ExtraData:
153
lora_request: LoRARequest | None
154
in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)
275
-
276
-
277
-def use_penalty(sampling_params: SamplingParams) -> bool:
278
- return (
279
- sampling_params.repetition_penalty != 1.0
280
- or sampling_params.frequency_penalty != 0.0
281
- or sampling_params.presence_penalty != 0.0
282
- )