DeepInfer-Metal

目的:在 Apple M 系列芯片上跑通 DeepInfer —— DeepSeek 开源的大规模推理框架。

Goal: get DeepInfer, DeepSeek's open-source large-scale inference framework, actually running on Apple M-series chips.

本仓库是 deepseek-ai/DeepInfer 的改造版本， 上游自带的说明保留在 README.upstream.md。Apache-2.0，随上游。

上游目前只支持 NVIDIA SM100 与昇腾 950 NPU。本仓库做 Apple Silicon 适配, 目标是真的把推理跑起来,而不是停在平台脚手架。 Upstream currently supports NVIDIA SM100 and Ascend 950 only. This repository adds Apple Silicon support, aiming for working inference rather than platform scaffolding.

验收标准 / Acceptance criteria

两个硬指标,缺一个都算没做完:

• 精度 Accuracy —— 与参考实现的等价性,不是「模型效果好不好」。判据用仓库自带的 deepinfer/testing.py:assert_close(cos_diff < 1e-4),逐 kernel 对照 tilelang_ops/ 顶层的纯 PyTorch 参考实现。
• 效率 Efficiency —— 吞吐 / 峰值内存要有 profiling 依据,不接受「感觉挺快」。

技术路线 / Approach

走 TileLang 的 Metal 后端生成真 kernel,而不是停在 PyTorch/MPS 算子回退。

DeepInfer 的 tilelang_ops/ 分两层:顶层是「接口 + 纯 PyTorch 参考实现 + 测试」, cuda/ 与 asc/ 是纯 kernel 目录,两者各导出 43 个一一对应的 *_tl 工厂。 新增后端的边界因此很清晰 —— 填一个 mps/ 目录。

关键前提:requirements.txt 钉死的 tilelang==0.1.15 已经自带完整的 Metal 后端 (codegen、独立 pass pipeline、metal_fragment_to_simdgroup、GEMM op; 执行走 torch.mps.compile_shader,无需完整 Xcode)。

现状 / Status

⚠️ 尚未跑通。 当前处于验证与搭基座阶段。

已跑通 / Working

第一个 TileLang Metal kernel 已在 mps/ 后端跑通：expand_swa_index （mps/swa_index.py）。

对照上游顶层自带的 expand_swa_index_ref，覆盖其测试的全部 7 个形状：

n=1      width=128   PASS      n=129    width=127   PASS
n=31     width=128   PASS      n=65     width=1025  PASS
n=33     width=128   PASS      n=257    width=4096  PASS
n=65536  width=128   PASS

带 4 组对照以排除假阳性：污染期望值必须判错、扰动输入输出必须跟着变、 预填哨兵必须被完全覆盖且不越界、输入张量必须不被改写。复现：

DEEPINFER_ACCELERATOR=apple python doc/metal/verify_expand_swa_index.py

已验证 / Verified

实测环境:Apple M1 Pro · 16 GB · macOS 27.0.1 · torch 2.13.0 · tilelang 0.1.15 · 仅 Command Line Tools(未装完整 Xcode)

项	结果
torch.mps.compile_shader	✅ 可用
TileLang target 自动探测	✅ metal(detectors: cuda / hip / ascend / metal)
TileLang Metal codegen + 执行	✅ elementwise 与 GEMM 均编译通过并跑出正确结果
生成的 kernel 是否用 simdgroup	✅ 含 simdgroup_multiply_accumulate
deepinfer._C 在 macOS 编译	✅ Apple clang 21,13 个 .o → Mach-O arm64 bundle
PLATFORM 解析	✅ PlatformMPS,_C.engine 导出 34 个符号
上游 #125 自带的 apple 测试	✅ 13 passed

GEMM 数值(M=N=K=256,对照 torch fp32 matmul):

A/B dtype	accum + 输出 dtype	结果
float16	float32	✅ rel_err = 0.00e+00
bfloat16	float32	✅ rel_err = 0.00e+00
float32	float32	✅ rel_err = 0.00e+00
float16	float16	✅ rel_err = 5.45e-03
bfloat16	bfloat16	❌ make_filled_simdgroup_matrix 无 bf16 重载

当前阻塞 / Current blockers

