Metal / MLX 与 CUDA / Triton：从硬件到 kernel 的对比

① 一张图看懂最大的区别

② 硬件对比

规格

编程模型里能看到的硬件

③ 软件栈对比

④ 硬件差异怎么塑造了软件设计

这是这一页的重点：两边的软件长得不一样，大多不是“谁的设计更好”，而是硬件不一样，最优的设计就不一样。

1. 统一内存 → 数组不属于任何设备

CUDA 程序的一大块工作是管理两块内存：cudaMalloc、cudaMemcpy、tensor.to("cuda")，以及想办法让拷贝和计算重叠。Apple silicon 上 CPU 和 GPU 读的是同一块物理内存，所以 MLX 的设计是：数组不分设备，你选的是在哪个 stream 上计算。

a = mx.random.normal((4096,))           # 不用说放在哪
mx.add(a, a, stream=mx.cpu)             # 在 CPU 上算
mx.add(a, a, stream=mx.gpu)             # 在 GPU 上算，同一块数据，没有拷贝

源码：MLX 分配的每一块内存都是 mlxMTL::ResourceStorageModeShared（CPU 和 GPU 都能直接访问）；文档 mlxunified_memory.rst。

代价：没有独立显存，“模型能不能放下”由整机内存决定；好处也在这里：一台 128 GB 的 Mac 能放下一张 32 GB 显卡放不下的模型（第 ⑥ 节）。

2. 带宽低一个数量级 → 量化是一等公民

LLM decode 每生成一个 token 都要把全部权重读一遍，速度 ≈ 带宽 ÷ 模型字节数（L2 讲过）。M2 的 100 GB/s 只有 RTX 5090 的约 1/18。在 Mac 上提升 decode 速度最直接的办法，就是让每个权重更小。

所以 MLX 把 4-bit / 8-bit 分组量化做成了一等公民：mx.quantize、QuantizedLinear，以及专门的 Metal kernel：decode 用的矩阵乘向量 mlxqmv_fast_impl、prefill 用的 mlxqmm_t_impl，反量化在 mlxdequantize 里和乘加融合。CUDA 这边当然也有量化（M3 的 FP8 block scaling），但数据中心 GPU 上的主要动机是算力（FP8 Tensor Core 快一倍），而 Mac 上的主要动机是带宽。

3. threadgroup 内存小、没有异步拷贝 → 更小的 tile、更简单的流水线

CUDA 上高性能 GEMM / attention 的套路（C4、T4、FlashAttention-3）是：大 tile（128×128 以上）放进上百 KB 的 shared memory，cp.async 或 TMA 异步预取下一块，Hopper 上再用 warp specialization 让一部分 warp 专门搬数据。

Apple GPU 上：

• threadgroup 内存最多 32 KB，tile 只能小得多。X3 的 simdgroup 版本用的是 32×32 的输出 tile 和 16 宽的 K 块。
• 没有公开的异步拷贝指令，“边搬边算”只能靠多个 simdgroup / threadgroup 交错执行来掩盖延迟。
• 所以 MLX 的 GEMM（steel）和 attention（mlxsteel_attention.h、decode 用的 mlxsdpa_vector.h）结构上比 FA3 简单，更接近 FA2 和 C4 的风格。

4. 矩阵单元是 8×8 → 用 simdgroup_matrix 拼 GEMM

X3 讲解里画过：一个 simdgroup 持有 2×2 个 8×8 累加器，做 2×2 的“块外积”。这和 C4 里一个线程持有 4×4 个标量做外积是同一个思想，只是粒度从标量变成了 8×8 的小矩阵。

5. 一个芯片、command buffer → 惰性求值 + 批量提交

CUDA 里每次 kernel launch 有几微秒的固定开销，decode 时大量小 kernel 的开销加起来很可观，所以有 CUDA graph（F1、F3 讲过）。Metal 的工作单位是 command buffer：把很多 kernel 编码进同一个 buffer，一次提交给 GPU。

MLX 用惰性求值配合这一点：你写 c = a + b 时什么都没算，只是记下了计算图；直到 mx.eval(c)（或需要结果时）才把整张图编码成 command buffer 提交。一个 command buffer 里攒多少个 op 有阈值：mlxneeds_commit：超过 max_ops 或 max_mb 才提交。文档：mlxlazy_evaluation.rst。mx.compile 还能把一串逐元素操作融合成一个 kernel（X1 的“测完之后”提到过）。

⑤ 同一个 kernel，三种写法

逐元素：add + GELU（C1 / T1 / X1）

__global__ void k(const float* x, float* y,
                  int64_t n) {
  int64_t i = (int64_t)blockIdx.x * blockDim.x
              + threadIdx.x;
  if (i < n) y[i] = gelu_tanh(x[i]);
}
k<<<(n + 255) / 256, 256, 0, stream>>>(x, y, n);

