[vLLM IR] 2/N fused_add_rms_norm and maybe_inplace overload (#36823)

Signed-off-by: Luka Govedič <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

docs/design/debug_vllm_compile.md

@@ -5,12 +5,14 @@ TL;DR:
 - use tlparse to acquire torch.compile logs. Include these logs in bug reports and/or support asks.
 - The vLLM-torch.compile integration is multiple pieces. vLLM exposes flags to turn off each piece:
 
-| Online Flag | Offline Flag | Result |
-| ----------- | ------------ | ------ |
-| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs |
-| -cc.mode=0 | mode=CompilationMode.NONE | Turn off torch.compile only |
-| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only |
-| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor |
+| Online Flag                    | Offline Flag                                                                   | Result                                               |
+|--------------------------------|--------------------------------------------------------------------------------|------------------------------------------------------|
+| --enforce-eager                | enforce_eager=True                                                             | Turn off torch.compile and CUDAGraphs                |
+| -cc.mode=0                     | compilation_config=CompilationConfig(mode=CompilationMode.NONE)                | Turn off torch.compile only                          |
+| -cc.mode=1                     | compilation_config=CompilationConfig(mode=CompilationMode.STOCK_TORCH_COMPILE) | Turn off vLLM-compile modifications to torch.compile |
+| -cc.cudagraph_mode=NONE        | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE)        | Turn off CUDAGraphs only                             |
+| -cc.backend=eager              | compilation_config=CompilationConfig(backend='eager')                          | Turn off TorchInductor                               |
+| -cc.ir_enable_torch_wrap=False | compilation_config=CompilationConfig(ir_enable_torch_wrap=False)               | Turn off vLLM IR wrapping                            |
 
 ## vLLM-torch.compile overview
 
@@ -22,7 +24,7 @@ Most notably, vLLM-compile is NOT torch.compile, it is a custom compiler built u
 
 - Given a model, we do a full graph capture via TorchDynamo that is dynamic on the batch size (number of tokens)
 - vLLM then optionally splits and/or specializes this graph and then uses TorchInductor to compile each graph into a compiled artifact.
-This step may use vLLM custom Inductor passes to further optimize the graph.
+This step may use vLLM custom Inductor passes to further optimize the graph. This includes vLLM IR lowering to remove dispatch overhead.
 - The compiled artifact is saved to vLLM's compile cache so that it can be loaded in the future.
 - vLLM applies CUDAGraphs to reduce CPU overheads.
 
