From 405bc86768d1267f05dd2d4139048eb30a84b6cb Mon Sep 17 00:00:00 2001 From: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:55:35 +0800 Subject: [PATCH] [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 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Nick Hill Co-authored-by: Claude Opus 5 (1M context) --- vllm/v1/sample/ops/topk_topp_triton.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index c284ff6187..7d593e69cd 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -931,12 +931,18 @@ def apply_top_k_top_p_triton( # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles. # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot # approximation; use smaller tiles for accurate top-p results. + launch_kwargs = {} if logits.device.type == "cpu": block_size, block_size_trunc = 256, 128 elif logits.device.type == "xpu": block_size, block_size_trunc = 4096, 2048 else: block_size, block_size_trunc = 8192, 4096 + # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so + # per-tile latency bounds kernel latency, and Triton's default of 4 + # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is + # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not. + launch_kwargs["num_warps"] = 8 _topk_topp_kernel[(NUM_PROGRAMS,)]( logits, @@ -953,6 +959,7 @@ def apply_top_k_top_p_triton( BLOCK_SIZE_TRUNC=block_size_trunc, TOPK_ENABLED=topk_enabled, TOPP_ENABLED=topp_enabled, + **launch_kwargs, ) return logits