English | 中文版

9. DeepSeek 推理:跨平台内核基准套件

摘要:Softmax 和 GEMM 作为微基准很有用,但一个真实的推理工作负载才是对内核工具链的诚实检验。我们把完成一次完整 DeepSeek-R1-Distill-Qwen-1.5B decode step 所需的 13 个内核打包成一个可移植套件,在四种生产级加速器上测量同一份 Rust 源码。头条结果:Ascend 910B2 上 168.9 tok/s,走联合的 mlir_to_cpp + mlir_to_pto 路径(是 aclnn-only 基线的 2.47×,是 3.7 tok/s CPU 参考的 45.6×)。跨厂商交叉验证:Google TPU v2-8 上 162.9 tok/s(经生成的 Pallas),Apple M2 Max 上 91.7 tok/s(经生成的 Metal,decode 打败 Apple 手工调优的 MLX),NVIDIA T4 上 53.7 tok/s(经生成的 CUDA)——全部来自同一份 13-kernel Rust 源码。本章剩下的部分描述这个套件、各平台的结果、两个不报告端到端 tok/s 的后端(AWS Trainium NKI 与 Vulkan/SPIR-V)及其原因,以及如何复现上面任何一个数字。

9.1 为什么选 DeepSeek?

DeepSeek-R1-Distill-Qwen-1.5B 小到能装进 8 GB 统一内存,大到在每一种现实中的加速器上都是 bandwidth-bound,而且架构上代表了现代 transformer 家族:

• 分组查询注意力(GQA) —— 12 个 Q-head 共享 2 个 KV-head。
• SwiGLU MLP —— 每层三个 matmul,可融合为一个内核。
• RMSNorm —— 处处替代 LayerNorm。
• 旋转位置编码(RoPE) —— 原地作用于 Q 和 K。

每个 token 的 decode 在 28 层上读取约 2.6 GB 权重。这让它成为一个带宽基准,而非 FLOPs 基准。硬件上限是 带宽 ÷ 每 token 字节数:

任何内核达到这个数字的 60% 就能与手工调优的生产代码竞争;达到 80% 是 memory-bound 内核的目标。同一模型上的 CPU 参考吞吐是 3.7 tok/s —— 这是每条加速器路径都必须跨过的地板。

9.2 13-Kernel 套件

decode 模式下的完整 transformer 层归结为 8 次 dispatch,加上 5 个模型级内核(embedding、两个 RMSNorm 变体、RoPE、argmax)。完整列表,对应 1.5B 模型的 shape(D=1536, NH=12, NKV=2, DH=128, INTER=8960, VOCAB=151936):

完整 Rust 源码在 crates/deepseek_metal/src/tile_kernels.rs,用的是安全的 tile.rs view API:

#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn rms_norm_1536(input: *const f32, gamma: *const f32, output: *mut f32) {
    let ctx = unsafe { GmDeviceCtx::new() };
    let in_v   = unsafe { ctx.view::<1, D, f32>(input) };
    let g_v    = unsafe { ctx.view::<1, D, f32>(gamma) };
    let out_v  = unsafe { ctx.view_mut::<1, D, f32>(output) };

    let x      = tile_load_view_f32(&in_v);
    let g      = tile_load_view_f32(&g_v);
    let normed = safe::tile_rms_norm_f32::<1, D>(x, 1e-6);
    let out    = safe::tile_mul_f32::<1, D>(normed, g);
    tile_store_view_f32(&out_v, out);
}
}

同一份源码编译到每一个 mlir_to_<target> 后端。各目标的参考内核签入在 benchmarks/deepseek_tile_kernels/templates/<target>/ 下。

9.3 Ascend 910B2 —— 头条结果

硬件:Huawei Ascend 910B2,CANN 8.5.0,bisheng 编译器,联合 mlir_to_cpp + mlir_to_pto codegen 路径。

设置:28 层 DeepSeek-R1-Distill-Qwen-1.5B,f16 权重,每次 forward pass 单条 ACL stream。decode 路径用 cpp-tile 内核跑 RMSNorm / RoPE / SiLU,用 PTO cube matmul 跑每层的 f16 projection,用 cached-executor 的 aclnnIncreFlashAttention 跑 attention。

168.9 是怎么到的

910B2 上的优化序列,每一步都对前一步测量:

贡献最大的两个自制内核(residual_add_rms_norm cpp-tile 融合版,以及 f16 PTO matmul 的 blocking)都由 rustc_codegen_mlir 从普通 Rust tile-API 源码生成 —— 没有手写 AscendC。逐算子计时见附录 I。

同一份二进制在 910C 上

同一份构建产物在 Ascend 910C(cube-only)上重新构建,ptoas 的 --cce-fatobj-link 路径负责 matmul 侧。在 910C 上的切分是 每层时间的 98.4% 在 NPU、1.6% 在 CPU —— 仍留在 host 的唯一内核是 RMSNorm,因为 910C 的 cube 单元对它没有加速(它是 memory-bound,DMA 拷贝占主导)。910C 的端到端 tok/s 暂不报告,等待 28 层在稳定的 910C 芯片分配上做更长的正确性验证。