@@ -34,6 +36,7 @@ For more details on the design, please see the following resources:
 
 - [Introduction to vLLM-torch.compile blogpost](https://blog.vllm.ai/2025/08/20/torch-compile.html)
 - [vLLM-torch.compile integration design](./torch_compile.md)
+- [vLLM IR design](./vllm_ir.md)
 - [vLLM Office Hours #26](https://www.youtube.com/live/xLyxc7hxCJc?si=Xulo9pe53C6ywf0V&t=561)
 - [Talk at PyTorch Conference 2025](https://youtu.be/1wV1ESbGrVQ?si=s1GqymUfwiwOrDTg&t=725)
 
@@ -117,6 +120,21 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode
 LLM(model, compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE))
 ```
 
+vLLM IR makes heavy use of the compilation pipeline, from functionalization, custom fusions, and lowering.
+To turn that off and capture eager-mode dispatching behavior of vLLM IR, run with `ir_enable_torch_wrap=False`.
+IR torch wrap is only enabled by default when using `mode=VLLM_COMPILE` and `backend="inductor"` (default).
+
+```sh
+# Online
+vllm serve -cc.ir_enable_torch_wrap=False
+```
+
+```py
+# Offline
+from vllm.config.compilation import CompilationConfig
+LLM(model, compilation_config=CompilationConfig(ir_enable_torch_wrap=False))
+```
+
 ## Debugging TorchDynamo
 
 vLLM requires model code be capturable into a full graph via TorchDynamo (torch.compile's frontend).

docs/design/vllm_ir.md

@@ -0,0 +1,615 @@
+# vLLM IR: Functional Intermediate Representation
+
+## Motivation
+
+vLLM IR is a **functional intermediate representation (IR)** that fills the gap between
+low-level `torch` ops and vLLM layers like `RMSNorm` and quantization operators,
+By separating operator **semantics** from the **implementation** and **dispatching**,
+vLLM IR simplifies both compilation and kernel registration & dispatching simultaneously.
+It operates as a **dialect** in the torch FX representation, allowing full interoperability
+with “regular” torch ops & custom torch ops/kernels, as well as a piecewise migration from
+the previous `CustomOp` approach.
+
+Key design principles:
+
+- **Eager-compile consistency**: identical behavior (barring minor numerics) in eager and compiled modes
+- **Simple, transparent, yet powerful kernel selection**: good visibility and control allowing easy debugging
+- **Convention over configuration**: near-zero boilerplate required to register ops and implementations
+- **Extensibility**: ops and implementations can be registered anywhere, in-tree or out-of-tree
+- **Interoperability**: fully compatible with “regular” torch ops & custom torch ops/kernels,
+reducing developer friction and allowing piecewise migration
+
+The clean semantics/implementation separation enables a unified and extensible dispatching mechanism,
+allowing multiple kernels per-platform and powerful kernel selection. The separation also facilitates
+cleaner testing and benchmarking, removing much of the boilerplate standard for legacy approaches.
+
+By delaying kernel selection until late in the compilation process, the compiler can operate on
+a higher-level representation, which has the following main benefits:
+
+- Pattern matching in fusion/transformation passes only requires a single, simple pattern per op
+- OOT compiler backends can lower from the higher-level representation (in-progress)
+- The compiler can autotune over available implementations (future feature)
+
+## Quick Overview
+
+### Declaring an IR Operation
+
+IR operations are declared using the `@register_op` decorator with a native PyTorch implementation that defines the op's semantics:
+
+```python
+# vllm/ir/ops/layernorm.py
+from torch import Tensor
+from vllm.ir import register_op
+
+@register_op
+def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
+    """Weighted root-mean-square layer normalization"""
+    orig_dtype = x.dtype
+    x = x.to(torch.float32)
+    x_var = x if variance_size is None else x[..., :variance_size]
+    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + epsilon)
+    x = x.to(orig_dtype)
+    if weight is not None:
+        x = x * weight
+    return x
+```
+
+The native implementation serves three purposes:
+
+1. **Semantic definition**: Specifies the exact semantics of the operation, including shapes and strides
+2. **Default implementation**: Used when no other (better) implementation is available
+3. **Reference for testing**: Other implementations must match these semantics
+
+### Registering Implementations
+
+Kernel implementations are registered using the `register_impl` decorator on the IR op object:
+
+```python
+# vllm/kernels/vllm_c.py
+from vllm import ir
+
+rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
+
+@ir.ops.rms_norm.register_impl("vllm_c", supports_args=rms_norm_no_var, supported=current_platform.is_cuda_alike())
+def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
+    output = torch.empty_like(x)
+    torch.ops._C.rms_norm(output, x, weight, epsilon)
+    return output
+```
+
+Implementations can specify:
+
+- `supported`: Static boolean indicating if this implementation is available
+- `supports_args`: Function checking if the implementation supports specific arguments
+- `inplace`: Whether this implementation reuses input memory for outputs
+
+### Using IR Operations in Models
+
+IR operations are imported and called directly in model code:
+
+```python
+# vllm/model_executor/layers/layernorm.py
+from vllm import ir
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x: Tensor, residual: Tensor | None = None):
+        if residual is None:
+            return ir.ops.rms_norm(x, self.weight, self.variance_epsilon)
+
+        # Use maybe_inplace overload to allow implementation to reuse input memory for outputs
+        # (using x or residual after this call is undefined behavior)
+        return ir.ops.fused_add_rms_norm.maybe_inplace(
+            x, residual, self.weight, self.variance_epsilon
+        )
+```
+
+### Configuring Kernel Selection
+
+Kernel selection is controlled via priority lists in the configuration.
+Priority lists specify the order in which implementations are considered,
+with the first supported implementation being selected.
+This includes the static support check (`supported=...`) and
+the dynamic arg support check (`supports_args=...`).
+
+#### Command Line Configuration
+
+Use `--ir-op-priority.<op_name>=<provider1>,<provider2>,...`:
+
+```bash
+# CUDA: Use vllm_c implementation for rms_norm
+vllm serve meta-llama/Llama-3.2-1B \
+  --ir-op-priority.rms_norm=vllm_c
+
+# ROCm: Try aiter first, fall back to vllm_c, then native
+vllm serve meta-llama/Llama-3.2-1B \
+  --ir-op-priority.rms_norm=aiter,vllm_c,native
+
+# Configure multiple operations
+vllm serve meta-llama/Llama-3.2-1B \
+  --ir-op-priority.rms_norm=vllm_c \
+  --ir-op-priority.fused_add_rms_norm=vllm_c
+```
+
+#### Python Configuration
+
+```python
+from vllm import LLM
+from vllm.config import VllmConfig, KernelConfig
+
+llm = LLM(
+    model="meta-llama/Llama-3.2-1B",
+    vllm_config=VllmConfig(
+        kernel_config=KernelConfig(
+            ir_op_priority={
+                "rms_norm": ["vllm_c", "native"],
+                "fused_add_rms_norm": ["vllm_c", "native"],
+            }
+        )
+    )
+)
+```
+
+#### Platform Defaults
+
+Each platform provides default priority lists that are automatically applied:
+
+```python
+# CUDA/XPU/ROCm platform defaults (when compiling with Inductor)
+{
+  "rms_norm": ["native"],  # Native torch is default
+  "fused_add_rms_norm": ["native"],
+}
+
+# CUDA platform defaults (eager or Dynamo-only)
+{
+  "rms_norm": ["vllm_c", "native"],
+  "fused_add_rms_norm": ["vllm_c", "native"],
+}
+
+# ROCm platform defaults (future - currently same as CUDA)
+{
+    "rms_norm": ["aiter", "vllm_c", "native"],
+    "fused_add_rms_norm": ["aiter", "vllm_c", "native"],
+}
+
+# XPU platform defaults (eager or Dynamo-only)
+{
+    "rms_norm": ["xpu_kernels", "native"],
+    "fused_add_rms_norm": ["xpu_kernels", "native"],
+}
+```
+
+User-specified priorities are prepended to platform defaults,
+so you only need to specify the out-of-order implementations,
+other implementations are appended automatically.
+
+## Compilation Pipeline
+
+vLLM IR heavily customizes the `torch.compile`-based compilation process to allow custom compile
+passes to operate on high-level IR while still producing efficient low-level code at the end.
+The compilation pipeline consists of several stages:
+
+### 1. Dynamo Tracing
+
+When `torch.compile` traces the model's forward pass, vLLM IR operations appear as custom operations
+in the `vllm_ir` torch library. These operations are opaque to Dynamo, meaning they appear directly
+in the FX graph without decomposition:
+
+```python
+# Python code (epsilon=1e-5)
+x1 = ir.ops.rms_norm(x, weight, epsilon)
+x2, residual_out = ir.ops.fused_add_rms_norm.maybe_inplace(x1, residual, weight, epsilon)
+
+# FX graph after Dynamo tracing
+x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
+out = torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace(x1, residual, weight, 1e-5); x1 = residual = None
+x2 = out[0]
+residual_out = out[1]
+```
+
+### 2. AOTAutograd and Functionalization
+
+AOTAutograd functionalizes the graph, converting any mutating operations to functional equivalents.
+For vLLM IR operations with `maybe_inplace` overloads, we perform this manually before AOTAutograd,
+converting them to the functional `default` overload using the pre-grad custom pass hook.
+
+```python
+# After functionalization
+x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
+out = torch.ops.vllm_ir.fused_add_rms_norm.default(x1, residual, weight, 1e-5); x1 = residual = None
+x2 = out[0]
+residual_out = out[1]
+```
+
+The pass also tracks which inputs were "donated" (passed to `maybe_inplace`),
+storing this information in vLLM's `PassContext` for later use in clone elimination.
+
+### 3. IR Fusion and Transformation Passes
+
+After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations.
+These passes can perform fusion, distribute operations for sequence parallelism, and other transformations:
+
+```python
+# Example: Sequence Parallelism (see SequenceParallelismPass)
+# Before SP pass
+
+all_reduce = torch.ops.vllm.all_reduce(x, "tp:0")
+rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
+
+# after SP pass
+reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0")
+rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
+all_gather = torch.ops.vllm.all_gather(x, "tp:0")
+```
+
+Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations,
+handle different kernel implementations separately, or deal with functionalization of custom kernels.
+
+### 4. IR Lowering
+
+The lowering pass (`VllmIRLoweringPass`) replaces each vLLM IR operation with its selected implementation.
+The implementation is chosen based on the priority list and support predicates,
+using the **fake tensors** in the graph's metadata in place of op arguments:
+
+```python
+# Implementation selection, same in eager dispatch and compile lowering
+def dispatch(*args) -> IrOpImpl:
+  for provider in priority_list:  # e.g., ["vllm_c", "native"]
+    impl = ir_op.impls[provider]
+    if not impl.supported:
+      continue
+    if impl.supports_args and not impl.supports_args(*args):
+      continue
+    return impl
+
+# make_fx uses torch.fx.symbolic_trace
+impl_graph = make_fx(selected_impl.impl_fn)
+# Replace IR op node with impl_graph's nodes
+match.replace_by_example(selected_impl.impl_fn, node.args)
+```
+
+For example, lowering `rms_norm` with the `vllm_c` implementation:
+
+```python
+# Before lowering (IR op)
+rms_norm = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5)
+
+# After lowering (vllm_c implementation traced)
+# Note: Lowering does not currently functionalize, this will likely change in the future.
+empty =  torch.ops.aten.empty.memory_format(x.shape, ...)
+rms_norm = torch.ops._C.rms_norm(empty, x, weight, 1e-5)
+```
+
+When lowering an implementation that mutates inputs (`inplace=True`),
+the lowering pass inserts clones to preserve functional semantics:
+
+```python
+# vllm_c implementation for fused_add_rms_norm mutates its first two arguments
+# Lowered with clones for safety
+clone_default = torch.ops.aten.clone.default(x)
+clone_default_1 = torch.ops.aten.clone.default(residual)
+fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(clone_default, clone_default_1, weight, 1e-5)
+```
+
+### 5. Clone Cleanup
+
+After lowering, the clone elimination pass (`UnsafeCloneEliminationPass`) removes unnecessary clones introduced during lowering.
+This pass is essential for achieving zero-copy behavior when using in-place kernels with `maybe_inplace`.
+The pass removes a clone if:
+
+- the cloned input is created in the graph and not used again in the graph
+- the cloned input is a graph parameter, marked as donated
+
+```python
+# After cleanup (donated inputs, no subsequent uses)
+fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(x, residual, weight, 1e-5)
+```
+
+The combination of inplace functionalization (tracking donated inputs) and clone cleanup enables the compiler to safely
+use in-place kernels without adding redundant copies or increasing the memory usage.
+
+### 6. Inductor Optimization and Codegen
+
+After IR lowering and cleanup, the graph contains only standard PyTorch operations and platform-specific custom ops.
+Inductor then performs its standard codegen:
+
+- **Inductor lowering and pointwise fusion**: Fusing element-wise operations, reductions, etc.
+- **Memory planning**: Determining buffer allocation and reuse
+- **Kernel generation**: Generating Triton or C++ code for fused operations
+- **Autotuning**: Selecting the best kernel configurations
+
+### Pipeline Summary
+
+```text
+Model Forward Pass
+    ↓
+[Dynamo Tracing] → FX Graph with vllm_ir.* ops
+    ↓
+[Pre-grad: Inplace Functionalization] → maybe_inplace → default, track donated inputs
+    ↓
+[AOTAutograd] → Functionalization
+    ↓
+[Post-grad: IR Fusion Passes] → Fuse high-level IR ops (e.g., rms_norm + quant)
+    ↓
+[Post-grad: IR Lowering] → vllm_ir.* ops → impl ops (with clones if needed)
+    ↓
+[Post-grad: Clone Cleanup] → Remove unnecessary clones using donated input info
+    ↓
+[Inductor] → Pattern matching, fusion, memory planning, codegen
+    ↓
+Compiled Code
+```
+
+## Core vLLM IR Concepts
+
+### Operation Declaration
+
+Operations are declared with the `@register_op` decorator, which creates an `IrOp` object:
+
+```python
+@register_op(
+    name=None,           # Operation name (defaults to function name)
+    activations=None,    # List of activation parameters (defaults to params starting with 'x')
+    allow_inplace=False, # Whether to create a maybe_inplace overload
+)
+def op_name(...):
+    ...
+```
+
+**Parameters:**
+
+- `activations`: List of parameter names considered "activations" (typically consumed by `maybe_inplace`). Defaults to parameters starting with `x`.
+- `allow_inplace`: Creates a `maybe_inplace` overload for memory-efficient execution (see below).
+
+### The `maybe_inplace` Overload
+
+The `maybe_inplace` overload is a critical feature for memory efficiency in LLM inference.
+It signals that the caller doesn't need to preserve the activation inputs after the operation,
+allowing in-place implementations to reuse input memory for outputs.
+
+#### Semantics and Usage
+
+```python
+# Standard usage: inputs are preserved
+out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
+# x and residual are unchanged, out and res_out are new tensors
+
+# maybe_inplace: inputs may be modified
+out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
+# x and residual may be modified (undefined behavior to use them after this)
+# out and res_out may alias x and residual
+```
+
+Using an activation input after passing it to `maybe_inplace` is **undefined behavior**:
+
+```python
+# WRONG: Using x after donating it
+out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
+result = out + x  # ERROR: x was donated!
+```
+
+If you need to preserve an input, either use the default overload or clone manually:
+
+```python
+# Option 1: Use default overload
+out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
+result = out + x  # OK: x is preserved
+
+# Option 2: Clone before maybe_inplace
+out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x.clone(), residual, weight, epsilon)
+result = out + x  # OK: x is preserved, clone was donated
+```
+
+#### Compilation Behavior
+
+During compilation, the inplace functionalization pass validates that donated inputs are
+not used again and converts `maybe_inplace` to the functional `default` overload:
+
+```python
+# Inplace functionalization pass (pre-grad)
+for node in graph.nodes:
+    if node.target == torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace:
+        # Check that activation inputs aren't used after this node
+        for activation_arg in activation_inputs:
+            for user in activation_arg.users:
+                if user appears after node:
+                    raise ValueError(f"Input {activation_arg} donated but used again")
+
+        # Convert to default overload
+        node.target = torch.ops.vllm_ir.fused_add_rms_norm.default
+
+        # Track donated graph inputs for later clone elimination
+        for i, arg in enumerate(node.args):
+            if arg.op == "placeholder" and i in activation_indices:
+                pass_context.donated_input_ids.add(node_to_idx[arg])
+```
+
+The donated input information is then used by the clone cleanup pass to eliminate
+unnecessary copies when in-place kernels are lowered.
+
+#### Eager Mode Behavior
+
+In eager mode (without `torch.compile`), `maybe_inplace` enables **maximally memory-efficient**
+execution by allowing the IR operation to dispatch directly to in-place implementations:
+
+```python
+# Eager dispatch logic for maybe_inplace
+impl: IrOpImpl = ir_op.dispatch(*args)
+return impl.impl_fn(*args)
+
+# Eager dispatch logic for default:
+impl: IrOpImpl = ir_op.dispatch(*args)
+if impl.inplace:
+  args = [
+    arg.clone() if i in ir_op.activations else arg
+    for i, arg in enumerate(args)
+  ]
+return impl.impl_fn(*args)
+```
+
+The combination of `maybe_inplace` in model code and in-place kernel implementations provides optimal memory efficiency
+in both eager and compiled modes, with identical semantics in both cases.
+
+#### Memory Savings Example
+
+Consider a transformer layer with residual connections:
+
+```python
+# Without maybe_inplace (2 allocations per layer)
+hidden_states = self.attention(input)
+normed, residual = ir.ops.fused_add_rms_norm(hidden_states, input, weight, eps)
+# Memory: input (preserved), hidden_states (preserved), normed (new), residual (new)
+
+# With maybe_inplace (0 allocations per layer when using in-place kernel)
+hidden_states = self.attention(input)
+normed, residual = ir.ops.fused_add_rms_norm.maybe_inplace(hidden_states, input, weight, eps)
+# Memory: normed (reuses hidden_states), residual (reuses input)
+```
+
+### Implementation Registration
+
+Implementations are registered using the `register_impl` method:
+
+```python
+@ir.ops.op_name.register_impl(
+    provider="provider_name",  # Unique identifier (e.g., "vllm_c", "aiter", "triton")
+    supported=True,            # Static availability check
+    supports_args=None,        # Dynamic argument support check
+)
+def impl_fn(...):
+    ...
+```
+
+**Provider naming conventions:**
+
+- `native`: Reserved for the native torch implementation (declared with `@register_op`)
+- `vllm_c`: C++/CUDA kernels via `torch.ops._C`
+- `aiter`: AMD AITER library
+- `xpu_kernels`: SYCL/SYCLTLA kernels implemented in `vllm-xpu-kernels`
+- `triton_*`: Triton kernels
+- Platform/library names for other implementations
+
+**Support checking:**
+
+- `supported`: Static boolean, checked once at import time (e.g., `HAS_TRITON`, `is_cuda_alike()`)
+- `supports_args`: Function `(*args, **kwargs) -> bool` checking argument compatibility
+    - Called with **fake tensors** during compilation for zero-cost checking
+    - Called with **real tensors** during eager mode dispatch
+    - Should NOT check batch sizes or add guards based on values
+
+Example support predicate:
+
+```python
+def aiter_rms_norm_supports(x, weight, epsilon, variance_size=None):
+    # Check dtype (OK: doesn't depend on batch size)
+    if x.dtype not in [torch.float16, torch.bfloat16]:
+        return False
+    # Check optional parameter (OK: static check)
+    if variance_size is not None:
+        return False
+    return True
+
+@ir.ops.rms_norm.register_impl("aiter", supports_args=aiter_rms_norm_supports)
+def rms_norm(...):
+    ...
+```
+
+Batch-invariant kernels are automatically selected when `VLLM_BATCH_INVARIANT=1` is set.
+
+### Eager Mode vs Compile Mode
+
+vLLM IR operations behave identically in eager and compile modes:
+
+**Eager mode:**
+
+- Direct dispatch to implementation based on priority list
+- Support checked with real tensor arguments
+- Minimal overhead (can be optimized further if needed)
+
+**Compile mode:**
+
+- IR ops appear in FX graph as `torch.ops.vllm_ir.*` custom ops
+- Lowering selects implementation using fake tensors
+- Full integration with Inductor optimizations
+
+This consistency enables:
+
+- Prototyping in eager mode with confidence
+- Debugging by disabling compilation
+- Gradual migration from eager to compiled execution
+
+## Other Topics
+
+### Out-of-Tree Implementations
+
+External platforms can register implementations without modifying vLLM:
+
+```python
+# In external package
+from vllm import ir
+
+@ir.ops.rms_norm.register_impl("my_platform", supported=is_my_platform())
+def rms_norm(x, weight, epsilon, variance_size=None):
+    return my_platform.rms_norm(x, weight, epsilon)
+```
+
+Then configure priority to use your implementation:
+
+```python
+class MyPlatform(Platform):
+  def get_default_ir_op_priority(self):
+    return IrOpPriorityConfig(rms_norm=['my_platform', 'native'])
+
+# Users can still override priority in the same way
+llm = LLM(ir_op_priority=IrOpPriorityConfig(rms_norm=['custom_oot_kernel']))
+```
+
+### Debugging and Observability
+
+!!! note
+    Please let us know how observability can be improved for your use-case!
+
+Enable debug logging to see kernel selection:
+
+```bash
+VLLM_LOGGING_LEVEL=DEBUG vllm serve ...
+```
+
+This logs:
+
+- Which implementations are selected for each operation
+- Why implementations were rejected (unsupported, args not supported)
+- Compilation cache hits/misses
+- IR lowering statistics
+
+Check selected implementations in compiled graphs:
+
+```python
+# After compilation, inspect the lowering pass
+lowering_pass = backend.lowering_pass
+print(lowering_pass.selected_impls)
+# Output: {'rms_norm': {'node_123': 'vllm_c', 'node_456': 'vllm_c'}}
+```
+
+## Migration from CustomOp
+
+vLLM IR is designed to coexist with and gradually replace `CustomOp`:
+
+1. **Op declaration**: Convert `CustomOp` class `PluggableLayer` and move `forward_native` to `@register_op` function
+2. **Implementation registration**: Use `@ir.ops.op_name.register_impl` instead of overriding methods
+3. **Layer usage**: Replace `self.op(...)` with `ir.ops.op_name(...)`
+4. **Configuration**: Migrate `--compilation-config.custom-ops` to `--ir-op-priority`
+
+The migration can be done incrementally, one operation at a time.
+
+## See Also
+
+- [torch.compile Integration](torch_compile.md) - General compilation infrastructure
+- [Fusions](fusions.md) - Custom fusion and transformation passes in vLLM
+- [Custom Operations](custom_op.md) - Legacy custom op system

docs/mkdocs/hooks/generate_argparse.py

@@ -7,7 +7,7 @@ import sys
 import textwrap
 import traceback
 from argparse import SUPPRESS, Action, HelpFormatter
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from importlib.machinery import ModuleSpec
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
@@ -48,6 +48,7 @@ class MockPluggableLayer:
 
 
 mock_if_no_torch("vllm._C", MagicMock())
+mock_if_no_torch("vllm._C_stable_libtorch", MagicMock())
 mock_if_no_torch(
     "vllm.model_executor.custom_op",
     MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
@@ -67,6 +68,31 @@ importlib.metadata.version = lambda name: VERSIONS.get(name) or "0.0.0"
 mock_if_no_torch("torch.nn", MagicMock(Parameter=object))
 
 
+# Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__
+# We need to return the corresponding number of inputs, as IR infra will assert it
+def get_outputs(native_fn: Callable) -> str:
+    """
+    Extract output schema from function's return type annotation,
+    e.g. 'Tensor' or 'Tensor, Tensor'.
+    """
+    import typing
+
+    return_type = typing.get_type_hints(native_fn)["return"]
+    origin = typing.get_origin(return_type)
+    arg_name = lambda a: a.__name__ if hasattr(a, "__name__") else str(a)
+    if origin is tuple:
+        args = typing.get_args(return_type)
+        return ", ".join(arg_name(arg) for arg in args)
+    else:
+        return f"{arg_name(return_type)}"
+
+
+mock_if_no_torch(
+    "torch.library",
+    MagicMock(infer_schema=lambda fn, **k: f"(Tensor x) -> {get_outputs(fn)}"),
+)
+
+
 class PydanticMagicMock(MagicMock):
     """`MagicMock` that's able to generate pydantic-core schemas."""
 

tests/compile/backend.py

@@ -12,10 +12,17 @@ from torch._ops import OpOverload, OpOverloadPacket
 from torch.fx._utils import lazy_format_graph_code
 
 from vllm.compilation.passes.fx_utils import find_op_nodes
-from vllm.compilation.passes.inductor_pass import InductorPass
+from vllm.compilation.passes.inductor_pass import (
+    InductorPass,
+    pass_context,
+)
+from vllm.compilation.passes.ir.inplace_functionalization import (
+    VllmIRInplaceFunctionalizationPass,
+)
 from vllm.compilation.passes.pass_manager import with_pattern_match_debug
 from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
 from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.config.utils import Range
 from vllm.logger import init_logger
 
 logger = init_logger("vllm.tests.compile.backend")
@@ -53,11 +60,17 @@ class TestBackend:
         self.custom_passes = list(passes)
         vllm_config = get_current_vllm_config()
         compile_config = vllm_config.compilation_config
+        self.range = Range(1, vllm_config.scheduler_config.max_num_batched_tokens)
         # Deepcopy to allow multiple TestBackend instances to use the same VllmConfig
         self.inductor_config = deepcopy(compile_config.inductor_compile_config)
         self.inductor_config["force_disable_caches"] = True
         self.inductor_config["post_grad_custom_post_pass"] = self.post_pass
 
+        # Add VllmIRInplaceFunctionalizationPass as pre-grad pass by default
+        self.inductor_config["pre_grad_custom_pass"] = (
+            VllmIRInplaceFunctionalizationPass(vllm_config)
+        )
+
         if debug_dump_path := vllm_config.compile_debug_dump_path():
             logger.debug("Dumping depyf output to %s", debug_dump_path)
             self.debug_ctx = depyf.prepare_debug(debug_dump_path.as_posix())
@@ -68,7 +81,7 @@ class TestBackend:
         self.graph_pre_compile = deepcopy(graph)
         from torch._inductor.compile_fx import compile_fx
 
-        with self.debug_ctx:
+        with self.debug_ctx, pass_context(self.range):
             return compile_fx(
                 graph, example_inputs, config_patches=self.inductor_config
             )

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -88,14 +88,10 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
         ]
 
     def ops_in_model(self):
-        return (
-            [torch.ops.vllm_ir.rms_norm]
-            + [
-                torch.ops._C.fused_add_rms_norm.default,
-            ]
-            if RMSNorm.enabled()
-            else []
-        )
+        return [
+            torch.ops.vllm_ir.rms_norm,
+            torch.ops.vllm_ir.fused_add_rms_norm,
+        ]
 
 
 class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
@@ -152,16 +148,17 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
     def ops_in_model(self):
         if self.vllm_config.compilation_config.pass_config.fuse_norm_quant:
             return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
-        elif RMSNorm.enabled():
-            return [
-                torch.ops._C.fused_add_rms_norm.default,
-            ]
-        elif any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers):
-            return [
-                torch.ops._C.static_scaled_fp8_quant.default,
-            ]
         else:
-            return []
+            quant_ops = (
+                [torch.ops._C.static_scaled_fp8_quant.default]
+                if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers)
+                else [torch.ops.aten.reciprocal]
+            )
+            return [
+                torch.ops.vllm_ir.rms_norm,
+                torch.ops.vllm_ir.fused_add_rms_norm,
+                *quant_ops,
+            ]
 
 
 @multi_gpu_test(num_gpus=2)

tests/compile/passes/ir/test_clone_cleanup.py

@@ -0,0 +1,412 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Comprehensive tests for UnsafeCloneEliminationPass.
+
+This test suite exercises all possible valid FX graph patterns involving clones:
+1. Clone with no users (dead code)
+2. Clone with read-only users
+3. Clone with mutation users
+4. Clone of graph input
+5. Clone with original used after mutation
+6. Clone chains
+"""
+
+import pytest
+import torch
+from torch import fx
+from torch.fx.experimental.proxy_tensor import make_fx
+
+from vllm.compilation.passes.fx_utils import find_op_nodes
+from vllm.compilation.passes.inductor_pass import get_pass_context, pass_context
+from vllm.compilation.passes.ir.clone_elimination import (
+    UnsafeCloneEliminationPass,
+    user_writes_to_node,
+)
+from vllm.config import VllmConfig
+from vllm.config.utils import Range
+
+
+def count_clones(graph: fx.Graph) -> int:
+    """Count clone nodes in a graph."""
+    return len(list(find_op_nodes(torch.ops.aten.clone.default, graph)))
+
+
+@pytest.fixture(scope="function")
+def clone_cleanup_pass():
+    return UnsafeCloneEliminationPass(VllmConfig())
+
+
+@pytest.fixture(autouse=True)
+def setup_pass_context():
+    """Set up pass context for each test."""
+    with pass_context(compile_range=Range(1, 8192)):
+        yield
+
+
+class TestCloneCleanup:
+    """Test UnsafeCloneEliminationPass behavior on various graph patterns."""
+
+    def test_remove_clone_readonly_users(self, clone_cleanup_pass):
+        """Clone with only read-only users should be removed."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            x_clone = x.clone()
+            return x_clone + 1
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 1
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        assert count_clones(graph_module.graph) == 0
+        torch.testing.assert_close(actual, expected)
+
+    def test_keep_clone_with_mutation_and_original_used_after(self, clone_cleanup_pass):
+        """Clone must be kept if it's mutated AND original is used after mutation."""
+
+        def f(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            x = x.relu()  # not a graph param
+            x_clone = x.clone()
+            x_clone.add_(1)
+            return x, x_clone
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 1
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        # Clone should be KEPT because original is used after mutation
+        assert count_clones(graph_module.graph) == 1
+        torch.testing.assert_close(actual[0], expected[0])
+        torch.testing.assert_close(actual[1], expected[1])
+
+    def test_remove_clone_with_mutation_no_original_use(self, clone_cleanup_pass):
+        """Clone can be removed if it's mutated but original is not used after."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            x = x.relu()  # not a graph param
+            x_clone = x.clone()
+            x_clone.add_(1)
+            return x_clone
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 1
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        assert count_clones(graph_module.graph) == 0
+        torch.testing.assert_close(actual, expected)
+
+    def test_clone_chain(self, clone_cleanup_pass):
+        """Test handling of clone chains: x -> clone1 -> clone2."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            x = x.relu()  # not a graph param
+            x1 = x.clone()
+            x2 = x1.clone()
+            return x2 + 1
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 2
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        # Both clones should be removed
+        assert count_clones(graph_module.graph) == 0
+        torch.testing.assert_close(actual, expected)
+
+    def test_multiple_clones_of_same_input(self, clone_cleanup_pass):
+        """Test multiple independent clones of the same input."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            x1 = x.clone()
+            x2 = x.clone()
+            return x1 + x2
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 2
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        # Both clones should be removed (only readonly uses)
+        assert count_clones(graph_module.graph) == 0
+        torch.testing.assert_close(actual, expected)
+
+    def test_no_clones_in_graph(self, clone_cleanup_pass):
+        """Test pass behavior when graph has no clones."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            return x + 1
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 0
+
+        expected = graph_module(inp)
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+        actual = graph_module(inp)
+
+        assert count_clones(graph_module.graph) == 0
+        torch.testing.assert_close(actual, expected)
+
+    def test_multiple_passes(self, clone_cleanup_pass):
+        """Test running the pass multiple times (should be idempotent)."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            x1 = x.clone()
+            return x1 + 1
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 1
+
+        expected = graph_module(inp)
+
+        clone_cleanup_pass(graph_module.graph)
+        assert count_clones(graph_module.graph) == 0
+        graph_module.recompile()
+        actual = graph_module(inp)
+        torch.testing.assert_close(actual, expected)
+
+        clone_cleanup_pass(graph_module.graph)
+        assert count_clones(graph_module.graph) == 0
+        graph_module.recompile()
+        actual = graph_module(inp)
+        torch.testing.assert_close(actual, expected)
+
+    def test_output_node_no_write(self):
+        """Output nodes never write to their inputs."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            return x
+
+        graph_module = make_fx(f)(torch.randn(2, 3))
+        x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0]
+        output_node = [n for n in graph_module.graph.nodes if n.op == "output"][0]
+
+        assert not user_writes_to_node(output_node, x_node)
+
+    def test_readonly_op_no_write(self):
+        """Readonly operations don't write to inputs."""
+
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            return x + y
+
+        graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3))
+        placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
+        add_node = [
+            n
+            for n in graph_module.graph.nodes
+            if n.op == "call_function" and n.target == torch.ops.aten.add.Tensor
+        ][0]
+
+        assert not user_writes_to_node(add_node, placeholders[0])
+        assert not user_writes_to_node(add_node, placeholders[1])
+
+    def test_inplace_op_writes(self):
+        """Inplace operations write to first argument."""
+
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            x.add_(y)
+            return x
+
+        graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3))
+        placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
+        add_node = [
+            n
+            for n in graph_module.graph.nodes
+            if n.op == "call_function" and "add_" in str(n.target)
+        ][0]
+
+        # add_ writes to first arg but not second
+        assert user_writes_to_node(add_node, placeholders[0])
+        assert not user_writes_to_node(add_node, placeholders[1])
+
+    def test_copy_writes(self):
+        """copy_ operation writes to first argument."""
+
+        def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+            x.copy_(y)
+            return x
+
+        graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3))
+        placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
+        copy_node = [
+            n
+            for n in graph_module.graph.nodes
+            if n.op == "call_function" and "copy_" in str(n.target)
+        ][0]
+
+        assert user_writes_to_node(copy_node, placeholders[0])
+        assert not user_writes_to_node(copy_node, placeholders[1])
+
+    def test_auto_functionalized_not_a_write(self):
+        """auto_functionalized ops are follow-up uses, not writes."""
+        from torch._higher_order_ops.auto_functionalize import auto_functionalized
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            return x
+
+        graph_module = make_fx(f)(torch.randn(2, 3))
+        x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0]
+
+        # Create an auto_functionalized node in the graph
+        with graph_module.graph.inserting_before(None):
+            af_node = graph_module.graph.call_function(
+                auto_functionalized, kwargs={"input": x_node}
+            )
+
+        # auto_functionalized should not be treated as a write
+        assert not user_writes_to_node(af_node, x_node)
+
+    def test_higher_order_op_conservatively_writes(self):
+        """Other higher-order operators are conservatively treated as writes."""
+        from torch._ops import HigherOrderOperator
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            return x
+
+        graph_module = make_fx(f)(torch.randn(2, 3))
+        x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0]
+
+        # Create a concrete higher-order operator subclass
+        class MockHigherOrderOp(HigherOrderOperator):
+            def __call__(self, *args, **kwargs):
+                return args[0] if args else None
+
+        mock_hoo = MockHigherOrderOp("mock_higher_order_op")
+
+        with graph_module.graph.inserting_before(None):
+            hoo_node = graph_module.graph.call_function(mock_hoo, args=(x_node,))
+
+        # Should be conservative and assume it could write
+        assert user_writes_to_node(hoo_node, x_node)
+
+
+class TestCloneCleanupWithDonatedInputs:
+    """Test UnsafeCloneEliminationPass with donated input tracking via PassContext."""
+
+    @pytest.fixture(autouse=True)
+    def setup_pass_context(self):
+        """Set up pass context for each test."""
+        with pass_context(compile_range=Range(1, 8192)):
+            yield
+
+    def test_donated_input_clone_removed(self, clone_cleanup_pass):
+        """Clone of donated input should be removed."""
+
+        def f(x: torch.Tensor) -> torch.Tensor:
+            x_clone = x.clone()
+            x_clone.add_(1)
+            return x_clone
+
+        inp = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp)
+        assert count_clones(graph_module.graph) == 1
+
+        # Mark first parameter as donated
+        get_pass_context().donated_input_ids = {0}
+
+        expected = graph_module(inp.clone())
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+
+        # Clone should be removed since input is donated
+        assert count_clones(graph_module.graph) == 0
+
+        # Input can be mutated (donated)
+        inp_copy = inp.clone()
+        actual = graph_module(inp_copy)
+        torch.testing.assert_close(actual, expected)
+
+    def test_non_donated_input_clone_kept(self, clone_cleanup_pass):
+        """Clone of non-donated input with mutation should be kept."""
+
+        def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            x_clone = x.clone()
+            x_clone.add_(1)
+            return x, x_clone
+
+        inp_x = torch.randn(2, 3)
+        inp_y = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp_x, inp_y)
+        assert count_clones(graph_module.graph) == 1
+
+        # No donated inputs
+        get_pass_context().donated_input_ids = set()
+
+        expected = graph_module(inp_x.clone(), inp_y.clone())
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+
+        # Clone should be kept since input is not donated and original is used
+        assert count_clones(graph_module.graph) == 1
+
+        # Verify inputs are not mutated
+        inp_x_before = inp_x.clone()
+        inp_y_before = inp_y.clone()
+        actual = graph_module(inp_x, inp_y)
+        torch.testing.assert_close(
+            inp_x, inp_x_before, msg="Input x should not be mutated"
+        )
+        torch.testing.assert_close(
+            inp_y, inp_y_before, msg="Input y should not be mutated"
+        )
+        torch.testing.assert_close(actual[0], expected[0])
+        torch.testing.assert_close(actual[1], expected[1])
+
+    def test_mixed_donated_inputs(self, clone_cleanup_pass):
+        """Test with some inputs donated and some not."""
+
+        def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+            x_clone = x.clone()
+            x_clone.add_(1)
+            y_clone = y.clone()
+            y_clone.add_(2)
+            return x_clone, y_clone
+
+        inp_x = torch.randn(2, 3)
+        inp_y = torch.randn(2, 3)
+        graph_module = make_fx(f)(inp_x, inp_y)
+        assert count_clones(graph_module.graph) == 2
+
+        # Only x is donated
+        get_pass_context().donated_input_ids = {0}
+
+        expected = graph_module(inp_x.clone(), inp_y.clone())
+        clone_cleanup_pass(graph_module.graph)
+        graph_module.recompile()
+
+        # x_clone removed (x is donated), y_clone kept (y is not donated)
+        assert count_clones(graph_module.graph) == 1
+
+        # Verify y is not mutated (x can be mutated since it's donated)
+        inp_y_before = inp_y.clone()
+        actual = graph_module(inp_x.clone(), inp_y)
+        torch.testing.assert_close(
+            inp_y, inp_y_before, msg="Input y should not be mutated"
+        )
+        torch.testing.assert_close(actual[0], expected[0])
+        torch.testing.assert_close(actual[1], expected[1])

tests/compile/passes/ir/test_inplace_functionalization.py

@@ -0,0 +1,465 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for IR inplace functionalization pass integration.
+
+This test suite verifies that the inplace functionalization pass, lowering pass,
+and clone cleanup pass work together correctly with donated buffer tracking.
+"""
+
+from collections.abc import Callable
+
+import pytest
+import torch
+import torch._dynamo.exc
+from torch import nn
+
+import vllm.kernels  # noqa: F401 to register kernels
+from vllm.compilation.passes.inductor_pass import InductorPass, get_pass_context
+from vllm.compilation.passes.ir.clone_elimination import (
+    UnsafeCloneEliminationPass,
+)
+from vllm.compilation.passes.ir.inplace_functionalization import (
+    VllmIRInplaceFunctionalizationPass,
+)
+from vllm.compilation.passes.ir.lowering_pass import VllmIRLoweringPass
+from vllm.config import VllmConfig
+from vllm.ir import ops
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+from ...backend import TestBackend
+
+
+class StoreDonationInfoPass(InductorPass):
+    def __init__(self):
+        self.donated_input_ids_sets: list[set[int]] = []
+
+    def __call__(self, *args, **kwargs):
+        ctx = get_pass_context()
+        self.donated_input_ids_sets += [ctx.donated_input_ids]
+
+
+class MaybeInplaceModel(nn.Module):
+    """Model using only maybe_inplace variants."""
+
+    def __init__(self, hidden_size=16):
+        super().__init__()
+        self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+        self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+
+    def forward(
+        self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor
+    ):
+        # First maybe_inplace - x & residual1 are donated
+        x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace(
+            x, residual1, self.weight1, 1e-5
+        )
+        # Second maybe_inplace - residual2 is donated
+        x_normed2, residual_out2 = ops.fused_add_rms_norm.maybe_inplace(
+            x_normed1, residual2, self.weight2, 1e-5
+        )
+        return x_normed2, residual_out1, residual_out2
+
+
+class FunctionalModel(nn.Module):
+    """Model using only functional (default) variants."""
+
+    def __init__(self, hidden_size=16):
+        super().__init__()
+        self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+        self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+
+    def forward(
+        self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor
+    ):
+        # First functional - no donation
+        x_normed1, residual_out1 = ops.fused_add_rms_norm(
+            x, residual1, self.weight1, 1e-5
+        )
+        # Second functional - no donation
+        x_normed2, residual_out2 = ops.fused_add_rms_norm(
+            x_normed1, residual2, self.weight2, 1e-5
+        )
+        return x_normed2, residual_out1, residual_out2
+
+
+class MixedModel(nn.Module):
+    """Model mixing maybe_inplace and functional variants."""
+
+    def __init__(self, hidden_size=16):
+        super().__init__()
+        self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+        self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+
+    def forward(
+        self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor
+    ):
+        # First maybe_inplace - x & residual1 are donated
+        x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace(
+            x, residual1, self.weight1, 1e-5
+        )
+        # Second functional - no donation, x_normed1 must be preserved as it's returned
+        x_normed2, residual_out2 = ops.fused_add_rms_norm(
+            x_normed1, residual2, self.weight2, 1e-5
+        )
+        # Return both to prevent x_normed1 from being optimized away
+        return x_normed1, x_normed2, residual_out1, residual_out2
+
+
+class ModelWithTritonAfterMaybeInplace(nn.Module):
+    """
+    Model using maybe_inplace followed by a Triton kernel.
+    Test clone elimination can handle Triton in the graph
+    """
+
+    def __init__(self, hidden_size=16):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+
+        @triton.jit
+        def _triton_add_kernel(
+            x_ptr,
+            y_ptr,
+            n_elements,
+            BLOCK_SIZE: tl.constexpr,
+        ):
+            pid = tl.program_id(axis=0)
+            block_start = pid * BLOCK_SIZE
+            offsets = block_start + tl.arange(0, BLOCK_SIZE)
+            mask = offsets < n_elements
+            x = tl.load(x_ptr + offsets, mask=mask)
+            y = x + 0.1
+            tl.store(y_ptr + offsets, y, mask=mask)
+
+        def triton_add(x: torch.Tensor) -> torch.Tensor:
+            """Simple Triton add kernel."""
+            y = torch.empty_like(x)
+            n_elements = x.numel()
+            grid = (triton.cdiv(n_elements, 256),)
+            _triton_add_kernel[grid](x, y, n_elements, BLOCK_SIZE=256)
+            return y
+
+        self.triton_add = triton_add
+
+    def forward(self, x: torch.Tensor, residual: torch.Tensor, residual2: torch.Tensor):
+        x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace(
+            x, residual, self.weight, 1e-5
+        )
+
+        x_processed = self.triton_add(x_normed)
+
+        # x_processed does not need to be cloned, residual2 does
+        x_normed2, residual_out2 = ops.fused_add_rms_norm(
+            x_processed, residual2, self.weight, 1e-5
+        )
+        return x_normed2, residual_out2
+
+
+skipif_no_triton = pytest.mark.skipif(not HAS_TRITON, reason="Requires Triton")
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on cuda and rocm platform",
+)
+@pytest.mark.parametrize(
+    "model_class,expected_functionalized,expected_donated,expected_clones",
+    [
+        # 2 inplace calls, all activations donated, all clones eliminated
+        (MaybeInplaceModel, 2, 3, 0),
+        # No inplace calls, no donations, 3 clones (one eliminated)
+        (FunctionalModel, 0, 0, 3),
+        # One inplace call, two donated activations, 2 clones
+        (MixedModel, 1, 2, 2),
+        # One inplace call, two donated, 1 clone remaining
+        pytest.param(ModelWithTritonAfterMaybeInplace, 1, 2, 1, marks=skipif_no_triton),
+    ],
+)
+def test_inplace_functionalization(
+    default_vllm_config: VllmConfig,
+    model_class,
+    expected_functionalized: int,
+    expected_clones: int,
+    expected_donated: int,
+):
+    """Test inplace functionalization, lowering, and clone cleanup."""
+    torch.set_default_device(current_platform.device_type)
+
+    # Use vllm_c so inplace path is triggered
+    default_vllm_config.kernel_config.ir_op_priority.fused_add_rms_norm = [
+        "vllm_c",
+        "native",
+    ]
+
+    # Create passes in order they run during compilation
+    functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config)
+    lowering_pass = VllmIRLoweringPass(default_vllm_config)
+    donated_info_pass = StoreDonationInfoPass()
+    cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config)
+
+    # Set up backend with pre-grad pass
+    backend = TestBackend(lowering_pass, donated_info_pass, cleanup_pass)
+    backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass
+
+    model = model_class()
+    x = torch.randn(8, 16, dtype=torch.bfloat16)
+    residual1 = torch.randn(8, 16, dtype=torch.bfloat16)
+    residual2 = torch.randn(8, 16, dtype=torch.bfloat16)
+
+    with default_vllm_config.kernel_config.ir_op_priority.set_priority():
+        # Reference output without optimization
+        ref_output = model(x.clone(), residual1.clone(), residual2.clone())
+
+        # Compile with inplace optimization
+        compiled_model = torch.compile(model, backend=backend, fullgraph=True)
+        output = compiled_model(x.clone(), residual1.clone(), residual2.clone())
+
+    # Verify correctness (relaxed tolerance for bfloat16)
+    for i in range(len(ref_output)):
+        torch.testing.assert_close(output[i], ref_output[i], rtol=1e-2, atol=1e-2)
+
+    # Verify expected number of ops were functionalized
+    func_ops = functionalization_pass.functionalized_ops
+    assert len(func_ops) == int(bool(expected_functionalized))
+    if expected_functionalized > 0:
+        assert "fused_add_rms_norm" in func_ops
+        assert func_ops["fused_add_rms_norm"] == expected_functionalized
+
+    # Verify lowering happened (2 ops in all cases)
+    assert "fused_add_rms_norm" in lowering_pass.selected_impls
+    assert len(lowering_pass.selected_impls["fused_add_rms_norm"]) == 2
+    assert all(
+        provider == "vllm_c"
+        for node, provider in lowering_pass.selected_impls["fused_add_rms_norm"].items()
+    ), lowering_pass.selected_impls
+
+    # Verify correct number of donated IDs
+    assert len(donated_info_pass.donated_input_ids_sets) == 1
+    assert len(donated_info_pass.donated_input_ids_sets[0]) == expected_donated
+
+    # Verify expected number of clones after cleanup
+    actual_clones = backend.op_count(torch.ops.aten.clone.default, before=False)
+    assert actual_clones == expected_clones, (
+        f"Expected {expected_clones} clones, got {actual_clones}:"
+        f"{backend.print_graphs()}"
+    )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on cuda and rocm platform",
+)
+def test_donated_buffer_context_propagation(default_vllm_config):
+    """Test that donated_input_ids propagates correctly through pass_context."""
+    torch.set_default_device(current_platform.device_type)
+
+    # Create a custom backend that inspects pass_context in cleanup pass
+    functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config)
+    lowering_pass = VllmIRLoweringPass(default_vllm_config)
+
+    donation_info_pass = StoreDonationInfoPass()
+    cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config)
+
+    backend = TestBackend(lowering_pass, donation_info_pass, cleanup_pass)
+    backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass
+
+    model = MaybeInplaceModel()
+    x = torch.randn(8, 16, dtype=torch.bfloat16)
+    residual1 = torch.randn(8, 16, dtype=torch.bfloat16)
+    residual2 = torch.randn(8, 16, dtype=torch.bfloat16)
+
+    compiled_model = torch.compile(model, backend=backend, fullgraph=True)
+    compiled_model(x.clone(), residual1.clone(), residual2.clone())
+
+    donated_ids_seen = donation_info_pass.donated_input_ids_sets
+    # Verify donated_input_ids was set and propagated
+    assert len(donated_ids_seen) == 1
+    # Should have donated inputs (exact indices depend on AOTAutograd)
+    assert len(donated_ids_seen[0]) == 3
+    # All donated ids should be valid non-negative integers
+    for idx in donated_ids_seen[0]:
+        assert isinstance(idx, int) and idx >= 0, f"Invalid donated index: {idx}"
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on cuda and rocm platform",
+)
+def test_maybe_inplace_reuse_error(default_vllm_config):
+    """Test that reusing a donated activation input raises ValueError."""
+    torch.set_default_device(current_platform.device_type)
+
+    class ReuseModel(nn.Module):
+        """Model that incorrectly reuses a donated activation input."""
+
+        def __init__(self, hidden_size=16):
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+
+        def forward(self, x: torch.Tensor, residual: torch.Tensor):
+            # x is donated to maybe_inplace
+            x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace(
+                x, residual, self.weight, 1e-5
+            )
+            # ERROR: x is used again after being donated
+            return x_normed + x  # This should raise ValueError
+
+    functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config)
+    lowering_pass = VllmIRLoweringPass(default_vllm_config)
+    cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config)
+
+    backend = TestBackend(lowering_pass, cleanup_pass)
+    backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass
+
+    model = ReuseModel()
+    x = torch.randn(8, 16, dtype=torch.bfloat16)
+    residual = torch.randn(8, 16, dtype=torch.bfloat16)
+
+    # Compilation should raise BackendCompilerFailed wrapping ValueError
+    with pytest.raises(
+        torch._dynamo.exc.BackendCompilerFailed,
+        match="is used again after the node",
+    ):
+        compiled_model = torch.compile(model, backend=backend, fullgraph=True)
+        compiled_model(x.clone(), residual.clone())
+
+
+# Piecewise compilation tests with graph splitting
+
+
+@torch.library.custom_op("vllm::test_split_marker", mutates_args=())
+def test_split_marker(x: torch.Tensor) -> torch.Tensor:
+    """Identity op that marks a split point for piecewise compilation."""
+    return x.clone()
+
+
+@test_split_marker.register_fake
+def _fake_split_marker(x: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(x)
+
+
+class TransformerBlockWithSplits(nn.Module):
+    """Transformer block with explicit split points for piecewise compilation."""
+
+    def __init__(self, hidden_size=32, intermediate_size=128):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+
+        # Attention-like projection
+        self.attn_proj = nn.Linear(
+            hidden_size, hidden_size, bias=False, dtype=torch.bfloat16
+        )
+
+        # Post-attention norm
+        self.post_attn_norm = nn.Parameter(
+            torch.ones(hidden_size, dtype=torch.bfloat16)
+        )
+
+        # MLP
+        self.gate_proj = nn.Linear(
+            hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16
+        )
+        self.up_proj = nn.Linear(
+            hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16
+        )
+        self.down_proj = nn.Linear(
+            intermediate_size, hidden_size, bias=False, dtype=torch.bfloat16
+        )
+
+        # Post-MLP norm
+        self.post_mlp_norm = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+
+    def forward(self, x: torch.Tensor):
+        # Attention block with residual
+        residual1 = x
+        attn_out = self.attn_proj(x)
+
+        # Fused add + norm (maybe_inplace: residual1 is donated)
+        normed1, residual1 = ops.fused_add_rms_norm.maybe_inplace(
+            attn_out, residual1, self.post_attn_norm, 1e-5
+        )
+
+        # Force a graph split here
+        normed1 = torch.ops.vllm.test_split_marker(normed1)
+
+        # MLP block
+        gate = self.gate_proj(normed1)
+        up = self.up_proj(normed1)
+        mlp_out = self.down_proj(gate * torch.nn.functional.silu(up))
+
+        # Fused add + norm (maybe_inplace: residual1 is donated)
+        normed2, residual2 = ops.fused_add_rms_norm.maybe_inplace(
+            mlp_out, residual1, self.post_mlp_norm, 1e-5
+        )
+
+        return normed2, residual2
+
+
+def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int):
+    def inner(*args):
+        torch._dynamo.mark_dynamic(args[arg_index], dim_index)
+        return fn(*args)
+
+    return inner
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on cuda and rocm platform",
+)
+def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache):
+    """
+    Test piecewise compilation with donated buffers across graph splits.
+    Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching.
+    """
+    torch.set_default_device(current_platform.device_type)
+
+    # Disable compilation cache to avoid serialization issues
+    monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
+
+    from vllm.compilation.backends import VllmBackend
+    from vllm.config import CompilationConfig, VllmConfig
+
+    # Create config with custom splitting op
+    store_donation_info = StoreDonationInfoPass()
+    vllm_config = VllmConfig(
+        compilation_config=CompilationConfig(
+            custom_ops=["all"],
+            splitting_ops=["vllm::test_split_marker"],
+            inductor_compile_config={"post_grad_custom_post_pass": store_donation_info},
+        )
+    )
+
+    backend = VllmBackend(vllm_config)
+
+    model = TransformerBlockWithSplits()
+    x = torch.randn(8, 32, dtype=torch.bfloat16)
+
+    # Reference output
+    ref_output = with_dyn_arg(model, 0, 0)(x.clone())
+
+    # Compile with piecewise compilation (graph will split at split_marker)
+    compiled_model = torch.compile(model, backend=backend, fullgraph=False)
+    output = with_dyn_arg(compiled_model, 0, 0)(x.clone())
+
+    # Verify correctness (relaxed tolerance for bfloat16)
+    torch.testing.assert_close(output[0], ref_output[0], rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(output[1], ref_output[1], rtol=1e-2, atol=1e-2)
+
+    # Verify the model was split into multiple submodules
+    assert hasattr(backend, "split_gm"), "Backend should have split graph module"
+
+    # Should have at least 2 submodules (split by test_split_marker op)
+    submodules = list(backend.split_gm.named_children())
+    num_submodules = len(submodules)
+    assert num_submodules >= 2, (
+        f"Expected at least 2 submodules (split), got {num_submodules}"
+    )
+
+    # Check that donation info was propagated correctly
+    donated_inputs_sets = store_donation_info.donated_input_ids_sets
+    assert len(donated_inputs_sets) == 2
+    assert len(donated_inputs_sets[0]) == 1
+    assert len(donated_inputs_sets[1]) == 1

tests/compile/passes/test_functionalization.py

@@ -126,7 +126,7 @@ class TestFusedAddRMSNorm(torch.nn.Module):
         if TEST_FP8 and do_fusion:
             return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
         else:
-            return [torch.ops._C.fused_add_rms_norm.default]
+            return []
 
     def ops_not_in_model(self):
         return []

tests/compile/passes/test_fuse_act_padding.py

@@ -59,7 +59,7 @@ class TestModel(torch.nn.Module):
 
     def ops_in_model_before(self):
         return [
-            rocm_aiter_ops.get_rmsnorm_fused_add_op(),
+            torch.ops.vllm_ir.fused_add_rms_norm,
             torch.ops.aten.constant_pad_nd,
         ]
 

tests/compile/passes/test_fusion.py

@@ -17,7 +17,6 @@ from vllm.compilation.passes.fusion.rms_quant_fusion import (
     FusedRMSQuantKey,
     RMSNormQuantFusionPass,
 )
-from vllm.compilation.passes.fx_utils import find_op_nodes
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
 from vllm.config import (
@@ -243,9 +242,10 @@ class TestModel(torch.nn.Module):
         ]
 
     def ops_in_model_before_partial(self):
-        return [torch.ops.vllm_ir.rms_norm] + (
-            [RMS_ADD_OP] if self.enable_rms_norm_custom_op else [torch.ops.aten.rsqrt]
-        )
+        return [
+            torch.ops.vllm_ir.rms_norm,
+            torch.ops.vllm_ir.fused_add_rms_norm.default,
+        ]
 
 
 def _run_fusion_test(
@@ -383,17 +383,6 @@ def test_fusion_rmsnorm_quant(
             model.ops_in_model_before_partial(), fully_replaced=False
         )
 
-        # If RMSNorm custom op is disabled (native/torch impl used),
-        # there's a risk that the fused add doesn't get included in the
-        # replacement and only the rms part gets fused with quant.
-        # Hence, we check only 2 add nodes are left (final fused rmsnorm add).
-        if not enable_rms_norm_custom_op:
-            n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))
-            # rms_norm is IR, not included
-            # 6 = 3x2 (3xRMS_ADD, 2 each)
-            assert n_add_nodes(backend.graph_pre_pass) == 6
-            assert n_add_nodes(backend.graph_post_pass) == 2
-
 
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("hidden_size", [256])

tests/ir/test_inplace_op.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+from torch import Tensor
+from torch.fx.experimental.proxy_tensor import make_fx
+
+import vllm.ir.op
+from vllm.ir.op import IrOp, IrOpInplaceOverload
+
+
+@vllm.ir.register_op(allow_inplace=True)
+def _custom_mm2(x: Tensor, w: Tensor) -> Tensor:
+    return x @ w
+
+
+@_custom_mm2.register_impl("regular")
+def _custom_mm2_regular(x: Tensor, w: Tensor) -> Tensor:
+    return x @ w + 1
+
+
+@_custom_mm2.register_impl("inplace", inplace=True)
+def _custom_mm2_inplace(x: Tensor, w: Tensor) -> Tensor:
+    x.copy_(x @ w + 2)
+    return x
+
+
+class TestInplaceOp:
+    def test_registration(self):
+        # Test that the inplace op is registered correctly.
+        assert "_custom_mm2" in IrOp.registry
+        assert IrOp.registry["_custom_mm2"] is _custom_mm2
+        assert _custom_mm2.torch_op is torch.ops.vllm_ir._custom_mm2.default
+        assert isinstance(_custom_mm2.maybe_inplace, IrOpInplaceOverload)
+        assert (
+            _custom_mm2.maybe_inplace.torch_op
+            is torch.ops.vllm_ir._custom_mm2.maybe_inplace
+        )
+
+    def test_inplace_dispatching(self):
+        # check that the correct implementation is dispatched based on priority,
+        # and inplace semantics hold
+        w = torch.randn(3, 3)
+        x = torch.randn(2, 3)
+        x1 = x.clone()
+
+        with _custom_mm2.set_priority(["regular"]):
+            result_regular = _custom_mm2.maybe_inplace(x, w)
+
+        # check that the regular op does not modify x
+        torch.testing.assert_close(x, x1, atol=0, rtol=0)
+
+        with _custom_mm2.set_priority(["inplace"]):
+            result_inplace: Tensor = _custom_mm2.maybe_inplace(x, w)
+
+        # check that the inplace op returns x directly
+        assert result_inplace.data_ptr() == x.data_ptr()
+
+        torch.testing.assert_close(result_inplace, x1 @ w + 2)
+        torch.testing.assert_close(result_regular, x1 @ w + 1)
+
+    def test_default_dispatching(self):
+        # check that the correct implementation is dispatched,
+        # and ops do not modify inputs when using the default overload
+        w = torch.randn(3, 3)
+        x = torch.randn(2, 3)
+        x1 = x.clone()
+
+        with _custom_mm2.set_priority(["regular"]):
+            result_regular = _custom_mm2(x, w)
+
+        with _custom_mm2.set_priority(["inplace"]):
+            result_inplace = _custom_mm2(x, w)
+
+        # check that x was not modified by either impl
+        torch.testing.assert_close(x, x1, atol=0, rtol=0)
+
+        torch.testing.assert_close(result_inplace, x1 @ w + 2)
+        torch.testing.assert_close(result_regular, x1 @ w + 1)
+
+    def test_trace(self):
+        # Test that the inplace op can be used in a graph.
+        def func(x: Tensor, y: Tensor) -> Tensor:
+            return _custom_mm2.maybe_inplace(x, y)
+
+        x = torch.randn(2, 3)
+        y = torch.randn(3, 4)
+        graph = make_fx(func)(x, y)
+        assert any(
+            node.target == torch.ops.vllm_ir._custom_mm2.maybe_inplace
+            for node in graph.graph.nodes
+        )

tests/ir/test_op.py

@@ -21,7 +21,7 @@ class CustomError(Exception):
     pass
 
 
-@vllm.ir.register_op
+@vllm.ir.register_op(allow_inplace=True)
 def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
     return x + y
 
@@ -129,11 +129,15 @@ class TestIrOpCustomAdd:
 
     @pytest.mark.parametrize("enable_torch_wrap", [True, False])
     @pytest.mark.parametrize("symbolic_trace", [True, False])
+    @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
     def test_trace_sees_single_custom_op(
-        self, symbolic_trace: bool, enable_torch_wrap: bool
+        self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str
     ):
+        op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
+        torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
+
         def fn(x, y):
-            return _custom_add(x, y)
+            return op_fn(x, y)
 
         def find_fn(target: Any, gm: fx.GraphModule):
             return gm.graph.find_nodes(op="call_function", target=target)
@@ -155,7 +159,7 @@ class TestIrOpCustomAdd:
         torch.testing.assert_close(out_fx, out_eager)
 
         # check that IR nodes only appear if enable_torch_wrap=True
-        ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm)
+        ir_nodes = find_fn(torch_op, gm)
         if enable_torch_wrap:
             assert len(ir_nodes) == 1, gm.code
         else:
@@ -167,7 +171,7 @@ class TestIrOpCustomAdd:
         else:
             gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2))
 
-        ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm)
+        ir_nodes = find_fn(torch_op, gm)
         assert len(ir_nodes) == 1, gm.code
 
 
@@ -176,9 +180,12 @@ def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
     return x + y + 10
 
 
-@_custom_add.register_impl("impl_b")
+@_custom_add.register_impl("impl_b", inplace=True)
 def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-    return x + y + 20
+    """Computes x+y+20"""
+    x.add_(y)
+    x.add_(20)
+    return x
 
 
 @_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0)
@@ -243,19 +250,23 @@ class TestIrOpImplDispatch:
         # Restored to empty
         assert _custom_add.get_priority() == []
 
-    def test_dispatch_priority_order(self):
+    @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
+    def test_dispatch_priority_order(self, overload: str):
+        op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
+        torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
+
         x = torch.tensor(1, dtype=torch.int32)
         y = torch.tensor(2, dtype=torch.int32)
 
         with _custom_add.set_priority(["impl_b", "impl_a"]):
             assert _custom_add.dispatch(x, y) is impl_b
-            out1 = _custom_add(x, y)
-            out2 = torch.ops.vllm_ir._custom_add(x, y)
+            out1 = op_fn(x.clone(), y)
+            out2 = torch_op(x.clone(), y)
 
             with _custom_add.set_priority(["impl_a"]):
                 assert _custom_add.dispatch(x, y) is impl_a
-                out3 = _custom_add(x, y)
-                out4 = torch.ops.vllm_ir._custom_add(x, y)
+                out3 = op_fn(x.clone(), y)
+                out4 = torch_op(x.clone(), y)
 
         # impl_b
         assert out1.item() == 1 + 2 + 20
@@ -265,18 +276,18 @@ class TestIrOpImplDispatch:
         assert out4.item() == 1 + 2 + 10
 
     def test_unsupported_impl_filtered(self):
-        @_custom_add.register_impl("unsupported", supported=False)
-        def impl_bad(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        @_custom_add.register_impl("impl_unsupported", supported=False)
+        def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             return x + y + 999
 
         x = torch.tensor(1, dtype=torch.int32)
         y = torch.tensor(2, dtype=torch.int32)
 
-        with _custom_add.set_priority(["unsupported", "impl_a"]):
+        with _custom_add.set_priority(["impl_unsupported", "impl_a"]):
             assert _custom_add.get_priority() == ["impl_a"]
             out = _custom_add(x, y)
 
-        # impl_bad skipped → impl_a
+        # impl_unsupported skipped → impl_a
         assert out.item() == 1 + 2 + 10
 
     def test_supports_args_runtime_dispatch_and_warning(

tests/kernels/ir/test_layernorm.py

@@ -28,7 +28,9 @@ def test_rms_norm_registration():
         "native": True,
         "vllm_c": current_platform.is_cuda_alike(),
         "aiter": current_platform.is_rocm(),
-        "oink": False,
+        "oink": current_platform.has_device_capability(100)
+        and hasattr(torch.ops, "oink")
+        and hasattr(torch.ops.oink, "rmsnorm"),
         "xpu_kernels": current_platform.is_xpu(),
     }
 
@@ -67,6 +69,14 @@ class TestRMSNorm:
         out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon)
         torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3)
 
+        # Mean square should be approximately 1 (ignoring epsilon and weight scaling)
+        combined_norm = out.float() / weight.float()
+        variance = combined_norm.pow(2).mean(dim=-1)
+        # After RMS normalization, variance should be close to 1
+        torch.testing.assert_close(
+            variance, torch.ones_like(variance), rtol=1e-2, atol=1e-2
+        )
+
         # Check behavior with and without weight
         weight1 = torch.ones_like(weight)
         out3 = rms_norm_native(x, weight1, epsilon=epsilon)
@@ -129,3 +139,197 @@ def test_aiter_rejects_unsupported_dtypes():
             num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
         )
         assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
+
+
+fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
+    reason="Currently only kernels on CUDA, ROCm and XPU",
+)
+def test_fused_add_rms_norm_registration():
+    expected = {
+        "native": True,
+        "vllm_c": current_platform.is_cuda_alike(),
+        "aiter": current_platform.is_rocm(),
+        "oink": current_platform.has_device_capability(100)
+        and hasattr(torch.ops, "oink")
+        and hasattr(torch.ops.oink, "fused_add_rms_norm"),
+        "xpu_kernels": current_platform.is_xpu(),
+    }
+
+    actual = {
+        provider: impl.supported
+        for provider, impl in ir.ops.fused_add_rms_norm.impls.items()
+    }
+
+    assert actual == expected
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
+@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
+    reason="Currently only kernels on CUDA, ROCm and XPU",
+)
+class TestFusedAddRMSNorm:
+    @classmethod
+    def setup_class(cls, **kwargs):
+        torch.set_default_device(current_platform.device_type)
+
+    def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
+        x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
+            num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
+        )
+        out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)
+
+        # Check shape, dtype, device
+        assert out.shape == x.shape
+        assert out.dtype == x.dtype
+        assert out.device == x.device
+        assert residual_out.shape == x_residual.shape
+        assert residual_out.dtype == x_residual.dtype
+        assert residual_out.device == x_residual.device
+
+        # Check that residual_out = x + x_residual
+        expected_residual = (x.float() + x_residual.float()).to(dtype)
+        torch.testing.assert_close(
+            residual_out, expected_residual, rtol=1e-3, atol=1e-3
+        )
+
+        # Verify that the output is RMS normalized version of (x + x_residual)
+        expected_out = rms_norm_native(expected_residual, weight, epsilon)
+        assert_close(
+            ir.ops.fused_add_rms_norm,
+            (out, residual_out),
+            (expected_out, expected_residual),
+        )
+
+        # Check the scaling property of rms norm
+        out1, _ = fused_add_rms_norm_native(
+            x, torch.zeros_like(x), weight, epsilon=epsilon
+        )
+        out2, _ = fused_add_rms_norm_native(
+            x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon
+        )
+        torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3)
+
+        # Check behavior with and without weight
+        weight1 = torch.ones_like(weight)
+        out3, _ = fused_add_rms_norm_native(x, x_residual, weight1, eps)
+        out4, _ = fused_add_rms_norm_native(x, x_residual, None, eps)
+        torch.testing.assert_close(out3, out4)
+
+    @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
+    def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
+        impl = ir.ops.fused_add_rms_norm.impls[provider]
+        x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
+            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
+        )
+        args = (x, x_residual, weight, eps, None)
+
+        if not impl.supports_args(*args):
+            pytest.skip(f"{provider} does not support args")
+
+        ref_output, ref_residual = fused_add_rms_norm_native(*clone_args(args))
+        output, residual = impl.impl_fn(*clone_args(args))
+        assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
+        assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)
+
+        # check that dispatched call matches direct call
+        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
+            out_dispatched, residual_dispatched = ir.ops.fused_add_rms_norm(*args[:4])
+        out_direct, residual_direct = impl.impl_fn(*clone_args(args))
+        torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
+        torch.testing.assert_close(
+            residual_dispatched, residual_direct, rtol=0.0, atol=0.0
+        )
+
+        # none of these support variance_size override
+        assert not impl.supports_args(x, x_residual, weight, epsilon, 4)
+        assert not impl.supports_args(x, x_residual, weight, epsilon, variance_size=4)
+
+        # test weight=None behavior
+        out_no_weight, residual_no_weight = impl.impl_fn(
+            x.clone(), x_residual.clone(), None, epsilon
+        )
+        out_unit_weight, residual_unit_weight = impl.impl_fn(
+            x.clone(), x_residual.clone(), torch.ones_like(weight), epsilon
+        )
+        assert_close(ir.ops.fused_add_rms_norm, out_no_weight, out_unit_weight)
+        assert_close(
+            ir.ops.fused_add_rms_norm, residual_no_weight, residual_unit_weight
+        )
+
+    @pytest.mark.parametrize("provider", ["vllm_c"])
+    def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider):
+        """Test that inplace implementations reuse inputs,
+        for maybe_inplace overload but not for default overload."""
+        impl = ir.ops.fused_add_rms_norm.impls[provider]
+        if not impl.supported:
+            pytest.skip(f"{provider} impl not supported on this platform")
+
+        x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
+            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
+        )
+
+        # Test default overload - should NOT modify inputs even with inplace impl
+        x_default = x.clone()
+        x_residual_default = x_residual.clone()
+        x_default_ptr = x_default.data_ptr()
+        x_residual_default_ptr = x_residual_default.data_ptr()
+
+        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
+            out_default, residual_default = ir.ops.fused_add_rms_norm(
+                x_default, x_residual_default, weight, eps
+            )
+
+        # Default should NOT be inplace (even with inplace implementation)
+        assert out_default.data_ptr() != x_default_ptr
+        assert residual_default.data_ptr() != x_residual_default_ptr
+        torch.testing.assert_close(x, x_default, rtol=0.0, atol=0.0)
+        torch.testing.assert_close(x_residual, x_residual_default, rtol=0.0, atol=0.0)
+
+        # Test maybe_inplace overload - should modify inputs with inplace impl
+        x_inplace = x.clone()
+        x_residual_inplace = x_residual.clone()
+        x_inplace_ptr = x_inplace.data_ptr()
+        x_residual_inplace_ptr = x_residual_inplace.data_ptr()
+
+        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
+            out_inplace, residual_inplace = ir.ops.fused_add_rms_norm.maybe_inplace(
+                x_inplace, x_residual_inplace, weight, eps
+            )
+
+        # maybe_inplace should be inplace
+        assert out_inplace.data_ptr() == x_inplace_ptr
+        assert residual_inplace.data_ptr() == x_residual_inplace_ptr
+
+        # Both should produce same results
+        torch.testing.assert_close(out_default, out_inplace, atol=0.0, rtol=0.0)
+        torch.testing.assert_close(
+            residual_default, residual_inplace, atol=0.0, rtol=0.0
+        )
+
+    @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
+    def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
+        args = ir.ops.fused_add_rms_norm.generate_inputs(
+            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
+        )
+        args = args + (None,)  # Add variance_size parameter
+
+        # When checking the torch op, we have to set priority and use dispatch
+        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
+            torch.library.opcheck(torch.ops.vllm_ir.fused_add_rms_norm.default, args)
+
+            # Only test maybe_inplace with non-inplace implementations
+            # Inplace implementations return aliases of inputs which is not allowed.
+            # We break this invariant, but we also convert maybe_inplace to the default
+            # overload during compilation, so maybe_inplace never reaches Inductor.
+            if not ir.ops.fused_add_rms_norm.impls[provider].inplace:
+                torch.library.opcheck(
+                    torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace, args
+                )

tests/model_executor/test_enabled_custom_ops.py

@@ -23,11 +23,7 @@ from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
     vllm_topk_sigmoid,
     vllm_topk_softmax,
 )
-from vllm.model_executor.layers.layernorm import (
-    RMSNorm,
-    dispatch_rocm_rmsnorm_func,
-    fused_add_rms_norm,
-)
+from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.platforms import current_platform
 
 RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
@@ -153,26 +149,3 @@ def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
         assert topk_func == rocm_aiter_ops.topk_sigmoid
     else:
         assert topk_func == vllm_topk_sigmoid
-
-
-@pytest.mark.parametrize("add_residual", [False])
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("use_rocm_aiter", [True, False])
-@pytest.mark.skipif(
-    not current_platform.is_rocm(), reason="AITER is a feature exclusive for ROCm"
-)
-def test_rms_norm_dispatch(
-    add_residual: bool, dtype: torch.dtype, use_rocm_aiter: bool
-):
-    rms_norm_func = dispatch_rocm_rmsnorm_func(dtype, use_rocm_aiter)
-
-    should_use_rocm_aiter = (
-        current_platform.is_rocm()
-        and use_rocm_aiter
-        and dtype in RMS_NORM_SUPPORTED_DTYPES
-    )
-
-    if should_use_rocm_aiter:
-        assert rms_norm_func == rocm_aiter_ops.rms_norm2d_with_add
-    else:
-        assert rms_norm_func == fused_add_rms_norm

tests/model_executor/test_oink_integration.py

@@ -1,60 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
+import multiprocessing
 import types
 
 import pytest
-import torch
+
+from vllm.platforms import current_platform
 
 
-def _load_oink_ops_module():
-    # Import the module normally (vllm is installed as an editable package in CI).
-    from vllm import _oink_ops
+def _test_oink_availability_impl(
+    device_capability: tuple[int, int],
+    has_rmsnorm: bool,
+    has_fused_add_rms_norm: bool,
+    expected_available: bool,
+    expected_fused: bool,
+) -> None:
+    """Test OINK support detection with mocked state."""
+    import torch
 
-    return _oink_ops
+    from vllm import platforms
+
+    # Mock device capability (class method, override on class)
+    dc = platforms.interface.DeviceCapability(*device_capability)
+    platforms.current_platform.__class__.get_device_capability = lambda device_id=0: dc
+
+    # Mock oink ops
+    oink_ops = types.SimpleNamespace()
+    if has_rmsnorm:
+        oink_ops.rmsnorm = lambda x, w, eps: x
+    if has_fused_add_rms_norm:
+        oink_ops.fused_add_rms_norm = lambda x, residual, w, eps: None
+
+    torch.ops.oink = oink_ops
+
+    # Now import vllm modules with mocks in place (fresh import with mocked platform)
+    import vllm.kernels.oink_ops  # noqa: F401
+    from vllm.ir.ops import fused_add_rms_norm, rms_norm
+
+    # Verify support checks
+    assert rms_norm.impls["oink"].supported is expected_available
+    assert fused_add_rms_norm.impls["oink"].supported is expected_fused
 
 
-def test_oink_availability_checks(monkeypatch: pytest.MonkeyPatch):
-    _oink_ops = _load_oink_ops_module()
+@pytest.mark.parametrize(
+    "device_capability,has_rmsnorm,has_fused_add_rms_norm,expected_available,expected_fused",
+    [
+        # Case 1: < SM100, ops not supported
+        ((9, 0), True, False, False, False),
+        # Case 2: CUDA available and SM100, rmsnorm op registered
+        ((10, 0), True, False, True, False),
+        # Case 3: SM100 with both rmsnorm and fused_add_rms_norm
+        ((10, 0), True, True, True, True),
+    ],
+)
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
+def test_oink_availability_checks(
+    device_capability: tuple[int, int],
+    has_rmsnorm: bool,
+    has_fused_add_rms_norm: bool,
+    expected_available: bool,
+    expected_fused: bool,
+):
+    """Test OINK support detection with clean import state for each parameter set."""
 
-    # Ensure the ops namespace exists and is mutable for tests.
-    monkeypatch.setattr(
-        torch.ops,
-        "oink",
-        types.SimpleNamespace(rmsnorm=lambda x, w, eps: x),
-        raising=False,
-    )
-
-    # Case 1: CUDA not available.
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
-    assert _oink_ops.is_oink_available_for_device(0) is False
-
-    # Case 2: CUDA available but < SM100.
-    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (9, 0))
-    assert _oink_ops.is_oink_available_for_device(0) is False
-
-    # Case 3: CUDA available and SM100, rmsnorm op registered.
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (10, 0))
-    assert _oink_ops.is_oink_available_for_device(0) is True
-
-    # fused op presence probe
-    assert _oink_ops.has_fused_add_rms_norm() is False
-    monkeypatch.setattr(
-        torch.ops,
-        "oink",
-        types.SimpleNamespace(
-            rmsnorm=lambda x, w, eps: x,
-            fused_add_rms_norm=lambda x, residual, w, eps: None,
+    # Use spawn to run function in fresh process with clean imports
+    # TODO migrate to spawn utility:
+    # https://github.com/vllm-project/vllm/issues/41415
+    ctx = multiprocessing.get_context("spawn")
+    process = ctx.Process(
+        target=_test_oink_availability_impl,
+        args=(
+            device_capability,
+            has_rmsnorm,
+            has_fused_add_rms_norm,
+            expected_available,
+            expected_fused,
         ),
-        raising=False,
     )
-    assert _oink_ops.has_fused_add_rms_norm() is True
+    process.start()
+    process.join()
+
+    if process.exitcode != 0:
+        raise AssertionError(
+            f"Subprocess test failed with exit code {process.exitcode}"
+        )
 
 
 def test_can_view_as_2d_stride_guard():
-    # Import the helper from the layernorm module.
-    from vllm.model_executor.layers.layernorm import _can_view_as_2d
+    # No global import
+    import torch
+
+    # Import the helper from the kernels module.
+    from vllm.kernels.oink_ops import _can_view_as_2d
 
     x = torch.zeros((2, 3, 4))
     assert _can_view_as_2d(x) is True

tests/test_config.py

@@ -1295,11 +1295,14 @@ def test_ir_op_priority_default():
     # Assert default is applied to ops
     priority_config = IrOpPriorityConfig.with_default(["vllm_c", "native"])
     assert priority_config.rms_norm == ["vllm_c", "native"]
+    assert priority_config.fused_add_rms_norm == ["vllm_c", "native"]
 
     # Assert single ops override the default
-    assert IrOpPriorityConfig.with_default(
-        ["vllm_c", "native"], rms_norm=["oink", "native"]
-    ) == IrOpPriorityConfig(rms_norm=["oink", "native"])
+    priority_config = IrOpPriorityConfig.with_default(
+        ["native"], rms_norm=["oink", "native"]
+    )
+    assert priority_config.rms_norm == ["oink", "native"]
+    assert priority_config.fused_add_rms_norm == ["native"]
 
 
 def test_ir_op_priority_str():
@@ -1318,3 +1321,34 @@ def test_ir_op_priority_str():
     with pytest.raises(pydantic.ValidationError):
         # must be list of only strings
         priority_config = IrOpPriorityConfig(rms_norm=["vllm_c", 4, "native"])
+
+
+def test_ir_op_priority_ctx():
+    """Test that the priority-setting context sets priority correctly."""
+    from vllm import ir
+    from vllm.config.kernel import IrOpPriorityConfig
+
+    priority = IrOpPriorityConfig.with_default(["native"], rms_norm=["vllm_c"])
+    priority2 = IrOpPriorityConfig.with_default(
+        ["native"], fused_add_rms_norm=["vllm_c"]
+    )
+    with priority.set_priority():
+        assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
+        assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]
+        with priority2.set_priority():
+            assert ir.ops.rms_norm.get_priority() == ["native"]
+            assert ir.ops.fused_add_rms_norm.get_priority() == ["vllm_c", "native"]
+
+        # context restored
+        assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
+        assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]
+
+        with pytest.raises(ValueError), priority2.set_priority():
+            assert ir.ops.rms_norm.get_priority() == ["native"]
+            assert ir.ops.fused_add_rms_norm.get_priority() == ["vllm_c", "native"]
+
+            raise ValueError
+
+        # context restored even after exception
+        assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
+        assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]

tests/v1/determinism/test_rms_norm_batch_invariant.py

@@ -12,7 +12,7 @@ import torch
 from utils import skip_unsupported
 
 from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
-from vllm.model_executor.layers.layernorm import RMSNorm, fused_add_rms_norm
+from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.platforms import current_platform
 
 DEVICE_TYPE = current_platform.device_type
@@ -105,6 +105,12 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
         dim=0,
     )
 
+    def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]:
+        import vllm._custom_ops as ops
+
+        ops.fused_add_rms_norm(x, residual, w, e)
+        return x, residual
+
     out_single, residual_out_single = fused_add_rms_norm(
         x_single.clone(),
         residual_single.clone(),

vllm/_aiter_ops.py

@@ -647,58 +647,6 @@ def _rocm_aiter_gemm_a8w8_blockscale_fake(
     return Y
 
 
-def _rocm_aiter_rms_norm_impl(
-    x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
-) -> torch.Tensor:
-    from aiter import rms_norm
-
-    if x.dim() > 2:
-        x_original_shape = x.shape
-        x = x.reshape(-1, x_original_shape[-1])
-        x = rms_norm(x, weight, variance_epsilon)
-        return x.reshape(x_original_shape)
-
-    return rms_norm(x, weight, variance_epsilon)
-
-
-def _rocm_aiter_rms_norm_fake(
-    x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
-) -> torch.Tensor:
-    return torch.empty_like(x)
-
-
-def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
-    x: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    variance_epsilon: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    from aiter import rmsnorm2d_fwd_with_add
-
-    residual_out = torch.empty_like(residual)
-    out = torch.empty_like(x)
-    rmsnorm2d_fwd_with_add(
-        out,  # output
-        x,  # input
-        residual,  # residual input
-        residual_out,  # residual output
-        weight,
-        variance_epsilon,
-    )
-    return out, residual_out
-
-
-def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
-    x: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    variance_epsilon: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    residual_out = torch.empty_like(residual)
-    out = torch.empty_like(x)
-    return out, residual_out
-
-
 def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(
     x: torch.Tensor,
     residual: torch.Tensor,
@@ -1229,10 +1177,9 @@ class rocm_aiter_ops:
 
         # Check if aiter is enabled before using operations
         if rocm_aiter_ops.is_enabled():
-            result = rocm_aiter_ops.rms_norm(x, weight, epsilon)
+            result = rocm_aiter_ops.per_token_quant(x, FP8_DTYPE)
 
     Operations:
-        - RMS normalization: rms_norm, rms_norm2d_with_add
         - GEMM operations: gemm_a8w8, gemm_a8w8_blockscale
         - Fused MoE: fused_moe, asm_moe_tkw1
         - Routing: topk_softmax, biased_grouped_topk, grouped_topk
@@ -1244,7 +1191,6 @@ class rocm_aiter_ops:
     # Check if the env variable is set
     _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
     _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
-    _RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
     _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
     _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
     _MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
@@ -1275,7 +1221,6 @@ class rocm_aiter_ops:
         """
         cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
         cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
-        cls._RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
         cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
         cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
         cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
@@ -1367,11 +1312,6 @@ class rocm_aiter_ops:
     def is_linear_fp8_enabled(cls) -> bool:
         return cls.is_linear_enabled()
 
-    @classmethod
-    @if_aiter_supported
-    def is_rmsnorm_enabled(cls) -> bool:
-        return cls._AITER_ENABLED and cls._RMSNORM_ENABLED
-
     @classmethod
     @if_aiter_supported
     def is_fused_moe_enabled(cls) -> bool:
@@ -1560,19 +1500,6 @@ class rocm_aiter_ops:
                 fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
             )
 
-            direct_register_custom_op(
-                op_name="rocm_aiter_rms_norm",
-                op_func=_rocm_aiter_rms_norm_impl,
-                fake_impl=_rocm_aiter_rms_norm_fake,
-            )
-
-            direct_register_custom_op(
-                op_name="rocm_aiter_rmsnorm2d_fwd_with_add",
-                op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
-                fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
-                dispatch_key=current_platform.dispatch_key,
-            )
-
             direct_register_custom_op(
                 op_name="rocm_aiter_rmsnorm_fused_dynamic_quant",
                 op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl,
@@ -1672,14 +1599,6 @@ class rocm_aiter_ops:
 
             _OPS_REGISTERED = True
 
-    @staticmethod
-    def get_rmsnorm_fused_add_op() -> OpOverload:
-        return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add.default
-
-    @staticmethod
-    def get_rmsnorm_op() -> OpOverload:
-        return torch.ops.vllm.rocm_aiter_rms_norm.default
-
     @staticmethod
     def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload:
         return torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant.default
@@ -1724,23 +1643,6 @@ class rocm_aiter_ops:
     def get_fused_mla_dual_rms_norm_op() -> OpOverload:
         return torch.ops.vllm.fused_mla_dual_rms_norm.default
 
-    @staticmethod
-    def rms_norm(
-        x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
-    ) -> torch.Tensor:
-        return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon)
-
-    @staticmethod
-    def rms_norm2d_with_add(
-        x: torch.Tensor,
-        residual: torch.Tensor,
-        weight: torch.Tensor,
-        variance_epsilon: float,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add(
-            x, residual, weight, variance_epsilon
-        )
-
     @staticmethod
     def w8a8_gemm(
         A: torch.Tensor,

vllm/_oink_ops.py

@@ -1,96 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Small helper wrappers for external Oink Blackwell custom ops.
-
-vLLM does not depend on the external Oink repository/package. When an external
-plugin registers torch.library.custom_op entrypoints under the `oink::`
-namespace (e.g. via vLLM's general_plugins mechanism) and
-`VLLM_USE_OINK_OPS=1` is set, vLLM can route eligible calls to those ops.
-
-This module provides:
-- A single place to probe Oink op availability at module init time
-  (outside torch.compile tracing), and
-- Thin wrappers around the torch.ops entrypoints for use in CUDA fast paths,
-  without introducing graph breaks.
-
-Important:
-  Do not call the availability helpers in a compiled region. They may call
-  functions decorated with `torch._dynamo.disable` to safely check
-  conditions that should not be traced.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Callable
-
-import torch
-
-try:
-    from torch._dynamo import disable as _dynamo_disable  # type: ignore[attr-defined]
-except Exception:  # pragma: no cover
-
-    def _dynamo_disable(fn: Callable):  # type: ignore[misc]
-        return fn
-
-
-def _has_oink_op(op_name: str) -> bool:
-    """Check if a specific oink op is registered."""
-    return hasattr(torch.ops, "oink") and hasattr(torch.ops.oink, op_name)
-
-
-@_dynamo_disable
-def is_oink_available_for_device(device_index: int) -> bool:
-    """Return True if Oink ops are registered and device is SM100+.
-
-    This function is intended to be called during module initialization
-    (e.g., in RMSNorm.__init__), not in the forward path.
-
-    External plugins are expected to gate registration on SM100+ and
-    VLLM_USE_OINK_OPS=1, so if the ops are present they should be usable.
-    """
-    if not torch.cuda.is_available():
-        return False
-
-    try:
-        major, minor = torch.cuda.get_device_capability(device_index)
-        sm = 10 * major + minor
-        if sm < 100:
-            return False
-    except Exception:
-        return False
-
-    return _has_oink_op("rmsnorm")
-
-
-def has_fused_add_rms_norm() -> bool:
-    """Return True if the in-place fused op is registered."""
-    return _has_oink_op("fused_add_rms_norm")
-
-
-def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
-    """Call `torch.ops.oink.rmsnorm`.
-
-    This wrapper is safe to call in torch.compile regions.
-    """
-    return torch.ops.oink.rmsnorm(x, weight, eps)
-
-
-def fused_add_rms_norm_(
-    x: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    eps: float,
-) -> None:
-    """Call `torch.ops.oink.fused_add_rms_norm` (mutates x and residual)."""
-    torch.ops.oink.fused_add_rms_norm(x, residual, weight, eps)
-
-
-def fused_add_rms_norm(
-    x: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    eps: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Convenience wrapper returning (x, residual) after in-place mutation."""
-    fused_add_rms_norm_(x, residual, weight, eps)
-    return x, residual

vllm/compilation/backends.py

@@ -50,6 +50,7 @@ from .partition_rules import (
     should_split,
 )
 from .passes.inductor_pass import InductorPass, pass_context
+from .passes.ir.inplace_functionalization import VllmIRInplaceFunctionalizationPass
 from .passes.pass_manager import PostGradPassManager
 
 logger = init_logger(__name__)
@@ -926,6 +927,24 @@ class VllmBackend:
         return standalone_compile_artifacts, sym_shape_indices_map, returns_tuple_map
 
     def configure_post_pass(self) -> None:
+        # TODO proper PassManager?
+        pre_grad_pass_key = "pre_grad_custom_pass"
+        assert self.pass_key != pre_grad_pass_key
+        assert pre_grad_pass_key not in self.inductor_config
+        self.inductor_config[pre_grad_pass_key] = VllmIRInplaceFunctionalizationPass(
+            self.vllm_config
+        )
+
+        # Make sure pre_grad_custom_pass is not pickled
+        # as part of AOTAutograd built-in cache key
+        # TODO(luka) is there a cleaner way to do this
+        import torch._inductor.config as inductor_config
+
+        ignore = inductor_config._cache_config_ignore_prefix + [pre_grad_pass_key]
+        assert "_cache_config_ignore_prefix" not in self.inductor_config
+        self.inductor_config["_cache_config_ignore_prefix"] = ignore
+
+        # Configure the (nominally post-grad) pass manager
         self.pass_manager.configure(self.vllm_config)
 
         # Post-grad custom passes are run using the post_grad_custom_post_pass

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -40,7 +40,7 @@ from ..vllm_inductor_pass import (
     VllmPatternMatcherPass,
     VllmPatternReplacement,
 )
-from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8
+from .matcher_utils import MatcherQuantFP8
 
 FP8_DTYPE = current_platform.fp8_dtype()
 
@@ -356,10 +356,11 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
         super().__init__(dtype, device)
         self.epsilon = epsilon
         self.allreduce_params = allreduce_params
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
 
     def get_inputs(self) -> list[torch.Tensor]:
-        input, residual, weight = self.rmsnorm_matcher.inputs()
+        input = self.empty(5, 16)
+        residual = self.empty(5, 16)
+        weight = self.empty(16)
 
         # input goes through allreduce first, always 16-bit
         return [residual, input.to(self.dtype), weight]
@@ -369,7 +370,9 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
             residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor]:
             allreduce_output = tensor_model_parallel_all_reduce(input)
-            rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
+            rms, residual = vllm.ir.ops.fused_add_rms_norm(
+                allreduce_output, residual, weight, self.epsilon
+            )
             return rms, residual
 
         def replacement(
@@ -503,11 +506,12 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
         self.allreduce_params = allreduce_params
         self.quant_dtype = torch.float8_e4m3fn
 
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
         self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
 
     def get_inputs(self) -> list[torch.Tensor]:
-        input, residual, weight = self.rmsnorm_matcher.inputs()
+        input = self.empty(5, 16)
+        residual = self.empty(5, 16)
+        weight = self.empty(16)
         _, scale = self.quant_matcher.inputs()
 
         # input goes through allreduce first, always 16-bit
@@ -521,7 +525,9 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
             scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             allreduce_output = tensor_model_parallel_all_reduce(input)
-            rms, res = self.rmsnorm_matcher(allreduce_output, weight, residual)
+            rms, res = vllm.ir.ops.fused_add_rms_norm(
+                allreduce_output, residual, weight, self.epsilon
+            )
             quant, _ = self.quant_matcher(rms, scale)
 
             return quant, res
@@ -668,7 +674,6 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
         super().__init__(dtype, device)
         self.epsilon = epsilon
         self.allreduce_params = allreduce_params
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
 
     def get_inputs(self) -> list[torch.Tensor]:
         input = torch.empty([16, 16], device=self.device, dtype=self.dtype)
@@ -700,7 +705,9 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
             input_global_scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             allreduce_output = tensor_model_parallel_all_reduce(input)
-            rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
+            rms, residual = vllm.ir.ops.fused_add_rms_norm(
+                allreduce_output, residual, weight, self.epsilon
+            )
             quant_out_tuple = auto_functionalized(
                 STATIC_FP4_QUANT_OP,
                 input=rms,
@@ -955,15 +962,11 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
         super().__init__(dtype, device)
         self.epsilon = epsilon
         self.dtype = dtype
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(
-            epsilon, match_rocm_aiter=use_aiter_rmsnorm
-        )
         self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()
 
     def get_inputs(self) -> list[torch.Tensor]:
-        input, residual, weight = self.rmsnorm_matcher.inputs()
-
-        return [residual, input.to(self.dtype), weight]
+        # input, residual, weight
+        return [self.empty(5, 16), self.empty(5, 16), self.empty(16)]
 
     @property
     def pattern(self):
@@ -971,8 +974,9 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
             residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor]:
             allreduce_output = tensor_model_parallel_all_reduce(input)
-            rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
-
+            rms, residual = vllm.ir.ops.fused_add_rms_norm(
+                allreduce_output, residual, weight, self.epsilon
+            )
             return rms, residual
 
         return _pattern

vllm/compilation/passes/fusion/matcher_utils.py

@@ -10,7 +10,6 @@ from torch._ops import OpOverload
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import get_current_vllm_config
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
@@ -159,72 +158,6 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
         return result
 
 
-class MatcherFusedAddRMSNorm(MatcherCustomOp):
-    def __init__(
-        self,
-        epsilon: float,
-        enabled: bool | None = None,
-        match_rocm_aiter: bool = False,
-    ) -> None:
-        if enabled is None:
-            enabled = RMSNorm.enabled()
-
-        super().__init__(enabled)
-        self.epsilon = epsilon
-        self.match_rocm_aiter = match_rocm_aiter
-
-        self._rmsnorm_op = RMS_ADD_OP
-
-        if match_rocm_aiter:
-            self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_fused_add_op()
-
-    def inputs(self) -> list[torch.Tensor]:
-        input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
-        weight = self.empty(16)
-        residual = self.empty(5, 16)
-        return [input, weight, residual]
-
-    def forward_rocm_aiter(
-        self,
-        input: torch.Tensor,
-        weight: torch.Tensor,
-        residual: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        return self._rmsnorm_op(  # type: ignore[no-any-return]
-            x=input, residual=residual, weight=weight, variance_epsilon=self.epsilon
-        )
-
-    def forward_custom(
-        self,
-        input: torch.Tensor,
-        weight: torch.Tensor,
-        residual: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if self.match_rocm_aiter:
-            return self.forward_rocm_aiter(input, weight, residual)
-
-        _, result, residual = auto_functionalized(
-            self._rmsnorm_op,
-            input=input,
-            residual=residual,
-            weight=weight,
-            epsilon=self.epsilon,
-        )
-
-        return result, residual
-
-    def forward_native(
-        self,
-        input: torch.Tensor,
-        weight: torch.Tensor,
-        residual: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        result: tuple[torch.Tensor, torch.Tensor] = RMSNorm.forward_static(
-            input, self.epsilon, input.size(-1), self.model_dtype, weight, residual
-        )
-        return result
-
-
 class MatcherQuantFP8(MatcherCustomOp):
     def __init__(
         self,

vllm/compilation/passes/fusion/rms_quant_fusion.py

@@ -29,7 +29,6 @@ from vllm.platforms import current_platform
 from ..inductor_pass import enable_fake_mode
 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
 from .matcher_utils import (
-    MatcherFusedAddRMSNorm,
     MatcherQuantFP8,
 )
 
@@ -146,9 +145,6 @@ class RMSNormQuantPattern:
         assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"
         self.FUSED_OP = FUSED_OPS[key]
 
-        if key.fused_add:
-            self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
-
         self.quant_matcher = MatcherQuantFP8(
             key.quant,
             has_col_major_scales=has_col_major_scales,
@@ -231,7 +227,9 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
             residual: torch.Tensor,
             scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
-            result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
+            result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
+                input, residual, weight, self.epsilon
+            )
             result, _ = self.quant_matcher(result_rms, scale)
 
             return result, residual
@@ -261,8 +259,9 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
             return at[1], at[2]
 
         inputs = [
-            # input, weight, residual
-            *self.rmsnorm_matcher.inputs(),
+            empty_bf16(5, 16),  # input
+            empty_bf16(16),  # weight
+            empty_bf16(5, 16),  # residual
             self.quant_matcher.inputs()[1],  # scale
         ]
 
@@ -311,7 +310,9 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
             residual: torch.Tensor,
             scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
+            result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
+                input, residual, weight, self.epsilon
+            )
             result = torch.empty(
                 result_rms.shape,
                 device=result_rms.device,
@@ -366,12 +367,17 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
             # result, residual, scale
             return at[1], at[3], at[2]
 
-        scale = self.quant_matcher.empty_f32(1, 1)
+        inputs = [
+            empty_bf16(5, 16),  # input
+            empty_bf16(16),  # weight
+            empty_bf16(5, 16),  # residual
+            self.quant_matcher.empty_f32(1, 1),  # scale
+        ]
 
         pm.register_replacement(
             pattern,
             replacement,
-            self.rmsnorm_matcher.inputs() + [scale],
+            inputs,
             pm.fwd_only,
             pm_pass,
             extra_check=_rms_input_weight_dtype_match,
@@ -552,7 +558,9 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
         def pattern(
             input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
+            result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
+                input, residual, weight, self.epsilon
+            )
             result, scale = self.quant_matcher(result_rms)
 
             return result, residual, scale
@@ -580,10 +588,16 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
             # result, residual, scale
             return at[1], at[3], at[2]
 
+        inputs = [
+            empty_bf16(5, 16),  # input
+            empty_bf16(16),  # weight
+            empty_bf16(5, 16),  # residual
+        ]
+
         pm.register_replacement(
             pattern,
             replacement,
-            self.rmsnorm_matcher.inputs(),
+            inputs,
             pm.fwd_only,
             pm_pass,
             extra_check=_rms_input_weight_dtype_match,

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

@@ -29,7 +29,6 @@ from ..vllm_inductor_pass import (
     VllmPatternReplacement,
 )
 from .matcher_utils import (
-    MatcherFusedAddRMSNorm,
     MatcherQuantFP8,
     MatcherSiluAndMul,
 )
@@ -49,10 +48,6 @@ class AiterRMSNormQuantPattern:
         self.quant_dtype = key.quant.dtype
         self.device = torch.device("cuda")
 
-        if key.fused_add:
-            self.rmsnorm_matcher = MatcherFusedAddRMSNorm(
-                epsilon, match_rocm_aiter=True
-            )
         self.quant_matcher = MatcherQuantFP8(
             key.quant,
             match_rocm_aiter=match_aiter_quant,
@@ -145,7 +140,9 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
             weight: torch.Tensor,
             residual: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
+            result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
+                input, residual, weight, self.epsilon
+            )
             result, scale = self.quant_matcher(result_rms)
 
             return result, residual_out, scale
@@ -163,10 +160,16 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
 
             return result[0], result[1], result[2]
 
+        inputs = [
+            self.empty(5, 16),  # input
+            self.empty(16),  # weight
+            self.empty(5, 16),  # residual
+        ]
+
         pm.register_replacement(
             pattern,
             replacement,
-            self.rmsnorm_matcher.inputs(),
+            inputs,
             pm.fwd_only,
             pm_pass,
         )
@@ -258,7 +261,9 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
             weight: torch.Tensor,
             residual: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
+            result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
+                input, residual, weight, self.epsilon
+            )
             result, scale = self.quant_matcher(result_rms)
 
             return result, residual_out, scale
@@ -279,9 +284,13 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
             # result, scale, residual
             return at[0], at[1], at[2]
 
-        pm.register_replacement(
-            pattern, replacement, self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass
-        )
+        inputs = [
+            self.empty(5, 16),  # input
+            self.empty(16),  # weight
+            self.empty(5, 16),  # residual
+        ]
+
+        pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
 
 
 class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
@@ -420,12 +429,15 @@ class AddAiterRMSNormPadPattern:
         self.epsilon = epsilon
         self.hidden_size = hidden_size
         self.x_pad_to_multiple = x_pad_to_multiple
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon, match_rocm_aiter=True)
 
     def get_inputs(self) -> list[torch.Tensor]:
-        input, weight, residual = self.rmsnorm_matcher.inputs()
-        router_weight = torch.empty([8, 16], dtype=weight.dtype, device=weight.device)
-        router_bias = torch.empty([8], dtype=weight.dtype, device=weight.device)
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
+        input = torch.empty(5, 16, dtype=dtype, device=device)
+        weight = torch.empty(16, dtype=dtype, device=device)
+        residual = torch.empty(5, 16, dtype=dtype, device=device)
+        router_weight = torch.empty([8, 16], dtype=dtype, device=device)
+        router_bias = torch.empty([8], dtype=dtype, device=device)
         return [input, weight, residual, router_weight, router_bias]
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
@@ -439,7 +451,9 @@ class AddAiterRMSNormPadPattern:
             pad_size = self.x_pad_to_multiple - (
                 self.hidden_size % self.x_pad_to_multiple
             )
-            result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
+            result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
+                input, residual, weight, self.epsilon
+            )
             router_logits = torch.ops.vllm.rocm_unquantized_gemm(
                 result_rms, router_weight, router_bias
             )

vllm/compilation/passes/fusion/sequence_parallelism.py

@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from ..inductor_pass import enable_fake_mode
 from ..utility.noop_elimination import NoOpEliminationPass
 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
-from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8
+from .matcher_utils import MatcherQuantFP8
 
 logger = init_logger(__name__)
 
@@ -174,7 +174,6 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
 class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
     def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
         super().__init__(epsilon, dtype, device)
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
 
     def get_inputs(self) -> list[torch.Tensor]:
         mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -195,7 +194,9 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
             rms_norm_weights: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             all_reduce = self._all_reduce(mm_1)
-            rmsnorm = self.rmsnorm_matcher(all_reduce, rms_norm_weights, residual)
+            rmsnorm = vllm.ir.ops.fused_add_rms_norm(
+                all_reduce, residual, rms_norm_weights, self.epsilon
+            )
             return rmsnorm[0], rmsnorm[1]
 
         def replacement(
@@ -208,7 +209,9 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
             # once the seqpar pattern with the previous rmsnorm is replaced
             reduce_scatter = self._reduce_scatter(mm_1)
             residual = residual[0 : reduce_scatter.size(0), ...]
-            rmsnorm = self.rmsnorm_matcher(reduce_scatter, rms_norm_weights, residual)
+            rmsnorm = vllm.ir.ops.fused_add_rms_norm(
+                reduce_scatter, residual, rms_norm_weights, self.epsilon
+            )
             all_gather = self._all_gather(rmsnorm[0])
             # shape of residual changes but that's fine,
             # next node is already slicing it, now becomes a noop
@@ -271,7 +274,6 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
 class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
     def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
         super().__init__(epsilon, dtype, device)
-        self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
         self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
 
     def get_inputs(self) -> list[torch.Tensor]:
@@ -290,8 +292,8 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
             scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             all_reduce = self._all_reduce(mm_1)
-            rms, residual_out = self.rmsnorm_matcher(
-                all_reduce, rms_norm_weights, residual
+            rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
+                all_reduce, residual, rms_norm_weights, self.epsilon
             )
             quant, _ = self.quant_matcher(rms, scale)
             return quant, residual_out
@@ -308,8 +310,8 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
             # once the seqpar pattern with the previous rmsnorm is replaced
             reduce_scatter = self._reduce_scatter(mm_1)
             residual = residual[0 : reduce_scatter.size(0), ...]
-            rms, residual_out = self.rmsnorm_matcher(
-                reduce_scatter, rms_norm_weights, residual
+            rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
+                reduce_scatter, residual, rms_norm_weights, self.epsilon
             )
             quant, _ = self.quant_matcher(rms, scale)
             all_gather = self._all_gather(quant)

vllm/compilation/passes/inductor_pass.py

@@ -30,6 +30,9 @@ class PassContext:
     def __init__(self, compile_range: Range):
         self.compile_range: Range = compile_range
 
+        # set of arg indices
+        self.donated_input_ids: set[int] = set()
+
 
 def get_pass_context() -> PassContext:
     """Get the current pass context."""

vllm/compilation/passes/ir/clone_elimination.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+from torch import fx
+from torch._higher_order_ops.auto_functionalize import auto_functionalized
+from torch._higher_order_ops.triton_kernel_wrap import TritonKernelWrapperFunctional
+from torch._ops import HigherOrderOperator, OpOverload
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from ..fx_utils import is_func
+from ..inductor_pass import get_pass_context
+from ..vllm_inductor_pass import VllmInductorPass
+
+logger = init_logger(__name__)
+
+
+def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool:
+    if user.op == "output":
+        return False
+
+    if is_func(user, auto_functionalized):
+        # While autofunc writes to the node,
+        # this is a follow-up use we're not interested in.
+        # It is also guaranteed to be the final use,
+        # as auto_functionalized returns the tensor back for follow-up use.
+        return False
+    elif user.op == "call_function" and isinstance(user.target, HigherOrderOperator):
+        # By default, be conservative, assume this could be a write
+        # (except functional HOPs)
+        return not isinstance(user.target, TritonKernelWrapperFunctional)
+
+    assert isinstance(user.target, OpOverload), (
+        f"{node=} {user=} {user.op=} {user.target=}"
+    )
+    schema = user.target._schema
+    assert len(user.args) <= len(schema.arguments)
+    for i, arg in enumerate(user.args):
+        # Only interested in writes to node
+        if arg is not node:
+            continue
+
+        # If not a write, next arg could be
+        if schema.arguments[i].is_write:
+            return True
+
+    # No writes found
+    return False
+
+
+class UnsafeCloneEliminationPass(VllmInductorPass):
+    """
+    This pass removes clone nodes that are no longer needed after vLLM IR lowering.
+    It uses donated_input_ids to eliminate clones of donated graph inputs, preserving
+    contents of non-donated graph inputs.
+
+    It is "unsafe" because it does not (yet) take aliasing into account. Solving
+    aliasing is an open problem, so this pass intends to support known vLLM cases
+    and not guarantee soundness on general graphs. In the future, this pass will likely
+    support basic forms of aliasing to handle simple views (e.g. qkv -> q,k,v).
+    """
+
+    def __init__(self, vllm_config: VllmConfig) -> None:
+        super().__init__(vllm_config)
+
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        count = 0
+        node_to_idx = {node: i for i, node in enumerate(graph.nodes)}
+        pass_context = get_pass_context()
+        donated_input_ids = pass_context.donated_input_ids
+        logger.debug("Donated input ids: %s", donated_input_ids)
+
+        for node in graph.nodes:
+            if not is_func(node, torch.ops.aten.clone.default):
+                continue
+
+            original_node = node.args[0]
+            assert isinstance(original_node, fx.Node)
+
+            # Clone needs to be preserved if node is getting written to and
+            # the old value is used again.
+            # This could only happen if an inplace implementation was lowered.
+            # Then node (the clone) will have one write.
+            # TODO(luka) hopefully this can be removed once we lower functional graphs.
+            write_idxs = [
+                node_to_idx[u] for u in node.users if user_writes_to_node(u, node)
+            ]
+            assert len(write_idxs) in (0, 1)
+            if write_idxs:
+                # Check if a user of original_node occurs after a write
+                write_idx = write_idxs[0]
+                if any(
+                    node_to_idx[orig_user] > write_idx
+                    for orig_user in original_node.users
+                ):
+                    logger.debug(
+                        "Clone removal not possible, "
+                        "original_node=%s used after mutation on node=%s",
+                        original_node,
+                        node,
+                    )
+                    continue
+
+                # Check if a node is a (non-donated) graph input
+                if (
+                    original_node.op == "placeholder"
+                    and node_to_idx[original_node] not in donated_input_ids
+                ):
+                    logger.debug(
+                        "Graph input %s not donated, cannot eliminate its clone",
+                        original_node,
+                    )
+                    continue
+
+            logger.debug(
+                "Node %s is a redundant clone node of %s, removing it",
+                node,
+                original_node,
+            )
+            node.replace_all_uses_with(original_node)
+            graph.erase_node(node)
+            count += 1
+
+        logger.debug("CloneCleanupPass removed %d clone nodes", count)

vllm/compilation/passes/ir/inplace_functionalization.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections import defaultdict
+
+from torch import fx
+from torch._inductor.pattern_matcher import (
+    PatternMatcherPass,
+)
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from ..inductor_pass import get_pass_context
+from ..vllm_inductor_pass import VllmInductorPass
+from .lowering_pass import get_ir_op
+from .utils import overload_or_default
+
+logger = init_logger(__name__)
+
+
+class VllmIRInplaceFunctionalizationPass(VllmInductorPass):
+    """
+    This pass functionalizes maybe_inplace vLLM IR ops to the default overload.
+    The maybe_inplace overloads have the same signature as the default overload
+    so the pass simply replaces the called overload.
+    That makes the graph properly functional.
+    The pass also validates that activations passed to maybe_inplace have no later
+    uses in the graph: they are donated to the maybe_inplace op call,
+    and their contents are not defined afterward.
+
+    This pass operates pre-AOTAutograd,
+    so it must handle non-normalized and non-functional IR.
+    """
+
+    def __init__(self, vllm_config: VllmConfig) -> None:
+        super().__init__(vllm_config)
+        self.patterns = PatternMatcherPass(self.pass_name)
+        self.functionalized_ops: dict[str, int] = defaultdict(lambda: 0)
+
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        # clear at the beginning instead of end, so that tests can inspect
+        self.functionalized_ops.clear()
+        assert graph.owning_module is not None
+        node_to_idx = {node: i for i, node in enumerate(graph.nodes)}
+
+        # Pass donated input via vLLM's pass context
+        pass_context = get_pass_context()
+        pass_context.donated_input_ids = set[int]()
+
+        for node in graph.nodes:
+            if (ir_op := get_ir_op(node)) is None:
+                continue
+
+            op_overload = overload_or_default(node.target)
+            overload_name = op_overload._overloadname
+            if overload_name != "maybe_inplace":
+                assert overload_name == "default", (
+                    f"Found overload {overload_name} for op {ir_op.name}, "
+                    f"expected maybe_inplace or default"
+                )
+                continue
+
+            # must have maybe_inplace overload and allow_inplace
+            assert ir_op.allow_inplace and hasattr(ir_op, "maybe_inplace")
+
+            # Check that activation inputs are not used after this op
+            for arg_idx in ir_op.activation_indices:
+                arg = node.args[arg_idx]
+                assert isinstance(arg, fx.Node), "Activation inputs must be fx.Node"
+                for user in arg.users:
+                    if node_to_idx[user] > node_to_idx[node]:
+                        raise ValueError(
+                            f"Input {arg} to maybe_inplace node {node} "
+                            f"is used again after the node. "
+                            f"This is not allowed; activation inputs to maybe_inplace "
+                            f"ops are donated to the op, meaning their memory may be "
+                            f"recycled for outputs.\n\n"
+                            f"To preserve the inputs, use the default overload or "
+                            f"clone them manually beforehand."
+                        )
+
+                if arg.op == "placeholder":
+                    # Graph input that maybe_inplace might modify.
+                    # Mark it so downstream passes know it's donated.
+                    # TODO(luka) store in placeholder node meta once supported
+                    pass_context.donated_input_ids.add(node_to_idx[arg])
+
+            # Same signature, just replace the overload that's called.
+            node.target = ir_op.torch_op
+            self.functionalized_ops[ir_op.name] += 1
+
+        count = sum(self.functionalized_ops.values())
+        ops = ",".join(self.functionalized_ops.keys())
+        logger.debug("Donated input IDs: %s", pass_context.donated_input_ids)
+        logger.debug(
+            "%s functionalized %d vLLM IR nodes for op(s) %s",
+            self.pass_name,
+            count,
+            ops,
+        )

vllm/compilation/passes/ir/lowering_pass.py

@@ -10,7 +10,6 @@ from torch._inductor.pattern_matcher import (
     PatternMatcherPass,
     register_graph_pattern,
 )
-from torch._ops import OpOverload, OpOverloadPacket
 
 from vllm.config import VllmConfig
 from vllm.ir.op import IrOp
@@ -18,41 +17,11 @@ from vllm.logger import init_logger
 from vllm.logging_utils import lazy
 
 from ..vllm_inductor_pass import VllmInductorPass
+from .utils import get_ir_op
 
 logger = init_logger(__name__)
 
 
-def get_default_overload(op: OpOverload | OpOverloadPacket) -> OpOverload:
-    if isinstance(op, OpOverloadPacket):
-        return op.default
-    assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket"
-    return op
-
-
-def get_ir_op(node: fx.Node) -> IrOp | None:
-    if node.op != "call_function":
-        return None
-
-    if not isinstance(node.target, (OpOverload, OpOverloadPacket)):
-        return None
-
-    op_overload = get_default_overload(node.target)
-    if op_overload.namespace != "vllm_ir":
-        return None
-
-    op_name = op_overload._opname
-    if op_name not in IrOp.registry:
-        logger.warning(
-            "Unknown vLLM IR op %s, there's likely an issue with torch registration, "
-            "or a torch custom op was registered in the vllm_ir namespace by mistake.",
-            op_name,
-        )
-        return None
-
-    ir_op = IrOp.registry[op_name]
-    return ir_op
-
-
 class VllmIRLoweringPass(VllmInductorPass):
     """
     This pass lowers vLLM IR ops to their implementations the priority list.
@@ -76,7 +45,7 @@ class VllmIRLoweringPass(VllmInductorPass):
 
         assert len(match.nodes) == 1, "Expected single node match"
         node = match.nodes[0]
-        ir_op = get_ir_op(node)
+        ir_op = get_ir_op(node)  # TODO is node.target always an overload?
         assert ir_op is not None, "Expected vLLM IR op"
         assert not node.kwargs  # I think there should never be kwargs here
 
@@ -86,13 +55,18 @@ class VllmIRLoweringPass(VllmInductorPass):
         self.selected_impls[ir_op.name][node.name] = ir_op_impl.provider
 
         # replace_by_example wants node args, not the fake tensors
+        # use func_impl_fn to properly handle in-place implementations
         # TODO(luka): Use aot_export_module to get functionalized graph
         # TODO(luka): Cache the fx_replacement to avoid re-tracing the same impl
 
         # 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)
+        # It is not safe to run functional passes (like DCE) on the replacements
+        # as they might not be functional.
+        match.replace_by_example(
+            ir_op_impl.func_impl_fn, bound_args.args, run_functional_passes=False
+        )
 
     @VllmInductorPass.time_and_log
     def __call__(self, graph: fx.Graph) -> None:
@@ -136,7 +110,7 @@ class VllmIRLoweringPass(VllmInductorPass):
 
         if failed_nodes or failed_ops:
             logger.warning("Failed to lower vLLM IR ops: %s", ",".join(failed_ops))
-            logger.warning("Full node list: %s", failed_nodes)
+            logger.warning("Full node list: %s", ",".join(str(n) for n in failed_nodes))
 
     def uuid(self) -> str:
         """

vllm/compilation/passes/ir/utils.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from torch import fx
+from torch._ops import OpOverload, OpOverloadPacket
+
+from vllm.ir.op import IrOp
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def overload_or_default(op: OpOverload | OpOverloadPacket) -> OpOverload:
+    if isinstance(op, OpOverloadPacket):
+        return op.default
+    assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket"
+    return op
+
+
+def get_ir_op(node: fx.Node) -> IrOp | None:
+    if node.op != "call_function":
+        return None
+
+    if not isinstance(node.target, (OpOverload, OpOverloadPacket)):
+        return None
+
+    op_overload = overload_or_default(node.target)
+    if op_overload.namespace != "vllm_ir":
+        return None
+
+    op_name = op_overload._opname
+    if op_name not in IrOp.registry:
+        logger.warning(
+            "Unknown vLLM IR op %s, there's likely an issue with torch registration, "
+            "or a torch custom op was registered in the vllm_ir namespace by mistake.",
+            op_name,
+        )
+        return None
+
+    ir_op = IrOp.registry[op_name]
+    return ir_op

vllm/compilation/passes/pass_manager.py

@@ -14,6 +14,7 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.system_utils import set_env_var
 
+from .ir.clone_elimination import UnsafeCloneEliminationPass
 from .ir.lowering_pass import VllmIRLoweringPass
 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
 
@@ -115,6 +116,8 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
         # DCE handles mutating ops correctly as well.
         self.ir_lowering(graph)
         VllmInductorPass.dump_prefix += 1
+        self.clone_elimination(graph)
+        VllmInductorPass.dump_prefix += 1
 
         # clean up after lowering again
         self.post_cleanup(graph)
@@ -149,11 +152,12 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
                 self.passes += [MiniMaxQKNormPass(config)]
 
             if self.pass_config.fuse_norm_quant:
-                self.passes += [RMSNormQuantFusionPass(config)]
                 if rocm_aiter_ops.is_enabled():
                     self.passes += [
                         RocmAiterRMSNormQuantFusionPass(config),
                     ]
+                self.passes += [RMSNormQuantFusionPass(config)]
+
             if self.pass_config.fuse_act_quant:
                 self.passes += [ActivationQuantFusionPass(config)]
                 if rocm_aiter_ops.is_enabled():
@@ -179,6 +183,7 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
                 self.passes += [QKNormRoPEFusionPass(config)]
 
             self.ir_lowering = VllmIRLoweringPass(config)
+            self.clone_elimination = UnsafeCloneEliminationPass(config)
             self.post_cleanup = PostCleanupPass(config)
             self.fix_functionalization = FixFunctionalizationPass(config)
 
@@ -200,6 +205,7 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
 
         passes.append(self.post_cleanup.uuid())
         passes.append(self.ir_lowering.uuid())
+        passes.append(self.clone_elimination.uuid())
         passes.append(self.post_cleanup.uuid())
         passes.append(self.fix_functionalization.uuid())
 

vllm/config/kernel.py

@@ -31,6 +31,9 @@ class IrOpPriorityConfig:
     rms_norm: list[str] = Field(default_factory=list)
     """Priority list for vllm.ir.ops.rms_norm"""
 
+    fused_add_rms_norm: list[str] = Field(default_factory=list)
+    """Priority list for vllm.ir.ops.fused_add_rms_norm"""
+
     def compute_hash(self) -> str:
         """
         Produces a hash unique to the pass configuration.

vllm/config/vllm.py

@@ -125,9 +125,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
         from vllm._aiter_ops import rocm_aiter_ops
 
         return (
-            rocm_aiter_ops.is_enabled()
-            and rocm_aiter_ops.is_rmsnorm_enabled()
-            and cfg.parallel_config.tensor_parallel_size > 1
+            rocm_aiter_ops.is_enabled() and cfg.parallel_config.tensor_parallel_size > 1
         )
 
     return (
@@ -159,10 +157,9 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:
 
 def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
     """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss."""
-    from vllm._aiter_ops import rocm_aiter_ops
 
     return (
-        rocm_aiter_ops.is_rmsnorm_enabled()
+        cfg.kernel_config.ir_op_priority.fused_add_rms_norm[0] == "aiter"
         and cfg.model_config is not None
         and cfg.model_config.get_hidden_size() == 2880
     )

vllm/envs.py

@@ -143,7 +143,7 @@ if TYPE_CHECKING:
     VLLM_DP_RANK_LOCAL: int = -1
     VLLM_DP_SIZE: int = 1
     VLLM_USE_STANDALONE_COMPILE: bool = True
-    VLLM_ENABLE_PREGRAD_PASSES: bool = False
+    VLLM_ENABLE_PREGRAD_PASSES: bool = True
     VLLM_DP_MASTER_IP: str = ""
     VLLM_DP_MASTER_PORT: int = 0
     VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
@@ -620,9 +620,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # The pre-grad passes get run even on cache-hit and negatively impact
     # vllm cold compile times by O(1s)
     # Can remove this after the following issue gets fixed
+    # TODO(luka): maybe_inplace requires this
     # https://github.com/pytorch/pytorch/issues/174502
     "VLLM_ENABLE_PREGRAD_PASSES": lambda: (
-        os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "0") == "1"
+        os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "1") == "1"
     ),
     # Debug pattern matching inside custom passes.
     # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').

vllm/ir/op.py

@@ -4,7 +4,7 @@ import contextlib
 import inspect
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, ClassVar, overload
+from typing import Any, ClassVar, Literal, overload
 
 import torch
 from torch.library import Library, infer_schema
@@ -46,35 +46,51 @@ def enable_torch_wrap(enable: bool = True):
         _ENABLE_TORCH_WRAP = old
 
 
-# 0-param decorator overload
+# 0-param decorator overload (no inplace)
 @overload
 def register_op(f: Callable[..., Any]) -> "IrOp": ...
 
 
-# parametrized decorator overload
+# parametrized decorator with allow_inplace=False (default)
 @overload
 def register_op(
     *,
     name: str | None = None,
+    activations: list[str] | None = None,
+    allow_inplace: Literal[False] = False,
 ) -> Callable[[Callable[..., Any]], "IrOp"]: ...
 
 
+# parametrized decorator with allow_inplace=True
+@overload
+def register_op(
+    *,
+    name: str | None = None,
+    activations: list[str] | None = None,
+    allow_inplace: Literal[True],
+) -> Callable[[Callable[..., Any]], "IrOpInplace"]: ...
+
+
 def register_op(
     f: Callable | None = None,
     *,
     name: str | None = None,
+    activations: list[str] | None = None,
+    allow_inplace: 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 activations: list of activation params, defaults to params starting with 'x'
+    :param allow_inplace: add a maybe_inplace overload that allows inplace impls
     :return: the IrOp object if f is provided, otherwise a decorator
 
     Example usage:
     ```python
     @vllm.ir.register_op