@triton.jit
def k(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    ...
    tl.store(out_ptr + offs, out, mask=mask)
k[(triton.cdiv(n, 1024),)](x, y, out, n, BLOCK=1024)

// 只写函数体，签名由 MLX 生成
uint i = thread_position_in_grid.x;
if (i < n) {
  out[i] = T(gelu_tanh(float(x[i]) + float(y[i])));
}
// Python：grid 是线程总数，不是 block 数
// kernel(inputs=[x, y, n], grid=(n,1,1),
//        threadgroup=(256,1,1), ...)

• CUDA 和 Metal 的 kernel 体几乎一样：都是“我是第几个线程、算一个元素”。区别在外面：CUDA 的 grid 是 block 数，Metal 的 grid 是线程总数；Metal 的签名由 MLX 生成。
• Triton 换了一个抽象层：你写的是一整块数据，线程怎么分由编译器决定。
• 源码：drillsC1 · drillsT1 · drillsX1

归约：一行求和（C2 / T2 / X2）

Metal 在这一层其实比 CUDA 更方便：simd_sum / simd_max 一个函数就是整个 warp 的归约，结果每个 lane 都有，所以 X2 的 tg_reduce 让每个 simdgroup 都做第二轮，省掉一次广播（X2 讲解第 ③ 节）。源码：drillsC2 warp_reduce_sum · drillsX2 tg_reduce。

GEMM 的内层（C4 / T4 / X3）

源码：drillsC4 外积 · drillsT4 主循环 · drillsX3 simdgroup MMA。

⑥ 性能对比

本仓库的实测：同一类 kernel

两点值得注意：

• 都打到了各自的上限。T1、T2 都到了 5090 峰值的 85% 左右；X1 的 MLX 自带操作也接近那台 M2 当时实测的带宽。所以“慢”不是 kernel 写得差，而是硬件上限低。
• 比例和硬件规格的比例差不多：带宽差约 18 倍（标称），GEMM 差约 25 倍。memory-bound 的 kernel 差多少，基本就是带宽差多少。

roofline 的拐点

一个有意思的结论：FP32 下两台机器的拐点差不多，所以同一个算法在两边是 compute-bound 还是 memory-bound，结论基本一样，只是绝对速度差了一个数量级。真正拉开差距的是 Tensor Core：它把 5090 的拐点推高了三倍多，而 M2 没有对应的东西。

LLM 推理：速度和容量的取舍

decode 速度的上限 ≈ 带宽 ÷ 每个 token 要读的字节数（约等于模型权重大小）。用这个公式估一下（只是上限，不含 KV cache 和其他开销）：

按“带宽 ÷ 权重字节数”估算的上限，实际会低一些。70B 那组的条长是相对 819 GB/s 画的。单张 5090 放不下 40 GB 的权重，需要多卡或把一部分放到 CPU 内存（那就受 PCIe 带宽限制了）。

• decode（生成）：memory-bound。NVIDIA 显卡带宽高，同一个模型快得多；但只要模型放得下，Mac 也能跑出可用的速度。
• 容量：Mac 的统一内存可以到 128–512 GB，一台机器就能放下单张消费级显卡放不下的大模型。这是 Mac 跑本地大模型最大的优势。
• prefill（处理 prompt）：compute-bound，靠矩阵算力。这里差距最大（5090 的 FP16 Tensor Core 实测 218 TFLOP/s，M2 的 FP32 GEMM 2.6 TFLOP/s），所以长 prompt 的首 token 延迟在 Mac 上明显更长。
• 并发服务：CUDA 生态有 paged KV cache、continuous batching、FlashInfer、vLLM / SGLang 这一整套（Track L、F）。Mac 上的工具更偏向单用户本地推理。

⑦ 框架和生态对比

MLX 本身也在往跨平台走：它已经有一个 CUDA 后端（mlxmlx/backend/cuda/，GEMM 走 mlxcuBLAS），同一份 MLX 代码可以跑在 NVIDIA GPU 上。MLX 的“快速”算子清单在 mlxmlx/fast.h：rms_norm、layer_norm、rope、scaled_dot_product_attention，以及你在 X1–X3 用的 metal_kernel。

⑧ 一句话总结与选择建议

核心结论：两边的编程模型几乎一一对应（grid / block / warp / thread），所以你学的 CUDA 思路可以直接用在 Metal 上。真正的区别来自硬件：Apple 用统一内存换来了容量和省去拷贝，但带宽和矩阵算力低一个数量级，所以软件把重点放在量化、惰性求值和简单的 kernel 结构上；NVIDIA 有 Tensor Core、大片上内存和异步拷贝，所以软件把重点放在复杂的流水线、tile 调度和多卡扩展上。

源码与资料索引

本仓库的对照题：C1 / T1 / X1（逐元素）· C2 / T2 / X2（归约）· C4 / T4 / X3（GEMM）· M3（FP8 量化）· L2（decode 是 memory-bound）

MLX：mlxunified_memory.rst · mlxlazy_evaluation.rst · mlxcompile.rst · mlxcommand buffer 提交策略 · mlx共享内存分配 · mlx量化 qmv · mlxsteel BlockMMA · mlxsdpa_vector · mlxMLX CUDA 后端

文档：MLX 文档 · MLX · Custom Metal Kernels · Metal Shading Language Specification · CUDA Programming Guide · Triton 文档