@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 615 lines 21.8 KB
Rendered Raw
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