vLLM 源码完全拆解（三十三）：算子抽象层与平台分派——一次 RMSNorm 调用如何落到 CUDA、ROCm、TPU 与 Ascend

为什么推理框架要自建算子层

PyTorch 已经提供了 torch.nn.functional.layer_norm、torch.nn.functional.silu、torch.bmm 这些算子，从语义上说，vLLM 完全可以拿它们拼出一个正确的模型。问题在于推理场景对每个算子的要求与训练框架不同：形状是动态的（batch 内序列长度从 1 到几万不等）、数据类型经常是 FP8 或 4-bit packed、访存模式围绕分页 KV 展开、单次调用可能只有一个 token。通用算子为了覆盖所有形状和 dtype，往往带多层分发和边界检查，在 decode 这种小算子里，框架开销会超过计算本身。

因此推理框架会为高频路径准备专用内核：RMSNorm 有按 token 一行的向量化实现，SwiGLU 把 gate、up、激活和乘法融成一个 kernel，MoE 有按专家分组的 GEMM，注意力有分页 KV 的专用内核。每多一种硬件，就要多一套内核；每多一种量化格式，就要多一组参数布局。如果模型层直接 import 某个 CUDA 扩展，代码立刻失去可移植性。

vLLM 的解法不是把硬件判断散落在模型代码里，而是抽出两层抽象：一层负责“同一个语义在不同平台上的 Python 方法选择”，另一层负责“同一个算子的语义与多个内核实现之间的选择”。前者是 CustomOp，后者是 vLLM IR。两者最终都把调用落到 torch.ops 上，由 PyTorch 的 dispatcher 把 C++ 函数送进对应的 CUDA、HIP 或 CPU 实现。

全景：从模型层到 torch.ops._C

先看一张自上而下的地图。下表是一次 RMSNorm 调用会经过的层次，以及每层负责做的决策。

这张表里有两个容易混淆的点，值得先说明。

第一，CustomOp 与 vLLM IR 不是替代关系，而是叠加关系。RMSNorm 本身仍然是 CustomOp 的子类，平台相关的 Python 分支仍然由 CustomOp.dispatch_forward 决定；但在 CUDA 分支内部，forward_native 调用的不是 torch.nn.functional，而是 ir.ops.rms_norm。IR 层再根据优先级挑选 vllm_c、native、oink、aiter 等实现。也就是说，一个算子可能先经过平台方法选择，再经过内核实现选择。

第二，最终的 torch.ops._C.rms_norm 是 C++ 扩展，不是一个 Python 函数。PyTorch 的 dispatcher 会根据张量所在的设备选择 dispatch key，vLLM 在注册时把 CUDA 实现挂在 CUDA key 上，把纯形状检查类函数挂在 CompositeExplicitAutograd 上。这个细节决定了 ROCm 为什么能复用同一份绑定代码。

图源：vLLM 官方文档 Architecture Overview，V1 Process Architecture（TP=4）。

从进程视角看，这张图解释了为什么平台分派必须是进程级的：EngineCore 进程负责调度，Worker 进程负责模型执行，每个 Worker 只看到自己那块设备。current_platform 在每个进程里独立初始化，out-of-tree 插件也必须在每个进程里被加载，否则就会出现调度进程认为自己在 CUDA 上、Worker 进程却在另一套后端上的错位。

平台单例：current_platform 是怎么选出来的

vllm.platforms.current_platform 是一个延迟初始化的单例。延迟的原因写在 platforms/__init__.py 的注释里：out-of-tree 平台插件需要 from vllm.platforms import Platform 来继承基类，如果模块导入时就把单例解析出来，插件还没有注册，解析必然失败；另外，vLLM 内部有些模块在 import 阶段就会访问 current_platform，必须保证插件已经加载完。

单例的实现依赖模块级的 __getattr__：

1
_current_platform = None
2

3
def __getattr__(name: str):
4
    if name == "current_platform":
5
        global _current_platform
6
        if _current_platform is None:
7
            platform_cls_qualname = resolve_current_platform_cls_qualname()
8
            _current_platform = resolve_obj_by_qualname(platform_cls_qualname)()
9
            global _init_trace
10
            _init_trace = "".join(traceback.format_stack())
11
        return _current_platform
12
    elif name in globals():
13
        return globals()[name]
14
    else:
15
        raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")

_init_trace 记录了第一次访问时的调用栈，用来排查“平台为什么被提前初始化”这类问题。这类细节在接入新硬件时很实用：如果插件里的某个 import 触发了 current_platform，栈里会直接显示出来。

真正的选择逻辑在 resolve_current_platform_cls_qualname。它先处理一个特例：VLLM_TARGET_DEVICE=cpu 时直接返回 CPU 平台，不再探测其他插件。注释解释得很清楚：CPU-only 的 CI 任务可能跑在带加速卡的机器上，复用加速卡 wheel，如果继续探测，就会同时激活 CPU 和主机加速卡两个插件。

之后，它把内置探测函数与通过 entry point 注册的 out-of-tree 插件放在同一条链上依次调用：

1
platform_plugins = load_plugins_by_group(PLATFORM_PLUGINS_GROUP)
2

3
for name, func in chain(builtin_platform_plugins.items(), platform_plugins.items()):
4
    platform_cls_qualname = func()
5
    if platform_cls_qualname is not None:
6
        activated_plugins.append(name)

内置插件有五个：tpu、cuda、rocm、xpu、cpu。它们的探测方式各不相同，但都遵循同一个约定：能识别出对应硬件就返回平台类的全限定名，否则返回 None。

