Disable Triton wrap by default, only enable in lowering pass

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

.buildkite/test_areas/misc.yaml

@@ -207,8 +207,8 @@ steps:
     - pip install pytest-timeout pytest-forked
     - pytest -v -s v1/determinism/test_batch_invariance.py
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
 
 - label: Batch Invariance (B200)
   timeout_in_minutes: 30
@@ -222,10 +222,10 @@ steps:
     - pip install pytest-timeout pytest-forked
     - pytest -v -s v1/determinism/test_batch_invariance.py
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN
     - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
-  
+
 - label: Acceptance Length Test (Large Models) # optional
   timeout_in_minutes: 25
   gpu: h100

tests/compile/passes/ir/test_lowering.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 from torch import nn
+from torch._library.triton import set_wrap_triton_enabled
 
 import vllm.kernels  # noqa: F401 to register kernels
 from vllm import ir
@@ -47,6 +48,7 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config):
     with (
         ops.rms_norm.set_priority([rms_provider, "native"]),
         ir.enable_torch_wrap(True),
+        set_wrap_triton_enabled(False),  # set by default in forward context
     ):
         compiled_model = torch.compile(model, backend=backend, fullgraph=True)
         compiled_unlowered_model = torch.compile(

tests/ir/test_op.py

@@ -51,7 +51,7 @@ def test_registration_overloads():
     def _custom_div(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return x / y
 
-    custom_div = IrOp("_custom_div", _custom_div)
+    custom_div = IrOp("_custom_div", _custom_div, False)
     assert custom_div.name == "_custom_div"
     assert "_custom_div" not in IrOp.registry
 
@@ -336,7 +336,7 @@ class TestIrOpImplDispatch:
         assert "priority not set" in message
 
 
-@vllm.ir.register_op
+@vllm.ir.register_op(has_reduction=True)
 def _custom_mm(
     x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
 ) -> torch.Tensor:
@@ -344,15 +344,16 @@ def _custom_mm(
     return tmp if bias is None else tmp + bias
 
 
+@_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)
+def impl_mm(
+    x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
+) -> torch.Tensor:
+    tmp = x @ y
+    return tmp + 50 if bias is None else tmp + bias + 100
+
+
 def test_default_args():
     # Test that default args are properly applied when dispatching and calling
-    @_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)
-    def impl_mm(
-        x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
-    ) -> torch.Tensor:
-        tmp = x @ y
-        return tmp + 50 if bias is None else tmp + bias + 100
-
     x1 = torch.tensor([1, 2], dtype=torch.int32)
     x2 = torch.tensor([3, 4], dtype=torch.int32)
 
@@ -362,6 +363,70 @@ def test_default_args():
         assert _custom_mm.dispatch(x1, x2) is impl_mm
 
 
+@_custom_add.register_impl(
+    "bv_impl_add", batch_invariant=False, supports_args=lambda x, y: True
+)
+def bv_impl_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    return x + y + 70
+
+
+@_custom_mm.register_impl(
+    "bi_impl_mm", batch_invariant=True, supports_args=lambda x, y, bias=None: True
+)
+def bi_impl_mm(
+    x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
+) -> torch.Tensor:
+    return x @ y + 20 + bias
+
+
+@_custom_mm.register_impl(
+    "bv_impl_mm", batch_invariant=False, supports_args=lambda x, y, bias=None: True
+)
+def bv_impl_mm(
+    x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
+) -> torch.Tensor:
+    return x @ y + 20 + bias
+
+
+def test_batch_invariant_defaults():
+    # _custom_add is batch invariant by default
+    assert _custom_add.impls["native"].batch_invariant
+    assert _custom_add.impls["impl_even"].batch_invariant
+    assert not _custom_add.impls["bv_impl_add"].batch_invariant
+
+    # _custom_mm is not batch invariant by default
+    assert _custom_mm.impls["native"].batch_invariant
+    assert _custom_mm.impls["bi_impl_mm"].batch_invariant
+    assert not _custom_mm.impls["impl_mm"].batch_invariant
+    assert not _custom_mm.impls["bv_impl_mm"].batch_invariant
+
+
+def test_batch_invariant_dispatching():
+    # batch invariance off, all ops remain
+    with _custom_add.set_priority(
+        ["bv_impl_add", "impl_even", "native"], batch_invariant_only=False
+    ):
+        assert _custom_add.get_priority() == ["bv_impl_add", "impl_even", "native"]
+
+    # batch invariance required, filter ops
+    with _custom_add.set_priority(
+        ["impl_even", "bv_impl_add", "native"], batch_invariant_only=True
+    ):
+        assert _custom_add.get_priority() == ["impl_even", "native"]
+
+    # batch invariance off, all ops remain
+    with _custom_mm.set_priority(
+        ["bv_impl_mm", "bi_impl_mm", "native"], batch_invariant_only=False
+    ):
+        assert _custom_mm.get_priority() == ["bv_impl_mm", "bi_impl_mm", "native"]
+
+    # batch invariance required, filter ops
+    with _custom_mm.set_priority(
+        ["bv_impl_mm", "bi_impl_mm", "native"], batch_invariant_only=True
+    ):
+        assert _custom_mm.get_priority() == ["bi_impl_mm", "native"]
+
+
 def test_bad_impl_registrations():
     # Check bad schema
     with pytest.raises(ValueError, match="does not match native schema"):
@@ -435,7 +500,12 @@ def test_bad_impl_registrations():
         ) -> torch.Tensor:
             return x @ y + 40
 