-    def my_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return x + y
 
 
@@ -85,7 +101,10 @@ 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)
+        if allow_inplace:
+            op: IrOp = IrOpInplace(op_name, _f, activations)
+        else:
+            op = IrOp(op_name, _f, activations)
         IrOp.registry[op_name] = op
         return op
 
@@ -100,8 +119,14 @@ class IrOp:
 
     name: str
     impls: dict[str, "IrOpImpl"]
+    allow_inplace: bool = False
 
-    def __init__(self, name: str, native_impl: Callable):
+    def __init__(
+        self,
+        name: str,
+        native_impl: Callable,
+        activations: list[str] | None = None,
+    ):
         self._py_signature = inspect.signature(native_impl)
         if any(
             p.kind == inspect.Parameter.KEYWORD_ONLY
@@ -112,8 +137,22 @@ class IrOp:
                 f"supported. That's because kwargs are not allowed during lowering."
             )
 
+        # By convention, we consider parameters starting with 'x' as activations.
+        if activations is None:
+            activations = [
+                p.name
+                for p in self._py_signature.parameters.values()
+                if p.name.startswith("x")
+            ]
+
         self.name = name
         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._schema_str = infer_schema(native_impl, mutates_args=[])
         self._input_generator: InputGenerator | None = None
@@ -121,7 +160,12 @@ class IrOp:
 
         # 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,
         )
 
         # By default, fake routes directly to native,
