Edge Inference

Edge inference 在用户设备（手机、笔记本电脑、IoT 传感器）上运行模型，无需将数据发送到 cloud。本文涵盖边缘约束、模型压缩管道、设备端运行时、编译器栈、硬件目标（NPU、Neural Engine）、设备端 LLM、federated learning 以及 latency 优化。

• Cloud inference 需要网络连接，增加 latency（往返 50-200 毫秒），每次请求需要付费，并将用户数据发送到第三方服务器。Edge inference 消除了这四个问题：模型在本地运行，响应即时，每次 inference 零成本，数据保持私密。

• 权衡：边缘设备的计算和内存比数据中心 GPU 少 100-1000 倍。在这些约束内运行模型需要在各个层面进行积极优化。

• Cactus（github.com/cactus-compute/cactus）是专为移动和可穿戴设备构建的低 latency AI 引擎。它在生产中演示了本文涵盖的许多技术：用于 attention 和矩阵操作的自定义 ARM SIMD kernel（第 16 章）、KV-cache 量化（第 17 章第 1 文件）、分块预填充、在 Apple 和 Qualcomm 芯片上的 NPU 加速 inference、零拷贝内存映射以降低 10 倍 RAM 使用、当设备端计算不足时自动回退到 cloud。Cactus 支持跨 iOS、Android、macOS 和嵌入式 Linux 的多模态 inference（LLM、视觉、语音），并提供 Swift、Kotlin、Python、Flutter、React Native 和 Rust SDK。其基准测试显示，在 M4 Pro 上 1.2B 模型 INT4 精度下解码速度达到 100 token/秒，在 iPhone 17 Pro 上达到 48 token/秒，这是优化的 edge inference 的具体示例。

边缘约束

资源	Cloud GPU（H100）	笔记本电脑（M4）	手机（Snapdragon 8 Gen 3）	IoT（ESP32）
RAM	80 GB HBM3	16-36 GB 统一内存	8-12 GB LPDDR5	520 KB
计算	989 TFLOPS（FP8）	38 TOPS（Neural Engine）	45 TOPS（NPU）	0.001 TOPS
功耗	700 W	15-30 W	5-10 W	0.1 W
存储	TB	256 GB-2 TB	128-512 GB	4 MB

• Cloud GPU 和手机 NPU 之间的计算差距约为 20 倍。GPU 和微控制器之间约为 1,000,000 倍。不同设备需要不同级别的压缩和不同的模型架构。

模型压缩管道

• 对于边缘 deployment，压缩不是单一技术——它是一个 顺序应用的互补技术管道：

完整模型（FP32，700 亿参数）
    ↓ Knowledge distillation → 更小的模型（70 亿参数）
    ↓ 结构化 pruning → 移除冗余头/层（40 亿有效参数）
    ↓ Quantisation（INT4）→ 缩小 4 倍（2 GB）
    ↓ 编译器优化 → 融合 kernel，优化内存布局
    ↓ 运行时 → 设备端执行

• 每一步都减少大小和 latency。顺序很重要：先 distillation（减少架构），再 prune（移除结构），再量化（降低精度），再编译（针对目标硬件优化）。在量化后 distillation 会尝试压缩一个已经有损的模型。

设备端运行时

• 运行时加载模型、分配内存并在目标硬件上执行 inference。每个平台有其首选运行时：

• ONNX Runtime：跨平台（Windows、Linux、macOS、iOS、Android）。支持 CPU、GPU（CUDA、DirectML、CoreML、NNAPI）和许多加速器后端。最便携的选项。模型从 PyTorch/TensorFlow 导出为 ONNX 格式。

• TensorFlow Lite（TFLite）：Google 的边缘运行时。针对 ARM CPU 和 Android NPU 优化。二进制文件极小（约 1 MB）。支持 INT8 和 float16。Android deployment 的标准。

• Core ML：Apple 面向 iOS/macOS 的运行时。根据模型特性自动使用 Neural Engine、GPU 或 CPU。模型使用 coremltools 从 PyTorch/TensorFlow 转换。与 Apple 硬件紧密集成（统一内存、Neural Engine）。

• ExecuTorch：Meta 针对设备端 PyTorch 的新运行时。设计用于边缘 deployment，具有提前编译和操作级委托给硬件加速器的功能。PyTorch Mobile 的继任者。

• TensorRT：NVIDIA 面向 GPU inference 优化的运行时（第 15 章）。融合层、选择最优 kernel，并自动量化。在 NVIDIA GPU 上比 PyTorch eager 模式快 2-5 倍。

• llama.cpp：用于 LLM 的单文件 C++ inference 引擎。支持 GGUF 量化（Q4、Q5、Q8）、CPU（AVX/NEON）、Metal（Apple GPU）、CUDA 和 Vulkan。在消费者硬件上运行 LLM 的首选。

编译器栈

• 在高级模型（PyTorch 计算图）和硬件（NPU 指令）之间有 编译器栈，它针对特定目标优化模型：

PyTorch 模型
    ↓ 导出（torch.export、ONNX、TorchScript）
