vLLM IR：功能性中间表示¶

动机¶

vLLM IR 是一种功能性中间表示 (IR)，它弥补了底层 torch 操作与 vLLM 层（如 RMSNorm 和量化算子）之间的空白。通过将算子语义与实现和调度分离，vLLM IR 同时简化了编译以及内核注册 & 调度。它在 torch FX 表示中作为一种方言运行，允许与“常规”torch 操作 & 自定义 torch 操作/内核完全互操作，并支持从以前的 CustomOp 方法逐步迁移。

核心设计原则

即时-编译一致性：在即时模式和编译模式下行为一致（除了细微的数值差异）
简单、透明而强大的内核选择：良好的可见性和控制，便于调试
约定优于配置：注册操作和实现几乎不需要样板代码
可扩展性：操作和实现可以在任何地方注册，无论是在代码树内还是外部
互操作性：与“常规”torch 操作 & 自定义 torch 操作/内核完全兼容，减少开发人员摩擦并允许分阶段迁移

清晰的语义/实现分离实现了统一且可扩展的调度机制，允许每个平台使用多个内核以及强大的内核选择。这种分离还有助于更清晰的测试和基准测试，消除了传统方法中大量样板标准。

通过将内核选择延迟到编译过程的后期，编译器可以在更高级别的表示上操作，这具有以下主要好处：

在融合/转换通道中的模式匹配每个操作只需要一个简单的模式
OOT 编译器后端可以从更高级别的表示进行降级（进行中）
编译器可以对可用实现进行自动调优（未来功能）

快速概览¶

声明一个IR操作¶

IR 操作使用 @register_op 装饰器声明，并附带一个定义操作语义的原生 PyTorch 实现

# vllm/ir/ops/layernorm.py
from torch import Tensor
from vllm.ir import register_op

@register_op
def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
    """Weighted root-mean-square layer normalization"""
    orig_dtype = x.dtype
    x = x.to(torch.float32)
    x_var = x if variance_size is None else x[..., :variance_size]
    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
    x = x * torch.rsqrt(variance + epsilon)
    x = x.to(orig_dtype)
    if weight is not None:
        x = x * weight
    return x

原生实现有三个目的

语义定义：指定操作的确切语义，包括形状和步幅
默认实现：当没有其他（更好）实现可用时使用
测试参考：其他实现必须与这些语义匹配

注册实现¶

内核实现使用 IR 操作对象上的 register_impl 装饰器进行注册

# vllm/kernels/vllm_c.py
from vllm import ir

rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None

@ir.ops.rms_norm.register_impl("vllm_c", supports_args=rms_norm_no_var, supported=current_platform.is_cuda_alike())
def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
    output = torch.empty_like(x)
    torch.ops._C.rms_norm(output, x, weight, epsilon)
    return output

实现可以指定

supported：静态布尔值，指示此实现是否可用
supports_args：函数，检查实现是否支持特定参数
inplace：此实现是否将输入内存重用于输出

在模型中使用IR操作¶

IR 操作在模型代码中直接导入和调用

# vllm/model_executor/layers/layernorm.py
from vllm import ir

class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x: Tensor, residual: Tensor | None = None):
        if residual is None:
            return ir.ops.rms_norm(x, self.weight, self.variance_epsilon)

        # Use maybe_inplace overload to allow implementation to reuse input memory for outputs
        # (using x or residual after this call is undefined behavior)
        return ir.ops.fused_add_rms_norm.maybe_inplace(
            x, residual, self.weight, self.variance_epsilon
        )

配置内核选择¶

内核选择通过配置中的优先级列表进行控制。优先级列表指定了实现被考虑的顺序，将选择第一个支持的实现。这包括静态支持检查 (supported=...) 和动态参数支持检查 (supports_args=...)。

命令行配置¶

使用 --ir-op-priority.<op_name>=<provider1>,<provider2>,...