@@ -161,12 +205,14 @@ class IrOp:
         *,
         supported: bool = True,
         supports_args: Callable[..., bool] | None = None,
+        inplace: 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
         :return: A decorator that registers the implementation.
 
         The decorated function must have the same semantics and signature as
@@ -193,7 +239,7 @@ class IrOp:
         )
 
         def _register_impl(f: Callable):
-            impl = IrOpImpl(self, provider, f, supported, supports_args)
+            impl = IrOpImpl(self, provider, f, supported, supports_args, inplace)
             self.impls[provider] = impl
 
             if self.get_priority():
@@ -213,7 +259,10 @@ class IrOp:
         __call__ routes straight here instead of going through torch op dispatching.
         """
         impl = self.dispatch(*args, **kwargs)
-        return impl.impl_fn(*args, **kwargs)
+
+        # Default overload must be functional,
+        # use func_impl_fn to correctly handle inplace impls.
+        return impl.func_impl_fn(*args, **kwargs)
 
     def apply_arg_defaults(self, args) -> tuple:
         """
@@ -314,6 +363,11 @@ class IrOp:
         old_priority_impls = self._priority_impls
         try:
             self._priority_impls = filter_priority_impls(priority)
+            logger.debug(
+                "Priority for vllm.ir.%s set to %s",
+                self.name,
+                lazy(lambda: [p.provider for p in self._priority_impls]),
+            )
             yield
         finally:
             self._priority_impls = old_priority_impls
