vLLM 自定义算子注册与分发
1. 概述
阅读 vLLM 源码时，最常见的困惑是：当代码中调用 torch.ops.vllm.unified_attention(...) 或 torch.ops._C.paged_attention_v1(...) 时，底层到底发生了什么？这个调用如何路由到 FlashAttention、Triton 或 CPU 的具体实现？
vLLM 使用了 三套并行的算子注册机制，分别服务于 C++/CUDA kernel、Python 注册的复合算子、以及纯 Triton kernel。理解这些机制是读懂整个代码库的前提。
本文以 attention 算子为锚点，逐层拆解整个注册和分发架构。
2. 整体架构
Python 上层代码
│
├── torch.ops._C.xxx ──→ C++/CUDA 原生算子（非稳定 ABI）
│ └── 在 _custom_ops.py 中有 Python 封装函数
│
├── torch.ops._C.xxx ──→ C++/CUDA 原生算子（稳定 ABI）
│ └── 同一个命名空间，但由 _C_stable_libtorch 提供
│
├── torch.ops.vllm.xxx ──→ Python 注册的复合算子
│ └── 内部 dispatch 到不同 backend 的实现
│
└── 直接 Python 函数调用 ──→ @triton.jit 编译的 Triton kernel
└── 不注册为 torch custom op，就是普通的 Python 函数调用
三种机制的核心差异：
| TORCH_LIBRARY_EXPAND | STABLE_TORCH_LIBRARY_FRAGMENT | torch.library.Library | |
|---|---|---|---|
| 注册位置 | C++ .cpp 文件 | C++ .cpp 文件 | Python .py 文件 |
| 命名空间 | torch.ops._C.* | torch.ops._C.*（同名！） | torch.ops.vllm.* |
| ABI 兼容性 | 绑死 PyTorch 版本 | 跨 PyTorch 版本兼容 (≥ 2.10) | Python 层面，天然兼容 |
| dispatch key | 在 m.impl() 中硬编码 | 在 TORCH_LIBRARY_IMPL 中硬编码 | 在 direct_register_custom_op() 中传参 |
| 典型算子 | paged_attention, cutlass_scaled_mm | layernorm, rotary_embedding, 量化 kernel | unified_attention, unified_kv_cache_update |
| 编译 | 需要 PyTorch C++ 头文件 | 需要 SABI 兼容的头文件 | 不需要编译 |
3. C++/CUDA 原生算子：TORCH_LIBRARY_EXPAND
3.1 注册入口
所有非稳定 ABI 的 C++/CUDA 算子都在 torch_bindings.cpp 中注册。这个文件约 450 行，是整个 C++ 层的注册总线。
注册模式：
// torch_bindings.cpp:39-49
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def(
"paged_attention_v1(Tensor! out, Tensor query, Tensor key_cache,"
" Tensor value_cache, int num_kv_heads, float scale,"
" Tensor block_tables, Tensor seq_lens, int block_size,"
" int max_seq_len, Tensor? alibi_slopes,"
" str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
" int tp_rank, int blocksparse_local_blocks,"
" int blocksparse_vert_stride, int blocksparse_block_size,"
" int blocksparse_head_sliding_step) -> ()");
ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
}
两步：
ops.def(schema_string)— 声明算子签名（输入/输出类型）ops.impl(name, dispatch_key, &function_ptr)— 绑定 C++ 函数指针
3.2 宏展开
TORCH_LIBRARY_EXPAND 定义在 registration.h:13：
#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
这一层间接宏展开是为了让 TORCH_EXTENSION_NAME（编译时确定）能被正确替换。TORCH_EXTENSION_NAME 在 CMake 构建时注入：
# cmake/utils.cmake:599
target_compile_definitions(${MOD_NAME} PRIVATE
"-DTORCH_EXTENSION_NAME=${MOD_NAME}")
3.3 多个扩展库
vLLM 编译产生多个独立的 .so 扩展库，各有自己的算子命名空间：
| 扩展库 | CMake 目标 | 命名空间 | 包含的算子类型 |
|---|---|---|---|
vllm._C | _C CMakeLists.txt:604 | torch.ops._C.* | paged_attention, cutlass_scaled_mm, rms_norm, rotary_embedding, silu_and_mul, 量化 kernel 等 |
vllm._C (子空间) | 同上 | torch.ops._C_cache_ops.* | reshape_and_cache, swap_blocks, copy_blocks |
vllm._C (子空间) | 同上 | torch.ops._C_cuda_utils.* | get_device_attribute |
vllm._C (子空间) | 同上 | torch.ops._C_custom_ar.* | 自定义 all-reduce |
vllm._moe_C | _moe_C CMakeLists.txt:1260 | torch.ops._moe_C.* | MoE 专用 kernel |
vllm._rocm_C | _rocm_C CMakeLists.txt:1281 | torch.ops._rocm_C.* | ROCm 专用 kernel |
3.4 Python 端触发加载
当 Python 执行 current_platform.import_kernels()（interface.py:242-249）时：
def import_kernels(self):
import vllm._C # ← 触发 REGISTER_EXTENSION → TORCH_LIBRARY 执行
import vllm._moe_C # ← 同上
REGISTER_EXTENSION 宏（registration.h:22-27）在模块导入时创建 PyInit_<name>() 函数，其执行会运行所有 TORCH_LIBRARY 块中注册的算子。
3.5 Python 封装层：_custom_ops.py
_custom_ops.py（约 3900 行）为所有 C++ 算子提供 Python 接口。它做以下事情：
- 类型转换与默认值：C++ 的
torch::kCUDA等 dispatch key 在 Python 侧不可见，封装函数提供 Python 友好的参数 - Tensor 格式处理：确保输入 tensor 的 dtype/device/contiguity 满足 kernel 要求
- 统一的错误处理：比直接调用
torch.ops._C.xxx(...)更友好的错误信息
典型模式（_custom_ops.py:2599）：
def reshape_and_cache(
key: torch.Tensor,
value: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
kv_cache_dtype: str,
k_scale: Optional[torch.Tensor] = None,
v_scale: Optional[torch.Tensor] = None,
) -> None:
torch.ops._C_cache_ops.reshape_and_cache(
key, value, key_cache, value_cache, slot_mapping,
kv_cache_dtype, k_scale, v_scale
)
3.6 函数声明的位置
C++ kernel 的实际实现（.cu 文件）通过头文件暴露函数签名：
- csrc/ops.h:18+ — paged_attention_v1, paged_attention_v2, rms_norm, rotary_embedding, silu_and_mul 等
- csrc/cache.h:9+ — reshape_and_cache, swap_blocks, copy_blocks 等
4. C++/CUDA 稳定 ABI 算子：STABLE_TORCH_LIBRARY_FRAGMENT
4.1 为什么需要稳定 ABI
PyTorch 2.x 的 C++ ABI 随着每个 minor 版本变化。如果 vLLM 为 PyTorch 2.5 编译的 .so，在 PyTorch 2.6 环境下可能无法加载。稳定 ABI（SABI, Stable Application Binary Interface）通过 STABLE_TORCH_LIBRARY_FRAGMENT 注册的算子与 PyTorch 版本解耦。
4.2 注册模式
libtorch_stable/torch_bindings.cpp:9-543 — 注册在 torch.ops._C.* 命名空间下（和非稳定 ABI 完全相同！）：
STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def(
"rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> ()");
}
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("rms_norm", TORCH_BOX(&rms_norm));
}
与非稳定版本的区别：
TORCH_LIBRARY_EXPAND | STABLE_TORCH_LIBRARY_FRAGMENT | |
|---|---|---|
| 宏 | TORCH_LIBRARY_EXPAND | STABLE_TORCH_LIBRARY_FRAGMENT |
| 实现宏 | ops.impl(name, key, &func) | ops.impl(name, TORCH_BOX(&func)) |
| dispatch 宏 | torch::kCUDA | STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) |
| 编译选项 | 标准 C++ ABI | -DUSE_SABI=3 |
| CMake 目标 | _C | _C_stable_libtorch |
| 算子范围 | 所有 kernel | 仅 layernorm、rotary embedding、激活函数、量化 kernel、能力检测 |
4.3 为什么同一个命名空间可以有两个提供者
PyTorch 的 dispatcher 支持多个库注册到同一命名空间。无论 _C 还是 _C_stable_libtorch 提供 torch.ops._C.rms_norm，调用方看到的是同一个算子。PT2 的 dispatcher 按注册顺序选择实现。
4.4 包含的算子类型
稳定 ABI 中注册的算子集中在经过充分验证、接口不会变化的模块：
torch.ops._C.rms_norm, fused_add_rms_norm # LayerNorm / RMSNorm
torch.ops._C.rotary_embedding # RoPE
torch.ops._C.silu_and_mul, gelu_and_mul, gelu_tanh_and_mul # 激活函数
torch.ops._C.gptq_gemm, gptq_shuffle, awq_gemm # GPTQ/AWQ 量化
torch.ops._C.ggml_mul_mat_a8, ggml_mul_mat, ... # GGML 量化
torch.ops._C.cutlass_scaled_mm, cutlass_w4a8_mm, ... # CUTLASS 量化
torch.ops._C.is_power_of_two_nvfp4, ... # 能力检测
5. Python 注册的复合算子：torch.ops.vllm.*
5.1 注册机制
与 C++ 注册不同，torch.ops.vllm.* 下的算子完全从 Python 侧注册，不需要编译 C++ 代码。
核心基础设施在 torch_utils.py:927-969：
vllm_lib = Library("vllm", "FRAGMENT") # 创建 torch.ops.vllm 命名空间
def direct_register_custom_op(
op_name: str,
op_func: Callable,
mutates_args: list[str],
fake_impl: Optional[Callable] = None,
dispatch_key: Optional[str] = None,
):
schema_str = infer_schema(op_func, mutates_args=...) # 自动推断签名
vllm_lib.define(op_name + schema_str) # 注册签名
vllm_lib.impl(op_name, op_func, dispatch_key=dispatch_key) # 绑定实现
if fake_impl is not None:
vllm_lib._register_fake(op_name, fake_impl) # 用于 torch.compile
dispatch_key 由平台决定：
- NVIDIA GPU →
"CUDA"（platforms/cuda.py:166） - CPU →
"CPU"（platforms/cpu.py:46） - XPU →
"XPU"
5.2 关键算子
这些算子都是 direct_register_custom_op() 注册的，表驱动如下：
| 算子 | 定义位置 | 行号 |
|---|---|---|
torch.ops.vllm.unified_attention_with_output | attention.py:734-761 | 注册在 L777-781 |
torch.ops.vllm.unified_kv_cache_update | attention.py:725-729 | 注册在 L725-729 |
torch.ops.vllm.maybe_calc_kv_scales | attention.py:615-628 | — |
torch.ops.vllm.fused_rope_and_unified_kv_cache_update | rope_kvcache_fusion.py:85-89 | 编译优化 pass |
torch.ops.vllm.rocm_aiter_* (30+ ops) | _aiter_ops.py | ROCm aiter 库的算子封装 |
5.3 为什么这些算子用 Python 注册
这些算子不是”原子”的 kernel 调用，而是包含 dispatch 逻辑的 Python 函数。以最核心的 unified_attention_with_output 为例：
# attention.py:734-761
def unified_attention_with_output(query, key, value, output, layer_name, ...):
# Step 1: 从上下文获取当前 attention 层的实现
attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
# Step 2: 委托给具体 backend 的 forward 方法
self.impl.forward(self, query, key, value, kv_cache, attn_metadata, output=output)
它不是直接调用 CUDA kernel，而是做了一个间接分发——self.impl 在模型初始化时已经根据 attention backend 选择好了。
5.4 Attention 算子的完整分发流程
这是理解”Python → backend 实现”的最关键路径：
模型层中调用:
torch.ops.vllm.unified_attention_with_output(query, key, value, output, layer_name, ...)
│
▼
实现函数 [attention.py:734]:
def unified_attention_with_output(...):
attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
self.impl.forward(self, query, key, value, kv_cache, attn_metadata, output)
│
├── FlashAttention backend [flash_attn.py:667]
│ └── flash_attn_varlen_func(query, key_cache, value_cache, ..., block_table=...)
│ └── flash-attn 库的 CUDA kernel
│
├── Triton backend [triton_attn.py:638-670]
│ └── unified_attention(q=query, k=key_cache, v=value_cache, out=output, ...)
│ └── @triton.jit kernel_unified_attention(grid, ...) (纯 Python 调用)
│
├── FlashInfer backend
│ └── flashinfer 库的 kernel
│
└── ROCm backend
└── aiter 库的 kernel
6. Triton Kernel 的分发方式
6.1 不注册为 torch custom op
Triton kernel 不是 torch.ops.vllm.* 或 torch.ops._C.*。它们就是普通的 Python 函数，用 @triton.jit 装饰后由 Triton 编译器处理。
# triton_unified_attention.py:179
@triton.jit
def kernel_unified_attention(...):
# GPU kernel 逻辑
# triton_unified_attention.py:763
def unified_attention(q, k, v, out, ...):
# Python wrapper: 计算 grid/block 配置后 launch kernel
grid = (num_heads, num_seqs)
kernel_unified_attention[grid](q, k, v, out, ...)
6.2 调用方式——就是普通函数调用
在 attention backend 中（如 triton_attn.py:638-670）：
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
# 直接调用，就是个普通的 Python 函数
unified_attention(q=query, k=key_cache, v=value_cache, out=output, ...)
没有 torch.ops.xxx 的包装，没有 m.impl() 的注册，就是 import + 调用。
6.3 主要 Triton Kernel
| Kernel 包装函数 | 文件 | 行号 | 用途 |
|---|---|---|---|
unified_attention | triton_unified_attention.py | L763 | 统一的 prefill + decode attention |
decode_attention_fwd | triton_decode_attention.py | L729 | 仅 decode 阶段的 attention |
context_attention_fwd | triton_prefill_attention.py | L191 | 仅 prefill 阶段的 attention |
triton_reshape_and_cache_flash | triton_reshape_and_cache_flash.py | L319 | KV cache 写入 |
6.4 部分 Triton Kernel 也注册为 torch.ops
triton_unified_attention.py 的 unified_attention() 是纯 Python 调用，但某些量化相关的 Triton kernel 被包装后注册为 torch.ops.vllm.*：
# 例如：w8a8 Triton kernel → 注册为 torch custom op
direct_register_custom_op(
"w8a8_triton_block_scaled_mm_func",
w8a8_triton_block_scaled_mm_func,
...
)
这样做的好处是：调用方可以像用 C++ kernel 一样使用 Triton kernel，不需要关心底层实现语言。
7. Rust 代码的角色
vLLM 有 Rust 代码，但不参与算子注册。Rust 代码是一个独立的前端二进制程序 vllm-rs（Cargo.toml），在 Python 侧由 v1/utils.py:240 启动。它包括：
chat,llm— LLM 前端server— HTTP 服务器tokenizer— 高性能 tokenizerreasoning-parser,tool-parser— 推理和工具调用解析
Rust 和 Python 的关系：Rust 做前端（接收请求、返回响应、tokenize），Python/CUDA 做推理引擎。它们通过进程间通信而非算子调用连接。
8. CustomOp 类——Platform Dispatch
8.1 用途
custom_op.py:103-353 — 对于跨平台有多个实现的高层算子（如 fused_moe、量化线性层），vLLM 提供了 CustomOp 基类：
class CustomOp:
def forward_cuda(self, *args, **kwargs): raise NotImplementedError
def forward_hip(self, *args, **kwargs): raise NotImplementedError
def forward_cpu(self, *args, **kwargs): raise NotImplementedError
def forward_xpu(self, *args, **kwargs): raise NotImplementedError
def dispatch_forward(self, *args, **kwargs): # [L174]
if current_platform.is_cuda():
return self.forward_cuda(*args, **kwargs)
elif current_platform.is_rocm():
return self.forward_hip(*args, **kwargs)
# ...
8.2 注册和使用
# 声明一个 custom op
@CustomOp.register("fused_moe")
class FusedMoE(CustomOp):
def forward_cuda(self, ...):
# CUDA 实现
# 被其他模块调用
op = CustomOp.op_registry["fused_moe"] # 获取已注册的实现
result = op.dispatch_forward(*args) # 自动选择平台实现
CustomOp 用于高层模型算子（layer 级别），而 torch.ops._C.* / torch.ops.vllm.* 用于底层 kernel（individual kernel 级别）。
9. 算子注册全景图
┌─────────────────────────────────────────────────────────────────┐
│ Python 调用层 │
│ │
│ model.forward() │
│ └─ Attention.forward() [attention.py:437] │
│ ├─ torch.ops.vllm.unified_attention_with_output(...) │
│ │ └──→ Python 注册的复合算子 [torch_utils.py:927] │
│ │ └─ attention.impl.forward() ─→ backend 分发 │
│ │ ├─ flash_attn_varlen_func() (FlashAttn) │
│ │ ├─ unified_attention() (Triton) │
│ │ └─ flashinfer_xxx() (FlashInfer) │
│ │ │
│ └─ torch.ops.vllm.unified_kv_cache_update(...) │
│ └──→ Python 注册 → do_kv_cache_update() │
│ ├─ reshape_and_cache_flash() (FlashAttn) │
│ │ └─ torch.ops._C_cache_ops.xxx (C++ op) │
│ └─ triton_reshape_and_cache_flash() (Triton) │
│ └─ @triton.jit kernel │
│ │
│ torch.ops._C.paged_attention_v1(...) ───→ C++ 非稳定 ABI │
│ torch.ops._C.rms_norm(...) ───→ C++ 稳定 ABI │
│ torch.ops._C.gptq_gemm(...) ───→ C++ 稳定 ABI │
│ torch.ops._C_cache_ops.reshape_and_cache(...) → C++ 非稳定 ABI │
├─────────────────────────────────────────────────────────────────┤
│ C++ 注册层 │
│ │
│ csrc/torch_bindings.cpp ──── TORCH_LIBRARY_EXPAND │
│ csrc/libtorch_stable/ ──── STABLE_TORCH_LIBRARY_FRAGMENT │
│ csrc/moe/torch_bindings ──── TORCH_LIBRARY_EXPAND (_moe_C) │
│ csrc/rocm/torch_bindings ──── TORCH_LIBRARY_EXPAND (_rocm_C) │
├─────────────────────────────────────────────────────────────────┤
│ CUDA Kernel 实现层 │
│ │
│ csrc/attention/ ── paged_attention_v1/v2, attention_kernels │
│ csrc/quantization/ ── gptq, awq, ggml, cutlass, marlin │
│ csrc/cache_kernels ── reshape_and_cache, swap_blocks │
│ csrc/layernorm/ ── rms_norm, fused_add_rms_norm │
│ csrc/activation/ ── silu_and_mul, gelu_and_mul │
│ csrc/rocm/ ── ROCm attention kernels │
│ csrc/cpu/ ── CPU fallback kernels │
├─────────────────────────────────────────────────────────────────┤
│ Triton Kernel 实现层 │
│ │
│ vllm/v1/attention/ops/triton_*.py ── @triton.jit kernels │
│ vllm/model_executor/kernels/linear/ ── 量化 linear 的 Triton │
└─────────────────────────────────────────────────────────────────┘
10. 从”阅读代码”的角度总结
当你在代码中看到以下模式，你应该立即知道它属于哪一层：
| 你看到的 | 它在做什么 | 去哪里找实现 |
|---|---|---|
torch.ops._C.xxx(...) | 调用 C++ 注册的 CUDA kernel | csrc/ 下的 .cu 文件 |
torch.ops._C_cache_ops.xxx(...) | 调用 KV cache 相关的 CUDA kernel | csrc/cache_kernels.cu |
torch.ops.vllm.unified_attention(...) | 调用 Python 注册的复合算子 → 内部 dispatch 到 backend | attention.py → flash_attn.py / triton_attn.py |
torch.ops.vllm.xxx(...) (其他) | 调用 Python 注册的算子 | grep "direct_register_custom_op.*xxx" 找到注册点 |
from triton_module import kernel_func 直接调用 | 调用 @triton.jit kernel | 对应的 triton_*.py 文件 |
CustomOp.op_registry["xxx"].dispatch_forward(...) | 平台分发的高层算子 | custom_op.py → grep "@CustomOp.register.*xxx" |
11. 关键文件索引
| 组件 | 文件 | 核心行号 |
|---|---|---|
| C++ op 注册总线 | torch_bindings.cpp | L18-445 |
| C++ registration 宏 | registration.h | L13 (EXPAND), L22 (REGISTER_EXTENSION) |
| C++ op 头文件 | ops.h | L18+ |
| C++ cache op 头文件 | cache.h | L9+ |
| 稳定 ABI 注册 | libtorch_stable/torch_bindings.cpp | L9-543 |
| Python op 封装 | _custom_ops.py | 约 3900 行 |
| Python op 注册工具 | torch_utils.py | L927-969 |
| Attention 算子分发点 | attention.py | L437 (forward), L725-729 (kv_cache_update), L734-761 (unified_attention) |
| FlashAttention backend | flash_attn.py | L667 (forward) |
| Triton backend | triton_attn.py | L638-670 |
| Triton unified attention kernel | triton_unified_attention.py | L179 (kernel), L763 (wrapper) |
| Platform import_kernels | interface.py | L242-249 |
| CustomOp 平台分发 | custom_op.py | L103-353 |
| CMake 构建 | CMakeLists.txt | L604 (_C), L1037 (_C_stable_libtorch), L1260 (_moe_C) |
| Rust 入口 | Cargo.toml | — |