9.4 Google TPU v2-8(Colab)—— 162.9 tok/s

硬件:Google Colab v2-8(Cloud TPU,8 核 × 8 MiB MXU,600 GB/s HBM),mlir_to_tpu codegen 发射 JAX Pallas。

设置:rms_norm 和 rope_inplace 走生成的 Pallas kernel;GQA attention 走生成的 Pallas;matvec 按内存层级切分 —— Pallas 跑 q/k/v/o projection(shape 小,VMEM 友好),XLA jnp.dot 跑 gate/up/down/lm_head(shape 大,受益于 XLA 的 HBM staging)。

在所有逐 op 对照测量上取平均,生成的 Pallas kernel 达到原生 JAX 基线的 0.98×。端到端做了 greedy-token 一致性验证:16 个生成 token 全部逐字节匹配 HuggingFace 参考实现。TPU 结果是整个套件里最重要的跨厂商交叉验证:它表明一个完全没有 C++ 出口的后端(Pallas 从 Python DSL 直接进 XLA)能从同一份 Rust 源码(这份源码本来是给 AscendC 写的)产出有竞争力的结果。

9.5 Apple M2 Max —— 91.7 tok/s(打败手工调优的 MLX)

硬件:Apple M2 Max,12 核 CPU,38 核 GPU,400 GB/s 统一内存带宽,macOS 14.5,Metal 3.1。

设置:28 层 DeepSeek-R1-Distill-Qwen-1.5B,bf16 权重直接以 Metal bfloat 上传到 GPU。每次 forward pass 单个 Metal command buffer。Repetition penalty 1.3,temperature 0.0(greedy)。

经过 rustc_codegen_mlir → mlir_to_msl 后,从 Rust 源码生成的内核在 decode 上超过了 Apple 手工调优的 MLX。在典型的推理会话里(一个 prompt,几百个生成 token),decode 是主导成本,所以这个数字对终端用户延迟最关键。

Apple M4(4P+6E CPU,10 核 GPU,120 GB/s):decode 33–35 tok/s vs MLX 32 tok/s —— Metal codegen 路径在这个更小的部分上也打败 MLX,但 prefill(9.3 vs MLX 72)还卡在重写 prefill matmul 使用 simdgroup_matrix_multiply。

91.7 是怎么到的

M2 Max 上的优化轮次(每步对前一步):

两个尝试的优化经过测量被回滚,因为会倒退:

Apple GPU 的 L1/L2 已经缓存了复用的激活,所以手动 threadgroup 缓存只有在(a)数据不在 cache 且(b)每线程计算大到能摊销 barrier 开销时才有用。对 K=1536(6 KB)的 decode matvec 来说两条都不成立。

9.6 NVIDIA Tesla T4(Colab)—— 53.7 tok/s

硬件:Google Colab 上的 NVIDIA Tesla T4,320 GB/s HBM2,CUDA 12.1,mlir_to_gpu codegen 发射 CUDA C,用 nvcc -arch=sm_75 -O3 编译。

设置:生成的 rms_norm_1536、matvec_f16(带 _bias 和 _add 变体覆盖融合情况)以及 GQA attention_decode_gqa 驱动 decode loop;权重加载和 tokenization 用 host 侧 Python 粘合。

53.7 tok/s 是 T4 理论带宽上限的 44%。剩下的 gap 分两块:次优的 matvec tiling(mlir_to_gpu 路径当前是每线程一个元素,没有走 warp-striped)和 matmul_f32 仍然临时走 cuBLAS。两件事都记录在第 13 章 §12.3.1 作为短期的 mlir_to_gpu + cudarc 集成工作。

每 token 内核的对齐情况同 Ascend 结果一致:13 个内核全部编译通过;发射的 .cu 源码是 2,001 行,由同一份 13-kernel tile_kernels.rs 生成。

9.7 时间去哪了 —— 逐内核分解(M2 Max)

M2 Max 上一个 decoded token(28 层 × 8 dispatch + 5 个模型级 dispatch = 229 次 kernel launch):

七个 matvec/MLP 内核 —— 来自 §9.2 套件的第 3、4、8、9、10 项 —— 占 decode 时间的 86%。优化精力花在这些内核上回报最大,这也是为什么 §9.5 列出的每一项优化都瞄准 matvec / attention 路径。Norm 和 RoPE 合起来每 token 不到 1 ms;把它们融合掉(我们试过)省不出可测量的带宽,还要加计算。

9.8 跨厂商状态

这份 13-kernel Rust 源码是每个 mlir_to_<target> 后端的共同输入。当前已测得的端到端状态(数字来自配套论文 Table 2):

NKI(AWS Trainium)。六个发射的内核(rms_norm_1536、matvec_f16 / _bias / _add、gate_up_silu、GQA attention)都编译并运行通过。端到端 tok/s 没有报告,因为 @nki.jit 用的是 eager dispatch 且没有跨调用 kernel 缓存 —— 每个 decoded token 要 370+ 次 kernel dispatch,每次要承担约 10 s 的建立开销。一个编译好的 torch-neuronx 图 wrapper 能把这折叠成一次 graph dispatch;那是后续工作,不是 codegen 缺口。