@@ -354,6 +408,66 @@ class IrOp:
         )
 
 
+class IrOpInplace(IrOp):
+    """IR op with inplace support via maybe_inplace."""
+
+    maybe_inplace: "IrOpInplaceOverload"
+    allow_inplace: bool = True
+
+    def __init__(
+        self,
+        name: str,
+        native_impl: Callable,
+        activations: list[str] | None = None,
+    ):
+        super().__init__(name, native_impl, activations)
+
+        # Create the inplace overload
+        self.maybe_inplace = IrOpInplaceOverload(self)
+
+
+class IrOpInplaceOverload:
+    def __init__(self, op: IrOp):
+        params, returns = op._schema_str.split(" -> ")
+        n_outputs = returns.count("Tensor")
+
+        assert returns.count("Tensor") == len(op.activations), (
+            "Inplace overload requires the same number of outputs as activations."
+        )
+
+        assert returns.count(",") == n_outputs - 1, (
+            "Inplace overload only supports Tensor outputs for now."
+        )
+
+        self.op = op
+        self.name = f"{op.name}.maybe_inplace"
+        self._schema_str = infer_schema(
+            op.impls["native"].impl_fn, mutates_args=op.activations
+        )
+
+        # torch registration
+        vllm_ir_lib.define(self.name + self._schema_str)
+        vllm_ir_lib.impl(
+            self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd"
+        )
+        # fake goes to default overload for now
+        vllm_ir_lib._register_fake(self.name, self.op._fake_call)
+
+        assert hasattr(getattr(torch.ops.vllm_ir, self.op.name), "maybe_inplace")
+        self.torch_op = getattr(torch.ops.vllm_ir, self.op.name).maybe_inplace
+
+    def __call__(self, *args, **kwargs) -> Any:
+        if not _ENABLE_TORCH_WRAP:
+            return self._inner_call(*args, **kwargs)
+
+        return self.torch_op(*args, **kwargs)
+
+    def _inner_call(self, *args, **kwargs) -> Any:
+        # Calling the maybe_inplace overload means we can use inplace impls directly.
+        impl = self.op.dispatch(*args, **kwargs)
+        return impl.impl_fn(*args, **kwargs)
+
+
 class IrOpImpl:
     def __init__(
         self,
@@ -362,6 +476,7 @@ class IrOpImpl:
         impl_fn: Callable,
         supported: bool,
         supports_args: Callable[..., bool] | None,
+        inplace: bool = False,
     ):
         assert provider not in op.impls, (
             f"Implementation for provider {provider} already registered."
@@ -420,11 +535,18 @@ class IrOpImpl:
                         f"native default {op_p.default}'"
                     )
 
