Deep Study: Apple CoreAI-Torch compare with vllm

问题

融合算子的本质矛盾：硬件专业化 vs 软件通用化

简单算子（add, mul, matmul）：各硬件都支持，抽象层能统一。

融合算子（flash attention, MoE gate, RMSNorm+quantize）：必须针对硬件写 kernel，因为：

内存访问模式决定了性能（shared memory / register 分配）
指令集差异巨大（tensor core / WGMMA / TMA）
线程块配置需要手工调优

矛盾：模型代码希望”写一次，到处跑”，但融合算子必须”一个硬件写一套”。

硬件碎片化导致 Dispatch 爆炸

GPU 架构	FP8	Tensor Core	特殊指令
H100 (Hopper)	✅	4th gen	WGMMA, TMA
A100 (Ampere)	❌	3rd gen	WMMA
RTX4090 (Ada)	❌	4th gen	WMMA
Apple M4	❌ (Metal 模拟)	ANE 专用	Metal 线程组

一个融合算子要支持 4 种硬件 → 需要 4 套 kernel + 运行时检测 + fallback 逻辑。

融合算子越多，硬件越多 → 组合爆炸。

vLLM 的困境（开放生态）

运行时 dispatch:
  检测 GPU 型号
  ├── H100 → 用 FP8 fused kernel
  ├── A100 → fallback 到 FP16 kernel
  └── RTX4090 → fallback 到另一套 kernel
  
问题：
  1. 代码复杂度随硬件数量线性增长
  2. 每个新硬件都要重写 kernel + 更新 dispatch 逻辑
  3. fallback 路径性能差，但必须支持
  4. CUDA Graph 捕获时必须固定 dispatch 路径，无法动态切换

Core AI 的解法（封闭生态）

编译时特化:
  to_coreai() → MLIR（不绑定硬件）
  executable() → 针对当前设备编译（知道是 M4 还是 M2）
  
优势：
  1. 不需要运行时 dispatch，一条路径走到底
  2. .aimodel 可以在任何 Apple 设备上加载，加载时编译
  3. 硬件数量有限（M1/M2/M3/M4 + ANE），可控
  4. 没有 fallback 路径，每个设备都有唯一最优实现

本质差异

两者通用性都差，差异在性能和可维护性的权衡：

维度	vLLM	Core AI
性能	最优（针对特定硬件手写/调优 kernel）	高层最优，底层受限（MLIR 图由编译器适配；手写 Metal kernel 则与 CUDA 相同）
可维护性	较差（多套算子实现或多硬件适配）	好（高层算子一套代码；手写 kernel 时维护成本相同）
通用性	好（官方支持 NVIDIA, AMD, Intel, TPU, Neuron, CPU）	差（只支持 Apple 生态）

融合算子的真正本质是内存层级优化

融合不是”把两个算子连起来”，而是”让中间结果不写回全局内存”。

以 Flash Attention 为例：

不融合：Q×K → 写回全局内存 → Softmax → 写回 → ×V → 写回
融合：Q×K → 存在 shared memory → Softmax → 存在 shared memory → ×V → 只写最终结果

融合的本质收益：减少全局内存访问次数。

不同 GPU 有不同内存层级：

NVIDIA: global memory → L2 → L1/shared memory → register → tensor core
Apple Metal: global memory → L2 → threadgroup memory → register → ANE

一个 fused kernel 的核心逻辑是：

哪些数据放 shared memory / threadgroup memory
每个线程块处理多少数据（tile size）
如何避免 bank conflict

这些全部是硬件相关的。MLIR 可以统一算子语义，但无法统一内存访问策略。

所以 fallback 的本质不是”兼容性差”，而是”内存层级无法抽象”。

vLLM 选择：为每种硬件手写最优 kernel → 性能优先，牺牲可维护性 Core AI 选择：一套代码适配所有 Apple 芯片 → 可维护优先，接受非最优性能

缺陷

vLLM 的缺陷与挑战：

GPU 碎片化：不同架构支持的操作和指令集不同，手写算子适配成本高。
CUDA Graph 的静态限制：Replay 时输入 Tensor 内存地址和 Shape 必须固定。虽然 vLLM 通过预设多尺寸 Capture 和内存拷贝克服了动态 Batch 限制，但带来了额外的内存与捕获开销。
冷启动开销：每次服务启动都需要 Warmup + Capture 不同的 Batch Size。
调试链路长（仅在编译模式下）：若启用 torch.compile 后端，链路涉及模型代码、FX Graph Pass、Inductor 编译和 CUDA Graph 捕获。而在默认的 Eager 模式下，主要依赖模型代码和 CUDA Graph。
内存缓存不持久化：CUDA Graph 存在内存中，重启后必须重新 Warmup + Capture。

