A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 08:45 UTC
[vLLM IR] 2/N fused_add_rms_norm and maybe_inplace overload (#36823)
Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Luka Govedič committed
May 1, 2026 at 23:41 UTC
d58c42e19cb792e24eb335b75164356a4f71bff0
44 files changed
+2837
-783
docs/design/debug_vllm_compile.md
+25
-7
@@ -5,12 +5,14 @@ TL;DR:
5
- use tlparse to acquire torch.compile logs. Include these logs in bug reports and/or support asks.
6
- The vLLM-torch.compile integration is multiple pieces. vLLM exposes flags to turn off each piece:
7
8
-| Online Flag | Offline Flag | Result |
9
-| ----------- | ------------ | ------ |
10
-| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs |
11
-| -cc.mode=0 | mode=CompilationMode.NONE | Turn off torch.compile only |
12
-| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only |
13
-| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor |
8
+| Online Flag | Offline Flag | Result |
9
+|--------------------------------|--------------------------------------------------------------------------------|------------------------------------------------------|
10
+| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs |
11
+| -cc.mode=0 | compilation_config=CompilationConfig(mode=CompilationMode.NONE) | Turn off torch.compile only |
12
+| -cc.mode=1 | compilation_config=CompilationConfig(mode=CompilationMode.STOCK_TORCH_COMPILE) | Turn off vLLM-compile modifications to torch.compile |
13
+| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only |
14
+| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor |
15
+| -cc.ir_enable_torch_wrap=False | compilation_config=CompilationConfig(ir_enable_torch_wrap=False) | Turn off vLLM IR wrapping |
16
17
## vLLM-torch.compile overview
18
@@ -22,7 +24,7 @@ Most notably, vLLM-compile is NOT torch.compile, it is a custom compiler built u
24
25
- Given a model, we do a full graph capture via TorchDynamo that is dynamic on the batch size (number of tokens)
26
- vLLM then optionally splits and/or specializes this graph and then uses TorchInductor to compile each graph into a compiled artifact.
25
-This step may use vLLM custom Inductor passes to further optimize the graph.
27
+This step may use vLLM custom Inductor passes to further optimize the graph. This includes vLLM IR lowering to remove dispatch overhead.
28
- The compiled artifact is saved to vLLM's compile cache so that it can be loaded in the future.
29
- vLLM applies CUDAGraphs to reduce CPU overheads.
30
@@ -34,6 +36,7 @@ For more details on the design, please see the following resources:
36
37
- [Introduction to vLLM-torch.compile blogpost](https://blog.vllm.ai/2025/08/20/torch-compile.html)
38
- [vLLM-torch.compile integration design](./torch_compile.md)
39
+- [vLLM IR design](./vllm_ir.md)
40
- [vLLM Office Hours #26](https://www.youtube.com/live/xLyxc7hxCJc?si=Xulo9pe53C6ywf0V&t=561)
41
- [Talk at PyTorch Conference 2025](https://youtu.be/1wV1ESbGrVQ?si=s1GqymUfwiwOrDTg&t=725)
42
@@ -117,6 +120,21 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode
120
LLM(model, compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE))
121
```
122
123
+vLLM IR makes heavy use of the compilation pipeline, from functionalization, custom fusions, and lowering.
124
+To turn that off and capture eager-mode dispatching behavior of vLLM IR, run with `ir_enable_torch_wrap=False`.
125
+IR torch wrap is only enabled by default when using `mode=VLLM_COMPILE` and `backend="inductor"` (default).
126
+
127
+```sh
128
+# Online
129
+vllm serve -cc.ir_enable_torch_wrap=False
130
+```
131
+
132
+```py
133
+# Offline
134
+from vllm.config.compilation import CompilationConfig
135
+LLM(model, compilation_config=CompilationConfig(ir_enable_torch_wrap=False))
136
+```
137
+
138
## Debugging TorchDynamo
139
140
vLLM requires model code be capturable into a full graph via TorchDynamo (torch.compile's frontend).
docs/design/vllm_ir.md
new
+615
@@ -0,0 +1,615 @@
1
+# vLLM IR: Functional Intermediate Representation
2
+
3
+## Motivation
4
+
5
+vLLM IR is a **functional intermediate representation (IR)** that fills the gap between
6
+low-level `torch` ops and vLLM layers like `RMSNorm` and quantization operators,
7
+By separating operator **semantics** from the **implementation** and **dispatching**,
8
+vLLM IR simplifies both compilation and kernel registration & dispatching simultaneously.
9
+It operates as a **dialect** in the torch FX representation, allowing full interoperability
10
+with “regular” torch ops & custom torch ops/kernels, as well as a piecewise migration from
11
+the previous `CustomOp` approach.
12
+
13
+Key design principles:
14
+
15
+- **Eager-compile consistency**: identical behavior (barring minor numerics) in eager and compiled modes
16
+- **Simple, transparent, yet powerful kernel selection**: good visibility and control allowing easy debugging
17
+- **Convention over configuration**: near-zero boilerplate required to register ops and implementations
18
+- **Extensibility**: ops and implementations can be registered anywhere, in-tree or out-of-tree
19
+- **Interoperability**: fully compatible with “regular” torch ops & custom torch ops/kernels,
20
+reducing developer friction and allowing piecewise migration
21
+
22
+The clean semantics/implementation separation enables a unified and extensible dispatching mechanism,
23
+allowing multiple kernels per-platform and powerful kernel selection. The separation also facilitates
24
+cleaner testing and benchmarking, removing much of the boilerplate standard for legacy approaches.
25
+
26
+By delaying kernel selection until late in the compilation process, the compiler can operate on
27
+a higher-level representation, which has the following main benefits:
28
+
29
+- Pattern matching in fusion/transformation passes only requires a single, simple pattern per op
30
+- OOT compiler backends can lower from the higher-level representation (in-progress)
31
+- The compiler can autotune over available implementations (future feature)
32
+
33
+## Quick Overview
34
+
35
+### Declaring an IR Operation
36
+
37
+IR operations are declared using the `@register_op` decorator with a native PyTorch implementation that defines the op's semantics:
38
+
39
+```python
40
+# vllm/ir/ops/layernorm.py
41
+from torch import Tensor
42
+from vllm.ir import register_op
43
+
44
+@register_op
45
+def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
46
+ """Weighted root-mean-square layer normalization"""
47
+ orig_dtype = x.dtype
48
+ x = x.to(torch.float32)
49
+ x_var = x if variance_size is None else x[..., :variance_size]
50
+ variance = x_var.pow(2).mean(dim=-1, keepdim=True)
51
+ x = x * torch.rsqrt(variance + epsilon)
52
+ x = x.to(orig_dtype)
53
+ if weight is not None:
54
+ x = x * weight
55
+ return x
56
+```
57
+
58
+The native implementation serves three purposes:
59
+
60
+1. **Semantic definition**: Specifies the exact semantics of the operation, including shapes and strides
61
+2. **Default implementation**: Used when no other (better) implementation is available
62
+3. **Reference for testing**: Other implementations must match these semantics
63
+
64
+### Registering Implementations
65
+
66
+Kernel implementations are registered using the `register_impl` decorator on the IR op object:
67
+
68
+```python
69
+# vllm/kernels/vllm_c.py
70
+from vllm import ir
71
+
72
+rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
73
+
74
+@ir.ops.rms_norm.register_impl("vllm_c", supports_args=rms_norm_no_var, supported=current_platform.is_cuda_alike())
75
+def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
76
+ output = torch.empty_like(x)
77
+ torch.ops._C.rms_norm(output, x, weight, epsilon)
78
+ return output
79
+```
80
+
81
+Implementations can specify:
82
+
83
+- `supported`: Static boolean indicating if this implementation is available
84
+- `supports_args`: Function checking if the implementation supports specific arguments
85
+- `inplace`: Whether this implementation reuses input memory for outputs
86
+
87
+### Using IR Operations in Models
88
+
89
+IR operations are imported and called directly in model code:
90
+
91
+```python
92
+# vllm/model_executor/layers/layernorm.py
93
+from vllm import ir
94
+
95
+class RMSNorm(nn.Module):
96
+ def __init__(self, hidden_size: int, eps: float = 1e-6):
97
+ super().__init__()
98
+ self.weight = nn.Parameter(torch.ones(hidden_size))
99
+ self.variance_epsilon = eps
100
+
101
+ def forward(self, x: Tensor, residual: Tensor | None = None):
102
+ if residual is None:
103
+ return ir.ops.rms_norm(x, self.weight, self.variance_epsilon)
104
+
105
+ # Use maybe_inplace overload to allow implementation to reuse input memory for outputs
106
+ # (using x or residual after this call is undefined behavior)
107
+ return ir.ops.fused_add_rms_norm.maybe_inplace(
108
+ x, residual, self.weight, self.variance_epsilon
109
+ )
110
+```
111
+
112
+### Configuring Kernel Selection
113
+
114
+Kernel selection is controlled via priority lists in the configuration.
115
+Priority lists specify the order in which implementations are considered,
116
+with the first supported implementation being selected.
117
+This includes the static support check (`supported=...`) and
118
+the dynamic arg support check (`supports_args=...`).
119
+
120
+#### Command Line Configuration
121
+
122
+Use `--ir-op-priority.<op_name>=<provider1>,<provider2>,...`:
123
+
124
+```bash
125
+# CUDA: Use vllm_c implementation for rms_norm
126
+vllm serve meta-llama/Llama-3.2-1B \
127
+ --ir-op-priority.rms_norm=vllm_c
128
+
129
+# ROCm: Try aiter first, fall back to vllm_c, then native
130
+vllm serve meta-llama/Llama-3.2-1B \
131
+ --ir-op-priority.rms_norm=aiter,vllm_c,native
132
+
133
+# Configure multiple operations
134
+vllm serve meta-llama/Llama-3.2-1B \
135
+ --ir-op-priority.rms_norm=vllm_c \
136
+ --ir-op-priority.fused_add_rms_norm=vllm_c
137
+```
138
+
139
+#### Python Configuration
140
+
141
+```python
142
+from vllm import LLM
143
+from vllm.config import VllmConfig, KernelConfig
144
+
145
+llm = LLM(
146
+ model="meta-llama/Llama-3.2-1B",
147
+ vllm_config=VllmConfig(
148
+ kernel_config=KernelConfig(
149
+ ir_op_priority={
150
+ "rms_norm": ["vllm_c", "native"],
151
+ "fused_add_rms_norm": ["vllm_c", "native"],
152
+ }
153
+ )
154
+ )
155
+)
156
+```
157
+
158
+#### Platform Defaults
159
+
160
+Each platform provides default priority lists that are automatically applied:
161
+
162
+```python
163
+# CUDA/XPU/ROCm platform defaults (when compiling with Inductor)
164
+{
165
+ "rms_norm": ["native"], # Native torch is default
166
+ "fused_add_rms_norm": ["native"],
167
+}
168
+
169
+# CUDA platform defaults (eager or Dynamo-only)
170
+{
171
+ "rms_norm": ["vllm_c", "native"],
172
+ "fused_add_rms_norm": ["vllm_c", "native"],
173
+}
174
+
175
+# ROCm platform defaults (future - currently same as CUDA)
176
+{
177
+ "rms_norm": ["aiter", "vllm_c", "native"],
178
+ "fused_add_rms_norm": ["aiter", "vllm_c", "native"],
179
+}
180
+
181
+# XPU platform defaults (eager or Dynamo-only)
182
+{
183
+ "rms_norm": ["xpu_kernels", "native"],
184
+ "fused_add_rms_norm": ["xpu_kernels", "native"],
185
+}
186
+```
187
+
188
+User-specified priorities are prepended to platform defaults,
189
+so you only need to specify the out-of-order implementations,
190
+other implementations are appended automatically.
191
+
192
+## Compilation Pipeline
193
+
194
+vLLM IR heavily customizes the `torch.compile`-based compilation process to allow custom compile
195
+passes to operate on high-level IR while still producing efficient low-level code at the end.
196
+The compilation pipeline consists of several stages:
197
+
198
+### 1. Dynamo Tracing
199
+
200
+When `torch.compile` traces the model's forward pass, vLLM IR operations appear as custom operations
201
+in the `vllm_ir` torch library. These operations are opaque to Dynamo, meaning they appear directly
202
+in the FX graph without decomposition:
203
+
204
+```python
205
+# Python code (epsilon=1e-5)
206
+x1 = ir.ops.rms_norm(x, weight, epsilon)
207
+x2, residual_out = ir.ops.fused_add_rms_norm.maybe_inplace(x1, residual, weight, epsilon)
208
+
209
+# FX graph after Dynamo tracing
210
+x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
211
+out = torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace(x1, residual, weight, 1e-5); x1 = residual = None
212
+x2 = out[0]
213
+residual_out = out[1]
214
+```
215
+
216
+### 2. AOTAutograd and Functionalization
217
+
218
+AOTAutograd functionalizes the graph, converting any mutating operations to functional equivalents.
219
+For vLLM IR operations with `maybe_inplace` overloads, we perform this manually before AOTAutograd,
220
+converting them to the functional `default` overload using the pre-grad custom pass hook.
221
+
222
+```python
223
+# After functionalization
224
+x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
225
+out = torch.ops.vllm_ir.fused_add_rms_norm.default(x1, residual, weight, 1e-5); x1 = residual = None
226
+x2 = out[0]
227
+residual_out = out[1]
228
+```
229
+
230
+The pass also tracks which inputs were "donated" (passed to `maybe_inplace`),
231
+storing this information in vLLM's `PassContext` for later use in clone elimination.
232
+
233
+### 3. IR Fusion and Transformation Passes
234
+
235
+After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations.
236
+These passes can perform fusion, distribute operations for sequence parallelism, and other transformations:
237
+
238
+```python
239
+# Example: Sequence Parallelism (see SequenceParallelismPass)
240
+# Before SP pass
241
+
242
+all_reduce = torch.ops.vllm.all_reduce(x, "tp:0")
243
+rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
244
+
245
+# after SP pass
246
+reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0")
247
+rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
248
+all_gather = torch.ops.vllm.all_gather(x, "tp:0")
249
+```
250
+
251
+Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations,
252
+handle different kernel implementations separately, or deal with functionalization of custom kernels.
253
+
254
+### 4. IR Lowering
255
+
256
+The lowering pass (`VllmIRLoweringPass`) replaces each vLLM IR operation with its selected implementation.
257
+The implementation is chosen based on the priority list and support predicates,
258
+using the **fake tensors** in the graph's metadata in place of op arguments:
259
+
260
+```python
261
+# Implementation selection, same in eager dispatch and compile lowering
262
+def dispatch(*args) -> IrOpImpl:
263
+ for provider in priority_list: # e.g., ["vllm_c", "native"]
264
+ impl = ir_op.impls[provider]
265
+ if not impl.supported:
266
+ continue
267
+ if impl.supports_args and not impl.supports_args(*args):
268
+ continue
269
+ return impl
270
+
271
+# make_fx uses torch.fx.symbolic_trace
272
+impl_graph = make_fx(selected_impl.impl_fn)
273
+# Replace IR op node with impl_graph's nodes
274
+match.replace_by_example(selected_impl.impl_fn, node.args)
275
+```
276
+
277
+For example, lowering `rms_norm` with the `vllm_c` implementation:
278
+
279
+```python
280
+# Before lowering (IR op)
281
+rms_norm = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5)
282
+
283
+# After lowering (vllm_c implementation traced)
284
+# Note: Lowering does not currently functionalize, this will likely change in the future.
285
+empty = torch.ops.aten.empty.memory_format(x.shape, ...)
286
+rms_norm = torch.ops._C.rms_norm(empty, x, weight, 1e-5)
287
+```
288
+
289
+When lowering an implementation that mutates inputs (`inplace=True`),
290
+the lowering pass inserts clones to preserve functional semantics:
291
+
292
+```python
293
+# vllm_c implementation for fused_add_rms_norm mutates its first two arguments
294
+# Lowered with clones for safety
295
+clone_default = torch.ops.aten.clone.default(x)
296
+clone_default_1 = torch.ops.aten.clone.default(residual)
297
+fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(clone_default, clone_default_1, weight, 1e-5)
298
+```
299
+
300
+### 5. Clone Cleanup
301
+
302
+After lowering, the clone elimination pass (`UnsafeCloneEliminationPass`) removes unnecessary clones introduced during lowering.
303
+This pass is essential for achieving zero-copy behavior when using in-place kernels with `maybe_inplace`.
304
+The pass removes a clone if:
305
+
306
+- the cloned input is created in the graph and not used again in the graph
307
+- the cloned input is a graph parameter, marked as donated
308
+
309
+```python
310
+# After cleanup (donated inputs, no subsequent uses)
311
+fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(x, residual, weight, 1e-5)
312
+```
313
+
314
+The combination of inplace functionalization (tracking donated inputs) and clone cleanup enables the compiler to safely
315
+use in-place kernels without adding redundant copies or increasing the memory usage.
316
+
317
+### 6. Inductor Optimization and Codegen
318
+
319
+After IR lowering and cleanup, the graph contains only standard PyTorch operations and platform-specific custom ops.
320
+Inductor then performs its standard codegen:
321
+
322
+- **Inductor lowering and pointwise fusion**: Fusing element-wise operations, reductions, etc.
323
+- **Memory planning**: Determining buffer allocation and reuse
324
+- **Kernel generation**: Generating Triton or C++ code for fused operations
325
+- **Autotuning**: Selecting the best kernel configurations
326
+
327
+### Pipeline Summary
328
+
329
+```text
330
+Model Forward Pass
331
+ ↓
332
+[Dynamo Tracing] → FX Graph with vllm_ir.* ops
333
+ ↓
334
+[Pre-grad: Inplace Functionalization] → maybe_inplace → default, track donated inputs
335
+ ↓
336
+[AOTAutograd] → Functionalization
337
+ ↓
338
+[Post-grad: IR Fusion Passes] → Fuse high-level IR ops (e.g., rms_norm + quant)
339
+ ↓
340
+[Post-grad: IR Lowering] → vllm_ir.* ops → impl ops (with clones if needed)
341
+ ↓
342
+[Post-grad: Clone Cleanup] → Remove unnecessary clones using donated input info
343
+ ↓
344
+[Inductor] → Pattern matching, fusion, memory planning, codegen
345
+ ↓
346
+Compiled Code
347
+```
348
+
349
+## Core vLLM IR Concepts
350
+
351
+### Operation Declaration
352
+
353
+Operations are declared with the `@register_op` decorator, which creates an `IrOp` object:
354
+
355
+```python
356
+@register_op(
357
+ name=None, # Operation name (defaults to function name)
358
+ activations=None, # List of activation parameters (defaults to params starting with 'x')
359
+ allow_inplace=False, # Whether to create a maybe_inplace overload
360
+)
361
+def op_name(...):
362
+ ...
363
+```
364
+
365
+**Parameters:**
366
+
367
+- `activations`: List of parameter names considered "activations" (typically consumed by `maybe_inplace`). Defaults to parameters starting with `x`.
368
+- `allow_inplace`: Creates a `maybe_inplace` overload for memory-efficient execution (see below).
369
+
370
+### The `maybe_inplace` Overload
371
+
372
+The `maybe_inplace` overload is a critical feature for memory efficiency in LLM inference.
373
+It signals that the caller doesn't need to preserve the activation inputs after the operation,
374
+allowing in-place implementations to reuse input memory for outputs.
375
+
376
+#### Semantics and Usage
377
+
378
+```python
379
+# Standard usage: inputs are preserved
380
+out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
381
+# x and residual are unchanged, out and res_out are new tensors
382
+
383
+# maybe_inplace: inputs may be modified
384
+out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
385
+# x and residual may be modified (undefined behavior to use them after this)
386
+# out and res_out may alias x and residual
387
+```
388
+
389
+Using an activation input after passing it to `maybe_inplace` is **undefined behavior**:
390
+
391
+```python
392
+# WRONG: Using x after donating it
393
+out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
394
+result = out + x # ERROR: x was donated!
395
+```
396
+
397
+If you need to preserve an input, either use the default overload or clone manually:
398
+
399
+```python
400
+# Option 1: Use default overload
401
+out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
402
+result = out + x # OK: x is preserved
403
+
404
+# Option 2: Clone before maybe_inplace
405
+out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x.clone(), residual, weight, epsilon)
406
+result = out + x # OK: x is preserved, clone was donated
407
+```
408
+
409
+#### Compilation Behavior
410
+
411
+During compilation, the inplace functionalization pass validates that donated inputs are
412
+not used again and converts `maybe_inplace` to the functional `default` overload:
413
+
414
+```python
415
+# Inplace functionalization pass (pre-grad)
416
+for node in graph.nodes:
417
+ if node.target == torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace:
418
+ # Check that activation inputs aren't used after this node
419
+ for activation_arg in activation_inputs:
420
+ for user in activation_arg.users:
421
+ if user appears after node:
422
+ raise ValueError(f"Input {activation_arg} donated but used again")
423
+
424
+ # Convert to default overload
425
+ node.target = torch.ops.vllm_ir.fused_add_rms_norm.default
426
+
427
+ # Track donated graph inputs for later clone elimination
428
+ for i, arg in enumerate(node.args):
429
+ if arg.op == "placeholder" and i in activation_indices:
430
+ pass_context.donated_input_ids.add(node_to_idx[arg])
431
+```
432
+
433
+The donated input information is then used by the clone cleanup pass to eliminate
434
+unnecessary copies when in-place kernels are lowered.
435
+
436
+#### Eager Mode Behavior
437
+
438
+In eager mode (without `torch.compile`), `maybe_inplace` enables **maximally memory-efficient**
439
+execution by allowing the IR operation to dispatch directly to in-place implementations:
440
+
441
+```python
442
+# Eager dispatch logic for maybe_inplace
443
+impl: IrOpImpl = ir_op.dispatch(*args)
444
+return impl.impl_fn(*args)
445
+
446
+# Eager dispatch logic for default:
447
+impl: IrOpImpl = ir_op.dispatch(*args)
448
+if impl.inplace:
449
+ args = [
450
+ arg.clone() if i in ir_op.activations else arg
451
+ for i, arg in enumerate(args)
452
+ ]
453
+return impl.impl_fn(*args)
454
+```
455
+
456
+The combination of `maybe_inplace` in model code and in-place kernel implementations provides optimal memory efficiency
457
+in both eager and compiled modes, with identical semantics in both cases.
458
+
459
+#### Memory Savings Example
460
+
461
+Consider a transformer layer with residual connections:
462
+
463
+```python
464
+# Without maybe_inplace (2 allocations per layer)
465
+hidden_states = self.attention(input)
466
+normed, residual = ir.ops.fused_add_rms_norm(hidden_states, input, weight, eps)
467
+# Memory: input (preserved), hidden_states (preserved), normed (new), residual (new)
468
+
469
+# With maybe_inplace (0 allocations per layer when using in-place kernel)
470
+hidden_states = self.attention(input)
471
+normed, residual = ir.ops.fused_add_rms_norm.maybe_inplace(hidden_states, input, weight, eps)
472
+# Memory: normed (reuses hidden_states), residual (reuses input)
473
+```
474
+
475
+### Implementation Registration
476
+
477
+Implementations are registered using the `register_impl` method:
478
+
479
+```python
480
+@ir.ops.op_name.register_impl(
481
+ provider="provider_name", # Unique identifier (e.g., "vllm_c", "aiter", "triton")
482
+ supported=True, # Static availability check
483
+ supports_args=None, # Dynamic argument support check
484
+)
485
+def impl_fn(...):
486
+ ...
487
+```
488
+
489
+**Provider naming conventions:**
490
+
491
+- `native`: Reserved for the native torch implementation (declared with `@register_op`)
492
+- `vllm_c`: C++/CUDA kernels via `torch.ops._C`
493
+- `aiter`: AMD AITER library
494
+- `xpu_kernels`: SYCL/SYCLTLA kernels implemented in `vllm-xpu-kernels`
495
+- `triton_*`: Triton kernels
496
+- Platform/library names for other implementations
497
+
498
+**Support checking:**
499
+
500
+- `supported`: Static boolean, checked once at import time (e.g., `HAS_TRITON`, `is_cuda_alike()`)
501
+- `supports_args`: Function `(*args, **kwargs) -> bool` checking argument compatibility
502
+ - Called with **fake tensors** during compilation for zero-cost checking
503
+ - Called with **real tensors** during eager mode dispatch
504
+ - Should NOT check batch sizes or add guards based on values
505
+
506
+Example support predicate:
507
+
508
+```python
509
+def aiter_rms_norm_supports(x, weight, epsilon, variance_size=None):
510
+ # Check dtype (OK: doesn't depend on batch size)
511
+ if x.dtype not in [torch.float16, torch.bfloat16]:
512
+ return False
513
+ # Check optional parameter (OK: static check)
514
+ if variance_size is not None:
515
+ return False
516
+ return True
517
+
518
+@ir.ops.rms_norm.register_impl("aiter", supports_args=aiter_rms_norm_supports)
519
+def rms_norm(...):
520
+ ...
521
+```
522
+
523
+Batch-invariant kernels are automatically selected when `VLLM_BATCH_INVARIANT=1` is set.
524
+
525
+### Eager Mode vs Compile Mode
526
+
527
+vLLM IR operations behave identically in eager and compile modes:
528
+
529
+**Eager mode:**
530
+
531
+- Direct dispatch to implementation based on priority list
532
+- Support checked with real tensor arguments
533
+- Minimal overhead (can be optimized further if needed)
534
+
535
+**Compile mode:**
536
+
537
+- IR ops appear in FX graph as `torch.ops.vllm_ir.*` custom ops
538
+- Lowering selects implementation using fake tensors
539
+- Full integration with Inductor optimizations
540
+
541
+This consistency enables:
542
+
543
+- Prototyping in eager mode with confidence
544
+- Debugging by disabling compilation
545
+- Gradual migration from eager to compiled execution
546
+
547
+## Other Topics
548
+
549
+### Out-of-Tree Implementations
550
+
551
+External platforms can register implementations without modifying vLLM:
552
+
553
+```python
554
+# In external package
555
+from vllm import ir
556
+
557
+@ir.ops.rms_norm.register_impl("my_platform", supported=is_my_platform())
558
+def rms_norm(x, weight, epsilon, variance_size=None):
559
+ return my_platform.rms_norm(x, weight, epsilon)
560
+```
561
+
562
+Then configure priority to use your implementation:
563
+
564
+```python
565
+class MyPlatform(Platform):
566
+ def get_default_ir_op_priority(self):
567
+ return IrOpPriorityConfig(rms_norm=['my_platform', 'native'])
568
+
569
+# Users can still override priority in the same way
570
+llm = LLM(ir_op_priority=IrOpPriorityConfig(rms_norm=['custom_oot_kernel']))
571
+```
572
+
573
+### Debugging and Observability
574
+
575
+!!! note
576
+ Please let us know how observability can be improved for your use-case!
577
+
578
+Enable debug logging to see kernel selection:
579
+
580
+```bash
581
+VLLM_LOGGING_LEVEL=DEBUG vllm serve ...
582
+```
583
+
584
+This logs:
585
+
586
+- Which implementations are selected for each operation
587
+- Why implementations were rejected (unsupported, args not supported)
588
+- Compilation cache hits/misses
589
+- IR lowering statistics
590
+
591
+Check selected implementations in compiled graphs:
592
+
593
+```python
594
+# After compilation, inspect the lowering pass
595
+lowering_pass = backend.lowering_pass
596
+print(lowering_pass.selected_impls)
597
+# Output: {'rms_norm': {'node_123': 'vllm_c', 'node_456': 'vllm_c'}}
598
+```
599
+
600
+## Migration from CustomOp
601
+
602
+vLLM IR is designed to coexist with and gradually replace `CustomOp`:
603
+
604
+1. **Op declaration**: Convert `CustomOp` class `PluggableLayer` and move `forward_native` to `@register_op` function
605
+2. **Implementation registration**: Use `@ir.ops.op_name.register_impl` instead of overriding methods
606
+3. **Layer usage**: Replace `self.op(...)` with `ir.ops.op_name(...)`
607
+4. **Configuration**: Migrate `--compilation-config.custom-ops` to `--ir-op-priority`
608
+
609
+The migration can be done incrementally, one operation at a time.
610
+
611
+## See Also
612
+
613
+- [torch.compile Integration](torch_compile.md) - General compilation infrastructure
614
+- [Fusions](fusions.md) - Custom fusion and transformation passes in vLLM
615
+- [Custom Operations](custom_op.md) - Legacy custom op system
docs/mkdocs/hooks/generate_argparse.py
+27
-1
@@ -7,7 +7,7 @@ import sys
7
import textwrap
8
import traceback
9
from argparse import SUPPRESS, Action, HelpFormatter
10
-from collections.abc import Iterable
10
+from collections.abc import Callable, Iterable
11
from importlib.machinery import ModuleSpec
12
from pathlib import Path
13
from typing import TYPE_CHECKING, Literal
@@ -48,6 +48,7 @@ class MockPluggableLayer:
48
49
50
mock_if_no_torch("vllm._C", MagicMock())
51
+mock_if_no_torch("vllm._C_stable_libtorch", MagicMock())
52
mock_if_no_torch(
53
"vllm.model_executor.custom_op",
54
MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
@@ -67,6 +68,31 @@ importlib.metadata.version = lambda name: VERSIONS.get(name) or "0.0.0"
68
mock_if_no_torch("torch.nn", MagicMock(Parameter=object))
69
70
71
+# Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__
72
+# We need to return the corresponding number of inputs, as IR infra will assert it
73
+def get_outputs(native_fn: Callable) -> str:
74
+ """
75
+ Extract output schema from function's return type annotation,
76
+ e.g. 'Tensor' or 'Tensor, Tensor'.
77
+ """
78
+ import typing
79
+
80
+ return_type = typing.get_type_hints(native_fn)["return"]
81
+ origin = typing.get_origin(return_type)
82
+ arg_name = lambda a: a.__name__ if hasattr(a, "__name__") else str(a)
83
+ if origin is tuple:
84
+ args = typing.get_args(return_type)
85
+ return ", ".join(arg_name(arg) for arg in args)
86
+ else:
87
+ return f"{arg_name(return_type)}"
88
+
89
+
90
+mock_if_no_torch(
91
+ "torch.library",
92
+ MagicMock(infer_schema=lambda fn, **k: f"(Tensor x) -> {get_outputs(fn)}"),
93
+)
94
+
95
+
96
class PydanticMagicMock(MagicMock):
97
"""`MagicMock` that's able to generate pydantic-core schemas."""
98
tests/compile/backend.py
+15
-2
@@ -12,10 +12,17 @@ from torch._ops import OpOverload, OpOverloadPacket
12
from torch.fx._utils import lazy_format_graph_code
13
14
from vllm.compilation.passes.fx_utils import find_op_nodes
15
-from vllm.compilation.passes.inductor_pass import InductorPass
15
+from vllm.compilation.passes.inductor_pass import (
16
+ InductorPass,
17
+ pass_context,
18
+)
19
+from vllm.compilation.passes.ir.inplace_functionalization import (
20
+ VllmIRInplaceFunctionalizationPass,
21
+)
22
from vllm.compilation.passes.pass_manager import with_pattern_match_debug
23
from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
24
from vllm.config import VllmConfig, get_current_vllm_config
25
+from vllm.config.utils import Range
26
from vllm.logger import init_logger
27
28
logger = init_logger("vllm.tests.compile.backend")
@@ -53,11 +60,17 @@ class TestBackend:
60
self.custom_passes = list(passes)
61
vllm_config = get_current_vllm_config()
62
compile_config = vllm_config.compilation_config
63
+ self.range = Range(1, vllm_config.scheduler_config.max_num_batched_tokens)
64
# Deepcopy to allow multiple TestBackend instances to use the same VllmConfig
65
self.inductor_config = deepcopy(compile_config.inductor_compile_config)
66
self.inductor_config["force_disable_caches"] = True
67
self.inductor_config["post_grad_custom_post_pass"] = self.post_pass
68
69
+ # Add VllmIRInplaceFunctionalizationPass as pre-grad pass by default
70
+ self.inductor_config["pre_grad_custom_pass"] = (
71
+ VllmIRInplaceFunctionalizationPass(vllm_config)
72
+ )
73
+
74
if debug_dump_path := vllm_config.compile_debug_dump_path():
75
logger.debug("Dumping depyf output to %s", debug_dump_path)
76
self.debug_ctx = depyf.prepare_debug(debug_dump_path.as_posix())
@@ -68,7 +81,7 @@ class TestBackend:
81
self.graph_pre_compile = deepcopy(graph)
82
from torch._inductor.compile_fx import compile_fx
83
71
- with self.debug_ctx:
84
+ with self.debug_ctx, pass_context(self.range):
85
return compile_fx(
86
graph, example_inputs, config_patches=self.inductor_config
87
)
tests/compile/passes/distributed/test_sequence_parallelism.py
+13
-16
@@ -88,14 +88,10 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
88
]
89
90
def ops_in_model(self):
91
- return (
92
- [torch.ops.vllm_ir.rms_norm]
93
- + [
94
- torch.ops._C.fused_add_rms_norm.default,
95
- ]
96
- if RMSNorm.enabled()
97
- else []
98
- )
91
+ return [
92
+ torch.ops.vllm_ir.rms_norm,
93
+ torch.ops.vllm_ir.fused_add_rms_norm,
94
+ ]
95
96
97
class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
@@ -152,16 +148,17 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
148
def ops_in_model(self):
149
if self.vllm_config.compilation_config.pass_config.fuse_norm_quant:
150
return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
155
- elif RMSNorm.enabled():
156
- return [
157
- torch.ops._C.fused_add_rms_norm.default,
158
- ]
159
- elif any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers):
151
+ else:
152
+ quant_ops = (
153
+ [torch.ops._C.static_scaled_fp8_quant.default]
154
+ if any(layer.is_quant_fp8_enabled() for layer in self.fp8_linear_layers)
155
+ else [torch.ops.aten.reciprocal]
156
+ )
157
return [
161
- torch.ops._C.static_scaled_fp8_quant.default,
158
+ torch.ops.vllm_ir.rms_norm,
159
+ torch.ops.vllm_ir.fused_add_rms_norm,
160
+ *quant_ops,
161
]
163
- else:
164
- return []
162
163
164
@multi_gpu_test(num_gpus=2)
tests/compile/passes/ir/test_clone_cleanup.py
new
+412
@@ -0,0 +1,412 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Comprehensive tests for UnsafeCloneEliminationPass.
5
+
6
+This test suite exercises all possible valid FX graph patterns involving clones:
7
+1. Clone with no users (dead code)
8
+2. Clone with read-only users
9
+3. Clone with mutation users
10
+4. Clone of graph input
11
+5. Clone with original used after mutation
12
+6. Clone chains
13
+"""
14
+
15
+import pytest
16
+import torch
17
+from torch import fx
18
+from torch.fx.experimental.proxy_tensor import make_fx
19
+
20
+from vllm.compilation.passes.fx_utils import find_op_nodes
21
+from vllm.compilation.passes.inductor_pass import get_pass_context, pass_context
22
+from vllm.compilation.passes.ir.clone_elimination import (
23
+ UnsafeCloneEliminationPass,
24
+ user_writes_to_node,
25
+)
26
+from vllm.config import VllmConfig
27
+from vllm.config.utils import Range
28
+
29
+
30
+def count_clones(graph: fx.Graph) -> int:
31
+ """Count clone nodes in a graph."""
32
+ return len(list(find_op_nodes(torch.ops.aten.clone.default, graph)))
33
+
34
+
35
+@pytest.fixture(scope="function")
36
+def clone_cleanup_pass():
37
+ return UnsafeCloneEliminationPass(VllmConfig())
38
+
39
+
40
+@pytest.fixture(autouse=True)
41
+def setup_pass_context():
42
+ """Set up pass context for each test."""
43
+ with pass_context(compile_range=Range(1, 8192)):
44
+ yield
45
+
46
+
47
+class TestCloneCleanup:
48
+ """Test UnsafeCloneEliminationPass behavior on various graph patterns."""
49
+
50
+ def test_remove_clone_readonly_users(self, clone_cleanup_pass):
51
+ """Clone with only read-only users should be removed."""
52
+
53
+ def f(x: torch.Tensor) -> torch.Tensor:
54
+ x_clone = x.clone()
55
+ return x_clone + 1
56
+
57
+ inp = torch.randn(2, 3)
58
+ graph_module = make_fx(f)(inp)
59
+ assert count_clones(graph_module.graph) == 1
60
+
61
+ expected = graph_module(inp)
62
+ clone_cleanup_pass(graph_module.graph)
63
+ graph_module.recompile()
64
+ actual = graph_module(inp)
65
+
66
+ assert count_clones(graph_module.graph) == 0
67
+ torch.testing.assert_close(actual, expected)
68
+
69
+ def test_keep_clone_with_mutation_and_original_used_after(self, clone_cleanup_pass):
70
+ """Clone must be kept if it's mutated AND original is used after mutation."""
71
+
72
+ def f(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
73
+ x = x.relu() # not a graph param
74
+ x_clone = x.clone()
75
+ x_clone.add_(1)
76
+ return x, x_clone
77
+
78
+ inp = torch.randn(2, 3)
79
+ graph_module = make_fx(f)(inp)
80
+ assert count_clones(graph_module.graph) == 1
81
+
82
+ expected = graph_module(inp)
83
+ clone_cleanup_pass(graph_module.graph)
84
+ graph_module.recompile()
85
+ actual = graph_module(inp)
86
+
87
+ # Clone should be KEPT because original is used after mutation
88
+ assert count_clones(graph_module.graph) == 1
89
+ torch.testing.assert_close(actual[0], expected[0])
90
+ torch.testing.assert_close(actual[1], expected[1])
91
+
92
+ def test_remove_clone_with_mutation_no_original_use(self, clone_cleanup_pass):
93
+ """Clone can be removed if it's mutated but original is not used after."""
94
+
95
+ def f(x: torch.Tensor) -> torch.Tensor:
96
+ x = x.relu() # not a graph param
97
+ x_clone = x.clone()
98
+ x_clone.add_(1)
99
+ return x_clone
100
+
101
+ inp = torch.randn(2, 3)
102
+ graph_module = make_fx(f)(inp)
103
+ assert count_clones(graph_module.graph) == 1
104
+
105
+ expected = graph_module(inp)
106
+ clone_cleanup_pass(graph_module.graph)
107
+ graph_module.recompile()
108
+ actual = graph_module(inp)
109
+
110
+ assert count_clones(graph_module.graph) == 0
111
+ torch.testing.assert_close(actual, expected)
112
+
113
+ def test_clone_chain(self, clone_cleanup_pass):
114
+ """Test handling of clone chains: x -> clone1 -> clone2."""
115
+
116
+ def f(x: torch.Tensor) -> torch.Tensor:
117
+ x = x.relu() # not a graph param
118
+ x1 = x.clone()
119
+ x2 = x1.clone()
120
+ return x2 + 1
121
+
122
+ inp = torch.randn(2, 3)
123
+ graph_module = make_fx(f)(inp)
124
+ assert count_clones(graph_module.graph) == 2
125
+
126
+ expected = graph_module(inp)
127
+ clone_cleanup_pass(graph_module.graph)
128
+ graph_module.recompile()
129
+ actual = graph_module(inp)
130
+
131
+ # Both clones should be removed
132
+ assert count_clones(graph_module.graph) == 0
133
+ torch.testing.assert_close(actual, expected)
134
+
135
+ def test_multiple_clones_of_same_input(self, clone_cleanup_pass):
136
+ """Test multiple independent clones of the same input."""
137
+
138
+ def f(x: torch.Tensor) -> torch.Tensor:
139
+ x1 = x.clone()
140
+ x2 = x.clone()
141
+ return x1 + x2
142
+
143
+ inp = torch.randn(2, 3)
144
+ graph_module = make_fx(f)(inp)
145
+ assert count_clones(graph_module.graph) == 2
146
+
147
+ expected = graph_module(inp)
148
+ clone_cleanup_pass(graph_module.graph)
149
+ graph_module.recompile()
150
+ actual = graph_module(inp)
151
+
152
+ # Both clones should be removed (only readonly uses)
153
+ assert count_clones(graph_module.graph) == 0
154
+ torch.testing.assert_close(actual, expected)
155
+
156
+ def test_no_clones_in_graph(self, clone_cleanup_pass):
157
+ """Test pass behavior when graph has no clones."""
158
+
159
+ def f(x: torch.Tensor) -> torch.Tensor:
160
+ return x + 1
161
+
162
+ inp = torch.randn(2, 3)
163
+ graph_module = make_fx(f)(inp)
164
+ assert count_clones(graph_module.graph) == 0
165
+
166
+ expected = graph_module(inp)
167
+ clone_cleanup_pass(graph_module.graph)
168
+ graph_module.recompile()
169
+ actual = graph_module(inp)
170
+
171
+ assert count_clones(graph_module.graph) == 0
172
+ torch.testing.assert_close(actual, expected)
173
+
174
+ def test_multiple_passes(self, clone_cleanup_pass):
175
+ """Test running the pass multiple times (should be idempotent)."""
176
+
177
+ def f(x: torch.Tensor) -> torch.Tensor:
178
+ x1 = x.clone()
179
+ return x1 + 1
180
+
181
+ inp = torch.randn(2, 3)
182
+ graph_module = make_fx(f)(inp)
183
+ assert count_clones(graph_module.graph) == 1
184
+
185
+ expected = graph_module(inp)
186
+
187
+ clone_cleanup_pass(graph_module.graph)
188
+ assert count_clones(graph_module.graph) == 0
189
+ graph_module.recompile()
190
+ actual = graph_module(inp)
191
+ torch.testing.assert_close(actual, expected)
192
+
193
+ clone_cleanup_pass(graph_module.graph)
194
+ assert count_clones(graph_module.graph) == 0
195
+ graph_module.recompile()
196
+ actual = graph_module(inp)
197
+ torch.testing.assert_close(actual, expected)
198
+
199
+ def test_output_node_no_write(self):
200
+ """Output nodes never write to their inputs."""
201
+
202
+ def f(x: torch.Tensor) -> torch.Tensor:
203
+ return x
204
+
205
+ graph_module = make_fx(f)(torch.randn(2, 3))
206
+ x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0]
207
+ output_node = [n for n in graph_module.graph.nodes if n.op == "output"][0]
208
+
209
+ assert not user_writes_to_node(output_node, x_node)
210
+
211
+ def test_readonly_op_no_write(self):
212
+ """Readonly operations don't write to inputs."""
213
+
214
+ def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
215
+ return x + y
216
+
217
+ graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3))
218
+ placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
219
+ add_node = [
220
+ n
221
+ for n in graph_module.graph.nodes
222
+ if n.op == "call_function" and n.target == torch.ops.aten.add.Tensor
223
+ ][0]
224
+
225
+ assert not user_writes_to_node(add_node, placeholders[0])
226
+ assert not user_writes_to_node(add_node, placeholders[1])
227
+
228
+ def test_inplace_op_writes(self):
229
+ """Inplace operations write to first argument."""
230
+
231
+ def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
232
+ x.add_(y)
233
+ return x
234
+
235
+ graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3))
236
+ placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
237
+ add_node = [
238
+ n
239
+ for n in graph_module.graph.nodes
240
+ if n.op == "call_function" and "add_" in str(n.target)
241
+ ][0]
242
+
243
+ # add_ writes to first arg but not second
244
+ assert user_writes_to_node(add_node, placeholders[0])
245
+ assert not user_writes_to_node(add_node, placeholders[1])
246
+
247
+ def test_copy_writes(self):
248
+ """copy_ operation writes to first argument."""
249
+
250
+ def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
251
+ x.copy_(y)
252
+ return x
253
+
254
+ graph_module = make_fx(f)(torch.randn(2, 3), torch.randn(2, 3))
255
+ placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
256
+ copy_node = [
257
+ n
258
+ for n in graph_module.graph.nodes
259
+ if n.op == "call_function" and "copy_" in str(n.target)
260
+ ][0]
261
+
262
+ assert user_writes_to_node(copy_node, placeholders[0])
263
+ assert not user_writes_to_node(copy_node, placeholders[1])
264
+
265
+ def test_auto_functionalized_not_a_write(self):
266
+ """auto_functionalized ops are follow-up uses, not writes."""
267
+ from torch._higher_order_ops.auto_functionalize import auto_functionalized
268
+
269
+ def f(x: torch.Tensor) -> torch.Tensor:
270
+ return x
271
+
272
+ graph_module = make_fx(f)(torch.randn(2, 3))
273
+ x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0]
274
+
275
+ # Create an auto_functionalized node in the graph
276
+ with graph_module.graph.inserting_before(None):
277
+ af_node = graph_module.graph.call_function(
278
+ auto_functionalized, kwargs={"input": x_node}
279
+ )
280
+
281
+ # auto_functionalized should not be treated as a write
282
+ assert not user_writes_to_node(af_node, x_node)
283
+
284
+ def test_higher_order_op_conservatively_writes(self):
285
+ """Other higher-order operators are conservatively treated as writes."""
286
+ from torch._ops import HigherOrderOperator
287
+
288
+ def f(x: torch.Tensor) -> torch.Tensor:
289
+ return x
290
+
291
+ graph_module = make_fx(f)(torch.randn(2, 3))
292
+ x_node = [n for n in graph_module.graph.nodes if n.op == "placeholder"][0]
293
+
294
+ # Create a concrete higher-order operator subclass
295
+ class MockHigherOrderOp(HigherOrderOperator):
296
+ def __call__(self, *args, **kwargs):
297
+ return args[0] if args else None
298
+
299
+ mock_hoo = MockHigherOrderOp("mock_higher_order_op")
300
+
301
+ with graph_module.graph.inserting_before(None):
302
+ hoo_node = graph_module.graph.call_function(mock_hoo, args=(x_node,))
303
+
304
+ # Should be conservative and assume it could write
305
+ assert user_writes_to_node(hoo_node, x_node)
306
+
307
+
308
+class TestCloneCleanupWithDonatedInputs:
309
+ """Test UnsafeCloneEliminationPass with donated input tracking via PassContext."""
310
+
311
+ @pytest.fixture(autouse=True)
312
+ def setup_pass_context(self):
313
+ """Set up pass context for each test."""
314
+ with pass_context(compile_range=Range(1, 8192)):
315
+ yield
316
+
317
+ def test_donated_input_clone_removed(self, clone_cleanup_pass):
318
+ """Clone of donated input should be removed."""
319
+
320
+ def f(x: torch.Tensor) -> torch.Tensor:
321
+ x_clone = x.clone()
322
+ x_clone.add_(1)
323
+ return x_clone
324
+
325
+ inp = torch.randn(2, 3)
326
+ graph_module = make_fx(f)(inp)
327
+ assert count_clones(graph_module.graph) == 1
328
+
329
+ # Mark first parameter as donated
330
+ get_pass_context().donated_input_ids = {0}
331
+
332
+ expected = graph_module(inp.clone())
333
+ clone_cleanup_pass(graph_module.graph)
334
+ graph_module.recompile()
335
+
336
+ # Clone should be removed since input is donated
337
+ assert count_clones(graph_module.graph) == 0
338
+
339
+ # Input can be mutated (donated)
340
+ inp_copy = inp.clone()
341
+ actual = graph_module(inp_copy)
342
+ torch.testing.assert_close(actual, expected)
343
+
344
+ def test_non_donated_input_clone_kept(self, clone_cleanup_pass):
345
+ """Clone of non-donated input with mutation should be kept."""
346
+
347
+ def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
348
+ x_clone = x.clone()
349
+ x_clone.add_(1)
350
+ return x, x_clone
351
+
352
+ inp_x = torch.randn(2, 3)
353
+ inp_y = torch.randn(2, 3)
354
+ graph_module = make_fx(f)(inp_x, inp_y)
355
+ assert count_clones(graph_module.graph) == 1
356
+
357
+ # No donated inputs
358
+ get_pass_context().donated_input_ids = set()
359
+
360
+ expected = graph_module(inp_x.clone(), inp_y.clone())
361
+ clone_cleanup_pass(graph_module.graph)
362
+ graph_module.recompile()
363
+
364
+ # Clone should be kept since input is not donated and original is used
365
+ assert count_clones(graph_module.graph) == 1
366
+
367
+ # Verify inputs are not mutated
368
+ inp_x_before = inp_x.clone()
369
+ inp_y_before = inp_y.clone()
370
+ actual = graph_module(inp_x, inp_y)
371
+ torch.testing.assert_close(
372
+ inp_x, inp_x_before, msg="Input x should not be mutated"
373
+ )
374
+ torch.testing.assert_close(
375
+ inp_y, inp_y_before, msg="Input y should not be mutated"
376
+ )
377
+ torch.testing.assert_close(actual[0], expected[0])
378
+ torch.testing.assert_close(actual[1], expected[1])
379
+
380
+ def test_mixed_donated_inputs(self, clone_cleanup_pass):
381
+ """Test with some inputs donated and some not."""
382
+
383
+ def f(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
384
+ x_clone = x.clone()
385
+ x_clone.add_(1)
386
+ y_clone = y.clone()
387
+ y_clone.add_(2)
388
+ return x_clone, y_clone
389
+
390
+ inp_x = torch.randn(2, 3)
391
+ inp_y = torch.randn(2, 3)
392
+ graph_module = make_fx(f)(inp_x, inp_y)
393
+ assert count_clones(graph_module.graph) == 2
394
+
395
+ # Only x is donated
396
+ get_pass_context().donated_input_ids = {0}
397
+
398
+ expected = graph_module(inp_x.clone(), inp_y.clone())
399
+ clone_cleanup_pass(graph_module.graph)
400
+ graph_module.recompile()
401
+
402
+ # x_clone removed (x is donated), y_clone kept (y is not donated)
403
+ assert count_clones(graph_module.graph) == 1
404
+
405
+ # Verify y is not mutated (x can be mutated since it's donated)
406
+ inp_y_before = inp_y.clone()
407
+ actual = graph_module(inp_x.clone(), inp_y)
408
+ torch.testing.assert_close(
409
+ inp_y, inp_y_before, msg="Input y should not be mutated"
410
+ )
411
+ torch.testing.assert_close(actual[0], expected[0])
412
+ torch.testing.assert_close(actual[1], expected[1])
tests/compile/passes/ir/test_inplace_functionalization.py
new
+465
@@ -0,0 +1,465 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for IR inplace functionalization pass integration.
5
+
6
+This test suite verifies that the inplace functionalization pass, lowering pass,
7
+and clone cleanup pass work together correctly with donated buffer tracking.
8
+"""
9
+
10
+from collections.abc import Callable
11
+
12
+import pytest
13
+import torch
14
+import torch._dynamo.exc
15
+from torch import nn
16
+
17
+import vllm.kernels # noqa: F401 to register kernels
18
+from vllm.compilation.passes.inductor_pass import InductorPass, get_pass_context
19
+from vllm.compilation.passes.ir.clone_elimination import (
20
+ UnsafeCloneEliminationPass,
21
+)
22
+from vllm.compilation.passes.ir.inplace_functionalization import (
23
+ VllmIRInplaceFunctionalizationPass,
24
+)
25
+from vllm.compilation.passes.ir.lowering_pass import VllmIRLoweringPass
26
+from vllm.config import VllmConfig
27
+from vllm.ir import ops
28
+from vllm.platforms import current_platform
29
+from vllm.triton_utils import HAS_TRITON, tl, triton
30
+
31
+from ...backend import TestBackend
32
+
33
+
34
+class StoreDonationInfoPass(InductorPass):
35
+ def __init__(self):
36
+ self.donated_input_ids_sets: list[set[int]] = []
37
+
38
+ def __call__(self, *args, **kwargs):
39
+ ctx = get_pass_context()
40
+ self.donated_input_ids_sets += [ctx.donated_input_ids]
41
+
42
+
43
+class MaybeInplaceModel(nn.Module):
44
+ """Model using only maybe_inplace variants."""
45
+
46
+ def __init__(self, hidden_size=16):
47
+ super().__init__()
48
+ self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
49
+ self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
50
+
51
+ def forward(
52
+ self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor
53
+ ):
54
+ # First maybe_inplace - x & residual1 are donated
55
+ x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace(
56
+ x, residual1, self.weight1, 1e-5
57
+ )
58
+ # Second maybe_inplace - residual2 is donated
59
+ x_normed2, residual_out2 = ops.fused_add_rms_norm.maybe_inplace(
60
+ x_normed1, residual2, self.weight2, 1e-5
61
+ )
62
+ return x_normed2, residual_out1, residual_out2
63
+
64
+
65
+class FunctionalModel(nn.Module):
66
+ """Model using only functional (default) variants."""
67
+
68
+ def __init__(self, hidden_size=16):
69
+ super().__init__()
70
+ self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
71
+ self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
72
+
73
+ def forward(
74
+ self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor
75
+ ):
76
+ # First functional - no donation
77
+ x_normed1, residual_out1 = ops.fused_add_rms_norm(
78
+ x, residual1, self.weight1, 1e-5
79
+ )
80
+ # Second functional - no donation
81
+ x_normed2, residual_out2 = ops.fused_add_rms_norm(
82
+ x_normed1, residual2, self.weight2, 1e-5
83
+ )
84
+ return x_normed2, residual_out1, residual_out2
85
+
86
+
87
+class MixedModel(nn.Module):
88
+ """Model mixing maybe_inplace and functional variants."""
89
+
90
+ def __init__(self, hidden_size=16):
91
+ super().__init__()
92
+ self.weight1 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
93
+ self.weight2 = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
94
+
95
+ def forward(
96
+ self, x: torch.Tensor, residual1: torch.Tensor, residual2: torch.Tensor
97
+ ):
98
+ # First maybe_inplace - x & residual1 are donated
99
+ x_normed1, residual_out1 = ops.fused_add_rms_norm.maybe_inplace(
100
+ x, residual1, self.weight1, 1e-5
101
+ )
102
+ # Second functional - no donation, x_normed1 must be preserved as it's returned
103
+ x_normed2, residual_out2 = ops.fused_add_rms_norm(
104
+ x_normed1, residual2, self.weight2, 1e-5
105
+ )
106
+ # Return both to prevent x_normed1 from being optimized away
107
+ return x_normed1, x_normed2, residual_out1, residual_out2
108
+
109
+
110
+class ModelWithTritonAfterMaybeInplace(nn.Module):
111
+ """
112
+ Model using maybe_inplace followed by a Triton kernel.
113
+ Test clone elimination can handle Triton in the graph
114
+ """
115
+
116
+ def __init__(self, hidden_size=16):
117
+ super().__init__()
118
+ self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
119
+
120
+ @triton.jit
121
+ def _triton_add_kernel(
122
+ x_ptr,
123
+ y_ptr,
124
+ n_elements,
125
+ BLOCK_SIZE: tl.constexpr,
126
+ ):
127
+ pid = tl.program_id(axis=0)
128
+ block_start = pid * BLOCK_SIZE
129
+ offsets = block_start + tl.arange(0, BLOCK_SIZE)
130
+ mask = offsets < n_elements
131
+ x = tl.load(x_ptr + offsets, mask=mask)
132
+ y = x + 0.1
133
+ tl.store(y_ptr + offsets, y, mask=mask)
134
+
135
+ def triton_add(x: torch.Tensor) -> torch.Tensor:
136
+ """Simple Triton add kernel."""
137
+ y = torch.empty_like(x)
138
+ n_elements = x.numel()
139
+ grid = (triton.cdiv(n_elements, 256),)
140
+ _triton_add_kernel[grid](x, y, n_elements, BLOCK_SIZE=256)
141
+ return y
142
+
143
+ self.triton_add = triton_add
144
+
145
+ def forward(self, x: torch.Tensor, residual: torch.Tensor, residual2: torch.Tensor):
146
+ x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace(
147
+ x, residual, self.weight, 1e-5
148
+ )
149
+
150
+ x_processed = self.triton_add(x_normed)
151
+
152
+ # x_processed does not need to be cloned, residual2 does
153
+ x_normed2, residual_out2 = ops.fused_add_rms_norm(
154
+ x_processed, residual2, self.weight, 1e-5
155
+ )
156
+ return x_normed2, residual_out2
157
+
158
+
159
+skipif_no_triton = pytest.mark.skipif(not HAS_TRITON, reason="Requires Triton")
160
+
161
+
162
+@pytest.mark.skipif(
163
+ not current_platform.is_cuda_alike(),
164
+ reason="Only test on cuda and rocm platform",
165
+)
166
+@pytest.mark.parametrize(
167
+ "model_class,expected_functionalized,expected_donated,expected_clones",
168
+ [
169
+ # 2 inplace calls, all activations donated, all clones eliminated
170
+ (MaybeInplaceModel, 2, 3, 0),
171
+ # No inplace calls, no donations, 3 clones (one eliminated)
172
+ (FunctionalModel, 0, 0, 3),
173
+ # One inplace call, two donated activations, 2 clones
174
+ (MixedModel, 1, 2, 2),
175
+ # One inplace call, two donated, 1 clone remaining
176
+ pytest.param(ModelWithTritonAfterMaybeInplace, 1, 2, 1, marks=skipif_no_triton),
177
+ ],
178
+)
179
+def test_inplace_functionalization(
180
+ default_vllm_config: VllmConfig,
181
+ model_class,
182
+ expected_functionalized: int,
183
+ expected_clones: int,
184
+ expected_donated: int,
185
+):
186
+ """Test inplace functionalization, lowering, and clone cleanup."""
187
+ torch.set_default_device(current_platform.device_type)
188
+
189
+ # Use vllm_c so inplace path is triggered
190
+ default_vllm_config.kernel_config.ir_op_priority.fused_add_rms_norm = [
191
+ "vllm_c",
192
+ "native",
193
+ ]
194
+
195
+ # Create passes in order they run during compilation
196
+ functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config)
197
+ lowering_pass = VllmIRLoweringPass(default_vllm_config)
198
+ donated_info_pass = StoreDonationInfoPass()
199
+ cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config)
200
+
201
+ # Set up backend with pre-grad pass
202
+ backend = TestBackend(lowering_pass, donated_info_pass, cleanup_pass)
203
+ backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass
204
+
205
+ model = model_class()
206
+ x = torch.randn(8, 16, dtype=torch.bfloat16)
207
+ residual1 = torch.randn(8, 16, dtype=torch.bfloat16)
208
+ residual2 = torch.randn(8, 16, dtype=torch.bfloat16)
209
+
210
+ with default_vllm_config.kernel_config.ir_op_priority.set_priority():
211
+ # Reference output without optimization
212
+ ref_output = model(x.clone(), residual1.clone(), residual2.clone())
213
+
214
+ # Compile with inplace optimization
215
+ compiled_model = torch.compile(model, backend=backend, fullgraph=True)
216
+ output = compiled_model(x.clone(), residual1.clone(), residual2.clone())
217
+
218
+ # Verify correctness (relaxed tolerance for bfloat16)
219
+ for i in range(len(ref_output)):
220
+ torch.testing.assert_close(output[i], ref_output[i], rtol=1e-2, atol=1e-2)
221
+
222
+ # Verify expected number of ops were functionalized
223
+ func_ops = functionalization_pass.functionalized_ops
224
+ assert len(func_ops) == int(bool(expected_functionalized))
225
+ if expected_functionalized > 0:
226
+ assert "fused_add_rms_norm" in func_ops
227
+ assert func_ops["fused_add_rms_norm"] == expected_functionalized
228
+
229
+ # Verify lowering happened (2 ops in all cases)
230
+ assert "fused_add_rms_norm" in lowering_pass.selected_impls
231
+ assert len(lowering_pass.selected_impls["fused_add_rms_norm"]) == 2
232
+ assert all(
233
+ provider == "vllm_c"
234
+ for node, provider in lowering_pass.selected_impls["fused_add_rms_norm"].items()
235
+ ), lowering_pass.selected_impls
236
+
237
+ # Verify correct number of donated IDs
238
+ assert len(donated_info_pass.donated_input_ids_sets) == 1
239
+ assert len(donated_info_pass.donated_input_ids_sets[0]) == expected_donated
240
+
241
+ # Verify expected number of clones after cleanup
242
+ actual_clones = backend.op_count(torch.ops.aten.clone.default, before=False)
243
+ assert actual_clones == expected_clones, (
244
+ f"Expected {expected_clones} clones, got {actual_clones}:"
245
+ f"{backend.print_graphs()}"
246
+ )
247
+
248
+
249
+@pytest.mark.skipif(
250
+ not current_platform.is_cuda_alike(),
251
+ reason="Only test on cuda and rocm platform",
252
+)
253
+def test_donated_buffer_context_propagation(default_vllm_config):
254
+ """Test that donated_input_ids propagates correctly through pass_context."""
255
+ torch.set_default_device(current_platform.device_type)
256
+
257
+ # Create a custom backend that inspects pass_context in cleanup pass
258
+ functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config)
259
+ lowering_pass = VllmIRLoweringPass(default_vllm_config)
260
+
261
+ donation_info_pass = StoreDonationInfoPass()
262
+ cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config)
263
+
264
+ backend = TestBackend(lowering_pass, donation_info_pass, cleanup_pass)
265
+ backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass
266
+
267
+ model = MaybeInplaceModel()
268
+ x = torch.randn(8, 16, dtype=torch.bfloat16)
269
+ residual1 = torch.randn(8, 16, dtype=torch.bfloat16)
270
+ residual2 = torch.randn(8, 16, dtype=torch.bfloat16)
271
+
272
+ compiled_model = torch.compile(model, backend=backend, fullgraph=True)
273
+ compiled_model(x.clone(), residual1.clone(), residual2.clone())
274
+
275
+ donated_ids_seen = donation_info_pass.donated_input_ids_sets
276
+ # Verify donated_input_ids was set and propagated
277
+ assert len(donated_ids_seen) == 1
278
+ # Should have donated inputs (exact indices depend on AOTAutograd)
279
+ assert len(donated_ids_seen[0]) == 3
280
+ # All donated ids should be valid non-negative integers
281
+ for idx in donated_ids_seen[0]:
282
+ assert isinstance(idx, int) and idx >= 0, f"Invalid donated index: {idx}"
283
+
284
+
285
+@pytest.mark.skipif(
286
+ not current_platform.is_cuda_alike(),
287
+ reason="Only test on cuda and rocm platform",
288
+)
289
+def test_maybe_inplace_reuse_error(default_vllm_config):
290
+ """Test that reusing a donated activation input raises ValueError."""
291
+ torch.set_default_device(current_platform.device_type)
292
+
293
+ class ReuseModel(nn.Module):
294
+ """Model that incorrectly reuses a donated activation input."""
295
+
296
+ def __init__(self, hidden_size=16):
297
+ super().__init__()
298
+ self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
299
+
300
+ def forward(self, x: torch.Tensor, residual: torch.Tensor):
301
+ # x is donated to maybe_inplace
302
+ x_normed, residual_out = ops.fused_add_rms_norm.maybe_inplace(
303
+ x, residual, self.weight, 1e-5
304
+ )
305
+ # ERROR: x is used again after being donated
306
+ return x_normed + x # This should raise ValueError
307
+
308
+ functionalization_pass = VllmIRInplaceFunctionalizationPass(default_vllm_config)
309
+ lowering_pass = VllmIRLoweringPass(default_vllm_config)
310
+ cleanup_pass = UnsafeCloneEliminationPass(default_vllm_config)
311
+
312
+ backend = TestBackend(lowering_pass, cleanup_pass)
313
+ backend.inductor_config["pre_grad_custom_pass"] = functionalization_pass
314
+
315
+ model = ReuseModel()
316
+ x = torch.randn(8, 16, dtype=torch.bfloat16)
317
+ residual = torch.randn(8, 16, dtype=torch.bfloat16)
318
+
319
+ # Compilation should raise BackendCompilerFailed wrapping ValueError
320
+ with pytest.raises(
321
+ torch._dynamo.exc.BackendCompilerFailed,
322
+ match="is used again after the node",
323
+ ):
324
+ compiled_model = torch.compile(model, backend=backend, fullgraph=True)
325
+ compiled_model(x.clone(), residual.clone())
326
+
327
+
328
+# Piecewise compilation tests with graph splitting
329
+
330
+
331
+@torch.library.custom_op("vllm::test_split_marker", mutates_args=())
332
+def test_split_marker(x: torch.Tensor) -> torch.Tensor:
333
+ """Identity op that marks a split point for piecewise compilation."""
334
+ return x.clone()
335
+
336
+
337
+@test_split_marker.register_fake
338
+def _fake_split_marker(x: torch.Tensor) -> torch.Tensor:
339
+ return torch.empty_like(x)
340
+
341
+
342
+class TransformerBlockWithSplits(nn.Module):
343
+ """Transformer block with explicit split points for piecewise compilation."""
344
+
345
+ def __init__(self, hidden_size=32, intermediate_size=128):
346
+ super().__init__()
347
+ self.hidden_size = hidden_size
348
+ self.intermediate_size = intermediate_size
349
+
350
+ # Attention-like projection
351
+ self.attn_proj = nn.Linear(
352
+ hidden_size, hidden_size, bias=False, dtype=torch.bfloat16
353
+ )
354
+
355
+ # Post-attention norm
356
+ self.post_attn_norm = nn.Parameter(
357
+ torch.ones(hidden_size, dtype=torch.bfloat16)
358
+ )
359
+
360
+ # MLP
361
+ self.gate_proj = nn.Linear(
362
+ hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16
363
+ )
364
+ self.up_proj = nn.Linear(
365
+ hidden_size, intermediate_size, bias=False, dtype=torch.bfloat16
366
+ )
367
+ self.down_proj = nn.Linear(
368
+ intermediate_size, hidden_size, bias=False, dtype=torch.bfloat16
369
+ )
370
+
371
+ # Post-MLP norm
372
+ self.post_mlp_norm = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
373
+
374
+ def forward(self, x: torch.Tensor):
375
+ # Attention block with residual
376
+ residual1 = x
377
+ attn_out = self.attn_proj(x)
378
+
379
+ # Fused add + norm (maybe_inplace: residual1 is donated)
380
+ normed1, residual1 = ops.fused_add_rms_norm.maybe_inplace(
381
+ attn_out, residual1, self.post_attn_norm, 1e-5
382
+ )
383
+
384
+ # Force a graph split here
385
+ normed1 = torch.ops.vllm.test_split_marker(normed1)
386
+
387
+ # MLP block
388
+ gate = self.gate_proj(normed1)
389
+ up = self.up_proj(normed1)
390
+ mlp_out = self.down_proj(gate * torch.nn.functional.silu(up))
391
+
392
+ # Fused add + norm (maybe_inplace: residual1 is donated)
393
+ normed2, residual2 = ops.fused_add_rms_norm.maybe_inplace(
394
+ mlp_out, residual1, self.post_mlp_norm, 1e-5
395
+ )
396
+
397
+ return normed2, residual2
398
+
399
+
400
+def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int):
401
+ def inner(*args):
402
+ torch._dynamo.mark_dynamic(args[arg_index], dim_index)
403
+ return fn(*args)
404
+
405
+ return inner
406
+
407
+
408
+@pytest.mark.skipif(
409
+ not current_platform.is_cuda_alike(),
410
+ reason="Only test on cuda and rocm platform",
411
+)
412
+def test_piecewise_compilation_with_donated_buffers(monkeypatch, fresh_vllm_cache):
413
+ """
414
+ Test piecewise compilation with donated buffers across graph splits.
415
+ Utilizes a custom splitting op. Uses fresh cache to avoid compilation caching.
416
+ """
417
+ torch.set_default_device(current_platform.device_type)
418
+
419
+ # Disable compilation cache to avoid serialization issues
420
+ monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
421
+
422
+ from vllm.compilation.backends import VllmBackend
423
+ from vllm.config import CompilationConfig, VllmConfig
424
+
425
+ # Create config with custom splitting op
426
+ store_donation_info = StoreDonationInfoPass()
427
+ vllm_config = VllmConfig(
428
+ compilation_config=CompilationConfig(
429
+ custom_ops=["all"],
430
+ splitting_ops=["vllm::test_split_marker"],
431
+ inductor_compile_config={"post_grad_custom_post_pass": store_donation_info},
432
+ )
433
+ )
434
+
435
+ backend = VllmBackend(vllm_config)
436
+
437
+ model = TransformerBlockWithSplits()
438
+ x = torch.randn(8, 32, dtype=torch.bfloat16)
439
+
440
+ # Reference output
441
+ ref_output = with_dyn_arg(model, 0, 0)(x.clone())
442
+
443
+ # Compile with piecewise compilation (graph will split at split_marker)
444
+ compiled_model = torch.compile(model, backend=backend, fullgraph=False)
445
+ output = with_dyn_arg(compiled_model, 0, 0)(x.clone())
446
+
447
+ # Verify correctness (relaxed tolerance for bfloat16)
448
+ torch.testing.assert_close(output[0], ref_output[0], rtol=1e-2, atol=1e-2)
449
+ torch.testing.assert_close(output[1], ref_output[1], rtol=1e-2, atol=1e-2)
450
+
451
+ # Verify the model was split into multiple submodules
452
+ assert hasattr(backend, "split_gm"), "Backend should have split graph module"
453
+
454
+ # Should have at least 2 submodules (split by test_split_marker op)
455
+ submodules = list(backend.split_gm.named_children())
456
+ num_submodules = len(submodules)
457
+ assert num_submodules >= 2, (
458
+ f"Expected at least 2 submodules (split), got {num_submodules}"
459
+ )
460
+
461
+ # Check that donation info was propagated correctly
462
+ donated_inputs_sets = store_donation_info.donated_input_ids_sets
463
+ assert len(donated_inputs_sets) == 2
464
+ assert len(donated_inputs_sets[0]) == 1
465
+ assert len(donated_inputs_sets[1]) == 1
tests/compile/passes/test_functionalization.py
+1
-1
@@ -126,7 +126,7 @@ class TestFusedAddRMSNorm(torch.nn.Module):
126
if TEST_FP8 and do_fusion:
127
return [torch.ops._C.fused_add_rms_norm_static_fp8_quant.default]
128
else:
129
- return [torch.ops._C.fused_add_rms_norm.default]
129
+ return []
130
131
def ops_not_in_model(self):
132
return []
tests/compile/passes/test_fuse_act_padding.py
+1
-1
@@ -59,7 +59,7 @@ class TestModel(torch.nn.Module):
59
60
def ops_in_model_before(self):
61
return [
62
- rocm_aiter_ops.get_rmsnorm_fused_add_op(),
62
+ torch.ops.vllm_ir.fused_add_rms_norm,
63
torch.ops.aten.constant_pad_nd,
64
]
65
tests/compile/passes/test_fusion.py
+4
-15
@@ -17,7 +17,6 @@ from vllm.compilation.passes.fusion.rms_quant_fusion import (
17
FusedRMSQuantKey,
18
RMSNormQuantFusionPass,
19
)
20
-from vllm.compilation.passes.fx_utils import find_op_nodes
20
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
21
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
22
from vllm.config import (
@@ -243,9 +242,10 @@ class TestModel(torch.nn.Module):
242
]
243
244
def ops_in_model_before_partial(self):
246
- return [torch.ops.vllm_ir.rms_norm] + (
247
- [RMS_ADD_OP] if self.enable_rms_norm_custom_op else [torch.ops.aten.rsqrt]
248
- )
245
+ return [
246
+ torch.ops.vllm_ir.rms_norm,
247
+ torch.ops.vllm_ir.fused_add_rms_norm.default,
248
+ ]
249
250
251
def _run_fusion_test(
@@ -383,17 +383,6 @@ def test_fusion_rmsnorm_quant(
383
model.ops_in_model_before_partial(), fully_replaced=False
384
)
385
386
- # If RMSNorm custom op is disabled (native/torch impl used),
387
- # there's a risk that the fused add doesn't get included in the
388
- # replacement and only the rms part gets fused with quant.
389
- # Hence, we check only 2 add nodes are left (final fused rmsnorm add).
390
- if not enable_rms_norm_custom_op:
391
- n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))
392
- # rms_norm is IR, not included
393
- # 6 = 3x2 (3xRMS_ADD, 2 each)
394
- assert n_add_nodes(backend.graph_pre_pass) == 6
395
- assert n_add_nodes(backend.graph_post_pass) == 2
396
-
386
387
@pytest.mark.parametrize("dtype", [torch.bfloat16])
388
@pytest.mark.parametrize("hidden_size", [256])
tests/ir/test_inplace_op.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+from torch import Tensor
5
+from torch.fx.experimental.proxy_tensor import make_fx
6
+
7
+import vllm.ir.op
8
+from vllm.ir.op import IrOp, IrOpInplaceOverload
9
+
10
+
11
+@vllm.ir.register_op(allow_inplace=True)
12
+def _custom_mm2(x: Tensor, w: Tensor) -> Tensor:
13
+ return x @ w
14
+
15
+
16
+@_custom_mm2.register_impl("regular")
17
+def _custom_mm2_regular(x: Tensor, w: Tensor) -> Tensor:
18
+ return x @ w + 1
19
+
20
+
21
+@_custom_mm2.register_impl("inplace", inplace=True)
22
+def _custom_mm2_inplace(x: Tensor, w: Tensor) -> Tensor:
23
+ x.copy_(x @ w + 2)
24
+ return x
25
+
26
+
27
+class TestInplaceOp:
28
+ def test_registration(self):
29
+ # Test that the inplace op is registered correctly.
30
+ assert "_custom_mm2" in IrOp.registry
31
+ assert IrOp.registry["_custom_mm2"] is _custom_mm2
32
+ assert _custom_mm2.torch_op is torch.ops.vllm_ir._custom_mm2.default
33
+ assert isinstance(_custom_mm2.maybe_inplace, IrOpInplaceOverload)
34
+ assert (
35
+ _custom_mm2.maybe_inplace.torch_op
36
+ is torch.ops.vllm_ir._custom_mm2.maybe_inplace
37
+ )
38
+
39
+ def test_inplace_dispatching(self):
40
+ # check that the correct implementation is dispatched based on priority,
41
+ # and inplace semantics hold
42
+ w = torch.randn(3, 3)
43
+ x = torch.randn(2, 3)
44
+ x1 = x.clone()
45
+
46
+ with _custom_mm2.set_priority(["regular"]):
47
+ result_regular = _custom_mm2.maybe_inplace(x, w)
48
+
49
+ # check that the regular op does not modify x
50
+ torch.testing.assert_close(x, x1, atol=0, rtol=0)
51
+
52
+ with _custom_mm2.set_priority(["inplace"]):
53
+ result_inplace: Tensor = _custom_mm2.maybe_inplace(x, w)
54
+
55
+ # check that the inplace op returns x directly
56
+ assert result_inplace.data_ptr() == x.data_ptr()
57
+
58
+ torch.testing.assert_close(result_inplace, x1 @ w + 2)
59
+ torch.testing.assert_close(result_regular, x1 @ w + 1)
60
+
61
+ def test_default_dispatching(self):
62
+ # check that the correct implementation is dispatched,
63
+ # and ops do not modify inputs when using the default overload
64
+ w = torch.randn(3, 3)
65
+ x = torch.randn(2, 3)
66
+ x1 = x.clone()
67
+
68
+ with _custom_mm2.set_priority(["regular"]):
69
+ result_regular = _custom_mm2(x, w)
70
+
71
+ with _custom_mm2.set_priority(["inplace"]):
72
+ result_inplace = _custom_mm2(x, w)
73
+
74
+ # check that x was not modified by either impl
75
+ torch.testing.assert_close(x, x1, atol=0, rtol=0)
76
+
77
+ torch.testing.assert_close(result_inplace, x1 @ w + 2)
78
+ torch.testing.assert_close(result_regular, x1 @ w + 1)
79
+
80
+ def test_trace(self):
81
+ # Test that the inplace op can be used in a graph.
82
+ def func(x: Tensor, y: Tensor) -> Tensor:
83
+ return _custom_mm2.maybe_inplace(x, y)
84
+
85
+ x = torch.randn(2, 3)
86
+ y = torch.randn(3, 4)
87
+ graph = make_fx(func)(x, y)
88
+ assert any(
89
+ node.target == torch.ops.vllm_ir._custom_mm2.maybe_inplace
90
+ for node in graph.graph.nodes
91
+ )
tests/ir/test_op.py
+27
-16
@@ -21,7 +21,7 @@ class CustomError(Exception):
21
pass
22
23
24
-@vllm.ir.register_op
24
+@vllm.ir.register_op(allow_inplace=True)
25
def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
26
return x + y
27
@@ -129,11 +129,15 @@ class TestIrOpCustomAdd:
129
130
@pytest.mark.parametrize("enable_torch_wrap", [True, False])
131
@pytest.mark.parametrize("symbolic_trace", [True, False])
132
+ @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
133
def test_trace_sees_single_custom_op(
133
- self, symbolic_trace: bool, enable_torch_wrap: bool
134
+ self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str
135
):
136
+ op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
137
+ torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
138
+
139
def fn(x, y):
136
- return _custom_add(x, y)
140
+ return op_fn(x, y)
141
142
def find_fn(target: Any, gm: fx.GraphModule):
143
return gm.graph.find_nodes(op="call_function", target=target)
@@ -155,7 +159,7 @@ class TestIrOpCustomAdd:
159
torch.testing.assert_close(out_fx, out_eager)
160
161
# check that IR nodes only appear if enable_torch_wrap=True
158
- ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm)
162
+ ir_nodes = find_fn(torch_op, gm)
163
if enable_torch_wrap:
164
assert len(ir_nodes) == 1, gm.code
165
else:
@@ -167,7 +171,7 @@ class TestIrOpCustomAdd:
171
else:
172
gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2))
173
170
- ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm)
174
+ ir_nodes = find_fn(torch_op, gm)
175
assert len(ir_nodes) == 1, gm.code
176
177
@@ -176,9 +180,12 @@ def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
180
return x + y + 10
181
182
179
-@_custom_add.register_impl("impl_b")
183
+@_custom_add.register_impl("impl_b", inplace=True)
184
def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
181
- return x + y + 20
185
+ """Computes x+y+20"""
186
+ x.add_(y)
187
+ x.add_(20)
188
+ return x
189
190
191
@_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0)
@@ -243,19 +250,23 @@ class TestIrOpImplDispatch:
250
# Restored to empty
251
assert _custom_add.get_priority() == []
252
246
- def test_dispatch_priority_order(self):
253
+ @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
254
+ def test_dispatch_priority_order(self, overload: str):
255
+ op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
256
+ torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
257
+
258
x = torch.tensor(1, dtype=torch.int32)
259
y = torch.tensor(2, dtype=torch.int32)
260
261
with _custom_add.set_priority(["impl_b", "impl_a"]):
262
assert _custom_add.dispatch(x, y) is impl_b
252
- out1 = _custom_add(x, y)
253
- out2 = torch.ops.vllm_ir._custom_add(x, y)
263
+ out1 = op_fn(x.clone(), y)
264
+ out2 = torch_op(x.clone(), y)
265
266
with _custom_add.set_priority(["impl_a"]):
267
assert _custom_add.dispatch(x, y) is impl_a
257
- out3 = _custom_add(x, y)
258
- out4 = torch.ops.vllm_ir._custom_add(x, y)
268
+ out3 = op_fn(x.clone(), y)
269
+ out4 = torch_op(x.clone(), y)
270
271
# impl_b
272
assert out1.item() == 1 + 2 + 20
@@ -265,18 +276,18 @@ class TestIrOpImplDispatch:
276
assert out4.item() == 1 + 2 + 10
277
278
def test_unsupported_impl_filtered(self):
268
- @_custom_add.register_impl("unsupported", supported=False)
269
- def impl_bad(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
279
+ @_custom_add.register_impl("impl_unsupported", supported=False)
280
+ def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
281
return x + y + 999
282
283
x = torch.tensor(1, dtype=torch.int32)
284
y = torch.tensor(2, dtype=torch.int32)
285
275
- with _custom_add.set_priority(["unsupported", "impl_a"]):
286
+ with _custom_add.set_priority(["impl_unsupported", "impl_a"]):
287
assert _custom_add.get_priority() == ["impl_a"]
288
out = _custom_add(x, y)
289
279
- # impl_bad skipped → impl_a
290
+ # impl_unsupported skipped → impl_a
291
assert out.item() == 1 + 2 + 10
292
293
def test_supports_args_runtime_dispatch_and_warning(
tests/kernels/ir/test_layernorm.py
+205
-1
@@ -28,7 +28,9 @@ def test_rms_norm_registration():
28
"native": True,
29
"vllm_c": current_platform.is_cuda_alike(),
30
"aiter": current_platform.is_rocm(),
31
- "oink": False,
31
+ "oink": current_platform.has_device_capability(100)
32
+ and hasattr(torch.ops, "oink")
33
+ and hasattr(torch.ops.oink, "rmsnorm"),
34
"xpu_kernels": current_platform.is_xpu(),
35
}
36
@@ -67,6 +69,14 @@ class TestRMSNorm:
69
out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon)
70
torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3)
71
72
+ # Mean square should be approximately 1 (ignoring epsilon and weight scaling)
73
+ combined_norm = out.float() / weight.float()
74
+ variance = combined_norm.pow(2).mean(dim=-1)
75
+ # After RMS normalization, variance should be close to 1
76
+ torch.testing.assert_close(
77
+ variance, torch.ones_like(variance), rtol=1e-2, atol=1e-2
78
+ )
79
+
80
# Check behavior with and without weight
81
weight1 = torch.ones_like(weight)
82
out3 = rms_norm_native(x, weight1, epsilon=epsilon)
@@ -129,3 +139,197 @@ def test_aiter_rejects_unsupported_dtypes():
139
num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
140
)
141
assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
142
+
143
+
144
+fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn
145
+
146
+
147
+@pytest.mark.skipif(
148
+ not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
149
+ reason="Currently only kernels on CUDA, ROCm and XPU",
150
+)
151
+def test_fused_add_rms_norm_registration():
152
+ expected = {
153
+ "native": True,
154
+ "vllm_c": current_platform.is_cuda_alike(),
155
+ "aiter": current_platform.is_rocm(),
156
+ "oink": current_platform.has_device_capability(100)
157
+ and hasattr(torch.ops, "oink")
158
+ and hasattr(torch.ops.oink, "fused_add_rms_norm"),
159
+ "xpu_kernels": current_platform.is_xpu(),
160
+ }
161
+
162
+ actual = {
163
+ provider: impl.supported
164
+ for provider, impl in ir.ops.fused_add_rms_norm.impls.items()
165
+ }
166
+
167
+ assert actual == expected
168
+
169
+
170
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
171
+@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
172
+@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
173
+@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
174
+@pytest.mark.skipif(
175
+ not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
176
+ reason="Currently only kernels on CUDA, ROCm and XPU",
177
+)
178
+class TestFusedAddRMSNorm:
179
+ @classmethod
180
+ def setup_class(cls, **kwargs):
181
+ torch.set_default_device(current_platform.device_type)
182
+
183
+ def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
184
+ x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
185
+ num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
186
+ )
187
+ out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)
188
+
189
+ # Check shape, dtype, device
190
+ assert out.shape == x.shape
191
+ assert out.dtype == x.dtype
192
+ assert out.device == x.device
193
+ assert residual_out.shape == x_residual.shape
194
+ assert residual_out.dtype == x_residual.dtype
195
+ assert residual_out.device == x_residual.device
196
+
197
+ # Check that residual_out = x + x_residual
198
+ expected_residual = (x.float() + x_residual.float()).to(dtype)
199
+ torch.testing.assert_close(
200
+ residual_out, expected_residual, rtol=1e-3, atol=1e-3
201
+ )
202
+
203
+ # Verify that the output is RMS normalized version of (x + x_residual)
204
+ expected_out = rms_norm_native(expected_residual, weight, epsilon)
205
+ assert_close(
206
+ ir.ops.fused_add_rms_norm,
207
+ (out, residual_out),
208
+ (expected_out, expected_residual),
209
+ )
210
+
211
+ # Check the scaling property of rms norm
212
+ out1, _ = fused_add_rms_norm_native(
213
+ x, torch.zeros_like(x), weight, epsilon=epsilon
214
+ )
215
+ out2, _ = fused_add_rms_norm_native(
216
+ x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon
217
+ )
218
+ torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3)
219
+
220
+ # Check behavior with and without weight
221
+ weight1 = torch.ones_like(weight)
222
+ out3, _ = fused_add_rms_norm_native(x, x_residual, weight1, eps)
223
+ out4, _ = fused_add_rms_norm_native(x, x_residual, None, eps)
224
+ torch.testing.assert_close(out3, out4)
225
+
226
+ @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
227
+ def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
228
+ impl = ir.ops.fused_add_rms_norm.impls[provider]
229
+ x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
230
+ num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
231
+ )
232
+ args = (x, x_residual, weight, eps, None)
233
+
234
+ if not impl.supports_args(*args):
235
+ pytest.skip(f"{provider} does not support args")
236
+
237
+ ref_output, ref_residual = fused_add_rms_norm_native(*clone_args(args))
238
+ output, residual = impl.impl_fn(*clone_args(args))
239
+ assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
240
+ assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)
241
+
242
+ # check that dispatched call matches direct call
243
+ with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
244
+ out_dispatched, residual_dispatched = ir.ops.fused_add_rms_norm(*args[:4])
245
+ out_direct, residual_direct = impl.impl_fn(*clone_args(args))
246
+ torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
247
+ torch.testing.assert_close(
248
+ residual_dispatched, residual_direct, rtol=0.0, atol=0.0
249
+ )
250
+
251
+ # none of these support variance_size override
252
+ assert not impl.supports_args(x, x_residual, weight, epsilon, 4)
253
+ assert not impl.supports_args(x, x_residual, weight, epsilon, variance_size=4)
254
+
255
+ # test weight=None behavior
256
+ out_no_weight, residual_no_weight = impl.impl_fn(
257
+ x.clone(), x_residual.clone(), None, epsilon
258
+ )
259
+ out_unit_weight, residual_unit_weight = impl.impl_fn(
260
+ x.clone(), x_residual.clone(), torch.ones_like(weight), epsilon
261
+ )
262
+ assert_close(ir.ops.fused_add_rms_norm, out_no_weight, out_unit_weight)
263
+ assert_close(
264
+ ir.ops.fused_add_rms_norm, residual_no_weight, residual_unit_weight
265
+ )
266
+
267
+ @pytest.mark.parametrize("provider", ["vllm_c"])
268
+ def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider):
269
+ """Test that inplace implementations reuse inputs,
270
+ for maybe_inplace overload but not for default overload."""
271
+ impl = ir.ops.fused_add_rms_norm.impls[provider]
272
+ if not impl.supported:
273
+ pytest.skip(f"{provider} impl not supported on this platform")
274
+
275
+ x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
276
+ num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
277
+ )
278
+
279
+ # Test default overload - should NOT modify inputs even with inplace impl
280
+ x_default = x.clone()
281
+ x_residual_default = x_residual.clone()
282
+ x_default_ptr = x_default.data_ptr()
283
+ x_residual_default_ptr = x_residual_default.data_ptr()
284
+
285
+ with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
286
+ out_default, residual_default = ir.ops.fused_add_rms_norm(
287
+ x_default, x_residual_default, weight, eps
288
+ )
289
+
290
+ # Default should NOT be inplace (even with inplace implementation)
291
+ assert out_default.data_ptr() != x_default_ptr
292
+ assert residual_default.data_ptr() != x_residual_default_ptr
293
+ torch.testing.assert_close(x, x_default, rtol=0.0, atol=0.0)
294
+ torch.testing.assert_close(x_residual, x_residual_default, rtol=0.0, atol=0.0)
295
+
296
+ # Test maybe_inplace overload - should modify inputs with inplace impl
297
+ x_inplace = x.clone()
298
+ x_residual_inplace = x_residual.clone()
299
+ x_inplace_ptr = x_inplace.data_ptr()
300
+ x_residual_inplace_ptr = x_residual_inplace.data_ptr()
301
+
302
+ with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
303
+ out_inplace, residual_inplace = ir.ops.fused_add_rms_norm.maybe_inplace(
304
+ x_inplace, x_residual_inplace, weight, eps
305
+ )
306
+
307
+ # maybe_inplace should be inplace
308
+ assert out_inplace.data_ptr() == x_inplace_ptr
309
+ assert residual_inplace.data_ptr() == x_residual_inplace_ptr
310
+
311
+ # Both should produce same results
312
+ torch.testing.assert_close(out_default, out_inplace, atol=0.0, rtol=0.0)
313
+ torch.testing.assert_close(
314
+ residual_default, residual_inplace, atol=0.0, rtol=0.0
315
+ )
316
+
317
+ @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
318
+ def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
319
+ args = ir.ops.fused_add_rms_norm.generate_inputs(
320
+ num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
321
+ )
322
+ args = args + (None,) # Add variance_size parameter
323
+
324
+ # When checking the torch op, we have to set priority and use dispatch
325
+ with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
326
+ torch.library.opcheck(torch.ops.vllm_ir.fused_add_rms_norm.default, args)
327
+
328
+ # Only test maybe_inplace with non-inplace implementations
329
+ # Inplace implementations return aliases of inputs which is not allowed.
330
+ # We break this invariant, but we also convert maybe_inplace to the default
331
+ # overload during compilation, so maybe_inplace never reaches Inductor.
332
+ if not ir.ops.fused_add_rms_norm.impls[provider].inplace:
333
+ torch.library.opcheck(
334
+ torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace, args
335
+ )
tests/model_executor/test_enabled_custom_ops.py
+1
-28
@@ -23,11 +23,7 @@ from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
23
vllm_topk_sigmoid,
24
vllm_topk_softmax,
25
)
26
-from vllm.model_executor.layers.layernorm import (
27
- RMSNorm,
28
- dispatch_rocm_rmsnorm_func,
29
- fused_add_rms_norm,
30
-)
26
+from vllm.model_executor.layers.layernorm import RMSNorm
27
from vllm.platforms import current_platform
28
29
RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
@@ -153,26 +149,3 @@ def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
149
assert topk_func == rocm_aiter_ops.topk_sigmoid
150
else:
151
assert topk_func == vllm_topk_sigmoid
156
-
157
-
158
-@pytest.mark.parametrize("add_residual", [False])
159
-@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
160
-@pytest.mark.parametrize("use_rocm_aiter", [True, False])
161
-@pytest.mark.skipif(
162
- not current_platform.is_rocm(), reason="AITER is a feature exclusive for ROCm"
163
-)
164
-def test_rms_norm_dispatch(
165
- add_residual: bool, dtype: torch.dtype, use_rocm_aiter: bool
166
-):
167
- rms_norm_func = dispatch_rocm_rmsnorm_func(dtype, use_rocm_aiter)
168
-
169
- should_use_rocm_aiter = (
170
- current_platform.is_rocm()
171
- and use_rocm_aiter
172
- and dtype in RMS_NORM_SUPPORTED_DTYPES
173
- )
174
-
175
- if should_use_rocm_aiter:
176
- assert rms_norm_func == rocm_aiter_ops.rms_norm2d_with_add
177
- else:
178
- assert rms_norm_func == fused_add_rms_norm
tests/model_executor/test_oink_integration.py
+83
-46
@@ -1,60 +1,97 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import multiprocessing
4
import types
5
6
import pytest
7
-import torch
8
-
9
-
10
-def _load_oink_ops_module():
11
- # Import the module normally (vllm is installed as an editable package in CI).
12
- from vllm import _oink_ops
13
-
14
- return _oink_ops
15
-
7
17
-def test_oink_availability_checks(monkeypatch: pytest.MonkeyPatch):
18
- _oink_ops = _load_oink_ops_module()
19
-
20
- # Ensure the ops namespace exists and is mutable for tests.
21
- monkeypatch.setattr(
22
- torch.ops,
23
- "oink",
24
- types.SimpleNamespace(rmsnorm=lambda x, w, eps: x),
25
- raising=False,
26
- )
27
-
28
- # Case 1: CUDA not available.
29
- monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
30
- assert _oink_ops.is_oink_available_for_device(0) is False
31
-
32
- # Case 2: CUDA available but < SM100.
33
- monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
34
- monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (9, 0))
35
- assert _oink_ops.is_oink_available_for_device(0) is False
36
-
37
- # Case 3: CUDA available and SM100, rmsnorm op registered.
38
- monkeypatch.setattr(torch.cuda, "get_device_capability", lambda idx: (10, 0))
39
- assert _oink_ops.is_oink_available_for_device(0) is True
40
-
41
- # fused op presence probe
42
- assert _oink_ops.has_fused_add_rms_norm() is False
43
- monkeypatch.setattr(
44
- torch.ops,
45
- "oink",
46
- types.SimpleNamespace(
47
- rmsnorm=lambda x, w, eps: x,
48
- fused_add_rms_norm=lambda x, residual, w, eps: None,
8
+from vllm.platforms import current_platform
9
+
10
+
11
+def _test_oink_availability_impl(
12
+ device_capability: tuple[int, int],
13
+ has_rmsnorm: bool,
14
+ has_fused_add_rms_norm: bool,
15
+ expected_available: bool,
16
+ expected_fused: bool,
17
+) -> None:
18
+ """Test OINK support detection with mocked state."""
19
+ import torch
20
+
21
+ from vllm import platforms
22
+
23
+ # Mock device capability (class method, override on class)
24
+ dc = platforms.interface.DeviceCapability(*device_capability)
25
+ platforms.current_platform.__class__.get_device_capability = lambda device_id=0: dc
26
+
27
+ # Mock oink ops
28
+ oink_ops = types.SimpleNamespace()
29
+ if has_rmsnorm:
30
+ oink_ops.rmsnorm = lambda x, w, eps: x
31
+ if has_fused_add_rms_norm:
32
+ oink_ops.fused_add_rms_norm = lambda x, residual, w, eps: None
33
+
34
+ torch.ops.oink = oink_ops
35
+
36
+ # Now import vllm modules with mocks in place (fresh import with mocked platform)
37
+ import vllm.kernels.oink_ops # noqa: F401
38
+ from vllm.ir.ops import fused_add_rms_norm, rms_norm
39
+
40
+ # Verify support checks
41
+ assert rms_norm.impls["oink"].supported is expected_available
42
+ assert fused_add_rms_norm.impls["oink"].supported is expected_fused
43
+
44
+
45
+@pytest.mark.parametrize(
46
+ "device_capability,has_rmsnorm,has_fused_add_rms_norm,expected_available,expected_fused",
47
+ [
48
+ # Case 1: < SM100, ops not supported
49
+ ((9, 0), True, False, False, False),
50
+ # Case 2: CUDA available and SM100, rmsnorm op registered
51
+ ((10, 0), True, False, True, False),
52
+ # Case 3: SM100 with both rmsnorm and fused_add_rms_norm
53
+ ((10, 0), True, True, True, True),
54
+ ],
55
+)
56
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
57
+def test_oink_availability_checks(
58
+ device_capability: tuple[int, int],
59
+ has_rmsnorm: bool,
60
+ has_fused_add_rms_norm: bool,
61
+ expected_available: bool,
62
+ expected_fused: bool,
63
+):
64
+ """Test OINK support detection with clean import state for each parameter set."""
65
+
66
+ # Use spawn to run function in fresh process with clean imports
67
+ # TODO migrate to spawn utility:
68
+ # https://github.com/vllm-project/vllm/issues/41415
69
+ ctx = multiprocessing.get_context("spawn")
70
+ process = ctx.Process(
71
+ target=_test_oink_availability_impl,
72
+ args=(
73
+ device_capability,
74
+ has_rmsnorm,
75
+ has_fused_add_rms_norm,
76
+ expected_available,
77
+ expected_fused,
78
),
50
- raising=False,
79
)
52
- assert _oink_ops.has_fused_add_rms_norm() is True
80
+ process.start()
81
+ process.join()
82
+
83
+ if process.exitcode != 0:
84
+ raise AssertionError(
85
+ f"Subprocess test failed with exit code {process.exitcode}"
86
+ )
87
88
89
def test_can_view_as_2d_stride_guard():
56
- # Import the helper from the layernorm module.
57
- from vllm.model_executor.layers.layernorm import _can_view_as_2d
90
+ # No global import
91
+ import torch
92
+
93
+ # Import the helper from the kernels module.
94
+ from vllm.kernels.oink_ops import _can_view_as_2d
95
96
x = torch.zeros((2, 3, 4))
97
assert _can_view_as_2d(x) is True
tests/test_config.py
+37
-3
@@ -1295,11 +1295,14 @@ def test_ir_op_priority_default():
1295
# Assert default is applied to ops
1296
priority_config = IrOpPriorityConfig.with_default(["vllm_c", "native"])
1297
assert priority_config.rms_norm == ["vllm_c", "native"]
1298
+ assert priority_config.fused_add_rms_norm == ["vllm_c", "native"]
1299
1300
# Assert single ops override the default
1300
- assert IrOpPriorityConfig.with_default(
1301
- ["vllm_c", "native"], rms_norm=["oink", "native"]
1302
- ) == IrOpPriorityConfig(rms_norm=["oink", "native"])
1301
+ priority_config = IrOpPriorityConfig.with_default(
1302
+ ["native"], rms_norm=["oink", "native"]
1303
+ )
1304
+ assert priority_config.rms_norm == ["oink", "native"]
1305
+ assert priority_config.fused_add_rms_norm == ["native"]
1306
1307
1308
def test_ir_op_priority_str():
@@ -1318,3 +1321,34 @@ def test_ir_op_priority_str():
1321
with pytest.raises(pydantic.ValidationError):
1322
# must be list of only strings
1323
priority_config = IrOpPriorityConfig(rms_norm=["vllm_c", 4, "native"])
1324
+
1325
+
1326
+def test_ir_op_priority_ctx():
1327
+ """Test that the priority-setting context sets priority correctly."""
1328
+ from vllm import ir
1329
+ from vllm.config.kernel import IrOpPriorityConfig
1330
+
1331
+ priority = IrOpPriorityConfig.with_default(["native"], rms_norm=["vllm_c"])
1332
+ priority2 = IrOpPriorityConfig.with_default(
1333
+ ["native"], fused_add_rms_norm=["vllm_c"]
1334
+ )
1335
+ with priority.set_priority():
1336
+ assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
1337
+ assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]
1338
+ with priority2.set_priority():
1339
+ assert ir.ops.rms_norm.get_priority() == ["native"]
1340
+ assert ir.ops.fused_add_rms_norm.get_priority() == ["vllm_c", "native"]
1341
+
1342
+ # context restored
1343
+ assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
1344
+ assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]
1345
+
1346
+ with pytest.raises(ValueError), priority2.set_priority():
1347
+ assert ir.ops.rms_norm.get_priority() == ["native"]
1348
+ assert ir.ops.fused_add_rms_norm.get_priority() == ["vllm_c", "native"]
1349
+
1350
+ raise ValueError
1351
+
1352
+ # context restored even after exception
1353
+ assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
1354
+ assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]
tests/v1/determinism/test_rms_norm_batch_invariant.py
+7
-1
@@ -12,7 +12,7 @@ import torch
12
from utils import skip_unsupported
13
14
from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
15
-from vllm.model_executor.layers.layernorm import RMSNorm, fused_add_rms_norm
15
+from vllm.model_executor.layers.layernorm import RMSNorm
16
from vllm.platforms import current_platform
17
18
DEVICE_TYPE = current_platform.device_type
@@ -105,6 +105,12 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
105
dim=0,
106
)
107
108
+ def fused_add_rms_norm(x, residual, w, e) -> tuple[torch.Tensor, torch.Tensor]:
109
+ import vllm._custom_ops as ops
110
+
111
+ ops.fused_add_rms_norm(x, residual, w, e)
112
+ return x, residual
113
+
114
out_single, residual_out_single = fused_add_rms_norm(
115
x_single.clone(),
116
residual_single.clone(),
vllm/_aiter_ops.py
+1
-99
@@ -647,58 +647,6 @@ def _rocm_aiter_gemm_a8w8_blockscale_fake(
647
return Y
648
649
650
-def _rocm_aiter_rms_norm_impl(
651
- x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
652
-) -> torch.Tensor:
653
- from aiter import rms_norm
654
-
655
- if x.dim() > 2:
656
- x_original_shape = x.shape
657
- x = x.reshape(-1, x_original_shape[-1])
658
- x = rms_norm(x, weight, variance_epsilon)
659
- return x.reshape(x_original_shape)
660
-
661
- return rms_norm(x, weight, variance_epsilon)
662
-
663
-
664
-def _rocm_aiter_rms_norm_fake(
665
- x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
666
-) -> torch.Tensor:
667
- return torch.empty_like(x)
668
-
669
-
670
-def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
671
- x: torch.Tensor,
672
- residual: torch.Tensor,
673
- weight: torch.Tensor,
674
- variance_epsilon: float,
675
-) -> tuple[torch.Tensor, torch.Tensor]:
676
- from aiter import rmsnorm2d_fwd_with_add
677
-
678
- residual_out = torch.empty_like(residual)
679
- out = torch.empty_like(x)
680
- rmsnorm2d_fwd_with_add(
681
- out, # output
682
- x, # input
683
- residual, # residual input
684
- residual_out, # residual output
685
- weight,
686
- variance_epsilon,
687
- )
688
- return out, residual_out
689
-
690
-
691
-def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
692
- x: torch.Tensor,
693
- residual: torch.Tensor,
694
- weight: torch.Tensor,
695
- variance_epsilon: float,
696
-) -> tuple[torch.Tensor, torch.Tensor]:
697
- residual_out = torch.empty_like(residual)
698
- out = torch.empty_like(x)
699
- return out, residual_out
700
-
701
-
650
def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(
651
x: torch.Tensor,
652
residual: torch.Tensor,
@@ -1229,10 +1177,9 @@ class rocm_aiter_ops:
1177
1178
# Check if aiter is enabled before using operations
1179
if rocm_aiter_ops.is_enabled():
1232
- result = rocm_aiter_ops.rms_norm(x, weight, epsilon)
1180
+ result = rocm_aiter_ops.per_token_quant(x, FP8_DTYPE)
1181
1182
Operations:
1235
- - RMS normalization: rms_norm, rms_norm2d_with_add
1183
- GEMM operations: gemm_a8w8, gemm_a8w8_blockscale
1184
- Fused MoE: fused_moe, asm_moe_tkw1
1185
- Routing: topk_softmax, biased_grouped_topk, grouped_topk
@@ -1244,7 +1191,6 @@ class rocm_aiter_ops:
1191
# Check if the env variable is set
1192
_AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
1193
_LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
1247
- _RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
1194
_FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
1195
_MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
1196
_MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
@@ -1275,7 +1221,6 @@ class rocm_aiter_ops:
1221
"""
1222
cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
1223
cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
1278
- cls._RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
1224
cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
1225
cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
1226
cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
@@ -1367,11 +1312,6 @@ class rocm_aiter_ops:
1312
def is_linear_fp8_enabled(cls) -> bool:
1313
return cls.is_linear_enabled()
1314
1370
- @classmethod
1371
- @if_aiter_supported
1372
- def is_rmsnorm_enabled(cls) -> bool:
1373
- return cls._AITER_ENABLED and cls._RMSNORM_ENABLED
1374
-
1315
@classmethod
1316
@if_aiter_supported
1317
def is_fused_moe_enabled(cls) -> bool:
@@ -1560,19 +1500,6 @@ class rocm_aiter_ops:
1500
fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
1501
)
1502
1563
- direct_register_custom_op(
1564
- op_name="rocm_aiter_rms_norm",
1565
- op_func=_rocm_aiter_rms_norm_impl,
1566
- fake_impl=_rocm_aiter_rms_norm_fake,
1567
- )
1568
-
1569
- direct_register_custom_op(
1570
- op_name="rocm_aiter_rmsnorm2d_fwd_with_add",
1571
- op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
1572
- fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
1573
- dispatch_key=current_platform.dispatch_key,
1574
- )
1575
-
1503
direct_register_custom_op(
1504
op_name="rocm_aiter_rmsnorm_fused_dynamic_quant",
1505
op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl,
@@ -1672,14 +1599,6 @@ class rocm_aiter_ops:
1599
1600
_OPS_REGISTERED = True
1601
1675
- @staticmethod
1676
- def get_rmsnorm_fused_add_op() -> OpOverload:
1677
- return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add.default
1678
-
1679
- @staticmethod
1680
- def get_rmsnorm_op() -> OpOverload:
1681
- return torch.ops.vllm.rocm_aiter_rms_norm.default
1682
-
1602
@staticmethod
1603
def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload:
1604
return torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant.default
@@ -1724,23 +1643,6 @@ class rocm_aiter_ops:
1643
def get_fused_mla_dual_rms_norm_op() -> OpOverload:
1644
return torch.ops.vllm.fused_mla_dual_rms_norm.default
1645
1727
- @staticmethod
1728
- def rms_norm(
1729
- x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
1730
- ) -> torch.Tensor:
1731
- return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon)
1732
-
1733
- @staticmethod
1734
- def rms_norm2d_with_add(
1735
- x: torch.Tensor,
1736
- residual: torch.Tensor,
1737
- weight: torch.Tensor,
1738
- variance_epsilon: float,
1739
- ) -> tuple[torch.Tensor, torch.Tensor]:
1740
- return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add(
1741
- x, residual, weight, variance_epsilon
1742
- )
1743
-
1646
@staticmethod
1647
def w8a8_gemm(
1648
A: torch.Tensor,
vllm/_oink_ops.py
deleted
-96
@@ -1,96 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Small helper wrappers for external Oink Blackwell custom ops.
4
-
5
-vLLM does not depend on the external Oink repository/package. When an external
6
-plugin registers torch.library.custom_op entrypoints under the `oink::`
7
-namespace (e.g. via vLLM's general_plugins mechanism) and
8
-`VLLM_USE_OINK_OPS=1` is set, vLLM can route eligible calls to those ops.
9
-
10
-This module provides:
11
-- A single place to probe Oink op availability at module init time
12
- (outside torch.compile tracing), and
13
-- Thin wrappers around the torch.ops entrypoints for use in CUDA fast paths,
14
- without introducing graph breaks.
15
-
16
-Important:
17
- Do not call the availability helpers in a compiled region. They may call
18
- functions decorated with `torch._dynamo.disable` to safely check
19
- conditions that should not be traced.
20
-"""
21
-
22
-from __future__ import annotations
23
-
24
-from collections.abc import Callable
25
-
26
-import torch
27
-
28
-try:
29
- from torch._dynamo import disable as _dynamo_disable # type: ignore[attr-defined]
30
-except Exception: # pragma: no cover
31
-
32
- def _dynamo_disable(fn: Callable): # type: ignore[misc]
33
- return fn
34
-
35
-
36
-def _has_oink_op(op_name: str) -> bool:
37
- """Check if a specific oink op is registered."""
38
- return hasattr(torch.ops, "oink") and hasattr(torch.ops.oink, op_name)
39
-
40
-
41
-@_dynamo_disable
42
-def is_oink_available_for_device(device_index: int) -> bool:
43
- """Return True if Oink ops are registered and device is SM100+.
44
-
45
- This function is intended to be called during module initialization
46
- (e.g., in RMSNorm.__init__), not in the forward path.
47
-
48
- External plugins are expected to gate registration on SM100+ and
49
- VLLM_USE_OINK_OPS=1, so if the ops are present they should be usable.
50
- """
51
- if not torch.cuda.is_available():
52
- return False
53
-
54
- try:
55
- major, minor = torch.cuda.get_device_capability(device_index)
56
- sm = 10 * major + minor
57
- if sm < 100:
58
- return False
59
- except Exception:
60
- return False
61
-
62
- return _has_oink_op("rmsnorm")
63
-
64
-
65
-def has_fused_add_rms_norm() -> bool:
66
- """Return True if the in-place fused op is registered."""
67
- return _has_oink_op("fused_add_rms_norm")
68
-
69
-
70
-def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
71
- """Call `torch.ops.oink.rmsnorm`.
72
-
73
- This wrapper is safe to call in torch.compile regions.
74
- """
75
- return torch.ops.oink.rmsnorm(x, weight, eps)
76
-
77
-
78
-def fused_add_rms_norm_(
79
- x: torch.Tensor,
80
- residual: torch.Tensor,
81
- weight: torch.Tensor,
82
- eps: float,
83
-) -> None:
84
- """Call `torch.ops.oink.fused_add_rms_norm` (mutates x and residual)."""
85
- torch.ops.oink.fused_add_rms_norm(x, residual, weight, eps)
86
-
87
-
88
-def fused_add_rms_norm(
89
- x: torch.Tensor,
90
- residual: torch.Tensor,
91
- weight: torch.Tensor,
92
- eps: float,
93
-) -> tuple[torch.Tensor, torch.Tensor]:
94
- """Convenience wrapper returning (x, residual) after in-place mutation."""
95
- fused_add_rms_norm_(x, residual, weight, eps)
96
- return x, residual
vllm/compilation/backends.py
+19
@@ -50,6 +50,7 @@ from .partition_rules import (
50
should_split,
51
)
52
from .passes.inductor_pass import InductorPass, pass_context
53
+from .passes.ir.inplace_functionalization import VllmIRInplaceFunctionalizationPass
54
from .passes.pass_manager import PostGradPassManager
55
56
logger = init_logger(__name__)
@@ -926,6 +927,24 @@ class VllmBackend:
927
return standalone_compile_artifacts, sym_shape_indices_map, returns_tuple_map
928
929
def configure_post_pass(self) -> None:
930
+ # TODO proper PassManager?
931
+ pre_grad_pass_key = "pre_grad_custom_pass"
932
+ assert self.pass_key != pre_grad_pass_key
933
+ assert pre_grad_pass_key not in self.inductor_config
934
+ self.inductor_config[pre_grad_pass_key] = VllmIRInplaceFunctionalizationPass(
935
+ self.vllm_config
936
+ )
937
+
938
+ # Make sure pre_grad_custom_pass is not pickled
939
+ # as part of AOTAutograd built-in cache key
940
+ # TODO(luka) is there a cleaner way to do this
941
+ import torch._inductor.config as inductor_config
942
+
943
+ ignore = inductor_config._cache_config_ignore_prefix + [pre_grad_pass_key]
944
+ assert "_cache_config_ignore_prefix" not in self.inductor_config
945
+ self.inductor_config["_cache_config_ignore_prefix"] = ignore
946
+
947
+ # Configure the (nominally post-grad) pass manager
948
self.pass_manager.configure(self.vllm_config)
949
950
# Post-grad custom passes are run using the post_grad_custom_post_pass
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+21
-17
@@ -40,7 +40,7 @@ from ..vllm_inductor_pass import (
40
VllmPatternMatcherPass,
41
VllmPatternReplacement,
42
)
43
-from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8
43
+from .matcher_utils import MatcherQuantFP8
44
45
FP8_DTYPE = current_platform.fp8_dtype()
46
@@ -356,10 +356,11 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
356
super().__init__(dtype, device)
357
self.epsilon = epsilon
358
self.allreduce_params = allreduce_params
359
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
359
360
def get_inputs(self) -> list[torch.Tensor]:
362
- input, residual, weight = self.rmsnorm_matcher.inputs()
361
+ input = self.empty(5, 16)
362
+ residual = self.empty(5, 16)
363
+ weight = self.empty(16)
364
365
# input goes through allreduce first, always 16-bit
366
return [residual, input.to(self.dtype), weight]
@@ -369,7 +370,9 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
370
residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
371
) -> tuple[torch.Tensor, torch.Tensor]:
372
allreduce_output = tensor_model_parallel_all_reduce(input)
372
- rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
373
+ rms, residual = vllm.ir.ops.fused_add_rms_norm(
374
+ allreduce_output, residual, weight, self.epsilon
375
+ )
376
return rms, residual
377
378
def replacement(
@@ -503,11 +506,12 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
506
self.allreduce_params = allreduce_params
507
self.quant_dtype = torch.float8_e4m3fn
508
506
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
509
self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
510
511
def get_inputs(self) -> list[torch.Tensor]:
510
- input, residual, weight = self.rmsnorm_matcher.inputs()
512
+ input = self.empty(5, 16)
513
+ residual = self.empty(5, 16)
514
+ weight = self.empty(16)
515
_, scale = self.quant_matcher.inputs()
516
517
# input goes through allreduce first, always 16-bit
@@ -521,7 +525,9 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):
525
scale: torch.Tensor,
526
) -> tuple[torch.Tensor, torch.Tensor]:
527
allreduce_output = tensor_model_parallel_all_reduce(input)
524
- rms, res = self.rmsnorm_matcher(allreduce_output, weight, residual)
528
+ rms, res = vllm.ir.ops.fused_add_rms_norm(
529
+ allreduce_output, residual, weight, self.epsilon
530
+ )
531
quant, _ = self.quant_matcher(rms, scale)
532
533
return quant, res
@@ -668,7 +674,6 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
674
super().__init__(dtype, device)
675
self.epsilon = epsilon
676
self.allreduce_params = allreduce_params
671
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
677
678
def get_inputs(self) -> list[torch.Tensor]:
679
input = torch.empty([16, 16], device=self.device, dtype=self.dtype)
@@ -700,7 +705,9 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
705
input_global_scale: torch.Tensor,
706
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
707
allreduce_output = tensor_model_parallel_all_reduce(input)
703
- rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
708
+ rms, residual = vllm.ir.ops.fused_add_rms_norm(
709
+ allreduce_output, residual, weight, self.epsilon
710
+ )
711
quant_out_tuple = auto_functionalized(
712
STATIC_FP4_QUANT_OP,
713
input=rms,
@@ -955,15 +962,11 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
962
super().__init__(dtype, device)
963
self.epsilon = epsilon
964
self.dtype = dtype
958
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(
959
- epsilon, match_rocm_aiter=use_aiter_rmsnorm
960
- )
965
self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()
966
967
def get_inputs(self) -> list[torch.Tensor]:
964
- input, residual, weight = self.rmsnorm_matcher.inputs()
965
-
966
- return [residual, input.to(self.dtype), weight]
968
+ # input, residual, weight
969
+ return [self.empty(5, 16), self.empty(5, 16), self.empty(16)]
970
971
@property
972
def pattern(self):
@@ -971,8 +974,9 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
974
residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
975
) -> tuple[torch.Tensor, torch.Tensor]:
976
allreduce_output = tensor_model_parallel_all_reduce(input)
974
- rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
975
-
977
+ rms, residual = vllm.ir.ops.fused_add_rms_norm(
978
+ allreduce_output, residual, weight, self.epsilon
979
+ )
980
return rms, residual
981
982
return _pattern
vllm/compilation/passes/fusion/matcher_utils.py
-67
@@ -10,7 +10,6 @@ from torch._ops import OpOverload
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import get_current_vllm_config
12
from vllm.model_executor.layers.activation import SiluAndMul
13
-from vllm.model_executor.layers.layernorm import RMSNorm
13
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
14
from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
GroupShape,
@@ -159,72 +158,6 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
158
return result
159
160
162
-class MatcherFusedAddRMSNorm(MatcherCustomOp):
163
- def __init__(
164
- self,
165
- epsilon: float,
166
- enabled: bool | None = None,
167
- match_rocm_aiter: bool = False,
168
- ) -> None:
169
- if enabled is None:
170
- enabled = RMSNorm.enabled()
171
-
172
- super().__init__(enabled)
173
- self.epsilon = epsilon
174
- self.match_rocm_aiter = match_rocm_aiter
175
-
176
- self._rmsnorm_op = RMS_ADD_OP
177
-
178
- if match_rocm_aiter:
179
- self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_fused_add_op()
180
-
181
- def inputs(self) -> list[torch.Tensor]:
182
- input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
183
- weight = self.empty(16)
184
- residual = self.empty(5, 16)
185
- return [input, weight, residual]
186
-
187
- def forward_rocm_aiter(
188
- self,
189
- input: torch.Tensor,
190
- weight: torch.Tensor,
191
- residual: torch.Tensor,
192
- ) -> tuple[torch.Tensor, torch.Tensor]:
193
- return self._rmsnorm_op( # type: ignore[no-any-return]
194
- x=input, residual=residual, weight=weight, variance_epsilon=self.epsilon
195
- )
196
-
197
- def forward_custom(
198
- self,
199
- input: torch.Tensor,
200
- weight: torch.Tensor,
201
- residual: torch.Tensor,
202
- ) -> tuple[torch.Tensor, torch.Tensor]:
203
- if self.match_rocm_aiter:
204
- return self.forward_rocm_aiter(input, weight, residual)
205
-
206
- _, result, residual = auto_functionalized(
207
- self._rmsnorm_op,
208
- input=input,
209
- residual=residual,
210
- weight=weight,
211
- epsilon=self.epsilon,
212
- )
213
-
214
- return result, residual
215
-
216
- def forward_native(
217
- self,
218
- input: torch.Tensor,
219
- weight: torch.Tensor,
220
- residual: torch.Tensor,
221
- ) -> tuple[torch.Tensor, torch.Tensor]:
222
- result: tuple[torch.Tensor, torch.Tensor] = RMSNorm.forward_static(
223
- input, self.epsilon, input.size(-1), self.model_dtype, weight, residual
224
- )
225
- return result
226
-
227
-
161
class MatcherQuantFP8(MatcherCustomOp):
162
def __init__(
163
self,
vllm/compilation/passes/fusion/rms_quant_fusion.py
+26
-12
@@ -29,7 +29,6 @@ from vllm.platforms import current_platform
29
from ..inductor_pass import enable_fake_mode
30
from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
31
from .matcher_utils import (
32
- MatcherFusedAddRMSNorm,
32
MatcherQuantFP8,
33
)
34
@@ -146,9 +145,6 @@ class RMSNormQuantPattern:
145
assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"
146
self.FUSED_OP = FUSED_OPS[key]
147
149
- if key.fused_add:
150
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
151
-
148
self.quant_matcher = MatcherQuantFP8(
149
key.quant,
150
has_col_major_scales=has_col_major_scales,
@@ -231,7 +227,9 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
227
residual: torch.Tensor,
228
scale: torch.Tensor,
229
) -> tuple[torch.Tensor, torch.Tensor]:
234
- result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
230
+ result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
231
+ input, residual, weight, self.epsilon
232
+ )
233
result, _ = self.quant_matcher(result_rms, scale)
234
235
return result, residual
@@ -261,8 +259,9 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
259
return at[1], at[2]
260
261
inputs = [
264
- # input, weight, residual
265
- *self.rmsnorm_matcher.inputs(),
262
+ empty_bf16(5, 16), # input
263
+ empty_bf16(16), # weight
264
+ empty_bf16(5, 16), # residual
265
self.quant_matcher.inputs()[1], # scale
266
]
267
@@ -311,7 +310,9 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
310
residual: torch.Tensor,
311
scale: torch.Tensor,
312
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
314
- result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
313
+ result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
314
+ input, residual, weight, self.epsilon
315
+ )
316
result = torch.empty(
317
result_rms.shape,
318
device=result_rms.device,
@@ -366,12 +367,17 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
367
# result, residual, scale
368
return at[1], at[3], at[2]
369
369
- scale = self.quant_matcher.empty_f32(1, 1)
370
+ inputs = [
371
+ empty_bf16(5, 16), # input
372
+ empty_bf16(16), # weight
373
+ empty_bf16(5, 16), # residual
374
+ self.quant_matcher.empty_f32(1, 1), # scale
375
+ ]
376
377
pm.register_replacement(
378
pattern,
379
replacement,
374
- self.rmsnorm_matcher.inputs() + [scale],
380
+ inputs,
381
pm.fwd_only,
382
pm_pass,
383
extra_check=_rms_input_weight_dtype_match,
@@ -552,7 +558,9 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
558
def pattern(
559
input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
560
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
555
- result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
561
+ result_rms, residual = vllm.ir.ops.fused_add_rms_norm(
562
+ input, residual, weight, self.epsilon
563
+ )
564
result, scale = self.quant_matcher(result_rms)
565
566
return result, residual, scale
@@ -580,10 +588,16 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
588
# result, residual, scale
589
return at[1], at[3], at[2]
590
591
+ inputs = [
592
+ empty_bf16(5, 16), # input
593
+ empty_bf16(16), # weight
594
+ empty_bf16(5, 16), # residual
595
+ ]
596
+
597
pm.register_replacement(
598
pattern,
599
replacement,
586
- self.rmsnorm_matcher.inputs(),
600
+ inputs,
601
pm.fwd_only,
602
pm_pass,
603
extra_check=_rms_input_weight_dtype_match,
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+30
-16
@@ -29,7 +29,6 @@ from ..vllm_inductor_pass import (
29
VllmPatternReplacement,
30
)
31
from .matcher_utils import (
32
- MatcherFusedAddRMSNorm,
32
MatcherQuantFP8,
33
MatcherSiluAndMul,
34
)
@@ -49,10 +48,6 @@ class AiterRMSNormQuantPattern:
48
self.quant_dtype = key.quant.dtype
49
self.device = torch.device("cuda")
50
52
- if key.fused_add:
53
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(
54
- epsilon, match_rocm_aiter=True
55
- )
51
self.quant_matcher = MatcherQuantFP8(
52
key.quant,
53
match_rocm_aiter=match_aiter_quant,
@@ -145,7 +140,9 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
140
weight: torch.Tensor,
141
residual: torch.Tensor,
142
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
148
- result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
143
+ result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
144
+ input, residual, weight, self.epsilon
145
+ )
146
result, scale = self.quant_matcher(result_rms)
147
148
return result, residual_out, scale
@@ -163,10 +160,16 @@ class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
160
161
return result[0], result[1], result[2]
162
163
+ inputs = [
164
+ self.empty(5, 16), # input
165
+ self.empty(16), # weight
166
+ self.empty(5, 16), # residual
167
+ ]
168
+
169
pm.register_replacement(
170
pattern,
171
replacement,
169
- self.rmsnorm_matcher.inputs(),
172
+ inputs,
173
pm.fwd_only,
174
pm_pass,
175
)
@@ -258,7 +261,9 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
261
weight: torch.Tensor,
262
residual: torch.Tensor,
263
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
261
- result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
264
+ result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
265
+ input, residual, weight, self.epsilon
266
+ )
267
result, scale = self.quant_matcher(result_rms)
268
269
return result, residual_out, scale
@@ -279,9 +284,13 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
284
# result, scale, residual
285
return at[0], at[1], at[2]
286
282
- pm.register_replacement(
283
- pattern, replacement, self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass
284
- )
287
+ inputs = [
288
+ self.empty(5, 16), # input
289
+ self.empty(16), # weight
290
+ self.empty(5, 16), # residual
291
+ ]
292
+
293
+ pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
294
295
296
class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
@@ -420,12 +429,15 @@ class AddAiterRMSNormPadPattern:
429
self.epsilon = epsilon
430
self.hidden_size = hidden_size
431
self.x_pad_to_multiple = x_pad_to_multiple
423
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon, match_rocm_aiter=True)
432
433
def get_inputs(self) -> list[torch.Tensor]:
426
- input, weight, residual = self.rmsnorm_matcher.inputs()
427
- router_weight = torch.empty([8, 16], dtype=weight.dtype, device=weight.device)
428
- router_bias = torch.empty([8], dtype=weight.dtype, device=weight.device)
434
+ device = torch.device("cuda")
435
+ dtype = torch.bfloat16
436
+ input = torch.empty(5, 16, dtype=dtype, device=device)
437
+ weight = torch.empty(16, dtype=dtype, device=device)
438
+ residual = torch.empty(5, 16, dtype=dtype, device=device)
439
+ router_weight = torch.empty([8, 16], dtype=dtype, device=device)
440
+ router_bias = torch.empty([8], dtype=dtype, device=device)
441
return [input, weight, residual, router_weight, router_bias]
442
443
def register(self, pm_pass: PatternMatcherPass) -> None:
@@ -439,7 +451,9 @@ class AddAiterRMSNormPadPattern:
451
pad_size = self.x_pad_to_multiple - (
452
self.hidden_size % self.x_pad_to_multiple
453
)
442
- result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
454
+ result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
455
+ input, residual, weight, self.epsilon
456
+ )
457
router_logits = torch.ops.vllm.rocm_unquantized_gemm(
458
result_rms, router_weight, router_bias
459
)
vllm/compilation/passes/fusion/sequence_parallelism.py
+11
-9
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
from ..inductor_pass import enable_fake_mode
24
from ..utility.noop_elimination import NoOpEliminationPass
25
from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
26
-from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8
26
+from .matcher_utils import MatcherQuantFP8
27
28
logger = init_logger(__name__)
29
@@ -174,7 +174,6 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
174
class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
175
def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
176
super().__init__(epsilon, dtype, device)
177
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
177
178
def get_inputs(self) -> list[torch.Tensor]:
179
mm_1 = torch.empty([4, 4], device=self.device, dtype=self.dtype)
@@ -195,7 +194,9 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
194
rms_norm_weights: torch.Tensor,
195
) -> tuple[torch.Tensor, torch.Tensor]:
196
all_reduce = self._all_reduce(mm_1)
198
- rmsnorm = self.rmsnorm_matcher(all_reduce, rms_norm_weights, residual)
197
+ rmsnorm = vllm.ir.ops.fused_add_rms_norm(
198
+ all_reduce, residual, rms_norm_weights, self.epsilon
199
+ )
200
return rmsnorm[0], rmsnorm[1]
201
202
def replacement(
@@ -208,7 +209,9 @@ class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
209
# once the seqpar pattern with the previous rmsnorm is replaced
210
reduce_scatter = self._reduce_scatter(mm_1)
211
residual = residual[0 : reduce_scatter.size(0), ...]
211
- rmsnorm = self.rmsnorm_matcher(reduce_scatter, rms_norm_weights, residual)
212
+ rmsnorm = vllm.ir.ops.fused_add_rms_norm(
213
+ reduce_scatter, residual, rms_norm_weights, self.epsilon
214
+ )
215
all_gather = self._all_gather(rmsnorm[0])
216
# shape of residual changes but that's fine,
217
# next node is already slicing it, now becomes a noop
@@ -271,7 +274,6 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
274
class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
275
def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
276
super().__init__(epsilon, dtype, device)
274
- self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
277
self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
278
279
def get_inputs(self) -> list[torch.Tensor]:
@@ -290,8 +292,8 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
292
scale: torch.Tensor,
293
) -> tuple[torch.Tensor, torch.Tensor]:
294
all_reduce = self._all_reduce(mm_1)
293
- rms, residual_out = self.rmsnorm_matcher(
294
- all_reduce, rms_norm_weights, residual
295
+ rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
296
+ all_reduce, residual, rms_norm_weights, self.epsilon
297
)
298
quant, _ = self.quant_matcher(rms, scale)
299
return quant, residual_out
@@ -308,8 +310,8 @@ class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
310
# once the seqpar pattern with the previous rmsnorm is replaced
311
reduce_scatter = self._reduce_scatter(mm_1)
312
residual = residual[0 : reduce_scatter.size(0), ...]
311
- rms, residual_out = self.rmsnorm_matcher(
312
- reduce_scatter, rms_norm_weights, residual
313
+ rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
314
+ reduce_scatter, residual, rms_norm_weights, self.epsilon
315
)
316
quant, _ = self.quant_matcher(rms, scale)
317
all_gather = self._all_gather(quant)
vllm/compilation/passes/inductor_pass.py
+3
@@ -30,6 +30,9 @@ class PassContext:
30
def __init__(self, compile_range: Range):
31
self.compile_range: Range = compile_range
32
33
+ # set of arg indices
34
+ self.donated_input_ids: set[int] = set()
35
+
36
37
def get_pass_context() -> PassContext:
38
"""Get the current pass context."""
vllm/compilation/passes/ir/clone_elimination.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+from torch import fx
5
+from torch._higher_order_ops.auto_functionalize import auto_functionalized
6
+from torch._higher_order_ops.triton_kernel_wrap import TritonKernelWrapperFunctional
7
+from torch._ops import HigherOrderOperator, OpOverload
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.logger import init_logger
11
+
12
+from ..fx_utils import is_func
13
+from ..inductor_pass import get_pass_context
14
+from ..vllm_inductor_pass import VllmInductorPass
15
+
16
+logger = init_logger(__name__)
17
+
18
+
19
+def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool:
20
+ if user.op == "output":
21
+ return False
22
+
23
+ if is_func(user, auto_functionalized):
24
+ # While autofunc writes to the node,
25
+ # this is a follow-up use we're not interested in.
26
+ # It is also guaranteed to be the final use,
27
+ # as auto_functionalized returns the tensor back for follow-up use.
28
+ return False
29
+ elif user.op == "call_function" and isinstance(user.target, HigherOrderOperator):
30
+ # By default, be conservative, assume this could be a write
31
+ # (except functional HOPs)
32
+ return not isinstance(user.target, TritonKernelWrapperFunctional)
33
+
34
+ assert isinstance(user.target, OpOverload), (
35
+ f"{node=} {user=} {user.op=} {user.target=}"
36
+ )
37
+ schema = user.target._schema
38
+ assert len(user.args) <= len(schema.arguments)
39
+ for i, arg in enumerate(user.args):
40
+ # Only interested in writes to node
41
+ if arg is not node:
42
+ continue
43
+
44
+ # If not a write, next arg could be
45
+ if schema.arguments[i].is_write:
46
+ return True
47
+
48
+ # No writes found
49
+ return False
50
+
51
+
52
+class UnsafeCloneEliminationPass(VllmInductorPass):
53
+ """
54
+ This pass removes clone nodes that are no longer needed after vLLM IR lowering.
55
+ It uses donated_input_ids to eliminate clones of donated graph inputs, preserving
56
+ contents of non-donated graph inputs.
57
+
58
+ It is "unsafe" because it does not (yet) take aliasing into account. Solving
59
+ aliasing is an open problem, so this pass intends to support known vLLM cases
60
+ and not guarantee soundness on general graphs. In the future, this pass will likely
61
+ support basic forms of aliasing to handle simple views (e.g. qkv -> q,k,v).
62
+ """
63
+
64
+ def __init__(self, vllm_config: VllmConfig) -> None:
65
+ super().__init__(vllm_config)
66
+
67
+ @VllmInductorPass.time_and_log
68
+ def __call__(self, graph: fx.Graph) -> None:
69
+ count = 0
70
+ node_to_idx = {node: i for i, node in enumerate(graph.nodes)}
71
+ pass_context = get_pass_context()
72
+ donated_input_ids = pass_context.donated_input_ids
73
+ logger.debug("Donated input ids: %s", donated_input_ids)
74
+
75
+ for node in graph.nodes:
76
+ if not is_func(node, torch.ops.aten.clone.default):
77
+ continue
78
+
79
+ original_node = node.args[0]
80
+ assert isinstance(original_node, fx.Node)
81
+
82
+ # Clone needs to be preserved if node is getting written to and
83
+ # the old value is used again.
84
+ # This could only happen if an inplace implementation was lowered.
85
+ # Then node (the clone) will have one write.
86
+ # TODO(luka) hopefully this can be removed once we lower functional graphs.
87
+ write_idxs = [
88
+ node_to_idx[u] for u in node.users if user_writes_to_node(u, node)
89
+ ]
90
+ assert len(write_idxs) in (0, 1)
91
+ if write_idxs:
92
+ # Check if a user of original_node occurs after a write
93
+ write_idx = write_idxs[0]
94
+ if any(
95
+ node_to_idx[orig_user] > write_idx
96
+ for orig_user in original_node.users
97
+ ):
98
+ logger.debug(
99
+ "Clone removal not possible, "
100
+ "original_node=%s used after mutation on node=%s",
101
+ original_node,
102
+ node,
103
+ )
104
+ continue
105
+
106
+ # Check if a node is a (non-donated) graph input
107
+ if (
108
+ original_node.op == "placeholder"
109
+ and node_to_idx[original_node] not in donated_input_ids
110
+ ):
111
+ logger.debug(
112
+ "Graph input %s not donated, cannot eliminate its clone",
113
+ original_node,
114
+ )
115
+ continue
116
+
117
+ logger.debug(
118
+ "Node %s is a redundant clone node of %s, removing it",
119
+ node,
120
+ original_node,
121
+ )
122
+ node.replace_all_uses_with(original_node)
123
+ graph.erase_node(node)
124
+ count += 1
125
+
126
+ logger.debug("CloneCleanupPass removed %d clone nodes", count)
vllm/compilation/passes/ir/inplace_functionalization.py
new
+101
@@ -0,0 +1,101 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections import defaultdict
4
+
5
+from torch import fx
6
+from torch._inductor.pattern_matcher import (
7
+ PatternMatcherPass,
8
+)
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.logger import init_logger
12
+
13
+from ..inductor_pass import get_pass_context
14
+from ..vllm_inductor_pass import VllmInductorPass
15
+from .lowering_pass import get_ir_op
16
+from .utils import overload_or_default
17
+
18
+logger = init_logger(__name__)
19
+
20
+
21
+class VllmIRInplaceFunctionalizationPass(VllmInductorPass):
22
+ """
23
+ This pass functionalizes maybe_inplace vLLM IR ops to the default overload.
24
+ The maybe_inplace overloads have the same signature as the default overload
25
+ so the pass simply replaces the called overload.
26
+ That makes the graph properly functional.
27
+ The pass also validates that activations passed to maybe_inplace have no later
28
+ uses in the graph: they are donated to the maybe_inplace op call,
29
+ and their contents are not defined afterward.
30
+
31
+ This pass operates pre-AOTAutograd,
32
+ so it must handle non-normalized and non-functional IR.
33
+ """
34
+
35
+ def __init__(self, vllm_config: VllmConfig) -> None:
36
+ super().__init__(vllm_config)
37
+ self.patterns = PatternMatcherPass(self.pass_name)
38
+ self.functionalized_ops: dict[str, int] = defaultdict(lambda: 0)
39
+
40
+ @VllmInductorPass.time_and_log
41
+ def __call__(self, graph: fx.Graph) -> None:
42
+ # clear at the beginning instead of end, so that tests can inspect
43
+ self.functionalized_ops.clear()
44
+ assert graph.owning_module is not None
45
+ node_to_idx = {node: i for i, node in enumerate(graph.nodes)}
46
+
47
+ # Pass donated input via vLLM's pass context
48
+ pass_context = get_pass_context()
49
+ pass_context.donated_input_ids = set[int]()
50
+
51
+ for node in graph.nodes:
52
+ if (ir_op := get_ir_op(node)) is None:
53
+ continue
54
+
55
+ op_overload = overload_or_default(node.target)
56
+ overload_name = op_overload._overloadname
57
+ if overload_name != "maybe_inplace":
58
+ assert overload_name == "default", (
59
+ f"Found overload {overload_name} for op {ir_op.name}, "
60
+ f"expected maybe_inplace or default"
61
+ )
62
+ continue
63
+
64
+ # must have maybe_inplace overload and allow_inplace
65
+ assert ir_op.allow_inplace and hasattr(ir_op, "maybe_inplace")
66
+
67
+ # Check that activation inputs are not used after this op
68
+ for arg_idx in ir_op.activation_indices:
69
+ arg = node.args[arg_idx]
70
+ assert isinstance(arg, fx.Node), "Activation inputs must be fx.Node"
71
+ for user in arg.users:
72
+ if node_to_idx[user] > node_to_idx[node]:
73
+ raise ValueError(
74
+ f"Input {arg} to maybe_inplace node {node} "
75
+ f"is used again after the node. "
76
+ f"This is not allowed; activation inputs to maybe_inplace "
77
+ f"ops are donated to the op, meaning their memory may be "
78
+ f"recycled for outputs.\n\n"
79
+ f"To preserve the inputs, use the default overload or "
80
+ f"clone them manually beforehand."
81
+ )
82
+
83
+ if arg.op == "placeholder":
84
+ # Graph input that maybe_inplace might modify.
85
+ # Mark it so downstream passes know it's donated.
86
+ # TODO(luka) store in placeholder node meta once supported
87
+ pass_context.donated_input_ids.add(node_to_idx[arg])
88
+
89
+ # Same signature, just replace the overload that's called.
90
+ node.target = ir_op.torch_op
91
+ self.functionalized_ops[ir_op.name] += 1
92
+
93
+ count = sum(self.functionalized_ops.values())
94
+ ops = ",".join(self.functionalized_ops.keys())
95
+ logger.debug("Donated input IDs: %s", pass_context.donated_input_ids)
96
+ logger.debug(
97
+ "%s functionalized %d vLLM IR nodes for op(s) %s",
98
+ self.pass_name,
99
+ count,
100
+ ops,
101
+ )
vllm/compilation/passes/ir/lowering_pass.py
+9
-35
@@ -10,7 +10,6 @@ from torch._inductor.pattern_matcher import (
10
PatternMatcherPass,
11
register_graph_pattern,
12
)
13
-from torch._ops import OpOverload, OpOverloadPacket
13
14
from vllm.config import VllmConfig
15
from vllm.ir.op import IrOp
@@ -18,41 +17,11 @@ from vllm.logger import init_logger
17
from vllm.logging_utils import lazy
18
19
from ..vllm_inductor_pass import VllmInductorPass
20
+from .utils import get_ir_op
21
22
logger = init_logger(__name__)
23
24
25
-def get_default_overload(op: OpOverload | OpOverloadPacket) -> OpOverload:
26
- if isinstance(op, OpOverloadPacket):
27
- return op.default
28
- assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket"
29
- return op
30
-
31
-
32
-def get_ir_op(node: fx.Node) -> IrOp | None:
33
- if node.op != "call_function":
34
- return None
35
-
36
- if not isinstance(node.target, (OpOverload, OpOverloadPacket)):
37
- return None
38
-
39
- op_overload = get_default_overload(node.target)
40
- if op_overload.namespace != "vllm_ir":
41
- return None
42
-
43
- op_name = op_overload._opname
44
- if op_name not in IrOp.registry:
45
- logger.warning(
46
- "Unknown vLLM IR op %s, there's likely an issue with torch registration, "
47
- "or a torch custom op was registered in the vllm_ir namespace by mistake.",
48
- op_name,
49
- )
50
- return None
51
-
52
- ir_op = IrOp.registry[op_name]
53
- return ir_op
54
-
55
-
25
class VllmIRLoweringPass(VllmInductorPass):
26
"""
27
This pass lowers vLLM IR ops to their implementations the priority list.
@@ -76,7 +45,7 @@ class VllmIRLoweringPass(VllmInductorPass):
45
46
assert len(match.nodes) == 1, "Expected single node match"
47
node = match.nodes[0]
79
- ir_op = get_ir_op(node)
48
+ ir_op = get_ir_op(node) # TODO is node.target always an overload?
49
assert ir_op is not None, "Expected vLLM IR op"
50
assert not node.kwargs # I think there should never be kwargs here
51
@@ -86,13 +55,18 @@ class VllmIRLoweringPass(VllmInductorPass):
55
self.selected_impls[ir_op.name][node.name] = ir_op_impl.provider
56
57
# replace_by_example wants node args, not the fake tensors
58
+ # use func_impl_fn to properly handle in-place implementations
59
# TODO(luka): Use aot_export_module to get functionalized graph
60
# TODO(luka): Cache the fx_replacement to avoid re-tracing the same impl
61
62
# Defaults not present on node.args but required for replacement tracing
63
bound_args = ir_op._py_signature.bind(*node.args)
64
bound_args.apply_defaults()
95
- match.replace_by_example(ir_op_impl.impl_fn, bound_args.args)
65
+ # It is not safe to run functional passes (like DCE) on the replacements
66
+ # as they might not be functional.
67
+ match.replace_by_example(
68
+ ir_op_impl.func_impl_fn, bound_args.args, run_functional_passes=False
69
+ )
70
71
@VllmInductorPass.time_and_log
72
def __call__(self, graph: fx.Graph) -> None:
@@ -136,7 +110,7 @@ class VllmIRLoweringPass(VllmInductorPass):
110
111
if failed_nodes or failed_ops:
112
logger.warning("Failed to lower vLLM IR ops: %s", ",".join(failed_ops))
139
- logger.warning("Full node list: %s", failed_nodes)
113
+ logger.warning("Full node list: %s", ",".join(str(n) for n in failed_nodes))
114
115
def uuid(self) -> str:
116
"""
vllm/compilation/passes/ir/utils.py
new
+40
@@ -0,0 +1,40 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from torch import fx
4
+from torch._ops import OpOverload, OpOverloadPacket
5
+
6
+from vllm.ir.op import IrOp
7
+from vllm.logger import init_logger
8
+
9
+logger = init_logger(__name__)
10
+
11
+
12
+def overload_or_default(op: OpOverload | OpOverloadPacket) -> OpOverload:
13
+ if isinstance(op, OpOverloadPacket):
14
+ return op.default
15
+ assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket"
16
+ return op
17
+
18
+
19
+def get_ir_op(node: fx.Node) -> IrOp | None:
20
+ if node.op != "call_function":
21
+ return None
22
+
23
+ if not isinstance(node.target, (OpOverload, OpOverloadPacket)):
24
+ return None
25
+
26
+ op_overload = overload_or_default(node.target)
27
+ if op_overload.namespace != "vllm_ir":
28
+ return None
29
+
30
+ op_name = op_overload._opname
31
+ if op_name not in IrOp.registry:
32
+ logger.warning(
33
+ "Unknown vLLM IR op %s, there's likely an issue with torch registration, "
34
+ "or a torch custom op was registered in the vllm_ir namespace by mistake.",
35
+ op_name,
36
+ )
37
+ return None
38
+
39
+ ir_op = IrOp.registry[op_name]
40
+ return ir_op
vllm/compilation/passes/pass_manager.py
+7
-1
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14
from vllm.platforms import current_platform
15
from vllm.utils.system_utils import set_env_var
16
17
+from .ir.clone_elimination import UnsafeCloneEliminationPass
18
from .ir.lowering_pass import VllmIRLoweringPass
19
from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
20
@@ -115,6 +116,8 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
116
# DCE handles mutating ops correctly as well.
117
self.ir_lowering(graph)
118
VllmInductorPass.dump_prefix += 1
119
+ self.clone_elimination(graph)
120
+ VllmInductorPass.dump_prefix += 1
121
122
# clean up after lowering again
123
self.post_cleanup(graph)
@@ -149,11 +152,12 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
152
self.passes += [MiniMaxQKNormPass(config)]
153
154
if self.pass_config.fuse_norm_quant:
152
- self.passes += [RMSNormQuantFusionPass(config)]
155
if rocm_aiter_ops.is_enabled():
156
self.passes += [
157
RocmAiterRMSNormQuantFusionPass(config),
158
]
159
+ self.passes += [RMSNormQuantFusionPass(config)]
160
+
161
if self.pass_config.fuse_act_quant:
162
self.passes += [ActivationQuantFusionPass(config)]
163
if rocm_aiter_ops.is_enabled():
@@ -179,6 +183,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
183
self.passes += [QKNormRoPEFusionPass(config)]
184
185
self.ir_lowering = VllmIRLoweringPass(config)
186
+ self.clone_elimination = UnsafeCloneEliminationPass(config)
187
self.post_cleanup = PostCleanupPass(config)
188
self.fix_functionalization = FixFunctionalizationPass(config)
189
@@ -200,6 +205,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
205
206
passes.append(self.post_cleanup.uuid())
207
passes.append(self.ir_lowering.uuid())
208
+ passes.append(self.clone_elimination.uuid())
209
passes.append(self.post_cleanup.uuid())
210
passes.append(self.fix_functionalization.uuid())
211
vllm/config/kernel.py
+3
@@ -31,6 +31,9 @@ class IrOpPriorityConfig:
31
rms_norm: list[str] = Field(default_factory=list)
32
"""Priority list for vllm.ir.ops.rms_norm"""
33
34
+ fused_add_rms_norm: list[str] = Field(default_factory=list)
35
+ """Priority list for vllm.ir.ops.fused_add_rms_norm"""
36
+
37
def compute_hash(self) -> str:
38
"""
39
Produces a hash unique to the pass configuration.
vllm/config/vllm.py
+2
-5
@@ -125,9 +125,7 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
125
from vllm._aiter_ops import rocm_aiter_ops
126
127
return (
128
- rocm_aiter_ops.is_enabled()
129
- and rocm_aiter_ops.is_rmsnorm_enabled()
130
- and cfg.parallel_config.tensor_parallel_size > 1
128
+ rocm_aiter_ops.is_enabled() and cfg.parallel_config.tensor_parallel_size > 1
129
)
130
131
return (
@@ -159,10 +157,9 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:
157
158
def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
159
"""Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss."""
162
- from vllm._aiter_ops import rocm_aiter_ops
160
161
return (
165
- rocm_aiter_ops.is_rmsnorm_enabled()
162
+ cfg.kernel_config.ir_op_priority.fused_add_rms_norm[0] == "aiter"
163
and cfg.model_config is not None
164
and cfg.model_config.get_hidden_size() == 2880
165
)
vllm/envs.py
+3
-2
@@ -143,7 +143,7 @@ if TYPE_CHECKING:
143
VLLM_DP_RANK_LOCAL: int = -1
144
VLLM_DP_SIZE: int = 1
145
VLLM_USE_STANDALONE_COMPILE: bool = True
146
- VLLM_ENABLE_PREGRAD_PASSES: bool = False
146
+ VLLM_ENABLE_PREGRAD_PASSES: bool = True
147
VLLM_DP_MASTER_IP: str = ""
148
VLLM_DP_MASTER_PORT: int = 0
149
VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
@@ -620,9 +620,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
620
# The pre-grad passes get run even on cache-hit and negatively impact
621
# vllm cold compile times by O(1s)
622
# Can remove this after the following issue gets fixed
623
+ # TODO(luka): maybe_inplace requires this
624
# https://github.com/pytorch/pytorch/issues/174502
625
"VLLM_ENABLE_PREGRAD_PASSES": lambda: (
625
- os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "0") == "1"
626
+ os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "1") == "1"
627
),
628
# Debug pattern matching inside custom passes.
629
# Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').
vllm/ir/op.py
+147
-9
@@ -4,7 +4,7 @@ import contextlib
4
import inspect
5
from collections.abc import Callable
6
from pathlib import Path
7
-from typing import Any, ClassVar, overload
7
+from typing import Any, ClassVar, Literal, overload
8
9
import torch
10
from torch.library import Library, infer_schema
@@ -46,35 +46,51 @@ def enable_torch_wrap(enable: bool = True):
46
_ENABLE_TORCH_WRAP = old
47
48
49
-# 0-param decorator overload
49
+# 0-param decorator overload (no inplace)
50
@overload
51
def register_op(f: Callable[..., Any]) -> "IrOp": ...
52
53
54
-# parametrized decorator overload
54
+# parametrized decorator with allow_inplace=False (default)
55
@overload
56
def register_op(
57
*,
58
name: str | None = None,
59
+ activations: list[str] | None = None,
60
+ allow_inplace: Literal[False] = False,
61
) -> Callable[[Callable[..., Any]], "IrOp"]: ...
62
63
64
+# parametrized decorator with allow_inplace=True
65
+@overload
66
+def register_op(
67
+ *,
68
+ name: str | None = None,
69
+ activations: list[str] | None = None,
70
+ allow_inplace: Literal[True],
71
+) -> Callable[[Callable[..., Any]], "IrOpInplace"]: ...
72
+
73
+
74
def register_op(
75
f: Callable | None = None,
76
*,
77
name: str | None = None,
78
+ activations: list[str] | None = None,
79
+ allow_inplace: bool = False,
80
) -> "IrOp | Callable[[Callable], IrOp]":
81
"""
82
Register a new vLLM IR op.
83
84
:param f: the native implementation of the op
85
:param name: the name of the op, defaults to the function name
86
+ :param activations: list of activation params, defaults to params starting with 'x'
87
+ :param allow_inplace: add a maybe_inplace overload that allows inplace impls
88
:return: the IrOp object if f is provided, otherwise a decorator
89
90
Example usage:
91
```python
92
@vllm.ir.register_op
77
- def my_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
93
+ def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
94
return x + y
95
96
@@ -85,7 +101,10 @@ def register_op(
101
def decorator(_f: Callable):
102
op_name: str = _f.__name__ if name is None else name
103
assert op_name not in IrOp.registry
88
- op = IrOp(op_name, _f)
104
+ if allow_inplace:
105
+ op: IrOp = IrOpInplace(op_name, _f, activations)
106
+ else:
107
+ op = IrOp(op_name, _f, activations)
108
IrOp.registry[op_name] = op
109
return op
110
@@ -100,8 +119,14 @@ class IrOp:
119
120
name: str
121
impls: dict[str, "IrOpImpl"]
122
+ allow_inplace: bool = False
123
104
- def __init__(self, name: str, native_impl: Callable):
124
+ def __init__(
125
+ self,
126
+ name: str,
127
+ native_impl: Callable,
128
+ activations: list[str] | None = None,
129
+ ):
130
self._py_signature = inspect.signature(native_impl)
131
if any(
132
p.kind == inspect.Parameter.KEYWORD_ONLY
@@ -112,8 +137,22 @@ class IrOp:
137
f"supported. That's because kwargs are not allowed during lowering."
138
)
139
140
+ # By convention, we consider parameters starting with 'x' as activations.
141
+ if activations is None:
142
+ activations = [
143
+ p.name
144
+ for p in self._py_signature.parameters.values()
145
+ if p.name.startswith("x")
146
+ ]
147
+
148
self.name = name
149
self.impls: dict[str, IrOpImpl] = {}
150
+ self.activations = activations
151
+ self.activation_indices = [
152
+ i
153
+ for i, p in enumerate(self._py_signature.parameters.values())
154
+ if p.name in activations
155
+ ]
156
self._priority_impls: list[IrOpImpl] = []
157
self._schema_str = infer_schema(native_impl, mutates_args=[])
158
self._input_generator: InputGenerator | None = None
@@ -121,7 +160,12 @@ class IrOp:
160
161
# native implementation
162
self.impls["native"] = IrOpImpl(
124
- self, "native", native_impl, supported=True, supports_args=None
163
+ self,
164
+ "native",
165
+ native_impl,
166
+ # always supported
167
+ supported=True,
168
+ supports_args=None,
169
)
170
171
# By default, fake routes directly to native,
@@ -161,12 +205,14 @@ class IrOp:
205
*,
206
supported: bool = True,
207
supports_args: Callable[..., bool] | None = None,
208
+ inplace: bool = False,
209
):
210
"""
211
Register an implementation for this custom op.
212
:param provider: The name of the provider, must be unique.
213
:param supported: Static support check, use this to check platform support.
214
:param supports_args: Dynamic arg support check, used for types and shapes.
215
+ :param inplace: Does this op reuse activation input memory for outputs
216
:return: A decorator that registers the implementation.
217
218
The decorated function must have the same semantics and signature as
@@ -193,7 +239,7 @@ class IrOp:
239
)
240
241
def _register_impl(f: Callable):
196
- impl = IrOpImpl(self, provider, f, supported, supports_args)
242
+ impl = IrOpImpl(self, provider, f, supported, supports_args, inplace)
243
self.impls[provider] = impl
244
245
if self.get_priority():
@@ -213,7 +259,10 @@ class IrOp:
259
__call__ routes straight here instead of going through torch op dispatching.
260
"""
261
impl = self.dispatch(*args, **kwargs)
216
- return impl.impl_fn(*args, **kwargs)
262
+
263
+ # Default overload must be functional,
264
+ # use func_impl_fn to correctly handle inplace impls.
265
+ return impl.func_impl_fn(*args, **kwargs)
266
267
def apply_arg_defaults(self, args) -> tuple:
268
"""
@@ -314,6 +363,11 @@ class IrOp:
363
old_priority_impls = self._priority_impls
364
try:
365
self._priority_impls = filter_priority_impls(priority)
366
+ logger.debug(
367
+ "Priority for vllm.ir.%s set to %s",
368
+ self.name,
369
+ lazy(lambda: [p.provider for p in self._priority_impls]),
370
+ )
371
yield
372
finally:
373
self._priority_impls = old_priority_impls
@@ -354,6 +408,66 @@ class IrOp:
408
)
409
410
411
+class IrOpInplace(IrOp):
412
+ """IR op with inplace support via maybe_inplace."""
413
+
414
+ maybe_inplace: "IrOpInplaceOverload"
415
+ allow_inplace: bool = True
416
+
417
+ def __init__(
418
+ self,
419
+ name: str,
420
+ native_impl: Callable,
421
+ activations: list[str] | None = None,
422
+ ):
423
+ super().__init__(name, native_impl, activations)
424
+
425
+ # Create the inplace overload
426
+ self.maybe_inplace = IrOpInplaceOverload(self)
427
+
428
+
429
+class IrOpInplaceOverload:
430
+ def __init__(self, op: IrOp):
431
+ params, returns = op._schema_str.split(" -> ")
432
+ n_outputs = returns.count("Tensor")
433
+
434
+ assert returns.count("Tensor") == len(op.activations), (
435
+ "Inplace overload requires the same number of outputs as activations."
436
+ )
437
+
438
+ assert returns.count(",") == n_outputs - 1, (
439
+ "Inplace overload only supports Tensor outputs for now."
440
+ )
441
+
442
+ self.op = op
443
+ self.name = f"{op.name}.maybe_inplace"
444
+ self._schema_str = infer_schema(
445
+ op.impls["native"].impl_fn, mutates_args=op.activations
446
+ )
447
+
448
+ # torch registration
449
+ vllm_ir_lib.define(self.name + self._schema_str)
450
+ vllm_ir_lib.impl(
451
+ self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd"
452
+ )
453
+ # fake goes to default overload for now
454
+ vllm_ir_lib._register_fake(self.name, self.op._fake_call)
455
+
456
+ assert hasattr(getattr(torch.ops.vllm_ir, self.op.name), "maybe_inplace")
457
+ self.torch_op = getattr(torch.ops.vllm_ir, self.op.name).maybe_inplace
458
+
459
+ def __call__(self, *args, **kwargs) -> Any:
460
+ if not _ENABLE_TORCH_WRAP:
461
+ return self._inner_call(*args, **kwargs)
462
+
463
+ return self.torch_op(*args, **kwargs)
464
+
465
+ def _inner_call(self, *args, **kwargs) -> Any:
466
+ # Calling the maybe_inplace overload means we can use inplace impls directly.
467
+ impl = self.op.dispatch(*args, **kwargs)
468
+ return impl.impl_fn(*args, **kwargs)
469
+
470
+
471
class IrOpImpl:
472
def __init__(
473
self,
@@ -362,6 +476,7 @@ class IrOpImpl:
476
impl_fn: Callable,
477
supported: bool,
478
supports_args: Callable[..., bool] | None,
479
+ inplace: bool = False,
480
):
481
assert provider not in op.impls, (
482
f"Implementation for provider {provider} already registered."
@@ -420,11 +535,18 @@ class IrOpImpl:
535
f"native default {op_p.default}'"
536
)
537
538
+ if inplace:
539
+ assert op.allow_inplace, (
540
+ f"Inplace implementation cannot be registered for op {op.name}"
541
+ f" that does not allow inplace."
542
+ )
543
+
544
self.op = op
545
self.provider = provider
546
self.impl_fn = impl_fn
547
self.supported = supported
548
self._supports_args = supports_args
549
+ self.inplace = inplace
550
551
@property
552
def supports_all_args(self) -> bool:
@@ -449,3 +571,19 @@ class IrOpImpl:
571
"""
572
sources = [Path(inspect.getfile(self.impl_fn))]
573
return hash_source(*sources)
574
+
575
+ def func_impl_fn(self, *args, **kwargs) -> Any:
576
+ """
577
+ Copy any inputs in activations if this is an inplace impl,
578
+ to ensure functional semantics.
579
+ """
580
+ if not self.inplace:
581
+ return self.impl_fn(*args, **kwargs)
582
+
583
+ # copy activations to ensure functional semantics
584
+ new_args = list(args)
585
+ for i in self.op.activation_indices:
586
+ assert isinstance(args[i], torch.Tensor)
587
+ new_args[i] = args[i].clone()
588
+
589
+ return self.impl_fn(*new_args, **kwargs)
vllm/ir/ops/__init__.py
+2
-2
@@ -1,5 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from .layernorm import rms_norm
3
+from .layernorm import fused_add_rms_norm, rms_norm
4
5
-__all__ = ["rms_norm"]
5
+__all__ = ["rms_norm", "fused_add_rms_norm"]
vllm/ir/ops/layernorm.py
+37
-1
@@ -27,10 +27,46 @@ def _rms_norm_input_generator(
27
) -> tuple:
28
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
29
weight = torch.randn(hidden_size, dtype=dtype)
30
- return (x, weight, epsilon)
30
+ return x, weight, epsilon
31
32
33
# Reductions in rms_norm accumulate rounding error at large shapes
34
# (e.g. 32768x16384), causing a few elements out of millions to exceed
35
# the default float16 tolerance.
36
rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
37
+
38
+
39
+@register_op(allow_inplace=True)
40
+def fused_add_rms_norm(
41
+ x: Tensor,
42
+ x_residual: Tensor,
43
+ weight: Tensor | None,
44
+ epsilon: float,
45
+ variance_size: int | None = None,
46
+) -> tuple[Tensor, Tensor]:
47
+ """Fused add and weighted root-mean-square layer normalization"""
48
+ orig_dtype = x.dtype
49
+ x = x.to(torch.float32)
50
+ x = x + x_residual.to(torch.float32)
51
+ x_residual = x.to(orig_dtype)
52
+
53
+ x_var = x if variance_size is None else x[..., :variance_size]
54
+ variance = x_var.pow(2).mean(dim=-1, keepdim=True)
55
+ x = x * torch.rsqrt(variance + epsilon)
56
+ if weight is not None:
57
+ x = x.to(weight.dtype) * weight
58
+ return x.to(orig_dtype), x_residual
59
+
60
+
61
+# fused_add_rms_norm has similar rounding error accumulation as rms_norm
62
+fused_add_rms_norm.override_tolerance(torch.float16, atol=1e-2, rtol=2e-3)
63
+
64
+
65
+@fused_add_rms_norm.register_input_generator
66
+def _fused_add_rms_norm_input_generator(
67
+ num_tokens: int, hidden_size: int, dtype: torch.dtype, epsilon: float = 1e-5
68
+) -> tuple:
69
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype)
70
+ x_residual = torch.randn(num_tokens, hidden_size, dtype=dtype)
71
+ weight = torch.randn(hidden_size, dtype=dtype)
72
+ return x, x_residual, weight, epsilon
vllm/kernels/aiter_ops.py
+69
@@ -75,3 +75,72 @@ def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor
75
direct_register_aiter_op(
76
op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake
77
)
78
+
79
+rms_add_no_var_16bit_only = (
80
+ lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
81
+ and x.dtype in (torch.float16, torch.bfloat16)
82
+ and (weight is None or weight.dtype == x.dtype)
83
+)
84
+"""
85
+AITER fused_add_rms_norm only supports 16-bit activations and no var_size override.
86
+Requires weight dtype to match x dtype.
87
+"""
88
+
89
+
90
+@ir.ops.fused_add_rms_norm.register_impl(
91
+ "aiter", supports_args=rms_add_no_var_16bit_only, supported=AITER_SUPPORTED
92
+)
93
+def fused_add_rms_norm(
94
+ x: Tensor,
95
+ x_residual: Tensor,
96
+ weight: Tensor | None,
97
+ epsilon: float,
98
+ variance_size: int | None = None,
99
+) -> tuple[Tensor, Tensor]:
100
+ assert variance_size is None
101
+ assert x.dtype in (torch.float16, torch.bfloat16)
102
+ if weight is None:
103
+ weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
104
+ return torch.ops.vllm_aiter.fused_add_rms_norm(x, x_residual, weight, epsilon)
105
+
106
+
107
+def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
108
+ x: torch.Tensor,
109
+ residual: torch.Tensor,
110
+ weight: torch.Tensor,
111
+ variance_epsilon: float,
112
+) -> tuple[torch.Tensor, torch.Tensor]:
113
+ from aiter import rmsnorm2d_fwd_with_add
114
+
115
+ # TODO can out = x and residual_out = residual to save memory?
116
+ # Need to check if the kernel supports in-place residual output
117
+ # (if yes set mutates_args and inplace)
118
+ residual_out = torch.empty_like(residual)
119
+ out = torch.empty_like(x)
120
+ rmsnorm2d_fwd_with_add(
121
+ out, # output
122
+ x, # input
123
+ residual, # residual input
124
+ residual_out, # residual output
125
+ weight,
126
+ variance_epsilon,
127
+ )
128
+ return out, residual_out
129
+
130
+
131
+def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
132
+ x: torch.Tensor,
133
+ residual: torch.Tensor,
134
+ weight: torch.Tensor,
135
+ variance_epsilon: float,
136
+) -> tuple[torch.Tensor, torch.Tensor]:
137
+ residual_out = torch.empty_like(residual)
138
+ out = torch.empty_like(x)
139
+ return out, residual_out
140
+
141
+
142
+direct_register_aiter_op(
143
+ op_name="fused_add_rms_norm",
144
+ op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
145
+ fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
146
+)
vllm/kernels/oink_ops.py
+55
-5
@@ -1,6 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""This file registers Oink implementations for vLLM IR ops.
4
+
5
+vLLM does not depend on the external Oink repository/package. When an external
6
+plugin registers torch.library.custom_op entrypoints under the `oink::`
7
+namespace (e.g. via vLLM's general_plugins mechanism), these ops will be marked
8
+ as supported. To dispatch to those ops, set kernel_config.ir_op_priority.<op> to oink.
9
+Alternatively, `VLLM_USE_OINK_OPS=1` will add this to priority by default.
10
+"""
11
+
12
import torch
13
+from torch import Tensor
14
15
from vllm import ir
16
from vllm.platforms import current_platform
@@ -15,7 +25,7 @@ def has_oink_op(name: str) -> bool:
25
return OINK_AVAILABLE and hasattr(torch.ops.oink, name)
26
27
18
-def _can_view_as_2d(x: torch.Tensor) -> bool:
28
+def _can_view_as_2d(x: Tensor) -> bool:
29
"""Return True if x.view(-1, x.shape[-1]) is viewable (no copy)."""
30
if x.dim() < 2:
31
return False
@@ -32,7 +42,7 @@ def _can_view_as_2d(x: torch.Tensor) -> bool:
42
return True
43
44
35
-def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
45
+def _is_oink_stride_compatible_2d(x_2d: Tensor) -> bool:
46
"""Return True if x_2d meets Oink's pointer-path stride constraints."""
47
if x_2d.dim() != 2:
48
return False
@@ -67,11 +77,51 @@ and no variance_size override.
77
"oink", supports_args=oink_rms_supported, supported=has_oink_op("rmsnorm")
78
)
79
def rms_norm(
70
- x: torch.Tensor,
71
- weight: torch.Tensor | None,
80
+ x: Tensor,
81
+ weight: Tensor | None,
82
epsilon: float,
83
variance_size: int | None = None,
74
-) -> torch.Tensor:
84
+) -> Tensor:
85
assert variance_size is None
86
x_2d = x.view(-1, x.shape[-1])
87
return torch.ops.oink.rmsnorm(x_2d, weight, epsilon).view_as(x)
88
+
89
+
90
+oink_add_rms_supported = (
91
+ lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
92
+ and weight is not None
93
+ and x.dim() >= 2
94
+ and x.dtype == weight.dtype
95
+ and weight.is_contiguous()
96
+ and _can_view_as_2d(x)
97
+ and _is_oink_stride_compatible_2d(x.view(-1, x.shape[-1]))
98
+ # residual must have 2d-compatible strides and match x shape/dtype
99
+ and x.dtype == x_residual.dtype
100
+ and x.shape == x_residual.shape
101
+ and _can_view_as_2d(x_residual)
102
+ and _is_oink_stride_compatible_2d(x_residual.view(-1, x_residual.shape[-1]))
103
+)
104
+"""
105
+Oink fused_add_rms_norm has the same constraints as rms_norm,
106
+and residual must be 2d-like with compatible strides.
107
+"""
108
+
109
+
110
+@ir.ops.fused_add_rms_norm.register_impl(
111
+ "oink",
112
+ supports_args=oink_add_rms_supported,
113
+ supported=has_oink_op("fused_add_rms_norm"),
114
+ inplace=True,
115
+)
116
+def fused_add_rms_norm(
117
+ x: Tensor,
118
+ x_residual: Tensor,
119
+ weight: Tensor | None,
120
+ epsilon: float,
121
+ variance_size: int | None = None,
122
+) -> tuple[Tensor, Tensor]:
123
+ assert variance_size is None
124
+ x_2d = x.view(-1, x.shape[-1])
125
+ residual_2d = x_residual.view(-1, x_residual.shape[-1])
126
+ torch.ops.oink.fused_add_rms_norm(x_2d, residual_2d, weight, epsilon)
127
+ return x, x_residual
vllm/kernels/vllm_c.py
+30
@@ -31,3 +31,33 @@ def rms_norm(
31
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
32
torch.ops._C.rms_norm(output, x, weight, epsilon)
33
return output
34
+
35
+
36
+rms_add_no_var_size = (
37
+ lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
38
+ and (weight is None or weight.dtype == x.dtype)
39
+)
40
+"""vLLM Kernel does not support variance_size parameter and requires
41
+matching input/weight dtype."""
42
+
43
+
44
+@ir.ops.fused_add_rms_norm.register_impl(
45
+ "vllm_c",
46
+ supports_args=rms_add_no_var_size,
47
+ supported=CUDA_ALIKE,
48
+ inplace=True,
49
+)
50
+def fused_add_rms_norm(
51
+ x: Tensor,
52
+ x_residual: Tensor,
53
+ weight: Tensor | None,
54
+ epsilon: float,
55
+ variance_size: int | None = None,
56
+) -> tuple[Tensor, Tensor]:
57
+ if weight is None:
58
+ # Kernel requires weight tensor, pass ones
59
+ weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
60
+
61
+ assert variance_size is None
62
+ torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
63
+ return x, x_residual
vllm/kernels/xpu_ops.py
+28
@@ -36,3 +36,31 @@ def rms_norm(
36
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
37
torch.ops._C.rms_norm(output, x, weight, epsilon)
38
return output
39
+
40
+
41
+rms_add_no_var_size = (
42
+ lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
43
+ and (weight is None or weight.dtype == x.dtype)
44
+)
45
+
46
+
47
+@ir.ops.fused_add_rms_norm.register_impl(
48
+ "xpu_kernels",
49
+ supports_args=rms_add_no_var_size,
50
+ supported=XPU_KERNELS_SUPPORTED,
51
+ inplace=True,
52
+)
53
+def fused_add_rms_norm(
54
+ x: Tensor,
55
+ x_residual: Tensor,
56
+ weight: Tensor | None,
57
+ epsilon: float,
58
+ variance_size: int | None = None,
59
+) -> tuple[Tensor, Tensor]:
60
+ if weight is None:
61
+ # Kernel requires weight tensor, pass ones
62
+ weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
63
+
64
+ assert variance_size is None
65
+ torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
66
+ return x, x_residual
vllm/model_executor/layers/layernorm.py
+29
-252
@@ -8,68 +8,15 @@ import torch.nn.functional as F
8
9
# Import kernels
10
import vllm.kernels # noqa: F401
11
-from vllm import _oink_ops, envs, ir
12
-from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm import envs, ir
12
+from vllm.config import get_current_vllm_config
13
from vllm.logger import init_logger
14
from vllm.model_executor.custom_op import CustomOp
15
-from vllm.model_executor.layers.batch_invariant import (
16
- rms_norm_batch_invariant,
17
-)
18
-from vllm.platforms import current_platform
15
+from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant
16
17
logger = init_logger(__name__)
18
19
23
-def _can_view_as_2d(x: torch.Tensor) -> bool:
24
- """Return True if x.view(-1, x.shape[-1]) is viewable (no copy)."""
25
- if x.dim() < 2:
26
- return False
27
- if x.dim() == 2:
28
- return True
29
- # For a view(-1, N) to be valid, all leading dims must be contiguous with
30
- # respect to each other (size-1 dims are ignored).
31
- for dim in range(x.dim() - 1):
32
- # Strides for size-1 dims are irrelevant and can be arbitrary.
33
- if x.size(dim + 1) != 1 and x.stride(dim) != x.stride(dim + 1) * x.size(
34
- dim + 1
35
- ):
36
- return False
37
- return True
38
-
39
-
40
-def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
41
- """Return True if x_2d meets Oink's pointer-path stride constraints."""
42
- if x_2d.dim() != 2:
43
- return False
44
- if x_2d.stride(1) != 1:
45
- return False
46
- # Match Oink's vectorization constraint: stride(0) divisible by 256b.
47
- if x_2d.dtype in (torch.float16, torch.bfloat16):
48
- divby = 16
49
- elif x_2d.dtype == torch.float32:
50
- divby = 8
51
- else:
52
- return False
53
- return (x_2d.stride(0) % divby) == 0
54
-
55
-
56
-def fused_add_rms_norm(
57
- x: torch.Tensor,
58
- residual: torch.Tensor,
59
- weight: torch.Tensor,
60
- variance_epsilon: float,
61
-) -> tuple[torch.Tensor, torch.Tensor]:
62
- from vllm import _custom_ops as ops
63
-
64
- ops.fused_add_rms_norm(
65
- x,
66
- residual,
67
- weight,
68
- variance_epsilon,
69
- )
70
- return x, residual
71
-
72
-
20
def poly_norm(
21
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, variance_epsilon: float
22
) -> torch.Tensor:
@@ -86,18 +33,6 @@ def poly_norm(
33
return out
34
35
89
-def dispatch_rocm_rmsnorm_func(dtype: torch.dtype, use_aiter: bool = False):
90
- use_aiter = use_aiter and dtype in [
91
- torch.float16,
92
- torch.bfloat16,
93
- ]
94
-
95
- if use_aiter:
96
- return rocm_aiter_ops.rms_norm2d_with_add
97
- else:
98
- return fused_add_rms_norm
99
-
100
-
36
# --8<-- [start:rms_norm]
37
@CustomOp.register("rms_norm")
38
class RMSNorm(CustomOp):
@@ -130,105 +65,19 @@ class RMSNorm(CustomOp):
65
if self.has_weight:
66
self.weight = nn.Parameter(self.weight)
67
133
- if current_platform.is_rocm():
134
- aiter_rmsnorm_enabled = rocm_aiter_ops.is_rmsnorm_enabled()
135
- self.rocm_norm_func_with_add = dispatch_rocm_rmsnorm_func(
136
- dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled
137
- )
138
-
139
- # Optional: enable Oink Blackwell RMSNorm custom-op fast path on
140
- # compatible CUDA devices (e.g., SM100) when the external Oink
141
- # package is available. This is detected once at construction time
142
- # to avoid per-call device queries in the hot path.
143
- self._use_oink_fused_add_rmsnorm = False
144
- if (
145
- not current_platform.is_rocm()
146
- and torch.cuda.is_available()
147
- and bool(getattr(envs, "VLLM_USE_OINK_OPS", False))
148
- ):
149
- # NOTE: vLLM disables custom ops by default when using Inductor.
150
- # If this op is disabled, CustomOp will dispatch to forward_native,
151
- # and the Oink path in forward_cuda will never run.
152
- if getattr(self._forward_method, "__func__", None) is getattr(
153
- self.forward_native, "__func__", None
154
- ):
155
- try:
156
- from vllm.config import get_cached_compilation_config
157
-
158
- custom_ops = get_cached_compilation_config().custom_ops
159
- except Exception:
160
- custom_ops = ["<unknown>"]
161
- logger.warning_once(
162
- "VLLM_USE_OINK_OPS=1 but the `rms_norm` custom op is "
163
- "disabled (CompilationConfig.custom_ops=%s). Enable it via "
164
- "`compilation_config={'custom_ops': ['none', '+rms_norm']}` "
165
- "(or `['all']`) to let vLLM call into torch.ops.oink.*.",
166
- custom_ops,
167
- )
168
- # Custom op disabled => forward_cuda won't run. Avoid doing any
169
- # external Oink initialization work in this case.
170
- else:
171
- try:
172
- device_index = torch.accelerator.current_device_index()
173
- if _oink_ops.is_oink_available_for_device(device_index):
174
- self._use_oink_fused_add_rmsnorm = (
175
- _oink_ops.has_fused_add_rms_norm()
176
- )
177
- except Exception as e:
178
- # If anything goes wrong (no Oink install, CPU-only env, etc.),
179
- # silently fall back to the built-in RMSNorm path.
180
- logger.warning_once(
181
- "VLLM_USE_OINK_OPS=1 but failed to initialize Oink "
182
- "RMSNorm; falling back to vLLM RMSNorm. Error: %s",
183
- e,
184
- )
185
- self._use_oink_fused_add_rmsnorm = False
186
-
187
- @staticmethod
188
- def forward_static(
189
- x: torch.Tensor,
190
- variance_epsilon: float,
191
- hidden_size: int,
192
- orig_dtype: torch.dtype,
193
- weight: torch.Tensor | None = None,
194
- residual: torch.Tensor | None = None,
195
- variance_size_override: int | None = None,
196
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
197
- """PyTorch-native implementation equivalent to forward()."""
198
- x = x.to(torch.float32)
199
- if residual is not None:
200
- # residual promoted f16->f32 automatically,
201
- # otherwise Inductor eliminates the casts to and from f16,
202
- # increasing memory usage (and complicating pattern matching)
203
- x = x + residual
204
- residual = x.to(orig_dtype)
205
-
206
- if x.shape[-1] != hidden_size:
207
- raise ValueError(
208
- f"Expected hidden_size to be {hidden_size}, but found: {x.shape[-1]}"
209
- )
210
-
211
- if variance_size_override is None:
212
- x_var = x
213
- else:
214
- if hidden_size < variance_size_override:
215
- raise ValueError(
216
- "Expected hidden_size to be at least "
217
- f"{variance_size_override}, but found: {hidden_size}"
218
- )
219
-
220
- x_var = x[:, :, :variance_size_override]
221
-
222
- variance = x_var.pow(2).mean(dim=-1, keepdim=True)
223
-
224
- x = x * torch.rsqrt(variance + variance_epsilon)
225
- x = x.to(orig_dtype)
226
- if weight is not None:
227
- x = x * weight
228
- if residual is None:
229
- return x
230
- else:
231
- return x, residual
68
+ # Do not pass identity weight to native implementation (causes issue on TPU).
69
+ # Other implementations require weight to be passed even if all ones.
70
+ # Cheat and predict if native will be dispatched to:
71
+ # 1) if native is first in priority list
72
+ # 2) if variance_size_override is given (only supported by native impl)
73
+ # TODO(luka): address weight passing inconsistency:
74
+ # https://github.com/vllm-project/vllm/issues/39370
75
+ priority = get_current_vllm_config().kernel_config.ir_op_priority
76
+ var_override = self.variance_size_override is not None
77
+ native_rms_norm = priority.rms_norm[0] == "native" or var_override
78
+ native_add_rms_norm = priority.fused_add_rms_norm[0] == "native" or var_override
79
+ self.pass_weight = self.has_weight or not native_rms_norm
80
+ self.pass_weight_add = self.has_weight or not native_add_rms_norm
81
82
def forward_native(
83
self,
@@ -237,106 +86,34 @@ class RMSNorm(CustomOp):
86
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
87
"""PyTorch-native implementation equivalent to forward()."""
88
if residual is None:
240
- # TODO(luka): address the weight=None passing issue more generally
89
return ir.ops.rms_norm(
90
x,
243
- self.weight.data if self.has_weight else None,
91
+ self.weight.data if self.pass_weight else None,
92
+ self.variance_epsilon,
93
+ self.variance_size_override,
94
+ )
95
+ else:
96
+ return ir.ops.fused_add_rms_norm.maybe_inplace(
97
+ x,
98
+ residual,
99
+ self.weight.data if self.pass_weight_add else None,
100
self.variance_epsilon,
101
self.variance_size_override,
102
)
247
-
248
- return self.forward_static(
249
- x,
250
- self.variance_epsilon,
251
- self.hidden_size,
252
- x.dtype,
253
- self.weight.data if self.has_weight else None,
254
- residual,
255
- self.variance_size_override,
256
- )
103
104
def forward_cuda(
105
self,
106
x: torch.Tensor,
107
residual: torch.Tensor | None = None,
108
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
263
- if residual is None and not envs.VLLM_BATCH_INVARIANT:
264
- return ir.ops.rms_norm(
265
- x, self.weight.data, self.variance_epsilon, self.variance_size_override
266
- )
267
-
268
- if self.variance_size_override is not None:
269
- return self.forward_native(x, residual)
270
-
271
- # Optional Oink SM100 fast path (fused residual-add + RMSNorm, in-place).
272
- # This mirrors vLLM's fused_add_rms_norm semantics by mutating both
273
- # `x` (normalized output) and `residual` (residual-out buffer).
109
if (
275
- residual is not None
276
- and getattr(self, "_use_oink_fused_add_rmsnorm", False)
277
- and x.is_cuda
278
- and residual.is_cuda
279
- and x.shape == residual.shape
280
- and x.dtype == residual.dtype
281
- and x.dim() >= 2
282
- and self.has_weight
283
- and not envs.VLLM_BATCH_INVARIANT
284
- and self.weight.data.dtype == x.dtype
285
- and self.weight.data.is_contiguous()
110
+ envs.VLLM_BATCH_INVARIANT
111
+ and residual is None
112
+ and self.variance_size_override is None
113
):
287
- orig_shape = x.shape
288
- hidden_size = orig_shape[-1]
289
- if _can_view_as_2d(x) and _can_view_as_2d(residual):
290
- x_2d = x.view(-1, hidden_size)
291
- res_2d = residual.view(-1, hidden_size)
292
-
293
- # The Oink in-place pointer path supports the common vLLM
294
- # layout where:
295
- # - `x` may be strided/padded row-major (stride(1) == 1), and
296
- # - `residual` is contiguous row-major ([M, N] with stride(0) == N).
297
- # If these conditions are not met, fall back to vLLM's built-in
298
- # fused kernel.
299
- if (
300
- _is_oink_stride_compatible_2d(x_2d)
301
- and _is_oink_stride_compatible_2d(res_2d)
302
- and res_2d.is_contiguous()
303
- ):
304
- _oink_ops.fused_add_rms_norm_(
305
- x_2d,
306
- res_2d,
307
- self.weight.data,
308
- self.variance_epsilon,
309
- )
310
- return x, residual
311
-
312
- if residual is not None:
313
- return fused_add_rms_norm(
314
- x, residual, self.weight.data, self.variance_epsilon
315
- )
316
- else:
317
- assert envs.VLLM_BATCH_INVARIANT
114
return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)
115
320
- def forward_hip(
321
- self,
322
- x: torch.Tensor,
323
- residual: torch.Tensor | None = None,
324
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
325
- if residual is None and not envs.VLLM_BATCH_INVARIANT:
326
- return ir.ops.rms_norm(
327
- x, self.weight.data, self.variance_epsilon, self.variance_size_override
328
- )
329
-
330
- if self.variance_size_override is not None:
331
- return self.forward_native(x, residual)
332
-
333
- if residual is not None:
334
- return self.rocm_norm_func_with_add(
335
- x, residual, self.weight.data, self.variance_epsilon
336
- )
337
- else:
338
- assert envs.VLLM_BATCH_INVARIANT
339
- return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)
116
+ return self.forward_native(x, residual)
117
118
def forward_xpu(
119
self,
vllm/platforms/cuda.py
+3
-1
@@ -578,7 +578,9 @@ class CudaPlatformBase(Platform):
578
if envs.VLLM_USE_OINK_OPS:
579
rms_norm = ["oink"] + default
580
581
- return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
581
+ return IrOpPriorityConfig.with_default(
582
+ default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
583
+ )
584
585
586
# NVML utils
vllm/platforms/rocm.py
+6
-16
@@ -693,21 +693,11 @@ class RocmPlatform(Platform):
693
@classmethod
694
def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
695
from vllm._aiter_ops import rocm_aiter_ops
696
- from vllm.config.compilation import CUDAGraphMode
696
697
compilation_config = vllm_config.compilation_config
699
- is_eager_execution = compilation_config.cudagraph_mode == CUDAGraphMode.NONE
698
use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled()
701
- use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()
699
use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled()
700
use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
704
- # Aiter rms norm perform best when CUDA Graph capture is enabled.
705
- if (
706
- use_aiter_rms_norm
707
- and not is_eager_execution
708
- and "-rms_norm" not in compilation_config.custom_ops
709
- ):
710
- compilation_config.custom_ops.append("+rms_norm")
701
702
if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops:
703
compilation_config.custom_ops.append("+quant_fp8")
@@ -939,7 +929,7 @@ class RocmPlatform(Platform):
929
def get_default_ir_op_priority(
930
cls, vllm_config: "VllmConfig"
931
) -> "IrOpPriorityConfig":
942
- from vllm.config.compilation import CompilationMode
932
+ from vllm.config.compilation import CompilationMode, CUDAGraphMode
933
from vllm.config.kernel import IrOpPriorityConfig
934
935
# Native used by default when compiling,
@@ -949,12 +939,10 @@ class RocmPlatform(Platform):
939
using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
940
default = ["native"] if using_inductor else ["vllm_c", "native"]
941
952
- # This (mostly) preserves previous CustomOp behavior
953
- # Necessary on ROCm because it's common that users
954
- # enable rms_norm to use the aiter kernel.
942
+ # Aiter rms norm perform best when CUDA Graph capture is enabled.
943
# TODO(luka/TJ) remove env vars completely
944
if (
957
- cc.is_custom_op_enabled("rms_norm")
945
+ cc.cudagraph_mode != CUDAGraphMode.NONE
946
and envs.VLLM_ROCM_USE_AITER
947
and envs.VLLM_ROCM_USE_AITER_RMSNORM
948
):
@@ -962,7 +950,9 @@ class RocmPlatform(Platform):
950
else:
951
rms_norm = default
952
965
- return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
953
+ return IrOpPriorityConfig.with_default(
954
+ default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
955
+ )
956
957
@classmethod
958
@with_amdsmi_context