+        if inplace:
+            assert op.allow_inplace, (
+                f"Inplace implementation cannot be registered for op {op.name}"
+                f" that does not allow inplace."
+            )
+
         self.op = op
         self.provider = provider
         self.impl_fn = impl_fn
         self.supported = supported
         self._supports_args = supports_args
+        self.inplace = inplace
 
     @property
     def supports_all_args(self) -> bool:
@@ -449,3 +571,19 @@ class IrOpImpl:
         """
         sources = [Path(inspect.getfile(self.impl_fn))]
         return hash_source(*sources)
+
+    def func_impl_fn(self, *args, **kwargs) -> Any:
+        """
+        Copy any inputs in activations if this is an inplace impl,
+        to ensure functional semantics.
+        """
+        if not self.inplace:
+            return self.impl_fn(*args, **kwargs)
+
+        # copy activations to ensure functional semantics
+        new_args = list(args)
+        for i in self.op.activation_indices:
+            assert isinstance(args[i], torch.Tensor)
+            new_args[i] = args[i].clone()
+
+        return self.impl_fn(*new_args, **kwargs)

vllm/ir/ops/__init__.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from .layernorm import rms_norm
+from .layernorm import fused_add_rms_norm, rms_norm
 
-__all__ = ["rms_norm"]
+__all__ = ["rms_norm", "fused_add_rms_norm"]

vllm/ir/ops/layernorm.py

@@ -27,10 +27,46 @@ def _rms_norm_input_generator(
 ) -> tuple:
     x = torch.randn(num_tokens, hidden_size, dtype=dtype)
     weight = torch.randn(hidden_size, dtype=dtype)
-    return (x, weight, epsilon)
+    return x, weight, epsilon
 
 
 # Reductions in rms_norm accumulate rounding error at large shapes
 # (e.g. 32768x16384), causing a few elements out of millions to exceed
 # the default float16 tolerance.
 rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
+
+
+@register_op(allow_inplace=True)
+def fused_add_rms_norm(
+    x: Tensor,
+    x_residual: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> tuple[Tensor, Tensor]:
+    """Fused add and weighted root-mean-square layer normalization"""
+    orig_dtype = x.dtype
+    x = x.to(torch.float32)
+    x = x + x_residual.to(torch.float32)
+    x_residual = x.to(orig_dtype)
+
+    x_var = x if variance_size is None else x[..., :variance_size]
+    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
+    x = x * torch.rsqrt(variance + epsilon)
+    if weight is not None:
+        x = x.to(weight.dtype) * weight
+    return x.to(orig_dtype), x_residual
+
+
+# fused_add_rms_norm has similar rounding error accumulation as rms_norm
+fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
+
+
+@fused_add_rms_norm.register_input_generator
+def _fused_add_rms_norm_input_generator(
+    num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
+) -> tuple:
+    x = torch.randn(num_tokens, hidden_size, dtype=dtype)
+    x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype)
+    weight = torch.randn(hidden_size, dtype=dtype)
+    return x, x_residual, weight, epsilon

vllm/kernels/aiter_ops.py

@@ -75,3 +75,72 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor
 direct_register_aiter_op(
     op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake
 )
+
+rms_add_no_var_16bit_only = (
+    lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
+    and x.dtype in (torch.float16, torch.bfloat16)
+    and (weight is None or weight.dtype == x.dtype)
+)
+"""
+AITER fused_add_rms_norm only supports 16-bit activations and no var_size override.
+Requires weight dtype to match x dtype.
+"""
+
+
+@ir.ops.fused_add_rms_norm.register_impl(
+    "aiter", supports_args=rms_add_no_var_16bit_only, supported=AITER_SUPPORTED
+)
+def fused_add_rms_norm(
+    x: Tensor,
+    x_residual: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> tuple[Tensor, Tensor]:
+    assert variance_size is None
+    assert x.dtype in (torch.float16, torch.bfloat16)
+    if weight is None:
+        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
+    return torch.ops.vllm_aiter.fused_add_rms_norm(x, x_residual, weight, epsilon)
+
+
+def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    variance_epsilon: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    from aiter import rmsnorm2d_fwd_with_add
+
+    # TODO can out = x and residual_out = residual to save memory?
+    #  Need to check if the kernel supports in-place residual output
+    #  (if yes set mutates_args and inplace)
+    residual_out = torch.empty_like(residual)
+    out = torch.empty_like(x)
+    rmsnorm2d_fwd_with_add(
+        out,  # output
+        x,  # input
+        residual,  # residual input
+        residual_out,  # residual output
+        weight,
+        variance_epsilon,
+    )
+    return out, residual_out
+
+
+def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    variance_epsilon: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    residual_out = torch.empty_like(residual)
+    out = torch.empty_like(x)
+    return out, residual_out
+
+
+direct_register_aiter_op(
+    op_name="fused_add_rms_norm",
+    op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
+    fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
+)

vllm/kernels/oink_ops.py

@@ -1,6 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""This file registers Oink implementations for vLLM IR ops.
+
+vLLM does not depend on the external Oink repository/package. When an external
+plugin registers torch.library.custom_op entrypoints under the `oink::`
+namespace (e.g. via vLLM's general_plugins mechanism), these ops will be marked
+ as supported. To dispatch to those ops, set kernel_config.ir_op_priority.<op> to oink.
+Alternatively, `VLLM_USE_OINK_OPS=1` will add this to priority by default.
+"""
+
 import torch