Vulkan(SPIR-V)。端到端 decode 需要一个暴露 shader-f16 特性的 adapter。我们能用到的、既支持 SPIR-V 又能跑 Colab notebook 的硬件只有 T4,而 Colab 的 T4 在 Vulkan 下只暴露 Mesa llvmpipe(CPU 光栅器)—— 这会让 decode loop 超时。Apple M2 Max 上经 Vulkan 后端跑的 per-kernel softmax 达到 90× CPU 加速(见附录 I)。

对代码树里其余的后端(mlir_to_musa、mlir_to_aie、mlir_to_bang、mlir_to_gaudi、mlir_to_csl、mlir_to_hexagon、mlir_to_linalg),13-kernel 套件都能干净编译;on-device decode 测量仅卡在各个 rig 的硬件时间分配。

9.9 复现结果

Apple M2 Max / M4:

git clone https://github.com/yijunyu/ascend-rs
cd ascend-rs
cargo run --release -p deepseek_metal -- \
    --prompt "The capital of France is" \
    --max-tokens 128

首次运行从 Hugging Face 下载 DeepSeek-R1-Distill-Qwen-1.5B(约 3 GB),缓存在 ~/.cache/huggingface/。后续运行会打印:

Loaded DeepSeek-R1-Distill-Qwen-1.5B on Metal
Prefill: 0.23s (26.1 tok/s)
[generated text]
Generated 128 tokens in 1.40s (91.43 tok/s)

MLX 对照基线:

pip install mlx mlx-lm
python -m mlx_lm.generate \
    --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
    --prompt "The capital of France is" \
    --max-tokens 128

Ascend 910B2(需要 CANN 8.5.0 和硬件访问):

source /usr/local/Ascend/cann-8.5.0/set_env.sh
export ACLRS_SOC_VERSION=Ascend910B2
cargo run --release -p deepseek_e2e -- --max-tokens 128

TPU v2-8(Colab) 和 NVIDIA T4(Colab):notebook 在 benchmarks/deepseek_tile_kernels/notebooks/ 下。每个 notebook 从 repo 拉取生成的 mlir_to_<target> 输出,对同一批 prompt 跑 decode loop。所有可复现运行都以 CSV 形式记录到 pu-rs.org 开放排行榜(截至 2026-04-23 跨所有后端和目标共 3,924 个数据点)。

9.10 为什么是套件,而不是单一内核

单内核基准(独立 softmax、GEMM、RMSNorm)对诊断某个具体瓶颈有用,但它们系统性地高估那些无法组合的优化的价值:

• 缓存激活在单独 matvec 基准里是明显赢,放到 transformer 层内部就明显输 —— 上一个 matvec 已经把 cache 暖了起来(§9.5)。
• 把 RMSNorm 融合进下一个 matvec 在融合内核微基准上赢,放到真实层里就输 —— 同一份 norm 输出被 Q、K、V 三个 matmul 消费。
• 一个忽略 KV cache 的“快 attention“内核毫无意义;decode 里,KV cache 就是 attention 的输入。

一个绑定到真实模型的 13-kernel 套件是能捕捉这些错误的最小基准。它也让厂商能诚实地对比后端:§9.8 里每个后端看到同一份 Rust 源码、同一批 shape、同一个内存流量预算。

9.11 关键要点

1. 一份 Rust 源码,四种生产加速器上端到端测量完成。Ascend 910B2 168.9 tok/s,Google TPU v2-8 162.9,Apple M2 Max 91.7,NVIDIA T4 53.7 —— 全部来自同一份 13-kernel tile_kernels.rs,经过不同的 mlir_to_<target> 后端编译。后端规模从 1,571 行(SPIR-V)到 11,383 行(mlir_to_cpp)不等,所以瞄准一个新厂商是一项有边界的工程,不是研究项目。

2. 910B2 上对 CPU 参考 45.6×,对 aclnn-only 基线 2.47×。Ascend 路径证明了一个安全优先的 Rust 内核工具链不会牺牲性能:头条数字来自编译器生成的内核流水线,而不是手写 AscendC。

3. Metal codegen 路径在 decode 上打败手工调优的 MLX。M2 Max 上 91.7 vs ≈ 88,M4 上 33–35 vs 32。Apple 的工程师是针对 Apple 自己的硬件手工调的 MLX;ascend-rs 从为另一家厂商写的 Rust 源码里产出了有竞争力的结果。

4. TPU Pallas 交叉验证达到原生 JAX 的 0.98×,与 HF 16/16 greedy-token 对齐。最干净的证据,表明 Rust → MLIR → Pallas 路径产出的内核是正确的,而不只是数值上近似。

5. 微基准在整流水线性能上撒谎。两个在孤立测量里显示为赢的优化(缓存、融合)在 M2 Max 的完整 decode 路径上倒退了 5–13 tok/s。套件级测量是唯一能抓到这种情况的办法。