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, 06:15 UTC
[Model Runner V2] Minor cleanup for Sampler (#34563)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 14, 2026 at 18:29 UTC
9ca768c7404ed8d8a42c5ea3279d804ae454a874
2 files changed
+49
-23
vllm/v1/worker/gpu/sample/sampler.py
+10
-17
@@ -7,12 +7,10 @@ import torch
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
10
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
12
-from vllm.v1.worker.gpu.sample.gumbel import apply_temperature, gumbel_sample
11
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
12
from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
13
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
15
-from vllm.v1.worker.gpu.sample.min_p import apply_min_p
14
from vllm.v1.worker.gpu.sample.output import SamplerOutput
15
from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
16
from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
@@ -127,20 +125,15 @@ class Sampler:
125
)
126
127
# Apply temperature in place.
130
- apply_temperature(logits, idx_mapping, self.sampling_states.temperature.gpu)
131
-
132
- # Apply min_p in place if any request has a non-zero min_p.
133
- do_min_p = self.sampling_states.do_min_p(idx_mapping_np)
134
- if do_min_p:
135
- apply_min_p(logits, idx_mapping, self.sampling_states.min_p.gpu)
136
-
137
- # Apply top_k and/or top_p. This might return a new tensor.
138
- do_top_k = self.sampling_states.do_top_k(idx_mapping_np)
139
- top_k = self.sampling_states.top_k.gpu[idx_mapping] if do_top_k else None
140
- do_top_p = self.sampling_states.do_top_p(idx_mapping_np)
141
- top_p = self.sampling_states.top_p.gpu[idx_mapping] if do_top_p else None
142
- if do_top_k or do_top_p:
143
- logits = apply_top_k_top_p(logits, top_k, top_p)
128
+ self.sampling_states.apply_temperature(logits, idx_mapping, idx_mapping_np)
129
+
130
+ # Apply min_p in place.
131
+ self.sampling_states.apply_min_p(logits, idx_mapping, idx_mapping_np)
132
+
133
+ # Apply top_k and/or top_p. This might or might not return a new tensor.
134
+ logits = self.sampling_states.apply_top_k_top_p(
135
+ logits, idx_mapping, idx_mapping_np
136
+ )
137
138
# Sample the next token.
139
sampled = gumbel_sample(
vllm/v1/worker/gpu/sample/states.py
+39
-6
@@ -4,7 +4,10 @@ import numpy as np
4
import torch
5
6
from vllm.sampling_params import SamplingParams
7
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
8
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
9
+from vllm.v1.worker.gpu.sample.gumbel import apply_temperature
10
+from vllm.v1.worker.gpu.sample.min_p import apply_min_p
11
12
NO_LOGPROBS = -1
13
_NP_INT64_MIN = np.iinfo(np.int64).min
@@ -58,14 +61,44 @@ class SamplingStates:
61
self.min_p.copy_to_uva()
62
self.seeds.copy_to_uva()
63
61
- def do_min_p(self, idx_mapping_np: np.ndarray) -> bool:
62
- return np.any(self.min_p.np[idx_mapping_np] != 0.0)
64
+ def apply_temperature(
65
+ self,
66
+ logits: torch.Tensor,
67
+ idx_mapping: torch.Tensor,
68
+ idx_mapping_np: np.ndarray,
69
+ ) -> None:
70
+ temp_np = self.temperature.np[idx_mapping_np]
71
+ if np.all((temp_np == 0.0) | (temp_np == 1.0)):
72
+ # No request requires temperature. Skip the kernel launch.
73
+ return
74
64
- def do_top_k(self, idx_mapping_np: np.ndarray) -> bool:
65
- return np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
75
+ apply_temperature(logits, idx_mapping, self.temperature.gpu)
76
67
- def do_top_p(self, idx_mapping_np: np.ndarray) -> bool:
68
- return np.any(self.top_p.np[idx_mapping_np] != 1.0)
77
+ def apply_min_p(
78
+ self,
79
+ logits: torch.Tensor,
80
+ idx_mapping: torch.Tensor,
81
+ idx_mapping_np: np.ndarray,
82
+ ) -> None:
83
+ if np.all(self.min_p.np[idx_mapping_np] == 0.0):
84
+ # No request uses min_p. Skip the kernel launch.
85
+ return
86
+ apply_min_p(logits, idx_mapping, self.min_p.gpu)
87
+
88
+ def apply_top_k_top_p(
89
+ self,
90
+ logits: torch.Tensor,
91
+ idx_mapping: torch.Tensor,
92
+ idx_mapping_np: np.ndarray,
93
+ ) -> torch.Tensor:
94
+ do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
95
+ do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0)
96
+ if not (do_top_k or do_top_p):
97
+ return logits
98
+
99
+ top_k = self.top_k.gpu[idx_mapping] if do_top_k else None
100
+ top_p = self.top_p.gpu[idx_mapping] if do_top_p else None
101
+ return apply_top_k_top_p(logits, top_k, top_p)
102
103
def max_num_logprobs(self, idx_mapping_np: np.ndarray) -> int:
104
return int(np.max(self.num_logprobs[idx_mapping_np]))