+from torch import Tensor
 
 from vllm import ir
 from vllm.platforms import current_platform
@@ -15,7 +25,7 @@ def has_oink_op(name: str) -> bool:
     return OINK_AVAILABLE and hasattr(torch.ops.oink, name)
 
 
-def _can_view_as_2d(x: torch.Tensor) -> bool:
+def _can_view_as_2d(x: Tensor) -> bool:
     """Return True if x.view(-1, x.shape[-1]) is viewable (no copy)."""
     if x.dim() < 2:
         return False
@@ -32,7 +42,7 @@ def _can_view_as_2d(x: torch.Tensor) -> bool:
     return True
 
 
-def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
+def _is_oink_stride_compatible_2d(x_2d: Tensor) -> bool:
     """Return True if x_2d meets Oink's pointer-path stride constraints."""
     if x_2d.dim() != 2:
         return False
@@ -67,11 +77,51 @@ and no variance_size override.
     "oink", supports_args=oink_rms_supported, supported=has_oink_op("rmsnorm")
 )
 def rms_norm(
-    x: torch.Tensor,
-    weight: torch.Tensor | None,
+    x: Tensor,
+    weight: Tensor | None,
     epsilon: float,
     variance_size: int | None = None,
-) -> torch.Tensor:
+) -> Tensor:
     assert variance_size is None
     x_2d = x.view(-1, x.shape[-1])
     return torch.ops.oink.rmsnorm(x_2d, weight, epsilon).view_as(x)