-    assert set(_custom_mm.impls.keys()) == {"impl_mm", "native"}
+    assert set(_custom_mm.impls.keys()) == {
+        "bi_impl_mm",
+        "bv_impl_mm",
+        "impl_mm",
+        "native",
+    }
 
 
 IMPL_OOT_SRC = """

tests/kernels/ir/test_layernorm.py

@@ -24,12 +24,15 @@ rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn
     reason="Currently only kernels on CUDA, ROCm and XPU",
 )
 def test_rms_norm_registration():
+    from vllm.triton_utils import HAS_TRITON
+
     expected = {
         "native": True,
         "vllm_c": current_platform.is_cuda_alike(),
         "aiter": current_platform.is_rocm(),
         "oink": False,
         "xpu_kernels": current_platform.is_xpu(),
+        "triton_batch_invariant": HAS_TRITON,
     }
 
     actual = {
@@ -71,11 +74,12 @@ class TestRMSNorm:
         out4 = rms_norm_native(x, None, epsilon=epsilon)
         torch.testing.assert_close(out3, out4)
 
-    @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels"])
+    @pytest.mark.parametrize("provider", vllm.ir.ops.rms_norm.supported_providers())
     def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
         impl = ir.ops.rms_norm.impls[provider]
-        if not impl.supported:
-            pytest.skip(f"{provider} impl not supported on this platform")
+        assert impl.supported, (
+            f"{provider} impl expected to be supported on this platform"
+        )
 
         x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
         args = (x, weight, epsilon, None)
@@ -102,9 +106,10 @@ class TestRMSNorm:
         # exact match
         torch.testing.assert_close(out_impl2, out_impl, rtol=0.0, atol=0.0)
 
-        # none of these support variance_size override
-        assert not impl.supports_args(x, weight, epsilon, 4)
-        assert not impl.supports_args(x, weight, epsilon, variance_size=4)
+        if impl.provider != "native":
+            # none of the kernels support variance_size override
+            assert not impl.supports_args(x, weight, epsilon, 4)
+            assert not impl.supports_args(x, weight, epsilon, variance_size=4)
 
         # test weight=None behavior
         out_impl_no_weight = impl.impl_fn(x, None, epsilon)

tests/v1/determinism/test_batch_invariance.py

@@ -23,12 +23,11 @@ IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()
 
 @skip_unsupported
 @pytest.mark.timeout(1000)
-@pytest.mark.parametrize(
-    "backend",
-    BACKENDS,
-)
+@pytest.mark.parametrize("backend", BACKENDS)
+@pytest.mark.parametrize("enforce_eager", [True, False])
 def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
     backend,
+    enforce_eager: bool,
 ):
     """
     Ensures that the same request (the 'needle' prompt) yields identical output
