技术穿刺验证项目 - 探索 Apple Silicon ANE + GPU 混合推理的可行性
榨干 Apple Silicon 性能，实现 ANE Prefill + GPU Decode 混合推理：
- ANE (Apple Neural Engine): 负责计算密集的 Prefill 阶段
- GPU (Metal): 负责内存密集的 Decode 阶段
端到端流程已跑通！
Prompt: "The capital of France is"
Output: "Paris, and the capital of Italy is Rome." ✅
Prompt: "2 + 2 equals"
Output: "4. The sum of the first" ✅
性能数据：
- ANE Prefill: 32-138ms
- GPU Decode: 42-45 tok/s
┌─────────────────────────────────────────────────────────────┐
│ 用户输入 │
│ "The capital of France is" │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ Tokenizer (CPU) │
│ HuggingFace Tokenizer → [1, 450, 7483, ...] │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ ANE Prefill (CoreML Model) │
│ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ Embedding Layer │ │
│ │ input_ids (1, 128) → hidden (1, 128, 2048) │ │
│ └──────────────────────────────────────────────────────┘ │
│ │ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ Transformer Layer × 22 │ │
│ │ │ │
│ │ ┌────────────────────────────────────────────────┐ │ │
│ │ │ Pre-LN (RMSNorm) │ │ │
│ │ └────────────────────────────────────────────────┘ │ │
│ │ │ │ │
│ │ ┌────────────────────────────────────────────────┐ │ │
│ │ │ Attention (Q, K, V Projection) │ │ │
│ │ │ Q: (1, 128, 2048) → (1, 128, 32, 64) │ │ │
│ │ │ K: (1, 128, 2048) → (1, 128, 4, 64) ──────┐ │ │
│ │ │ V: (1, 128, 2048) → (1, 128, 4, 64) ──────┤ │ │
│ │ └────────────────────────────────────────────────┘ │ │
│ │ │ │ │
│ │ ┌────────────────────────────────────────────────┐ │ │
│ │ │ MLP (Gate + Up + Down Projection) │ │ │
│ │ └────────────────────────────────────────────────┘ │ │
│ │ │ │ │
│ │ ┌────────────────────────────────────────────────┐ │ │
│ │ │ Residual Connection │ │ │
│ │ └────────────────────────────────────────────────┘ │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ 输出: K_cache (1, 2816, 4, 64), V_cache (1, 2816, 4, 64) │
│ = 22 层 × 128 tokens × 4 heads × 64 dim │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ KV Cache Preparation (GPU - MLX) │
│ │
│ 1. 提取每层的 KV: k_layer = k_raw[:, i*128:(i+1)*128, :, :] │
│ 2. 格式转换: (1, seq, heads, dim) → (1, heads, seq, dim) │
│ 3. 应用 RoPE 到 K cache: │
│ │
│ positions = [0, 1, 2, ..., seq_len-1] │
│ freqs = 1.0 / (θ ^ (i/d)) for i in [0, 2, 4, ...] │
│ angles = positions × freqs │
│ k_rotated = k × cos(angles) - k × sin(angles) │
│ │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ GPU Decode (MLX - Metal) │
│ │
│ for step in range(num_tokens): │
│ │ │
│ ├─ Embedding: token → hidden (1, 1, 2048) │
│ │ │
│ ├─ Transformer Layers: │
│ │ │ │
│ │ ├─ Pre-LN │
│ │ │ │
│ │ ├─ Q Projection: hidden → q (1, 1, 32, 64) │
│ │ │ │
│ │ ├─ Apply RoPE to Q: q = RoPE(q, pos=seq_len+step) │
│ │ │ │
│ │ ├─ K, V Projection: hidden → k, v (1, 1, 4, 64) │
│ │ │ │
│ │ ├─ Apply RoPE to K: k = RoPE(k, pos=seq_len+step) │
│ │ │ │
│ │ ├─ Extend KV Cache: │
│ │ │ k_cache = concat([k_cache, k], axis=seq) │
│ │ │ v_cache = concat([v_cache, v], axis=seq) │
│ │ │ │
│ │ ├─ GQA: expand KV (4 heads → 32 heads) │
│ │ │ │
│ │ ├─ Attention: scores = Q @ K.T / sqrt(64) │
│ │ │ weights = softmax(scores) │
│ │ │ output = weights @ V │
│ │ │ │
│ │ ├─ MLP + Residual │
│ │ │ │
│ ├─ Final LayerNorm │
│ │ │
│ ├─ LM Head: hidden → logits (1, 1, 32000) │
│ │ │
│ └─ Sample: next_token = argmax(logits) │
│ │
│ 输出: generated_tokens [3681, 11, 278, ...] │
│ │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ Tokenizer (CPU) │
│ [3681, 11, 278, ...] → "Paris, and the..." │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 用户输出 │
│ "Paris, and the capital of Italy is Rome." │
└─────────────────────────────────────────────────────────────┘
为什么需要完整 Transformer？
❌ 错误方案：只提取 K/V projection
for layer in layers:
k = layer.k_proj(x) # 只做 projection
v = layer.v_proj(x)
x = x # hidden state 不更新！问题：每层的 hidden state 相同，输出错误。
✅ 正确方案：运行完整 Transformer 层
for layer in layers:
# Pre-LN
x_normed = layer.input_layernorm(x)
# Q, K, V projection
q = layer.q_proj(x_normed)
k = layer.k_proj(x_normed) # 保存到 KV cache
v = layer.v_proj(x_normed) # 保存到 KV cache
# Attention (更新 hidden)
attn_out = attention(q, k, v)
x = x + attn_out
# MLP (更新 hidden)
x = x + mlp(layer.post_layernorm(x))关键技术点：
- ANE 提取原始 K/V（不应用 RoPE）
- GPU Decode 时对 K cache 应用 RoPE
问题：动态形状导致 int op 错误
ERROR - converting 'int' op: TypeError
原因：
batch, seq_len = input_ids.shape # ❌ 生成 int op解决方案：硬编码形状
# 固定形状: (1, 128, num_heads, head_dim)
k.reshape(1, 128, 4, 64) # ✅ 不生成 int opCoreML 输出：
K: (1, 2816, 4, 64) # 22 layers × 128 tokens
MLX 需要：
K: (1, 4, 128, 64) # (batch, heads, seq, dim)
转换：
# 提取单层
k_layer = k_raw[:, layer_idx*128 : (layer_idx+1)*128, :, :]
# 转置
k_mx = k_layer.transpose(0, 2, 1, 3) # (1, 4, 128, 64)RoPE 公式：
f(x, pos) = x × cos(pos × θ^(-2i/d)) - x_rotated × sin(...)
MLX 实现：
def apply_rope_mlx(x, positions, theta=10000.0):
# x: (batch, heads, seq, dim)
# positions: (1, seq)
freqs = 1.0 / (theta ** (arange(0, dim, 2) / dim))
angles = positions[:, :, None] * freqs[None, None, :]
cos = cos(angles)[:, None, :, :] # (1, 1, seq, dim//2)
sin = sin(angles)[:, None, :, :]
x1 = x[..., :dim//2]
x2 = x[..., dim//2:]
return concat([x1*cos - x2*sin, x2*cos + x1*sin], axis=-1)关键：ANE 不应用 RoPE，GPU 在使用 K cache 时应用。
pd-engine/
├── README.md # 本文档
├── requirements.txt # 依赖
├── .gitignore # Git 忽略文件
│
├── convert_to_ane_full_transformer.py # CoreML 转换脚本
├── rope.py # RoPE 实现
├── test_e2e.py # 端到端测试
│
└── model_ane_full.mlpackage/ # ANE 模型（需单独生成）
└── ...
# 创建虚拟环境
python3 -m venv venv
source venv/bin/activate
# 安装依赖
pip install -r requirements.txt# 下载 TinyLlama 模型
huggingface-cli download tinyllama/tinyllama-1.1b-chat-v1.0 \
--local-dir ./tinyllama
# 转换为 MLX 格式
python -m mlx_lm.convert \
--hf-path ./tinyllama \
--mlx-path ./tinyllama-mlx
# 转换为 ANE 格式
python convert_to_ane_full_transformer.pypython test_e2e.py| 阶段 | 时间 | 说明 |
|---|---|---|
| ANE Prefill (6 tokens) | 138ms | 固定处理 128 tokens |
| GPU Prefill (6 tokens) | 0.7ms | MLX 高度优化 |
| GPU Decode | 163 tok/s | 使用 generate_step |
性能测试结论：
| Tokens | ANE Prefill | GPU Prefill | GPU 加速 |
|---|---|---|---|
| 10 | 136ms | 0.83ms | 165x |
| 64 | 36ms | 0.47ms | 77x |
| 128 | 34ms | 0.43ms | 79x |
| 512 | 135ms | 0.46ms | 293x |
结论：GPU Prefill 比 ANE Prefill 快 67-293 倍
# ❌ 错误
batch, seq_len = input_ids.shape # 生成 int op
# ✅ 正确
k.reshape(1, 128, 4, 64) # 硬编码# ❌ 错误：只提取 K/V，不更新 hidden
k = k_proj(x)
v = v_proj(x)
x = x # hidden 不变！
# ✅ 正确：运行完整 attention + MLP
x = x + attention(x)
x = x + mlp(x)# ANE: 不应用 RoPE（提取原始 K）
# GPU: 使用 K cache 时应用 RoPE
k_with_rope = apply_rope(k_cache, positions)- CoreML: ANE 模型转换与推理
- MLX: Apple Silicon GPU 推理框架
- MLX-LM: MLX 大语言模型库
- Transformers: HuggingFace Tokenizer
- PyTorch: 模型加载与转换
经过完整测试，ANE Prefill 在当前实现下不可行：
| 问题 | 说明 |
|---|---|
| 固定开销 | CoreML 每次调用 ~33ms 启动延迟 |
| Padding 浪费 | 固定处理 128 tokens，短 prompt 浪费计算 |
| 多次调用 | 长 prompt 需要多次调用，延迟累积 |
GPU Prefill 更优：
- 时间：0.4-0.8ms
- 速度：12K - 1.1M tok/s
- 比 ANE 快 67-293 倍
建议：放弃 ANE Prefill，专注 GPU Decode 优化。
MIT License
本项目受到以下项目启发：
最后更新: 2026-04-09