中文 | English

10. DeepSeek 推理：跨平台内核基准套件

概述：Softmax 与 GEMM 是有用的微基准，但只有真实推理工作负载才能诚实评判一个内核工具链。我们将完成一次 DeepSeek-R1-Distill-Qwen-1.5B 解码所需的 13 个内核打包成可移植套件，将 Rust 源码经 mlir_to_msl 编译，并在 Apple 芯片上测量结果。生成的 Metal 内核在 M2 Max 上达到 91.7 tok/s（占 400 GB/s 内存带宽上限的 60%），在 M4 上达到 33–35 tok/s，在解码上击败了 Apple 手工调优的 MLX 运行时。同一份 Rust 源码可面向另外九个后端;本章记录该套件,以便在任一后端上复现。

10.1 为什么选 DeepSeek？

DeepSeek-R1-Distill-Qwen-1.5B 既小到可以装入 8 GB 统一内存，又大到在每一种现实加速器上都受带宽限制，且架构上代表了现代 transformer 家族：

• 分组查询注意力 (GQA) —— 12 个 Q 头共享 2 个 KV 头。
• SwiGLU MLP —— 每层 3 个 matmul，可融合为一个 kernel。
• RMSNorm —— 全面取代 LayerNorm。
• 旋转位置编码 (RoPE) —— 对 Q 与 K 原地应用。

每生成一个 token，解码会跨 28 层读取约 2.6 GB 权重。这使其成为一个 带宽 基准，而不是 FLOPs 基准。硬件上限为 带宽 ÷ 每 token 字节数：

任何 kernel 达到该数字的 60%，就已与手工调优的生产代码相当;达到 80% 是访存受限 kernel 的目标线。

10.2 13 个内核组成的套件

解码模式下完整的一个 transformer 层归结为 8 次 dispatch，加上 5 个模型级 kernel（embedding、两种 RMSNorm 调用点、RoPE、argmax）。完整列表（按 1.5B 模型的形状：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);
}
}

同一份源码可编译到全部 10 个 mlir_to_<target> 后端。每个目标的参考 kernel 已在 benchmarks/deepseek_tile_kernels/templates/<target>/ 提交。

10.3 Apple M2 Max —— 头条结果

硬件：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 使用单个 Metal command buffer。重复惩罚 1.3，温度 0.0（贪婪解码）。

经由 rustc_codegen_mlir → mlir_to_msl 生成的 Rust 源 kernel 在解码上击败了 Apple 手工调优的 MLX。解码是典型推理会话的主要成本（一个 prompt，数百个生成 token），因此这是对最终用户延迟最重要的指标。

91.7 是怎样达到的

M2 Max 上的优化轮次（每一步相对前一步测量）：

另有两个尝试经测量后被 回滚，因为它们造成了倒退：

教训记录在 crates/deepseek_metal/templates/ 与优化日志中;简短的总结是：Apple GPU 的 L1/L2 已经能很好地缓存被复用的激活，所以手工 threadgroup 缓存只有当 (a) 数据装不下 cache 且 (b) 单线程计算量足以摊销 barrier 时才有意义。对于 K = 1536 (6 KB) 的解码 matvec，两个条件都不成立。

10.4 Apple M4 —— 小内存机型结果

硬件：Apple M4，4 性能 + 6 能效核心 CPU，10 核 GPU，120 GB/s 内存带宽，macOS 14.5。

M4 的结果对解码佐证了 M2 Max 的故事：codegen 路径优于 MLX (33–35 vs 32)。Prefill 则是另一种情况 —— MLX 使用 Apple 的 simdgroup_matrix_multiply 原语，非常契合 prefill 的计算受限特性（大 matmul，M ≫ 1）。ascend-rs 的 prefill 路径使用 tiled matmul，达到 9.3 tok/s;弥合 prefill 差距已列入下一轮迭代（templates/matmul_simd.metal 是进行中的替代实现）。