@@ -51,6 +50,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
       seed.
     - Keep max_tokens and max_model_len bounded for speed and memory use.
     """
+    if not enforce_eager and IS_DEVICE_CAPABILITY_BELOW_90:
+        pytest.skip("enforce_eager required for <sm90")
+
     seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
     random.seed(seed)
 
@@ -90,6 +92,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
             max_num_seqs=max_batch_size,
             gpu_memory_utilization=gpu_mem_util,
             max_model_len=max_model_len,
+            enforce_eager=enforce_eager,
             attention_config=attention_config,
         )
 
@@ -146,13 +149,15 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
 
 
 @skip_unsupported
-@pytest.mark.parametrize(
-    "backend",
-    BACKENDS,
-)
+@pytest.mark.parametrize("backend", BACKENDS)
+@pytest.mark.parametrize("enforce_eager", [True, False])
 def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
     backend,
+    enforce_eager: bool,
 ):
+    if not enforce_eager and IS_DEVICE_CAPABILITY_BELOW_90:
+        pytest.skip("enforce_eager required for <sm90")
+
     seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
     random.seed(seed)
     tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
@@ -175,7 +180,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
         max_model_len=8192,
         dtype="auto",  # not everything is supported
         gpu_memory_utilization=0.9,
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
+        enforce_eager=enforce_eager,
         attention_config={"backend": backend},
     )
 
@@ -906,6 +911,7 @@ def LLM_with_max_seqs(
     max_num_seqs: int,
     gpu_memory_utilization: float,
     max_model_len: int,
+    enforce_eager: bool,
     attention_config: dict | None = None,
 ) -> LLM:
     """
@@ -920,7 +926,7 @@ def LLM_with_max_seqs(
         dtype="auto",
         tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
         enable_prefix_caching=False,
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
+        enforce_eager=enforce_eager,
         attention_config=attention_config,
         # Enable for MOE models
         # enable_expert_parallel=True,

vllm/compilation/passes/ir/lowering_pass.py

@@ -10,6 +10,7 @@ from torch._inductor.pattern_matcher import (
     PatternMatcherPass,
     register_graph_pattern,
 )
+from torch._library.triton import set_wrap_triton_enabled
 from torch._ops import OpOverload, OpOverloadPacket
 
 from vllm.config import VllmConfig
@@ -92,14 +93,18 @@ class VllmIRLoweringPass(VllmInductorPass):
         # Defaults not present on node.args but required for replacement tracing
         bound_args = ir_op._py_signature.bind(*node.args)
         bound_args.apply_defaults()
-        match.replace_by_example(ir_op_impl.impl_fn, bound_args.args)
+        match.replace_by_example(
+            ir_op_impl.impl_fn, bound_args.args, run_functional_passes=False
+        )
 
     @VllmInductorPass.time_and_log
     def __call__(self, graph: fx.Graph) -> None:
         # clear at the beginning instead of end, so that tests can inspect
         self.selected_impls.clear()
 
-        count = self.patterns.apply(graph)
+        # Triton wrap is disabled in the forward context, enable it during lowering.
+        with set_wrap_triton_enabled(True):
+            count = self.patterns.apply(graph)
         logger.debug("VllmIRLoweringPass lowered %d vLLM IR nodes", count)
 
         # TODO write self.selected_impls to depyf/tlparse dir

vllm/compilation/passes/pass_manager.py

@@ -113,8 +113,8 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
         VllmInductorPass.dump_prefix += 1
 
         # clean up after lowering again
-        self.post_cleanup(graph)
-        VllmInductorPass.dump_prefix += 1
+        # self.post_cleanup(graph)
+        # VllmInductorPass.dump_prefix += 1
 
         # always run fix_functionalization last
         self.fix_functionalization(graph)
@@ -190,7 +190,7 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
 
         passes.append(self.post_cleanup.uuid())
         passes.append(self.ir_lowering.uuid())
-        passes.append(self.post_cleanup.uuid())
+        # passes.append(self.post_cleanup.uuid())
         passes.append(self.fix_functionalization.uuid())
 
         # Include the compile range in the uuid to ensure that inductor

vllm/config/kernel.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import Field, field_validator
 
+from vllm import envs
 from vllm.config.utils import config, get_hash_factors, hash_factors
 from vllm.logger import init_logger
 
@@ -86,7 +87,11 @@ class IrOpPriorityConfig:
                     "Setting IR op priority for %s to %s", field.name, op_priority
                 )
                 ir_op = IrOp.registry[field.name]