• cuda_platform_plugin 通过 NVML 查询设备数量，同时检查 vLLM 是否是用 CPU 版本编译的。后者是为了避免“在 GPU 机器上装了 CPU 版 vLLM”时误判。NVML 不可用时还会检查 Jetson 的 /etc/nv_tegra_release。
• rocm_platform_plugin 通过 amdsmi 查询处理器句柄。WSL 环境下如果 amdsmi 不可用，会退回到 PyTorch 的 torch.version.hip 与 torch.accelerator.is_available()。
• xpu_platform_plugin 检查 torch.xpu.is_available()，并顺便设置分布式后端为 xccl。
• tpu_platform_plugin 检查 Pathways 代理或 libtpu 是否可导入。
• cpu_platform_plugin 检查构建标签、VLLM_TARGET_DEVICE 与 macOS；在 AMD Zen 且安装了 zentorch 时返回 ZenCpuPlatform。

选择规则写得很严格：out-of-tree 插件优先。只要激活了一个 OOT 插件，就用它；激活两个 OOT 插件直接抛 RuntimeError，因为一个进程只能有一套硬件语义。没有 OOT 插件时，内置插件也必须恰好激活一个，多了同样报错；一个都没有时落到 UnspecifiedPlatform，后续大部分算子会走 native 路径。

这条规则带来一个实际后果：在昇腾机器上同时装了 CUDA 版 vLLM 和 vllm-ascend 插件时，OOT 插件会赢，即便 NVML 也报告了 GPU。对多硬件共存的开发机，这避免了“插件装了但没生效”的困惑。

Platform 基类定义了一组供上层查询的属性和钩子。最常用的几个是：

• dispatch_key：默认 "CPU"，CudaPlatform 覆写为 "CUDA"。这个字符串会直接传给 torch.library.Library.impl，决定自定义算子注册到哪个 dispatch key。
• device_type 与 device_name：用于创建张量、选择设备。
• supported_dtypes：不同计算能力支持的 dtype 列表，CUDA 平台会按 has_device_capability(80) 返回 BF16/FP16/FP32，老卡则去掉 BF16。
• import_kernels()：默认实现导入 vllm._C，CUDA 平台覆写为导入 vllm._C_stable_libtorch、_moe_C_stable_libtorch，并可选导入 _qutlass_C。
• import_ir_kernels()：默认导入 vllm.kernels，把内置 IR 实现注册进 op registry；OOT 平台应覆写它来注册自己的内核模块。
• get_attn_backend_cls(...)：按选中的注意力后端枚举返回后端类路径，这是注意力层真正的平台分派入口。
• get_default_ir_op_priority(vllm_config)：返回该平台的 IR 优先级默认值。
• apply_config_platform_defaults(vllm_config) 与 check_and_update_config(vllm_config)：在配置初始化阶段修改默认值或做兼容性检查。

dispatch_key 这个属性值得多说一句。它的默认值是 "CPU"，注释写得很明确：对于没有在 PyTorch 里注册 dispatch key 的平台，用 CPU 作为兜底。这个值会被 direct_register_custom_op 使用——当调用方没有显式指定 dispatch key 时，它取 current_platform.dispatch_key，把 Python 函数注册到该 key 上。如果一块新硬件在 PyTorch 里对应 PrivateUse1，但平台类没有覆写 dispatch_key，那么通过 torch.library 注册的算子会被挂到 CPU 上，设备张量调用时可能根本命中不了预期实现。相反，torch.ops._C.rms_norm 这类 C++ 扩展是在编译期按 CUDA key 注册的，不受这个属性影响，两者不要混为一谈。

get_default_ir_op_priority 在基类里返回 ["native"]，CudaPlatform 覆写为：使用 inductor 编译时返回 ["native"]，eager 或只做 Dynamo trace 时返回 ["vllm_c", "native"]。这条规则的动因会在后面讲编译边界时展开。

图源：vLLM 官方文档 Architecture Overview，Class Hierarchy。

类层次图里可以看到 Platform 位于最上层，向下依次影响 Executor、Worker 与 ModelRunner。算子层的分派不是孤立的一行 if，而是这条层次结构把“当前设备是什么”一路传到了每一次内核选择。

CustomOp：按平台挑 forward，按编译开关挑实现

CustomOp 是 vLLM 最早的算子抽象，解决的问题很直接：一个 nn.Module 在不同平台上要有不同的 forward 实现，但模型层不应该出现平台判断。

先看它的注册表：

1
op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
2
op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}

op_registry 里是 vLLM 自带的算子类，op_registry_oot 里是 out-of-tree 平台注册的替换类，两者都按注册名索引。注册靠装饰器：

1
@classmethod
2
def register(cls, name: str, dynamic_arg_dims=None):
3
    def decorator(op_cls):
4
        assert name not in op_registry, f"Duplicate op name: {name}"
5
        op_cls.name = name
6
        op_cls._dynamic_arg_dims = dynamic_arg_dims
7
        op_registry[name] = op_cls
8
        return op_cls
9
    return decorator

OOT 替换则用 register_oot，注册名默认取被装饰类的类名：

1
@classmethod
2
def register_oot(cls, _decorated_op_cls=None, name: str | None = None):
3
    def decorator(op_cls):
4
        reg_name = name if name is not None else cls.__name__
5
        assert reg_name not in op_registry_oot
6
        op_cls.name = reg_name
7
        op_registry_oot[reg_name] = op_cls
8
        return op_cls

关键的替换发生在 __new__，而不是 forward 里：

1
def __new__(cls, *args, **kwargs):
2
    op_name = cls.__name__
3
    if op_name not in op_registry_oot:
4
        op_cls_to_instantiate = cls
5
    else:
6
        op_cls_to_instantiate = op_registry_oot[op_name]
7
    return super().__new__(op_cls_to_instantiate)

这意味着模型代码写 RMSNorm(hidden_size) 时，拿到的可能已经是 OOT 平台提供的子类实例。构造完成后，__init__ 再把 forward 绑定到一个具体方法：

1
def __init__(self, *, enforce_enable=False, compile_native=False):
2
    super().__init__()
3
    self._enforce_enable = enforce_enable
4
    self._forward_method = self.dispatch_forward(compile_native=compile_native)
5

