[CPU] Skip Triton kernel monkey-patches when Triton-CPU is available (#44991)
Signed-off-by: jmamou <jonathan.mamou@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Jonathan Mamou committed
Jun 18, 2026 at 13:59 UTC
351c72d6e5d43148f16d67b11a613d14dafbf6a4
3 files changed
+26
-6
vllm/v1/sample/rejection_sampler.py
+17
-5
@@ -732,7 +732,11 @@ def rejection_greedy_sample_kernel(
732
# Early exit for non-greedy sampling requests.
733
return
734
735
- start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
735
+ start_idx = (
736
+ tl.zeros([], dtype=cu_num_draft_tokens_ptr.dtype.element_ty)
737
+ if req_idx == 0
738
+ else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
739
+ )
740
end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
741
num_draft_tokens = end_idx - start_idx
742
@@ -788,7 +792,11 @@ def rejection_random_sample_kernel(
792
# Early exit for greedy sampling requests.
793
return
794
791
- start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
795
+ start_idx = (
796
+ tl.zeros([], dtype=cu_num_draft_tokens_ptr.dtype.element_ty)
797
+ if req_idx == 0
798
+ else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
799
+ )
800
end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
801
num_draft_tokens = end_idx - start_idx
802
@@ -844,8 +852,8 @@ def expand_kernel(
852
MAX_NUM_TOKENS: tl.constexpr,
853
):
854
req_idx = tl.program_id(0)
847
- if req_idx == 0: # noqa: SIM108
848
- start_idx = 0
855
+ if req_idx == 0:
856
+ start_idx = tl.zeros([], dtype=cu_num_tokens_ptr.dtype.element_ty)
857
else:
858
start_idx = tl.load(cu_num_tokens_ptr + req_idx - 1)
859
end_idx = tl.load(cu_num_tokens_ptr + req_idx)
@@ -871,7 +879,11 @@ def sample_recovered_tokens_kernel(
879
USE_FP64_GUMBEL: tl.constexpr,
880
):
881
req_idx = tl.program_id(0)
874
- start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
882
+ start_idx = (
883
+ tl.zeros([], dtype=cu_num_draft_tokens_ptr.dtype.element_ty)
884
+ if req_idx == 0
885
+ else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
886
+ )
887
end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
888
num_draft_tokens = end_idx - start_idx
889
vllm/v1/spec_decode/utils.py
-1
@@ -156,7 +156,6 @@ def eagle_prepare_inputs_padded_kernel(
156
# cumulative sum (first entry is the first value, not zero).
157
cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx)
158
159
- num_draft_tokens = 0
159
if req_idx == 0:
160
num_draft_tokens = cu_draft_curr
161
else:
vllm/v1/worker/cpu_model_runner.py
+9
@@ -60,6 +60,15 @@ class CPUModelRunner(GPUModelRunner):
60
v.gpu = v.cpu
61
62
def _postprocess_triton(self) -> None:
63
+ from vllm.triton_utils import HAS_TRITON
64
+
65
+ if HAS_TRITON:
66
+ logger.info(
67
+ "Triton-CPU backend is available; skipping C++ monkey-patches "
68
+ "for Triton kernels."
69
+ )
70
+ return
71
+
72
import vllm.v1.worker.block_table
73
74
vllm.v1.worker.block_table._compute_slot_mapping_kernel = (