@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a238cbd89d

Fork 0

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 11, 22:30 UTC

[Model Runner V2] Support min-p sampling (#30171)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Dec 5, 2025 at 21:42 UTC a238cbd89d07b4b0ed8fb3dff3c219a3ee3a1651

4 files changed +77

vllm/v1/worker/gpu/sample/metadata.py

+13

@@ -13,6 +13,7 @@ class SamplingMetadata:

14 top_p: torch.Tensor | None

15 top_k: torch.Tensor | None

16 + min_p: torch.Tensor | None

18 repetition_penalty: torch.Tensor

19 frequency_penalty: torch.Tensor

@@ -44,6 +45,7 @@ class SamplingMetadata:

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)

@@ -64,6 +66,7 @@ class SamplingMetadata:

66 temperature=temperature,

67 top_p=top_p,

68 top_k=top_k,

69 + min_p=min_p,

70 repetition_penalty=repetition_penalty,

71 frequency_penalty=frequency_penalty,

72 presence_penalty=presence_penalty,

@@ -85,6 +88,8 @@ def _expand_sampling_metadata_kernel(

88 expanded_top_p_ptr,

89 top_k_ptr,

90 expanded_top_k_ptr,

91 + min_p_ptr,

92 + expanded_min_p_ptr,

93 rep_penalty_ptr,

94 expanded_rep_penalty_ptr,

95 freq_penalty_ptr,

@@ -115,6 +120,10 @@ def _expand_sampling_metadata_kernel(

120 top_k = tl.load(top_k_ptr + req_idx)

121 tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)

122

123 + if min_p_ptr is not None:

124 + min_p = tl.load(min_p_ptr + req_idx)

125 + tl.store(expanded_min_p_ptr + start_idx + block, min_p, mask=mask)

126 +

127 rep_penalty = tl.load(rep_penalty_ptr + req_idx)

128 tl.store(expanded_rep_penalty_ptr + start_idx + block, rep_penalty, mask=mask)

129

@@ -138,6 +147,7 @@ def expand_sampling_metadata(

147 expanded_temp = create_empty(sampling_metadata.temperature)

148 expanded_top_p = create_empty(sampling_metadata.top_p)

149 expanded_top_k = create_empty(sampling_metadata.top_k)

150 + expanded_min_p = create_empty(sampling_metadata.min_p)

151 expanded_repetition_penalty = create_empty(sampling_metadata.repetition_penalty)

152 expanded_frequency_penalty = create_empty(sampling_metadata.frequency_penalty)

153 expanded_presence_penalty = create_empty(sampling_metadata.presence_penalty)

@@ -151,6 +161,8 @@ def expand_sampling_metadata(

161 expanded_top_p,

162 sampling_metadata.top_k,

163 expanded_top_k,

164 + sampling_metadata.min_p,

165 + expanded_min_p,

166 sampling_metadata.repetition_penalty,

167 expanded_repetition_penalty,

168 sampling_metadata.frequency_penalty,

@@ -166,6 +178,7 @@ def expand_sampling_metadata(

178 temperature=expanded_temp,

179 top_p=expanded_top_p,

180 top_k=expanded_top_k,

181 + min_p=expanded_min_p,

182 seeds=expanded_seeds,

183 repetition_penalty=expanded_repetition_penalty,

184 frequency_penalty=expanded_frequency_penalty,

vllm/v1/worker/gpu/sample/min_p.py new

+53

@@ -0,0 +1,53 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import torch

4 +

5 +from vllm.triton_utils import tl, triton

6 +

7 +

8 +@triton.jit

9 +def _min_p_kernel(

10 + logits_ptr,

11 + logits_stride,

12 + min_p_ptr,

13 + vocab_size,

14 + BLOCK_SIZE: tl.constexpr,

15 +):

16 + req_idx = tl.program_id(0)

17 + min_p = tl.load(min_p_ptr + req_idx).to(tl.float32)

18 + if min_p == 0.0:

19 + return

20 +

21 + max_val = float("-inf")

22 + for i in range(0, vocab_size, BLOCK_SIZE):

23 + block = i + tl.arange(0, BLOCK_SIZE)

24 + mask = block < vocab_size

25 + logits = tl.load(

26 + logits_ptr + req_idx * logits_stride + block, mask=mask, other=float("-inf")

27 + )

28 + max_val = tl.max(tl.maximum(logits, max_val))

29 + max_val = max_val.to(tl.float32) # type: ignore

30 +

31 + threshold = max_val + tl.log(min_p)

32 + for i in range(0, vocab_size, BLOCK_SIZE):

33 + block = i + tl.arange(0, BLOCK_SIZE)

34 + mask = block < vocab_size

35 + logits = tl.load(

36 + logits_ptr + req_idx * logits_stride + block, mask=mask, other=float("-inf")

37 + )

38 + logits = tl.where(logits < threshold, float("-inf"), logits)

39 + tl.store(logits_ptr + req_idx * logits_stride + block, logits, mask=mask)

40 +

41 +

42 +def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor | None) -> None:

43 + if min_p is None:

44 + return

45 + num_reqs, vocab_size = logits.shape

46 + BLOCK_SIZE = 1024

47 + _min_p_kernel[(num_reqs,)](

48 + logits,

49 + logits.stride(0),

50 + min_p,

51 + vocab_size,

52 + BLOCK_SIZE=BLOCK_SIZE,

53 + )

vllm/v1/worker/gpu/sample/sampler.py

@@ -9,6 +9,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p

9 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

10 from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs

11 from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata

12 +from vllm.v1.worker.gpu.sample.min_p import apply_min_p

13 from vllm.v1.worker.gpu.sample.penalties import apply_penalties_and_temperature

@@ -61,6 +62,9 @@ class Sampler:

63 # Apply penalties and temperature in place.

64 apply_penalties_and_temperature(logits, sampling_metadata)

65 + # Apply min_p in place.

66 + apply_min_p(logits, sampling_metadata.min_p)

67 + # Apply top_k and/or top_p. This might return a new tensor.

68 logits = apply_top_k_top_p(

69 logits, sampling_metadata.top_k, sampling_metadata.top_p

70 )

vllm/v1/worker/gpu/states.py

@@ -87,6 +87,7 @@ class RequestState:

87 self.temperature = self._make_param(self.max_num_reqs, torch.float32)

88 self.top_p = self._make_param(self.max_num_reqs, torch.float32)

89 self.top_k = self._make_param(self.max_num_reqs, torch.int32)

90 + self.min_p = self._make_param(self.max_num_reqs, torch.float32)

91 self.repetition_penalty = self._make_param(self.max_num_reqs, torch.float32)

92 self.frequency_penalty = self._make_param(self.max_num_reqs, torch.float32)

93 self.presence_penalty = self._make_param(self.max_num_reqs, torch.float32)

@@ -162,6 +163,7 @@ class RequestState:

163 else:

164 top_k = self.vocab_size

165 self.top_k.np[req_idx] = top_k

166 + self.min_p.np[req_idx] = sampling_params.min_p

167 self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty

168 self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty

169 self.presence_penalty.np[req_idx] = sampling_params.presence_penalty

@@ -217,6 +219,10 @@ class RequestState:

219 no_top_k = np.all(top_k == self.vocab_size)

220 top_k = self.top_k.copy_np_to_gpu(top_k) 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_np_to_gpu(min_p) if not no_min_p else None

225 +

226 rep_penalty = self.repetition_penalty.np[idx_mapping_np]

227 rep_penalty = self.repetition_penalty.copy_np_to_gpu(rep_penalty)

228 freq_penalty = self.frequency_penalty.np[idx_mapping_np]

@@ -236,6 +242,7 @@ class RequestState:

242 temperature=temperature,

243 top_p=top_p,

244 top_k=top_k,

245 + min_p=min_p,

246 repetition_penalty=rep_penalty,

247 frequency_penalty=freq_penalty,

248 presence_penalty=pres_penalty,