6
def forward(self, *args, **kwargs):
7
    return self._forward_method(*args, **kwargs)

dispatch_forward 做两件事：先决定这个算子是否启用，再按平台选方法。

1
enabled = self._enforce_enable or self.enabled()
2
if enabled:
3
    compilation_config.enabled_custom_ops.update([self.__class__.name])
4
else:
5
    compilation_config.disabled_custom_ops.update([self.__class__.name])
6

7
if not enabled:
8
    return self.maybe_compile(self.forward_native, enable=compile_native)
9

10
if current_platform.is_rocm():
11
    return self.forward_hip
12
elif current_platform.is_cpu():
13
    return self.forward_cpu
14
elif current_platform.is_tpu():
15
    return self.forward_tpu
16
elif current_platform.is_xpu():
17
    return self.forward_xpu
18
elif current_platform.is_out_of_tree():
19
    return self.forward_oot
20
else:
21
    return self.forward_cuda

几个默认回退值得记住：

• forward_hip 默认调用 forward_cuda，因为 HIP 算子通常与 CUDA 兼容；
• forward_xpu、forward_cpu、forward_tpu、forward_oot 默认调用 forward_native；
• 最后落到 forward_cuda，所以未识别的平台也会先尝试 CUDA 实现，基类的 forward_cuda 会抛 NotImplementedError，迫使算子作者至少实现一个平台方法。

enabled() 的语义与 CompilationConfig.custom_ops 绑定。这个配置支持 all、none 两种基础模式，以及 +op、-op 形式的逐个开关：

1
enabled = f"+{cls.name}" in custom_ops
2
disabled = f"-{cls.name}" in custom_ops
3
if enabled and disabled:
4
    raise ValueError("custom_ops cannot both enable and disable ...")
5
return (CustomOp.default_on() or enabled) and not disabled

default_on() 要求列表中恰好有一个基础模式：包含 none 则默认关闭，包含 all 则默认开启。vLLM 在配置初始化时按编译后端补齐这个基础模式：使用 inductor 且编译模式不是 NONE 时默认追加 none，否则追加 all。这条规则非常关键，它解释了为什么同一个模型在 eager 模式下走厂商内核、在 torch.compile 模式下却走 PyTorch 原生算子：inductor 需要看见算子内部的计算，才能做融合；一个不透明的自定义算子会切断融合。

maybe_compile 是为另一种情况准备的：一个 CustomOp 被包在更大的不透明自定义算子内部（例如 fused_moe、unified_attention），模型级编译看不到它内部。此时如果该算子处于“关闭自定义实现”的状态，forward_native 仍然需要被编译，才能避免 eager 下逐个小算子发射。它根据 _dynamic_arg_dims 对指定维度调用 torch._dynamo.mark_dynamic，再以 dynamic=True 编译，减少重编译。代码注释也提醒了边界：这种编译不跨算子融合，因此能用普通算子表达的地方仍应尽量展开。

一个更具体的例子是 SiluAndMul：

1
@CustomOp.register("silu_and_mul")
2
class SiluAndMul(CustomOp):
3
    def __init__(self, *, compile_native: bool = True):
4
        super().__init__(compile_native=compile_native)
5
        if (current_platform.is_cuda_alike()
6
                or current_platform.is_xpu()
7
                or (current_platform.is_cpu()
8
                    and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC)):
9
            self.op = torch.ops._C.silu_and_mul
10

11
    @staticmethod
12
    def forward_native(x):
13
        d = x.shape[-1] // 2
14
        return F.silu(x[..., :d]) * x[..., d:]
15

16
    def forward_cuda(self, x):
17
        d = x.shape[-1] // 2
18
        out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device)
19
        self.op(out, x)
20
        return out

它的内核句柄在构造期就选好了，而不是每次 forward 再判断平台。这是 CustomOp 的典型用法：把平台判断移到初始化阶段，让热路径只剩下一次函数调用。

vLLM 里的层大致有三种形态，选择哪一种取决于需要替换的范围。

• 普通 nn.Module：没有平台差异，或者平台差异已经在内部通过别的机制处理，例如只做张量拼接、形状变换的工具层。
• CustomOp：单个算子的 forward 在不同平台上有不同实现，或者需要在编译模式下退回 native。典型例子是 RMSNorm、SiluAndMul、RotaryEmbedding。
• PluggableLayer：整个模块的结构都可能被替换，包括子模块的组合方式、参数初始化、forward 的组织。典型例子是注意力层、MoE 层这类需要厂商重写内部结构的模块。

PluggableLayer 不提供 forward_* 方法分派，它的替换发生在 __new__：只要 OOT 表里有同名类，实例化的就是 OOT 类。这个设计把“换实现”从方法级提升到类级，代价是 OOT 类必须自己保证与基类的接口兼容。对于要改注意力层内部子模块结构的平台，这比在 forward_cuda 里塞一堆条件判断更清晰。

两种替换机制的选择可以归纳成一张对照：

实际工程里最常见的组合是：核心计算算子用 IR provider，少数语义差异大的算子用 register_oot，注意力和 MoE 用 PluggableLayer。这样既能复用 vLLM 的大部分模型代码，又能让厂商专注于真正不同的部分。

一次 RMSNorm 的完整链路

把上面的抽象串起来，最直观的方式是完整走一遍 RMSNorm。vLLM 当前的 layernorm.py 里，这个类同时用了 CustomOp 和 vLLM IR 两套机制，正好能展示叠加关系。

1
@CustomOp.register("rms_norm")
2
class RMSNorm(CustomOp):
3
    def __init__(self, hidden_size, eps=1e-6, var_hidden_size=None,
4
                 has_weight=True, dtype=None):
5
        super().__init__()
6
        self.hidden_size = hidden_size
7
        self.variance_epsilon = eps
8
        self.variance_size_override = (
9
            None if var_hidden_size == hidden_size else var_hidden_size
10
        )