# CUDA: Use vllm_c implementation for rms_norm
vllm serve meta-llama/Llama-3.2-1B \
  --ir-op-priority.rms_norm=vllm_c

# ROCm: Try aiter first, fall back to vllm_c, then native
vllm serve meta-llama/Llama-3.2-1B \
  --ir-op-priority.rms_norm=aiter,vllm_c,native

# Configure multiple operations
vllm serve meta-llama/Llama-3.2-1B \
  --ir-op-priority.rms_norm=vllm_c \
  --ir-op-priority.fused_add_rms_norm=vllm_c

Python 配置¶

from vllm import LLM
from vllm.config import VllmConfig, KernelConfig

llm = LLM(
    model="meta-llama/Llama-3.2-1B",
    vllm_config=VllmConfig(
        kernel_config=KernelConfig(
            ir_op_priority={
                "rms_norm": ["vllm_c", "native"],
                "fused_add_rms_norm": ["vllm_c", "native"],
            }
        )
    )
)

平台默认值¶

每个平台都提供自动应用的默认优先级列表

# CUDA/XPU/ROCm platform defaults (when compiling with Inductor)
{
  "rms_norm": ["native"],  # Native torch is default
  "fused_add_rms_norm": ["native"],
}

# CUDA platform defaults (eager or Dynamo-only)
{
  "rms_norm": ["vllm_c", "native"],
  "fused_add_rms_norm": ["vllm_c", "native"],
}

# ROCm platform defaults (future - currently same as CUDA)
{
    "rms_norm": ["aiter", "vllm_c", "native"],
    "fused_add_rms_norm": ["aiter", "vllm_c", "native"],
}

# XPU platform defaults (eager or Dynamo-only)
{
    "rms_norm": ["xpu_kernels", "native"],
    "fused_add_rms_norm": ["xpu_kernels", "native"],
}

用户指定的优先级会添加到平台默认优先级的前面，因此您只需指定无序的实现，其他实现将自动追加。

编译管线¶

vLLM IR 大幅定制了基于 torch.compile 的编译过程，以允许自定义编译通道在高层 IR 上操作，同时最终生成高效的低层代码。编译管线包含以下几个阶段：

1. Dynamo 追踪¶

当 torch.compile 追踪模型的正向传播时，vLLM IR 操作在 vllm_ir torch 库中以自定义操作的形式出现。这些操作对 Dynamo 是不透明的，这意味着它们直接出现在 FX 图中，而无需分解。

# Python code (epsilon=1e-5)
x1 = ir.ops.rms_norm(x, weight, epsilon)
x2, residual_out = ir.ops.fused_add_rms_norm.maybe_inplace(x1, residual, weight, epsilon)

# FX graph after Dynamo tracing
x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
out = torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace(x1, residual, weight, 1e-5); x1 = residual = None
x2 = out[0]
residual_out = out[1]

2. AOTAutograd 和功能化¶

AOTAutograd 对图进行功能化，将任何可变操作转换为功能等效操作。对于带有 maybe_inplace 重载的 vLLM IR 操作，我们在 AOTAutograd 之前手动执行此操作，使用预梯度自定义通道钩子将它们转换为功能性的 default 重载。

# After functionalization
x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
out = torch.ops.vllm_ir.fused_add_rms_norm.default(x1, residual, weight, 1e-5); x1 = residual = None
x2 = out[0]
residual_out = out[1]

该通道还跟踪哪些输入被“捐赠”（传递给 maybe_inplace），并将此信息存储在 vLLM 的 PassContext 中，以便稍后在克隆消除中使用。

3. IR 融合与转换通道¶

功能化之后，自定义 vLLM 通道对包含高级 IR 操作的功能性 FX 图进行操作。这些通道可以执行融合、为序列并行性分配操作以及其他转换。

# Example: Sequence Parallelism (see SequenceParallelismPass)
# Before SP pass

