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
[Bugfix][V1] Fix flashinfer sampling (#14815)
DefTruth committed
Mar 15, 2025 at 11:42 UTC
acaea3bb07883c80b71643ebee1cd08d555797bc
1 file changed
+18
-1
vllm/v1/sample/ops/topk_topp_sampler.py
+18
-1
@@ -24,7 +24,24 @@ class TopKTopPSampler(nn.Module):
24
super().__init__()
25
if current_platform.is_cuda():
26
if is_flashinfer_available:
27
- if envs.VLLM_USE_FLASHINFER_SAMPLER is not False:
27
+ flashinfer_version = flashinfer.__version__
28
+ if flashinfer_version >= "0.2.3":
29
+ # FIXME(DefTruth): Currently, we have errors when using
30
+ # FlashInfer>=v0.2.3 for top-p & top-k sampling. As a
31
+ # workaround, we disable FlashInfer for top-p & top-k
32
+ # sampling by default while FlashInfer>=v0.2.3.
33
+ # The sampling API removes the success return value
34
+ # of all sampling API, which is not compatible with
35
+ # earlier design.
36
+ # https://github.com/flashinfer-ai/flashinfer/releases/
37
+ # tag/v0.2.3
38
+ logger.info(
39
+ "Currently, FlashInfer top-p & top-k sampling sampler "
40
+ "is disabled because FlashInfer>=v0.2.3 is not "
41
+ "backward compatible. Falling back to the PyTorch-"
42
+ "native implementation of top-p & top-k sampling.")
43
+ self.forward = self.forward_native
44
+ elif envs.VLLM_USE_FLASHINFER_SAMPLER is not False:
45
# NOTE(woosuk): The V0 sampler doesn't use FlashInfer for
46
# sampling unless VLLM_USE_FLASHINFER_SAMPLER=1 (i.e., by
47
# default it is unused). For backward compatibility, we set