11
        self.pass_weight = has_weight
12
        self.pass_weight_add = has_weight
13
        ...

1
def forward_cuda(self, x, residual=None):
2
    if envs.VLLM_BATCH_INVARIANT:
3
        ...
4
        return rms_norm_batch_invariant(
5
            x, self.weight.data if pass_weight else None,
6
            self.variance_epsilon, residual=residual,
7
        )
8
    return self.forward_native(x, residual)

1
def forward_native(self, x, residual=None):
2
    if residual is None:
3
        return ir.ops.rms_norm(
4
            x, self.weight.data if self.pass_weight else None,
5
            self.variance_epsilon, self.variance_size_override,
6
        )
7
    return ir.ops.fused_add_rms_norm.maybe_inplace(
8
        x, residual, self.weight.data if self.pass_weight_add else None,
9
        self.variance_epsilon, self.variance_size_override,
10
    )

1
@register_op
2
def rms_norm(x, weight, epsilon, variance_size=None):
3
    orig_dtype = x.dtype
4
    x = x.to(torch.float32)
5
    x_var = x if variance_size is None else x[..., :variance_size]
6
    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
7
    x = x * torch.rsqrt(variance + epsilon)
8
    if weight is not None:
9
        x = x.to(weight.dtype) * weight
10
    return x.to(orig_dtype)

1
rms_no_var_size = lambda x, weight, epsilon, variance_size=None: (
2
    variance_size is None and (weight is None or weight.dtype == x.dtype)
3
)
4

5
@ir.ops.rms_norm.register_impl(
6
    "vllm_c", supports_args=rms_no_var_size, supported=GPGPU_DEVICE
7
)
8
def rms_norm(x, weight, epsilon, variance_size=None):
9
    assert variance_size is None
10
    if IS_ROCM and (x.dim() > 2 or not x.is_contiguous()):
11
        original_shape = x.shape
12
        x = x.reshape(-1, original_shape[-1])
13
        output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
14
        torch.ops._C.rms_norm(output, x, weight, epsilon)
15
        return output.reshape(original_shape)
16

17
    output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
18
    torch.ops._C.rms_norm(output, x, weight, epsilon)
19
    return output

1
STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
2
  ops.def(
3
      "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) "
4
      "-> ()");
5
  ops.def(
6
      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
7
      "float epsilon) -> ()");
8
}

1
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
2
  ops.impl("rms_norm", TORCH_BOX(&rms_norm));
3
  ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm));
4
}

1
void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
2
              std::optional<torch::stable::Tensor> weight, double epsilon) {
3
  STD_TORCH_CHECK(out.is_contiguous());
4
  if (input.stride(-1) != 1) {
5
    input = torch::stable::contiguous(input);
6
  }
7

8
  int hidden_size = input.size(-1);
9
  int num_tokens = input.numel() / hidden_size;
10
  ...
11
  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
12
  const int max_block_size = batch_invariant_launch
13
      ? 1024 : ((num_tokens < 256) ? 1024 : 256);
14
  dim3 grid(num_tokens);
15
  ...
16
  const int calculated_vec_size = std::gcd(16 / sizeof(scalar_t), hidden_size);
17
  const int block_size = std::min(hidden_size / calculated_vec_size, max_block_size);
18
  ...
19
  vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, true>
20
      <<<grid, block, 0, stream>>>(out..., input..., strides..., weight_ptr,
21
                                   weight_stride, epsilon, num_tokens, hidden_size);
22
}

vLLM IR：语义与实现分离的第二代设计

CustomOp 的分派粒度是“平台上的 Python 方法”。它解决了一个算子要在 CUDA 与 CPU 上有不同写法的问题，但没有解决另一个问题：CUDA 上也可能有多个内核实现，谁快用谁；同一个平台的不同编译模式下，最优实现也可能不同；out-of-tree 厂商想给现有算子加一个新实现，不应该被迫替换整个类。

vLLM IR 就是为这组需求设计的。它把算子拆成“一个语义定义”和“多个具名实现”，并在调用时按优先级挑选。

1
@register_op(allow_inplace=True)
2
def fused_add_rms_norm(x, x_residual, weight, epsilon, variance_size=None):
3
    orig_dtype = x.dtype
4
    x = x.to(torch.float32)
5
    x = x + x_residual.to(torch.float32)
6
    x_residual = x.to(orig_dtype)
7

8
    x_var = x if variance_size is None else x[..., :variance_size]
9
    variance = x_var.pow(2).mean(dim=-1, keepdim=True)
10
    x = x * torch.rsqrt(variance + epsilon)
11
    if weight is not None:
12
        x = x.to(weight.dtype) * weight
13
    return x.to(orig_dtype), x_residual

1
schema = infer_schema(impl_fn, mutates_args=[])
2
if schema != op._schema_str:
3
    raise ValueError(
4
        f"Implementation for provider {provider} has schema '{schema}' "
5
        f"which does not match native schema '{op._schema_str}' ..."
6
    )

1
@ir.ops.rms_norm.register_impl(
2
    "vllm_c", supports_args=rms_no_var_size, supported=GPGPU_DEVICE
3
)

1
def set_default(self) -> None:
2
    """Permanently set the IR op priority for all op members."""
3
    for ir_op, op_priority in self._iter_op_priorities():
4
        ir_op.set_default(op_priority)

1
def func_impl_fn(self, *args, **kwargs):
2
    if not self.inplace:
3
        return self.impl_fn(*args, **kwargs)
4

5
    new_args = list(args)
6
    for i in self.op.activation_indices:
7
        new_args[i] = args[i].clone()
8
    return self.impl_fn(*new_args, **kwargs)

第三个维度：量化格式如何进入分派

平台与编译模式并不是内核选择的全部。同一个 CUDA 平台上，一个 Linear 可能走 BF16 的 cuBLAS、FP8 的 CUTLASS scaled GEMM、INT4 的 Marlin、AWQ 的 dequant 内核，或者某套 Triton 实现。决定这件事的是权重格式，而权重格式来自 checkpoint。