Core AI 的缺陷：

固定 shape：torch.export 导出时 shape 就固定了，不像 vLLM 可以处理动态 token 数
需要写翻译规则：每个自定义算子都要额外实现 lowering 函数（但提供了更多上下文）
运行时编译延迟：executable() 第一次调用时需要编译，有冷启动延迟
封闭生态：只能跑在 Apple 设备上，无法跨平台

Apple 做法

1. get_decomp_table + _COMPOSITE_OPS

定义了一组 op 不分解，效果：模型可以用一些融合算子表达，但是不分解，保证后续能够匹配到这些融合算子，这个方法解决了 vLLM 新做法中”不能采用 compile”的问题。

原理：

def get_decomp_table():
    table = torch.export.default_decompositions()  # PyTorch 默认分解字典
    for op in _COMPOSITE_OPS:
        table.pop(op, None)  # 移除需要保留的算子
    return table

保留的算子：

复合算子（保持语义边界）：hardsigmoid, instance_norm, pixel_shuffle, scaled_dot_product_attention
直接降级（有自定义实现）：hardswish, silu

2. torch.export.export

vLLM 是用 compile 的时候操作 fx graph，复杂度高一些，运行时和 inductor 集成到一起。Apple 是纯粹的离线处理，api 简单。

底层逻辑一样，都需要固定的输入（用于 tracing 决定 shape 和分支）。

3. TorchConverter()

3.1 输入阶段

add_exported_program / add_pytorch_module

接收 ExportedProgram 或 nn.Module
放入 _staged 队列，支持多模型合并

3.2 预处理

inject_subbyte_tensors 标记量化的 tensor 的类型，之后恢复的时候使用。相当于 trace 的时候有的信息（如 nbits）丢失了，通过图模式匹配逆向推导并打补丁。

validate_exported_program 检查两项：

应该分解但没分解：算子在 PyTorch 默认分解字典里，但用户忘了调 run_decompositions()
转换器不支持：算子不在 _aten_to_core_resolver 里，需要用户用 register_torch_lowering() 自定义翻译

3.3 排队

_staged.append 把模型加入队列，to_coreai() 时遍历队列，每个 entry 生成一个 coreai.graph。

3.4 转换核心

to_coreai()

创建空 MLIR Module（容器）
遍历 _staged，为每个 entry 调用 _get_graph_op()
返回 AIProgram._from_mlir_module(module)

_get_graph_op()

清空状态（_clean）
注册权重（_register_constants）：从 graph_signature 提取 parameters/buffers/constants
注册输入输出（_register_io）
创建 coreai.GraphOp（MLIR 函数声明）
遍历 FX Graph 节点，调用 _get_operation(node)
设置返回值和 buffer mutation 属性

_get_operation(node) 分发器：

placeholder → _handle_placeholder_op（处理常量/输入）
call_function → _handle_call_function_op（算子翻译）
get_attr / output → 跳过（在其他地方处理）

_handle_placeholder_op 遍历 node，从 _constants_map / _parameters_map / _buffers_map 取 tensor，调用 _constant_from_tensor 创建常量，存入 _values_map。

_constant_from_tensor 把 PyTorch Tensor 转化成 MLIR ConstantOp（Attribute）：

识别真实类型（future_dtype）
解包 sub-byte 量化数据（tensor.elem）
搬到 CPU 并连续化（detach().cpu().contiguous()）
处理特殊类型（float8/bfloat16 → ml_dtypes，int64/float64 → 降级到 32 位）
创建 DenseElementsAttr（数据序列化进 IR）
返回 coreai.ConstantOp(value=attr).result

结果不是保存地址，而是把数据”焊”进了图的定义里（Attribute）。

_handle_call_function_op 按优先级查找翻译函数（5 种来源）：

外部化子模块（_externalized_lowerings）
用户自定义翻译（_user_defined_torch_lowering）
标准 ATen 算子（_aten_to_core_resolver）—— 90% 走这里
Core AI 自定义算子（_custom_to_core_resolver）
高阶算子（_higher_order_resolver）