计算图 IR（中间表示）
    ↓ 图优化
        - 常量折叠（在编译时计算常量表达式）
        - 死代码消除（移除未使用的操作）
        - 算子融合（conv + bn + relu → 单个融合操作）
        - 布局变换（NCHW → NHWC 针对 ARM，channels-last）
    ↓ 降级
硬件特定 IR
    ↓ 后端优化
        - 分块和循环排序（缓存友好的访问模式）
        - 向量化（SIMD，第 16 章）
        - 内存规划（重用缓冲区以最小化峰值内存）
        - Kernel 选择（为每个操作选择最佳实现）
    ↓ 代码生成
机器码 / NPU 指令

• 算子融合是影响最大的优化。一个 transformer 块有约 20 个操作（matmul、add、layernorm、softmax 等）。没有融合，每个操作将其输出写入内存，下一个操作再从中读取。通过融合，多个操作合并为一个 kernel，数据保存在寄存器/缓存中。这可以快 2-5 倍（第 16 章，roofline model）。

• 内存规划：编译器分析模型计算图以确定哪些 tensor 在生命周期上有重叠，可以共享同一内存缓冲区。有 100 个中间 tensor 的模型可能只需要 10 个的内存，因为大多数在其他 tensor 被创建之前就已被消费和释放。这在 RAM 有限的设备上至关重要。

硬件目标

移动 GPU

• Qualcomm Adreno（Android）：支持 OpenCL、Vulkan 计算（第 16 章）以及 Qualcomm 的专有 SNPE（Snapdragon Neural Processing Engine）。Adreno GPU 有 256-1024 个 ALU，支持 FP16 和 INT8。

• ARM Mali（Android）：支持 OpenCL 和 Vulkan。Mali GPU 使用基于分块的架构（与桌面 GPU 不同），影响最优内存访问模式。

• Apple GPU（iOS/macOS）：通过 Metal（Apple 的 GPU API）访问。统一内存架构意味着没有 CPU↔GPU 复制开销。Metal Performance Shaders（MPS）提供优化的 ML 原语。

神经处理单元（NPU）

• NPU 是专为 ML inference 设计的固定功能加速器。对于标准 ML 操作（matmul、卷积、激活），它们比 GPU 更节能。

• Apple Neural Engine：16 核，约 38 TOPS（INT8）。通过 Core ML 访问。擅长视觉模型和设备端扩散。无法运行任意代码——只能运行 Core ML 支持的操作。

• Qualcomm Hexagon NPU：集成在 Snapdragon SoC 中。支持 INT8 和 INT4 inference。通过 SNPE 或带 QNN 后端的 ONNX Runtime 访问。驱动背景模糊、语音识别和实时翻译等设备端功能。

• Google Edge TPU：cloud TPU 的小型低功耗版本。4 TOPS，2W。用于 Coral 设备的设备端 inference。只支持 INT8 量化的 TFLite 模型。

• 委托模式：运行时将模型计算图分割为 NPU（支持的操作）和 CPU（不支持的操作）。最大化在 NPU 上运行的部分是性能和能效的关键。

设备端 LLM

• 在手机和笔记本电脑上运行 LLM 已经通过小型模型和激进量化变得可行：

模型	参数量	量化大小	目标设备	性能
Phi-3 Mini	3.8B	约 2 GB（Q4）	手机/笔记本电脑	iPhone 15 约 15 token/秒
Gemma 2B	2B	约 1.5 GB（Q4）	手机	Pixel 8 约 20 token/秒
Llama 3.2 1B	1B	约 700 MB（Q4）	手机	约 30 token/秒
Llama 3.2 3B	3B	约 2 GB（Q4）	手机/笔记本电脑	约 15 token/秒
Llama 3.1 8B	8B	约 4.5 GB（Q4）	笔记本电脑	M2 约 20 token/秒

• 挑战：

  • 内存：3B Q4 模型占用 2 GB，但长对话的 KV-cache 会显著增加。上下文长度在手机上通常限制为 2-4K token。
  • 热降频：持续 inference 会使手机发热。连续生成 30 秒后，SoC 降低时钟频率以防止过热，性能下降 30-50%。
  • 电池：以 15 token/秒运行 3B 模型消耗约 3-5W。30 分钟的对话消耗典型手机电池的约 5%。偶尔使用是可接受的，对于始终在线的应用则有问题。

• llama.cpp 是设备端 LLM 的标准。它在 CPU（AVX2、NEON、I8MM）、Apple GPU（Metal）、NVIDIA GPU（CUDA）、AMD GPU（ROCm/Vulkan），甚至手机（通过 Android 上的 Termux）上运行。

Federated Learning

• Federated learning 在许多设备上训练模型，而不集中数据。每台设备在其本地数据上训练，计算梯度更新，并将只有更新（而不是数据）发送给聚合更新的中央服务器。

• 算法（FedAvg）：

  1. 服务器将当前模型发送给选定的 \(K\) 台设备。
  2. 每台设备在其本地数据上微调模型几步。
  3. 每台设备将其更新后的模型（或差异）发送回服务器。
  4. 服务器平均更新：\(W_{\text{new}} = \frac{1}{K} \sum_{k=1}^{K} W_k\)。
  5. 重复。