这条链路从 LinearBase.__init__ 开始：

1
if quant_config is None:
2
    self.quant_method = UnquantizedLinearMethod()
3
elif quant_method := resolve_quant_method(quant_config, self, prefix=prefix):
4
    self.quant_method = quant_method

resolve_quant_method 先按模块前缀向 QuantizationConfig 询问“这一层用什么方法”，再叠加在线量化配置。方法一旦确定，后面三件事都交给它：

• create_weights(...)：声明这一层有哪些参数。BF16 层只有 weight；FP8 层多出 per-tensor 或 block 的 weight_scale_inv；GPTQ/AWQ 还有 qzeros、scales。参数对象上带着各自的 weight_loader，决定按 TP rank 切哪一段。
• process_weights_after_loading(...)：把 checkpoint 布局整理成内核布局。FP8 需要转置成列主序、合并融合层的多个 scale；4-bit 权重需要按 Marlin 的 tile 重排。
• apply(...)：真正的前向分派，按格式与运行时形状选择 cutlass_scaled_mm、marlin_gemm、awq_gemm、per_token_group_quant 等内核。

这条轴与 IR 的优先级轴很像，但粒度不同。IR 是按“算子”维护一个全局优先级列表；量化是按“层”在构造期绑定一个方法对象，方法内部再根据格式、dtype、group size 与平台支持列表选内核。平台在这里的作用是给出 supported_quantization 列表，提前拒绝不支持的格式；真正的内核选择仍然发生在量化方法内部。

因此读 vLLM 的算子分派时，至少要区分三种选择：平台方法选择（CustomOp）、算子实现选择（IR 或量化方法内部）、编译路径选择（custom_ops 开关）。三者可以同时生效，例如一个 FP8 的 RMSNorm 在 CUDA 上可以先经过 forward_cuda，再走 vllm_c provider，最后因为启用了 inductor 而被替换成可融合的原生表达式。

初始化顺序：算子为什么必须在建壳前注册

算子的注册与绑定发生在多个 import 阶段，顺序错了会出现“类存在但内核没有”“平台是 CUDA 但 _C 没导入”这类问题。把关键节点按时间排开，会更容易排查。

有两个细节值得单独说。

layernorm.py 顶部那行 import vllm.kernels # noqa: F401 看起来像无用导入，实际上是注册副作用：导入 vllm.kernels 会依次导入 vllm_c、aiter_ops、oink_ops 与 Triton 相关模块，这些模块在 import 时执行 register_impl。如果没有这行，模型层第一次调用 ir.ops.rms_norm 时，provider 表里可能只有 native。

vllm/kernels/vllm_c.py 顶部的 current_platform.import_kernels() 同理。它确保在注册 vllm_c provider 之前，torch.ops._C.rms_norm 已经存在；否则 provider 虽然注册成功，第一次调用仍会在 torch.ops._C 上失败。把“导入扩展”和“注册 provider”放在同一个模块的导入期，顺序天然正确，也避免了运行时的懒加载分支。

还有一条容易忽视的约束：CustomOp 的 forward 方法在构造期绑定，enabled() 也在构造期求值。模型建好之后再修改 compilation_config.custom_ops，不会改变已经实例化的层。做 A/B 对比时要重新建模型，而不是在同一个进程里改配置再跑第二次。

平台抽象管不到的地方

Platform 是一层粗粒度的门，它回答“这台机器属于哪类硬件”“支持哪些 dtype 与量化格式”“默认注意力后端候选是谁”。它不回答“这一层该用哪个内核”。后者由各自的 selector 决定，平台只提供候选与校验。

注意力层是一个典型例子。Platform.get_attn_backend_cls(selected_backend, attn_selector_config, num_heads) 返回后端类路径，但选择逻辑要考虑用户显式指定的后端、head size、dtype、是否使用 MLA、是否启用分页、KV cache dtype、是否要求确定性。同一个 CUDA 平台上可能同时存在 FlashAttention、FlashInfer、Triton 与 FlexAttention 后端，具体选哪个由 AttentionSelectorConfig 与后端自身的支持矩阵决定。平台做的是“把不可能的后端排除掉”，而不是“替每一层做最终选择”。

MoE 层类似。moe_backend 配置可以在 Triton、CUTLASS、FlashInfer、DeepGEMM 等实现之间切换，平台负责校验该后端在当前设备上是否可用，层内部再按专家数、top-k、dtype、EP 规模与 token 分布选择分组 GEMM 或 fused kernel。

量化层也一样：平台给出 supported_quantization，量化方法内部再按 group size、是否有 zero point、是否需要 act quant 选择具体 kernel。

理解这个分工，能避免一个常见误解：以为把平台插件接好、supported_quantization 列表填对，所有算子就会自动用上厂商最优内核。实际上每个子系统都有自己的选择器，平台只是第一道闸门。vLLM IR 的优先级列表是试图把这类选择统一起来的一种尝试，但它只覆盖了 rms_norm、fused_add_rms_norm、gelu_and_mul_sparse 等少数算子，注意力和 MoE 仍然沿用各自的机制。

编译边界：custom op、Triton 与 inductor

推理框架的算子分派与 torch.compile 之间存在直接冲突。自定义算子是图上的一个黑盒，inductor 看不到内部，就无法把 RMSNorm 与后面的量化、把 SiLU 与后面的乘法融合成一个 Triton kernel。vLLM 的取舍写在 CompilationConfig.custom_ops 的默认值里：使用 inductor 时默认关闭 CustomOp，让 forward_native 里的 PyTorch 表达式暴露给编译器；eager 时默认开启，走厂商内核。

vLLM IR 试图把这件事做得更细。官方设计文档描述的流水线是：

1
模型前向
2
  → Dynamo 追踪：vllm_ir.* 作为不透明 custom op 进入 FX 图
