From c72e3600b9073cfaf785cf1bed23b1d1bb7c1754 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Mon, 18 May 2026 17:16:18 +0800 Subject: [PATCH] [Feature] Triton dispatcher Signed-off-by: wangxiyuan --- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 13 +- .../model_executor/test_triton_dispatcher.py | 222 ++++++++++++++++++ vllm/platforms/cpu.py | 54 +++++ vllm/platforms/interface.py | 13 + vllm/triton_utils/dispatcher.py | 175 ++++++++++++++ vllm/utils/cpu_triton_utils.py | 36 --- vllm/v1/worker/cpu_model_runner.py | 61 ----- vllm/v1/worker/worker_base.py | 6 + 8 files changed, 482 insertions(+), 98 deletions(-) create mode 100644 tests/model_executor/test_triton_dispatcher.py create mode 100644 vllm/triton_utils/dispatcher.py diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index c239f831c614..b2480dd39227 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -1024,7 +1024,18 @@ def test_batch_memcpy_cpu_fallback() -> None: copy each src into its dst, validating the (src_ptrs, dst_ptrs, sizes) argument order against ctypes.memmove(dst, src, size). """ - from vllm.utils.cpu_triton_utils import batch_memcpy_kernel + from vllm.triton_utils.dispatcher import ( + KernelOverride, + register_kernels, + ) + from vllm.utils.cpu_triton_utils import _batch_memcpy_impl + + register_kernels( + {"vllm.v1.worker.mamba_utils.batch_memcpy_kernel": _batch_memcpy_impl} + ) + from vllm.v1.worker.mamba_utils import batch_memcpy_kernel + + assert isinstance(batch_memcpy_kernel, KernelOverride) # Varied byte sizes, including a non-power-of-two run. sizes_bytes = [256, 1024, 17 * 4, 4096] diff --git a/tests/model_executor/test_triton_dispatcher.py b/tests/model_executor/test_triton_dispatcher.py new file mode 100644 index 000000000000..909230c40f5c --- /dev/null +++ b/tests/model_executor/test_triton_dispatcher.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the apply-time Triton kernel override registry. + +The registry must let out-of-tree platforms replace kernels defined as +plain ``@triton.jit`` functions, patching every hold-site (defining +module, from-import copies, and JIT warmup owners) so that existing call +sites dispatch to the override without any core-side decoration. +""" + +import sys +import types +from typing import Any + +import pytest + +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + _DecoratedTritonJitKernel, +) +from vllm.triton_utils import tl, triton +from vllm.triton_utils.dispatcher import ( + KernelOverride, + register_kernels, +) + + +class _FakeTritonKernel: + """Stands in for a Triton JITFunction (works without Triton installed).""" + + arg_names = ("output_ptr", "input_ptr", "replace_from", "MAX_NUM_TOKENS") + + def __getitem__(self, grid: Any) -> Any: + def launch(*args: Any, **kwargs: Any) -> Any: + return ("default", grid, kwargs) + + return launch + + +@pytest.fixture +def _fake_module(monkeypatch: pytest.MonkeyPatch) -> Any: + """Expose a fake defining module holding the kernel and its warmup owner.""" + import types + + kernel = _FakeTritonKernel() + owner = _DecoratedTritonJitKernel( + kernel, + lambda: {}, + lambda output, input, replace_from: ( # noqa: ARG005 + (1,), + dict(MAX_NUM_TOKENS=8), + ), + ) + module = types.ModuleType("vllm.tests.fake_kernel_home") + module.expand_kernel = kernel # type: ignore[attr-defined] + module._expand = owner # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, module.__name__, module) + return module + + +def test_override_patches_defining_module_and_warmup_owner( + _fake_module: Any, +) -> None: + calls: list[tuple[Any, ...]] = [] + + def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None): + calls.append((output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS)) + return "cpu" + + register_kernels( + { + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl, + } + ) + + overridden = _fake_module.expand_kernel + assert isinstance(overridden, KernelOverride) + # Launch with kernel-arg names like the warmup launch path does. + assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu" + assert calls == [(1, 2, 0, None)] + + # The warmup owner's captured kernel is patched too, and its cached + # argument-name introspection is invalidated. + assert _fake_module._expand.kernel is overridden + + +def test_override_forwards_positionally_for_mismatched_names( + _fake_module: Any, +) -> None: + calls: list[tuple[Any, ...]] = [] + + def cpu_impl(output, input_val, replace_from, MAX_NUM_TOKENS=None): + calls.append((output, input_val, replace_from, MAX_NUM_TOKENS)) + return "cpu" + + register_kernels( + { + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl, + } + ) + overridden = _fake_module.expand_kernel + assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu" + assert calls == [(1, 2, 0, None)] + + +def test_override_covers_from_import_copies(_fake_module: Any) -> None: + import types + + importer = types.ModuleType("vllm.tests.fake_kernel_importer") + importer.expand_kernel = _fake_module.expand_kernel # type: ignore[attr-defined] + sys.modules["vllm.tests.fake_kernel_importer"] = importer + + def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None): + return "cpu" + + register_kernels( + { + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl, + } + ) + assert isinstance(importer.expand_kernel, KernelOverride) + del sys.modules["vllm.tests.fake_kernel_importer"] + + +def test_override_owner_launch_goes_through_impl(_fake_module: Any) -> None: + launched: list[Any] = [] + + def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None): + launched.append((output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS)) + return "cpu" + + register_kernels( + { + "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl, + } + ) + owner = _fake_module._expand + # In-place kernels return None from launch; assert on the recorded call. + owner(1, 2, 0) + assert launched == [(1, 2, 0, 8)] + + +def test_register_kernel_auto_imports_defining_module( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Registering a kernel whose defining module is not imported yet must + import it on the spot and apply the override.""" + kernel_name = "vllm.v1.sample.rejection_sampler.expand_kernel" + module_name = kernel_name.rpartition(".")[0] + calls: list[tuple[Any, ...]] = [] + + def cpu_impl( + output_ptr, + input_ptr, + cu_num_tokens_ptr, + replace_from, + replace_to, + MAX_NUM_TOKENS=None, + ): + calls.append((output_ptr, replace_from, MAX_NUM_TOKENS)) + return "cpu" + + # Ensure the module is NOT imported before registration. Drop it from + # sys.modules; a fresh import creates new kernel/owner objects, which + # is what a cold-start platform process would see. + saved = sys.modules.pop(module_name, None) + # Re-run the registry patch against the original module afterwards so + # other tests that already imported it keep working. + try: + assert module_name not in sys.modules + register_kernels({kernel_name: cpu_impl}) + assert module_name in sys.modules + finally: + if saved is None: + # First import in the whole session: keep the fresh module. + saved = sys.modules[module_name] + else: + # Restore the original module object and re-apply the override + # to its (pre-existing) kernel so both generations agree. + sys.modules[module_name] = saved + register_kernels({kernel_name: cpu_impl}) + + overridden = saved.expand_kernel + assert isinstance(overridden, KernelOverride) + overridden[(2,)]( + output_ptr=1, + input_ptr=2, + cu_num_tokens_ptr=3, + replace_from=0, + replace_to=0, + MAX_NUM_TOKENS=8, + ) + assert calls == [(1, 0, 8)] + # The warmup owner captured at the fresh import must be patched too + # (when the module existed before, there are two module generations). + fresh = sys.modules[module_name] + if fresh is not saved: + assert fresh._expand.kernel is overridden + + +def test_real_triton_kernel_override_end_to_end() -> None: + """Override a real @triton.jit kernel through the registry.""" + if not triton.__class__.__module__.startswith("triton"): + pytest.skip("Triton is not installed") + + @triton.jit + def _add_one_kernel(x_ptr, out_ptr, BLOCK: tl.constexpr): # pragma: no cover + offsets = tl.arange(0, BLOCK) + tl.store(out_ptr + offsets, tl.load(x_ptr + offsets) + 1) + + module = types.ModuleType("vllm.tests.real_jit_kernel_home") + module._add_one_kernel = _add_one_kernel # type: ignore[attr-defined] + sys.modules[module.__name__] = module + + def cpu_impl(x_ptr, out_ptr, BLOCK=None): # noqa: ARG001 + return "cpu" + + register_kernels({"vllm.tests.real_jit_kernel_home._add_one_kernel": cpu_impl}) + assert module._add_one_kernel.arg_names == ( # type: ignore[attr-defined] + "x_ptr", + "out_ptr", + "BLOCK", + ) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 390c0dbfede7..405d5c70a5c6 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -626,6 +626,60 @@ def import_kernels(cls) -> None: except ImportError as e: logger.warning_once("Failed to import from vllm._C: %s", repr(e)) + @classmethod + def register_triton_kernel_overrides(cls) -> None: + """Fallback C++ implementations for core Triton kernels. + + Used when the Triton-CPU backend is unavailable. When Triton-CPU is + present, the Triton kernels can run natively and the overrides + would only shadow them, so registration is skipped. + """ + from vllm.triton_utils import HAS_TRITON + + if HAS_TRITON: + return + + from vllm.triton_utils.dispatcher import register_kernels + from vllm.utils import cpu_triton_utils as cpu_tl + + register_kernels( + { + "vllm.v1.worker.block_table.ComputeSlotMappingKernel.kernel": ( + cpu_tl._compute_slot_mapping_kernel_impl + ), + "vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel": ( + cpu_tl._eagle_step_slot_mapping_metadata_kernel_impl + ), + "vllm.v1.spec_decode.utils.eagle_prepare_inputs_padded_kernel": ( + cpu_tl._eagle_prepare_inputs_padded_kernel_impl + ), + "vllm.v1.spec_decode.utils.eagle_prepare_next_token_padded_kernel": ( + cpu_tl._eagle_prepare_next_token_padded_kernel_impl + ), + "vllm.v1.spec_decode.utils.copy_and_expand_eagle_inputs_kernel": ( + cpu_tl._copy_and_expand_eagle_inputs_kernel_impl + ), + "vllm.v1.spec_decode.utils.copy_and_expand_dflash_inputs_kernel": ( + cpu_tl._copy_and_expand_dflash_inputs_kernel_impl + ), + "vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel": ( + cpu_tl._rejection_greedy_sample_kernel_impl + ), + "vllm.v1.sample.rejection_sampler.rejection_random_sample_kernel": ( + cpu_tl._rejection_random_sample_kernel_impl + ), + "vllm.v1.sample.rejection_sampler.expand_kernel": ( + cpu_tl._expand_kernel_impl + ), + "vllm.v1.sample.rejection_sampler.sample_recovered_tokens_kernel": ( + cpu_tl._sample_recovered_tokens_kernel_impl + ), + "vllm.v1.worker.mamba_utils.batch_memcpy_kernel": ( + cpu_tl._batch_memcpy_impl + ), + } + ) + @classmethod def pack_kv_cache( cls, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 13e5b55402c6..0711d0eda8f1 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -364,6 +364,19 @@ def import_kernels(cls) -> None: with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 + @classmethod + def register_triton_kernel_overrides(cls) -> None: + """Override core Triton kernels with platform implementations. + + Platforms whose Triton kernels need replacement implementations + (e.g. CPU fallbacks when the Triton-CPU backend is unavailable) + override this method and call + :func:`vllm.triton_utils.dispatcher.register_kernels` + with their overrides. It runs during model runner init, before any + kernel launch. The default implementation does nothing. + """ + return + @classmethod def get_attn_backend_cls( cls, diff --git a/vllm/triton_utils/dispatcher.py b/vllm/triton_utils/dispatcher.py new file mode 100644 index 000000000000..021ed97214e9 --- /dev/null +++ b/vllm/triton_utils/dispatcher.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Registry for overriding Triton kernels with platform implementations. + +Non-CUDA platforms can replace Triton kernels defined in vLLM core +without modifying core code: + + from vllm.triton_utils.dispatcher import register_kernels + + register_kernels({ + "vllm.v1.sample.rejection_sampler.expand_kernel": my_expand_impl, + "vllm.v1.worker.mamba_utils.batch_memcpy_kernel": my_memcpy_impl, + }) +""" + +import importlib +import inspect +import sys +from collections.abc import Callable, Mapping +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_registry: dict[str, tuple[Callable[..., Any], Callable[..., Any]]] = {} + + +def _kernel_arg_names(kernel: Callable[..., Any]) -> tuple[str, ...]: + """Return the argument names of a Triton kernel or its fallback.""" + names = getattr(kernel, "arg_names", None) + if names is None: + # No real Triton (TritonPlaceholder): read the plain function. + fn = getattr(kernel, "func", kernel) + names = tuple(inspect.signature(fn).parameters) + return tuple(names) + + +class KernelOverride: + """Stands in for a Triton kernel and routes launches to ``impl``. + + The JIT warmup infrastructure launches kernels as + ``kernel[grid](**kwargs)`` where ``kwargs`` are keyed by the original + kernel's argument names. This wrapper forwards those launch arguments + to the platform implementation: by keyword when the implementation's + parameter names match the kernel's, otherwise positionally in the + kernel's parameter order. + """ + + def __init__(self, kernel: Callable[..., Any], impl: Callable[..., Any]) -> None: + self._impl = impl + # Mirror the original kernel's names so launch binding and warmup + # introspection keep working against this wrapper. + self.arg_names = _kernel_arg_names(kernel) + self.constexprs = getattr(kernel, "constexprs", None) + self.func = impl + self.__name__ = getattr(impl, "__name__", "kernel_override") + self.__module__ = getattr(kernel, "__module__", "") + self._forward_by_name = ( + tuple(inspect.signature(impl).parameters) == self.arg_names + ) + + def __getitem__(self, grid: Any) -> Callable[..., Any]: + return self._launch + + def _launch(self, *args: Any, **kwargs: Any) -> Any: + if args: + return self._impl(*args, **kwargs) + if self._forward_by_name: + return self._impl(**kwargs) + unexpected = [name for name in kwargs if name not in self.arg_names] + if unexpected: + raise RuntimeError( + f"Override for {self.__name__!r} received unexpected launch " + f"arguments: {unexpected}" + ) + return self._impl(*(kwargs[name] for name in self.arg_names if name in kwargs)) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self._impl(*args, **kwargs) + + +def _resolve_kernel(name: str) -> tuple[Any, str]: + """Return the (host object, attribute) holding the kernel `name`.""" + module_name, attr = name.rsplit(".", 1) + parts = module_name.split(".") + if parts[-1][:1].isupper(): + class_module = importlib.import_module(".".join(parts[:-1])) + return getattr(class_module, parts[-1]), attr + importlib.import_module(module_name) + module = sys.modules[module_name] + if hasattr(module, attr): + return module, attr + raise ValueError(f"Kernel {name!r} not found for override") + + +def _rebind_kernels( + overrides: list[tuple[Callable[..., Any], KernelOverride]], +) -> None: + """Point every reference to each original kernel at its wrapper. + + A kernel object can be captured in two kinds of places besides its + defining module, and both must be patched for the overrides to take + effect everywhere: + + * plain attributes of other modules (``from mod import kernel`` + copies), rebound to the wrapper; + * JIT warmup owners, whose ``kernel`` attribute is the kernel they + launch; swapped, and their cached ``_kernel_arg_names`` invalidated + so the launch binding re-derives from the wrapper. + + One pass over ``sys.modules`` handles all kernels in `overrides`. + """ + originals = {original: wrapper for original, wrapper in overrides} + + def lookup(value: Any) -> KernelOverride | None: + # Only compare by identity: this scan runs over every module + # namespace, and some attribute values (e.g. PlaceholderModule + # sentinels) trigger imports or errors on hash/eq. + for original in originals: + if value is original: + return originals[original] + return None + + for module in list(sys.modules.values()): + namespace = getattr(module, "__dict__", None) + if not namespace: + continue + for attr, value in list(namespace.items()): + wrapper = lookup(value) + if wrapper is not None: + setattr(module, attr, wrapper) + continue + # JIT warmup owners (e.g. _DecoratedTritonJitKernel) hold the + # kernel as an instance attribute; read the instance dict directly + # so module-level __getattr__ hooks are never triggered, and so + # this scan needs no dependency on the warmup machinery. + if hasattr(value, "__dict__"): + wrapper = lookup(vars(value).get("kernel")) + if wrapper is None: + continue + value.kernel = wrapper + if not isinstance(value, type): + value.__dict__.pop("_kernel_arg_names", None) + + +def register_kernels(overrides: Mapping[str, Callable[..., Any]]) -> None: + """Register platform-specific implementations for Triton kernels. + + Args: + overrides: Mapping from fully qualified kernel name to the + platform implementation. Kernel names look like + "vllm.v1.sample.rejection_sampler.expand_kernel"; for kernels + owned by a JIT warmup class, include the attribute path, e.g. + "vllm.v1.worker.block_table.ComputeSlotMappingKernel.kernel". + Implementations are invoked like the kernels themselves, with + the launch arguments in the original kernel's parameter order + (or by keyword when their parameter names match the kernel's). + + """ + resolved: list[tuple[str, Any, str, Callable[..., Any], KernelOverride]] = [] + for name, impl in overrides.items(): + if name in _registry: + logger.warning("Kernel %s is already registered. Overwriting.", name) + host, attr = _resolve_kernel(name) + original = getattr(host, attr) + wrapper = KernelOverride(original, impl) + _registry[name] = (original, impl) + resolved.append((name, host, attr, original, wrapper)) + # Rebind other holders first: once the defining attributes are swapped + # below, the scan could no longer find the original kernels. + _rebind_kernels([(original, wrapper) for *_, original, wrapper in resolved]) + for name, host, attr, _, wrapper in resolved: + setattr(host, attr, wrapper) + logger.debug("Registered kernel override %s", name) diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index d1696317aab8..d8e4cd389d30 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -4,19 +4,10 @@ """Contains replacement functions to fallback Triton usages in CPU backend.""" import ctypes -from collections.abc import Callable import torch -class _FuncWrapper: - def __init__(self, func: Callable) -> None: - self.func = func - - def __getitem__(self, *args, **kwargs) -> Callable: - return self.func - - # For _compute_slot_mapping_kernel in vllm/v1/worker/block_table.py def _compute_slot_mapping_kernel_impl( num_tokens: int, @@ -47,9 +38,6 @@ def _compute_slot_mapping_kernel_impl( ) -compute_slot_mapping_kernel = _FuncWrapper(_compute_slot_mapping_kernel_impl) - - def _ensure_int64(t: torch.Tensor) -> torch.Tensor: return t if t.dtype == torch.int64 else t.to(torch.int64) @@ -454,31 +442,7 @@ def _sample_recovered_tokens_kernel_impl( output_token_ids.copy_(output_i64.to(orig_dtype)) -eagle_prepare_inputs_padded_kernel = _FuncWrapper( - _eagle_prepare_inputs_padded_kernel_impl -) -eagle_prepare_next_token_padded_kernel = _FuncWrapper( - _eagle_prepare_next_token_padded_kernel_impl -) -copy_and_expand_eagle_inputs_kernel = _FuncWrapper( - _copy_and_expand_eagle_inputs_kernel_impl -) -copy_and_expand_dflash_inputs_kernel = _FuncWrapper( - _copy_and_expand_dflash_inputs_kernel_impl -) -eagle_step_slot_mapping_metadata_kernel = _FuncWrapper( - _eagle_step_slot_mapping_metadata_kernel_impl -) -rejection_greedy_sample_kernel = _FuncWrapper(_rejection_greedy_sample_kernel_impl) -rejection_random_sample_kernel = _FuncWrapper(_rejection_random_sample_kernel_impl) -expand_kernel = _FuncWrapper(_expand_kernel_impl) -sample_recovered_tokens_kernel = _FuncWrapper(_sample_recovered_tokens_kernel_impl) - - def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None): # BLOCK_SIZE is unused; kept for signature parity with the Triton kernel. for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()): ctypes.memmove(dst, src, size) - - -batch_memcpy_kernel = _FuncWrapper(_batch_memcpy_impl) diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 04ca2b9194e4..32fa01bd27ec 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -1,13 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import sys from contextlib import AbstractContextManager, contextmanager from typing import Any import torch import torch.nn as nn -import vllm.utils.cpu_triton_utils as cpu_tl from vllm.config import ( CompilationMode, VllmConfig, @@ -37,7 +35,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.cascade_attn_enabled = False self._postprocess_tensors() - self._postprocess_triton() def _postprocess_tensors(self) -> None: # Note: replace device tensors with cpu tensors @@ -62,64 +59,6 @@ def replace_tensor(obj: Any, cpu_attr_name: str, device_attr_name) -> None: if isinstance(v, CpuGpuBuffer): v.gpu = v.cpu - def _postprocess_triton(self) -> None: - from vllm.triton_utils import HAS_TRITON - - if HAS_TRITON: - logger.info( - "Triton-CPU backend is available; skipping C++ monkey-patches " - "for Triton kernels." - ) - return - - import vllm.v1.worker.block_table - - vllm.v1.worker.block_table._COMPUTE_SLOT_MAPPING_KERNEL.kernel = ( - cpu_tl.compute_slot_mapping_kernel - ) - - # Speculative decoding fallbacks - import vllm.v1.sample.rejection_sampler - import vllm.v1.spec_decode.utils as spec_decode_utils - - spec_decode_utils._eagle_prepare_inputs_padded.kernel = ( - cpu_tl.eagle_prepare_inputs_padded_kernel - ) - spec_decode_utils._eagle_prepare_next_token_padded.kernel = ( - cpu_tl.eagle_prepare_next_token_padded_kernel - ) - spec_decode_utils._copy_and_expand_eagle_inputs.kernel = ( - cpu_tl.copy_and_expand_eagle_inputs_kernel - ) - spec_decode_utils.copy_and_expand_dflash_inputs_kernel = ( - cpu_tl.copy_and_expand_dflash_inputs_kernel - ) - dflash_module = sys.modules.get("vllm.v1.spec_decode.dflash") - if dflash_module is not None: - dflash_kernel_name = "copy_and_expand_dflash_inputs_kernel" - setattr( - dflash_module, - dflash_kernel_name, - cpu_tl.copy_and_expand_dflash_inputs_kernel, - ) - spec_decode_utils.eagle_step_slot_mapping_metadata_kernel = ( - cpu_tl.eagle_step_slot_mapping_metadata_kernel - ) - vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = ( - cpu_tl.rejection_greedy_sample_kernel - ) - vllm.v1.sample.rejection_sampler.rejection_random_sample_kernel = ( - cpu_tl.rejection_random_sample_kernel - ) - vllm.v1.sample.rejection_sampler.expand_kernel = cpu_tl.expand_kernel - vllm.v1.sample.rejection_sampler.sample_recovered_tokens_kernel = ( - cpu_tl.sample_recovered_tokens_kernel - ) - - import vllm.v1.worker.mamba_utils - - vllm.v1.worker.mamba_utils.batch_memcpy_kernel = cpu_tl.batch_memcpy_kernel - @instrument(span_name="Loading (CPU)") def load_model(self, load_dummy_weights: bool = False) -> None: if load_dummy_weights: diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index b0839bc4cae2..017b1230cdac 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -268,6 +268,12 @@ def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: load_general_plugins() + # Let the platform replace core Triton kernels (e.g. CPU fallback + # implementations) before any worker or model code can launch them. + from vllm.platforms import current_platform + + current_platform.register_triton_kernel_overrides() + parallel_config = vllm_config.parallel_config if isinstance(parallel_config.worker_cls, str): worker_class: type[WorkerBase] = resolve_obj_by_qualname(