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, 08:30 UTC
[V1][Sampler] Improve performance of FlashInfer sampling by sampling logits instead of probs (#18608)
Lukas Geiger committed
May 26, 2025 at 16:49 UTC
e7523c2e031bc96740723ab63833d1cf94229ab4
1 file changed
+9
-8
vllm/v1/sample/ops/topk_topp_sampler.py
+9
-8
@@ -89,18 +89,18 @@ class TopKTopPSampler(nn.Module):
89
p: Optional[torch.Tensor],
90
) -> torch.Tensor:
91
"""More optimized implementation for top-k and top-p sampling."""
92
- probs = logits.softmax(dim=-1, dtype=torch.float32)
92
if k is None and p is None:
93
# We prefer `random_sample` over `flashinfer_sample` when sorting is
94
# not needed. This is because `random_sample` does not require
95
# CPU-GPU synchronization while `flashinfer_sample` does.
96
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
97
return random_sample(probs, generators)
98
if generators:
99
logger.warning("FlashInfer 0.2.3+ does not support "
100
"per-request generators. Falling back to "
101
"PyTorch-native implementation.")
102
return self.forward_native(logits, generators, k, p)
103
- return flashinfer_sample(probs, k, p, generators)
103
+ return flashinfer_sample(logits, k, p, generators)
104
105
def forward_tpu(
106
self,
@@ -254,17 +254,17 @@ def random_sample(
254
255
256
def flashinfer_sample(
257
- probs: torch.Tensor,
257
+ logits: torch.Tensor,
258
k: Optional[torch.Tensor],
259
p: Optional[torch.Tensor],
260
generators: dict[int, torch.Generator],
261
) -> torch.Tensor:
262
- """Sample from the probabilities using FlashInfer.
262
+ """Sample from the logits using FlashInfer.
263
264
Statistically, this function is equivalent to the `random_sample` function.
265
However, this function is faster because it avoids sorting the logits tensor
266
via rejection sampling.
267
-
267
+
268
NOTE: The outputs of this function do not necessarily match the outputs of
269
the `random_sample` function. It only guarantees that the outputs are
270
statistically equivalent.
@@ -274,18 +274,19 @@ def flashinfer_sample(
274
the synchronization overhead.
275
"""
276
assert not (k is None and p is None)
277
-
277
if k is None:
278
# Top-p only.
279
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
280
next_token_ids = flashinfer.sampling.top_p_sampling_from_probs(
281
probs, p, deterministic=True)
282
elif p is None:
283
# Top-k only.
284
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
285
next_token_ids = flashinfer.sampling.top_k_sampling_from_probs(
286
probs, k, deterministic=True)
287
else:
288
# Both top-k and top-p.
288
- next_token_ids = (flashinfer.sampling.top_k_top_p_sampling_from_probs(
289
- probs, k, p, deterministic=True))
289
+ next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(
290
+ logits, k, p, deterministic=True)
291
292
return next_token_ids.view(-1)