-                stack.enter_context(ir_op.set_priority(op_priority))
+                stack.enter_context(
+                    ir_op.set_priority(
+                        op_priority, batch_invariant_only=envs.VLLM_BATCH_INVARIANT
+                    )
+                )
 
             yield
 

vllm/config/vllm.py

@@ -94,13 +94,14 @@ IS_DENSE = False
 
 def enable_norm_fusion(cfg: "VllmConfig") -> bool:
     """Enable if either RMS norm or quant FP8 custom op is active;
-    otherwise Inductor handles fusion."""
+    otherwise Inductor handles fusion. Also disable for batch invariant
+    as the custom fused kernels are not currently batch invariant."""
 
     return (
         cfg.compilation_config.is_custom_op_enabled("rms_norm")
         or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
         or cfg.kernel_config.ir_op_priority.rms_norm[0] != "native"
-    )
+    ) and not envs.VLLM_BATCH_INVARIANT
 
 
 def enable_act_fusion(cfg: "VllmConfig") -> bool:
@@ -117,7 +118,9 @@ def enable_act_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
-    """Enable if TP > 1 and Hopper/Blackwell and flashinfer installed."""
+    """Enable if TP > 1 and Hopper/Blackwell and flashinfer installed.
+    Disable if batch invariance is enabled.
+    """
     from vllm.platforms import current_platform
     from vllm.utils.flashinfer import has_flashinfer
 
@@ -135,6 +138,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
         # tp-pp combination broken:
         # https://github.com/vllm-project/vllm/issues/35426
         and cfg.parallel_config.pipeline_parallel_size == 1
+        and not envs.VLLM_BATCH_INVARIANT
     )
 
 

vllm/forward_context.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
 from typing import Any
 
 import torch
+from torch._library.triton import set_wrap_triton_enabled
 
 import vllm.envs as envs
 import vllm.ir
@@ -326,6 +327,7 @@ def set_forward_context(
             vllm.ir.enable_torch_wrap(
                 vllm_config.compilation_config.ir_enable_torch_wrap
             ),
+            set_wrap_triton_enabled(False),
         ):
             yield
     finally:

vllm/ir/op.py

@@ -53,6 +53,7 @@ def register_op(f: Callable[..., Any]) -> "IrOp": ...
 def register_op(
     *,
     name: str | None = None,
+    has_reduction: bool = False,
 ) -> Callable[[Callable[..., Any]], "IrOp"]: ...
 
 
@@ -60,12 +61,14 @@ def register_op(
     f: Callable | None = None,
     *,
     name: str | None = None,
+    has_reduction: bool = False,
 ) -> "IrOp | Callable[[Callable], IrOp]":
     """
     Register a new vLLM IR op.
 
     :param f: the native implementation of the op
     :param name: the name of the op, defaults to the function name
+    :param has_reduction: is this op is a reduction op, which affects batch-invariance
     :return: the IrOp object if f is provided, otherwise a decorator
 
     Example usage:
@@ -82,7 +85,7 @@ def register_op(
     def decorator(_f: Callable):
         op_name: str = _f.__name__ if name is None else name
         assert op_name not in IrOp.registry
-        op = IrOp(op_name, _f)
+        op = IrOp(op_name, _f, has_reduction)
         IrOp.registry[op_name] = op
         return op
 
@@ -96,9 +99,10 @@ class IrOp:
     registry: ClassVar[dict[str, "IrOp"]] = {}
 
     name: str
+    has_reduction: bool
     impls: dict[str, "IrOpImpl"]
 
-    def __init__(self, name: str, native_impl: Callable):
+    def __init__(self, name: str, native_impl: Callable, has_reduction: bool):
         self._py_signature = inspect.signature(native_impl)
         if any(
             p.kind == inspect.Parameter.KEYWORD_ONLY
@@ -110,13 +114,22 @@ class IrOp:
             )
 
         self.name = name
+        self.has_reduction = has_reduction
         self.impls: dict[str, IrOpImpl] = {}
         self._priority_impls: list[IrOpImpl] = []
         self._schema_str = infer_schema(native_impl, mutates_args=[])
 
         # native implementation
         self.impls["native"] = IrOpImpl(
-            self, "native", native_impl, supported=True, supports_args=None
+            self,
+            "native",
+            native_impl,
+            # always supported
+            supported=True,
+            supports_args=None,
+            # Native implementation is always batch-invariant
+            # (batch invariance is controlled at the torch level)
+            batch_invariant=True,
         )
 
         # By default, fake routes directly to native,
@@ -156,12 +169,14 @@ class IrOp:
         *,
         supported: bool = True,
         supports_args: Callable[..., bool] | None = None,
+        batch_invariant: bool | None = None,
     ):
         """
         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 batch_invariant: is this implementation is batch-invariant.
         :return: A decorator that registers the implementation.
 
         The decorated function must have the same semantics and signature as
@@ -182,13 +197,26 @@ class IrOp:
         def my_provider_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: ...
         ```
 
+        Default behavior of batch_invariant depends on op.has_reduction:
+        - op.has_reduction == True: batch_invariant = False
+        - op.has_reduction == False: batch_invariant = True
+
+        This is because ops without reductions are always batch-invariant
+        (unless explicitly opting out).
+        Ops with reductions have to opt in, as they are not batch-invariant by default.
+
         """
         assert provider not in RESERVED_PROVIDERS, (
             f"Provider name {provider} is reserved."
         )
 
+        if batch_invariant is None:
+            batch_invariant = not self.has_reduction
+
         def _register_impl(f: Callable):
-            impl = IrOpImpl(self, provider, f, supported, supports_args)
+            impl = IrOpImpl(
+                self, provider, f, supported, supports_args, batch_invariant
+            )
             self.impls[provider] = impl
 
             if self.get_priority():
@@ -274,12 +302,13 @@ 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], *, batch_invariant_only: bool = False):
         """
         Context manager to set the dispatch priority for implementations for this op.
         """
         assert all(p in self.impls for p in priority), (
-            "All providers in priority must be registered implementations."
+            f"All providers in priority must be registered implementations, missing "
+            f"{','.join(p for p in priority if p not in self.impls)}"
         )
 
         def filter_priority_impls(p_list: list[str]) -> list[IrOpImpl]:
@@ -290,6 +319,10 @@ class IrOp:
                     # Skip unsupported implementations
                     continue
 
+                if batch_invariant_only and not impl.batch_invariant:
+                    # Skip non-batch-invariant implementations
+                    continue
+
                 filtered_impls.append(impl)
 
                 # If all args are supported, skip other implementations
@@ -318,6 +351,12 @@ class IrOp:
 
 
 class IrOpImpl:
+    op: IrOp
+    provider: str
+    impl_fn: Callable
+    supported: bool
+    batch_invariant: bool
+
     def __init__(
         self,
         op: IrOp,
@@ -325,6 +364,7 @@ class IrOpImpl:
         impl_fn: Callable,
         supported: bool,
         supports_args: Callable[..., bool] | None,
+        batch_invariant: bool,
     ):
         assert provider not in op.impls, (
             f"Implementation for provider {provider} already registered."
@@ -388,6 +428,7 @@ class IrOpImpl:
         self.impl_fn = impl_fn
         self.supported = supported
         self._supports_args = supports_args
+        self.batch_invariant = batch_invariant
 
     @property
     def supports_all_args(self) -> bool:

vllm/ir/ops/layernorm.py

@@ -6,7 +6,7 @@ from torch import Tensor
 from ..op import register_op
 
 
-@register_op
+@register_op(has_reduction=True)
 def rms_norm(
     x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
 ) -> Tensor:

vllm/kernels/__init__.py

@@ -2,6 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Kernel implementations for vLLM."""
 
-from . import aiter_ops, oink_ops, vllm_c, xpu_ops
+from . import aiter_ops, oink_ops, triton, vllm_c, xpu_ops
 
-__all__ = ["vllm_c", "aiter_ops", "oink_ops", "xpu_ops"]
+__all__ = ["vllm_c", "aiter_ops", "oink_ops", "xpu_ops", "triton"]

vllm/kernels/triton/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from .layernorm_batch_invariant import *

vllm/kernels/triton/layernorm_batch_invariant.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+from torch import Tensor
+
+from vllm import ir
+from vllm.triton_utils import HAS_TRITON
+
+rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
+"""Variance size override not supported"""
+
+
+@ir.ops.rms_norm.register_impl(
+    "triton_batch_invariant",
+    supported=HAS_TRITON,
+    supports_args=rms_norm_no_var,
+    batch_invariant=True,
+)
+def rms_norm(
+    x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
+) -> Tensor:
+    assert variance_size is None
+    if weight is None:
+        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
+
+    # TODO move kernel here
+    from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant
+
+    return rms_norm_batch_invariant(x, weight, epsilon)

vllm/model_executor/layers/batch_invariant.py

@@ -792,6 +792,7 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None =
     return result
 
 
+@torch.library.wrap_triton  # needed for make_fx lowering to work
 @triton.jit
 def _rms_norm_kernel(
     input_ptr,

vllm/model_executor/layers/layernorm.py

@@ -264,7 +264,7 @@ class RMSNorm(CustomOp):
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        if residual is None and not envs.VLLM_BATCH_INVARIANT:
+        if residual is None:
             return ir.ops.rms_norm(
                 x, self.weight.data, self.variance_epsilon, self.variance_size_override
             )
@@ -276,8 +276,7 @@ class RMSNorm(CustomOp):
         # This mirrors vLLM's fused_add_rms_norm semantics by mutating both
         # `x` (normalized output) and `residual` (residual-out buffer).
         if (
-            residual is not None
-            and getattr(self, "_use_oink_fused_add_rmsnorm", False)
+            getattr(self, "_use_oink_fused_add_rmsnorm", False)
             and x.is_cuda
             and residual.is_cuda
             and x.shape == residual.shape
@@ -313,20 +312,14 @@ class RMSNorm(CustomOp):
                     )
                     return x, residual
 
-        if residual is not None:
-            return fused_add_rms_norm(
-                x, residual, self.weight.data, self.variance_epsilon
-            )
-        else:
-            assert envs.VLLM_BATCH_INVARIANT
-            return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)
+        return fused_add_rms_norm(x, residual, self.weight.data, self.variance_epsilon)
 
     def forward_hip(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        if residual is None and not envs.VLLM_BATCH_INVARIANT:
+        if residual is None:
             return ir.ops.rms_norm(
                 x, self.weight.data, self.variance_epsilon, self.variance_size_override
             )
@@ -334,13 +327,9 @@ class RMSNorm(CustomOp):
         if self.variance_size_override is not None:
             return self.forward_native(x, residual)
 
-        if residual is not None:
-            return self.rocm_norm_func_with_add(
-                x, residual, self.weight.data, self.variance_epsilon
-            )
-        else:
-            assert envs.VLLM_BATCH_INVARIANT
-            return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)
+        return self.rocm_norm_func_with_add(
+            x, residual, self.weight.data, self.variance_epsilon
+        )
 
     def forward_xpu(
         self,

vllm/platforms/cuda.py

@@ -576,12 +576,19 @@ class CudaPlatformBase(Platform):
         using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
         default = ["native"] if using_inductor else ["vllm_c", "native"]
 
+        # triton_batch_invariant available even when VLLM_BATCH_INVARIANT=0,
+        # but it won't be selected. vllm_c is skipped when VLLM_BATCH_INVARIANT=1.
+        rms_norm = (
+            ["native"]
+            if using_inductor
+            else ["vllm_c", "triton_batch_invariant", "native"]
+        )
+
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,
         #  users can just use IR op priority directly
-        rms_norm = default
         if envs.VLLM_USE_OINK_OPS:
-            rms_norm = ["oink"] + default
+            rms_norm = ["oink"] + rms_norm
 
         return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)