[Perf] Reduce redundant Triton sampler warmup specializations (#58605)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Misha Goin committed
Sep 29, 2026 at 17:23 UTC
1ee7f78e806b5e9b7476c74edf453b58fe108157
1 file changed
+2
-2
vllm/v1/sample/ops/topk_topp_triton.py
+2
-2
@@ -103,7 +103,7 @@ def _update_min_larger_stats(data, above_mask, min_larger, num_min_larger, senti
103
return min_larger, num_min_larger
104
105
106
-@triton.jit
106
+@triton.jit(do_not_specialize_on_alignment=["BATCH_SIZE"])
107
def _topk_topp_kernel(
108
LOGITS,
109
LOGITS_STRIDE_0,
@@ -1225,7 +1225,7 @@ def _topp_sb_combine(
1225
)
1226
1227
1228
-@triton.jit
1228
+@triton.jit(do_not_specialize_on_alignment=["ROUND"])
1229
def _topp_sb_step_kernel(
1230
LOGITS,
1231
LOGITS_STRIDE_0,