• 隐私：原始数据从不离开设备。服务器只看到聚合的模型更新。差分隐私在更新中添加噪声，使得个体数据点无法从梯度中逆向工程。

• 通信效率：模型更新很大（与模型大小相同）。压缩技术减少这一点：梯度量化（发送 INT8 梯度而非 FP32）、稀疏化（只发送最大的梯度）和 梯度累积（进行更多本地步骤，减少发送频率）。

• 应用：Google 的键盘预测（Gboard）、Apple 的语音识别、健康监测（在不集中敏感健康数据的情况下训练）。

Latency 优化

• 除了压缩，几种技术还能减少端到端 inference latency：

• 早退：在中间层添加分类头。如果模型在第 6 层（共 24 层）已经有把握，则返回预测而无需运行第 7-24 层。简单输入提前退出，困难输入使用完整模型。对于包含简单和困难输入混合的任务，平均 latency 显著下降。

• 模型分区：将模型在 NPU（矩阵乘法高效）、GPU（不规则操作高效）和 CPU（处理其他一切）之间分割。编译器根据性能分析决定哪些操作去哪里。

• 缓存：对于有重复查询的应用（自动完成、代码补全），缓存最近的计算。如果用户输入"如何"，而模型最近已经为"如何"生成了补全，可以重用缓存的 KV-cache，完全跳过 prefill 阶段。

• 投机性预取：预测用户接下来会做什么，并在他们询问之前开始 inference。聊天应用可能在用户阅读当前答案时开始生成可能的后续问题的响应。

编程任务（使用 CoLab 或 notebook）

1. 模拟模型压缩管道。从 float32 模型开始，应用 distillation（模拟）、pruning 和量化，并追踪每一步的大小。

   def compression_pipeline(original_params_M, original_bits=32):
       size_mb = original_params_M * 1e6 * original_bits / 8 / 1e6
   
       print(f"原始: {original_params_M}M 参数，{original_bits} 位 → {size_mb:.0f} MB")
   
       # 第 1 步：Knowledge distillation（减少参数）
       distilled_params = original_params_M * 0.15  # 700 亿 → 约 100 亿等效
       size_mb = distilled_params * 1e6 * original_bits / 8 / 1e6
       print(f"Distillation 后（{distilled_params:.0f}M 参数）: {size_mb:.0f} MB")
   
       # 第 2 步：结构化 pruning（移除剩余的 30%）
       pruned_params = distilled_params * 0.7
       size_mb = pruned_params * 1e6 * original_bits / 8 / 1e6
       print(f"Pruning 后（{pruned_params:.0f}M 参数）: {size_mb:.0f} MB")
   
       # 第 3 步：INT4 量化
       size_mb = pruned_params * 1e6 * 4 / 8 / 1e6
       print(f"INT4 量化后: {size_mb:.0f} MB")
   
       print(f"总压缩比: {original_params_M * 1e6 * original_bits / 8 / 1e6 / size_mb:.0f} 倍")
   
   print("=== 从 700 亿参数模型开始 ===")
   compression_pipeline(70000)
   
   print("\n=== 从 70 亿参数模型开始 ===")
   compression_pipeline(7000)
   

2. 估计设备端 inference latency。给定模型的操作数和硬件规格，计算是否满足 latency 目标。

   def estimate_latency(model_name, params_M, bits, compute_tops, mem_bw_gbs, seq_len=256):
       """估计内存带宽受限模型的 token 生成 latency。"""
       # 模型大小（字节）
       model_bytes = params_M * 1e6 * bits / 8
   
       # Decode 受内存限制：每个 token 必须加载整个模型
       time_per_token_ms = model_bytes / (mem_bw_gbs * 1e9) * 1000
   
       # 每秒 token 数
       tokens_per_sec = 1000 / time_per_token_ms
   
       print(f"{model_name}: {params_M/1000:.1f}B 参数 @ {bits} 位 = {model_bytes/1e9:.1f} GB")
       print(f"  内存带宽: {mem_bw_gbs} GB/s")
       print(f"  每 token 时间: {time_per_token_ms:.1f} 毫秒")
       print(f"  Token/秒: {tokens_per_sec:.0f}")
       print()
   
   # Apple M2 Pro：200 GB/s 统一内存带宽
   print("=== Apple M2 Pro（200 GB/s）===")
   estimate_latency("Llama-7B Q4", 7000, 4, 15.8, 200)
   estimate_latency("Llama-7B Q8", 7000, 8, 15.8, 200)
   estimate_latency("Llama-70B Q4", 70000, 4, 15.8, 200)
   
   # 手机（Snapdragon 8 Gen 3）：约 50 GB/s LPDDR5
   print("=== Snapdragon 8 Gen 3（50 GB/s）===")
   estimate_latency("Phi-3 Mini Q4", 3800, 4, 45, 50)
   estimate_latency("Llama-3B Q4", 3000, 4, 45, 50)