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 2, 03:30 UTC
[Perf] Apply single-pass min_larger finding and binary search in Triton Top-p path. (#42191)
Signed-off-by: js_park <cakeng@naver.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Jongseok Park committed
Jun 2, 2026 at 17:57 UTC
ca17b6b17d988bc848d5ca629e677513d3ce6c4a
1 file changed
+68
-178
vllm/v1/sample/ops/topk_topp_triton.py
+68
-178
@@ -186,7 +186,7 @@ def _topk_topp_kernel(
186
# max so the search converges to -inf (no masking).
187
min_logit = tl.minimum(min_logit, max_logit)
188
189
- # Second passes: Ternary search for pivots
189
+ # Second passes: Ternary search for pivot
190
num_iters = 0
191
k_pivot = float("inf")
192
k_pivots_num = tl.zeros((), dtype=tl.uint32)
@@ -279,6 +279,8 @@ def _topk_topp_kernel(
279
num_iters += 1
280
if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9:
281
k_pivot = (max_range + min_range) / 2.0
282
+ min_larger = min_larger_0
283
+ num_min_larger = num_min_larger_0
284
found_pivot = 1
285
else:
286
# If top-k outlier gathering failed, search whole logit space
@@ -286,12 +288,12 @@ def _topk_topp_kernel(
288
min_range = min_logit
289
found_pivot = 0
290
while found_pivot == 0:
289
- k_pivot_0 = (max_range - min_range) * 1.0 / 4.0 + min_range
291
+ k_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
292
k_pivots_num_0 = tl.zeros((), dtype=tl.uint32)
293
min_larger_0 = float("inf")
294
num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
295
294
- k_pivot_1 = (max_range - min_range) * 2.0 / 4.0 + min_range
296
+ k_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
297
k_pivots_num_1 = tl.zeros((), dtype=tl.uint32)
298
min_larger_1 = float("inf")
299
num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
@@ -359,6 +361,8 @@ def _topk_topp_kernel(
361
num_iters += 1
362
if num_iters >= 18 or tl.abs(min_range - max_range) < 1e-9:
363
k_pivot = (max_range + min_range) / 2.0
364
+ min_larger = min_larger_0
365
+ num_min_larger = num_min_larger_0
366
found_pivot = 1
367
368
duplicate_logit = min_larger
@@ -520,17 +524,15 @@ def _topk_topp_kernel(
524
# Fifth passes: Search for p_pivot
525
found_pivot = 0
526
while found_pivot == 0:
523
- p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
527
+ p_pivot_0 = (max_range - min_range) * 0.5 + min_range
528
p_pivots_sum_0 = 0.0
529
min_larger_0 = 1.0
530
num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
531
528
- p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
529
- p_pivots_sum_1 = 0.0
530
- min_larger_1 = 1.0
531
- num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
532
-
533
- # First pass: Calculate p_pivots_sum and min_larger
532
+ # Single fused pass: compute p_pivots_sum,
533
+ # min_larger, and num_min_larger together.
534
+ # See _update_min_larger_stats for the
535
+ # tile-level merge logic.
536
for i in range(0, search_iters):
537
offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
538
0, BLOCK_SIZE_TRUNC
@@ -540,52 +542,20 @@ def _topk_topp_kernel(
542
BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
543
)
544
543
- p_pivots_sum_0 += tl.sum(
544
- probs_blk * (probs_blk > p_pivot_0)
545
- )
546
- masked_larger_0 = tl.where(
547
- probs_blk > p_pivot_0, probs_blk, 1.0
548
- )
549
- min_larger_0 = tl.minimum(
550
- min_larger_0, tl.min(masked_larger_0)
551
- )
552
-
553
- p_pivots_sum_1 += tl.sum(
554
- probs_blk * (probs_blk > p_pivot_1)
555
- )
556
- masked_larger_1 = tl.where(
557
- probs_blk > p_pivot_1, probs_blk, 1.0
558
- )
559
- min_larger_1 = tl.minimum(
560
- min_larger_1, tl.min(masked_larger_1)
561
- )
562
-
563
- # Second pass: Calculate num_min_larger
564
- for i in range(0, search_iters):
565
- offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
566
- 0, BLOCK_SIZE_TRUNC
567
- )
568
- mask_n_2 = offs_n < search_range
569
- probs_blk = tl.load(
570
- BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
571
- )
545
+ above_0 = probs_blk > p_pivot_0
546
+ p_pivots_sum_0 += tl.sum(probs_blk * above_0)
547
573
- num_min_larger_0 += tl.sum(
574
- tl.abs(probs_blk - min_larger_0) < 1e-9
575
- )
576
- num_min_larger_1 += tl.sum(
577
- tl.abs(probs_blk - min_larger_1) < 1e-9
548
+ min_larger_0, num_min_larger_0 = (
549
+ _update_min_larger_stats(
550
+ probs_blk,
551
+ above_0,
552
+ min_larger_0,
553
+ num_min_larger_0,
554
+ 1.0,
555
+ )
556
)
557
580
- # Check if any of the pivots satisfy termination condition
581
- if p_pivots_sum_1 >= p and (
582
- p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p
583
- ):
584
- p_pivot = p_pivot_1
585
- min_larger_prob = min_larger_1
586
- num_min_larger = num_min_larger_1
587
- p_pivots_sum = p_pivots_sum_1
588
- found_pivot = 1
558
+ # Check if the pivot satisfies termination condition
559
if p_pivots_sum_0 >= p and (
560
p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p
561
):
@@ -596,19 +566,17 @@ def _topk_topp_kernel(
566
found_pivot = 1
567
568
# Update range
599
- if p_pivots_sum_1 > p:
600
- min_range = p_pivot_1
601
- elif p_pivots_sum_0 > p:
569
+ if p_pivots_sum_0 > p:
570
min_range = p_pivot_0
603
-
604
- if p_pivots_sum_0 < p:
571
+ elif p_pivots_sum_0 < p:
572
max_range = p_pivot_0
606
- elif p_pivots_sum_1 < p:
607
- max_range = p_pivot_1
573
574
num_iters += 1
575
if (max_range - min_range) < 1e-9 or num_iters >= 18:
576
p_pivot = (max_range + min_range) / 2.0
577
+ min_larger_prob = min_larger_0
578
+ num_min_larger = num_min_larger_0
579
+ p_pivots_sum = p_pivots_sum_0
580
found_pivot = 1
581
582
duplicate_logit = (
@@ -725,17 +693,15 @@ def _topk_topp_kernel(
693
694
found_pivot = 0
695
while found_pivot == 0:
728
- p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
696
+ p_pivot_0 = (max_range - min_range) * 0.5 + min_range
697
p_pivots_sum_0 = 0.0
698
min_larger_0 = 1.0
699
num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
700
733
- p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
734
- p_pivots_sum_1 = 0.0
735
- min_larger_1 = 1.0
736
- num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
737
-
738
- # First pass: Calculate p_pivots_sum and min_larger
701
+ # Single fused pass: compute p_pivots_sum,
702
+ # min_larger, and num_min_larger together.
703
+ # See _update_min_larger_stats for the
704
+ # tile-level merge logic.
705
for i in range(0, search_iters):
706
offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
707
0, BLOCK_SIZE_TRUNC
@@ -745,53 +711,18 @@ def _topk_topp_kernel(
711
BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
712
)
713
748
- p_pivots_sum_0 += tl.sum(
749
- probs_blk * (probs_blk > p_pivot_0)
750
- )
751
- masked_larger_0 = tl.where(
752
- probs_blk > p_pivot_0, probs_blk, 1.0
753
- )
754
- min_larger_0 = tl.minimum(
755
- min_larger_0, tl.min(masked_larger_0)
756
- )
714
+ above_0 = probs_blk > p_pivot_0
715
+ p_pivots_sum_0 += tl.sum(probs_blk * above_0)
716
758
- p_pivots_sum_1 += tl.sum(
759
- probs_blk * (probs_blk > p_pivot_1)
760
- )
761
- masked_larger_1 = tl.where(
762
- probs_blk > p_pivot_1, probs_blk, 1.0
763
- )
764
- min_larger_1 = tl.minimum(
765
- min_larger_1, tl.min(masked_larger_1)
766
- )
767
-
768
- # Second pass: Calculate num_min_larger
769
- for i in range(0, search_iters):
770
- offs_n = i * BLOCK_SIZE_TRUNC + tl.arange(
771
- 0, BLOCK_SIZE_TRUNC
772
- )
773
- mask_n_2 = offs_n < search_range
774
- probs_blk = tl.load(
775
- BUFFER_ROW + offs_n, mask=mask_n_2, other=0.0
776
- )
777
-
778
- num_min_larger_0 += tl.sum(
779
- tl.abs(probs_blk - min_larger_0) < 1e-9
780
- )
781
- num_min_larger_1 += tl.sum(
782
- tl.abs(probs_blk - min_larger_1) < 1e-9
717
+ min_larger_0, num_min_larger_0 = _update_min_larger_stats(
718
+ probs_blk,
719
+ above_0,
720
+ min_larger_0,
721
+ num_min_larger_0,
722
+ 1.0,
723
)
724
785
- # Check if any of the pivots satisfy termination condition
786
- if (
787
- p_pivots_sum_1 >= p
788
- and p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p
789
- ):
790
- p_pivot = p_pivot_1
791
- min_larger_prob = min_larger_1
792
- num_min_larger = num_min_larger_1
793
- p_pivots_sum = p_pivots_sum_1
794
- found_pivot = 1
725
+ # Check if the pivot satisfies termination condition
726
if (
727
p_pivots_sum_0 >= p
728
and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p
@@ -803,19 +734,17 @@ def _topk_topp_kernel(
734
found_pivot = 1
735
736
# Update range
806
- if p_pivots_sum_1 > p:
807
- min_range = p_pivot_1
808
- elif p_pivots_sum_0 > p:
737
+ if p_pivots_sum_0 > p:
738
min_range = p_pivot_0
810
-
811
- if p_pivots_sum_0 < p:
739
+ elif p_pivots_sum_0 < p:
740
max_range = p_pivot_0
813
- elif p_pivots_sum_1 < p:
814
- max_range = p_pivot_1
741
742
num_iters += 1
743
if (max_range - min_range) < 1e-9 or num_iters >= 18:
744
p_pivot = (max_range + min_range) / 2.0
745
+ min_larger_prob = min_larger_0
746
+ num_min_larger = num_min_larger_0
747
+ p_pivots_sum = p_pivots_sum_0
748
found_pivot = 1
749
else:
750
# Re-populate the buffer with full softmax probabilities
@@ -832,17 +761,15 @@ def _topk_topp_kernel(
761
762
found_pivot = 0
763
while found_pivot == 0:
835
- p_pivot_0 = (max_range - min_range) * 1.0 / 3.0 + min_range
764
+ p_pivot_0 = (max_range - min_range) * 0.5 + min_range
765
p_pivots_sum_0 = 0.0
766
min_larger_0 = 1.0
767
num_min_larger_0 = tl.zeros((), dtype=tl.uint32)
768
840
- p_pivot_1 = (max_range - min_range) * 2.0 / 3.0 + min_range
841
- p_pivots_sum_1 = 0.0
842
- min_larger_1 = 1.0
843
- num_min_larger_1 = tl.zeros((), dtype=tl.uint32)
844
-
845
- # First pass: Calculate p_pivots_sum and min_larger
769
+ # Single fused pass: compute p_pivots_sum,
770
+ # min_larger, and num_min_larger together.
771
+ # See _update_min_larger_stats for the
772
+ # tile-level merge logic.
773
for i in range(0, NUM_TILES):
774
offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
775
mask_n = offs_n < VOCAB_SIZE
@@ -850,51 +777,18 @@ def _topk_topp_kernel(
777
BUFFER_ROW + offs_n, mask=mask_n, other=0.0
778
)
779
853
- p_pivots_sum_0 += tl.sum(
854
- probs_blk * (probs_blk > p_pivot_0)
855
- )
856
- masked_larger_0 = tl.where(
857
- probs_blk > p_pivot_0, probs_blk, 1.0
858
- )
859
- min_larger_0 = tl.minimum(
860
- min_larger_0, tl.min(masked_larger_0)
861
- )
862
-
863
- p_pivots_sum_1 += tl.sum(
864
- probs_blk * (probs_blk > p_pivot_1)
865
- )
866
- masked_larger_1 = tl.where(
867
- probs_blk > p_pivot_1, probs_blk, 1.0
868
- )
869
- min_larger_1 = tl.minimum(
870
- min_larger_1, tl.min(masked_larger_1)
871
- )
872
-
873
- # Second pass: Calculate num_min_larger
874
- for i in range(0, NUM_TILES):
875
- offs_n = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
876
- mask_n = offs_n < VOCAB_SIZE
877
- probs_blk = tl.load(
878
- BUFFER_ROW + offs_n, mask=mask_n, other=0.0
879
- )
780
+ above_0 = probs_blk > p_pivot_0
781
+ p_pivots_sum_0 += tl.sum(probs_blk * above_0)
782
881
- num_min_larger_0 += tl.sum(
882
- tl.abs(probs_blk - min_larger_0) < 1e-9
883
- )
884
- num_min_larger_1 += tl.sum(
885
- tl.abs(probs_blk - min_larger_1) < 1e-9
783
+ min_larger_0, num_min_larger_0 = _update_min_larger_stats(
784
+ probs_blk,
785
+ above_0,
786
+ min_larger_0,
787
+ num_min_larger_0,
788
+ 1.0,
789
)
790
888
- # Check if any of the pivots satisfy termination condition
889
- if (
890
- p_pivots_sum_1 >= p
891
- and p_pivots_sum_1 - (min_larger_1 * num_min_larger_1) < p
892
- ):
893
- p_pivot = p_pivot_1
894
- min_larger_prob = min_larger_1
895
- num_min_larger = num_min_larger_1
896
- p_pivots_sum = p_pivots_sum_1
897
- found_pivot = 1
791
+ # Check if the pivot satisfies termination condition
792
if (
793
p_pivots_sum_0 >= p
794
and p_pivots_sum_0 - (min_larger_0 * num_min_larger_0) < p
@@ -906,22 +800,20 @@ def _topk_topp_kernel(
800
found_pivot = 1
801
802
# Update range
909
- if p_pivots_sum_1 > p:
910
- min_range = p_pivot_1
911
- elif p_pivots_sum_0 > p:
803
+ if p_pivots_sum_0 > p:
804
min_range = p_pivot_0
913
-
914
- if p_pivots_sum_0 < p:
805
+ elif p_pivots_sum_0 < p:
806
max_range = p_pivot_0
916
- elif p_pivots_sum_1 < p:
917
- max_range = p_pivot_1
807
808
num_iters += 1
809
if (max_range - min_range) < 1e-9 or num_iters >= 18:
810
p_pivot = (max_range + min_range) / 2.0
811
+ min_larger_prob = min_larger_0
812
+ num_min_larger = num_min_larger_0
813
+ p_pivots_sum = p_pivots_sum_0
814
found_pivot = 1
815
924
- duplicate_logit = tl.log(min_larger_prob * sum_exp_logits) + max_logit
816
+ duplicate_logit = tl.log(min_larger_prob * sum_exp_logits) + max_sample
817
num_duplicate_logit = num_min_larger
818
num_keep = num_duplicate_logit - tl.cast(
819
(p_pivots_sum - p) / min_larger_prob, tl.uint32
@@ -952,9 +844,7 @@ def _topk_topp_kernel(
844
tl.abs(logits_blk - duplicate_logit) < 1e-9
845
) & mask_n
846
duplicate_count = tl.cumsum(duplicate_mask) + num_kept
955
- duplicate_keep_mask = (
956
- duplicate_count <= num_duplicate_logit
957
- ) & duplicate_mask
847
+ duplicate_keep_mask = (duplicate_count <= num_keep) & duplicate_mask
848
duplicate_remove_mask = duplicate_mask & ~duplicate_keep_mask
849
num_kept += tl.sum(duplicate_keep_mask)
850
keep_mask = keep_mask & (~duplicate_remove_mask)