all_reduce = torch.ops.vllm.all_reduce(x, "tp:0")
rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)

# after SP pass
reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0")
rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
all_gather = torch.ops.vllm.all_gather(x, "tp:0")

融合通道受益于高级表示：它们不需要匹配底层 PyTorch 操作，无需单独处理不同的内核实现，也无需处理自定义内核的功能化。

4. IR 降级¶

降级通道 (VllmIRLoweringPass) 将每个 vLLM IR 操作替换为其选定的实现。该实现根据优先级列表和支持谓词进行选择，使用图中元数据中的伪张量代替操作参数。

# Implementation selection, same in eager dispatch and compile lowering
def dispatch(*args) -> IrOpImpl:
  for provider in priority_list:  # e.g., ["vllm_c", "native"]
    impl = ir_op.impls[provider]
    if not impl.supported:
      continue
    if impl.supports_args and not impl.supports_args(*args):
      continue
    return impl

# make_fx uses torch.fx.symbolic_trace
impl_graph = make_fx(selected_impl.impl_fn)
# Replace IR op node with impl_graph's nodes
match.replace_by_example(selected_impl.impl_fn, node.args)

例如，使用 vllm_c 实现降级 rms_norm

# Before lowering (IR op)
rms_norm = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5)

# After lowering (vllm_c implementation traced)
# Note: Lowering does not currently functionalize, this will likely change in the future.
empty =  torch.ops.aten.empty.memory_format(x.shape, ...)
rms_norm = torch.ops._C.rms_norm(empty, x, weight, 1e-5)

当降级一个会改变输入（inplace=True）的实现时，降级通道会插入克隆以保留功能语义。

# vllm_c implementation for fused_add_rms_norm mutates its first two arguments
# Lowered with clones for safety
clone_default = torch.ops.aten.clone.default(x)
clone_default_1 = torch.ops.aten.clone.default(residual)
fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(clone_default, clone_default_1, weight, 1e-5)

5. 克隆清理¶

降级后，克隆消除通道 (UnsafeCloneEliminationPass) 会移除降级过程中引入的不必要的克隆。当使用带有 maybe_inplace 的原地内核时，此通道对于实现零拷贝行为至关重要。如果满足以下条件，该通道将移除一个克隆：

克隆的输入在图中创建且未在图中再次使用
克隆的输入是图参数，并被标记为已捐赠

# After cleanup (donated inputs, no subsequent uses)
fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(x, residual, weight, 1e-5)

原地功能化（跟踪捐赠输入）和克隆清理的结合使编译器能够安全地使用原地内核，而无需添加冗余副本或增加内存使用量。

6. Inductor 优化和代码生成¶

在 IR 降级和清理之后，图只包含标准的 PyTorch 操作和平台特定的自定义操作。然后 Inductor 执行其标准代码生成。

Inductor 降级和逐点融合：融合逐元素操作、归约等。
内存规划：确定缓冲区分配和重用
内核生成：为融合操作生成 Triton 或 C++ 代码
自动调优：选择最佳内核配置

管线摘要¶

Model Forward Pass
    ↓
[Dynamo Tracing] → FX Graph with vllm_ir.* ops
    ↓
[Pre-grad: Inplace Functionalization] → maybe_inplace → default, track donated inputs
    ↓
[AOTAutograd] → Functionalization
    ↓
[Post-grad: IR Fusion Passes] → Fuse high-level IR ops (e.g., rms_norm + quant)
    ↓
[Post-grad: IR Lowering] → vllm_ir.* ops → impl ops (with clones if needed)
    ↓
[Post-grad: Clone Cleanup] → Remove unnecessary clones using donated input info
    ↓
[Inductor] → Pattern matching, fusion, memory planning, codegen
    ↓
Compiled Code

vLLM IR 核心概念¶

操作声明¶

操作使用 @register_op 装饰器声明，它会创建一个 IrOp 对象

