[Bugfix] Forward launch grid to Triton kernel overrides (#63999)

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

tests/model_executor/test_triton_dispatcher.py

@@ -10,7 +10,8 @@ sites dispatch to the override without any core-side decoration.
 
 import sys
 import types
-from typing import Any
+from collections.abc import Callable
+from typing import Any, cast
 
 import pytest
 
@@ -139,6 +140,172 @@ def test_override_owner_launch_goes_through_impl(_fake_module: Any) -> None:
     assert launched == [(1, 2, 0, 8)]
 
 
+def test_override_forwards_grid_to_kernel_like_impl(_fake_module: Any) -> None:
+    """A Triton-kernel override takes over the launch natively via __getitem__.
+
+    The override mirrors a real ``@triton.jit`` kernel: it binds the grid
+    in ``__getitem__`` and is launched by keyword under the original
+    kernel's parameter names, which the KernelOverride forwards unchanged.
+    """
+    seen: list[Any] = []
+
+    class NpuKernel:
+        def __getitem__(self, grid: Any) -> Any:
+            seen.append(grid)
+
+            def launch(
+                output_ptr,
+                input_ptr,
+                replace_from,
+                MAX_NUM_TOKENS=None,  # noqa: ARG001
+            ):
+                return ("launched", output_ptr, input_ptr, replace_from)
+
+            return launch
+
+    register_kernels(
+        {
+            "vllm.tests.fake_kernel_home.expand_kernel": cast(
+                Callable[..., Any], NpuKernel()
+            ),
+        }
+    )
+    overridden = _fake_module.expand_kernel
+    assert overridden[(8,)](output_ptr=1, input_ptr=2, replace_from=0) == (
+        "launched",
+        1,
+        2,
+        0,
+    )
+    assert seen == [(8,)]
+
+
+def test_override_plain_impl_never_sees_grid(_fake_module: Any) -> None:
+    """Plain-callable implementations keep receiving launch args without grid."""
+    calls: list[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 overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu"
+    assert calls == [(1, 2, 0, None)]
+
+
+def test_override_delegates_warmup_to_impl(_fake_module: Any) -> None:
+    """JIT warmup of an overridden kernel warms the implementation."""
+    warmed: list[Any] = []
+
+    class FakeNpuKernel:
+        arg_names = ("output_ptr", "input_ptr", "replace_from", "MAX_NUM_TOKENS")
+
+        def __getitem__(self, grid: Any) -> Any:
+            return lambda *args: None
+
+        def warmup(self, grid=None, **kwargs: Any) -> None:
+            warmed.append((grid, kwargs))
+
+    register_kernels(
+        {
+            "vllm.tests.fake_kernel_home.expand_kernel": cast(
+                Callable[..., Any], FakeNpuKernel()
+            ),
+        }
+    )
+    overridden = _fake_module.expand_kernel
+    assert overridden.warmup(grid=(1,), MAX_NUM_TOKENS=8) is None
+    assert warmed == [((1,), {"MAX_NUM_TOKENS": 8})]
+
+
+def test_override_warmup_noop_without_impl_warmup(_fake_module: Any) -> None:
+    """Plain-callable overrides have no Triton kernel to warm up."""
+    calls: list[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 overridden.warmup(grid=(1,)) is None
+    assert calls == []
+
+
+def test_register_triton_kernel_swaps_raw_for_device_function(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Triton-kernel impls (device functions included) are swapped raw.
+
+    A ``KernelOverride`` is not a Triton JITFunction, so wrapping the
+    callee would break an outer kernel's compilation; a Triton kernel
+    implementation must replace the original by identity in the defining
+    module and in from-import copies, where the inliner looks.
+    """
+    module = types.ModuleType("vllm.tests.fake_device_fn_home")
+
+    class FakeTritonKernel:
+        arg_names = ("a", "b", "BLOCK")
+
+        def __getitem__(self, grid: Any) -> Any:
+            return lambda *args, **kwargs: None
+
+        def run(self, *args: Any, **kwargs: Any) -> Any:
+            return None
+
+        def warmup(self, grid=None, **kwargs: Any) -> None:
+            return None
+
+        def __call__(self, *args: Any, **kwargs: Any) -> Any:
+            return None
+
+    replacement = FakeTritonKernel()
+    original = object()
+    module._device_fn = original  # type: ignore[attr-defined]
+    importer = types.ModuleType("vllm.tests.fake_device_fn_importer")
+    importer._device_fn = original  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    monkeypatch.setitem(sys.modules, importer.__name__, importer)
+
+    register_kernels(
+        {
+            "vllm.tests.fake_device_fn_home._device_fn": replacement,
+        }
+    )
+
+    assert module._device_fn is replacement  # type: ignore[attr-defined]
+    assert importer._device_fn is replacement  # type: ignore[attr-defined]
+    assert not isinstance(module._device_fn, KernelOverride)
+
+
+def test_register_non_triton_impl_is_wrapped(_fake_module: Any) -> None:
+    """Non-Triton impls (plain callables) still route through KernelOverride."""
+    calls: list[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)
+    assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu"
+    assert calls == [(1, 2, 0, None)]
+
+
 def test_register_kernel_auto_imports_defining_module(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:

vllm/triton_utils/dispatcher.py

@@ -11,6 +11,13 @@ without modifying core code:
         "vllm.v1.sample.rejection_sampler.expand_kernel": my_expand_impl,
         "vllm.v1.worker.mamba_utils.batch_memcpy_kernel": my_memcpy_impl,
     })
+
+The same call covers Triton *device functions* (``@triton.jit`` helpers
+called from inside other kernels and inlined at compile time): an
+implementation that is itself a Triton kernel is swapped in directly, so
+outer kernels keep compiling and launchable kernels run natively. Only
+non-Triton implementations (plain callables, grid-remapping launchers)
+are wrapped in a launch adapter.
 """
 
 import importlib
@@ -36,15 +43,35 @@ def _kernel_arg_names(kernel: Callable[..., Any]) -> tuple[str, ...]:
     return tuple(names)
 
 
+def _is_triton_kernel(obj: Any) -> bool:
+    """Whether ``obj`` is a real Triton kernel (JITFunction).
+
+    Triton kernels (both launchable kernels and device functions called
+    from inside other kernels) expose the JIT launch protocol; plain
+    callables and grid-remapping launcher objects do not.
+    """
+    return callable(obj) and all(
+        hasattr(obj, name) for name in ("__getitem__", "run", "warmup", "arg_names")
+    )
+
+
 class KernelOverride:
-    """Stands in for a Triton kernel and routes launches to ``impl``.
+    """Stand in for a Triton kernel, routing launches to a non-Triton ``impl``.
+
+    Used by :func:`register_kernels` for implementations that are not
+    themselves Triton kernels (plain callables, grid-remapping launchers);
+    Triton kernel implementations are swapped in directly and never wrapped.
 
     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.
+    kernel's parameter order. An implementation that reimplements the
+    ``kernel[grid]`` protocol (a grid-remapping launcher) receives the
+    caller's grid. Compile-only JIT warmup (``kernel.warmup(grid,
+    **kwargs)``) is delegated to the implementation's own ``warmup`` when
+    present, and is a no-op otherwise.
     """
 
     def __init__(self, kernel: Callable[..., Any], impl: Callable[..., Any]) -> None:
@@ -56,13 +83,47 @@ class KernelOverride:
         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
-        )
+        try:
+            impl_names = _kernel_arg_names(impl)
+        except (TypeError, ValueError):
+            # Kernel-like implementations (e.g. a grid-rewriting launcher)
+            # are not introspectable; their launch arguments are forwarded
+            # positionally in the original kernel's parameter order.
+            impl_names = ()
+        self._forward_by_name = impl_names == self.arg_names
 
     def __getitem__(self, grid: Any) -> Callable[..., Any]:
+        """Bind a launch grid, mirroring Triton's ``kernel[grid]``.
+
+        A platform override that is itself a Triton kernel (or any object
+        implementing ``__getitem__``) takes over the launch natively:
+        ``impl[grid](*args, **kwargs)`` binds against the override's own
+        signature, so the caller's grid reaches the override unchanged.
+        Such an override must mirror the original kernel's parameter names
+        (or order), since launch sites key arguments by the original names.
+        Plain-callable overrides are invoked with the launch arguments
+        directly.
+        """
+        impl_getitem = getattr(self._impl, "__getitem__", None)
+        if impl_getitem is not None:
+            return impl_getitem(grid)
         return self._launch
 
+    def warmup(self, grid: Any = None, **kwargs: Any) -> Any:
+        """Delegate compile-only warmup to the implementation.
+
+        The JIT warmup machinery compiles warmup-owned kernels as
+        ``kernel.warmup(grid, **kwargs)``. Forward that to the
+        implementation's own kernel so the override is what gets compiled.
+        Implementations without a warmable Triton kernel (plain callables,
+        grid-remapping launchers) have nothing to warm up and compile
+        lazily on first launch.
+        """
+        impl_warmup = getattr(self._impl, "warmup", None)
+        if impl_warmup is None:
+            return None
+        return impl_warmup(grid=grid, **kwargs)
+
     def _launch(self, *args: Any, **kwargs: Any) -> Any:
         if args:
             return self._impl(*args, **kwargs)
@@ -94,32 +155,34 @@ def _resolve_kernel(name: str) -> tuple[Any, str]:
     raise ValueError(f"Kernel {name!r} not found for override")
 
 
-def _rebind_kernels(
-    overrides: list[tuple[Callable[..., Any], KernelOverride]],
+def _rebind_impls(
+    replacements: list[tuple[Callable[..., Any], Callable[..., Any]]],
 ) -> None:
-    """Point every reference to each original kernel at its wrapper.
+    """Point every reference to each original object at its replacement.
 
-    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:
+    A kernel or device-function 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;
+      copies), rebound to the replacement;
     * 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.
+      so the launch binding re-derives from the replacement.
 
-    One pass over ``sys.modules`` handles all kernels in `overrides`.
+    One pass over ``sys.modules`` handles all objects in `replacements`.
     """
-    originals = {original: wrapper for original, wrapper in overrides}
+    replacements_by_original = {
+        original: replacement for original, replacement in replacements
+    }
 
-    def lookup(value: Any) -> KernelOverride | None:
+    def lookup(value: Any) -> Callable[..., Any] | 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:
+        for original in replacements_by_original:
             if value is original:
-                return originals[original]
+                return replacements_by_original[original]
         return None
 
     for module in list(sys.modules.values()):
@@ -127,19 +190,19 @@ def _rebind_kernels(
         if not namespace:
             continue
         for attr, value in list(namespace.items()):
-            wrapper = lookup(value)
-            if wrapper is not None:
-                setattr(module, attr, wrapper)
+            replacement = lookup(value)
+            if replacement is not None:
+                setattr(module, attr, replacement)
                 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:
+                replacement = lookup(vars(value).get("kernel"))
+            if replacement is None:
                 continue
-            value.kernel = wrapper
+            value.kernel = replacement
             if not isinstance(value, type):
                 value.__dict__.pop("_kernel_arg_names", None)
 
@@ -153,23 +216,39 @@ def register_kernels(overrides: Mapping[str, Callable[..., Any]]) -> None:
             "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).
+
+            An implementation that is itself a Triton kernel (also the
+            only kind that works for Triton *device functions*, which are
+            inlined into other kernels at compile time) is swapped in
+            directly and launches natively: ``impl[grid](*args, **kwargs)``
+            binds against the implementation's own signature and the
+            caller's grid reaches it unchanged. Such an implementation must
+            mirror the original kernel's parameter names (or order).
+
+            Any other implementation (a plain callable, or a launcher that
+            rewrites the grid for a non-CUDA device) is invoked with the
+            launch arguments in the original kernel's parameter order (or
+            by keyword when its parameter names match the kernel's).
 
     """
-    resolved: list[tuple[str, Any, str, Callable[..., Any], KernelOverride]] = []
+    resolved: list[tuple[str, Any, str, Callable[..., Any], Callable[..., Any]]] = []
     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)
+        # Triton kernel implementations are swapped in raw: they launch
+        # natively at call sites and keep compiling when an outer kernel
+        # inlines them as device functions. Wrapping such an object in
+        # KernelOverride would break the outer kernel's compilation.
+        replacement = (
+            impl if _is_triton_kernel(impl) else KernelOverride(original, impl)
+        )
         _registry[name] = (original, impl)
-        resolved.append((name, host, attr, original, wrapper))
+        resolved.append((name, host, attr, original, replacement))
     # 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)
+    _rebind_impls([(original, replacement) for *_, original, replacement in resolved])
+    for name, host, attr, _, replacement in resolved:
+        setattr(host, attr, replacement)
         logger.debug("Registered kernel override %s", name)