TileLang 0.1.15 的 Metal 执行路径有两个缺陷挡在 kernel 移植前面(均已实测复现):

• A · buffer 绑定按位置转发,但 codegen 可能重排 —— MetalKernelAdapter 假定生成的 Metal kernel 的 [[buffer(N)]] 顺序等于 T.prim_func 签名顺序,实际不保证。 不一致时 kernel 写进输入张量、从输出张量读,不报错,静默返回错误结果并踩踏相邻显存。

  kernel	签名	生成的 buffer 顺序
  GEMM	A, B, C	0=A, 1=B, 2=C ✅
  elementwise	p, o	0=o, 1=p ❌

  已验证可修(按 buffer_map 名字重排实参)。

• B · 动态 shape 不支持 —— adapter 在编译期读 thread_extent,T.dynamic() 的 grid 保持符号态(grid_info == [(n + 3) // 4, 1, 1]),启动时抛 pybind cast 错误。 这是关键路径:DeepInfer cuda/ 的 29 个 kernel 文件里 28 个用了 T.dynamic,共 162 处。

两条均已查重,tile-ai/tilelang 无同类 issue,准备上报。

其他已知限制 / Other limitations

1. accum dtype 必须等于输出 buffer 的 dtype —— epilogue 的 simdgroup_store 不插类型转换,fp16 输入 + fp32 accum + fp16 输出会编译失败 (deduced conflicting types for parameter 'T' ('float' vs. 'half'))。
2. bfloat16 累加器不支持 —— make_filled_simdgroup_matrix 无 bf16 重载。 bf16 输入 + fp32 累加正常。
3. MetalKernelAdapter 不支持 out_idx —— Metal kernel 原地写 buffer 并返回 None,调用方须自行分配输出张量。

Apple GPU 侧硬约束(TileLang 探测值,M1 Pro):

约束	Metal	CUDA 参照
max_threads_per_block	256	1024
max_shared_memory_per_block	32 KB	48–228 KB
thread_warp_size	16	32
Metal 4 cooperative tensor	不可用(需 Apple M5+)	—

现有 cuda/、asc/ kernel 的 tile 尺寸多半超出前两项上限,移植时需重新切分。

构建 / Build

uv venv --python 3.12 .venv && source .venv/bin/activate
DEEPINFER_ACCELERATOR=apple ./build.sh

build.sh 在 apple 分支下会跳过 vendor 子模块（无 MPS 版本）， 只编 deepinfer._C（Apple clang，13 个编译单元）并链进 deepinfer/_so/apple-arm64/。

验证：

DEEPINFER_ACCELERATOR=apple python -m pytest tests/test_hardware_apple.py tests/test_platform_mps.py -q

tests/ 全量目前 35 passed / 7 failed / 7 errors —— 失败项全部需要真实 checkpoint， 错误项全部是缺少 CUDA vendor 包（flash_mla、deep_ep），均非平台问题。

复现脚本 / Reproducers

脚本	内容
doc/metal/verify_expand_swa_index.py	第一个 Metal kernel 的正确性 + 4 组对照
doc/metal/repro_buffer_order.py	缺陷 A 的最小复现与候选修复
doc/metal/repro_gemm_dtype_matrix.py	GEMM 的 dtype 组合矩阵

与上游的关系 / Relationship to upstream

本仓库基于上游 PR deepseek-ai/DeepInfer#125 ("add Apple silicon support", by @Eddy0623)的平台脚手架: accelerator kind apple、PlatformMPS、macOS/clang 构建、C++ 可移植性修复。 其路线图见该 PR 内的 doc/apple-silicon.md。

本仓库做的是该路线图中 P3(原生 Metal kernel)的部分,并带回一处修正:P3 原文写的是 「TileLang 后端如支持 Metal 目标」—— 实测 0.1.15 已经自带 Metal 后端,codegen 可用且 GEMM 数值精确;真正挡路的不是「支不支持」,而是上面 A、B 两个执行路径的缺陷。

上游 TileLang 仓库 tile-ai/tilelang 的 Metal 支持 在路线图 #56 上且在活跃推进。

许可 / License

Apache-2.0,随上游。