@register_op(
    name=None,           # Operation name (defaults to function name)
    activations=None,    # List of activation parameters (defaults to params starting with 'x')
    allow_inplace=False, # Whether to create a maybe_inplace overload
)
def op_name(...):
    ...

参数

activations：被视为“激活”的参数名称列表（通常由 maybe_inplace 消耗）。默认为以 x 开头的参数。
allow_inplace：为内存高效执行创建 maybe_inplace 重载（见下文）。

maybe_inplace 重载¶

maybe_inplace 重载是 LLM 推理中内存效率的关键特性。它表明调用者在操作后不需要保留激活输入，从而允许原地实现将输入内存重用于输出。

语义与用法¶

# Standard usage: inputs are preserved
out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
# x and residual are unchanged, out and res_out are new tensors

# maybe_inplace: inputs may be modified
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
# x and residual may be modified (undefined behavior to use them after this)
# out and res_out may alias x and residual

在将激活输入传递给 maybe_inplace 后再使用它将导致未定义行为

# WRONG: Using x after donating it
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
result = out + x  # ERROR: x was donated!

如果您需要保留输入，请使用默认重载或手动克隆

# Option 1: Use default overload
out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
result = out + x  # OK: x is preserved

# Option 2: Clone before maybe_inplace
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x.clone(), residual, weight, epsilon)
result = out + x  # OK: x is preserved, clone was donated

编译行为¶

在编译期间，原地功能化通道会验证捐赠的输入是否未再次使用，并将 maybe_inplace 转换为功能性的 default 重载。

# Inplace functionalization pass (pre-grad)
for node in graph.nodes:
    if node.target == torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace:
        # Check that activation inputs aren't used after this node
        for activation_arg in activation_inputs:
            for user in activation_arg.users:
                if user appears after node:
                    raise ValueError(f"Input {activation_arg} donated but used again")

        # Convert to default overload
        node.target = torch.ops.vllm_ir.fused_add_rms_norm.default

        # Track donated graph inputs for later clone elimination
        for i, arg in enumerate(node.args):
            if arg.op == "placeholder" and i in activation_indices:
                pass_context.donated_input_ids.add(node_to_idx[arg])

然后，克隆清理通道使用捐赠的输入信息，在原地内核降级时消除不必要的复制。

即时模式行为¶

在即时模式下（不使用 torch.compile），maybe_inplace 通过允许 IR 操作直接调度到原地实现，从而实现最大限度的内存高效执行。

# Eager dispatch logic for maybe_inplace
impl: IrOpImpl = ir_op.dispatch(*args)
return impl.impl_fn(*args)

# Eager dispatch logic for default:
impl: IrOpImpl = ir_op.dispatch(*args)
if impl.inplace:
  args = [
    arg.clone() if i in ir_op.activations else arg
    for i, arg in enumerate(args)
  ]
return impl.impl_fn(*args)

模型代码中的 maybe_inplace 和原地内核实现的结合，在即时模式和编译模式下都提供了最佳的内存效率，并且两种情况下的语义是相同的。

内存节省示例¶

考虑一个带有残差连接的 Transformer 层

# Without maybe_inplace (2 allocations per layer)
hidden_states = self.attention(input)
normed, residual = ir.ops.fused_add_rms_norm(hidden_states, input, weight, eps)
# Memory: input (preserved), hidden_states (preserved), normed (new), residual (new)

# With maybe_inplace (0 allocations per layer when using in-place kernel)
hidden_states = self.attention(input)
normed, residual = ir.ops.fused_add_rms_norm.maybe_inplace(hidden_states, input, weight, eps)
# Memory: normed (reuses hidden_states), residual (reuses input)

实现注册¶

实现使用 register_impl 方法注册

@ir.ops.op_name.register_impl(
    provider="provider_name",  # Unique identifier (e.g., "vllm_c", "aiter", "triton")
    supported=True,            # Static availability check
    supports_args=None,        # Dynamic argument support check
)
def impl_fn(...):
    ...

