[Perf] Launch the top-k/top-p Triton sampler kernel with 8 warps (#51507)
Signed-off-by: BabyDrangoner <148877251+BabyDrangoner@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
kiroxu committed
Aug 11, 2026 at 01:55 UTC
405bc86768d1267f05dd2d4139048eb30a84b6cb
1 file changed
+7
vllm/v1/sample/ops/topk_topp_triton.py
+7
@@ -931,12 +931,18 @@ def apply_top_k_top_p_triton(
931
# Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles.
932
# On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot
933
# approximation; use smaller tiles for accurate top-p results.
934
+ launch_kwargs = {}
935
if logits.device.type == "cpu":
936
block_size, block_size_trunc = 256, 128
937
elif logits.device.type == "xpu":
938
block_size, block_size_trunc = 4096, 2048
939
else:
940
block_size, block_size_trunc = 8192, 4096
941
+ # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so
942
+ # per-tile latency bounds kernel latency, and Triton's default of 4
943
+ # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is
944
+ # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not.
945
+ launch_kwargs["num_warps"] = 8
946
947
_topk_topp_kernel[(NUM_PROGRAMS,)](
948
logits,
@@ -953,6 +959,7 @@ def apply_top_k_top_p_triton(
959
BLOCK_SIZE_TRUNC=block_size_trunc,
960
TOPK_ENABLED=topk_enabled,
961
TOPP_ENABLED=topp_enabled,
962
+ **launch_kwargs,
963
)
964
965
return logits