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 15, 19:45 UTC
[V1] Temporarily disable FlashInfer Rejection Sampler (#14788)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 13, 2025 at 20:40 UTC
32ef4983cd029d613172dbcf1edf91e62920bbc8
2 files changed
+13
-4
vllm/v1/sample/ops/topk_topp_sampler.py
+1
-1
@@ -22,7 +22,7 @@ class TopKTopPSampler(nn.Module):
22
23
def __init__(self):
24
super().__init__()
25
- if current_platform.is_cuda:
25
+ if current_platform.is_cuda():
26
if is_flashinfer_available:
27
if envs.VLLM_USE_FLASHINFER_SAMPLER is not False:
28
# NOTE(woosuk): The V0 sampler doesn't use FlashInfer for
vllm/v1/sample/rejection_sampler.py
+12
-3
@@ -24,9 +24,18 @@ class RejectionSampler(nn.Module):
24
25
def __init__(self):
26
super().__init__()
27
- if current_platform.is_cuda:
27
+ if current_platform.is_cuda():
28
if is_flashinfer_available:
29
if envs.VLLM_USE_FLASHINFER_SAMPLER is not False:
30
+ # FIXME(woosuk): Currently, we have errors when using
31
+ # FlashInfer for rejection sampling. As a workaround, we
32
+ # disable FlashInfer for rejection sampling by default.
33
+ logger.info("Currently, FlashInfer rejection sampler is "
34
+ "disabled because of a bug. Falling back to "
35
+ "the PyTorch-native implementation of "
36
+ "rejection sampling.")
37
+ self.forward_method = self.forward_native
38
+
39
# NOTE(woosuk): The V0 sampler doesn't use FlashInfer for
40
# sampling unless VLLM_USE_FLASHINFER_SAMPLER=1 (i.e., by
41
# default it is unused). For backward compatibility, we set
@@ -35,8 +44,8 @@ class RejectionSampler(nn.Module):
44
# None means False, while in V1, None means True. This is
45
# why we use the condition
46
# `envs.VLLM_USE_FLASHINFER_SAMPLER is not False` here.
38
- logger.info("Using FlashInfer for rejection sampling.")
39
- self.forward_method = self.flashinfer_sample
47
+ # logger.info("Using FlashInfer for rejection sampling.")
48
+ # self.forward_method = self.flashinfer_sample
49
else:
50
logger.warning(
51
"FlashInfer is available, but it is not enabled. "