[Perf][Spec Decode] Avoid triton recompiles in the acceptance estimator (#57107)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Giancarlo Delfin committed
Sep 28, 2026 at 13:46 UTC
3d5f4d4cd5d86788151664597bcebea666c34700
1 file changed
+5
-4
vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py
+5
-4
@@ -20,7 +20,7 @@ _INIT_SLOPE = 0.5
20
_INIT_BIAS = -0.2
21
22
23
-@triton.jit
23
+@triton.jit(do_not_specialize=["num_reqs"])
24
def _accumulate_kernel(
25
info_ptr,
26
grad_ptr,
@@ -160,7 +160,7 @@ def _refit_kernel(
160
tl.store(counts_ptr + k, 0.0, mask=mask)
161
162
163
-@triton.jit
163
+@triton.jit(do_not_specialize=["num_tokens"])
164
def _local_max_sumexp_kernel(
165
local_max_ptr,
166
local_max_stride,
@@ -216,7 +216,7 @@ def _local_max_sumexp_kernel(
216
)
217
218
219
-@triton.jit
219
+@triton.jit(do_not_specialize=["num_tokens"])
220
def _predict_kernel(
221
features_ptr,
222
features_stride,
@@ -422,7 +422,8 @@ class OnlineAcceptanceEstimator:
422
self.predictions.stride(0),
423
self.counts,
424
num_reqs,
425
- BLOCK_R=triton.next_power_of_2(max(num_reqs, 1)),
425
+ BLOCK_R=max(256, triton.next_power_of_2(num_reqs)),
426
+ num_warps=1 if num_reqs <= 256 else 4,
427
)
428
429
self._steps_since_refit += 1