找到 lowering function 后：

收集输入：get_operands(values_map, node) → 从 _values_map 找已翻译的 MLIR Value
执行翻译：lowering_fn(values_map, node, location) → 调用 Core AI 算子
存储结果：_values_map[node.name] = result

_values_map 的作用 它是连接 FX Graph (Python) 和 MLIR (编译器) 的桥梁。前一个节点的输出存在 map 里，后一个节点去 map 里取，翻译完再放进去。循环结束后 map 被丢弃，真正的产物是通过 map 连起来的那一整张网（MLIR Graph）。

最终产物

AIProgram

不是 Python 函数，不能直接调用
是”图纸”（数据结构），装着 MLIR 计算图
后续流程：optimize() → save_asset(.aimodel) → executable() → 设备端推理

custom Op

custom torch Op

定义算子：用 torch.library.custom_op 等 API 定义 shape、mutation，和 vLLM 基本一样
模型调用：在模型里调用 torch.ops.my_lib.scaled_add，和 vLLM 基本一样
注册翻译：通过 converter.register_torch_lowering("my_lib::scaled_add") 注册翻译函数，输入 (values_map, node, location)，实现翻译逻辑并返回 MLIR Value

核心差异： vLLM 的自定义 op 只需实现算法本身（CUDA kernel 或 PyTorch 实现）。Core AI 还要实现翻译规则——告诉转换器”看到这个 FX 节点时，用什么 Core AI 算子替换”。

这个设计更高明，因为提供了更多上下文（values_map 可以访问所有已翻译的值，location 用于调试），让翻译逻辑可以做出更智能的决策。

custom metal kernel

创建 TorchMetalKernel：继承 CustomMetalKernel，增加 PyTorch 集成层。包含：名字、参数、Metal 着色器代码（src）、参考 PyTorch 实现（torch_defn 用于 shape inference）
自动验证：输入只能是 torch.Tensor/int/float/bool，返回只能是 Tensor/list[Tensor]/tuple[Tensor, ...]
模型调用：在模型里直接调用，需要传入 threads_per_grid、threads_per_thread_group、result_shapes（自动注入到 torch signature 中）
Converter 注册：通过 converter.register_custom_kernels([kernel]) 自动为每个 kernel 注册 lowering 到 coreai.metal4_kernel，绕过 PyTorch

本质上是完全一模一样的困境： 虽然 Core AI 提供了 Python 的结构化封装，但对于手写 Metal Kernel (MSL) 来说，它面临与 vLLM 手写 CUDA/Triton Kernel 几乎一模一样的硬件碎片化挑战：

调优无法免除：threads_per_thread_group 和 MSL 代码内部的访存与循环划分是与具体芯片规格绑定的。在 M4 上达到最优的线程配置和共享内存（Threadgroup memory）块大小，在 M1 上可能会因硬件资源受限导致低效甚至无法运行。
代码碎片化：要使同一个自定义算子在不同代际（M1/M2/M3/M4）都达到最优，开发者仍需为不同硬件开发多套 Metal Kernel（或包含分支逻辑），这与 vLLM 针对 A100 vs H100 编写多套 CUDA Kernel 在本质上是一致的。

编译与加载/保存

vLLM (CUDA) 运行链路

vLLM 支持两种主要模式：

Eager 模式 (主流/默认)：直接执行手写的 CUDA/Triton Kernel，并在启动时对其进行 CUDA Graph Capture。
Compile + Capture 模式 (较新/优化)：先使用 torch.compile (Inductor 后端) 对模型分片（非 Attention 子图）进行编译优化和算子融合，然后再在编译产物的基础上进行 CUDA Graph Capture。

1. Model 代码
      │
      ├──────────────────────────────┐
      │ (默认 Eager 模式)            │ (Compile 模式)
      ▼                              ▼  torch.compile()
   PyTorch Eager 执行             Inductor 编译子图 → PTX/cubin
   (调用 custom CUDA/Triton)         │
      │                              ▼  缓存在磁盘
      │                           运行时加载编译缓存
      │                              │
      └──────────────┬───────────────┘
                     ▼  Warmup 阶段
               CUDA Graph Capture
                     │
                     ▼  推理阶段
               CUDAGraph.replay()

关键点：