3
  → pre-grad：把 maybe_inplace 改写成 default，记录被捐赠的输入
4
  → AOTAutograd：函数化
5
  → post-grad：高层 IR 融合 pass（如 allreduce + rms_norm）
6
  → post-grad：IR lowering，把每个 IR op 换成选中的实现
7
  → clone 清理：根据捐赠信息删掉多余克隆
8
  → inductor：常规的模式匹配、融合、内存规划与代码生成

这条流水线的关键在于 lowering 发生的时机。Dynamo 阶段保留高层语义，融合 pass 可以按一个 rms_norm 节点做匹配；到了 lowering 阶段，再用 fake tensor 作为参数代入选中的实现，通过 make_fx 把实现函数展开成一组底层算子。展开后插入的克隆，会在 clone 清理阶段被删掉。

编译路径下 supports_args 的判断依据是 fake tensor，而不是真实张量。这意味着支持判断只能依赖形状、dtype、stride 这类元数据，不能依赖张量的实际数值。写 provider 的 supports_args 时需要遵守这个约束，否则 eager 与编译两条路径会给出不同的选择。

图源：vLLM 官方博客 vLLM V1: A Major Upgrade to vLLM’s Core Architecture。

另外两个细节也影响分派行为。

CustomOp.maybe_compile 只编译被包在不透明算子内部的原生实现。它不能跨算子融合，所以注释明确说“能用普通算子表达的地方应尽量展开”。如果一个融合算子内部又塞了十几个小算子，编译只能减少每个小算子的发射开销，拿不到跨算子融合的收益。

IrOp.uuid 与 KernelConfig.ir_op_priority.compute_hash 一起参与编译缓存键。优先级列表变了、某个 provider 的实现源码变了，缓存都会失效。这条机制保证了“换了内核实现但复用了旧编译产物”这类难以排查的问题不会发生。

out-of-tree 平台如何接管

vLLM 的硬件插件机制在 2025 年 5 月由社区与昇腾团队推动成型。它的目标是把 Executor、Worker、ModelRunner、Attention Backend 与 Communicator 的实现从主仓库里解耦出去，让新硬件以独立 Python 包的形式接入。

接入点是一个标准 Python entry point，组名固定为 vllm.platform_plugins：

1
setup(
2
    name="vllm_ascend",
3
    entry_points={
4
        "vllm.platform_plugins": [
5
            "ascend = vllm_ascend:register",
6
        ],
7
    },
8
)

register 函数返回平台类的全限定名。vLLM 启动时，load_plugins_by_group 用 importlib.metadata.entry_points 枚举该组下的所有插件，逐个调用注册函数；随后 resolve_current_platform_cls_qualname 按前面说的规则选中唯一的 OOT 平台。

插件会被加载到 vLLM 创建的每个进程里，包括 EngineCore 与各个 Worker。这一点很关键：调度进程要用平台信息决定块大小、注意力后端与可用 dtype，Worker 进程要用平台信息导入内核。VLLM_PLUGINS 环境变量可以按名字过滤插件，便于在同一台机器上切换后端。

一个 OOT 平台通常要覆写这些钩子：

• import_kernels()：导入厂商编译的 C/C++ 扩展，让 torch.ops 下出现对应算子；
• import_ir_kernels()：导入注册了 IR provider 的模块，或直接调用 ir.ops.<name>.register_impl；
• get_attn_backend_cls(...)：返回厂商的注意力后端类路径；
• get_default_ir_op_priority(...)：给出厂商内核与 native 的优先级顺序；
• apply_config_platform_defaults(...) 与 check_and_update_config(...)：改写默认配置、校验不兼容项；
• pre_register_and_update(parser)：在配置初始化之前注册 out-of-tree 量化方法或修改命令行参数；
• register_oot 装饰器：替换单个 CustomOp 或 PluggableLayer 类。

PluggableLayer 是 CustomOp 的模块级对应物。它不提供 forward_* 的平台分派，而是在实例化时替换整个层类，适用于需要自定义子模块组合与初始化的场景，例如更换注意力层的内部结构、替换 MoE 层的专家组织方式。CustomOp.__new__ 与 PluggableLayer.__new__ 都查同一张 op_registry_oot 表，但键分别是算子名与层类名。

一个最小骨架大致是这样。它不依赖 vLLM 主仓库的任何私有接口，只用到前面读过的基类钩子：

1
from vllm.platforms import Platform, PlatformEnum
2

3
class MyPlatform(Platform):
4
    _enum = PlatformEnum.OOT
5
    device_name = "myaccel"
6
    device_type = "myaccel"
7
    dispatch_key = "PrivateUse1"   # PyTorch 已注册该后端时才这样写
8
    supported_quantization = ["fp8", "compressed-tensors"]
9

10
    @classmethod
11
    def import_kernels(cls):
12
        import myaccel_vllm._C  # 注册 torch.ops.myaccel_C
13

14
    @classmethod
15
    def import_ir_kernels(cls):
16
        from . import kernels   # 模块内调用 register_impl
17
        return kernels
18

19
def register():
20
    return "my_platform:MyPlatform"

对应的 entry point 写 my_platform = my_platform:register 即可。register 返回的是平台类路径，不是类对象，这样 vLLM 可以在确定插件唯一性之后再导入它，避免两个插件在导入期互相干扰。

如果厂商只提供了少量算子，也可以不写完整的 Platform 子类，而是在现有算子上注册 OOT 实现。例如用 @CustomOp.register_oot(name="RMSNorm") 替换整类，或调用 ir.ops.rms_norm.register_impl("myaccel", supported=True) 增加一个 provider。前一种方式适合算子语义本身不同、需要额外参数或子模块的场景；后一种方式适合语义一致、只是内核不同的场景。选择哪一种，取决于厂商实现是否需要改变层的构造行为。

