benchmark changes

Signed-off-by: Luka Govedič <lgovedic@redhat.com>

benchmarks/kernels/ir/bench_ir_ops.py

@@ -145,9 +145,6 @@ def _bench_one(fn, args, cfg: BenchConfig) -> float:
     return ms * 1000
 
 
-# TODO(gmagogsfm): Once compiled native implementation lands (#38775),
-# the benchmark baseline should be the compiled native (what vLLM runs by
-# default) rather than the uncompiled native implementation.
 def collect_timings(
     op: IrOp, shape_configs: list[dict], cfg: BenchConfig
 ) -> tuple[list[str], list[str], dict[str, dict[str, float]]]:
@@ -158,11 +155,15 @@ def collect_timings(
         "_".join(f"{k}={fmt(v)}" for k, v in kwargs.items()) for kwargs in shape_configs
     ]
     providers = [n for n, impl in op.impls.items() if impl.supported]
+    provider_strs = [f"{p}-compiled" if op.impls[p].compiled else p for p in providers]
 
     results: dict[str, dict[str, float]] = {c: {} for c in case_names}
-    for provider in providers:
+    for provider, p_str in zip(providers, provider_strs):
         impl = op.impls[provider]
-        desc = f"{op.name} / {provider}"
+        if impl.compiled:
+            impl = impl.compile()
+
+        desc = f"{op.name} / {p_str}"
         for case_name, kwargs in tqdm(
             zip(case_names, shape_configs),
             desc=desc,
@@ -171,11 +172,11 @@ def collect_timings(
         ):
             args = op.generate_inputs(**kwargs)
             if impl.supports_args(*args):
-                results[case_name][provider] = _bench_one(impl.impl_fn, args, cfg)
+                results[case_name][p_str] = _bench_one(impl.impl_fn, args, cfg)
             else:
-                results[case_name][provider] = float("nan")
+                results[case_name][p_str] = float("nan")
 
-    return case_names, providers, results
+    return case_names, provider_strs, results
 
 
 def analyze_results(
@@ -184,7 +185,7 @@ def analyze_results(
     providers: list[str],
     results: dict[str, dict[str, float]],
 ) -> tuple[list[dict[str, str]], list[dict[str, str]], list[str]]:
-    native_col = "native"
+    native_col = "native-compiled" if "native-compiled" in providers else "native"
     non_native = [p for p in providers if p != native_col]
 
     header_cols = ["case"]
@@ -229,7 +230,9 @@ def analyze_results(
         losses = sum(1 for s in speedups if s < 1.0)
         total = len(speedups)
 
-        print(f"\n{p} vs native ({wins}/{total} faster, {losses}/{total} slower):")
+        print(
+            f"\n{p} vs {native_col} ({wins}/{total} faster, {losses}/{total} slower):"
+        )
         print(f"  geomean speedup: {geomean:.2f}x")
         print(f"  best:            {best_val:.2f}x  ({best_case})")
         print(f"  worst:           {worst_val:.2f}x  ({worst_case})")
@@ -288,10 +291,10 @@ def save_results(
 def parse_args():
     parser = argparse.ArgumentParser(description="Benchmark vLLM IR ops")
     parser.add_argument(
-        "--ops",
+        "op",
         type=str,
-        default=None,
-        help="Comma-separated list of op names to benchmark (substring match)",
+        nargs="*",
+        help="Op name(s) to benchmark",
     )
     parser.add_argument(
         "--no-cuda-graph",
@@ -339,12 +342,11 @@ def main():
     )
     os.makedirs(save_dir, exist_ok=True)
 
-    op_filters = [f.strip() for f in args.ops.split(",")] if args.ops else None
+    op_names = [f.strip() for f in args.ops] if args.op else IrOp.registry.keys()
     all_summary_rows: list[dict[str, str]] = []
 
-    for op in IrOp.registry.values():
-        if op_filters and not any(f in op.name for f in op_filters):
-            continue
+    for op_name in op_names:
+        op = IrOp.registry[op_name]
         if not op.has_input_generator:
             print(f"Skipping op '{op.name}': no input generator registered")
             continue
@@ -354,6 +356,11 @@ def main():
                 f"Add it to benchmarks/kernels/ir/shapes.py"
             )
 
+        # Sort descending, so that torch.compile first sees the largest shape
+        # (simulates vLLM usage)
+        shape_configs = SHAPE_CONFIGS[op.name]
+        shape_configs.sort(key=lambda x: x["num_tokens"], reverse=True)
+
         case_names, providers, results = collect_timings(
             op, SHAPE_CONFIGS[op.name], cfg
         )

benchmarks/kernels/ir/shapes.py

@@ -26,4 +26,10 @@ SHAPE_CONFIGS: dict[str, list[dict]] = {
         for d in COMMON_HIDDEN_SIZES
         for n in NUM_TOKENS
     ],
+    "fused_add_rms_norm": [
+        {"num_tokens": n, "hidden_size": d, "dtype": dtype}
+        for dtype in [torch.float16, torch.bfloat16, torch.float32]
+        for d in COMMON_HIDDEN_SIZES
+        for n in NUM_TOKENS
+    ],
 }

tests/ir/test_compile.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+import torch._dynamo
+
+import vllm.ir.op
+from vllm.ir.op import IrOpImplCompiledWrapper
+
+
+# Create a simple IR op for testing
+@vllm.ir.register_op
+def _test_add_mul(
+    x_a: torch.Tensor, x_b: torch.Tensor, scale: float = 2.0
+) -> torch.Tensor:
+    """Simple op: (a + b) * scale"""
+    return (x_a + x_b) * scale
+
+
+@pytest.fixture(autouse=True)
+def clear_compile_wrapper():
+    """Clear the compile wrapper before every test."""
+    _test_add_mul.impls["native"].compile_clear()
+
+
+def test_compile_context_manager_semantics():
+    """Test that compile=True in set_priority produces correct results"""
+    a = torch.randn(4, 5)
+    b = torch.randn(4, 5)
+    scale = 3.0
+
+    # Without compilation
+    with _test_add_mul.set_priority(["native"], compile=False):
+        assert not isinstance(_test_add_mul.dispatch(a, b), IrOpImplCompiledWrapper)
+        out_no_compile = _test_add_mul(a, b, scale)
+
+    # With compilation
+    with _test_add_mul.set_priority(["native"], compile=True):
+        assert isinstance(_test_add_mul.dispatch(a, b), IrOpImplCompiledWrapper)
+        out_compile = _test_add_mul(a, b, scale)
+
+    # Both should produce the same result
+    torch.testing.assert_close(out_no_compile, out_compile)
+    torch.testing.assert_close(out_compile, (a + b) * scale)
+
+
+def test_no_recompile():
+    """Test that dynamic shape is marked correctly and no recompilation happens."""
+    torch._dynamo.reset()
+
+    a = torch.randn(4, 5)
+    b = torch.randn(4, 5)
+    scale = 3.0
+
+    # Without compilation
+    with _test_add_mul.set_priority(["native"], compile=True):
+        out1 = _test_add_mul(a, b, scale)
+
+    torch.testing.assert_close(out1, (a + b) * scale)
+
+    a = torch.randn(10, 5)
+    b = torch.randn(10, 5)
+
+    with (
+        torch.compiler.set_stance("fail_on_recompile"),
+        _test_add_mul.set_priority(["native"], compile=True),
+    ):
+        out2 = _test_add_mul(a, b, scale)
+
+    torch.testing.assert_close(out2, (a + b) * scale)
+
+
+def test_compile_inside_custom_op():
+    """Test that compiled impl works when called inside a custom op"""
+
+    # Create a custom op that calls our IR op internally
+    @torch.library.custom_op("mylib::wrapped_add_mul", mutates_args=())
+    def wrapped_add_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return _test_add_mul(a, b, scale=3.0)
+
+    # Use compile=True to optimize even though this custom op is opaque
+
+    @wrapped_add_mul.register_fake
+    def _(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return torch.empty_like(x)
+
+    # Test the wrapped op
+    a = torch.randn(4, 5)
+    b = torch.randn(4, 5)
+
+    assert _test_add_mul.impls["native"]._compiled_wrapper is None
+    with _test_add_mul.set_priority(["native"], compile=True):
+        result = wrapped_add_mul(a, b)
+
+    torch.testing.assert_close(result, (a + b) * 3.0)
+    assert _test_add_mul.impls["native"]._compiled_wrapper is not None
+
+
+def test_compile_with_custom_impl():
+    """Test that compile=True works with custom implementations"""
+
+    @_test_add_mul.register_impl("optimized", compiled=True)
+    def optimized_impl(
+        x_a: torch.Tensor, x_b: torch.Tensor, scale: float = 2.0
+    ) -> torch.Tensor:
+        # Alternative implementation
+        return torch.mul(torch.add(x_a, x_b), scale) + 4
+
+    a = torch.randn(3, 3)
+    b = torch.randn(3, 3)
+
+    # Use the custom implementation with compile=True
+    with _test_add_mul.set_priority(["optimized"], compile=True):
+        result = _test_add_mul(a, b, scale=1.5)
+
+    torch.testing.assert_close(result, (a + b) * 1.5 + 4)
+
+
+def test_nested_priority_contexts():
+    """Test that nested set_priority contexts work correctly with compile"""
+    a = torch.randn(4, 5)
+    b = torch.randn(4, 5)
+
+    # Outer context with compile=False
+    with _test_add_mul.set_priority(["native"], compile=False):
+        assert not isinstance(_test_add_mul.dispatch(a, b), IrOpImplCompiledWrapper)
+        out1 = _test_add_mul(a, b)
+
+        # Inner context with compile=True
+        with _test_add_mul.set_priority(["native"], compile=True):
+            assert isinstance(_test_add_mul.dispatch(a, b), IrOpImplCompiledWrapper)
+            out2 = _test_add_mul(a, b)
+
+        # Back to compile=False
+        assert not isinstance(_test_add_mul.dispatch(a, b), IrOpImplCompiledWrapper)
+        out3 = _test_add_mul(a, b)
+
+    # All should produce the same result
+    torch.testing.assert_close(out1, out2)
+    torch.testing.assert_close(out2, out3)
+
+
+def test_compiled_flag():
+    """Test that compiled flag controls whether impl can be compiled"""
+    # Native impl should have compiled=True
+    assert _test_add_mul.impls["native"].compiled is True
+    assert _test_add_mul.impls["native"].uncompiled_impl_fn is not None
+
+    # Register impl without compiled flag (defaults to False)
+    @_test_add_mul.register_impl("no_compile")
+    def no_compile_impl(
+        x_a: torch.Tensor, x_b: torch.Tensor, scale: float = 2.0
+    ) -> torch.Tensor:
+        return (x_a + x_b) * scale
+
+    assert _test_add_mul.impls["no_compile"].compiled is False
+
+    a = torch.randn(4, 5)
+    b = torch.randn(4, 5)
+
+    # compile=True with no_compile impl should NOT compile it
+    with _test_add_mul.set_priority(["no_compile"], compile=True):
+        # Should use the original impl, not compiled wrapper
+        assert not isinstance(_test_add_mul.dispatch(a, b), IrOpImplCompiledWrapper)
+        out = _test_add_mul(a, b)
+
+    torch.testing.assert_close(out, (a + b) * 2.0)

vllm/ir/op.py

@@ -5,7 +5,7 @@ import inspect
 import traceback
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, ClassVar, Literal, overload
+from typing import Any, ClassVar, Literal, Protocol, overload
 
 import regex as re
 import torch
@@ -141,6 +141,35 @@ def register_op(
     return decorator
 
 
+class CallableOpImpl(Protocol):
+    """
+    Protocol for callable op implementations.
+
+    Both IrOpImpl and IrOpImplCompiledWrapper implement this protocol,
+    allowing them to be used interchangeably in priority lists.
+    """
+
+    impl_fn: Callable
+    """The implementation function to call."""
+
+    uncompiled_impl_fn: Callable
+    """The uncompiled version of the implementation function."""
+
+    supported: bool
+    """Is this implementation supported? Asserted during dispatch"""
+
+    provider: str
+    """Implementation provider"""
+
+    def supports_args(self, *args, **kwargs) -> bool:
+        """Check if this implementation supports the given args."""
+        ...
+
+    def func_impl_fn(self, *args, **kwargs) -> Any:
+        """Call impl_fn with functional semantics (copying for inplace impls)."""
+        ...
+
+
 class IrOp:
     registry: ClassVar[dict[str, "IrOp"]] = {}
 
@@ -176,14 +205,14 @@ class IrOp:
         self.name = name
         self._docstring = inspect.getdoc(native_impl) or ""
         self._registration_stack = registration_stack or []
-        self.impls: dict[str, IrOpImpl] = {}
         self.activations = activations
         self.activation_indices = [
             i
             for i, p in enumerate(self._py_signature.parameters.values())
             if p.name in activations
         ]
-        self._priority_impls: list[IrOpImpl] = []
+        self.impls: dict[str, IrOpImpl] = {}
+        self._priority_impls: list[CallableOpImpl] = []
         self._schema_str = infer_schema(native_impl, mutates_args=[])
         self._input_generator: InputGenerator | None = None
         self._tolerance_overrides: ToleranceSpec = {}
@@ -196,6 +225,8 @@ class IrOp:
             # always supported
             supported=True,
             supports_args=None,
+            # Native can be compiled
+            compiled=True,
             registration_stack=self._registration_stack,
         )
 
@@ -237,13 +268,15 @@ class IrOp:
         supported: bool = True,
         supports_args: Callable[..., bool] | None = None,
         inplace: bool = False,
+        compiled: bool = False,
     ):
         """
         Register an implementation for this custom op.
         :param provider: The name of the provider, must be unique.
         :param supported: Static support check, use this to check platform support.
         :param supports_args: Dynamic arg support check, used for types and shapes.
-        :param inplace: Does this op reuse activation input memory for outputs
+        :param inplace: Does this impl reuse activation input memory for outputs
+        :param compiled: Does this impl get compiled if set_priority(..., compile=True)
         :return: A decorator that registers the implementation.
 
         The decorated function must have the same semantics and signature as
@@ -258,6 +291,9 @@ class IrOp:
         compatible with the implementation.
         For custom enablement logic, set op impl priority.
 
+        If compiled is True, that means that eager-mode dispatch will dispatch to a
+        torch.compile-decorated version of this implementation (with
+
         Example:
         ```python
         @my_op.register_impl("my_provider", supported=torch.cuda.is_available())
@@ -273,7 +309,9 @@ class IrOp:
         def _register_impl(f: Callable):
             # Slice out the decorator function from the stack
             stack = traceback.format_stack()[:-1]
-            impl = IrOpImpl(self, provider, f, supported, supports_args, inplace, stack)
+            impl = IrOpImpl(
+                self, provider, f, supported, supports_args, inplace, compiled, stack
+            )
             self.impls[provider] = impl
 
             if self.get_priority():
@@ -310,10 +348,10 @@ class IrOp:
         bound_args.apply_defaults()
         return bound_args.args
 
-    def dispatch(self, *args, **kwargs) -> "IrOpImpl":
+    def dispatch(self, *args, **kwargs) -> CallableOpImpl:
         """
         Dispatch to the appropriate implementation based on current priority
-        and argument support checks. Returns the selected IrOpImpl.
+        and argument support checks. Returns the selected CallableOpImpl.
 
         THIS FUNCTION IS ON THE HOT PATH (OP DISPATCH), MUST BE FAST.
         """
@@ -373,23 +411,31 @@ class IrOp:
         return [p.provider for p in self._priority_impls]
 
     @contextlib.contextmanager
-    def set_priority(self, priority: list[str]):
+    def set_priority(self, priority: list[str], *, compile: bool = False):
         """
         Context manager to set the dispatch priority for implementations for this op.
+
+        :param priority: List of provider names in priority order
+        :param compile: Wrap implementations with torch.compile for better performance
+                       when called inside opaque custom ops.
         """
         assert all(p in self.impls for p in priority), (
             "All providers in priority must be registered implementations."
         )
 
-        def filter_priority_impls(p_list: list[str]) -> list[IrOpImpl]:
-            filtered_impls = []
+        # If compile=True and impl allows compilation, use compiled wrapper
+        def maybe_compile(impl: IrOpImpl) -> CallableOpImpl:
+            return impl.compile() if compile and impl.compiled else impl
+
+        def filter_priority_impls(p_list: list[str]) -> list[CallableOpImpl]:
+            filtered_impls: list[CallableOpImpl] = []
             for p in p_list:
                 impl = self.impls[p]
                 if not impl.supported:
                     # Skip unsupported implementations
                     continue
 
-                filtered_impls.append(impl)
+                filtered_impls.append(maybe_compile(impl))
 
                 # If all args are supported, skip other implementations
                 if impl.supports_all_args:
@@ -401,7 +447,7 @@ class IrOp:
                 "explicitly add 'native' to the end of the priority list",
                 self.name,
             )
-            filtered_impls.append(self.impls["native"])
+            filtered_impls.append(maybe_compile(self.impls["native"]))
             return filtered_impls
 
         # Temporarily set priority
@@ -409,9 +455,10 @@ class IrOp:
         try:
             self._priority_impls = filter_priority_impls(priority)
             logger.debug(
-                "Priority for vllm.ir.%s set to %s",
+                "Priority for vllm.ir.%s set to %s%s",
                 self.name,
                 lazy(lambda: [p.provider for p in self._priority_impls]),
+                " (compiled)" if compile else "",
             )
             yield
         finally:
@@ -523,6 +570,7 @@ class IrOpImpl:
         supported: bool,
         supports_args: Callable[..., bool] | None,
         inplace: bool = False,
+        compiled: bool = False,
         registration_stack: list[str] | None = None,
     ):
         assert provider not in op.impls, (
@@ -593,9 +641,12 @@ class IrOpImpl:
         self.op = op
         self.provider = provider
         self.impl_fn = impl_fn
+        self.uncompiled_impl_fn = impl_fn  # Always the uncompiled version
         self.supported = supported
         self._supports_args = supports_args
         self.inplace = inplace
+        self.compiled = compiled  # Whether this impl can be compiled
+        self._compiled_wrapper: IrOpImplCompiledWrapper | None = None
         self._registration_stack = registration_stack or []
 
     @property
@@ -609,6 +660,30 @@ class IrOpImpl:
 
         return self._supports_args(*args, **kwargs)
 
+    def compile(self) -> "IrOpImplCompiledWrapper":
+        """
+        Get a compiled wrapper for this implementation.
+
+        This helps when the IR op is either called inside an opaque torch custom op and
+        hence invisible to model-level compilation, or outside any compilation context.
+        By pre-compiling the implementation, we can still get performance benefits even
+        when the op is not compiled externally.
+
+        The wrapper also marks activation tensors with dynamic batch dimensions
+        using torch._dynamo.mark_dynamic(t, 0) to ensure proper dynamic shapes.
+
+        The compiled wrapper is cached, so multiple calls return the same wrapper.
+
+        :return: Compiled wrapper following the CallableOpImpl protocol
+        """
+        if self._compiled_wrapper is None:
+            self._compiled_wrapper = IrOpImplCompiledWrapper(self)
+        return self._compiled_wrapper
+
+    def compile_clear(self):
+        """Clear the cached compile wrapper for the implementation."""
+        self._compiled_wrapper = None
+
     @weak_cache
     def uuid(self):
         """
@@ -637,3 +712,47 @@ class IrOpImpl:
             new_args[i] = args[i].clone()
 
         return self.impl_fn(*new_args, **kwargs)
+
+
+class IrOpImplCompiledWrapper:
+    """
+    Wrapper for IrOpImpl that provides a torch.compile-wrapped implementation.
+
+    This wrapper implements CallableOpImpl protocol so it can be used in
+    priority lists, but wraps the implementation with torch.compile and
+    marks dynamic dimensions on activations.
+    """
+
+    def __init__(self, base_impl: IrOpImpl, **compile_kwargs):
+        self.base_impl = base_impl
+        self.supported = base_impl.supported
+        self.provider = f"{base_impl.provider} (compiled)"
+        self.activation_indices = base_impl.op.activation_indices
+        self.uncompiled_impl_fn = base_impl.impl_fn
+
+        # Compile the implementation
+        compile_kwargs = {"dynamic": False, **compile_kwargs}
+        self.compiled_impl_fn = torch.compile(base_impl.impl_fn, **compile_kwargs)
+
+        # Create impl_fn that marks dynamic dims and calls compiled implementation
+        def impl_fn(*args, **kwargs) -> Any:
+            # Mark batch dimension (dim 0) as dynamic for activation inputs
+            for idx in self.activation_indices:
+                if idx < len(args) and isinstance(args[idx], torch.Tensor):
+                    torch._dynamo.mark_dynamic(args[idx], 0)
+
+            return self.compiled_impl_fn(*args, **kwargs)
+
+        self.impl_fn = impl_fn
+
+    def supports_args(self, *args, **kwargs) -> bool:
+        return self.base_impl.supports_args(*args, **kwargs)
+
+    def func_impl_fn(self, *args, **kwargs) -> Any:
+        """
+        Call impl_fn with functional semantics.
+
+        TODO: Implement compiled version that handles inplace impls correctly.
+        For now, delegate to the base implementation's func_impl_fn.
+        """
+        return self.base_impl.func_impl_fn(*args, **kwargs)