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 6, 13:15 UTC
[Perf] Optimize top-k search in apply_top_k_top_p_triton sampler (#37225)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Mar 17, 2026 at 19:35 UTC
51b2333be19000db7d03b76ccf1b842972c98541
1 file changed
+63
-45
vllm/v1/sample/ops/topk_topp_triton.py
+63
-45
@@ -67,6 +67,29 @@ _PERCENTILE_TO_STD_TABLE = [
67
# fmt: on
68
69
70
+@triton.jit
71
+def _update_min_larger_stats(data, above_mask, min_larger, num_min_larger, sentinel):
72
+ """Update running (min, count) of values above a pivot across tiles.
73
+
74
+ Tracks the smallest value strictly above a pivot and how many times
75
+ it occurs. Called once per tile per pivot; the running state is
76
+ carried across tiles via `min_larger` / `num_min_larger`.
77
+
78
+ Merge rule:
79
+ - tile min < running min → replace both
80
+ - tile min == running min → accumulate count
81
+ - tile min > running min → keep running values
82
+ """
83
+ tile_min = tl.min(tl.where(above_mask, data, sentinel))
84
+ tile_eq = above_mask & (tl.abs(data - tile_min) < 1e-9)
85
+ tile_cnt = tl.sum(tile_eq)
86
+ is_new = tile_min < min_larger
87
+ is_same = tl.abs(tile_min - min_larger) < 1e-9
88
+ num_min_larger = tl.where(is_new, tile_cnt, num_min_larger + tile_cnt * is_same)
89
+ min_larger = tl.minimum(min_larger, tile_min)
90
+ return min_larger, num_min_larger
91
+
92
+
93
@triton.jit
94
def _topk_topp_kernel(
95
LOGITS,
@@ -188,7 +211,10 @@ def _topk_topp_kernel(
211
min_larger_1 = float("inf")
212
num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
213
191
- # First pass: Calculate k_pivots_num and min_larger
214
+ # Single fused pass: compute k_pivots_num,
215
+ # min_larger, and num_min_larger together to avoid
216
+ # a second data scan. See _update_min_larger_stats
217
+ # for the tile-level merge logic.
218
for i in range(0, search_iters):
219
offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
220
0, BLOCK_SIZE_TRUNC
@@ -198,27 +224,24 @@ def _topk_topp_kernel(
224
BUFFER_ROW + offs_n, mask=mask_n_2, other=-float("inf")
225
)
226
201
- k_pivots_num_0 += tl.sum(logits_blk2 > k_pivot_0)
202
- k_pivots_num_1 += tl.sum(logits_blk2 > k_pivot_1)
203
-
204
- min_larger_0 = tl.minimum(min_larger_0, tl.min(logits_blk2))
205
- min_larger_1 = tl.minimum(min_larger_1, tl.min(logits_blk2))
206
-
207
- # Second pass: Calculate num_min_larger
208
- for i in range(0, search_iters):
209
- offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
210
- 0, BLOCK_SIZE_TRUNC
227
+ above_0 = logits_blk2 > k_pivot_0
228
+ above_1 = logits_blk2 > k_pivot_1
229
+ k_pivots_num_0 += tl.sum(above_0)
230
+ k_pivots_num_1 += tl.sum(above_1)
231
+
232
+ min_larger_0, num_min_larger_0 = _update_min_larger_stats(
233
+ logits_blk2,
234
+ above_0,
235
+ min_larger_0,
236
+ num_min_larger_0,
237
+ float("inf"),
238
)
212
- mask_n_2 = offs_n < search_range
213
- logits_blk2 = tl.load(
214
- BUFFER_ROW + offs_n, mask=mask_n_2, other=-float("inf")
215
- )
216
-
217
- num_min_larger_0 += tl.sum(
218
- tl.abs(logits_blk2 - min_larger_0) < 1e-9
219
- )
220
- num_min_larger_1 += tl.sum(
221
- tl.abs(logits_blk2 - min_larger_1) < 1e-9
239
+ min_larger_1, num_min_larger_1 = _update_min_larger_stats(
240
+ logits_blk2,
241
+ above_1,
242
+ min_larger_1,
243
+ num_min_larger_1,
244
+ float("inf"),
245
)
246
247
# Check if any of the pivots satisfy termination condition
@@ -272,26 +295,8 @@ def _topk_topp_kernel(
295
min_larger_1 = float("inf")
296
num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
297
275
- # First pass: Calculate k_pivots_num and min_larger
276
- for i in range(0, NUM_TILES):
277
- offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
278
- mask_n = offs_n < VOCAB_SIZE
279
- logits_blk2 = tl.load(
280
- LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
281
- )
282
-
283
- k_pivots_num_0 += tl.sum(logits_blk2 > k_pivot_0)
284
- k_pivots_num_1 += tl.sum(logits_blk2 > k_pivot_1)
285
-
286
- # Exclude -inf from min_larger to avoid
287
- # poisoning the convergence check.
288
- finite_blk2 = tl.where(
289
- logits_blk2 > -float("inf"), logits_blk2, float("inf")
290
- )
291
- min_larger_0 = tl.minimum(min_larger_0, tl.min(finite_blk2))
292
- min_larger_1 = tl.minimum(min_larger_1, tl.min(finite_blk2))
293
-
294
- # Second pass: Calculate num_min_larger
298
+ # Single fused pass over full vocab (same approach
299
+ # as the buffer path above).
300
for i in range(0, NUM_TILES):
301
offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
302
mask_n = offs_n < VOCAB_SIZE
@@ -299,11 +304,24 @@ def _topk_topp_kernel(
304
LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")
305
)
306
302
- num_min_larger_0 += tl.sum(
303
- tl.abs(logits_blk2 - min_larger_0) < 1e-9
307
+ above_0 = logits_blk2 > k_pivot_0
308
+ above_1 = logits_blk2 > k_pivot_1
309
+ k_pivots_num_0 += tl.sum(above_0)
310
+ k_pivots_num_1 += tl.sum(above_1)
311
+
312
+ min_larger_0, num_min_larger_0 = _update_min_larger_stats(
313
+ logits_blk2,
314
+ above_0,
315
+ min_larger_0,
316
+ num_min_larger_0,
317
+ float("inf"),
318
)
305
- num_min_larger_1 += tl.sum(
306
- tl.abs(logits_blk2 - min_larger_1) < 1e-9
319
+ min_larger_1, num_min_larger_1 = _update_min_larger_stats(
320
+ logits_blk2,
321
+ above_1,
322
+ min_larger_1,
323
+ num_min_larger_1,
324
+ float("inf"),
325
)
326
327
# Check if any of the pivots satisfy termination condition