@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 44257c2e35

[Feature] Triton kernel dispatcher (#43048)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

wangxiyuan committed Sep 25, 2026 at 12:07 UTC 44257c2e358c4dbd648f99273a21a3f496f23148
8 files changed +482 -98
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+12 -1
@@ -1071,7 +1071,18 @@ def test_batch_memcpy_cpu_fallback() -> None:
1071 copy each src into its dst, validating the (src_ptrs, dst_ptrs, sizes)
1072 argument order against ctypes.memmove(dst, src, size).
1073 """
1074 - from vllm.utils.cpu_triton_utils import batch_memcpy_kernel
1074 + from vllm.triton_utils.dispatcher import (
1075 + KernelOverride,
1076 + register_kernels,
1077 + )
1078 + from vllm.utils.cpu_triton_utils import _batch_memcpy_impl
1079 +
1080 + register_kernels(
1081 + {"vllm.v1.worker.mamba_utils.batch_memcpy_kernel": _batch_memcpy_impl}
1082 + )
1083 + from vllm.v1.worker.mamba_utils import batch_memcpy_kernel
1084 +
1085 + assert isinstance(batch_memcpy_kernel, KernelOverride)
1086
1087 # Varied byte sizes, including a non-power-of-two run.
1088 sizes_bytes = [256, 1024, 17 * 4, 4096]
tests/model_executor/test_triton_dispatcher.py new
+222
@@ -0,0 +1,222 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the apply-time Triton kernel override registry.
4 +
5 +The registry must let out-of-tree platforms replace kernels defined as
6 +plain ``@triton.jit`` functions, patching every hold-site (defining
7 +module, from-import copies, and JIT warmup owners) so that existing call
8 +sites dispatch to the override without any core-side decoration.
9 +"""
10 +
11 +import sys
12 +import types
13 +from typing import Any
14 +
15 +import pytest
16 +
17 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
18 + _DecoratedTritonJitKernel,
19 +)
20 +from vllm.triton_utils import tl, triton
21 +from vllm.triton_utils.dispatcher import (
22 + KernelOverride,
23 + register_kernels,
24 +)
25 +
26 +
27 +class _FakeTritonKernel:
28 + """Stands in for a Triton JITFunction (works without Triton installed)."""
29 +
30 + arg_names = ("output_ptr", "input_ptr", "replace_from", "MAX_NUM_TOKENS")
31 +
32 + def __getitem__(self, grid: Any) -> Any:
33 + def launch(*args: Any, **kwargs: Any) -> Any:
34 + return ("default", grid, kwargs)
35 +
36 + return launch
37 +
38 +
39 +@pytest.fixture
40 +def _fake_module(monkeypatch: pytest.MonkeyPatch) -> Any:
41 + """Expose a fake defining module holding the kernel and its warmup owner."""
42 + import types
43 +
44 + kernel = _FakeTritonKernel()
45 + owner = _DecoratedTritonJitKernel(
46 + kernel,
47 + lambda: {},
48 + lambda output, input, replace_from: ( # noqa: ARG005
49 + (1,),
50 + dict(MAX_NUM_TOKENS=8),
51 + ),
52 + )
53 + module = types.ModuleType("vllm.tests.fake_kernel_home")
54 + module.expand_kernel = kernel # type: ignore[attr-defined]
55 + module._expand = owner # type: ignore[attr-defined]
56 + monkeypatch.setitem(sys.modules, module.__name__, module)
57 + return module
58 +
59 +
60 +def test_override_patches_defining_module_and_warmup_owner(
61 + _fake_module: Any,
62 +) -> None:
63 + calls: list[tuple[Any, ...]] = []
64 +
65 + def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None):
66 + calls.append((output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS))
67 + return "cpu"
68 +
69 + register_kernels(
70 + {
71 + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
72 + }
73 + )
74 +
75 + overridden = _fake_module.expand_kernel
76 + assert isinstance(overridden, KernelOverride)
77 + # Launch with kernel-arg names like the warmup launch path does.
78 + assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu"
79 + assert calls == [(1, 2, 0, None)]
80 +
81 + # The warmup owner's captured kernel is patched too, and its cached
82 + # argument-name introspection is invalidated.
83 + assert _fake_module._expand.kernel is overridden
84 +
85 +
86 +def test_override_forwards_positionally_for_mismatched_names(
87 + _fake_module: Any,
88 +) -> None:
89 + calls: list[tuple[Any, ...]] = []
90 +
91 + def cpu_impl(output, input_val, replace_from, MAX_NUM_TOKENS=None):
92 + calls.append((output, input_val, replace_from, MAX_NUM_TOKENS))
93 + return "cpu"
94 +
95 + register_kernels(
96 + {
97 + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
98 + }
99 + )
100 + overridden = _fake_module.expand_kernel
101 + assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu"
102 + assert calls == [(1, 2, 0, None)]
103 +
104 +
105 +def test_override_covers_from_import_copies(_fake_module: Any) -> None:
106 + import types
107 +
108 + importer = types.ModuleType("vllm.tests.fake_kernel_importer")
109 + importer.expand_kernel = _fake_module.expand_kernel # type: ignore[attr-defined]
110 + sys.modules["vllm.tests.fake_kernel_importer"] = importer
111 +
112 + def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None):
113 + return "cpu"
114 +
115 + register_kernels(
116 + {
117 + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
118 + }
119 + )
120 + assert isinstance(importer.expand_kernel, KernelOverride)
121 + del sys.modules["vllm.tests.fake_kernel_importer"]
122 +
123 +
124 +def test_override_owner_launch_goes_through_impl(_fake_module: Any) -> None:
125 + launched: list[Any] = []
126 +
127 + def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None):
128 + launched.append((output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS))
129 + return "cpu"
130 +
131 + register_kernels(
132 + {
133 + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
134 + }
135 + )
136 + owner = _fake_module._expand
137 + # In-place kernels return None from launch; assert on the recorded call.
138 + owner(1, 2, 0)
139 + assert launched == [(1, 2, 0, 8)]
140 +
141 +
142 +def test_register_kernel_auto_imports_defining_module(
143 + monkeypatch: pytest.MonkeyPatch,
144 +) -> None:
145 + """Registering a kernel whose defining module is not imported yet must
146 + import it on the spot and apply the override."""
147 + kernel_name = "vllm.v1.sample.rejection_sampler.expand_kernel"
148 + module_name = kernel_name.rpartition(".")[0]
149 + calls: list[tuple[Any, ...]] = []
150 +
151 + def cpu_impl(
152 + output_ptr,
153 + input_ptr,
154 + cu_num_tokens_ptr,
155 + replace_from,
156 + replace_to,
157 + MAX_NUM_TOKENS=None,
158 + ):
159 + calls.append((output_ptr, replace_from, MAX_NUM_TOKENS))
160 + return "cpu"
161 +
162 + # Ensure the module is NOT imported before registration. Drop it from
163 + # sys.modules; a fresh import creates new kernel/owner objects, which
164 + # is what a cold-start platform process would see.
165 + saved = sys.modules.pop(module_name, None)
166 + # Re-run the registry patch against the original module afterwards so
167 + # other tests that already imported it keep working.
168 + try:
169 + assert module_name not in sys.modules
170 + register_kernels({kernel_name: cpu_impl})
171 + assert module_name in sys.modules
172 + finally:
173 + if saved is None:
174 + # First import in the whole session: keep the fresh module.
175 + saved = sys.modules[module_name]
176 + else:
177 + # Restore the original module object and re-apply the override
178 + # to its (pre-existing) kernel so both generations agree.
179 + sys.modules[module_name] = saved
180 + register_kernels({kernel_name: cpu_impl})
181 +
182 + overridden = saved.expand_kernel
183 + assert isinstance(overridden, KernelOverride)
184 + overridden[(2,)](
185 + output_ptr=1,
186 + input_ptr=2,
187 + cu_num_tokens_ptr=3,
188 + replace_from=0,
189 + replace_to=0,
190 + MAX_NUM_TOKENS=8,
191 + )
192 + assert calls == [(1, 0, 8)]
193 + # The warmup owner captured at the fresh import must be patched too
194 + # (when the module existed before, there are two module generations).
195 + fresh = sys.modules[module_name]
196 + if fresh is not saved:
197 + assert fresh._expand.kernel is overridden
198 +
199 +
200 +def test_real_triton_kernel_override_end_to_end() -> None:
201 + """Override a real @triton.jit kernel through the registry."""
202 + if not triton.__class__.__module__.startswith("triton"):
203 + pytest.skip("Triton is not installed")
204 +
205 + @triton.jit
206 + def _add_one_kernel(x_ptr, out_ptr, BLOCK: tl.constexpr): # pragma: no cover
207 + offsets = tl.arange(0, BLOCK)
208 + tl.store(out_ptr + offsets, tl.load(x_ptr + offsets) + 1)
209 +
210 + module = types.ModuleType("vllm.tests.real_jit_kernel_home")
211 + module._add_one_kernel = _add_one_kernel # type: ignore[attr-defined]
212 + sys.modules[module.__name__] = module
213 +
214 + def cpu_impl(x_ptr, out_ptr, BLOCK=None): # noqa: ARG001
215 + return "cpu"
216 +
217 + register_kernels({"vllm.tests.real_jit_kernel_home._add_one_kernel": cpu_impl})
218 + assert module._add_one_kernel.arg_names == ( # type: ignore[attr-defined]
219 + "x_ptr",
220 + "out_ptr",
221 + "BLOCK",
222 + )
vllm/platforms/cpu.py
+54
@@ -647,6 +647,60 @@ class CpuPlatform(Platform):
647 except ImportError as e:
648 logger.warning_once("Failed to import from vllm._C: %s", repr(e))
649
650 + @classmethod
651 + def register_triton_kernel_overrides(cls) -> None:
652 + """Fallback C++ implementations for core Triton kernels.
653 +
654 + Used when the Triton-CPU backend is unavailable. When Triton-CPU is
655 + present, the Triton kernels can run natively and the overrides
656 + would only shadow them, so registration is skipped.
657 + """
658 + from vllm.triton_utils import HAS_TRITON
659 +
660 + if HAS_TRITON:
661 + return
662 +
663 + from vllm.triton_utils.dispatcher import register_kernels
664 + from vllm.utils import cpu_triton_utils as cpu_tl
665 +
666 + register_kernels(
667 + {
668 + "vllm.v1.worker.block_table.ComputeSlotMappingKernel.kernel": (
669 + cpu_tl._compute_slot_mapping_kernel_impl
670 + ),
671 + "vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel": (
672 + cpu_tl._eagle_step_slot_mapping_metadata_kernel_impl
673 + ),
674 + "vllm.v1.spec_decode.utils.eagle_prepare_inputs_padded_kernel": (
675 + cpu_tl._eagle_prepare_inputs_padded_kernel_impl
676 + ),
677 + "vllm.v1.spec_decode.utils.eagle_prepare_next_token_padded_kernel": (
678 + cpu_tl._eagle_prepare_next_token_padded_kernel_impl
679 + ),
680 + "vllm.v1.spec_decode.utils.copy_and_expand_eagle_inputs_kernel": (
681 + cpu_tl._copy_and_expand_eagle_inputs_kernel_impl
682 + ),
683 + "vllm.v1.spec_decode.utils.copy_and_expand_dflash_inputs_kernel": (
684 + cpu_tl._copy_and_expand_dflash_inputs_kernel_impl
685 + ),
686 + "vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel": (
687 + cpu_tl._rejection_greedy_sample_kernel_impl
688 + ),
689 + "vllm.v1.sample.rejection_sampler.rejection_random_sample_kernel": (
690 + cpu_tl._rejection_random_sample_kernel_impl
691 + ),
692 + "vllm.v1.sample.rejection_sampler.expand_kernel": (
693 + cpu_tl._expand_kernel_impl
694 + ),
695 + "vllm.v1.sample.rejection_sampler.sample_recovered_tokens_kernel": (
696 + cpu_tl._sample_recovered_tokens_kernel_impl
697 + ),
698 + "vllm.v1.worker.mamba_utils.batch_memcpy_kernel": (
699 + cpu_tl._batch_memcpy_impl
700 + ),
701 + }
702 + )
703 +
704 @classmethod
705 def pack_kv_cache(
706 cls,
vllm/platforms/interface.py
+13
@@ -366,6 +366,19 @@ class Platform:
366 with contextlib.suppress(ImportError):
367 import vllm._moe_C_stable_libtorch # noqa: F401
368
369 + @classmethod
370 + def register_triton_kernel_overrides(cls) -> None:
371 + """Override core Triton kernels with platform implementations.
372 +
373 + Platforms whose Triton kernels need replacement implementations
374 + (e.g. CPU fallbacks when the Triton-CPU backend is unavailable)
375 + override this method and call
376 + :func:`vllm.triton_utils.dispatcher.register_kernels`
377 + with their overrides. It runs during model runner init, before any
378 + kernel launch. The default implementation does nothing.
379 + """
380 + return
381 +
382 @classmethod
383 def get_attn_backend_cls(
384 cls,
vllm/triton_utils/dispatcher.py new
+175
@@ -0,0 +1,175 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Registry for overriding Triton kernels with platform implementations.
4 +
5 +Non-CUDA platforms can replace Triton kernels defined in vLLM core
6 +without modifying core code:
7 +
8 + from vllm.triton_utils.dispatcher import register_kernels
9 +
10 + register_kernels({
11 + "vllm.v1.sample.rejection_sampler.expand_kernel": my_expand_impl,
12 + "vllm.v1.worker.mamba_utils.batch_memcpy_kernel": my_memcpy_impl,
13 + })
14 +"""
15 +
16 +import importlib
17 +import inspect
18 +import sys
19 +from collections.abc import Callable, Mapping
20 +from typing import Any
21 +
22 +from vllm.logger import init_logger
23 +
24 +logger = init_logger(__name__)
25 +
26 +_registry: dict[str, tuple[Callable[..., Any], Callable[..., Any]]] = {}
27 +
28 +
29 +def _kernel_arg_names(kernel: Callable[..., Any]) -> tuple[str, ...]:
30 + """Return the argument names of a Triton kernel or its fallback."""
31 + names = getattr(kernel, "arg_names", None)
32 + if names is None:
33 + # No real Triton (TritonPlaceholder): read the plain function.
34 + fn = getattr(kernel, "func", kernel)
35 + names = tuple(inspect.signature(fn).parameters)
36 + return tuple(names)
37 +
38 +
39 +class KernelOverride:
40 + """Stands in for a Triton kernel and routes launches to ``impl``.
41 +
42 + The JIT warmup infrastructure launches kernels as
43 + ``kernel[grid](**kwargs)`` where ``kwargs`` are keyed by the original
44 + kernel's argument names. This wrapper forwards those launch arguments
45 + to the platform implementation: by keyword when the implementation's
46 + parameter names match the kernel's, otherwise positionally in the
47 + kernel's parameter order.
48 + """
49 +
50 + def __init__(self, kernel: Callable[..., Any], impl: Callable[..., Any]) -> None:
51 + self._impl = impl
52 + # Mirror the original kernel's names so launch binding and warmup
53 + # introspection keep working against this wrapper.
54 + self.arg_names = _kernel_arg_names(kernel)
55 + self.constexprs = getattr(kernel, "constexprs", None)
56 + self.func = impl
57 + self.__name__ = getattr(impl, "__name__", "kernel_override")
58 + self.__module__ = getattr(kernel, "__module__", "")
59 + self._forward_by_name = (
60 + tuple(inspect.signature(impl).parameters) == self.arg_names
61 + )
62 +
63 + def __getitem__(self, grid: Any) -> Callable[..., Any]:
64 + return self._launch
65 +
66 + def _launch(self, *args: Any, **kwargs: Any) -> Any:
67 + if args:
68 + return self._impl(*args, **kwargs)
69 + if self._forward_by_name:
70 + return self._impl(**kwargs)
71 + unexpected = [name for name in kwargs if name not in self.arg_names]
72 + if unexpected:
73 + raise RuntimeError(
74 + f"Override for {self.__name__!r} received unexpected launch "
75 + f"arguments: {unexpected}"
76 + )
77 + return self._impl(*(kwargs[name] for name in self.arg_names if name in kwargs))
78 +
79 + def __call__(self, *args: Any, **kwargs: Any) -> Any:
80 + return self._impl(*args, **kwargs)
81 +
82 +
83 +def _resolve_kernel(name: str) -> tuple[Any, str]:
84 + """Return the (host object, attribute) holding the kernel `name`."""
85 + module_name, attr = name.rsplit(".", 1)
86 + parts = module_name.split(".")
87 + if parts[-1][:1].isupper():
88 + class_module = importlib.import_module(".".join(parts[:-1]))
89 + return getattr(class_module, parts[-1]), attr
90 + importlib.import_module(module_name)
91 + module = sys.modules[module_name]
92 + if hasattr(module, attr):
93 + return module, attr
94 + raise ValueError(f"Kernel {name!r} not found for override")
95 +
96 +
97 +def _rebind_kernels(
98 + overrides: list[tuple[Callable[..., Any], KernelOverride]],
99 +) -> None:
100 + """Point every reference to each original kernel at its wrapper.
101 +
102 + A kernel object can be captured in two kinds of places besides its
103 + defining module, and both must be patched for the overrides to take
104 + effect everywhere:
105 +
106 + * plain attributes of other modules (``from mod import kernel``
107 + copies), rebound to the wrapper;
108 + * JIT warmup owners, whose ``kernel`` attribute is the kernel they
109 + launch; swapped, and their cached ``_kernel_arg_names`` invalidated
110 + so the launch binding re-derives from the wrapper.
111 +
112 + One pass over ``sys.modules`` handles all kernels in `overrides`.
113 + """
114 + originals = {original: wrapper for original, wrapper in overrides}
115 +
116 + def lookup(value: Any) -> KernelOverride | None:
117 + # Only compare by identity: this scan runs over every module
118 + # namespace, and some attribute values (e.g. PlaceholderModule
119 + # sentinels) trigger imports or errors on hash/eq.
120 + for original in originals:
121 + if value is original:
122 + return originals[original]
123 + return None
124 +
125 + for module in list(sys.modules.values()):
126 + namespace = getattr(module, "__dict__", None)
127 + if not namespace:
128 + continue
129 + for attr, value in list(namespace.items()):
130 + wrapper = lookup(value)
131 + if wrapper is not None:
132 + setattr(module, attr, wrapper)
133 + continue
134 + # JIT warmup owners (e.g. _DecoratedTritonJitKernel) hold the
135 + # kernel as an instance attribute; read the instance dict directly
136 + # so module-level __getattr__ hooks are never triggered, and so
137 + # this scan needs no dependency on the warmup machinery.
138 + if hasattr(value, "__dict__"):
139 + wrapper = lookup(vars(value).get("kernel"))
140 + if wrapper is None:
141 + continue
142 + value.kernel = wrapper
143 + if not isinstance(value, type):
144 + value.__dict__.pop("_kernel_arg_names", None)
145 +
146 +
147 +def register_kernels(overrides: Mapping[str, Callable[..., Any]]) -> None:
148 + """Register platform-specific implementations for Triton kernels.
149 +
150 + Args:
151 + overrides: Mapping from fully qualified kernel name to the
152 + platform implementation. Kernel names look like
153 + "vllm.v1.sample.rejection_sampler.expand_kernel"; for kernels
154 + owned by a JIT warmup class, include the attribute path, e.g.
155 + "vllm.v1.worker.block_table.ComputeSlotMappingKernel.kernel".
156 + Implementations are invoked like the kernels themselves, with
157 + the launch arguments in the original kernel's parameter order
158 + (or by keyword when their parameter names match the kernel's).
159 +
160 + """
161 + resolved: list[tuple[str, Any, str, Callable[..., Any], KernelOverride]] = []
162 + for name, impl in overrides.items():
163 + if name in _registry:
164 + logger.warning("Kernel %s is already registered. Overwriting.", name)
165 + host, attr = _resolve_kernel(name)
166 + original = getattr(host, attr)
167 + wrapper = KernelOverride(original, impl)
168 + _registry[name] = (original, impl)
169 + resolved.append((name, host, attr, original, wrapper))
170 + # Rebind other holders first: once the defining attributes are swapped
171 + # below, the scan could no longer find the original kernels.
172 + _rebind_kernels([(original, wrapper) for *_, original, wrapper in resolved])
173 + for name, host, attr, _, wrapper in resolved:
174 + setattr(host, attr, wrapper)
175 + logger.debug("Registered kernel override %s", name)
vllm/utils/cpu_triton_utils.py
-36
@@ -4,19 +4,10 @@
4 """Contains replacement functions to fallback Triton usages in CPU backend."""
5
6 import ctypes
7 -from collections.abc import Callable
7
8 import torch
9
10
12 -class _FuncWrapper:
13 - def __init__(self, func: Callable) -> None:
14 - self.func = func
15 -
16 - def __getitem__(self, *args, **kwargs) -> Callable:
17 - return self.func
18 -
19 -
11 # For _compute_slot_mapping_kernel in vllm/v1/worker/block_table.py
12 def _compute_slot_mapping_kernel_impl(
13 num_tokens: int,
@@ -47,9 +38,6 @@ def _compute_slot_mapping_kernel_impl(
38 )
39
40
50 -compute_slot_mapping_kernel = _FuncWrapper(_compute_slot_mapping_kernel_impl)
51 -
52 -
41 def _ensure_int64(t: torch.Tensor) -> torch.Tensor:
42 return t if t.dtype == torch.int64 else t.to(torch.int64)
43
@@ -454,31 +442,7 @@ def _sample_recovered_tokens_kernel_impl(
442 output_token_ids.copy_(output_i64.to(orig_dtype))
443
444
457 -eagle_prepare_inputs_padded_kernel = _FuncWrapper(
458 - _eagle_prepare_inputs_padded_kernel_impl
459 -)
460 -eagle_prepare_next_token_padded_kernel = _FuncWrapper(
461 - _eagle_prepare_next_token_padded_kernel_impl
462 -)
463 -copy_and_expand_eagle_inputs_kernel = _FuncWrapper(
464 - _copy_and_expand_eagle_inputs_kernel_impl
465 -)
466 -copy_and_expand_dflash_inputs_kernel = _FuncWrapper(
467 - _copy_and_expand_dflash_inputs_kernel_impl
468 -)
469 -eagle_step_slot_mapping_metadata_kernel = _FuncWrapper(
470 - _eagle_step_slot_mapping_metadata_kernel_impl
471 -)
472 -rejection_greedy_sample_kernel = _FuncWrapper(_rejection_greedy_sample_kernel_impl)
473 -rejection_random_sample_kernel = _FuncWrapper(_rejection_random_sample_kernel_impl)
474 -expand_kernel = _FuncWrapper(_expand_kernel_impl)
475 -sample_recovered_tokens_kernel = _FuncWrapper(_sample_recovered_tokens_kernel_impl)
476 -
477 -
445 def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None):
446 # BLOCK_SIZE is unused; kept for signature parity with the Triton kernel.
447 for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()):
448 ctypes.memmove(dst, src, size)
482 -
483 -
484 -batch_memcpy_kernel = _FuncWrapper(_batch_memcpy_impl)
vllm/v1/worker/cpu_model_runner.py
-61
@@ -1,13 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import sys
3 from contextlib import AbstractContextManager, contextmanager
4 from typing import Any
5
6 import torch
7 import torch.nn as nn
8
10 -import vllm.utils.cpu_triton_utils as cpu_tl
9 from vllm.config import (
10 CompilationMode,
11 VllmConfig,
@@ -37,7 +35,6 @@ class CPUModelRunner(GPUModelRunner):
35 self.cascade_attn_enabled = False
36
37 self._postprocess_tensors()
40 - self._postprocess_triton()
38
39 def _postprocess_tensors(self) -> None:
40 # Note: replace device tensors with cpu tensors
@@ -62,64 +59,6 @@ class CPUModelRunner(GPUModelRunner):
59 if isinstance(v, CpuGpuBuffer):
60 v.gpu = v.cpu
61
65 - def _postprocess_triton(self) -> None:
66 - from vllm.triton_utils import HAS_TRITON
67 -
68 - if HAS_TRITON:
69 - logger.info(
70 - "Triton-CPU backend is available; skipping C++ monkey-patches "
71 - "for Triton kernels."
72 - )
73 - return
74 -
75 - import vllm.v1.worker.block_table
76 -
77 - vllm.v1.worker.block_table._COMPUTE_SLOT_MAPPING_KERNEL.kernel = (
78 - cpu_tl.compute_slot_mapping_kernel
79 - )
80 -
81 - # Speculative decoding fallbacks
82 - import vllm.v1.sample.rejection_sampler
83 - import vllm.v1.spec_decode.utils as spec_decode_utils
84 -
85 - spec_decode_utils._eagle_prepare_inputs_padded.kernel = (
86 - cpu_tl.eagle_prepare_inputs_padded_kernel
87 - )
88 - spec_decode_utils._eagle_prepare_next_token_padded.kernel = (
89 - cpu_tl.eagle_prepare_next_token_padded_kernel
90 - )
91 - spec_decode_utils._copy_and_expand_eagle_inputs.kernel = (
92 - cpu_tl.copy_and_expand_eagle_inputs_kernel
93 - )
94 - spec_decode_utils.copy_and_expand_dflash_inputs_kernel = (
95 - cpu_tl.copy_and_expand_dflash_inputs_kernel
96 - )
97 - dflash_module = sys.modules.get("vllm.v1.spec_decode.dflash")
98 - if dflash_module is not None:
99 - dflash_kernel_name = "copy_and_expand_dflash_inputs_kernel"
100 - setattr(
101 - dflash_module,
102 - dflash_kernel_name,
103 - cpu_tl.copy_and_expand_dflash_inputs_kernel,
104 - )
105 - spec_decode_utils.eagle_step_slot_mapping_metadata_kernel = (
106 - cpu_tl.eagle_step_slot_mapping_metadata_kernel
107 - )
108 - vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = (
109 - cpu_tl.rejection_greedy_sample_kernel
110 - )
111 - vllm.v1.sample.rejection_sampler.rejection_random_sample_kernel = (
112 - cpu_tl.rejection_random_sample_kernel
113 - )
114 - vllm.v1.sample.rejection_sampler.expand_kernel = cpu_tl.expand_kernel
115 - vllm.v1.sample.rejection_sampler.sample_recovered_tokens_kernel = (
116 - cpu_tl.sample_recovered_tokens_kernel
117 - )
118 -
119 - import vllm.v1.worker.mamba_utils
120 -
121 - vllm.v1.worker.mamba_utils.batch_memcpy_kernel = cpu_tl.batch_memcpy_kernel
122 -
62 @instrument(span_name="Loading (CPU)")
63 def load_model(self, load_dummy_weights: bool = False) -> None:
64 if load_dummy_weights:
vllm/v1/worker/worker_base.py
+6
@@ -268,6 +268,12 @@ class WorkerWrapperBase:
268
269 load_general_plugins()
270
271 + # Let the platform replace core Triton kernels (e.g. CPU fallback
272 + # implementations) before any worker or model code can launch them.
273 + from vllm.platforms import current_platform
274 +
275 + current_platform.register_triton_kernel_overrides()
276 +
277 parallel_config = vllm_config.parallel_config
278 if isinstance(parallel_config.worker_cls, str):
279 worker_class: type[WorkerBase] = resolve_obj_by_qualname(