diff --git a/docs/design/debug_vllm_compile.md b/docs/design/debug_vllm_compile.md index fbee9f4c3e..7edda6fa64 100644 --- a/docs/design/debug_vllm_compile.md +++ b/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). diff --git a/docs/design/vllm_ir.md b/docs/design/vllm_ir.md new file mode 100644 index 0000000000..82628f3762 --- /dev/null +++ b/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.=,,...`: + +```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 diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 3266b80e5d..2c19dc1763 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/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.""" diff --git a/tests/compile/backend.py b/tests/compile/backend.py index d61c128a59..87f98946a8 100644 --- a/tests/compile/backend.py +++ b/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 ) diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py index 1f1eeb8b47..c40d75f675 100644 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ b/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) diff --git a/tests/compile/passes/ir/test_clone_cleanup.py b/tests/compile/passes/ir/test_clone_cleanup.py new file mode 100644 index 0000000000..9fedb5fc91 --- /dev/null +++ b/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]) diff --git a/tests/compile/passes/ir/test_inplace_functionalization.py b/tests/compile/passes/ir/test_inplace_functionalization.py new file mode 100644 index 0000000000..1e8d566216 --- /dev/null +++ b/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 diff --git a/tests/compile/passes/test_functionalization.py b/tests/compile/passes/test_functionalization.py index 9a03a69887..31bf225d41 100644 --- a/tests/compile/passes/test_functionalization.py +++ b/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 [] diff --git a/tests/compile/passes/test_fuse_act_padding.py b/tests/compile/passes/test_fuse_act_padding.py index f3f3bda472..bfbe762abd 100644 --- a/tests/compile/passes/test_fuse_act_padding.py +++ b/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, ] diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 32803aad8c..2feb0bc4f7 100644 --- a/tests/compile/passes/test_fusion.py +++ b/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]) diff --git a/tests/ir/test_inplace_op.py b/tests/ir/test_inplace_op.py new file mode 100644 index 0000000000..decc4f51c7 --- /dev/null +++ b/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 + ) diff --git a/tests/ir/test_op.py b/tests/ir/test_op.py index 524497916b..3576e5aef8 100644 --- a/tests/ir/test_op.py +++ b/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( diff --git a/tests/kernels/ir/test_layernorm.py b/tests/kernels/ir/test_layernorm.py index 7510ae5010..e9661f5202 100644 --- a/tests/kernels/ir/test_layernorm.py +++ b/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 + ) diff --git a/tests/model_executor/test_enabled_custom_ops.py b/tests/model_executor/test_enabled_custom_ops.py index fc4f6f6b63..490284f439 100644 --- a/tests/model_executor/test_enabled_custom_ops.py +++ b/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 diff --git a/tests/model_executor/test_oink_integration.py b/tests/model_executor/test_oink_integration.py index d7f38fdd51..2f37472b73 100644 --- a/tests/model_executor/test_oink_integration.py +++ b/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 diff --git a/tests/test_config.py b/tests/test_config.py index 41d34a6cb0..02e4d1d5d7 100644 --- a/tests/test_config.py +++ b/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"] diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 5c036c1b38..2e9f778812 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/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(), diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index b11fc21975..45defc6926 100644 --- a/vllm/_aiter_ops.py +++ b/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, diff --git a/vllm/_oink_ops.py b/vllm/_oink_ops.py deleted file mode 100644 index c7a055410b..0000000000 --- a/vllm/_oink_ops.py +++ /dev/null @@ -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 diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 353567fd96..4e8a485c01 100644 --- a/vllm/compilation/backends.py +++ b/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 diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index e683b1dfa6..6cb0c8f49f 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/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 diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index c2490d8a21..e5130c19c3 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/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, diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index 850e434a3e..cc986595d4 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/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, diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index cdd0e23773..28159dbe08 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/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 ) diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index 35885eeb0b..2c7a1390bd 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/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) diff --git a/vllm/compilation/passes/inductor_pass.py b/vllm/compilation/passes/inductor_pass.py index b54c7bfa14..8a0d5326dd 100644 --- a/vllm/compilation/passes/inductor_pass.py +++ b/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.""" diff --git a/vllm/compilation/passes/ir/clone_elimination.py b/vllm/compilation/passes/ir/clone_elimination.py new file mode 100644 index 0000000000..61ba750a6c --- /dev/null +++ b/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) diff --git a/vllm/compilation/passes/ir/inplace_functionalization.py b/vllm/compilation/passes/ir/inplace_functionalization.py new file mode 100644 index 0000000000..e69351075b --- /dev/null +++ b/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, + ) diff --git a/vllm/compilation/passes/ir/lowering_pass.py b/vllm/compilation/passes/ir/lowering_pass.py index 02acdd1a29..f34f1c64b7 100644 --- a/vllm/compilation/passes/ir/lowering_pass.py +++ b/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: """ diff --git a/vllm/compilation/passes/ir/utils.py b/vllm/compilation/passes/ir/utils.py new file mode 100644 index 0000000000..50b4773ce5 --- /dev/null +++ b/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 diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 3dc0d7b096..5d4355a5b2 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/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()) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 93fb4c54b7..e5dcc20b50 100644 --- a/vllm/config/kernel.py +++ b/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. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8d2c2608e5..0146ee4c14 100644 --- a/vllm/config/vllm.py +++ b/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 ) diff --git a/vllm/envs.py b/vllm/envs.py index 24ec92c3d7..8378c9762a 100755 --- a/vllm/envs.py +++ b/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'). diff --git a/vllm/ir/op.py b/vllm/ir/op.py index 5d7c01be1b..841df0f9ad 100644 --- a/vllm/ir/op.py +++ b/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) diff --git a/vllm/ir/ops/__init__.py b/vllm/ir/ops/__init__.py index 25ad27c8a0..d4d71afef7 100644 --- a/vllm/ir/ops/__init__.py +++ b/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"] diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py index 981d5e3bd8..33a71b8f85 100644 --- a/vllm/ir/ops/layernorm.py +++ b/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 diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 14c2b87fbb..273bc58935 100644 --- a/vllm/kernels/aiter_ops.py +++ b/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, +) diff --git a/vllm/kernels/oink_ops.py b/vllm/kernels/oink_ops.py index e8e3cb91f8..835cd062d0 100644 --- a/vllm/kernels/oink_ops.py +++ b/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. 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 diff --git a/vllm/kernels/vllm_c.py b/vllm/kernels/vllm_c.py index 124b02e4e2..5c602c3984 100644 --- a/vllm/kernels/vllm_c.py +++ b/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 diff --git a/vllm/kernels/xpu_ops.py b/vllm/kernels/xpu_ops.py index c680c542c1..5e7f90f708 100644 --- a/vllm/kernels/xpu_ops.py +++ b/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 diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index d9184bb770..a5d4e4db79 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/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 = [""] - 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, diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 4f9b9d7bf2..9f04bf1166 100644 --- a/vllm/platforms/cuda.py +++ b/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 diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 866b9ffd1a..7200a7698d 100644 --- a/vllm/platforms/rocm.py +++ b/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