10.5 时间花在哪里 —— 单 kernel 拆解

M2 Max 上单个解码 token（28 层 × 8 dispatch + 5 个模型级 dispatch = 229 次内核启动）：

7 个 matvec/MLP kernel —— 即 §10.2 中的第 3、4、8、9、10 项 —— 占解码时间的 86%。优化投入回报最高的是这些 kernel，这也是为什么 §10.3 列出的所有提升都瞄准了 matvec / attention 路径。Norm 与 RoPE 加起来每 token 不到 1 ms;像我们尝试过的那样去融合它们既无可量化的带宽节省，又增加了计算。

10.6 跨厂商进度

crates/deepseek_metal/src/tile_kernels.rs 这同一份 Rust 源码是全部 10 个 codegen 后端的输入。截至本文成稿：

“可编译” 指 kernel 通过 mlir_to_<target> 后被对应厂商的编译器接受;“端到端运行” 指其在真实硬件上对照已知正确参考能产生正确 logits。

“待运行” 一栏的多寡并不反映各后端距离完成有多远 —— 它反映的是我们在每台设备上分配给驱动测试套件的硬件时间。10 个后端的 codegen 表面均已完整，并由 crates/mlir_to_<target>_tests/ 单元测试覆盖。

10.7 复现 Apple 上的结果

# 克隆公开的产物 + 基准仓库。
git clone https://github.com/yijunyu/ascend-rs
cd ascend-rs

# 在装有 Xcode command-line tools，以及环境变量中带 Hugging Face token 的 Mac 上：
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

两次运行使用相同的模型权重和相同的 prompt;唯一不同的是 kernel 实现。

10.8 为什么选套件，而不是单 kernel

单 kernel 基准（孤立的 softmax、GEMM、RMSNorm）有助于诊断特定瓶颈，但它们会系统性地高估那些 无法组合 的优化的价值：

• “缓存激活” 在独立 matvec 基准上是明显的胜利，在 transformer 层内部却是明显的失败 —— 因为上一个 matvec 已经把 cache 预热了。
• “把 RMSNorm 融合进下一个 matvec” 在融合 kernel 微基准上是胜利，在真实层内部却是失败 —— 同一个 norm 输出会被 Q、K、V 三个 matvec 消费。
• 一个忽略 KV 缓存的 “快 attention” kernel 是无关紧要的;在解码中，KV 缓存 就是 attention 的输入。

绑定到真实模型的 13 个 kernel 套件，是能捕捉这些错误的最小基准。它也使各厂商可以诚实地比较后端：10 个后端看到的是同一份 Rust 源码、同样的形状、同样的内存通信预算。

10.9 关键结论

1. Rust 到 Metal 的 codegen 路径在解码上追平或超过手工调优的 MLX。 M2 Max 上 91.7 tok/s（vs MLX ≈ 88），M4 上 33–35 tok/s（vs MLX 32），证明内存安全的 kernel 工具链不必为交互推理最关键的解码路径让出性能。

2. 解码受带宽限制；该套件达到峰值的 60%。 剩下的 40% 分别来自 dispatch 开销（每 token 约 229 次启动）以及尚未使用 Apple simdgroup_matrix_multiply 原语的 matmul kernel。两者都有已知的修复方向。

3. 微基准会就完整流水线性能撒谎。 两个在孤立测量中显得是赢家（缓存、融合）的优化，让完整解码路径退化了 5–13 tok/s。套件级测量是捕捉这类错误的唯一方法。

4. 一份 Rust 源码，十个后端。 同一份 tile_kernels.rs 经 mlir_to_<target> 可面向 Metal、CUDA、MUSA、AscendC、PTO、NKI、AIE、BANG、Gaudi、SPIR-V 编译。Apple 是首个达到生产保真度并完成端到端测量的后端;其余后端的 codegen 表面已就绪，仅缺硬件运行时间。