PD-Engine: ANE Prefill + GPU Decode Hybrid Inference

技术穿刺验证项目 - 探索 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."         │
└─────────────────────────────────────────────────────────────┘

🔧 实现原理

1. ANE Prefill (CoreML)

为什么需要完整 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

2. CoreML 转换陷阱

问题：动态形状导致 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 op

3. KV Cache 格式转换

CoreML 输出：

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)

4. RoPE 实现细节

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 模型（需单独生成）
    └── ...

🚀 快速开始

1. 环境准备

# 创建虚拟环境
python3 -m venv venv
source venv/bin/activate

# 安装依赖
pip install -r requirements.txt

2. 准备模型

# 下载 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.py

3. 运行测试

python 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 倍

🎓 技术教训

1. CoreML 不支持动态形状

# ❌ 错误
batch, seq_len = input_ids.shape  # 生成 int op

# ✅ 正确
k.reshape(1, 128, 4, 64)  # 硬编码

2. Transformer 需要完整计算

# ❌ 错误：只提取 K/V，不更新 hidden
k = k_proj(x)
v = v_proj(x)
x = x  # hidden 不变！

# ✅ 正确：运行完整 attention + MLP
x = x + attention(x)
x = x + mlp(x)

3. RoPE 需要在正确位置应用

# 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 优化。

📄 License

MIT License

🙏 致谢

本项目受到以下项目启发：

ANEMLL - ANE 模型转换参考
MLX-LM - MLX 推理框架
llama.cpp - 混合推理思路

最后更新: 2026-04-09

Navigation Menu

Repository navigation

README.md

README.md

PD-Engine: ANE Prefill + GPU Decode Hybrid Inference

🎯 项目目标

✅ 技术可行性验证

🏗️ 技术架构

🔧 实现原理

1. ANE Prefill (CoreML)

2. CoreML 转换陷阱

3. KV Cache 格式转换

4. RoPE 实现细节

📁 项目结构

🚀 快速开始

1. 环境准备

2. 准备模型

3. 运行测试

📊 性能数据

🎓 技术教训

1. CoreML 不支持动态形状

2. Transformer 需要完整计算

3. RoPE 需要在正确位置应用

🛠️ 技术栈

📝 参考文献

⚠️ 技术穿刺结论

📄 License

🙏 致谢

FilesExpand file tree

README.md

Latest commit

History

README.md

File metadata and controls

PD-Engine: ANE Prefill + GPU Decode Hybrid Inference

🎯 项目目标

✅ 技术可行性验证

🏗️ 技术架构

🔧 实现原理

1. ANE Prefill (CoreML)

2. CoreML 转换陷阱

3. KV Cache 格式转换

4. RoPE 实现细节

📁 项目结构

🚀 快速开始

1. 环境准备

2. 准备模型

3. 运行测试

📊 性能数据

🎓 技术教训

1. CoreML 不支持动态形状

2. Transformer 需要完整计算

3. RoPE 需要在正确位置应用

🛠️ 技术栈

📝 参考文献

⚠️ 技术穿刺结论

📄 License

🙏 致谢