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 4, 12:30 UTC
[Perf] Enable FlashInfer top-k/top-p sampler by default (#40376)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
Apr 29, 2026 at 18:10 UTC
b92ef9ec5a041b538f44d9584bef0e34bfbeecd1
5 files changed
+336
-23
.buildkite/test_areas/samplers.yaml
+3
-1
@@ -10,7 +10,9 @@ steps:
10
- tests/samplers
11
- tests/conftest.py
12
commands:
13
- - pytest -v -s samplers
13
+ # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
14
+ # values explicitly to still cover the PyTorch-native (Triton) path.
15
+ - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
16
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
17
mirror:
18
amd:
tests/models/language/generation/test_hybrid.py
+1
-1
@@ -881,7 +881,7 @@ def test_apc_common_prefix_same_batch(
881
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
882
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501
883
]
884
- sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=20)
884
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=20)
885
outputs = llm.generate(prompts, sampling_params)
886
for output in outputs:
887
assert "two" in output.outputs[0].text
tests/v1/sample/test_topk_topp_sampler.py
+301
@@ -13,6 +13,30 @@ BATCH_SIZE = 1024
13
VOCAB_SIZE = 128 * 1024
14
15
16
+def _flashinfer_topk_topp_supported() -> bool:
17
+ """True iff the FlashInfer top-k/top-p sampler is usable on this host.
18
+
19
+ Mirrors the gate in `TopKTopPSampler.__init__`: CUDA + flashinfer
20
+ importable + GPU compute capability supported by the FlashInfer
21
+ backend.
22
+ """
23
+ if not current_platform.is_cuda():
24
+ return False
25
+ try:
26
+ import flashinfer # noqa: F401
27
+
28
+ from vllm.v1.attention.backends.flashinfer import FlashInferBackend
29
+ except ImportError:
30
+ return False
31
+ capability = current_platform.get_device_capability()
32
+ if capability is None:
33
+ return False
34
+ return FlashInferBackend.supports_compute_capability(capability)
35
+
36
+
37
+FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
38
+
39
+
40
@pytest.fixture(autouse=True)
41
def reset_default_device():
42
"""
@@ -568,3 +592,280 @@ class TestTritonTopkTopp:
592
finite_in = (logits[i] > float("-inf")).sum().item()
593
if finite_in > 0:
594
assert kept > 0, f"Row {i}: no tokens kept"
595
+
596
+
597
+# =============================================================================
598
+# FlashInfer top-k/top-p robustness tests
599
+# =============================================================================
600
+
601
+
602
+@pytest.mark.skipif(
603
+ not FLASHINFER_TOPK_TOPP_SUPPORTED,
604
+ reason="FlashInfer top-k/top-p sampler requires CUDA "
605
+ "and a GPU with FlashInfer support.",
606
+)
607
+class TestFlashInferTopkToppRobustness:
608
+ """Robustness of FlashInfer top-k / top-p sampling to NaN / Inf logits.
609
+
610
+ The FlashInfer sampler is enabled by default on supported GPUs. A
611
+ single poisoned request (NaN / +Inf / -Inf in row 0) must not:
612
+
613
+ 1. crash or hang the process;
614
+ 2. produce out-of-range token ids (anything outside ``[0, vocab)``);
615
+ 3. corrupt other batch rows — neighbours of a poisoned row must
616
+ still receive valid token ids (regression for cross-row
617
+ corruption in a DP batch where one bad request would otherwise
618
+ poison its peers).
619
+
620
+ The reference is "no crash + valid token ids", not bit-exact equality
621
+ against the PyTorch-native path.
622
+ """
623
+
624
+ BATCH = 8
625
+ VOCAB = 32768
626
+ TOPK = 50
627
+ TOPP = 0.9
628
+
629
+ @pytest.fixture(autouse=True)
630
+ def setup(self):
631
+ torch.set_default_device(DEVICE_TYPE)
632
+ self.generator = Generator(device=DEVICE_TYPE).manual_seed(1234)
633
+
634
+ def _make_logits(self, pattern: str) -> torch.Tensor:
635
+ """Build (BATCH, VOCAB) logits with `pattern` applied to row 0
636
+ (rows 1..B-1 stay clean so we can detect cross-row corruption)."""
637
+ logits = (
638
+ torch.randn(
639
+ self.BATCH,
640
+ self.VOCAB,
641
+ generator=self.generator,
642
+ dtype=torch.float32,
643
+ )
644
+ * 5.0
645
+ )
646
+ if pattern == "clean":
647
+ return logits
648
+ if pattern == "nan_one_row":
649
+ logits[0, :] = float("nan")
650
+ elif pattern == "nan_few":
651
+ # Scatter 16 NaNs across row 0, keep the rest finite.
652
+ idx = torch.randperm(self.VOCAB, generator=self.generator)[:16]
653
+ logits[0, idx] = float("nan")
654
+ elif pattern == "nan_at_top":
655
+ # Poison the top-32 highest-scoring positions of row 0 — worst
656
+ # case for top-k since these are exactly the tokens that would
657
+ # otherwise be selected. Use argsort instead of topk to avoid
658
+ # a known compute-sanitizer false positive in mbtopk.
659
+ top_idx = logits[0].argsort(descending=True)[:32]
660
+ logits[0, top_idx] = float("nan")
661
+ elif pattern == "nan_all_rows":
662
+ logits[:, :] = float("nan")
663
+ elif pattern == "pos_inf_one_row":
664
+ logits[0, :] = float("inf")
665
+ elif pattern == "neg_inf_one_row":
666
+ logits[0, :] = float("-inf")
667
+ elif pattern == "mixed_inf_nan":
668
+ assert self.BATCH >= 3
669
+ logits[0, :] = float("nan")
670
+ logits[1, :] = float("inf")
671
+ logits[2, :] = float("-inf")
672
+ elif pattern == "degenerate_flat":
673
+ logits[:, :] = 1.0
674
+ else:
675
+ raise ValueError(f"unknown pattern: {pattern}")
676
+ return logits
677
+
678
+ def _check_tokens(self, tokens: torch.Tensor, ctx: str):
679
+ assert tokens.dim() == 1, f"{ctx}: expected 1-D output, got {tokens.shape}"
680
+ assert tokens.shape[0] == self.BATCH, (
681
+ f"{ctx}: expected batch size {self.BATCH}, got {tokens.shape[0]}"
682
+ )
683
+ ids = tokens.tolist()
684
+ min_id, max_id = min(ids), max(ids)
685
+ assert 0 <= min_id < self.VOCAB and 0 <= max_id < self.VOCAB, (
686
+ f"{ctx}: token id(s) outside [0, {self.VOCAB}): min={min_id}, max={max_id}"
687
+ )
688
+
689
+ @pytest.mark.parametrize(
690
+ "pattern",
691
+ [
692
+ "clean",
693
+ "nan_one_row",
694
+ "nan_few",
695
+ "nan_at_top",
696
+ "nan_all_rows",
697
+ "pos_inf_one_row",
698
+ "neg_inf_one_row",
699
+ "mixed_inf_nan",
700
+ "degenerate_flat",
701
+ ],
702
+ )
703
+ @pytest.mark.parametrize("path", ["topk_only", "topp_only", "topk_topp"])
704
+ def test_flashinfer_handles_pathological_logits(self, pattern: str, path: str):
705
+ """flashinfer_sample must return valid ids even on poisoned logits.
706
+
707
+ Direct call into ``flashinfer_sample`` — exactly the code path
708
+ ``TopKTopPSampler.forward_cuda`` takes when FI is enabled.
709
+ """
710
+ from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sample
711
+
712
+ logits = self._make_logits(pattern)
713
+ k = (
714
+ torch.full(
715
+ (self.BATCH,),
716
+ self.TOPK,
717
+ device=DEVICE_TYPE,
718
+ dtype=torch.int32,
719
+ )
720
+ if path in ("topk_only", "topk_topp")
721
+ else None
722
+ )
723
+ p = (
724
+ torch.full(
725
+ (self.BATCH,),
726
+ self.TOPP,
727
+ device=DEVICE_TYPE,
728
+ dtype=torch.float32,
729
+ )
730
+ if path in ("topp_only", "topk_topp")
731
+ else None
732
+ )
733
+
734
+ # flashinfer_sample may mutate its input in-place; pass a clone so
735
+ # the parametrize iterations stay independent.
736
+ tokens = flashinfer_sample(logits.clone().contiguous(), k, p, {})
737
+ # Surface any async CUDA error synchronously (e.g. illegal memory
738
+ # access from a malformed FlashInfer call) so it's attributed to
739
+ # this test rather than a later, unrelated GPU op.
740
+ torch.accelerator.synchronize()
741
+ self._check_tokens(tokens, ctx=f"pattern={pattern}, path={path}")
742
+
743
+
744
+# =============================================================================
745
+# FlashInfer top-k/top-p distribution-match tests
746
+# =============================================================================
747
+
748
+
749
+@pytest.mark.skipif(
750
+ not FLASHINFER_TOPK_TOPP_SUPPORTED,
751
+ reason="FlashInfer top-k/top-p sampler requires CUDA "
752
+ "and a GPU with FlashInfer support.",
753
+)
754
+class TestFlashInferDistributionMatch:
755
+ """Chi-square goodness-of-fit: FlashInfer and PyTorch-native samplers
756
+ both reproduce the expected token distribution after top-k / top-p.
757
+
758
+ Regression guard against historical FlashInfer distribution-shift.
759
+ Each impl is compared to the theoretical distribution (softmax of
760
+ filtered logits); if both pass they are statistically equivalent
761
+ to each other by transitivity.
762
+ """
763
+
764
+ VOCAB = 32
765
+ N_SAMPLES = 50_000
766
+ ALPHA = 1e-6
767
+ SEED = 0
768
+
769
+ @pytest.mark.parametrize(
770
+ "topk,topp",
771
+ [
772
+ (8, None),
773
+ (16, None),
774
+ (None, 0.5),
775
+ (None, 0.7),
776
+ (None, 0.99),
777
+ (8, 0.9),
778
+ (4, 0.5),
779
+ ],
780
+ )
781
+ def test_distribution_matches_theoretical(self, topk, topp):
782
+ from scipy.stats import chisquare
783
+
784
+ from vllm.v1.sample.ops.topk_topp_sampler import (
785
+ apply_top_k_top_p,
786
+ flashinfer_sample,
787
+ random_sample,
788
+ )
789
+
790
+ torch.set_default_device(DEVICE_TYPE)
791
+ torch.manual_seed(self.SEED)
792
+
793
+ # Same logits row used for both impls so the comparison is fair.
794
+ logits_one = (
795
+ torch.randn(
796
+ (1, self.VOCAB),
797
+ dtype=torch.float32,
798
+ )
799
+ * 2.0
800
+ )
801
+
802
+ # Theoretical expected distribution from PyTorch-native filter.
803
+ k_one = torch.tensor([topk], dtype=torch.int32) if topk is not None else None
804
+ p_one = torch.tensor([topp], dtype=torch.float32) if topp is not None else None
805
+ masked = apply_top_k_top_p_pytorch(logits_one.clone(), k_one, p_one)
806
+ expected_probs = masked.softmax(dim=-1).flatten().cpu().numpy()
807
+ expected_counts = expected_probs * self.N_SAMPLES
808
+
809
+ # Build a batch of N identical rows for both impls.
810
+ batch = logits_one.expand(self.N_SAMPLES, self.VOCAB).contiguous()
811
+ k_batch = (
812
+ torch.full((self.N_SAMPLES,), topk, dtype=torch.int32)
813
+ if topk is not None
814
+ else None
815
+ )
816
+ p_batch = (
817
+ torch.full((self.N_SAMPLES,), topp, dtype=torch.float32)
818
+ if topp is not None
819
+ else None
820
+ )
821
+
822
+ # FlashInfer dispatch path.
823
+ fi_tokens = flashinfer_sample(batch.contiguous(), k_batch, p_batch, {})
824
+ fi_counts = torch.bincount(fi_tokens, minlength=self.VOCAB).cpu().numpy()
825
+ self._chi2_check(
826
+ fi_counts,
827
+ expected_counts,
828
+ chisquare,
829
+ label=f"flashinfer top-k={topk} top-p={topp}",
830
+ )
831
+
832
+ # PyTorch-native dispatch path (Triton-routed filter + Gumbel sample).
833
+ processed = apply_top_k_top_p(batch.clone(), k_batch, p_batch)
834
+ probs = processed.softmax(dim=-1, dtype=torch.float32)
835
+ pt_tokens = random_sample(probs, {})
836
+ pt_counts = torch.bincount(pt_tokens, minlength=self.VOCAB).cpu().numpy()
837
+ self._chi2_check(
838
+ pt_counts,
839
+ expected_counts,
840
+ chisquare,
841
+ label=f"native top-k={topk} top-p={topp}",
842
+ )
843
+
844
+ def _chi2_check(self, empirical, expected, chisquare_fn, *, label):
845
+ import numpy as np
846
+
847
+ # Hard check: the sampler must never produce a token outside the
848
+ # expected support (zero theoretical probability).
849
+ outside = (expected == 0) & (empirical > 0)
850
+ assert not outside.any(), (
851
+ f"{label}: sampled out-of-support tokens "
852
+ f"(zero expected prob): indices={outside.nonzero()[0].tolist()}"
853
+ )
854
+ # Skip chi-square in the degenerate case where the support
855
+ # collapses to a single token (e.g. very restrictive joint
856
+ # top-k + top-p): all samples must land there and the hard
857
+ # check above already verified they do.
858
+ in_support = expected > 0
859
+ if int(in_support.sum()) <= 1:
860
+ return
861
+ # Soft check: chi-square goodness-of-fit on in-support tokens.
862
+ # Cast to float64 so the rescaling step below stays within
863
+ # scipy.chisquare's strict 1.5e-8 sum-equality tolerance.
864
+ emp = empirical[in_support].astype(np.float64)
865
+ exp = expected[in_support].astype(np.float64)
866
+ exp = exp * (emp.sum() / exp.sum())
867
+ chi2, p_value = chisquare_fn(emp, exp)
868
+ assert p_value > self.ALPHA, (
869
+ f"{label}: distribution differs from theoretical: "
870
+ f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
871
+ )
vllm/envs.py
+5
-3
@@ -45,7 +45,7 @@ if TYPE_CHECKING:
45
NO_COLOR: bool = False
46
VLLM_LOG_STATS_INTERVAL: float = 10.0
47
VLLM_TRACE_FUNCTION: int = 0
48
- VLLM_USE_FLASHINFER_SAMPLER: bool | None = None
48
+ VLLM_USE_FLASHINFER_SAMPLER: bool = True
49
VLLM_PP_LAYER_PARTITION: str | None = None
50
VLLM_CPU_KVCACHE_SPACE: int | None = 0
51
VLLM_CPU_OMP_THREADS_BIND: str = "auto"
@@ -712,11 +712,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
712
# If set to 1, vllm will trace function calls
713
# Useful for debugging
714
"VLLM_TRACE_FUNCTION": lambda: int(os.getenv("VLLM_TRACE_FUNCTION", "0")),
715
- # If set, vllm will use flashinfer sampler
715
+ # Whether to use the FlashInfer top-k / top-p sampler on CUDA. Enabled
716
+ # by default when the hardware supports it — set to 0 to opt out
717
+ # explicitly, which forces the PyTorch-native (Triton for bs>=8) path.
718
"VLLM_USE_FLASHINFER_SAMPLER": lambda: (
719
bool(int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"]))
720
if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ
719
- else None
721
+ else True
722
),
723
# Pipeline stage partition strategy
724
"VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None),
vllm/v1/sample/ops/topk_topp_sampler.py
+26
-18
@@ -41,23 +41,35 @@ class TopKTopPSampler(nn.Module):
41
42
capability = current_platform.get_device_capability()
43
assert capability is not None
44
- if not FlashInferBackend.supports_compute_capability(capability):
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
)
50
- # Users must opt in explicitly via VLLM_USE_FLASHINFER_SAMPLER=1.
51
- logger.info_once(
52
- "Using FlashInfer for top-p & top-k sampling.",
53
- scope="global",
54
- )
55
- self.forward = self.forward_cuda
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:
57
- logger.debug_once(
58
- "FlashInfer top-p/top-k sampling is available but disabled "
59
- "by default. Set VLLM_USE_FLASHINFER_SAMPLER=1 to opt in "
60
- "after verifying accuracy for your workloads."
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
@@ -120,9 +132,9 @@ class TopKTopPSampler(nn.Module):
132
p: torch.Tensor | None,
133
) -> tuple[torch.Tensor, torch.Tensor | None]:
134
"""More optimized implementation for top-k and top-p sampling."""
123
- # We prefer `random_sample` over `flashinfer_sample` when sorting is
124
- # not needed. This is because `random_sample` does not require
125
- # CPU-GPU synchronization while `flashinfer_sample` does.
135
+ # Fall back to the PyTorch-native path when FlashInfer has nothing
136
+ # to do (no top-k / top-p filter) or when per-request generators
137
+ # are present (unsupported by FlashInfer 0.2.3+).
138
if (k is None and p is None) or generators:
139
if generators:
140
logger.debug_once(
@@ -361,10 +373,6 @@ def flashinfer_sample(
373
NOTE: The outputs of this function do not necessarily match the outputs of
374
the `random_sample` function. It only guarantees that the outputs are
375
statistically equivalent.
364
-
365
- NOTE: This function includes CPU-GPU synchronization, while `random_sample`
366
- does not. Call this function at the end of the forward pass to minimize
367
- the synchronization overhead.
376
"""
377
import flashinfer
378