@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c656ba3b4d
Fork 0

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 8, 03:45 UTC

[Kernel] Triton-based Top-k and Top-p sampler kernels (#33538)

Signed-off-by: js_park <cakeng@naver.com> Signed-off-by: Jongseok Park <37990712+cakeng@users.noreply.github.com> Signed-off-by: Sunga Kim <sunga.kim@berkeley.edu> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Sunga Kim <sunga.kim@berkeley.edu> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Jongseok Park committed Feb 17, 2026 at 15:14 UTC c656ba3b4d2cda82ca753eefde4b10cbf04c0a3f
6 files changed +2002 -32
benchmarks/benchmark_topk_topp.py new
+471
@@ -0,0 +1,471 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""
5 +Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
6 +
7 +Compares:
8 +- apply_top_k_top_p_triton (Triton binary search)
9 +- apply_top_k_top_p (PyTorch sort-based)
10 +
11 +Scenarios:
12 +- top_k only (whole batch, partial batch)
13 +- top_p only (whole batch, partial batch)
14 +- mix of top_k and top_p
15 +"""
16 +
17 +import argparse
18 +import gc
19 +from dataclasses import dataclass
20 +
21 +import torch
22 +
23 +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
24 +from vllm.v1.sample.ops.topk_topp_triton import (
25 + apply_top_k_top_p_triton,
26 + reset_buffer_cache,
27 +)
28 +
29 +
30 +@dataclass
31 +class BenchmarkConfig:
32 + """Configuration for a benchmark run."""
33 +
34 + name: str
35 + batch_size: int
36 + vocab_size: int
37 + # k and p can be tensors or None
38 + k_values: torch.Tensor | None # [batch_size] or None
39 + p_values: torch.Tensor | None # [batch_size] or None
40 + description: str
41 + ops_pct: float = 0.0 # Percentage of ops relative to batch size
42 +
43 +
44 +def calculate_ops_pct(
45 + k_values: torch.Tensor | None,
46 + p_values: torch.Tensor | None,
47 + vocab_size: int,
48 + batch_size: int,
49 +) -> float:
50 + """
51 + Calculate the percentage of active top-k and top-p operations.
52 +
53 + Returns percentage where 100% = batch_size ops.
54 + E.g., if all rows have both top-k and top-p active, returns 200%.
55 + """
56 + active_ops = 0
57 +
58 + if k_values is not None:
59 + # Count rows where k < vocab_size (active top-k filtering)
60 + active_ops += (k_values < vocab_size).sum().item()
61 +
62 + if p_values is not None:
63 + # Count rows where p < 1.0 (active top-p filtering)
64 + active_ops += (p_values < 1.0).sum().item()
65 +
66 + return (active_ops / batch_size) * 100 if batch_size > 0 else 0.0
67 +
68 +
69 +def create_logits(
70 + batch_size: int, vocab_size: int, device: str = "cuda"
71 +) -> torch.Tensor:
72 + """Create random logits mimicking a realistic LLM distribution.
73 +
74 + Uses a Zipf-like probability distribution (rank^-1.1) converted to logits
75 + via log, then randomly permuted per row. This produces a peaked distribution
76 + where a small number of tokens capture most probability mass, similar to
77 + real model outputs.
78 + """
79 + # Create Zipf-like probabilities: p(rank) ~ rank^(-alpha)
80 + ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device)
81 + probs = ranks.pow(-1.1)
82 + probs = probs / probs.sum()
83 +
84 + # Convert to logits (log-probabilities, unnormalized is fine)
85 + base_logits = probs.log()
86 +
87 + # Broadcast to batch and randomly permute each row
88 + logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone()
89 + for i in range(batch_size):
90 + logits[i] = logits[i, torch.randperm(vocab_size, device=device)]
91 +
92 + return logits
93 +
94 +
95 +def measure_memory() -> tuple[int, int]:
96 + """Return (allocated, reserved) memory in bytes."""
97 + torch.cuda.synchronize()
98 + return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
99 +
100 +
101 +def reset_memory_stats():
102 + """Reset peak memory statistics."""
103 + reset_buffer_cache()
104 + torch.cuda.reset_peak_memory_stats()
105 + torch.cuda.empty_cache()
106 + gc.collect()
107 +
108 +
109 +def benchmark_function(
110 + func,
111 + logits: torch.Tensor,
112 + k: torch.Tensor | None,
113 + p: torch.Tensor | None,
114 + warmup_iters: int = 5,
115 + benchmark_iters: int = 20,
116 +) -> tuple[float, int]:
117 + """
118 + Benchmark a function and return (avg_time_ms, peak_memory_bytes).
119 +
120 + Returns average time in milliseconds and peak memory usage.
121 + """
122 + # Warmup
123 + for _ in range(warmup_iters):
124 + logits_copy = logits.clone()
125 + func(logits_copy, k, p)
126 + torch.cuda.synchronize()
127 +
128 + # Reset memory stats before benchmark
129 + reset_memory_stats()
130 +
131 + # Benchmark
132 + start_events = [
133 + torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
134 + ]
135 + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
136 +
137 + for i in range(benchmark_iters):
138 + logits_copy = logits.clone()
139 + start_events[i].record()
140 + func(logits_copy, k, p)
141 + end_events[i].record()
142 +
143 + torch.cuda.synchronize()
144 +
145 + # Calculate timing
146 + times = [
147 + start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters)
148 + ]
149 + avg_time = sum(times) / len(times)
150 +
151 + # Get peak memory
152 + _, peak_memory = measure_memory()
153 +
154 + return avg_time, peak_memory
155 +
156 +
157 +def create_benchmark_configs(
158 + batch_sizes: list[int],
159 + vocab_sizes: list[int],
160 + device: str = "cuda",
161 +) -> list[BenchmarkConfig]:
162 + """Create all benchmark configurations."""
163 + configs = []
164 +
165 + for vocab_size in vocab_sizes:
166 + for batch_size in batch_sizes:
167 + # 1. Top-k only - whole batch (all rows have k < vocab_size)
168 + k_all = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
169 + configs.append(
170 + BenchmarkConfig(
171 + name=f"topk_whole_b{batch_size}_v{vocab_size // 1000}k",
172 + batch_size=batch_size,
173 + vocab_size=vocab_size,
174 + k_values=k_all,
175 + p_values=None,
176 + description=f"Top-k only (whole batch, k=50), "
177 + f"batch={batch_size}, vocab={vocab_size}",
178 + ops_pct=calculate_ops_pct(k_all, None, vocab_size, batch_size),
179 + )
180 + )
181 +
182 + # 2. Top-k only - partial batch (half have k=50, half have k=vocab_size)
183 + k_partial = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
184 + k_partial[batch_size // 2 :] = vocab_size # No filtering for second half
185 + configs.append(
186 + BenchmarkConfig(
187 + name=f"topk_partial_b{batch_size}_v{vocab_size // 1000}k",
188 + batch_size=batch_size,
189 + vocab_size=vocab_size,
190 + k_values=k_partial,
191 + p_values=None,
192 + description=f"Top-k only (partial batch, 50% k=50, 50% k=vocab), "
193 + f"batch={batch_size}, vocab={vocab_size}",
194 + ops_pct=calculate_ops_pct(k_partial, None, vocab_size, batch_size),
195 + )
196 + )
197 +
198 + # 3. Top-p only - whole batch (all rows have p < 1.0)
199 + p_all = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
200 + configs.append(
201 + BenchmarkConfig(
202 + name=f"topp_whole_b{batch_size}_v{vocab_size // 1000}k",
203 + batch_size=batch_size,
204 + vocab_size=vocab_size,
205 + k_values=None,
206 + p_values=p_all,
207 + description=f"Top-p only (whole batch, p=0.9), "
208 + f"batch={batch_size}, vocab={vocab_size}",
209 + ops_pct=calculate_ops_pct(None, p_all, vocab_size, batch_size),
210 + )
211 + )
212 +
213 + # 4. Top-p only - partial batch (half have p=0.9, half have p=1.0)
214 + p_partial = torch.full(
215 + (batch_size,), 0.9, dtype=torch.float32, device=device
216 + )
217 + p_partial[batch_size // 2 :] = 1.0 # No filtering for second half
218 + configs.append(
219 + BenchmarkConfig(
220 + name=f"topp_partial_b{batch_size}_v{vocab_size // 1000}k",
221 + batch_size=batch_size,
222 + vocab_size=vocab_size,
223 + k_values=None,
224 + p_values=p_partial,
225 + description=f"Top-p only (partial batch, 50% p=0.9, 50% p=1.0), "
226 + f"batch={batch_size}, vocab={vocab_size}",
227 + ops_pct=calculate_ops_pct(None, p_partial, vocab_size, batch_size),
228 + )
229 + )
230 +
231 + # 5. Mix of top-k and top-p (both applied to whole batch)
232 + k_mix = torch.full((batch_size,), 100, dtype=torch.int32, device=device)
233 + p_mix = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
234 + configs.append(
235 + BenchmarkConfig(
236 + name=f"topk_topp_whole_b{batch_size}_v{vocab_size // 1000}k",
237 + batch_size=batch_size,
238 + vocab_size=vocab_size,
239 + k_values=k_mix,
240 + p_values=p_mix,
241 + description=f"Top-k + Top-p (whole batch, k=100, p=0.9), "
242 + f"batch={batch_size}, vocab={vocab_size}",
243 + ops_pct=calculate_ops_pct(k_mix, p_mix, vocab_size, batch_size),
244 + )
245 + )
246 +
247 + # 6. Mix with partial application (some rows k only, some p only, some both)
248 + k_mixed = torch.full(
249 + (batch_size,), vocab_size, dtype=torch.int32, device=device
250 + )
251 + p_mixed = torch.full((batch_size,), 1.0, dtype=torch.float32, device=device)
252 + # First third: k only
253 + third = batch_size // 3
254 + k_mixed[:third] = 50
255 + # Second third: p only
256 + p_mixed[third : 2 * third] = 0.5
257 + # Last third: both k and p
258 + k_mixed[2 * third :] = 100
259 + p_mixed[2 * third :] = 0.9
260 + configs.append(
261 + BenchmarkConfig(
262 + name=f"mixed_partial_b{batch_size}_v{vocab_size // 1000}k",
263 + batch_size=batch_size,
264 + vocab_size=vocab_size,
265 + k_values=k_mixed,
266 + p_values=p_mixed,
267 + description=f"Mixed partial (1/3 k=50, 1/3 p=0.9, 1/3 both), "
268 + f"batch={batch_size}, vocab={vocab_size}",
269 + ops_pct=calculate_ops_pct(k_mixed, p_mixed, vocab_size, batch_size),
270 + )
271 + )
272 +
273 + return configs
274 +
275 +
276 +def format_memory(bytes_val: int) -> str:
277 + """Format memory in human-readable form."""
278 + if bytes_val >= 1024**3:
279 + return f"{bytes_val / (1024**3):.2f} GB"
280 + elif bytes_val >= 1024**2:
281 + return f"{bytes_val / (1024**2):.2f} MB"
282 + elif bytes_val >= 1024:
283 + return f"{bytes_val / 1024:.2f} KB"
284 + return f"{bytes_val} B"
285 +
286 +
287 +def run_benchmark(
288 + configs: list[BenchmarkConfig],
289 + warmup_iters: int = 5,
290 + benchmark_iters: int = 20,
291 + verbose: bool = True,
292 +):
293 + """Run all benchmarks and print results."""
294 + results = []
295 +
296 + print("=" * 100)
297 + print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based")
298 + print("=" * 100)
299 + print()
300 +
301 + for config in configs:
302 + if verbose:
303 + print(f"Running: {config.description}")
304 +
305 + # Create fresh logits for this config
306 + logits = create_logits(config.batch_size, config.vocab_size)
307 +
308 + # Benchmark Triton
309 + reset_memory_stats()
310 + triton_time, triton_mem = benchmark_function(
311 + apply_top_k_top_p_triton,
312 + logits,
313 + config.k_values,
314 + config.p_values,
315 + warmup_iters,
316 + benchmark_iters,
317 + )
318 +
319 + # Benchmark PyTorch
320 + reset_memory_stats()
321 + pytorch_time, pytorch_mem = benchmark_function(
322 + apply_top_k_top_p_pytorch,
323 + logits,
324 + config.k_values,
325 + config.p_values,
326 + warmup_iters,
327 + benchmark_iters,
328 + )
329 +
330 + speedup = pytorch_time / triton_time if triton_time > 0 else float("inf")
331 + mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf")
332 +
333 + result = {
334 + "config": config,
335 + "triton_time_ms": triton_time,
336 + "pytorch_time_ms": pytorch_time,
337 + "triton_mem": triton_mem,
338 + "pytorch_mem": pytorch_mem,
339 + "speedup": speedup,
340 + "mem_ratio": mem_ratio,
341 + }
342 + results.append(result)
343 +
344 + if verbose:
345 + print(f" Triton: {triton_time:.3f} ms, {format_memory(triton_mem)}")
346 + print(f" PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}")
347 + print(f" Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x")
348 + print()
349 +
350 + # Clean up
351 + del logits
352 + reset_memory_stats()
353 +
354 + return results
355 +
356 +
357 +def print_summary_table(results: list[dict]):
358 + """Print a summary table of results."""
359 + print()
360 + print("=" * 130)
361 + print("SUMMARY TABLE")
362 + print("=" * 130)
363 + print()
364 +
365 + # Header
366 + header = (
367 + f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} "
368 + f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} "
369 + f"{'Tri Mem':>10} {'Pyt Mem':>10}"
370 + )
371 + print(header)
372 + print("-" * 130)
373 +
374 + # Group by scenario type
375 + current_vocab = None
376 + for result in results:
377 + config = result["config"]
378 +
379 + # Add separator between vocab sizes
380 + if current_vocab != config.vocab_size:
381 + if current_vocab is not None:
382 + print("-" * 130)
383 + current_vocab = config.vocab_size
384 +
385 + scenario = config.name.split("_b")[0] # Extract scenario name
386 + print(
387 + f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} "
388 + f"{config.ops_pct:>5.0f}% "
389 + f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} "
390 + f"{result['speedup']:>7.2f}x "
391 + f"{format_memory(result['triton_mem']):>10} "
392 + f"{format_memory(result['pytorch_mem']):>10}"
393 + )
394 +
395 + print("=" * 130)
396 +
397 +
398 +def main():
399 + parser = argparse.ArgumentParser(
400 + description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations"
401 + )
402 + parser.add_argument(
403 + "--batch-sizes",
404 + type=int,
405 + nargs="+",
406 + default=[1, 4, 16, 64, 128, 512, 1024, 2048],
407 + help="Batch sizes to test (default: 1 4 16 64)",
408 + )
409 + parser.add_argument(
410 + "--vocab-sizes",
411 + type=int,
412 + nargs="+",
413 + default=[32768, 131072], # 32k, 128k
414 + help="Vocabulary sizes to test (default: 32768 131072)",
415 + )
416 + parser.add_argument(
417 + "--warmup-iters",
418 + type=int,
419 + default=5,
420 + help="Number of warmup iterations (default: 5)",
421 + )
422 + parser.add_argument(
423 + "--benchmark-iters",
424 + type=int,
425 + default=20,
426 + help="Number of benchmark iterations (default: 20)",
427 + )
428 + parser.add_argument(
429 + "--quiet",
430 + action="store_true",
431 + help="Only print summary table",
432 + )
433 +
434 + args = parser.parse_args()
435 +
436 + # Print configuration
437 + print(f"Batch sizes: {args.batch_sizes}")
438 + print(f"Vocab sizes: {args.vocab_sizes}")
439 + print(f"Warmup iterations: {args.warmup_iters}")
440 + print(f"Benchmark iterations: {args.benchmark_iters}")
441 + print()
442 +
443 + # Check CUDA
444 + if not torch.cuda.is_available():
445 + print("ERROR: CUDA is not available. This benchmark requires a GPU.")
446 + return
447 +
448 + device_name = torch.cuda.get_device_name(0)
449 + print(f"GPU: {device_name}")
450 + print()
451 +
452 + # Create configs
453 + configs = create_benchmark_configs(
454 + args.batch_sizes,
455 + args.vocab_sizes,
456 + )
457 +
458 + # Run benchmarks
459 + results = run_benchmark(
460 + configs,
461 + warmup_iters=args.warmup_iters,
462 + benchmark_iters=args.benchmark_iters,
463 + verbose=not args.quiet,
464 + )
465 +
466 + # Print summary
467 + print_summary_table(results)
468 +
469 +
470 +if __name__ == "__main__":
471 + main()
tests/entrypoints/instrumentator/test_basic.py
+2 -1
@@ -145,6 +145,7 @@ async def test_request_cancellation(server: RemoteOpenAIServer):
145 model=MODEL_NAME,
146 max_tokens=10000,
147 extra_body={"min_tokens": 10000},
148 + temperature=0.0,
149 )
150 )
151 tasks.append(task)
@@ -163,7 +164,7 @@ async def test_request_cancellation(server: RemoteOpenAIServer):
164 # be able to respond to this one within the timeout
165 client = server.get_async_client(timeout=5)
166 response = await client.chat.completions.create(
166 - messages=chat_input, model=MODEL_NAME, max_tokens=10
167 + messages=chat_input, model=MODEL_NAME, max_tokens=10, temperature=0.0
168 )
169
170 assert len(response.choices) == 1
tests/v1/sample/test_topk_topp_sampler.py
+453 -4
@@ -5,8 +5,9 @@ import torch
5 from torch import Generator
6
7 from vllm.platforms import current_platform
8 -from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
8 +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
9
10 +CUDA_DEVICE = "cuda" if current_platform.is_cuda() else None
11 DEVICE = current_platform.device_type
12
13 BATCH_SIZE = 1024
@@ -39,11 +40,11 @@ def test_topk_impl_equivalence():
40 )
41
42 # Top-k only implementation
42 - result1 = apply_top_k_top_p(logits=logits.clone(), k=k, p=None)
43 + result1 = apply_top_k_top_p_pytorch(logits=logits.clone(), k=k, p=None)
44
45 # Top-p + top-k
46 no_op_top_p = torch.tensor([1.0])
46 - result2 = apply_top_k_top_p(logits=logits.clone(), k=k, p=no_op_top_p)
47 + result2 = apply_top_k_top_p_pytorch(logits=logits.clone(), k=k, p=no_op_top_p)
48
49 assert torch.allclose(result1, result2)
50
@@ -98,7 +99,7 @@ def test_flashinfer_sampler():
99 torch.randint(0, 2, (BATCH_SIZE,), generator=generator, dtype=torch.bool), 1.0
100 )
101
101 - python_logits = apply_top_k_top_p(
102 + python_logits = apply_top_k_top_p_pytorch(
103 logits=logits.clone(),
104 k=k_values,
105 p=p_values,
@@ -120,3 +121,451 @@ def test_flashinfer_sampler():
121 assert torch.allclose(python_probs, flashinfer_probs, atol=2e-2), (
122 "FlashInfer and Python sampling implementations do not match!"
123 )
124 +
125 +
126 +# =============================================================================
127 +# Triton kernel tests
128 +# =============================================================================
129 +
130 +
131 +@pytest.mark.skipif(CUDA_DEVICE is None, reason="CUDA not available")
132 +class TestTritonTopkTopp:
133 + """Tests for the Triton top-k/top-p kernel."""
134 +
135 + @pytest.fixture(autouse=True)
136 + def setup(self):
137 + """Set up test fixtures."""
138 + torch.set_default_device(CUDA_DEVICE)
139 + self.generator = Generator(device=CUDA_DEVICE).manual_seed(42)
140 +
141 + def _compare_results(
142 + self,
143 + logits: torch.Tensor,
144 + k: torch.Tensor | None,
145 + p: torch.Tensor | None,
146 + ):
147 + """Compare Triton kernel results with PyTorch sorting implementation.
148 +
149 + For top-k only, we expect exact match.
150 + For top-p (with or without top-k), we allow small differences due to
151 + floating-point precision in probability sum calculations.
152 + """
153 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
154 +
155 + # Clone logits for both implementations
156 + logits_pytorch = logits.clone()
157 + logits_triton = logits.clone().to(torch.float32)
158 +
159 + # Apply PyTorch sorting implementation
160 + result_pytorch = apply_top_k_top_p_pytorch(logits_pytorch, k, p)
161 +
162 + # Apply Triton kernel
163 + k_i32 = k.to(torch.int32) if k is not None else None
164 + p_f32 = p.to(torch.float32) if p is not None else None
165 + result_triton = apply_top_k_top_p_triton(logits_triton, k_i32, p_f32)
166 +
167 + # Compare kept counts per row
168 + pytorch_kept = (result_pytorch != float("-inf")).sum(dim=-1)
169 + triton_kept = (result_triton != float("-inf")).sum(dim=-1)
170 +
171 + if p is None:
172 + # Top-k only: expect exact match
173 + assert torch.equal(pytorch_kept, triton_kept), (
174 + f"Top-k mask mismatch: PyTorch kept {pytorch_kept.tolist()}, "
175 + f"Triton kept {triton_kept.tolist()}"
176 + )
177 + else:
178 + # Top-p involved: allow small differences
179 + # Either < 1% of kept values OR < 5 values absolute
180 + max_diff = (pytorch_kept - triton_kept).abs().max().item()
181 + max_kept = pytorch_kept.max().item()
182 + if max_kept > 0 and max_diff > 3:
183 + diff_pct = max_diff / max_kept * 100
184 + assert diff_pct < 0.5, (
185 + f"Top-p mask difference too large: {diff_pct:.2f}% "
186 + f"(max diff {max_diff} values out of {max_kept})"
187 + )
188 +
189 + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
190 + @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
191 + def test_topk_only(self, batch_size: int, vocab_size: int):
192 + """Test top-k only (p=None)."""
193 + logits = torch.randn(
194 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
195 + )
196 + k = torch.randint(
197 + 1, min(100, vocab_size), (batch_size,), generator=self.generator
198 + )
199 + # Randomly disable top-k for some rows (~25%)
200 + disable_mask = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
201 + k.masked_fill_(disable_mask, vocab_size)
202 +
203 + self._compare_results(logits, k, p=None)
204 +
205 + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
206 + @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
207 + def test_topp_only(self, batch_size: int, vocab_size: int):
208 + """Test top-p only (k=None)."""
209 + logits = torch.randn(
210 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
211 + )
212 + p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 # [0.1, 1.0]
213 + # Randomly disable top-p for some rows (~25%)
214 + disable_mask = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
215 + p.masked_fill_(disable_mask, 1.0)
216 +
217 + self._compare_results(logits, k=None, p=p)
218 +
219 + @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024])
220 + @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256])
221 + def test_topk_and_topp(self, batch_size: int, vocab_size: int):
222 + """Test combined top-k and top-p."""
223 + logits = torch.randn(
224 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
225 + )
226 + k = torch.randint(
227 + 1, min(100, vocab_size), (batch_size,), generator=self.generator
228 + )
229 + p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 # [0.1, 1.0]
230 +
231 + # Randomly disable top-k for some rows (~25%)
232 + disable_k = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
233 + k.masked_fill_(disable_k, vocab_size)
234 + # Randomly disable top-p for some rows (~25%)
235 + disable_p = torch.randint(0, 4, (batch_size,), generator=self.generator) == 0
236 + p.masked_fill_(disable_p, 1.0)
237 +
238 + self._compare_results(logits, k, p)
239 +
240 + def test_both_disabled(self):
241 + """Test when both k and p are None (should be no-op)."""
242 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
243 +
244 + logits = torch.randn(32, 1024, generator=self.generator, dtype=torch.float32)
245 + logits_clone = logits.clone()
246 +
247 + result = apply_top_k_top_p_triton(logits_clone, k=None, p=None)
248 +
249 + assert torch.equal(result, logits), "Should be no-op when both k and p are None"
250 +
251 + def test_extreme_k_values(self):
252 + """Test edge cases for k values."""
253 + batch_size, vocab_size = 16, 1024
254 + logits = torch.randn(
255 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
256 + )
257 +
258 + # k=1 (keep only top 1)
259 + k = torch.ones(batch_size, dtype=torch.int32)
260 + self._compare_results(logits.clone(), k, p=None)
261 +
262 + # k=vocab_size (keep all)
263 + k = torch.full((batch_size,), vocab_size, dtype=torch.int32)
264 + self._compare_results(logits.clone(), k, p=None)
265 +
266 + # Mixed extreme values
267 + k = torch.tensor([1, vocab_size, 2, vocab_size - 1] * 4, dtype=torch.int32)
268 + self._compare_results(logits.clone(), k, p=None)
269 +
270 + def test_extreme_p_values(self):
271 + """Test edge cases for p values."""
272 + batch_size, vocab_size = 16, 1024
273 + logits = torch.randn(
274 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
275 + )
276 +
277 + # p close to 0 (very restrictive)
278 + p = torch.full((batch_size,), 0.01, dtype=torch.float32)
279 + self._compare_results(logits.clone(), k=None, p=p)
280 +
281 + # p=1.0 (keep all)
282 + p = torch.ones(batch_size, dtype=torch.float32)
283 + self._compare_results(logits.clone(), k=None, p=p)
284 +
285 + # Mixed values
286 + p = torch.tensor([0.1, 0.5, 0.9, 1.0] * 4, dtype=torch.float32)
287 + self._compare_results(logits.clone(), k=None, p=p)
288 +
289 + def test_large_batch(self):
290 + """Test with a large batch size."""
291 + batch_size, vocab_size = 512, 32000
292 + logits = torch.randn(
293 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
294 + )
295 + k = torch.randint(1, 50, (batch_size,), generator=self.generator)
296 + p = torch.rand(batch_size, generator=self.generator) * 0.5 + 0.5
297 +
298 + self._compare_results(logits, k, p)
299 +
300 + # -----------------------------------------------------------------
301 + # Tests for -inf logits (e.g. from grammar / structured output masks)
302 + # -----------------------------------------------------------------
303 +
304 + @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99])
305 + def test_topk_with_neginf_logits(self, inf_fraction: float):
306 + """Top-k with many -inf logits (simulating grammar bitmask).
307 +
308 + The kernel must not produce NaN when most logits are -inf, which
309 + can happen when structured-output grammar masks are applied before
310 + sampling.
311 + """
312 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
313 +
314 + batch_size, vocab_size = 32, 128256
315 + logits = torch.randn(
316 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
317 + )
318 + # Mask a fraction of logits to -inf.
319 + mask = (
320 + torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction
321 + )
322 + logits[mask] = float("-inf")
323 +
324 + k = torch.randint(
325 + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
326 + )
327 + result = apply_top_k_top_p_triton(logits.clone(), k, None)
328 +
329 + assert not result.isnan().any(), "NaN found in top-k result with -inf logits"
330 + for i in range(batch_size):
331 + kept = (result[i] > float("-inf")).sum().item()
332 + assert kept <= k[i].item(), f"Row {i}: kept {kept} > k={k[i].item()}"
333 + # At least one value should survive unless the row was all -inf.
334 + finite_in = (logits[i] > float("-inf")).sum().item()
335 + if finite_in > 0:
336 + assert kept > 0, f"Row {i}: no tokens kept despite finite input"
337 +
338 + @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99])
339 + def test_topp_with_neginf_logits(self, inf_fraction: float):
340 + """Top-p with many -inf logits."""
341 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
342 +
343 + batch_size, vocab_size = 32, 128256
344 + logits = torch.randn(
345 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
346 + )
347 + mask = (
348 + torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction
349 + )
350 + logits[mask] = float("-inf")
351 +
352 + p = (
353 + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9
354 + + 0.1
355 + )
356 + result = apply_top_k_top_p_triton(logits.clone(), None, p)
357 +
358 + assert not result.isnan().any(), "NaN found in top-p result with -inf logits"
359 + for i in range(batch_size):
360 + finite_in = (logits[i] > float("-inf")).sum().item()
361 + kept = (result[i] > float("-inf")).sum().item()
362 + if finite_in > 0:
363 + assert kept > 0, f"Row {i}: no tokens kept despite finite input"
364 +
365 + @pytest.mark.parametrize("inf_fraction", [0.5, 0.9, 0.99])
366 + def test_topk_topp_with_neginf_logits(self, inf_fraction: float):
367 + """Combined top-k + top-p with many -inf logits."""
368 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
369 +
370 + batch_size, vocab_size = 32, 128256
371 + logits = torch.randn(
372 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
373 + )
374 + mask = (
375 + torch.rand(batch_size, vocab_size, generator=self.generator) < inf_fraction
376 + )
377 + logits[mask] = float("-inf")
378 +
379 + k = torch.randint(
380 + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
381 + )
382 + p = (
383 + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9
384 + + 0.1
385 + )
386 + result = apply_top_k_top_p_triton(logits.clone(), k, p)
387 +
388 + assert not result.isnan().any(), (
389 + "NaN found in top-k+top-p result with -inf logits"
390 + )
391 + for i in range(batch_size):
392 + kept = (result[i] > float("-inf")).sum().item()
393 + assert kept <= k[i].item(), f"Row {i}: kept {kept} > k={k[i].item()}"
394 +
395 + def test_all_neginf_logits(self):
396 + """All logits are -inf (fully masked). Kernel should be a no-op."""
397 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
398 +
399 + batch_size, vocab_size = 16, 128256
400 + logits = torch.full(
401 + (batch_size, vocab_size), float("-inf"), dtype=torch.float32
402 + )
403 +
404 + k = torch.randint(
405 + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
406 + )
407 + p = torch.full((batch_size,), 0.9, dtype=torch.float32)
408 +
409 + # top-k only
410 + result = apply_top_k_top_p_triton(logits.clone(), k, None)
411 + assert not result.isnan().any(), "NaN from all-inf top-k"
412 + assert (result == float("-inf")).all(), "Expected all -inf unchanged"
413 +
414 + # top-p only
415 + result = apply_top_k_top_p_triton(logits.clone(), None, p)
416 + assert not result.isnan().any(), "NaN from all-inf top-p"
417 + assert (result == float("-inf")).all(), "Expected all -inf unchanged"
418 +
419 + # top-k + top-p
420 + result = apply_top_k_top_p_triton(logits.clone(), k, p)
421 + assert not result.isnan().any(), "NaN from all-inf top-k+top-p"
422 + assert (result == float("-inf")).all(), "Expected all -inf unchanged"
423 +
424 + def test_few_valid_tokens_with_neginf(self):
425 + """Only a handful of tokens are finite per row (strict grammar)."""
426 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
427 +
428 + batch_size, vocab_size = 32, 128256
429 + logits = torch.full(
430 + (batch_size, vocab_size), float("-inf"), dtype=torch.float32
431 + )
432 + # Allow only 5 random tokens per row to be finite.
433 + for i in range(batch_size):
434 + indices = torch.randperm(vocab_size, generator=self.generator)[:5]
435 + logits[i, indices] = torch.randn(
436 + 5, generator=self.generator, dtype=torch.float32
437 + )
438 +
439 + k = torch.full((batch_size,), 50, dtype=torch.int32)
440 + p = torch.full((batch_size,), 0.9, dtype=torch.float32)
441 +
442 + # top-k only (k=50 but only 5 finite → keep all 5)
443 + result = apply_top_k_top_p_triton(logits.clone(), k, None)
444 + assert not result.isnan().any()
445 + for i in range(batch_size):
446 + kept = (result[i] > float("-inf")).sum().item()
447 + assert kept == 5, f"Row {i}: expected 5 kept, got {kept}"
448 +
449 + # top-k with k < num_finite
450 + k_small = torch.full((batch_size,), 3, dtype=torch.int32)
451 + result = apply_top_k_top_p_triton(logits.clone(), k_small, None)
452 + assert not result.isnan().any()
453 + for i in range(batch_size):
454 + kept = (result[i] > float("-inf")).sum().item()
455 + assert kept <= 3, f"Row {i}: expected <=3 kept, got {kept}"
456 +
457 + # top-p only
458 + result = apply_top_k_top_p_triton(logits.clone(), None, p)
459 + assert not result.isnan().any()
460 + for i in range(batch_size):
461 + kept = (result[i] > float("-inf")).sum().item()
462 + assert kept > 0, f"Row {i}: no tokens kept"
463 +
464 + @pytest.mark.parametrize("num_valid", [1, 2, 5, 10, 50])
465 + @pytest.mark.parametrize(
466 + "mode",
467 + ["topk_only", "topp_only", "topk_and_topp"],
468 + )
469 + def test_equal_logits_few_valid(self, num_valid: int, mode: str):
470 + """Few valid tokens all sharing the same logit value.
471 +
472 + This is the pattern produced by grammar bitmask filtering when
473 + the model assigns similar scores to the few allowed tokens.
474 + The ternary search can converge to a pivot equal to max_logit,
475 + causing the strict `>` keep_mask to exclude everything.
476 + Regression test for the `final_pivot >= max_logit` guard.
477 + """
478 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
479 +
480 + batch_size, vocab_size = 32, 128256
481 + logits = torch.full(
482 + (batch_size, vocab_size), float("-inf"), dtype=torch.float32
483 + )
484 + # Set exactly `num_valid` tokens per row to the SAME finite value.
485 + for i in range(batch_size):
486 + indices = torch.randperm(vocab_size, generator=self.generator)[:num_valid]
487 + logits[i, indices] = 1.0 # all equal
488 +
489 + k: torch.Tensor | None = None
490 + p: torch.Tensor | None = None
491 + if mode in ("topk_only", "topk_and_topp"):
492 + k = torch.full((batch_size,), max(1, num_valid - 1), dtype=torch.int32)
493 + if mode in ("topp_only", "topk_and_topp"):
494 + p = torch.full((batch_size,), 0.95, dtype=torch.float32)
495 +
496 + result = apply_top_k_top_p_triton(logits.clone(), k, p)
497 +
498 + assert not result.isnan().any(), "NaN in equal-logit result"
499 + for i in range(batch_size):
500 + kept = (result[i] > float("-inf")).sum().item()
501 + # The key invariant: at least one token must survive.
502 + # With all-equal logits the pivot search can't differentiate
503 + # tokens, so the guard may keep more than k — that is the
504 + # intended safe fallback.
505 + assert kept > 0, (
506 + f"Row {i}: all tokens masked with {num_valid} equal-valued "
507 + f"finite logits ({mode})"
508 + )
509 +
510 + @pytest.mark.parametrize("num_valid", [2, 5, 10])
511 + def test_nearly_equal_logits_topp(self, num_valid: int):
512 + """Few valid tokens with very similar (but not identical) logits.
513 +
514 + Ensures the kernel handles near-degenerate probability
515 + distributions where the ternary search range collapses.
516 + """
517 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
518 +
519 + batch_size, vocab_size = 32, 128256
520 + logits = torch.full(
521 + (batch_size, vocab_size), float("-inf"), dtype=torch.float32
522 + )
523 + for i in range(batch_size):
524 + indices = torch.randperm(vocab_size, generator=self.generator)[:num_valid]
525 + # Tiny spread: values in [1.0, 1.0 + 1e-6]
526 + logits[i, indices] = (
527 + 1.0
528 + + torch.rand(num_valid, generator=self.generator, dtype=torch.float32)
529 + * 1e-6
530 + )
531 +
532 + p = torch.full((batch_size,), 0.95, dtype=torch.float32)
533 + result = apply_top_k_top_p_triton(logits.clone(), None, p)
534 +
535 + assert not result.isnan().any(), "NaN in nearly-equal-logit result"
536 + for i in range(batch_size):
537 + kept = (result[i] > float("-inf")).sum().item()
538 + assert kept > 0, (
539 + f"Row {i}: all tokens masked with {num_valid} "
540 + f"nearly-equal finite logits"
541 + )
542 +
543 + def test_mixed_neginf_and_normal_rows(self):
544 + """Batch with a mix of normal rows and heavily-masked rows."""
545 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
546 +
547 + batch_size, vocab_size = 32, 32000
548 + logits = torch.randn(
549 + batch_size, vocab_size, generator=self.generator, dtype=torch.float32
550 + )
551 + # Mask even rows heavily (99% -inf), leave odd rows normal.
552 + for i in range(0, batch_size, 2):
553 + mask = torch.rand(vocab_size, generator=self.generator) < 0.99
554 + logits[i][mask] = float("-inf")
555 +
556 + k = torch.randint(
557 + 1, 50, (batch_size,), generator=self.generator, dtype=torch.int32
558 + )
559 + p = (
560 + torch.rand(batch_size, generator=self.generator, dtype=torch.float32) * 0.9
561 + + 0.1
562 + )
563 +
564 + result = apply_top_k_top_p_triton(logits.clone(), k, p)
565 + assert not result.isnan().any(), "NaN in mixed normal/-inf batch"
566 + for i in range(batch_size):
567 + kept = (result[i] > float("-inf")).sum().item()
568 + assert kept <= k[i].item()
569 + finite_in = (logits[i] > float("-inf")).sum().item()
570 + if finite_in > 0:
571 + assert kept > 0, f"Row {i}: no tokens kept"
vllm/utils/math_utils.py
+2 -6
@@ -14,16 +14,12 @@ def cdiv(a: int, b: int) -> int:
14
15 def next_power_of_2(n: int) -> int:
16 """The next power of 2 (inclusive)"""
17 - if n < 1:
18 - return 1
19 - return 1 << (n - 1).bit_length()
17 + return 1 if n < 1 else 1 << (n - 1).bit_length()
18
19
20 def prev_power_of_2(n: int) -> int:
21 """The previous power of 2 (inclusive)"""
24 - if n <= 0:
25 - return 0
26 - return 1 << (n.bit_length() - 1)
22 + return 0 if n <= 0 else 1 << (n.bit_length() - 1)
23
24
25 def round_up(x: int, y: int) -> int:
vllm/v1/sample/ops/topk_topp_sampler.py
+35 -21
@@ -11,6 +11,10 @@ from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config.model import LogprobsMode
12 from vllm.logger import init_logger
13 from vllm.platforms import CpuArchEnum, current_platform
14 +from vllm.triton_utils import HAS_TRITON
15 +
16 +if HAS_TRITON:
17 + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
18
19 logger = init_logger(__name__)
20
@@ -87,8 +91,6 @@ class TopKTopPSampler(nn.Module):
91 else:
92 self.forward = self.forward_native
93
90 - self.apply_top_k_top_p = apply_top_k_top_p
91 -
94 def forward_native(
95 self,
96 logits: torch.Tensor,
@@ -101,7 +103,7 @@ class TopKTopPSampler(nn.Module):
103
104 The logits tensor may be updated in-place.
105 """
104 - logits = self.apply_top_k_top_p(logits, k, p)
106 + logits = apply_top_k_top_p(logits, k, p)
107 logits_to_return = None
108 if self.logprobs_mode == "processed_logits":
109 logits_to_return = logits
@@ -149,7 +151,7 @@ class TopKTopPSampler(nn.Module):
151
152 The logits tensor may be updated in-place.
153 """
152 - logits = self.apply_top_k_top_p(logits, k, p)
154 + logits = apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)
155 logits_to_return = None
156 if self.logprobs_mode == "processed_logits":
157 logits_to_return = logits
@@ -158,14 +160,14 @@ class TopKTopPSampler(nn.Module):
160
161 if len(generators) != logits.shape[0]:
162 return compiled_random_sample(logits), logits_to_return
161 - else:
162 - probs = logits.softmax(dim=-1, dtype=torch.float32)
163 - q = torch.empty_like(probs)
164 - q.exponential_()
165 - for i, generator in generators.items():
166 - q[i].exponential_(generator=generator)
163
168 - return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return
164 + probs = logits.softmax(dim=-1, dtype=torch.float32)
165 + q = torch.empty_like(probs)
166 + q.exponential_()
167 + for i, generator in generators.items():
168 + q[i].exponential_(generator=generator)
169 +
170 + return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return
171
172 def forward_hip(
173 self,
@@ -241,9 +243,23 @@ def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor:
243
244
245 def apply_top_k_top_p(
246 + logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None
247 +) -> torch.Tensor:
248 + if p is None and k is None:
249 + return logits
250 +
251 + if HAS_TRITON and logits.shape[0] >= 8:
252 + return apply_top_k_top_p_triton(logits, k, p)
253 +
254 + # Use pytorch sort implementation for small batch sizes.
255 + return apply_top_k_top_p_pytorch(logits, k, p)
256 +
257 +
258 +def apply_top_k_top_p_pytorch(
259 logits: torch.Tensor,
260 k: torch.Tensor | None,
261 p: torch.Tensor | None,
262 + allow_cpu_sync: bool = False,
263 ) -> torch.Tensor:
264 """Apply top-k and top-p masks to the logits.
265
@@ -256,8 +272,9 @@ def apply_top_k_top_p(
272 if k is None:
273 return logits
274
259 - # Avoid sorting vocab for top-k only case.
260 - return apply_top_k_only(logits, k)
275 + if allow_cpu_sync:
276 + # Avoid sorting vocab for top-k only case.
277 + return apply_top_k_only(logits, k)
278
279 logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
280
@@ -279,18 +296,16 @@ def apply_top_k_top_p(
296 logits_sort.masked_fill_(top_p_mask, -float("inf"))
297
298 # Re-sort the probabilities.
282 - logits = logits_sort.scatter(dim=-1, index=logits_idx, src=logits_sort)
283 - return logits
299 + return logits.scatter_(dim=-1, index=logits_idx, src=logits_sort)
300
301
286 -def apply_top_k_only(
287 - logits: torch.Tensor,
288 - k: torch.Tensor,
289 -) -> torch.Tensor:
302 +def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
303 """
304 Apply top-k mask to the logits.
305
306 This implementation doesn't involve sorting the entire vocab.
307 + Note however that it involves a GPU->CPU sync which can be detrimental for
308 + async scheduling performance.
309
310 The logits tensor may be updated in-place.
311 """
@@ -304,8 +319,7 @@ def apply_top_k_only(
319 top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index.long())
320 # Handle non-topk rows.
321 top_k_mask.masked_fill_(no_top_k_mask.unsqueeze(1), -float("inf"))
307 - logits.masked_fill_(logits < top_k_mask, -float("inf"))
308 - return logits
322 + return logits.masked_fill_(logits < top_k_mask, -float("inf"))
323
324
325 def random_sample(
vllm/v1/sample/ops/topk_topp_triton.py new
+1039
@@ -0,0 +1,1039 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Combined Top-K and Top-P Triton kernels.
5 +
6 +Based on the paper "Qrita: High-performance Top-k and Top-p Algorithm for GPUs
7 +using Pivot-based Truncation and Selection" By Park et al.
8 +(https://arxiv.org/abs/2602.01518)
9 +
10 +"""
11 +
12 +import torch
13 +
14 +from vllm.triton_utils import tl, triton
15 +from vllm.utils.math_utils import next_power_of_2
16 +
17 +_TRITON_TABLE_CACHE: dict[tuple[torch.device], tuple[torch.Tensor, torch.Tensor]] = {}
18 +_TRITON_BUFFER_CACHE: dict[tuple[torch.device, torch.dtype, int], torch.Tensor] = {}
19 +
20 +# fmt: off
21 +_NORMAL_CDF_TO_SIGMA_TABLE = [
22 + 3.656, 3.650, 3.650, 3.650, 3.626, 3.626, 3.626, 3.514, 3.514, 3.503,
23 + 3.503, 3.434, 3.434, 3.428, 3.428, 3.387, 3.380, 3.380, 3.376, 3.373,
24 + 3.373, 3.356, 3.354, 3.354, 3.291, 3.249, 3.234, 3.214, 3.198, 3.198,
25 + 3.185, 3.177, 3.177, 3.165, 3.164, 3.161, 3.138, 3.120, 3.115, 3.113,
26 + 3.093, 3.066, 3.054, 3.043, 3.037, 3.023, 2.993, 2.991, 2.976, 2.970,
27 + 2.952, 2.946, 2.932, 2.908, 2.902, 2.895, 2.886, 2.874, 2.861, 2.844,
28 + 2.836, 2.810, 2.801, 2.790, 2.784, 2.779, 2.767, 2.757, 2.745, 2.733,
29 + 2.723, 2.716, 2.693, 2.678, 2.671, 2.656, 2.649, 2.629, 2.611, 2.595,
30 + 2.592, 2.585, 2.574, 2.550, 2.543, 2.534, 2.521, 2.518, 2.497, 2.485,
31 + 2.468, 2.450, 2.441, 2.430, 2.412, 2.402, 2.389, 2.383, 2.377, 2.364,
32 + 2.349, 2.338, 2.332, 2.319, 2.310, 2.301, 2.282, 2.274, 2.266, 2.250,
33 + 2.242, 2.236, 2.226, 2.215, 2.207, 2.196, 2.179, 2.171, 2.162, 2.147,
34 + 2.135, 2.121, 2.109, 2.095, 2.085, 2.073, 2.063, 2.045, 2.030, 2.016,
35 + 2.003, 1.992, 1.983, 1.972, 1.960, 1.949, 1.940, 1.928, 1.912, 1.897,
36 + 1.881, 1.869, 1.854, 1.838, 1.824, 1.807, 1.792, 1.779, 1.764, 1.751,
37 + 1.739, 1.726, 1.711, 1.697, 1.685, 1.668, 1.652, 1.636, 1.622, 1.603,
38 + 1.585, 1.568, 1.551, 1.534, 1.513, 1.499, 1.480, 1.464, 1.441, 1.422,
39 + 1.394, 1.373, 1.347, 1.320, 1.296, 1.270, 1.246, 1.219, 1.190, 1.163,
40 + 1.135, 1.104, 1.073, 1.041, 1.006, 0.969, 0.931, 0.894, 0.851, 0.806,
41 + 0.757, 0.702, 0.643, 0.574, 0.498, 0.405, 0.288, 0.134, -0.110, -3.813
42 +]
43 +
44 +_PERCENTILE_TO_STD_TABLE = [
45 + 2.576, 2.319, 2.178, 2.064, 1.968, 1.892, 1.819, 1.757, 1.708, 1.659,
46 + 1.616, 1.568, 1.526, 1.492, 1.456, 1.420, 1.382, 1.342, 1.309, 1.280,
47 + 1.249, 1.221, 1.193, 1.169, 1.145, 1.121, 1.095, 1.073, 1.050, 1.030,
48 + 1.008, 0.987, 0.966, 0.945, 0.926, 0.910, 0.891, 0.871, 0.854, 0.837,
49 + 0.819, 0.803, 0.784, 0.767, 0.753, 0.734, 0.719, 0.702, 0.690, 0.675,
50 + 0.658, 0.640, 0.625, 0.609, 0.595, 0.578, 0.564, 0.550, 0.537, 0.521,
51 + 0.509, 0.495, 0.481, 0.466, 0.453, 0.439, 0.424, 0.410, 0.397, 0.383,
52 + 0.370, 0.356, 0.343, 0.330, 0.316, 0.302, 0.289, 0.274, 0.261, 0.247,
53 + 0.235, 0.223, 0.209, 0.196, 0.184, 0.172, 0.159, 0.149, 0.137, 0.124,
54 + 0.112, 0.100, 0.086, 0.074, 0.062, 0.050, 0.035, 0.023, 0.009, -0.003,
55 + -0.015, -0.027, -0.039, -0.052, -0.063, -0.074, -0.085, -0.097, -0.109, -0.122,
56 + -0.134, -0.147, -0.158, -0.171, -0.184, -0.196, -0.210, -0.223, -0.235, -0.248,
57 + -0.261, -0.275, -0.289, -0.302, -0.317, -0.328, -0.341, -0.353, -0.368, -0.382,
58 + -0.396, -0.410, -0.426, -0.439, -0.452, -0.465, -0.480, -0.493, -0.507, -0.521,
59 + -0.537, -0.551, -0.568, -0.582, -0.597, -0.614, -0.628, -0.643, -0.658, -0.673,
60 + -0.691, -0.706, -0.721, -0.738, -0.754, -0.769, -0.789, -0.808, -0.824, -0.838,
61 + -0.857, -0.877, -0.893, -0.912, -0.929, -0.947, -0.965, -0.983, -1.003, -1.027,
62 + -1.050, -1.070, -1.092, -1.117, -1.139, -1.162, -1.189, -1.216, -1.241, -1.272,
63 + -1.300, -1.330, -1.367, -1.404, -1.441, -1.485, -1.523, -1.564, -1.607, -1.658,
64 + -1.710, -1.778, -1.832, -1.901, -1.978, -2.068, -2.174, -2.325, -2.577, -3.813
65 +]
66 +# fmt: on
67 +
68 +
69 +@triton.jit
70 +def _topk_topp_kernel(
71 + LOGITS,
72 + BUFFER,
73 + PERCENTILE_TO_STD_TABLE,
74 + NORMAL_CDF_TO_SIGMA_TABLE,
75 + K,
76 + P,
77 + BATCH_SIZE,
78 + VOCAB_SIZE: tl.constexpr,
79 + MASK_VALUE: tl.constexpr,
80 + BLOCK_SIZE: tl.constexpr,
81 + BLOCK_SIZE_TRUNC: tl.constexpr,
82 + TOPK_ENABLED: tl.constexpr,
83 + TOPP_ENABLED: tl.constexpr,
84 +):
85 + NUM_TILES: tl.constexpr = (VOCAB_SIZE + BLOCK_SIZE - 1) // BLOCK_SIZE
86 + pid = tl.program_id(0)
87 + num_programs = tl.num_programs(0)
88 + for row_id in tl.range(pid, BATCH_SIZE, num_programs):
89 + LOGITS_ROW = LOGITS + row_id * VOCAB_SIZE
90 + BUFFER_ROW = BUFFER + pid * VOCAB_SIZE
91 +
92 + final_pivot = -float("inf")
93 + duplicate_logit = float("inf")
94 + num_duplicate_logit = tl.zeros((), dtype=tl.uint32)
95 + num_keep = tl.zeros((), dtype=tl.uint32)
96 + num_kept = tl.zeros((), dtype=tl.uint32)
97 +
98 + max_logit = -float("inf")
99 + min_logit = float("inf")
100 +
101 + if TOPK_ENABLED:
102 + k = tl.load(K + row_id)
103 + if k < VOCAB_SIZE:
104 + # Zeroth pass: Compute avg and std from a sample block
105 + offs = tl.arange(0, BLOCK_SIZE)
106 + mask_n = offs < VOCAB_SIZE
107 + logits_blk0 = tl.load(
108 + LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
109 + )
110 + # Exclude -inf values (e.g. from grammar bitmasks) from
111 + # statistics to avoid NaN in pivot computation.
112 + finite_mask = (logits_blk0 > -float("inf")) & mask_n
113 + num_finite = tl.sum(finite_mask)
114 + finite_logits = tl.where(finite_mask, logits_blk0, 0.0)
115 + avg_logit = tl.where(
116 + num_finite > 0, tl.sum(finite_logits) / num_finite, 0.0
117 + )
118 + sq_avg_logit = tl.where(
119 + num_finite > 0,
120 + tl.sum(finite_logits * finite_logits) / num_finite,
121 + 0.0,
122 + )
123 + std_logit = tl.sqrt(
124 + tl.maximum(sq_avg_logit - avg_logit * avg_logit, 0.0)
125 + )
126 +
127 + # Calculate outlier pivot t for Gaussian sigma-truncation
128 + percentile = tl.cast(k / VOCAB_SIZE * 200, tl.uint32)
129 + percentile = tl.minimum(percentile, 199)
130 + sigma = tl.load(PERCENTILE_TO_STD_TABLE + percentile)
131 + sigma = sigma + tl.abs(sigma) * -0.15
132 + outlier_pivot = avg_logit + std_logit * sigma
133 + num_outliers = tl.zeros((), dtype=tl.uint32)
134 +
135 + # First pass: compute max and min logits and gather outliers
136 + num_finite_total = tl.zeros((), dtype=tl.uint32)
137 + for i in range(0, NUM_TILES):
138 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
139 + mask_n = offs_n < VOCAB_SIZE
140 + logits_blk = tl.load(
141 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
142 + )
143 +
144 + max_logit = tl.maximum(max_logit, tl.max(logits_blk))
145 + # Exclude -inf from min to keep binary search bounds
146 + # finite (avoids NaN pivots).
147 + finite_blk_mask = logits_blk > -float("inf")
148 + finite_blk = tl.where(finite_blk_mask, logits_blk, float("inf"))
149 + min_logit = tl.minimum(min_logit, tl.min(finite_blk))
150 + num_finite_total += tl.sum(finite_blk_mask & mask_n)
151 +
152 + outlier_mask = (logits_blk > outlier_pivot) & mask_n
153 + cumulative_pos = tl.cast(
154 + tl.cumsum(outlier_mask) - 1 + num_outliers, tl.int32
155 + )
156 + num_outliers += tl.sum(outlier_mask)
157 + write_pos = tl.where(outlier_mask, cumulative_pos, -1)
158 + tl.store(BUFFER_ROW + write_pos, logits_blk, mask=outlier_mask)
159 +
160 + # If no finite logits exist (all -inf), clamp min to
161 + # max so the search converges to -inf (no masking).
162 + min_logit = tl.minimum(min_logit, max_logit)
163 +
164 + # Second passes: Ternary search for pivots
165 + num_iters = 0
166 + k_pivot = float("inf")
167 + k_pivots_num = tl.zeros((), dtype=tl.uint32)
168 + min_larger = float("inf")
169 + num_min_larger = tl.zeros((), dtype=tl.uint32)
170 + if num_outliers > k:
171 + max_range = max_logit
172 + min_range = outlier_pivot
173 + search_range = tl.cast(num_outliers, tl.int32)
174 + search_iters = tl.cast(
175 + (num_outliers + BLOCK_SIZE_TRUNC - 1) // BLOCK_SIZE_TRUNC,
176 + tl.int32,
177 + )
178 + found_pivot = 0
179 + while found_pivot == 0:
180 + k_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
181 + k_pivots_num_0 = tl.zeros((), dtype=tl.uint32)
182 + min_larger_0 = float("inf")
183 + num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
184 +
185 + k_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
186 + k_pivots_num_1 = tl.zeros((), dtype=tl.uint32)
187 + min_larger_1 = float("inf")
188 + num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
189 +
190 + # First pass: Calculate k_pivots_num and min_larger
191 + for i in range(0, search_iters):
192 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
193 + 0, BLOCK_SIZE_TRUNC
194 + )
195 + mask_n_2 = offs_n < search_range
196 + logits_blk2 = tl.load(
197 + BUFFER_ROW + offs_n, mask=mask_n_2, other=-float("inf")
198 + )
199 +
200 + k_pivots_num_0 += tl.sum(logits_blk2 > k_pivot_0)
201 + k_pivots_num_1 += tl.sum(logits_blk2 > k_pivot_1)
202 +
203 + min_larger_0 = tl.minimum(min_larger_0, tl.min(logits_blk2))
204 + min_larger_1 = tl.minimum(min_larger_1, tl.min(logits_blk2))
205 +
206 + # Second pass: Calculate num_min_larger
207 + for i in range(0, search_iters):
208 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
209 + 0, BLOCK_SIZE_TRUNC
210 + )
211 + mask_n_2 = offs_n < search_range
212 + logits_blk2 = tl.load(
213 + BUFFER_ROW + offs_n, mask=mask_n_2, other=-float("inf")
214 + )
215 +
216 + num_min_larger_0 += tl.sum(
217 + tl.abs(logits_blk2 - min_larger_0) < 1e-9
218 + )
219 + num_min_larger_1 += tl.sum(
220 + tl.abs(logits_blk2 - min_larger_1) < 1e-9
221 + )
222 +
223 + # Check if any of the pivots satisfy termination condition
224 + if (
225 + k_pivots_num_0 >= k
226 + and k_pivots_num_0 - num_min_larger_0 < k
227 + ):
228 + k_pivot = k_pivot_0
229 + k_pivots_num = k_pivots_num_0
230 + min_larger = min_larger_0
231 + num_min_larger = num_min_larger_0
232 + found_pivot = 1
233 + if (
234 + k_pivots_num_1 >= k
235 + and k_pivots_num_1 - num_min_larger_1 < k
236 + ):
237 + k_pivot = k_pivot_1
238 + k_pivots_num = k_pivots_num_1
239 + min_larger = min_larger_1
240 + num_min_larger = num_min_larger_1
241 + found_pivot = 1
242 +
243 + # Update range
244 + if k_pivots_num_1 > k:
245 + min_range = k_pivot_1
246 + elif k_pivots_num_0 > k:
247 + min_range = k_pivot_0
248 +
249 + if k_pivots_num_0 < k:
250 + max_range = k_pivot_0
251 + elif k_pivots_num_1 < k:
252 + max_range = k_pivot_1
253 +
254 + num_iters += 1
255 + if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9:
256 + k_pivot = (max_range + min_range) / 2.0
257 + found_pivot = 1
258 + else:
259 + # If top-k outlier gathering failed, search whole logit space
260 + max_range = max_logit
261 + min_range = min_logit
262 + found_pivot = 0
263 + while found_pivot == 0:
264 + k_pivot_0 = (max_range - min_range) * 1.0 / 4.0 + min_range
265 + k_pivots_num_0 = tl.zeros((), dtype=tl.uint32)
266 + min_larger_0 = float("inf")
267 + num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
268 +
269 + k_pivot_1 = (max_range - min_range) * 2.0 / 4.0 + min_range
270 + k_pivots_num_1 = tl.zeros((), dtype=tl.uint32)
271 + min_larger_1 = float("inf")
272 + num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
273 +
274 + # First pass: Calculate k_pivots_num and min_larger
275 + for i in range(0, NUM_TILES):
276 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
277 + mask_n = offs_n < VOCAB_SIZE
278 + logits_blk2 = tl.load(
279 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
280 + )
281 +
282 + k_pivots_num_0 += tl.sum(logits_blk2 > k_pivot_0)
283 + k_pivots_num_1 += tl.sum(logits_blk2 > k_pivot_1)
284 +
285 + # Exclude -inf from min_larger to avoid
286 + # poisoning the convergence check.
287 + finite_blk2 = tl.where(
288 + logits_blk2 > -float("inf"), logits_blk2, float("inf")
289 + )
290 + min_larger_0 = tl.minimum(min_larger_0, tl.min(finite_blk2))
291 + min_larger_1 = tl.minimum(min_larger_1, tl.min(finite_blk2))
292 +
293 + # Second pass: Calculate num_min_larger
294 + for i in range(0, NUM_TILES):
295 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
296 + mask_n = offs_n < VOCAB_SIZE
297 + logits_blk2 = tl.load(
298 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
299 + )
300 +
301 + num_min_larger_0 += tl.sum(
302 + tl.abs(logits_blk2 - min_larger_0) < 1e-9
303 + )
304 + num_min_larger_1 += tl.sum(
305 + tl.abs(logits_blk2 - min_larger_1) < 1e-9
306 + )
307 +
308 + # Check if any of the pivots satisfy termination condition
309 + if (
310 + k_pivots_num_0 >= k
311 + and k_pivots_num_0 - num_min_larger_0 < k
312 + ):
313 + k_pivot = k_pivot_0
314 + k_pivots_num = k_pivots_num_0
315 + min_larger = min_larger_0
316 + num_min_larger = num_min_larger_0
317 + found_pivot = 1
318 + if (
319 + k_pivots_num_1 >= k
320 + and k_pivots_num_1 - num_min_larger_1 < k
321 + ):
322 + k_pivot = k_pivot_1
323 + k_pivots_num = k_pivots_num_1
324 + min_larger = min_larger_1
325 + num_min_larger = num_min_larger_1
326 + found_pivot = 1
327 +
328 + # Update range
329 + if k_pivots_num_1 > k:
330 + min_range = k_pivot_1
331 + elif k_pivots_num_0 > k:
332 + min_range = k_pivot_0
333 +
334 + if k_pivots_num_0 < k:
335 + max_range = k_pivot_0
336 + elif k_pivots_num_1 < k:
337 + max_range = k_pivot_1
338 +
339 + num_iters += 1
340 + if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9:
341 + k_pivot = (max_range + min_range) / 2.0
342 + found_pivot = 1
343 +
344 + duplicate_logit = min_larger
345 + num_duplicate_logit = num_min_larger
346 + num_keep = num_duplicate_logit - (k_pivots_num - k)
347 + num_kept = tl.zeros((), dtype=tl.uint32)
348 +
349 + # Top-k only path. If there are fewer finite values
350 + # than k (e.g. grammar mask), keep everything.
351 + final_pivot = k_pivot if num_finite_total > k else -float("inf")
352 +
353 + if TOPP_ENABLED and num_finite_total > k:
354 + #### TOP-P SAMPLING AFTER TOP-K ####
355 + p = tl.load(P + row_id)
356 + if p < 1.0:
357 + min_logit = k_pivot
358 + sum_exp_logits = 0.0
359 + num_outliers_2 = tl.zeros((), dtype=tl.uint32)
360 + search_range = tl.cast(num_outliers, tl.int32)
361 + search_iters = tl.cast(
362 + (num_outliers + BLOCK_SIZE_TRUNC - 1) // BLOCK_SIZE_TRUNC,
363 + tl.int32,
364 + )
365 +
366 + # Third pass: Calculate exp logits and sum, gather outliers
367 + if num_outliers > k:
368 + for i in range(0, search_iters):
369 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
370 + 0, BLOCK_SIZE_TRUNC
371 + )
372 + mask_n_2 = offs_n < search_range
373 +
374 + probs_blk = tl.load(
375 + BUFFER_ROW + offs_n,
376 + mask=mask_n_2,
377 + other=-float("inf"),
378 + )
379 +
380 + outlier_mask = (probs_blk > min_logit) & mask_n_2
381 +
382 + # Duplicate logit handling for Top-k
383 + if num_keep < num_duplicate_logit:
384 + duplicate_mask = (
385 + tl.abs(probs_blk - duplicate_logit) < 1e-9
386 + )
387 + duplicate_count = (
388 + tl.cumsum(duplicate_mask) + num_kept
389 + )
390 + duplicate_keep_mask = (
391 + duplicate_count <= num_keep
392 + ) & duplicate_mask
393 + duplicate_remove_mask = (
394 + duplicate_mask & ~duplicate_keep_mask
395 + )
396 + outlier_mask = outlier_mask & (
397 + ~duplicate_remove_mask
398 + )
399 + num_kept += tl.sum(duplicate_keep_mask)
400 +
401 + probs_blk = tl.where(
402 + outlier_mask, probs_blk, -float("inf")
403 + )
404 + probs_blk = probs_blk - max_logit
405 + probs_blk = tl.exp(probs_blk)
406 + sum_exp_logits += tl.sum(probs_blk)
407 +
408 + # Fourth pass: Calculate BUFFER and get outliers
409 + for i in range(0, search_iters):
410 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
411 + 0, BLOCK_SIZE_TRUNC
412 + )
413 + mask_n_2 = offs_n < search_range
414 +
415 + probs_blk = tl.load(
416 + BUFFER_ROW + offs_n,
417 + mask=mask_n_2,
418 + other=-float("inf"),
419 + )
420 +
421 + probs_blk = probs_blk - max_logit
422 + probs_blk = tl.exp(probs_blk)
423 + probs_blk = probs_blk / sum_exp_logits
424 + tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n_2)
425 + else:
426 + # If top-k outlier gathering failed,
427 + # retry gathering using top-k pivot
428 + for i in range(0, NUM_TILES):
429 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
430 + mask_n = offs_n < VOCAB_SIZE
431 +
432 + probs_blk = tl.load(
433 + LOGITS_ROW + offs_n,
434 + mask=mask_n,
435 + other=-float("inf"),
436 + )
437 +
438 + outlier_mask = (probs_blk > min_logit) & mask_n
439 +
440 + # Duplicate logit handling for Top-k
441 + duplicate_mask = (
442 + tl.abs(probs_blk - duplicate_logit) < 1e-9
443 + )
444 + duplicate_count = tl.cumsum(duplicate_mask) + num_kept
445 + duplicate_keep_mask = (
446 + duplicate_count <= num_keep
447 + ) & duplicate_mask
448 + duplicate_remove_mask = (
449 + duplicate_mask & ~duplicate_keep_mask
450 + )
451 + outlier_mask = outlier_mask & (~duplicate_remove_mask)
452 + num_kept += tl.sum(duplicate_keep_mask)
453 +
454 + probs_blk = tl.where(
455 + outlier_mask, probs_blk, -float("inf")
456 + )
457 + probs_blk = probs_blk - max_logit
458 + probs_blk = tl.exp(probs_blk)
459 + sum_exp_logits += tl.sum(probs_blk)
460 +
461 + cumulative_pos = tl.cast(
462 + tl.cumsum(outlier_mask) - 1 + num_outliers_2,
463 + tl.int32,
464 + )
465 + num_outliers_2 += tl.sum(outlier_mask)
466 + write_pos = tl.where(outlier_mask, cumulative_pos, -1)
467 + tl.store(
468 + BUFFER_ROW + write_pos, probs_blk, mask=outlier_mask
469 + )
470 +
471 + search_range = tl.cast(num_outliers_2, tl.int32)
472 + search_iters = tl.cast(
473 + (num_outliers_2 + BLOCK_SIZE_TRUNC - 1)
474 + // BLOCK_SIZE_TRUNC,
475 + tl.int32,
476 + )
477 +
478 + # Fourth pass: Calculate BUFFER and get outliers
479 + for i in range(0, search_iters):
480 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
481 + 0, BLOCK_SIZE_TRUNC
482 + )
483 + mask_n_2 = offs_n < search_range
484 +
485 + probs_blk = tl.load(
486 + BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
487 + )
488 + probs_blk = probs_blk / sum_exp_logits
489 + tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n_2)
490 +
491 + max_range = tl.exp(max_logit - max_logit) / sum_exp_logits
492 + min_range = tl.exp(min_logit - max_logit) / sum_exp_logits
493 +
494 + p_pivot = 1.0
495 + num_iters = 0
496 + min_larger_prob = 1.0
497 + num_min_larger = tl.zeros((), dtype=tl.uint32)
498 + p_pivots_sum = 0.0
499 +
500 + # Fifth passes: Search for p_pivot
501 + found_pivot = 0
502 + while found_pivot == 0:
503 + p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
504 + p_pivots_sum_0 = 0.0
505 + min_larger_0 = 1.0
506 + num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
507 +
508 + p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
509 + p_pivots_sum_1 = 0.0
510 + min_larger_1 = 1.0
511 + num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
512 +
513 + # First pass: Calculate p_pivots_sum and min_larger
514 + for i in range(0, search_iters):
515 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
516 + 0, BLOCK_SIZE_TRUNC
517 + )
518 + mask_n_2 = offs_n < search_range
519 + probs_blk = tl.load(
520 + BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
521 + )
522 +
523 + p_pivots_sum_0 += tl.sum(
524 + probs_blk * (probs_blk > p_pivot_0)
525 + )
526 + masked_larger_0 = tl.where(
527 + probs_blk > p_pivot_0, probs_blk, 1.0
528 + )
529 + min_larger_0 = tl.minimum(
530 + min_larger_0, tl.min(masked_larger_0)
531 + )
532 +
533 + p_pivots_sum_1 += tl.sum(
534 + probs_blk * (probs_blk > p_pivot_1)
535 + )
536 + masked_larger_1 = tl.where(
537 + probs_blk > p_pivot_1, probs_blk, 1.0
538 + )
539 + min_larger_1 = tl.minimum(
540 + min_larger_1, tl.min(masked_larger_1)
541 + )
542 +
543 + # Second pass: Calculate num_min_larger
544 + for i in range(0, search_iters):
545 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
546 + 0, BLOCK_SIZE_TRUNC
547 + )
548 + mask_n_2 = offs_n < search_range
549 + probs_blk = tl.load(
550 + BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
551 + )
552 +
553 + num_min_larger_0 += tl.sum(
554 + tl.abs(probs_blk - min_larger_0) < 1e-9
555 + )
556 + num_min_larger_1 += tl.sum(
557 + tl.abs(probs_blk - min_larger_1) < 1e-9
558 + )
559 +
560 + # Check if any of the pivots satisfy termination condition
561 + if p_pivots_sum_1 >= p and (
562 + p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p
563 + ):
564 + p_pivot = p_pivot_1
565 + min_larger_prob = min_larger_1
566 + num_min_larger = num_min_larger_1
567 + p_pivots_sum = p_pivots_sum_1
568 + found_pivot = 1
569 + if p_pivots_sum_0 >= p and (
570 + p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p
571 + ):
572 + p_pivot = p_pivot_0
573 + min_larger_prob = min_larger_0
574 + num_min_larger = num_min_larger_0
575 + p_pivots_sum = p_pivots_sum_0
576 + found_pivot = 1
577 +
578 + # Update range
579 + if p_pivots_sum_1 > p:
580 + min_range = p_pivot_1
581 + elif p_pivots_sum_0 > p:
582 + min_range = p_pivot_0
583 +
584 + if p_pivots_sum_0 < p:
585 + max_range = p_pivot_0
586 + elif p_pivots_sum_1 < p:
587 + max_range = p_pivot_1
588 +
589 + num_iters += 1
590 + if (max_range - min_range) < 1e-9 or num_iters >= 18:
591 + p_pivot = (max_range + min_range) / 2.0
592 + found_pivot = 1
593 +
594 + duplicate_logit = (
595 + tl.log(min_larger_prob * sum_exp_logits) + max_logit
596 + )
597 + num_duplicate_logit = num_min_larger
598 + num_keep = num_duplicate_logit - tl.cast(
599 + (p_pivots_sum - p) / min_larger_prob, tl.uint32
600 + )
601 + num_kept = tl.zeros((), dtype=tl.uint32)
602 +
603 + # Top-k + Top-p path
604 + final_pivot = tl.log(p_pivot * sum_exp_logits) + max_logit
605 +
606 + if TOPP_ENABLED and final_pivot == -float("inf"):
607 + #### STANDALONE TOP-P SAMPLING ####
608 + p = tl.load(P + row_id)
609 + if p < 1.0:
610 + # Zeroth pass: Compute avg and std from a sample block
611 + offs = tl.arange(0, BLOCK_SIZE)
612 + mask_n = offs < VOCAB_SIZE
613 + logits_blk0 = tl.load(
614 + LOGITS_ROW + offs, mask=mask_n, other=-float("inf")
615 + )
616 + # Exclude -inf values (e.g. from grammar bitmasks) from
617 + # statistics to avoid NaN in pivot computation.
618 + finite_mask = (logits_blk0 > -float("inf")) & mask_n
619 + num_finite = tl.sum(finite_mask)
620 + finite_logits = tl.where(finite_mask, logits_blk0, 0.0)
621 + avg_logit = tl.where(
622 + num_finite > 0, tl.sum(finite_logits) / num_finite, 0.0
623 + )
624 + sq_avg_logit = tl.where(
625 + num_finite > 0,
626 + tl.sum(finite_logits * finite_logits) / num_finite,
627 + 0.0,
628 + )
629 + std_logit = tl.sqrt(
630 + tl.maximum(sq_avg_logit - avg_logit * avg_logit, 0.0)
631 + )
632 + max_sample = avg_logit + std_logit * 10.0
633 + sum_exp_logits = 0.0
634 +
635 + # First pass: compute max and min logits and sum_exp_logits
636 + for i in range(0, NUM_TILES):
637 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
638 + mask_n = offs_n < VOCAB_SIZE
639 + logits_blk = tl.load(
640 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
641 + )
642 + max_logit = tl.maximum(max_logit, tl.max(logits_blk))
643 + # Exclude -inf from min to keep binary search bounds
644 + # finite (avoids NaN pivots).
645 + finite_blk = tl.where(
646 + logits_blk > -float("inf"), logits_blk, float("inf")
647 + )
648 + min_logit = tl.minimum(min_logit, tl.min(finite_blk))
649 +
650 + probs_blk = tl.exp(logits_blk - max_sample)
651 + probs_blk = tl.where(mask_n, probs_blk, 0.0)
652 + sum_exp_logits += tl.sum(probs_blk)
653 +
654 + # If no finite logits exist (all -inf), clamp min to
655 + # max so the search converges to -inf (no masking).
656 + min_logit = tl.minimum(min_logit, max_logit)
657 +
658 + idx = tl.cast(p * 200, tl.int32)
659 + idx = tl.maximum(0, tl.minimum(idx, 199))
660 + sigma = tl.load(NORMAL_CDF_TO_SIGMA_TABLE + idx)
661 + sigma = sigma + tl.abs(sigma) * -0.25
662 + outlier_pivot = avg_logit + std_logit * sigma
663 +
664 + outlier_prob = tl.exp(outlier_pivot - max_sample) / sum_exp_logits
665 + sum_outlier_probs = 0.0
666 + num_outliers = tl.zeros((), dtype=tl.uint32)
667 +
668 + # Second pass: Calculate softmax and gather outliers
669 + for i in range(0, NUM_TILES):
670 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
671 + mask_n = offs_n < VOCAB_SIZE
672 +
673 + probs_blk = tl.load(
674 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
675 + )
676 + probs_blk = tl.exp(probs_blk - max_sample)
677 + probs_blk = probs_blk / sum_exp_logits
678 +
679 + outlier_mask = (probs_blk > outlier_prob) & mask_n
680 + sum_outlier_probs += tl.sum(outlier_mask * probs_blk)
681 + cumulative_pos = tl.cast(
682 + tl.cumsum(outlier_mask) - 1 + num_outliers, tl.int32
683 + )
684 + num_outliers += tl.sum(outlier_mask)
685 + write_pos = tl.where(outlier_mask, cumulative_pos, -1)
686 + tl.store(BUFFER_ROW + write_pos, probs_blk, mask=outlier_mask)
687 +
688 + max_range = tl.exp(max_logit - max_sample) / sum_exp_logits
689 + min_range = tl.exp(min_logit - max_sample) / sum_exp_logits
690 +
691 + p_pivot = 1.0
692 + num_iters = 0
693 + min_larger_prob = 1.0
694 + num_min_larger = tl.zeros((), dtype=tl.uint32)
695 + p_pivots_sum = 0.0
696 +
697 + # Third pass: Search for p_pivot
698 + if sum_outlier_probs > p:
699 + min_range = outlier_prob
700 + search_range = tl.cast(num_outliers, tl.int32)
701 + search_iters = tl.cast(
702 + (num_outliers + BLOCK_SIZE_TRUNC - 1) // BLOCK_SIZE_TRUNC,
703 + tl.int32,
704 + )
705 +
706 + found_pivot = 0
707 + while found_pivot == 0:
708 + p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
709 + p_pivots_sum_0 = 0.0
710 + min_larger_0 = 1.0
711 + num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
712 +
713 + p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
714 + p_pivots_sum_1 = 0.0
715 + min_larger_1 = 1.0
716 + num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
717 +
718 + # First pass: Calculate p_pivots_sum and min_larger
719 + for i in range(0, search_iters):
720 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
721 + 0, BLOCK_SIZE_TRUNC
722 + )
723 + mask_n_2 = offs_n < search_range
724 + probs_blk = tl.load(
725 + BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
726 + )
727 +
728 + p_pivots_sum_0 += tl.sum(
729 + probs_blk * (probs_blk > p_pivot_0)
730 + )
731 + masked_larger_0 = tl.where(
732 + probs_blk > p_pivot_0, probs_blk, 1.0
733 + )
734 + min_larger_0 = tl.minimum(
735 + min_larger_0, tl.min(masked_larger_0)
736 + )
737 +
738 + p_pivots_sum_1 += tl.sum(
739 + probs_blk * (probs_blk > p_pivot_1)
740 + )
741 + masked_larger_1 = tl.where(
742 + probs_blk > p_pivot_1, probs_blk, 1.0
743 + )
744 + min_larger_1 = tl.minimum(
745 + min_larger_1, tl.min(masked_larger_1)
746 + )
747 +
748 + # Second pass: Calculate num_min_larger
749 + for i in range(0, search_iters):
750 + offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
751 + 0, BLOCK_SIZE_TRUNC
752 + )
753 + mask_n_2 = offs_n < search_range
754 + probs_blk = tl.load(
755 + BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
756 + )
757 +
758 + num_min_larger_0 += tl.sum(
759 + tl.abs(probs_blk - min_larger_0) < 1e-9
760 + )
761 + num_min_larger_1 += tl.sum(
762 + tl.abs(probs_blk - min_larger_1) < 1e-9
763 + )
764 +
765 + # Check if any of the pivots satisfy termination condition
766 + if (
767 + p_pivots_sum_1 >= p
768 + and p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p
769 + ):
770 + p_pivot = p_pivot_1
771 + min_larger_prob = min_larger_1
772 + num_min_larger = num_min_larger_1
773 + p_pivots_sum = p_pivots_sum_1
774 + found_pivot = 1
775 + if (
776 + p_pivots_sum_0 >= p
777 + and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p
778 + ):
779 + p_pivot = p_pivot_0
780 + min_larger_prob = min_larger_0
781 + num_min_larger = num_min_larger_0
782 + p_pivots_sum = p_pivots_sum_0
783 + found_pivot = 1
784 +
785 + # Update range
786 + if p_pivots_sum_1 > p:
787 + min_range = p_pivot_1
788 + elif p_pivots_sum_0 > p:
789 + min_range = p_pivot_0
790 +
791 + if p_pivots_sum_0 < p:
792 + max_range = p_pivot_0
793 + elif p_pivots_sum_1 < p:
794 + max_range = p_pivot_1
795 +
796 + num_iters += 1
797 + if (max_range - min_range) < 1e-9 or num_iters >= 18:
798 + p_pivot = (max_range + min_range) / 2.0
799 + found_pivot = 1
800 + else:
801 + # Re-populate the buffer with full softmax probabilities
802 + for i in range(0, NUM_TILES):
803 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
804 + mask_n = offs_n < VOCAB_SIZE
805 +
806 + probs_blk = tl.load(
807 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
808 + )
809 + probs_blk = tl.exp(probs_blk - max_sample)
810 + probs_blk = probs_blk / sum_exp_logits
811 + tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n)
812 +
813 + found_pivot = 0
814 + while found_pivot == 0:
815 + p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
816 + p_pivots_sum_0 = 0.0
817 + min_larger_0 = 1.0
818 + num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
819 +
820 + p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
821 + p_pivots_sum_1 = 0.0
822 + min_larger_1 = 1.0
823 + num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
824 +
825 + # First pass: Calculate p_pivots_sum and min_larger
826 + for i in range(0, NUM_TILES):
827 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
828 + mask_n = offs_n < VOCAB_SIZE
829 + probs_blk = tl.load(
830 + BUFFER_ROW + offs_n, mask=mask_n, other=0.0
831 + )
832 +
833 + p_pivots_sum_0 += tl.sum(
834 + probs_blk * (probs_blk > p_pivot_0)
835 + )
836 + masked_larger_0 = tl.where(
837 + probs_blk > p_pivot_0, probs_blk, 1.0
838 + )
839 + min_larger_0 = tl.minimum(
840 + min_larger_0, tl.min(masked_larger_0)
841 + )
842 +
843 + p_pivots_sum_1 += tl.sum(
844 + probs_blk * (probs_blk > p_pivot_1)
845 + )
846 + masked_larger_1 = tl.where(
847 + probs_blk > p_pivot_1, probs_blk, 1.0
848 + )
849 + min_larger_1 = tl.minimum(
850 + min_larger_1, tl.min(masked_larger_1)
851 + )
852 +
853 + # Second pass: Calculate num_min_larger
854 + for i in range(0, NUM_TILES):
855 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
856 + mask_n = offs_n < VOCAB_SIZE
857 + probs_blk = tl.load(
858 + BUFFER_ROW + offs_n, mask=mask_n, other=0.0
859 + )
860 +
861 + num_min_larger_0 += tl.sum(
862 + tl.abs(probs_blk - min_larger_0) < 1e-9
863 + )
864 + num_min_larger_1 += tl.sum(
865 + tl.abs(probs_blk - min_larger_1) < 1e-9
866 + )
867 +
868 + # Check if any of the pivots satisfy termination condition
869 + if (
870 + p_pivots_sum_1 >= p
871 + and p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p
872 + ):
873 + p_pivot = p_pivot_1
874 + min_larger_prob = min_larger_1
875 + num_min_larger = num_min_larger_1
876 + p_pivots_sum = p_pivots_sum_1
877 + found_pivot = 1
878 + if (
879 + p_pivots_sum_0 >= p
880 + and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p
881 + ):
882 + p_pivot = p_pivot_0
883 + min_larger_prob = min_larger_0
884 + num_min_larger = num_min_larger_0
885 + p_pivots_sum = p_pivots_sum_0
886 + found_pivot = 1
887 +
888 + # Update range
889 + if p_pivots_sum_1 > p:
890 + min_range = p_pivot_1
891 + elif p_pivots_sum_0 > p:
892 + min_range = p_pivot_0
893 +
894 + if p_pivots_sum_0 < p:
895 + max_range = p_pivot_0
896 + elif p_pivots_sum_1 < p:
897 + max_range = p_pivot_1
898 +
899 + num_iters += 1
900 + if (max_range - min_range) < 1e-9 or num_iters >= 18:
901 + p_pivot = (max_range + min_range) / 2.0
902 + found_pivot = 1
903 +
904 + duplicate_logit = tl.log(min_larger_prob * sum_exp_logits) + max_logit
905 + num_duplicate_logit = num_min_larger
906 + num_keep = num_duplicate_logit - tl.cast(
907 + (p_pivots_sum - p) / min_larger_prob, tl.uint32
908 + )
909 + num_kept = tl.zeros((), dtype=tl.uint32)
910 +
911 + # Top-p only path
912 + final_pivot = tl.log(p_pivot * sum_exp_logits) + max_sample
913 +
914 + # Sixth pass: Apply mask and store final output.
915 + # If the pivot >= max logit (or is NaN), no token would
916 + # survive the strict `>` keep_mask. Skip masking.
917 + # Using `not <` instead of `>=` so that NaN is also caught.
918 + if not (final_pivot < max_logit):
919 + final_pivot = -float("inf")
920 + elif final_pivot != -float("inf"):
921 + for i in range(0, NUM_TILES):
922 + offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
923 + mask_n = offs_n < VOCAB_SIZE
924 + logits_blk = tl.load(
925 + LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
926 + )
927 + keep_mask = (logits_blk > final_pivot) & mask_n
928 +
929 + # Duplicate logit handling
930 + if num_keep < num_duplicate_logit:
931 + duplicate_mask = (
932 + tl.abs(logits_blk - duplicate_logit) < 1e-9
933 + ) & mask_n
934 + duplicate_count = tl.cumsum(duplicate_mask) + num_kept
935 + duplicate_keep_mask = (
936 + duplicate_count <= num_duplicate_logit
937 + ) & duplicate_mask
938 + duplicate_remove_mask = duplicate_mask & ~duplicate_keep_mask
939 + num_kept += tl.sum(duplicate_keep_mask)
940 + keep_mask = keep_mask & (~duplicate_remove_mask)
941 +
942 + logits_blk = tl.where(keep_mask, logits_blk, MASK_VALUE)
943 + tl.store(LOGITS_ROW + offs_n, logits_blk, mask=mask_n)
944 +
945 +
946 +def apply_top_k_top_p_triton(
947 + logits: torch.Tensor,
948 + k: torch.Tensor | None,
949 + p: torch.Tensor | None,
950 + mask_value: float = float("-inf"),
951 +) -> torch.Tensor:
952 + """
953 + Apply combined top-k and top-p masking using Triton.
954 +
955 + Top-k is applied first (by logit value), then top-p is applied
956 + to the remaining k values (by probability).
957 +
958 + Args:
959 + logits: [batch_size, vocab_size] float32 tensor, modified in-place
960 + k: [batch_size] int32 tensor of top-k values per row, or None to disable top-k
961 + p: [batch_size] float32 tensor of top-p values per row (0 to 1),
962 + or None to disable top-p
963 + mask_value: Value for masked positions (default: -inf)
964 +
965 + Returns:
966 + The logits tensor (modified in-place)
967 + """
968 + assert logits.ndim == 2
969 + assert logits.dtype == torch.float32
970 + assert logits.is_cuda
971 +
972 + batch_size, vocab_size = logits.shape
973 +
974 + topk_enabled = k is not None
975 + topp_enabled = p is not None
976 +
977 + if batch_size == 0 or not (topk_enabled or topp_enabled):
978 + return logits
979 +
980 + if k is not None:
981 + assert k.ndim == 1 and k.shape[0] == batch_size and k.is_cuda
982 + k_ptr = k.to(torch.int32)
983 + else:
984 + k_ptr = logits # Dummy pointer (won't be read)
985 +
986 + if p is not None:
987 + assert p.ndim == 1 and p.shape[0] == batch_size and p.is_cuda
988 + p_ptr = p.to(torch.float32)
989 + else:
990 + p_ptr = logits # Dummy pointer (won't be read)
991 +
992 + num_sm = torch.cuda.get_device_properties(logits.device).multi_processor_count
993 + NUM_PROGRAMS = min(num_sm, batch_size)
994 +
995 + # Cache per-Triton Program buffer on each device.
996 + buf_key = (logits.device, logits.dtype, vocab_size)
997 + buffer = _TRITON_BUFFER_CACHE.get(buf_key)
998 + if buffer is None or buffer.shape[0] < NUM_PROGRAMS:
999 + size = min(next_power_of_2(NUM_PROGRAMS), num_sm)
1000 + buffer = logits.new_empty((size, vocab_size))
1001 + _TRITON_BUFFER_CACHE[buf_key] = buffer
1002 + if buffer.shape[0] > NUM_PROGRAMS:
1003 + buffer = buffer[:NUM_PROGRAMS]
1004 +
1005 + # Cache lookup table entries on each device.
1006 + tables = _TRITON_TABLE_CACHE.get(logits.device)
1007 + if tables is None:
1008 + normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)
1009 + percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)
1010 + _TRITON_TABLE_CACHE[logits.device] = (
1011 + normal_cdf_to_sigma_table,
1012 + percentile_to_std_table,
1013 + )
1014 + else:
1015 + normal_cdf_to_sigma_table, percentile_to_std_table = tables
1016 +
1017 + _topk_topp_kernel[(NUM_PROGRAMS,)](
1018 + logits,
1019 + buffer,
1020 + percentile_to_std_table,
1021 + normal_cdf_to_sigma_table,
1022 + k_ptr,
1023 + p_ptr,
1024 + BATCH_SIZE=batch_size,
1025 + MASK_VALUE=mask_value,
1026 + VOCAB_SIZE=vocab_size,
1027 + BLOCK_SIZE=8192,
1028 + BLOCK_SIZE_TRUNC=4096,
1029 + TOPK_ENABLED=topk_enabled,
1030 + TOPP_ENABLED=topp_enabled,
1031 + )
1032 +
1033 + return logits
1034 +
1035 +
1036 +def reset_buffer_cache():
1037 + _TRITON_BUFFER_CACHE.clear()
1038 + _TRITON_TABLE_CACHE.clear()
1039 + torch.cuda.empty_cache()