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 14, 05:30 UTC
[Kernel] Lazy import FlashInfer (#26977)
Jee Jee Li committed
Oct 17, 2025 at 12:48 UTC
fec2b341ad76d322ffe956fc9581e958bd6887bc
2 files changed
+25
-38
tests/v1/sample/test_topk_topp_sampler.py
+9
-8
@@ -5,20 +5,13 @@ import torch
5
from torch import Generator
6
7
from vllm.platforms import current_platform
8
-from vllm.v1.sample.ops.topk_topp_sampler import (
9
- apply_top_k_top_p,
10
- is_flashinfer_available,
11
-)
8
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
9
10
DEVICE = current_platform.device_type
11
12
BATCH_SIZE = 1024
13
VOCAB_SIZE = 128 * 1024
14
18
-FLASHINFER_ENABLED = current_platform.is_cuda() and is_flashinfer_available
19
-if is_flashinfer_available:
20
- from flashinfer.sampling import top_k_renorm_probs, top_p_renorm_probs
21
-
15
16
@pytest.fixture(autouse=True)
17
def reset_default_device():
@@ -65,6 +58,14 @@ def test_flashinfer_sampler():
58
sampling results due to randomness), so we will compare the probability
59
renormed consequently by top-k and then top-p of FlashInfer implementation.
60
"""
61
+ try:
62
+ from flashinfer.sampling import top_k_renorm_probs, top_p_renorm_probs
63
+
64
+ is_flashinfer_available = True
65
+ except ImportError:
66
+ is_flashinfer_available = False
67
+
68
+ FLASHINFER_ENABLED = current_platform.is_cuda() and is_flashinfer_available
69
70
if not FLASHINFER_ENABLED:
71
pytest.skip("FlashInfer not installed or not available on this platform.")
vllm/v1/sample/ops/topk_topp_sampler.py
+16
-30
@@ -13,13 +13,6 @@ from vllm.platforms import CpuArchEnum, current_platform
13
14
logger = init_logger(__name__)
15
16
-try:
17
- import flashinfer.sampling
18
-
19
- is_flashinfer_available = True
20
-except ImportError:
21
- is_flashinfer_available = False
22
-
16
17
class TopKTopPSampler(nn.Module):
18
"""
@@ -38,32 +31,18 @@ class TopKTopPSampler(nn.Module):
31
logprobs_mode not in ("processed_logits", "processed_logprobs")
32
and current_platform.is_cuda()
33
):
41
- if is_flashinfer_available:
42
- flashinfer_version = flashinfer.__version__
43
- if version.parse(flashinfer_version) < version.parse("0.2.3"):
44
- logger.warning_once(
45
- "FlashInfer version >= 0.2.3 required. "
46
- "Falling back to default sampling implementation."
47
- )
48
- self.forward = self.forward_native
49
- elif envs.VLLM_USE_FLASHINFER_SAMPLER:
50
- # Users must opt in explicitly via VLLM_USE_FLASHINFER_SAMPLER=1.
51
- logger.info_once("Using FlashInfer for top-p & top-k sampling.")
52
- self.forward = self.forward_cuda
53
- else:
54
- logger.debug_once(
55
- "FlashInfer top-p/top-k sampling is available but disabled "
56
- "by default. Set VLLM_USE_FLASHINFER_SAMPLER=1 to opt in "
57
- "after verifying accuracy for your workloads."
58
- )
59
- self.forward = self.forward_native
34
+ if envs.VLLM_USE_FLASHINFER_SAMPLER:
35
+ # Users must opt in explicitly via VLLM_USE_FLASHINFER_SAMPLER=1.
36
+ logger.info_once("Using FlashInfer for top-p & top-k sampling.")
37
+ self.forward = self.forward_cuda
38
else:
61
- logger.warning_once(
62
- "FlashInfer is not available. Falling back to the PyTorch-"
63
- "native implementation of top-p & top-k sampling. For the "
64
- "best performance, please install FlashInfer."
39
+ logger.debug_once(
40
+ "FlashInfer top-p/top-k sampling is available but disabled "
41
+ "by default. Set VLLM_USE_FLASHINFER_SAMPLER=1 to opt in "
42
+ "after verifying accuracy for your workloads."
43
)
44
self.forward = self.forward_native
45
+
46
elif current_platform.is_cpu():
47
arch = current_platform.get_cpu_architecture()
48
# Fall back to native implementation for POWERPC and RISCV.
@@ -278,6 +257,13 @@ def flashinfer_sample(
257
does not. Call this function at the end of the forward pass to minimize
258
the synchronization overhead.
259
"""
260
+ import flashinfer
261
+
262
+ if version.parse(flashinfer.__version__) < version.parse("0.2.3"):
263
+ raise ImportError(
264
+ "FlashInfer version >= 0.2.3 required for top-k and top-p sampling. "
265
+ )
266
+
267
assert not (k is None and p is None)
268
if k is None:
269
# Top-p only.