需要提醒的是基类默认值：forward_oot 默认调用 forward_native。这意味着 OOT 平台如果只是注册了插件、没有覆写具体算子，运行时会静默落到 PyTorch 原生实现，性能可能远低于预期，但结果仍然正确。排查性能问题时，先确认厂商扩展是否真的被 import_kernels() 导入、对应算子是否真的注册了 forward_oot 或 IR provider，比盯着调度器更有效。

一条 decode 路径上的账：分派结果为什么会差出几毫秒

前面的机制听起来只是“选哪个函数”，落到 decode 阶段会变成实实在在的时间。举一个只做数量级估算的例子，假设模型 61 层、hidden size 7168、decode batch 为 32、每层做两次 RMSNorm，且不使用 torch.compile。

forward_native 的表达式会被 eager 模式拆成一组 ATen 算子：转 FP32、平方、按行求均值、加 epsilon、开方取倒数、乘回、按权重缩放、转回原 dtype。即使编译器能消掉一部分，实际执行的 kernel 数量仍在 6 到 8 个之间。vllm_c 的实现是单个内核，一次发射完成全部计算。

按每层节省 6 次发射、61 层、两次 RMSNorm 计算，一个 decode step 少掉的 kernel 发射约为：

单次小内核的端到端间隔按 1.5 到 4 微秒估算，对应的墙钟差异在 1.1 到 2.9 毫秒之间。这个数字对 decode 来说并不小：一个 32 并发、每步输出一个 token 的服务，如果单步总耗时是 20 毫秒，2 毫秒就是 10% 的吞吐差异。

但这笔账有两个前提。第一，发射开销能否被隐藏。GPU 队列很深时，CPU 端的发射可以与上一批计算重叠；当 CPU 成为瓶颈，节省的发射次数直接转化为吞吐。第二，batch 很小时计算量极小：32 行、7168 列的 RMSNorm 单程约 458KB 的 BF16 数据，读加写接近 0.9MB，在 3TB/s 带宽下只要约 0.3 微秒，内核时间几乎全被固定开销占据。反过来，prefill 阶段一个 batch 处理上万 token，计算与访存占主导，少几次发射的收益就被摊薄。

这也解释了编译路径的取舍。torch.compile 打开后，forward_native 里的那串算子会被 inductor 融合成一到两个 Triton kernel，发射次数与 vllm_c 相当，还能和相邻的量化、残差加法继续融合。此时继续启用 CustomOp 反而会切断融合，所以 vLLM 默认在 inductor 模式下关闭它。

验证方法比估算更可靠：用 torch.profiler 抓一个 decode step，按 kernel 名统计发射次数与 GPU 时间；或者在 nsys 里看每层之间的空隙。如果空隙主要来自 aten:: 系列小算子，说明当前走的是 native 路径；如果每层只有一个 rms_norm_kernel，说明厂商内核生效；如果看到大量 triton_ 前缀的融合内核，说明编译路径接管了这段计算。

调试与验证

算子分派出问题时，现象往往很分散：模型能跑但慢、某种 dtype 下结果不对、某个平台启动即报 AttributeError。下面几条路径可以快速定位。

先确认平台：

1
python -c "from vllm.platforms import current_platform as p; print(p, p.device_type, p.dispatch_key)"

如果要看平台是在哪里被初始化的，可以读取 vllm.platforms._init_trace，它保存了第一次访问时的调用栈。插件没有生效时，这个栈通常能指出是哪个模块提前触发了初始化。

再确认算子注册：

1
python - <<'PY'
2
from vllm.model_executor.custom_op import op_registry, op_registry_oot
3
print(sorted(op_registry)[:10])
4
print(sorted(op_registry_oot))
5
PY

op_registry_oot 为空，说明 OOT 类替换没有发生；op_registry 里没有目标算子，说明对应模块没有被导入，常见于只导入了模型壳、没有触发 layers/__init__ 的场景。

然后确认编译开关。想在 eager 下强制使用厂商内核，可以把 custom_ops 设成只启用目标算子：

1
vllm serve <model> --compilation-config '{"custom_ops": ["none", "+rms_norm"]}'

反过来，怀疑某个内核有问题时，可以禁用单个算子，让它退回 native：

1
vllm serve <model> --compilation-config '{"custom_ops": ["all", "-rms_norm"]}'

注意 all 与 none 只能出现一个，同时出现 +op 与 -op 会直接报错。

常见的报错与含义：

• Failed to import from vllm._C_stable_libtorch：扩展没有编译、ABI 与当前 PyTorch 不匹配，或动态库依赖缺失。CUDA 平台导入的是 stable ABI 版本，ROCm 构建会走另一条源码编译路径。
• module 'torch.ops' has no attribute '_C'：扩展没有被提前导入。Platform.import_kernels() 的职责就是在模型构建前完成这件事；自定义启动脚本如果绕过了 vLLM 的初始化流程，容易出现这个问题。
• Duplicate op name 或 Duplicate layer name：同一个注册名被注册两次，常见于插件被重复加载，或两个 OOT 包注册了同名类。
• does not match native schema：IR 实现签名与语义定义不一致，多一个参数、少一个默认值都会触发。
• Priority set incorrectly: the last implementation must support all args：优先级列表里最后一个实现也不能保证支持当前参数，且没有 native 兜底。修复方式是把支持面最广的实现放到末尾，或显式加入 native。
• Implementation ... not supported：优先级列表里包含了一个静态 supported=False 的实现。_filter_priority_impls 通常会在安装时过滤掉它，但手工 set_priority 或平台默认值写错时仍可能遇到。

IR 还提供了一个对内核开发很有用的测试设施：register_input_generator 与 override_tolerance。每个 IR op 可以注册一个按 num_tokens、hidden_size、dtype 生成随机输入的生成器，并覆盖特定 dtype 的容差。rms_norm 就把 FP16 的容差放宽到 atol=1e-2, rtol=2e-3，因为大形状下的归约会累积舍入误差，默认容差会让几百万个元素里的少数几个误报。这套设施让“新 provider 是否与 native 语义一致”可以自动化验证，而不必靠人眼比对。