提供者命名约定

native：保留给原生 torch 实现（用 @register_op 声明）
vllm_c：通过 torch.ops._C 的 C++/CUDA 内核
aiter：AMD AITER 库
xpu_kernels：在 vllm-xpu-kernels 中实现的 SYCL/SYCLTLA 内核
triton_*：Triton 内核
其他实现的平台/库名称

支持检查

supported：静态布尔值，在导入时检查一次（例如，HAS_TRITON, is_cuda_alike()）
supports_args：函数 (*args, **kwargs) -> bool，检查参数兼容性
- 在编译期间使用伪张量调用以进行零成本检查
- 在即时模式调度期间使用真实张量调用
- 不应检查批大小或根据值添加守卫

示例支持谓词

def aiter_rms_norm_supports(x, weight, epsilon, variance_size=None):
    # Check dtype (OK: doesn't depend on batch size)
    if x.dtype not in [torch.float16, torch.bfloat16]:
        return False
    # Check optional parameter (OK: static check)
    if variance_size is not None:
        return False
    return True

@ir.ops.rms_norm.register_impl("aiter", supports_args=aiter_rms_norm_supports)
def rms_norm(...):
    ...

当设置 VLLM_BATCH_INVARIANT=1 时，批次无关的内核会自动选择。

即时模式对比编译模式¶

vLLM IR 操作在即时模式和编译模式下行为一致

即时模式

基于优先级列表直接调度到实现
使用真实张量参数检查支持
最小开销（如果需要可以进一步优化）

编译模式

IR 操作在 FX 图中显示为 torch.ops.vllm_ir.* 自定义操作
降级使用伪张量选择实现
与 Inductor 优化完全集成

这种一致性使得

自信地在即时模式下进行原型设计
通过禁用编译进行调试
从即时执行到编译执行的逐步迁移

其他主题¶

外部实现¶

外部平台无需修改 vLLM 即可注册实现

# In external package
from vllm import ir

@ir.ops.rms_norm.register_impl("my_platform", supported=is_my_platform())
def rms_norm(x, weight, epsilon, variance_size=None):
    return my_platform.rms_norm(x, weight, epsilon)

然后配置优先级以使用您的实现

class MyPlatform(Platform):
  def get_default_ir_op_priority(self):
    return IrOpPriorityConfig(rms_norm=['my_platform', 'native'])

# Users can still override priority in the same way
llm = LLM(ir_op_priority=IrOpPriorityConfig(rms_norm=['custom_oot_kernel']))

调试与可观测性¶

注意

请告诉我们如何为您的用例改进可观测性！

启用调试日志以查看内核选择

VLLM_LOGGING_LEVEL=DEBUG vllm serve ...

这会记录

每个操作选择了哪些实现
为什么实现被拒绝（不支持，参数不支持）
编译缓存命中/未命中
IR 降级统计

检查编译图中选定的实现

# After compilation, inspect the lowering pass
lowering_pass = backend.lowering_pass
print(lowering_pass.selected_impls)
# Output: {'rms_norm': {'node_123': 'vllm_c', 'node_456': 'vllm_c'}}

从 CustomOp 迁移¶

vLLM IR 旨在与 CustomOp 共存并逐步取代它

操作声明：将 CustomOp 类 PluggableLayer 转换为 @register_op 函数，并将 forward_native 移到该函数中
实现注册：使用 @ir.ops.op_name.register_impl 而不是重写方法
层使用：将 self.op(...) 替换为 ir.ops.op_name(...)
配置：将 --compilation-config.custom-ops 迁移到 --ir-op-priority

迁移可以逐步进行，每次一个操作。

参见¶

torch.compile 集成 - 通用编译基础设施
融合 - vLLM 中的自定义融合和转换通道
自定义操作 - 遗留自定义操作系统