@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,090 lines 42.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5 from torch import Generator
6
7 from tests.utils import large_gpu_mark
8 from vllm.platforms import current_platform
9 from vllm.triton_utils import HAS_TRITON
10 from vllm.utils.torch_utils import set_random_seed
11 from vllm.v1.sample.ops.topk_topp_sampler import (
12 apply_top_k_top_p_pytorch,
13 random_sample,
14 )
15 from vllm.v1.sample.sampler import Sampler
16
17 DEVICE_TYPE = current_platform.device_type
18
19 BATCH_SIZE = 1024
20 VOCAB_SIZE = 128 * 1024
21
22
23 def _flashinfer_topk_topp_supported() -> bool:
24 """True iff the FlashInfer top-k/top-p sampler is usable on this host.
25
26 Mirrors the gate in `TopKTopPSampler.__init__`: CUDA + flashinfer
27 importable + GPU compute capability supported by the FlashInfer
28 backend + more than 16 SMs.
29 """
30 if not current_platform.is_cuda():
31 return False
32 try:
33 import flashinfer # noqa: F401
34
35 from vllm.v1.attention.backends.flashinfer import FlashInferBackend
36 except ImportError:
37 return False
38 capability = current_platform.get_device_capability()
39 if capability is None:
40 return False
41 return FlashInferBackend.supports_compute_capability(capability) and (
42 current_platform.num_compute_units(torch.accelerator.current_device_index())
43 > 16
44 )
45
46
47 FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
48
49
50 def _seed_default_generator(seed: int) -> None:
51 set_random_seed(seed)
52
53
54 @pytest.fixture(autouse=True)
55 def reset_default_device():
56 """Explicitly set the default device, which can affect subsequent tests.
57 Adding this fixture helps avoid this problem.
58 """
59 original_device = torch.get_default_device()
60 yield
61 torch.set_default_device(original_device)
62
63
64 def test_sampler_threads_fp64_gumbel_to_topk_topp_sampler():
65 sampler = Sampler(use_fp64_gumbel=True)
66
67 assert sampler.topk_topp_sampler.use_fp64_gumbel
68
69
70 @pytest.mark.skipif(
71 not current_platform.is_rocm(),
72 reason="ROCm aiter sampler test only runs on ROCm",
73 )
74 def test_rocm_aiter_sampler_defers_import_when_generators_force_native(
75 monkeypatch: pytest.MonkeyPatch,
76 ):
77 from vllm.v1.sample.ops import topk_topp_sampler
78
79 class MockPlatform:
80 @staticmethod
81 def is_cuda():
82 return False
83
84 @staticmethod
85 def is_cuda_alike():
86 return True
87
88 @staticmethod
89 def is_cpu():
90 return False
91
92 @staticmethod
93 def is_xpu():
94 return False
95
96 class MockRocmAiterOps:
97 @staticmethod
98 def is_enabled():
99 return True
100
101 real_import = __import__
102
103 def guard_aiter_sampling_import(name, *args, **kwargs):
104 if name == "aiter.ops.sampling":
105 raise AssertionError("aiter sampling import should be deferred")
106 return real_import(name, *args, **kwargs)
107
108 monkeypatch.setattr(topk_topp_sampler, "current_platform", MockPlatform())
109 monkeypatch.setattr(topk_topp_sampler, "rocm_aiter_ops", MockRocmAiterOps())
110 monkeypatch.setattr("builtins.__import__", guard_aiter_sampling_import)
111
112 sampler = topk_topp_sampler.TopKTopPSampler()
113 logits = torch.randn(2, 8)
114 k = torch.full((2,), 2, dtype=torch.int32)
115 generators = {0: torch.Generator(device=logits.device).manual_seed(0)}
116
117 token_ids, logits_to_return = sampler(logits, generators, k, None)
118
119 assert token_ids.shape == (2,)
120 assert logits_to_return is None
121
122
123 def test_random_sample_uses_fp64_exponential_race_when_requested():
124 torch.set_default_device(DEVICE_TYPE)
125 probs = torch.tensor(
126 [
127 [0.70, 0.20, 0.10],
128 [0.05, 0.15, 0.80],
129 [0.25, 0.25, 0.50],
130 ],
131 dtype=torch.float32,
132 device=DEVICE_TYPE,
133 )
134
135 _seed_default_generator(12345)
136 q = torch.empty(probs.shape, dtype=torch.float64, device=probs.device)
137 q.exponential_()
138 expected = q.reciprocal_().mul_(probs).argmax(dim=-1).view(-1)
139
140 _seed_default_generator(12345)
141 actual = random_sample(probs.clone(), {}, use_fp64_gumbel=True)
142
143 assert torch.equal(actual, expected)
144
145
146 def test_topk_impl_equivalence():
147 torch.set_default_device(DEVICE_TYPE)
148 generator = Generator(device=DEVICE_TYPE).manual_seed(33)
149
150 logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator)
151
152 # Random top-k values between 1 and 9.
153 k = torch.randint(1, 10, (BATCH_SIZE,), generator=generator)
154
155 # Set k=vocab_size for ~50% of requests in the batch (top-k disabled).
156 k.masked_fill_(
157 torch.randint(0, 2, (BATCH_SIZE,), generator=generator, dtype=bool), VOCAB_SIZE
158 )
159
160 # Top-k only implementation
161 result1 = apply_top_k_top_p_pytorch(logits=logits.clone(), k=k, p=None)
162
163 # Top-p + top-k
164 no_op_top_p = torch.tensor([1.0])
165 result2 = apply_top_k_top_p_pytorch(logits=logits.clone(), k=k, p=no_op_top_p)
166
167 assert torch.allclose(result1, result2)
168
169
170 @pytest.mark.skipif(
171 not FLASHINFER_TOPK_TOPP_SUPPORTED,
172 reason="Requires CUDA and flashinfer",
173 )
174 def test_flashinfer_sampler_fallback_when_jit_cannot_target_gpu(monkeypatch):
175 """FlashInfer swallows arch-detection errors (e.g. SM 12.x with a CUDA
176 toolkit older than 12.9), leaving an empty target-arch set that makes
177 every JIT spec fail with a misleading "requires sm75 or higher" error at
178 first use, killing the engine during startup profiling (issue #42393).
179 The sampler gate must detect this and fall back to native sampling, and
180 surface the real reason when the user explicitly opted in.
181 """
182 import flashinfer.jit.core
183 from flashinfer.compilation_context import CompilationContext
184
185 from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sampler_supported
186
187 def fake_check_cuda_arch():
188 raise RuntimeError("FlashInfer requires GPUs with sm75 or higher")
189
190 def fake_normalize_cuda_arch(major, minor):
191 raise RuntimeError("SM 12.x requires CUDA >= 12.9")
192
193 monkeypatch.setattr(flashinfer.jit.core, "check_cuda_arch", fake_check_cuda_arch)
194 monkeypatch.setattr(
195 CompilationContext,
196 "_normalize_cuda_arch",
197 staticmethod(fake_normalize_cuda_arch),
198 )
199
200 monkeypatch.delenv("VLLM_USE_FLASHINFER_SAMPLER", raising=False)
201 assert flashinfer_sampler_supported() is False
202
203 monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "1")
204 with pytest.raises(RuntimeError, match="SM 12.x requires CUDA >= 12.9"):
205 flashinfer_sampler_supported()
206
207
208 # =============================================================================
209 # Triton kernel tests
210 # =============================================================================
211
212
213 @pytest.mark.skipif(not HAS_TRITON, reason="Triton not available on this platform")
214 class TestTritonTopkTopp:
215 """Tests for the Triton top-k/top-p kernel."""
216
217 @pytest.fixture(autouse=True)
218 def setup(self):
219 """Set up test fixtures."""
220 torch.set_default_device(DEVICE_TYPE)
221 self.generator = Generator(device=DEVICE_TYPE).manual_seed(42)
222
223 def _compare_results(
224 self,
225 logits: torch.Tensor,
226 k: torch.Tensor | None,
227 p: torch.Tensor | None,
228 ):
229 """Compare Triton kernel results with PyTorch sorting implementation.
230
231 For top-k only, we expect exact match.
232 For top-p (with or without top-k), we allow small differences due to
233 floating-point precision in probability sum calculations.
234 """
235 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
236
237 # Clone logits for both implementations
238 logits_pytorch = logits.clone()
239 logits_triton = logits.clone().to(torch.float32)
240
241 # Apply PyTorch sorting implementation
242 result_pytorch = apply_top_k_top_p_pytorch(logits_pytorch, k, p)
243
244 # Apply Triton kernel
245 k_i32 = k.to(torch.int32) if k is not None else None
246 p_f32 = p.to(torch.float32) if p is not None else None
247 result_triton = apply_top_k_top_p_triton(logits_triton, k_i32, p_f32)
248
249 # Compare kept counts per row
250 pytorch_kept = (result_pytorch != float("-inf")).sum(dim=-1)
251 triton_kept = (result_triton != float("-inf")).sum(dim=-1)
252
253 if p is None:
254 # Top-k only: expect exact match
255 assert torch.equal(pytorch_kept, triton_kept), (
256 f"Top-k mask mismatch: PyTorch kept {pytorch_kept.tolist()}, "
257 f"Triton kept {triton_kept.tolist()}"
258 )
259 else:
260 # Top-p involved: allow small differences
261 # Either < 1% of kept values OR < 5 values absolute
262 max_diff = (pytorch_kept - triton_kept).abs().max().item()
263 max_kept = pytorch_kept.max().item()
264 if max_kept > 0 and max_diff > 3:
265 diff_pct = max_diff / max_kept * 100
266 assert diff_pct < 0.5, (
267 f"Top-p mask difference too large: {diff_pct:.2f}% "
268 f"(max diff {max_diff} values out of {max_kept})"
269 )
270
271 @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
272 @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
273 def test_topk_only(self, batch_size: int, vocab_size: int):
274 """Test top-k only (p=None)."""
275 logits = torch.randn(
276 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
277 )
278 k = torch.randint(
279 1, min(100, vocab_size), (batch_size,), generator=self.generator
280 )
281 # Randomly disable top-k for some rows (~25%)
282 disable_mask = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
283 k.masked_fill_(disable_mask, vocab_size)
284
285 self._compare_results(logits, k, p=None)
286
287 @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
288 @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
289 def test_topp_only(self, batch_size: int, vocab_size: int):
290 """Test top-p only (k=None)."""
291 logits = torch.randn(
292 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
293 )
294 p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 # [0.1, 1.0]
295 # Randomly disable top-p for some rows (~25%)
296 disable_mask = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
297 p.masked_fill_(disable_mask, 1.0)
298
299 self._compare_results(logits, k=None, p=p)
300
301 @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
302 @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
303 def test_topk_and_topp(self, batch_size: int, vocab_size: int):
304 """Test combined top-k and top-p."""
305 logits = torch.randn(
306 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
307 )
308 k = torch.randint(
309 1, min(100, vocab_size), (batch_size,), generator=self.generator
310 )
311 p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 # [0.1, 1.0]
312
313 # Randomly disable top-k for some rows (~25%)
314 disable_k = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
315 k.masked_fill_(disable_k, vocab_size)
316 # Randomly disable top-p for some rows (~25%)
317 disable_p = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
318 p.masked_fill_(disable_p, 1.0)
319
320 self._compare_results(logits, k, p)
321
322 def test_both_disabled(self):
323 """Test when both k and p are None (should be no-op)."""
324 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
325
326 logits = torch.randn(32, 1024, generator=self.generator, dtype=torch.float32)
327 logits_clone = logits.clone()
328
329 result = apply_top_k_top_p_triton(logits_clone, k=None, p=None)
330
331 assert torch.equal(result, logits), "Should be no-op when both k and p are None"
332
333 def test_extreme_k_values(self):
334 """Test edge cases for k values."""
335 batch_size, vocab_size = 16, 1024
336 logits = torch.randn(
337 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
338 )
339
340 # k=1 (keep only top 1)
341 k = torch.ones(batch_size, dtype=torch.int32)
342 self._compare_results(logits.clone(), k, p=None)
343
344 # k=vocab_size (keep all)
345 k = torch.full((batch_size,), vocab_size, dtype=torch.int32)
346 self._compare_results(logits.clone(), k, p=None)
347
348 # Mixed extreme values
349 k = torch.tensor([1, vocab_size, 2, vocab_size - 1] * 4, dtype=torch.int32)
350 self._compare_results(logits.clone(), k, p=None)
351
352 def test_extreme_p_values(self):
353 """Test edge cases for p values."""
354 batch_size, vocab_size = 16, 1024
355 logits = torch.randn(
356 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
357 )
358
359 # p close to 0 (very restrictive)
360 p = torch.full((batch_size,), 0.01, dtype=torch.float32)
361 self._compare_results(logits.clone(), k=None, p=p)
362
363 # p=1.0 (keep all)
364 p = torch.ones(batch_size, dtype=torch.float32)
365 self._compare_results(logits.clone(), k=None, p=p)
366
367 # Mixed values
368 p = torch.tensor([0.1, 0.5, 0.9, 1.0] * 4, dtype=torch.float32)
369 self._compare_results(logits.clone(), k=None, p=p)
370
371 def test_large_batch(self):
372 """Test with a large batch size."""
373 batch_size, vocab_size = 512, 32000
374 logits = torch.randn(
375 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
376 )
377 k = torch.randint(1, 50, (batch_size,), generator=self.generator)
378 p = torch.rand(batch_size, generator=self.generator) * 0.5 + 0.5
379
380 self._compare_results(logits, k, p)
381
382 @large_gpu_mark(min_gb=24)
383 def test_large_batch_int64_row_offset(self):
384 """Regression: per-row offset (row * vocab_size) must not overflow int32.
385
386 Speculative decoding expands the logits batch (e.g. DFlash drafts K
387 tokens per request), so batch_size * vocab_size can exceed 2**31. With
388 int32 offset arithmetic the per-row pointer wraps to a negative address
389 and the kernel hits a CUDA illegal memory access. Use a batch where
390 batch_size * vocab_size > 2**31 and give the highest-offset row the same
391 logits as row 0: an overflow there would read a different row and change
392 the kept set.
393 """
394 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
395
396 if not current_platform.is_cuda():
397 pytest.skip("int32 row-offset overflow is a CUDA kernel issue")
398 vocab_size = 131072
399 batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31
400 # logits is modified in place; the only extra device memory is the
401 # per-SM scratch buffer (~num_sm * vocab), so allow ~1 GB of headroom.
402 required_bytes = batch_size * vocab_size * 4 + (1 << 30)
403 if torch.accelerator.get_memory_info()[0] < required_bytes:
404 pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")
405
406 logits = torch.randn(
407 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
408 )
409 logits[batch_size - 1] = logits[0]
410 k = torch.full((batch_size,), 5, dtype=torch.int32)
411 result = apply_top_k_top_p_triton(logits, k, None)
412 torch.accelerator.synchronize() # surface any async illegal memory access
413 kept_first = (result[0] > float("-inf")).nonzero(as_tuple=True)[0]
414 kept_last = (result[batch_size - 1] > float("-inf")).nonzero(as_tuple=True)[0]
415 assert kept_first.numel() == 5, f"row 0 kept {kept_first.numel()}, expected 5"
416 assert torch.equal(kept_first, kept_last), (
417 "highest-offset row produced a different top-k mask than the "
418 "identical row 0 (int32 row-offset overflow)"
419 )
420
421 @pytest.mark.parametrize(
422 "mode",
423 ["topk_only", "topp_only", "topk_and_topp"],
424 )
425 def test_noncontiguous_logits_match_contiguous(self, mode: str):
426 """Non-contiguous logits views should behave like contiguous inputs."""
427 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
428
429 device = torch.device(DEVICE_TYPE)
430 batch_size, vocab_size, pad = 16, 4096, 8
431 backing = torch.full(
432 (batch_size, vocab_size + pad),
433 -1000.0,
434 device=device,
435 dtype=torch.float32,
436 )
437 base = torch.linspace(
438 10.0, -10.0, vocab_size, device=device, dtype=torch.float32
439 )
440 source = base[None, :] + (
441 torch.arange(batch_size, device=device, dtype=torch.float32)[:, None]
442 / 1000.0
443 )
444
445 logits = backing[:, :vocab_size]
446 logits.copy_(source)
447 contig_logits = source.clone()
448 pytorch_logits = source.clone()
449
450 assert logits.shape == (batch_size, vocab_size)
451 assert logits.stride() == (vocab_size + pad, 1)
452 assert not logits.is_contiguous()
453
454 k: torch.Tensor | None = None
455 p: torch.Tensor | None = None
456 if mode in ("topk_only", "topk_and_topp"):
457 k = torch.full((batch_size,), 154, device=device, dtype=torch.int32)
458 if mode in ("topp_only", "topk_and_topp"):
459 p = torch.full((batch_size,), 0.95, device=device, dtype=torch.float32)
460
461 noncontig_out = apply_top_k_top_p_triton(logits, k, p)
462 contig_out = apply_top_k_top_p_triton(contig_logits, k, p)
463 pytorch_out = apply_top_k_top_p_pytorch(pytorch_logits, k, p)
464
465 assert noncontig_out.data_ptr() == logits.data_ptr()
466 assert not noncontig_out.is_contiguous()
467 assert torch.equal(logits, noncontig_out)
468 assert torch.equal(torch.isfinite(noncontig_out), torch.isfinite(contig_out))
469 assert torch.equal(torch.isfinite(noncontig_out), torch.isfinite(pytorch_out))
470
471 # -----------------------------------------------------------------
472 # Tests for -inf logits (e.g. from grammar / structured output masks)
473 # -----------------------------------------------------------------
474
475 @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99])
476 def test_topk_with_neginf_logits(self, inf_fraction: float):
477 """Top-k with many -inf logits (simulating grammar bitmask).
478
479 The kernel must not produce NaN when most logits are -inf, which
480 can happen when structured-output grammar masks are applied before
481 sampling.
482 """
483 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
484
485 batch_size, vocab_size = 32, 128256
486 logits = torch.randn(
487 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
488 )
489 # Mask a fraction of logits to -inf.
490 mask = (
491 torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction
492 )
493 logits[mask] = float("-inf")
494
495 k = torch.randint(
496 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
497 )
498 result = apply_top_k_top_p_triton(logits.clone(), k, None)
499
500 assert not result.isnan().any(), "NaN found in top-k result with -inf logits"
501 for i in range(batch_size):
502 kept = (result[i] > float("-inf")).sum().item()
503 assert kept <= k[i].item(), f"Row {i}: kept {kept} > k={k[i].item()}"
504 # At least one value should survive unless the row was all -inf.
505 finite_in = (logits[i] > float("-inf")).sum().item()
506 if finite_in > 0:
507 assert kept > 0, f"Row {i}: no tokens kept despite finite input"
508
509 @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99])
510 def test_topp_with_neginf_logits(self, inf_fraction: float):
511 """Top-p with many -inf logits."""
512 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
513
514 batch_size, vocab_size = 32, 128256
515 logits = torch.randn(
516 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
517 )
518 mask = (
519 torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction
520 )
521 logits[mask] = float("-inf")
522
523 p = (
524 torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9
525 + 0.1
526 )
527 result = apply_top_k_top_p_triton(logits.clone(), None, p)
528
529 assert not result.isnan().any(), "NaN found in top-p result with -inf logits"
530 for i in range(batch_size):
531 finite_in = (logits[i] > float("-inf")).sum().item()
532 kept = (result[i] > float("-inf")).sum().item()
533 if finite_in > 0:
534 assert kept > 0, f"Row {i}: no tokens kept despite finite input"
535
536 @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99])
537 def test_topk_topp_with_neginf_logits(self, inf_fraction: float):
538 """Combined top-k + top-p with many -inf logits."""
539 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
540
541 batch_size, vocab_size = 32, 128256
542 logits = torch.randn(
543 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
544 )
545 mask = (
546 torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction
547 )
548 logits[mask] = float("-inf")
549
550 k = torch.randint(
551 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
552 )
553 p = (
554 torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9
555 + 0.1
556 )
557 result = apply_top_k_top_p_triton(logits.clone(), k, p)
558
559 assert not result.isnan().any(), (
560 "NaN found in top-k+top-p result with -inf logits"
561 )
562 for i in range(batch_size):
563 kept = (result[i] > float("-inf")).sum().item()
564 assert kept <= k[i].item(), f"Row {i}: kept {kept} > k={k[i].item()}"
565
566 def test_all_neginf_logits(self):
567 """All logits are -inf (fully masked). Kernel should be a no-op."""
568 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
569
570 batch_size, vocab_size = 16, 128256
571 logits = torch.full(
572 (batch_size, vocab_size), float("-inf"), dtype=torch.float32
573 )
574
575 k = torch.randint(
576 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
577 )
578 p = torch.full((batch_size,), 0.9, dtype=torch.float32)
579
580 # top-k only
581 result = apply_top_k_top_p_triton(logits.clone(), k, None)
582 assert not result.isnan().any(), "NaN from all-inf top-k"
583 assert (result == float("-inf")).all(), "Expected all -inf unchanged"
584
585 # top-p only
586 result = apply_top_k_top_p_triton(logits.clone(), None, p)
587 assert not result.isnan().any(), "NaN from all-inf top-p"
588 assert (result == float("-inf")).all(), "Expected all -inf unchanged"
589
590 # top-k + top-p
591 result = apply_top_k_top_p_triton(logits.clone(), k, p)
592 assert not result.isnan().any(), "NaN from all-inf top-k+top-p"
593 assert (result == float("-inf")).all(), "Expected all -inf unchanged"
594
595 def test_few_valid_tokens_with_neginf(self):
596 """Only a handful of tokens are finite per row (strict grammar)."""
597 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
598
599 batch_size, vocab_size = 32, 128256
600 logits = torch.full(
601 (batch_size, vocab_size), float("-inf"), dtype=torch.float32
602 )
603 # Allow only 5 random tokens per row to be finite.
604 for i in range(batch_size):
605 indices = torch.randperm(vocab_size, generator=self.generator)[:5]
606 logits[i, indices] = torch.randn(
607 5, generator=self.generator, dtype=torch.float32
608 )
609
610 k = torch.full((batch_size,), 50, dtype=torch.int32)
611 p = torch.full((batch_size,), 0.9, dtype=torch.float32)
612
613 # top-k only (k=50 but only 5 finite → keep all 5)
614 result = apply_top_k_top_p_triton(logits.clone(), k, None)
615 assert not result.isnan().any()
616 for i in range(batch_size):
617 kept = (result[i] > float("-inf")).sum().item()
618 assert kept == 5, f"Row {i}: expected 5 kept, got {kept}"
619
620 # top-k with k < num_finite
621 k_small = torch.full((batch_size,), 3, dtype=torch.int32)
622 result = apply_top_k_top_p_triton(logits.clone(), k_small, None)
623 assert not result.isnan().any()
624 for i in range(batch_size):
625 kept = (result[i] > float("-inf")).sum().item()
626 assert kept <= 3, f"Row {i}: expected <=3 kept, got {kept}"
627
628 # top-p only
629 result = apply_top_k_top_p_triton(logits.clone(), None, p)
630 assert not result.isnan().any()
631 for i in range(batch_size):
632 kept = (result[i] > float("-inf")).sum().item()
633 assert kept > 0, f"Row {i}: no tokens kept"
634
635 @pytest.mark.parametrize("num_valid", [1, 2, 5, 10, 50])
636 @pytest.mark.parametrize(
637 "mode",
638 ["topk_only", "topp_only", "topk_and_topp"],
639 )
640 def test_equal_logits_few_valid(self, num_valid: int, mode: str):
641 """Few valid tokens all sharing the same logit value.
642
643 This is the pattern produced by grammar bitmask filtering when
644 the model assigns similar scores to the few allowed tokens.
645 The ternary search can converge to a pivot equal to max_logit,
646 causing the strict `>` keep_mask to exclude everything.
647 Regression test for the `final_pivot >= max_logit` guard.
648 """
649 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
650
651 batch_size, vocab_size = 32, 128256
652 logits = torch.full(
653 (batch_size, vocab_size), float("-inf"), dtype=torch.float32
654 )
655 # Set exactly `num_valid` tokens per row to the SAME finite value.
656 for i in range(batch_size):
657 indices = torch.randperm(vocab_size, generator=self.generator)[:num_valid]
658 logits[i, indices] = 1.0 # all equal
659
660 k: torch.Tensor | None = None
661 p: torch.Tensor | None = None
662 if mode in ("topk_only", "topk_and_topp"):
663 k = torch.full((batch_size,), max(1, num_valid - 1), dtype=torch.int32)
664 if mode in ("topp_only", "topk_and_topp"):
665 p = torch.full((batch_size,), 0.95, dtype=torch.float32)
666
667 result = apply_top_k_top_p_triton(logits.clone(), k, p)
668
669 assert not result.isnan().any(), "NaN in equal-logit result"
670 for i in range(batch_size):
671 kept = (result[i] > float("-inf")).sum().item()
672 # The key invariant: at least one token must survive.
673 # With all-equal logits the pivot search can't differentiate
674 # tokens, so the guard may keep more than k — that is the
675 # intended safe fallback.
676 assert kept > 0, (
677 f"Row {i}: all tokens masked with {num_valid} equal-valued "
678 f"finite logits ({mode})"
679 )
680
681 @pytest.mark.parametrize("num_valid", [2, 5, 10])
682 def test_nearly_equal_logits_topp(self, num_valid: int):
683 """Few valid tokens with very similar (but not identical) logits.
684
685 Ensures the kernel handles near-degenerate probability
686 distributions where the ternary search range collapses.
687 """
688 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
689
690 batch_size, vocab_size = 32, 128256
691 logits = torch.full(
692 (batch_size, vocab_size), float("-inf"), dtype=torch.float32
693 )
694 for i in range(batch_size):
695 indices = torch.randperm(vocab_size, generator=self.generator)[:num_valid]
696 # Tiny spread: values in [1.0, 1.0 + 1e-6]
697 logits[i, indices] = (
698 1.0
699 + torch.rand(num_valid, generator=self.generator, dtype=torch.float32)
700 * 1e-6
701 )
702
703 p = torch.full((batch_size,), 0.95, dtype=torch.float32)
704 result = apply_top_k_top_p_triton(logits.clone(), None, p)
705
706 assert not result.isnan().any(), "NaN in nearly-equal-logit result"
707 for i in range(batch_size):
708 kept = (result[i] > float("-inf")).sum().item()
709 assert kept > 0, (
710 f"Row {i}: all tokens masked with {num_valid} "
711 f"nearly-equal finite logits"
712 )
713
714 def test_mixed_neginf_and_normal_rows(self):
715 """Batch with a mix of normal rows and heavily-masked rows."""
716 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
717
718 batch_size, vocab_size = 32, 32000
719 logits = torch.randn(
720 batch_size, vocab_size, generator=self.generator, dtype=torch.float32
721 )
722 # Mask even rows heavily (99% -inf), leave odd rows normal.
723 for i in range(0, batch_size, 2):
724 mask = torch.rand(vocab_size, generator=self.generator) < 0.99
725 logits[i][mask] = float("-inf")
726
727 k = torch.randint(
728 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
729 )
730 p = (
731 torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9
732 + 0.1
733 )
734
735 result = apply_top_k_top_p_triton(logits.clone(), k, p)
736 assert not result.isnan().any(), "NaN in mixed normal/-inf batch"
737 for i in range(batch_size):
738 kept = (result[i] > float("-inf")).sum().item()
739 assert kept <= k[i].item()
740 finite_in = (logits[i] > float("-inf")).sum().item()
741 if finite_in > 0:
742 assert kept > 0, f"Row {i}: no tokens kept"
743
744 def test_mixed_batch_grammar_rows_get_topp(self):
745 """Grammar-masked rows in a mixed top-k batch must still get top-p.
746
747 Regression: with the split top-p pipeline active (batch <= 64), a
748 row with k < vocab but <= k finite logits was covered by neither the
749 monolithic kernel (standalone top-p disabled) nor the split pipeline
750 (skips rows with k < vocab), so its top-p was silently dropped.
751 """
752 batch_size, vocab_size = 8, 32000
753 logits = torch.full((batch_size, vocab_size), float("-inf"))
754 # Rows 0-2: grammar rows (30 finite tokens); k=50 is a no-op for them.
755 for i in range(3):
756 idx = torch.randperm(vocab_size, generator=self.generator)[:30]
757 logits[i, idx] = torch.randn(30, generator=self.generator)
758 # Rows 3-4: p-only rows (k = vocab); rows 5-7: normal top-k+top-p.
759 logits[3:5] = torch.randn(2, vocab_size, generator=self.generator)
760 logits[5:] = torch.randn(3, vocab_size, generator=self.generator)
761 k = torch.tensor(
762 [50, 50, 50, vocab_size, vocab_size, 50, 50, 50], dtype=torch.int32
763 )
764 p = torch.tensor([0.5, 0.5, 0.5, 0.9, 0.9, 0.9, 0.9, 0.9])
765
766 self._compare_results(logits, k, p)
767
768 # Top-p must actually trim the grammar rows beyond the (no-op) top-k.
769 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
770
771 kept_topk_only = (
772 apply_top_k_top_p_triton(logits.clone(), k, None)[:3] != float("-inf")
773 ).sum(-1)
774 kept_topk_topp = (
775 apply_top_k_top_p_triton(logits.clone(), k, p)[:3] != float("-inf")
776 ).sum(-1)
777 assert (kept_topk_topp < kept_topk_only).all(), (
778 f"top-p did not trim grammar rows: {kept_topk_only.tolist()} -> "
779 f"{kept_topk_topp.tolist()}"
780 )
781
782 def test_topp_tie_break_deterministic(self):
783 """Boundary-duplicate trimming must be deterministic across calls.
784
785 Regression: the split top-p mask kernel handed out the tie budget in
786 cross-program atomic order, so which copies of a tied boundary value
787 survived could vary between identical calls.
788 """
789 from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
790
791 batch_size, vocab_size = 4, 128256
792 # 1000 tied top tokens; the p=0.5 boundary lands inside the tie
793 # group (each tied token has mass e/(1000e+1000) ~ 0.073%).
794 logits = torch.full((batch_size, vocab_size), float("-inf"))
795 logits[:, :1000] = 1.0
796 logits[:, 1000:2000] = 0.0
797 p = torch.full((batch_size,), 0.5, dtype=torch.float32)
798
799 results = [apply_top_k_top_p_triton(logits.clone(), None, p) for _ in range(5)]
800 for r in results[1:]:
801 assert torch.equal(results[0], r), "non-deterministic tie break"
802
803 # The split pipeline now keeps the first `numkeep` copies in index
804 # order, exactly like the monolithic kernel and the PyTorch impl.
805 kept_idx = (results[0][0] != float("-inf")).nonzero().flatten()
806 assert kept_idx.min().item() == 0
807 assert kept_idx.max().item() == len(kept_idx) - 1
808 pytorch_kept = (
809 (apply_top_k_top_p_pytorch(logits.clone(), None, p)[0] != float("-inf"))
810 .sum()
811 .item()
812 )
813 assert len(kept_idx) == pytorch_kept
814
815
816 # =============================================================================
817 # FlashInfer top-k/top-p robustness tests
818 # =============================================================================
819
820
821 @pytest.mark.skipif(
822 not FLASHINFER_TOPK_TOPP_SUPPORTED,
823 reason="FlashInfer top-k/top-p sampler requires CUDA "
824 "and a GPU with FlashInfer support.",
825 )
826 class TestFlashInferTopkToppRobustness:
827 """Robustness of FlashInfer top-k / top-p sampling to NaN / Inf logits.
828
829 The FlashInfer sampler is enabled by default on supported GPUs. A
830 single poisoned request (NaN / +Inf / -Inf in row 0) must not:
831
832 1. crash or hang the process;
833 2. produce out-of-range token ids (anything outside ``[0, vocab)``);
834 3. corrupt other batch rows — neighbours of a poisoned row must
835 still receive valid token ids (regression for cross-row
836 corruption in a DP batch where one bad request would otherwise
837 poison its peers).
838
839 The reference is "no crash + valid token ids", not bit-exact equality
840 against the PyTorch-native path.
841 """
842
843 BATCH = 8
844 VOCAB = 32768
845 TOPK = 50
846 TOPP = 0.9
847
848 @pytest.fixture(autouse=True)
849 def setup(self):
850 torch.set_default_device(DEVICE_TYPE)
851 self.generator = Generator(device=DEVICE_TYPE).manual_seed(1234)
852
853 def _make_logits(self, pattern: str) -> torch.Tensor:
854 """Build (BATCH, VOCAB) logits with `pattern` applied to row 0
855 (rows 1..B-1 stay clean so we can detect cross-row corruption)."""
856 logits = (
857 torch.randn(
858 self.BATCH,
859 self.VOCAB,
860 generator=self.generator,
861 dtype=torch.float32,
862 )
863 * 5.0
864 )
865 if pattern == "clean":
866 return logits
867 if pattern == "nan_one_row":
868 logits[0, :] = float("nan")
869 elif pattern == "nan_few":
870 # Scatter 16 NaNs across row 0, keep the rest finite.
871 idx = torch.randperm(self.VOCAB, generator=self.generator)[:16]
872 logits[0, idx] = float("nan")
873 elif pattern == "nan_at_top":
874 # Poison the top-32 highest-scoring positions of row 0 — worst
875 # case for top-k since these are exactly the tokens that would
876 # otherwise be selected. Use argsort instead of topk to avoid
877 # a known compute-sanitizer false positive in mbtopk.
878 top_idx = logits[0].argsort(descending=True)[:32]
879 logits[0, top_idx] = float("nan")
880 elif pattern == "nan_all_rows":
881 logits[:, :] = float("nan")
882 elif pattern == "pos_inf_one_row":
883 logits[0, :] = float("inf")
884 elif pattern == "neg_inf_one_row":
885 logits[0, :] = float("-inf")
886 elif pattern == "mixed_inf_nan":
887 assert self.BATCH >= 3
888 logits[0, :] = float("nan")
889 logits[1, :] = float("inf")
890 logits[2, :] = float("-inf")
891 elif pattern == "degenerate_flat":
892 logits[:, :] = 1.0
893 else:
894 raise ValueError(f"unknown pattern: {pattern}")
895 return logits
896
897 def _check_tokens(self, tokens: torch.Tensor, ctx: str):
898 assert tokens.dim() == 1, f"{ctx}: expected 1-D output, got {tokens.shape}"
899 assert tokens.shape[0] == self.BATCH, (
900 f"{ctx}: expected batch size {self.BATCH}, got {tokens.shape[0]}"
901 )
902 ids = tokens.tolist()
903 min_id, max_id = min(ids), max(ids)
904 assert 0 <= min_id < self.VOCAB and 0 <= max_id < self.VOCAB, (
905 f"{ctx}: token id(s) outside [0, {self.VOCAB}): min={min_id}, max={max_id}"
906 )
907
908 @pytest.mark.parametrize(
909 "pattern",
910 [
911 "clean",
912 "nan_one_row",
913 "nan_few",
914 "nan_at_top",
915 "nan_all_rows",
916 "pos_inf_one_row",
917 "neg_inf_one_row",
918 "mixed_inf_nan",
919 "degenerate_flat",
920 ],
921 )
922 @pytest.mark.parametrize("path", ["topk_only", "topp_only", "topk_topp"])
923 def test_flashinfer_handles_pathological_logits(self, pattern: str, path: str):
924 """flashinfer_sample must return valid ids even on poisoned logits.
925
926 Direct call into ``flashinfer_sample`` — exactly the code path
927 ``TopKTopPSampler.forward_cuda`` takes when FI is enabled.
928 """
929 from vllm.v1.sample.ops.topk_topp_sampler import flashinfer_sample
930
931 logits = self._make_logits(pattern)
932 k = (
933 torch.full(
934 (self.BATCH,),
935 self.TOPK,
936 device=DEVICE_TYPE,
937 dtype=torch.int32,
938 )
939 if path in ("topk_only", "topk_topp")
940 else None
941 )
942 p = (
943 torch.full(
944 (self.BATCH,),
945 self.TOPP,
946 device=DEVICE_TYPE,
947 dtype=torch.float32,
948 )
949 if path in ("topp_only", "topk_topp")
950 else None
951 )
952
953 # flashinfer_sample may mutate its input in-place; pass a clone so
954 # the parametrize iterations stay independent.
955 tokens = flashinfer_sample(logits.clone().contiguous(), k, p, {})
956 # Surface any async CUDA error synchronously (e.g. illegal memory
957 # access from a malformed FlashInfer call) so it's attributed to
958 # this test rather than a later, unrelated GPU op.
959 torch.accelerator.synchronize()
960 self._check_tokens(tokens, ctx=f"pattern={pattern}, path={path}")
961
962
963 # =============================================================================
964 # FlashInfer top-k/top-p distribution-match tests
965 # =============================================================================
966
967
968 @pytest.mark.skipif(
969 not FLASHINFER_TOPK_TOPP_SUPPORTED,
970 reason="FlashInfer top-k/top-p sampler requires CUDA "
971 "and a GPU with FlashInfer support.",
972 )
973 class TestFlashInferDistributionMatch:
974 """Chi-square goodness-of-fit: FlashInfer and PyTorch-native samplers
975 both reproduce the expected token distribution after top-k / top-p.
976
977 Regression guard against historical FlashInfer distribution-shift.
978 Each impl is compared to the theoretical distribution (softmax of
979 filtered logits); if both pass they are statistically equivalent
980 to each other by transitivity.
981 """
982
983 VOCAB = 32
984 N_SAMPLES = 50_000
985 ALPHA = 1e-6
986 SEED = 0
987
988 @pytest.mark.parametrize(
989 "topk,topp",
990 [
991 (8, None),
992 (16, None),
993 (None, 0.5),
994 (None, 0.7),
995 (None, 0.99),
996 (8, 0.9),
997 (4, 0.5),
998 ],
999 )
1000 def test_distribution_matches_theoretical(self, topk, topp):
1001 from scipy.stats import chisquare
1002
1003 from vllm.v1.sample.ops.topk_topp_sampler import (
1004 apply_top_k_top_p,
1005 flashinfer_sample,
1006 random_sample,
1007 )
1008
1009 torch.set_default_device(DEVICE_TYPE)
1010 torch.manual_seed(self.SEED)
1011
1012 # Same logits row used for both impls so the comparison is fair.
1013 logits_one = (
1014 torch.randn(
1015 (1, self.VOCAB),
1016 dtype=torch.float32,
1017 )
1018 * 2.0
1019 )
1020
1021 # Theoretical expected distribution from PyTorch-native filter.
1022 k_one = torch.tensor([topk], dtype=torch.int32) if topk is not None else None
1023 p_one = torch.tensor([topp], dtype=torch.float32) if topp is not None else None
1024 masked = apply_top_k_top_p_pytorch(logits_one.clone(), k_one, p_one)
1025 expected_probs = masked.softmax(dim=-1).flatten().cpu().numpy()
1026 expected_counts = expected_probs * self.N_SAMPLES
1027
1028 # Build a batch of N identical rows for both impls.
1029 batch = logits_one.expand(self.N_SAMPLES, self.VOCAB).contiguous()
1030 k_batch = (
1031 torch.full((self.N_SAMPLES,), topk, dtype=torch.int32)
1032 if topk is not None
1033 else None
1034 )
1035 p_batch = (
1036 torch.full((self.N_SAMPLES,), topp, dtype=torch.float32)
1037 if topp is not None
1038 else None
1039 )
1040
1041 # FlashInfer dispatch path.
1042 fi_tokens = flashinfer_sample(batch.contiguous(), k_batch, p_batch, {})
1043 fi_counts = torch.bincount(fi_tokens, minlength=self.VOCAB).cpu().numpy()
1044 self._chi2_check(
1045 fi_counts,
1046 expected_counts,
1047 chisquare,
1048 label=f"flashinfer top-k={topk} top-p={topp}",
1049 )
1050
1051 # PyTorch-native dispatch path (Triton-routed filter + Gumbel sample).
1052 processed = apply_top_k_top_p(batch.clone(), k_batch, p_batch)
1053 probs = processed.softmax(dim=-1, dtype=torch.float32)
1054 pt_tokens = random_sample(probs, {})
1055 pt_counts = torch.bincount(pt_tokens, minlength=self.VOCAB).cpu().numpy()
1056 self._chi2_check(
1057 pt_counts,
1058 expected_counts,
1059 chisquare,
1060 label=f"native top-k={topk} top-p={topp}",
1061 )
1062
1063 def _chi2_check(self, empirical, expected, chisquare_fn, *, label):
1064 import numpy as np
1065
1066 # Hard check: the sampler must never produce a token outside the
1067 # expected support (zero theoretical probability).
1068 outside = (expected == 0) & (empirical > 0)
1069 assert not outside.any(), (
1070 f"{label}: sampled out-of-support tokens "
1071 f"(zero expected prob): indices={outside.nonzero()[0].tolist()}"
1072 )
1073 # Skip chi-square in the degenerate case where the support
1074 # collapses to a single token (e.g. very restrictive joint
1075 # top-k + top-p): all samples must land there and the hard
1076 # check above already verified they do.
1077 in_support = expected > 0
1078 if int(in_support.sum()) <= 1:
1079 return
1080 # Soft check: chi-square goodness-of-fit on in-support tokens.
1081 # Cast to float64 so the rescaling step below stays within
1082 # scipy.chisquare's strict 1.5e-8 sum-equality tolerance.
1083 emp = empirical[in_support].astype(np.float64)
1084 exp = expected[in_support].astype(np.float64)
1085 exp = exp * (emp.sum() / exp.sum())
1086 chi2, p_value = chisquare_fn(emp, exp)
1087 assert p_value > self.ALPHA, (
1088 f"{label}: distribution differs from theoretical: "
1089 f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"
1090 )