Eager 模式 不涉及 Inductor 编译，只捕获原生 PyTorch + Custom Kernels。
Compile 模式 通过 torch.compile 对可编译子图进行融合以提升性能，但仍需在外面包裹 CUDA Graph 以消除运行时 Launch 开销。
CUDA Graph 依然是运行时捕获并存在内存中，重启后丢失，需要重新 Warmup 和 Capture。
cudagraph_capture_sizes 预设多个 batch size 分别捕获。

Core AI (Metal) 完整链路

1. Model 代码
      │
      ▼  to_coreai()
2. 生成 MLIR → coreai.metal4_kernel { src = "Metal代码字符串" }
      │
      ▼  save_asset()
3. 序列化到 .aimodel 目录 (metadata.json + main.mlirb + hash)
      │
      ▼  运行时 executable()
4. Core AI 编译器读取 .mlirb → 编译 Metal kernel 为设备代码
      │
      ▼  function(inputs)
5. 在 GPU/ANE 上执行

关键点：

Metal 代码在 to_coreai() 时只是字符串，没编译
真正的编译发生在运行时 executable() 调用时
一切都在 .aimodel 文件里，持久化，跨重启可用
不需要 warmup/capture，加载即可用

内存数据处理

vLLM (CUDA)

权重：常驻 GPU VRAM，不序列化
CUDA Graph 限制：replay 时输入 tensor 的内存地址必须与捕获时完全一致
```
# cuda_graph.py line 347
assert new_input_addresses == entry.input_addresses
```
内存管理：用 graph_pool 管理捕获时内存，weak_ref_tensors 减少占用
KV cache：预分配在 GPU 内存，不随模型序列化
输入数据：运行时直接传 GPU tensor，无需额外传输

Core AI (Metal)

权重：序列化在 .aimodel 的 DenseElementsAttr 里（数据焊进图定义）
无地址限制：每次加载都重新编译，不要求内存地址一致
输入数据：通过 NDArray(data, backing=StorageKind.METAL) 传输到 GPU
模型状态：完全在 .aimodel 文件中，可分发、可加载到任何兼容设备

vLLM 可以借鉴 Core AI 的内容

1. 显式翻译规则注册（lowering function 模式）

Core AI 做法：

@converter.register_torch_lowering("my_lib::fused_op")
def lowering(values_map, node, location) -> Value:
    # values_map 提供完整的上游上下文
    # location 提供精确的调试信息
    ...

vLLM 现状：

子图匹配替换（FX Graph pass）更隐式，依赖 pattern 匹配
调试困难：匹配失败时难以定位

借鉴方案： 引入显式的算子翻译规则注册机制，替代部分隐式子图匹配。翻译函数可以访问完整的 _values_map 上下文和调试位置信息，降低维护成本。

2. 结构化 kernel 封装

Core AI 做法： TorchMetalKernel 统一封装了：

Metal 着色器代码（src）
线程配置（threads_per_grid, threads_per_threadgroup）
参考 Python 实现（torch_defn 用于 shape inference）
自动验证（输入类型、返回类型检查）

vLLM 现状：

Triton kernel 和 CUDA custom op 相对松散
配置分散在多处，缺少统一验证

借鉴方案： 提供统一的 kernel 注册接口，将 kernel 代码、线程配置、shape 推导函数、输入验证封装在一起。降低新硬件适配时的维护成本。

3. 编译时硬件特化替代运行时 dispatch

Core AI 做法：

to_coreai()    → MLIR（不绑定硬件）
executable()   → 针对当前设备编译（知道是 M4 还是 M2）

编译时知道目标硬件特性，直接生成最优代码，没有运行时分支。

vLLM 现状：

运行时检测 GPU 型号 → 选 kernel → fallback
CUDA Graph 捕获时必须固定 dispatch 路径

借鉴方案： 服务启动 warmup 阶段，根据当前 GPU 型号生成特化图，编译时确定所有 kernel 选择，避免运行时 dispatch。类似 Core AI 的 executable() 机制。

4. 图翻译而非图修改的思路

Core AI 做法：

FX Graph → TorchConverter → MLIR

前端解析和后端优化清晰分离，TorchConverter 只做翻译，不做原地修改。

vLLM 现状： 原地修改 FX Graph，pass 之间耦合紧密，限制多。

借鉴方案： 引入中间表示层，将图解析（FX → IR）和图优化（IR → 编译产物）分离。每个 pass 对 IR 进行操作，而不是直接修改 FX Graph。降低 pass 之间的耦合度，提高可维护性。