除了这些运行时手段，还有几条编码层面的检查值得在接入新算子时固化下来。

第一，provider 的 supports_args 必须能在 fake tensor 上求值。它经常会被写成 x.dtype == torch.float8_e4m3fn and x.is_contiguous() 这类判断，这没有问题；但如果写成读取张量元素、调用 .item()、依赖设备实际可用的动态查询，eager 和编译两条路径就会给出不同的结论，甚至在 lowering 阶段直接失败。

第二，supported 只能在注册时求值一次，不要在里面放随请求变化的判断。平台是否支持、第三方库是否可导入，属于静态信息，适合放在 supported；dtype、形状、是否使用 variance_size 覆盖，属于动态信息，适合放在 supports_args。两者混用会让优先级列表的过滤结果与实际分派不一致。

第三，注册 provider 的模块必须保证在模型构造之前被导入。如果厂商把 provider 放在一个懒加载模块里，import_ir_kernels 的覆写实现必须显式导入它。否则会出现“代码写了、单独测试能过、服务启动后永远走 native”的情况，而且日志里通常没有提示，因为 native 永远是兜底。

第四，OOT 类替换与 IR provider 不要同时用于同一个算子。两种机制都能改变执行路径，叠加使用会让排查变难：op_registry_oot 决定用哪个类，IR 优先级决定类内部走哪个 kernel，两层都要配置正确才能得到预期效果。厂商通常选择其中一种作为主路径，另一种只用于灰度对照。

第五，数值对账要覆盖所有可能被选中的 provider。不同实现的归约顺序、累加精度与向量化宽度不同，override_tolerance 必须按 provider 与 dtype 分别设置。它经常会被写成 x.dtype == torch.float8_e4m3fn and x.is_contiguous() 这类判断，这没有问题；但如果写成读取张量元素、调用 .item()、依赖全局 CUDA 状态的逻辑，eager 和编译两条路径就会给出不同的结果，甚至在 lowering 阶段直接失败。

第二，supported 只能在注册时求值一次，不要在里面放随请求变化的判断。平台是否支持、第三方库是否可导入，属于静态信息，适合放在 supported；dtype、形状、是否使用 variance_size 覆盖，属于动态信息，适合放在 supports_args。两者混用会让优先级列表的过滤结果与实际分派不一致。

第三，注册 provider 的模块必须保证在模型构造之前被导入。如果厂商把 provider 放在一个懒加载模块里，import_ir_kernels 的覆写实现必须显式导入它。否则会出现“代码写了、测试单独跑能过、服务启动后永远走 native”的情况，而且日志里通常没有任何提示，因为 native 永远是兜底。

第四，OOT 类替换与 IR provider 不要同时用于同一个算子。两种机制都能改变执行路径，叠加使用会让排查变得困难：op_registry_oot 决定用哪个类，IR 优先级决定类内部走哪个 kernel，两层都要配置正确才能得到预期效果。厂商通常选择其中一种作为主路径，另一种只用于灰度对照。

第五，数值对账要覆盖所有会选中的 provider。不同实现的归约顺序、累加精度与向量化宽度不同，override_tolerance 必须按 provider 与 dtype 分别设置。rms_norm 对 FP16 放宽容差就是例子：在大形状下，native 的 mean 归约与 vllm_c 的块内归约会产生不同的舍入累积，默认容差会让少数元素误报。

最后，VLLM_BATCH_INVARIANT 环境变量会改变 forward_cuda 与内核发射参数。开启后，RMSNorm 走 rms_norm_batch_invariant，内核固定 1024 线程块，保证同一行在不同 batch 大小下逐位一致。做数值对账或复现实验时，先确认这个开关的状态，否则会看到“同一个请求、不同 batch、结果最后几位不同”的现象。

小结

把这篇的链路压缩成一句话：模型层只声明语义，平台单例决定这套语义在哪个设备上执行，CustomOp 在构造期选出该平台的 Python 方法，IR 在调用期选出该平台的具体内核，最终由 torch.ops 把调用交给 C++ 扩展或厂商算子库。

这条设计里有三个值得带走的判断。

平台选择必须是进程级、延迟初始化、且 OOT 优先的。推理框架的调度、显存预算、注意力后端、量化支持列表都依赖平台信息，一旦不同进程对平台的判断不一致，问题会以非常隐蔽的方式出现。

算子分派的粒度要与编译策略匹配。CustomOp 的类替换适合“整套实现都不同”的硬件，IR 的 provider 优先级适合“同一平台上多个内核竞争”的场景。两者叠加并不冗余，因为一个解决可移植性，一个解决内核选择。

编译开关是算子层的一部分，而不是外部配置。custom_ops 的 all / none 默认值直接改变了 RMSNorm 走厂商内核还是走可融合的原生表达式。读 vLLM 的性能数据时，先确认编译模式与算子开关，再比较不同内核的差异，否则很容易把“编译路径不同”误判为“内核更快”。

站内延伸阅读可以沿着这条链路继续往下钻：注意力后端的平台选择见 vLLM 源码完全拆解（六）：Attention Backend 如何统一 FlashAttention 与 FlashInfer；分页注意力的内核实现见 vLLM 源码完全拆解（七）：PagedAttention Kernel 如何遍历分页 KV 并完成稳定 Softmax；量化层如何参与建壳与后处理见 vLLM 源码完全拆解（二十）：权重装载——safetensors 分片、TP 切分与量化层的参数清单；编译期融合 pass 如何匹配这些算子见 vLLM 源码完全拆解（三十一）：编译期的融合 pass；另一套引擎如何用 Platform 抽象同时支撑 CUDA 与 Ascend，见 DeepInfer 源码拆解（三十三）：Platform 抽象与 CUDA/Ascend 双后端。