+
+
+oink_add_rms_supported = (
+    lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
+    and weight is not None
+    and x.dim() >= 2
+    and x.dtype == weight.dtype
+    and weight.is_contiguous()
+    and _can_view_as_2d(x)
+    and _is_oink_stride_compatible_2d(x.view(-1, x.shape[-1]))
+    # residual must have 2d-compatible strides and match x shape/dtype
+    and x.dtype == x_residual.dtype
+    and x.shape == x_residual.shape
+    and _can_view_as_2d(x_residual)
+    and _is_oink_stride_compatible_2d(x_residual.view(-1, x_residual.shape[-1]))
+)
+"""
+Oink fused_add_rms_norm has the same constraints as rms_norm,
+and residual must be 2d-like with compatible strides.
+"""
+
+
+@ir.ops.fused_add_rms_norm.register_impl(
+    "oink",
+    supports_args=oink_add_rms_supported,
+    supported=has_oink_op("fused_add_rms_norm"),
+    inplace=True,
+)
+def fused_add_rms_norm(
+    x: Tensor,
+    x_residual: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> tuple[Tensor, Tensor]:
+    assert variance_size is None
+    x_2d = x.view(-1, x.shape[-1])
+    residual_2d = x_residual.view(-1, x_residual.shape[-1])
+    torch.ops.oink.fused_add_rms_norm(x_2d, residual_2d, weight, epsilon)
+    return x, x_residual

vllm/kernels/vllm_c.py

@@ -31,3 +31,33 @@ def rms_norm(
     output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
     torch.ops._C.rms_norm(output, x, weight, epsilon)
     return output
+
+
+rms_add_no_var_size = (
+    lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
+    and (weight is None or weight.dtype == x.dtype)
+)
+"""vLLM Kernel does not support variance_size parameter and requires
+matching input/weight dtype."""
+
+
+@ir.ops.fused_add_rms_norm.register_impl(
+    "vllm_c",
+    supports_args=rms_add_no_var_size,
+    supported=CUDA_ALIKE,
+    inplace=True,
+)
+def fused_add_rms_norm(
+    x: Tensor,
+    x_residual: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> tuple[Tensor, Tensor]:
+    if weight is None:
+        # Kernel requires weight tensor, pass ones
+        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
+
+    assert variance_size is None
+    torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
+    return x, x_residual

vllm/kernels/xpu_ops.py

@@ -36,3 +36,31 @@ def rms_norm(
     output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
     torch.ops._C.rms_norm(output, x, weight, epsilon)
     return output
+
+
+rms_add_no_var_size = (
+    lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
+    and (weight is None or weight.dtype == x.dtype)
+)
+
+
+@ir.ops.fused_add_rms_norm.register_impl(
+    "xpu_kernels",
+    supports_args=rms_add_no_var_size,
+    supported=XPU_KERNELS_SUPPORTED,
+    inplace=True,
+)
+def fused_add_rms_norm(
+    x: Tensor,
+    x_residual: Tensor,
+    weight: Tensor | None,
+    epsilon: float,
+    variance_size: int | None = None,
+) -> tuple[Tensor, Tensor]:
+    if weight is None:
+        # Kernel requires weight tensor, pass ones
+        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
+
+    assert variance_size is None
+    torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
+    return x, x_residual

vllm/model_executor/layers/layernorm.py

@@ -8,68 +8,15 @@ import torch.nn.functional as F
 
 # Import kernels
 import vllm.kernels  # noqa: F401
-from vllm import _oink_ops, envs, ir
-from vllm._aiter_ops import rocm_aiter_ops
+from vllm import envs, ir
+from vllm.config import get_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
-from vllm.model_executor.layers.batch_invariant import (
-    rms_norm_batch_invariant,
-)
-from vllm.platforms import current_platform
+from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant
 
 logger = init_logger(__name__)
 
 
-def _can_view_as_2d(x: torch.Tensor) -> bool:
-    """Return True if x.view(-1, x.shape[-1]) is viewable (no copy)."""
-    if x.dim() < 2:
-        return False
-    if x.dim() == 2:
-        return True
-    # For a view(-1, N) to be valid, all leading dims must be contiguous with
-    # respect to each other (size-1 dims are ignored).
-    for dim in range(x.dim() - 1):
-        # Strides for size-1 dims are irrelevant and can be arbitrary.
-        if x.size(dim + 1) != 1 and x.stride(dim) != x.stride(dim + 1) * x.size(
-            dim + 1
-        ):
-            return False
-    return True
-
-
-def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
-    """Return True if x_2d meets Oink's pointer-path stride constraints."""
-    if x_2d.dim() != 2:
-        return False
-    if x_2d.stride(1) != 1:
-        return False
-    # Match Oink's vectorization constraint: stride(0) divisible by 256b.
-    if x_2d.dtype in (torch.float16, torch.bfloat16):
-        divby = 16
-    elif x_2d.dtype == torch.float32:
-        divby = 8
-    else:
-        return False
-    return (x_2d.stride(0) % divby) == 0
-
-
-def fused_add_rms_norm(
-    x: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    variance_epsilon: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    from vllm import _custom_ops as ops
-
-    ops.fused_add_rms_norm(
-        x,
-        residual,
-        weight,
-        variance_epsilon,
-    )
-    return x, residual
-
-
 def poly_norm(
     x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, variance_epsilon: float
 ) -> torch.Tensor:
@@ -86,18 +33,6 @@ def poly_norm(
     return out
 
 
-def dispatch_rocm_rmsnorm_func(dtype: torch.dtype, use_aiter: bool = False):
-    use_aiter = use_aiter and dtype in [
-        torch.float16,
-        torch.bfloat16,
-    ]
-
-    if use_aiter:
-        return rocm_aiter_ops.rms_norm2d_with_add
-    else:
-        return fused_add_rms_norm
-
-
 # --8<-- [start:rms_norm]
 @CustomOp.register("rms_norm")
 class RMSNorm(CustomOp):
@@ -130,105 +65,19 @@ class RMSNorm(CustomOp):
         if self.has_weight:
             self.weight = nn.Parameter(self.weight)
 
-        if current_platform.is_rocm():
-            aiter_rmsnorm_enabled = rocm_aiter_ops.is_rmsnorm_enabled()
-            self.rocm_norm_func_with_add = dispatch_rocm_rmsnorm_func(
-                dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled
-            )
-
-        # Optional: enable Oink Blackwell RMSNorm custom-op fast path on
-        # compatible CUDA devices (e.g., SM100) when the external Oink
-        # package is available. This is detected once at construction time
-        # to avoid per-call device queries in the hot path.
-        self._use_oink_fused_add_rmsnorm = False
-        if (
-            not current_platform.is_rocm()
-            and torch.cuda.is_available()
-            and bool(getattr(envs, "VLLM_USE_OINK_OPS", False))
-        ):
-            # NOTE: vLLM disables custom ops by default when using Inductor.
-            # If this op is disabled, CustomOp will dispatch to forward_native,
-            # and the Oink path in forward_cuda will never run.
-            if getattr(self._forward_method, "__func__", None) is getattr(
-                self.forward_native, "__func__", None
-            ):
-                try:
-                    from vllm.config import get_cached_compilation_config
-
-                    custom_ops = get_cached_compilation_config().custom_ops
-                except Exception:
-                    custom_ops = ["<unknown>"]
-                logger.warning_once(
-                    "VLLM_USE_OINK_OPS=1 but the `rms_norm` custom op is "
-                    "disabled (CompilationConfig.custom_ops=%s). Enable it via "
-                    "`compilation_config={'custom_ops': ['none', '+rms_norm']}` "
-                    "(or `['all']`) to let vLLM call into torch.ops.oink.*.",
-                    custom_ops,
-                )
-                # Custom op disabled => forward_cuda won't run. Avoid doing any
-                # external Oink initialization work in this case.
-            else:
-                try:
-                    device_index = torch.accelerator.current_device_index()
-                    if _oink_ops.is_oink_available_for_device(device_index):
-                        self._use_oink_fused_add_rmsnorm = (
-                            _oink_ops.has_fused_add_rms_norm()
-                        )
-                except Exception as e:
-                    # If anything goes wrong (no Oink install, CPU-only env, etc.),
-                    # silently fall back to the built-in RMSNorm path.
-                    logger.warning_once(
-                        "VLLM_USE_OINK_OPS=1 but failed to initialize Oink "
-                        "RMSNorm; falling back to vLLM RMSNorm. Error: %s",
-                        e,
-                    )
-                    self._use_oink_fused_add_rmsnorm = False
-
-    @staticmethod
-    def forward_static(
-        x: torch.Tensor,
-        variance_epsilon: float,
-        hidden_size: int,
-        orig_dtype: torch.dtype,
-        weight: torch.Tensor | None = None,
-        residual: torch.Tensor | None = None,
-        variance_size_override: int | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """PyTorch-native implementation equivalent to forward()."""
-        x = x.to(torch.float32)
-        if residual is not None:
-            # residual promoted f16->f32 automatically,
-            # otherwise Inductor eliminates the casts to and from f16,
-            # increasing memory usage (and complicating pattern matching)
-            x = x + residual
-            residual = x.to(orig_dtype)
-
-        if x.shape[-1] != hidden_size:
-            raise ValueError(
-                f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}"
-            )
-
-        if variance_size_override is None:
-            x_var = x
-        else:
-            if hidden_size < variance_size_override:
-                raise ValueError(
-                    "Expected hidden_size to be at least "
-                    f"{variance_size_override}, but found: {hidden_size}"
-                )
-
-            x_var = x[:, :, :variance_size_override]
-
-        variance = x_var.pow(2).mean(dim=-1, keepdim=True)
-
-        x = x * torch.rsqrt(variance + variance_epsilon)
-        x = x.to(orig_dtype)
-        if weight is not None:
-            x = x * weight
-        if residual is None:
-            return x
-        else:
-            return x, residual
+        # Do not pass identity weight to native implementation (causes issue on TPU).
+        # Other implementations require weight to be passed even if all ones.
+        # Cheat and predict if native will be dispatched to:
+        #  1) if native is first in priority list
+        #  2) if variance_size_override is given (only supported by native impl)
+        # TODO(luka): address weight passing inconsistency:
+        # https://github.com/vllm-project/vllm/issues/39370
+        priority = get_current_vllm_config().kernel_config.ir_op_priority
+        var_override = self.variance_size_override is not None
+        native_rms_norm = priority.rms_norm[0] == "native" or var_override
+        native_add_rms_norm = priority.fused_add_rms_norm[0] == "native" or var_override
+        self.pass_weight = self.has_weight or not native_rms_norm
+        self.pass_weight_add = self.has_weight or not native_add_rms_norm
 
     def forward_native(
         self,
@@ -237,106 +86,34 @@ class RMSNorm(CustomOp):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """PyTorch-native implementation equivalent to forward()."""
         if residual is None:
-            # TODO(luka): address the weight=None passing issue more generally
             return ir.ops.rms_norm(
                 x,
-                self.weight.data if self.has_weight else None,
+                self.weight.data if self.pass_weight else None,
+                self.variance_epsilon,
+                self.variance_size_override,
+            )
+        else:
+            return ir.ops.fused_add_rms_norm.maybe_inplace(
+                x,
+                residual,
+                self.weight.data if self.pass_weight_add else None,
                 self.variance_epsilon,
                 self.variance_size_override,
             )
-
-        return self.forward_static(
-            x,
-            self.variance_epsilon,
-            self.hidden_size,
-            x.dtype,
-            self.weight.data if self.has_weight else None,
-            residual,
-            self.variance_size_override,
-        )
 
     def forward_cuda(
         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:
-            return ir.ops.rms_norm(
-                x, self.weight.data, self.variance_epsilon, self.variance_size_override
-            )
-
-        if self.variance_size_override is not None:
-            return self.forward_native(x, residual)
-
-        # Optional Oink SM100 fast path (fused residual-add + RMSNorm, in-place).
-        # 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)
-            and x.is_cuda
-            and residual.is_cuda
-            and x.shape == residual.shape
-            and x.dtype == residual.dtype
-            and x.dim() >= 2
-            and self.has_weight
-            and not envs.VLLM_BATCH_INVARIANT
-            and self.weight.data.dtype == x.dtype
-            and self.weight.data.is_contiguous()
+            envs.VLLM_BATCH_INVARIANT
+            and residual is None
+            and self.variance_size_override is None
         ):
-            orig_shape = x.shape
-            hidden_size = orig_shape[-1]
-            if _can_view_as_2d(x) and _can_view_as_2d(residual):
-                x_2d = x.view(-1, hidden_size)
-                res_2d = residual.view(-1, hidden_size)
-
-                # The Oink in-place pointer path supports the common vLLM
-                # layout where:
-                # - `x` may be strided/padded row-major (stride(1) == 1), and
-                # - `residual` is contiguous row-major ([M, N] with stride(0) == N).
-                # If these conditions are not met, fall back to vLLM's built-in
-                # fused kernel.
-                if (
-                    _is_oink_stride_compatible_2d(x_2d)
-                    and _is_oink_stride_compatible_2d(res_2d)
-                    and res_2d.is_contiguous()
-                ):
-                    _oink_ops.fused_add_rms_norm_(
-                        x_2d,
-                        res_2d,
-                        self.weight.data,
-                        self.variance_epsilon,
-                    )
-                    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)
 
-    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:
-            return ir.ops.rms_norm(
-                x, self.weight.data, self.variance_epsilon, self.variance_size_override
-            )
-
-        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.forward_native(x, residual)
 
     def forward_xpu(
         self,

vllm/platforms/cuda.py

@@ -578,7 +578,9 @@ class CudaPlatformBase(Platform):
         if envs.VLLM_USE_OINK_OPS:
             rms_norm = ["oink"] + default
 
-        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
+        return IrOpPriorityConfig.with_default(
+            default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
+        )
 
 
 # NVML utils

vllm/platforms/rocm.py

@@ -693,21 +693,11 @@ class RocmPlatform(Platform):
     @classmethod
     def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
         from vllm._aiter_ops import rocm_aiter_ops
-        from vllm.config.compilation import CUDAGraphMode
 
         compilation_config = vllm_config.compilation_config
-        is_eager_execution = compilation_config.cudagraph_mode == CUDAGraphMode.NONE
         use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled()
-        use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()
         use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled()
         use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
-        #  Aiter rms norm perform best when CUDA Graph capture is enabled.
-        if (
-            use_aiter_rms_norm
-            and not is_eager_execution
-            and "-rms_norm" not in compilation_config.custom_ops
-        ):
-            compilation_config.custom_ops.append("+rms_norm")
 
         if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops:
             compilation_config.custom_ops.append("+quant_fp8")
@@ -939,7 +929,7 @@ class RocmPlatform(Platform):
     def get_default_ir_op_priority(
         cls, vllm_config: "VllmConfig"
     ) -> "IrOpPriorityConfig":
-        from vllm.config.compilation import CompilationMode
+        from vllm.config.compilation import CompilationMode, CUDAGraphMode
         from vllm.config.kernel import IrOpPriorityConfig
 
         # Native used by default when compiling,
@@ -949,12 +939,10 @@ class RocmPlatform(Platform):
         using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
         default = ["native"] if using_inductor else ["vllm_c", "native"]
 
-        # This (mostly) preserves previous CustomOp behavior
-        # Necessary on ROCm because it's common that users
-        # enable rms_norm to use the aiter kernel.
+        #  Aiter rms norm perform best when CUDA Graph capture is enabled.
         # TODO(luka/TJ) remove env vars completely
         if (
-            cc.is_custom_op_enabled("rms_norm")
+            cc.cudagraph_mode != CUDAGraphMode.NONE
             and envs.VLLM_ROCM_USE_AITER
             and envs.VLLM_ROCM_USE_AITER_RMSNORM
         ):
@@ -962,7 +950,9 @@ class RocmPlatform(Platform):
         else:
             rms_norm = default
 
-        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
+        return IrOpPriorityConfig.with_default(
+            default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
+        )
 
     @classmethod
     @with_amdsmi_context