| 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) |