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 2, 01:00 UTC
[Model Runner V2] Use FlashInfer sampler (#42472)
Nick Hill committed
Jun 3, 2026 at 07:59 UTC
ec8d60bea8658b715370543e785ed6a866d4c93f
3 files changed
+141
-69
vllm/v1/sample/ops/topk_topp_sampler.py
+74
-49
@@ -19,6 +19,69 @@ if HAS_TRITON:
19
logger = init_logger(__name__)
20
21
22
+_FLASHINFER_MIN_VERSION = "0.2.3"
23
+
24
+
25
+def flashinfer_sampler_supported() -> bool:
26
+ """Decide whether FlashInfer's top-p/top-k sampler can be used.
27
+
28
+ Returns False (with appropriate logging) when ``VLLM_USE_FLASHINFER_SAMPLER``
29
+ is 0, when the platform isn't CUDA, when the GPU's compute capability is
30
+ unsupported, or when the installed flashinfer is missing or too old. Raises
31
+ ``RuntimeError`` if the user explicitly opted in via the env var but
32
+ FlashInfer is unavailable.
33
+
34
+ Note: callers must additionally ensure ``logprobs_mode`` doesn't require
35
+ post-top-k/top-p logits/logprobs for any request whose logprobs will be
36
+ returned in this step, since FlashInfer doesn't expose those.
37
+ """
38
+ if not current_platform.is_cuda():
39
+ return False
40
+ if not envs.VLLM_USE_FLASHINFER_SAMPLER:
41
+ logger.info_once(
42
+ "FlashInfer top-p/top-k sampling disabled via "
43
+ "VLLM_USE_FLASHINFER_SAMPLER=0."
44
+ )
45
+ return False
46
+ from vllm.v1.attention.backends.flashinfer import FlashInferBackend
47
+
48
+ capability = current_platform.get_device_capability()
49
+ assert capability is not None
50
+ unsupported_reason: str | None = None
51
+ if not FlashInferBackend.supports_compute_capability(capability):
52
+ unsupported_reason = (
53
+ f"unsupported compute capability {capability.as_version_str()}"
54
+ )
55
+ else:
56
+ try:
57
+ import flashinfer
58
+
59
+ if version.parse(flashinfer.__version__) < version.parse(
60
+ _FLASHINFER_MIN_VERSION
61
+ ):
62
+ unsupported_reason = (
63
+ f"flashinfer {flashinfer.__version__} is too old "
64
+ f"(>={_FLASHINFER_MIN_VERSION} required)"
65
+ )
66
+ except ImportError:
67
+ unsupported_reason = "flashinfer is not installed"
68
+
69
+ if unsupported_reason is None:
70
+ logger.info_once("Using FlashInfer for top-p & top-k sampling.", scope="global")
71
+ return True
72
+ if envs.is_set("VLLM_USE_FLASHINFER_SAMPLER"):
73
+ raise RuntimeError(
74
+ f"FlashInfer top-p/top-k sampling unavailable: {unsupported_reason}. "
75
+ "Unset VLLM_USE_FLASHINFER_SAMPLER=1."
76
+ )
77
+ logger.warning_once(
78
+ "FlashInfer top-p/top-k sampling unavailable: %s; falling back. "
79
+ "Set VLLM_USE_FLASHINFER_SAMPLER=0 to silence.",
80
+ unsupported_reason,
81
+ )
82
+ return False
83
+
84
+
85
class TopKTopPSampler(nn.Module):
86
"""
87
Module that performs optional top-k and top-p filtering followed by
@@ -30,49 +93,16 @@ class TopKTopPSampler(nn.Module):
93
def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs") -> None:
94
super().__init__()
95
self.logprobs_mode = logprobs_mode
33
- # flashinfer optimization does not apply if intermediate
34
- # logprobs/logits after top_k/top_p need to be returned
35
- if (
36
- logprobs_mode not in ("processed_logits", "processed_logprobs")
37
- and current_platform.is_cuda()
38
- ):
39
- if envs.VLLM_USE_FLASHINFER_SAMPLER:
40
- from vllm.v1.attention.backends.flashinfer import FlashInferBackend
41
-
42
- capability = current_platform.get_device_capability()
43
- assert capability is not None
44
- if FlashInferBackend.supports_compute_capability(capability):
45
- logger.info_once(
46
- "Using FlashInfer for top-p & top-k sampling.",
47
- scope="global",
48
- )
49
- self.forward = self.forward_cuda
50
- elif envs.is_set("VLLM_USE_FLASHINFER_SAMPLER"):
51
- # User explicitly opted in but the GPU can't run FlashInfer.
52
- capability_str = capability.as_version_str()
53
- raise RuntimeError(
54
- "FlashInfer does not support compute capability "
55
- f"{capability_str}, unset VLLM_USE_FLASHINFER_SAMPLER=1."
56
- )
57
- else:
58
- # Default-on path; hardware can't run FlashInfer →
59
- # quietly fall back to the PyTorch-native sampler
60
- # instead of failing server startup.
61
- logger.warning_once(
62
- "FlashInfer top-p/top-k sampling not supported on "
63
- "compute capability %s; falling back to PyTorch-native "
64
- "sampler. Set VLLM_USE_FLASHINFER_SAMPLER=0 to silence.",
65
- capability.as_version_str(),
66
- )
67
- self.forward = self.forward_native
68
- else:
69
- # User explicitly set VLLM_USE_FLASHINFER_SAMPLER=0.
70
- logger.info_once(
71
- "FlashInfer top-p/top-k sampling disabled via "
72
- "VLLM_USE_FLASHINFER_SAMPLER=0; using PyTorch-native sampler."
73
- )
74
- self.forward = self.forward_native
75
-
96
+ if current_platform.is_cuda():
97
+ # FlashInfer doesn't expose post-top-k/top-p logits/logprobs,
98
+ # so it can't be used when the configured mode requires them.
99
+ can_use_flashinfer = (
100
+ logprobs_mode not in ("processed_logits", "processed_logprobs")
101
+ and flashinfer_sampler_supported()
102
+ )
103
+ self.forward = (
104
+ self.forward_cuda if can_use_flashinfer else self.forward_native
105
+ )
106
elif current_platform.is_cpu():
107
arch = current_platform.get_cpu_architecture()
108
# Fall back to native implementation for POWERPC and RISCV.
@@ -417,7 +447,7 @@ def flashinfer_sample(
447
logits: torch.Tensor,
448
k: torch.Tensor | None,
449
p: torch.Tensor | None,
420
- generators: dict[int, torch.Generator],
450
+ generators: dict[int, torch.Generator] = {}, # noqa
451
) -> torch.Tensor:
452
"""Sample from the logits using FlashInfer.
453
@@ -431,11 +461,6 @@ def flashinfer_sample(
461
"""
462
import flashinfer
463
434
- if version.parse(flashinfer.__version__) < version.parse("0.2.3"):
435
- raise ImportError(
436
- "FlashInfer version >= 0.2.3 required for top-k and top-p sampling. "
437
- )
438
-
464
assert not (k is None and p is None)
465
if k is None:
466
# Top-p only.
vllm/v1/worker/gpu/sample/sampler.py
+46
-14
@@ -7,6 +7,11 @@ 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 (
11
+ apply_top_k_top_p,
12
+ flashinfer_sample,
13
+ flashinfer_sampler_supported,
14
+)
15
from vllm.v1.worker.gpu.input_batch import InputBatch
16
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
17
from vllm.v1.worker.gpu.sample.bad_words import BadWordsState
@@ -45,6 +50,7 @@ class Sampler:
50
self.bad_words_state = BadWordsState(req_states)
51
self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
52
self.num_speculative_tokens = num_speculative_tokens
53
+ self.use_flashinfer = flashinfer_sampler_supported()
54
55
def add_request(
56
self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
@@ -77,6 +83,13 @@ class Sampler:
83
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
84
# that num_nans is computed before applying penalties and temperature.
85
num_nans = get_num_nans(logits) if self.compute_nans else None
86
+
87
+ max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
88
+ max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(
89
+ idx_mapping_np
90
+ )
91
+ return_logprobs = max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0
92
+
93
sampled, processed_logits = self.sample(
94
logits,
95
expanded_idx_mapping,
@@ -84,13 +97,10 @@ class Sampler:
97
pos,
98
input_ids,
99
expanded_local_pos,
100
+ return_logprobs=return_logprobs,
101
)
102
89
- max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
90
- max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(
91
- idx_mapping_np
92
- )
93
- if max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0:
103
+ if return_logprobs:
104
if self.logprobs_mode == "processed_logprobs":
105
logits = processed_logits
106
expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
@@ -128,6 +138,7 @@ class Sampler:
138
pos: torch.Tensor,
139
input_ids: torch.Tensor,
140
expanded_local_pos: torch.Tensor,
141
+ skip_top_k_top_p: bool = False,
142
) -> torch.Tensor:
143
# Copy logits to a new FP32 tensor.
144
logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
@@ -163,6 +174,9 @@ class Sampler:
174
# Apply min_p in place.
175
self.sampling_states.apply_min_p(logits, expanded_idx_mapping, idx_mapping_np)
176
177
+ if skip_top_k_top_p:
178
+ return logits
179
+
180
# Apply top_k and/or top_p. This might or might not return a new tensor.
181
return self.sampling_states.apply_top_k_top_p(
182
logits, expanded_idx_mapping, idx_mapping_np
@@ -176,6 +190,7 @@ class Sampler:
190
pos: torch.Tensor,
191
input_ids: torch.Tensor,
192
expanded_local_pos: torch.Tensor,
193
+ return_logprobs: bool = False,
194
) -> tuple[torch.Tensor, torch.Tensor]:
195
processed_logits = self.apply_sampling_params(
196
logits,
@@ -184,16 +199,33 @@ class Sampler:
199
pos,
200
input_ids,
201
expanded_local_pos,
202
+ skip_top_k_top_p=True,
203
+ )
204
+ top_k, top_p = self.sampling_states.get_top_k_top_p(
205
+ expanded_idx_mapping, idx_mapping_np
206
+ )
207
+ use_flashinfer = self.use_flashinfer and not (
208
+ # Don't use FI sampler if no requests use top_k/top_p, if there are
209
+ # any greedy requests or per-request seeds, or if post-processed
210
+ # logprobs need to be returned for any requests.
211
+ (top_k is None and top_p is None)
212
+ or (return_logprobs and self.logprobs_mode == "processed_logprobs")
213
+ or self.sampling_states.any_greedy(idx_mapping_np)
214
+ or self.sampling_states.any_explicit_seed(idx_mapping_np)
215
)
216
217
# Sample the next token.
190
- sampled = gumbel_sample(
191
- processed_logits,
192
- expanded_idx_mapping,
193
- self.sampling_states.temperature.gpu,
194
- self.sampling_states.seeds.gpu,
195
- pos,
196
- apply_temperature=False,
197
- use_fp64=self.use_fp64_gumbel,
198
- )
218
+ if use_flashinfer:
219
+ sampled = flashinfer_sample(processed_logits, top_k, top_p).to(torch.int64)
220
+ else:
221
+ processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
222
+ sampled = gumbel_sample(
223
+ processed_logits,
224
+ expanded_idx_mapping,
225
+ self.sampling_states.temperature.gpu,
226
+ self.sampling_states.seeds.gpu,
227
+ pos,
228
+ apply_temperature=False,
229
+ use_fp64=self.use_fp64_gumbel,
230
+ )
231
return sampled, processed_logits
vllm/v1/worker/gpu/sample/states.py
+21
-6
@@ -24,6 +24,9 @@ class SamplingStates:
24
self.top_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
25
self.min_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
26
self.seeds = UvaBackedTensor(max_num_reqs, dtype=torch.int64)
27
+ # Tracks whether `seed` was set explicitly by the user, so callers
28
+ # can fall back from RNG paths that don't honor per-request seeds.
29
+ self.seeds_set = np.zeros(max_num_reqs, dtype=bool)
30
31
# Initialize top_k and top_p manually because 0 is an invalid value for them.
32
self.top_k.np.fill(self.vocab_size)
@@ -45,6 +48,7 @@ class SamplingStates:
48
self.min_p.np[req_idx] = sampling_params.min_p
49
50
seed = sampling_params.seed
51
+ self.seeds_set[req_idx] = seed is not None
52
if seed is None:
53
seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)
54
self.seeds.np[req_idx] = seed
@@ -85,20 +89,31 @@ class SamplingStates:
89
return
90
apply_min_p(logits, expanded_idx_mapping, self.min_p.gpu)
91
92
+ def get_top_k_top_p(
93
+ self, expanded_idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray
94
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
95
+ do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size)
96
+ do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0)
97
+ top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None
98
+ top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None
99
+ return top_k, top_p
100
+
101
def apply_top_k_top_p(
102
self,
103
logits: torch.Tensor,
104
expanded_idx_mapping: torch.Tensor,
105
idx_mapping_np: np.ndarray,
106
) -> 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):
107
+ top_k, top_p = self.get_top_k_top_p(expanded_idx_mapping, idx_mapping_np)
108
+ if top_k is None and top_p is None:
109
return logits
98
-
99
- top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None
100
- top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None
110
return apply_top_k_top_p(logits, top_k, top_p)
111
112
+ def any_greedy(self, idx_mapping_np: np.ndarray) -> bool:
113
+ return bool(np.any(self.temperature.np[idx_mapping_np] == 0.0))
114
+
115
+ def any_explicit_seed(self, idx_mapping_np: np.ndarray) -> bool:
116
+ return bool(np.any(self.seeds_set[idx_mapping_np]))
117
+
118
def max_num_logprobs(self, idx_mapping_np: np.ndarray) -> int:
119
return int(np.max(self.num_logprobs[idx_mapping_np]))