端侧推理

端侧推理

端侧推理（edge inference）把模型跑在用户设备上（手机、笔记本、IoT 传感器），不把数据发到云端。本文件涵盖端侧约束、模型压缩流水线、端侧运行时、编译器栈、硬件目标（NPU、神经引擎）、端侧 LLM、联邦学习，以及延迟优化

• 云端推理需要联网、增加延迟（往返 50-200 ms）、每次请求都要花钱，还会把用户数据发给第三方服务器。端侧推理一举消除这四样：模型在本地跑、即时响应、每次推理零成本、数据隐私得以保留。

• 权衡：端侧设备的算力和内存比数据中心 GPU 少 100-1000 倍。要让模型在这些约束下跑起来，需要在每一层做激进的优化。

• Cactus（开源项目，仓库位于 github.com/cactus-compute/cactus）是一个为移动和可穿戴设备量身打造的低延迟 AI 引擎。它在生产中演示了本文件讲到的很多技术：用于注意力和矩阵运算的自定义 ARM SIMD kernel（第 16 章）、KV-cache 量化（第 17 章文件 01）、分块预填充、在 Apple 和高通芯片上的 NPU 加速推理、用零拷贝内存映射把内存占用降到十分之一，以及在端侧算力不足时自动回退到云端。Cactus 支持 iOS、Android、macOS 和嵌入式 Linux 上的多模态推理（LLM、视觉、语音），并提供 Swift、Kotlin、Python、Flutter、React Native 和 Rust 的 SDK。它的基准测试显示：在 M4 Pro 上对 1.2B 模型做 INT4 推理，解码速度可达 100 tokens/s；在 iPhone 17 Pro 上可达 48 tokens/s——这是优化过的端侧推理具体长什么样。

端侧约束

• 云端 GPU 和手机 NPU 之间的算力差距约 20 倍。GPU 和单片机之间的差距约 100 万倍。不同设备需要不同程度的压缩和不同的模型架构。

模型压缩流水线

• 对端侧部署来说，压缩不是单一技术——它是一条由互补技术按顺序应用的流水线：

完整模型（FP32，70B 参数）
    ↓ 知识蒸馏 → 更小的模型（7B 参数）
    ↓ 结构化剪枝 → 移除冗余头/层（有效 4B）
    ↓ 量化（INT4）→ 缩小 4 倍（2 GB）
    ↓ 编译器优化 → 融合 kernel、优化内存布局
    ↓ 运行时 → 端侧执行
​

• 每一步都减小体积和延迟。顺序很重要：先蒸馏（缩减架构）、再剪枝（移除结构）、再量化（降精度）、再编译（针对目标硬件优化）。如果在量化后再蒸馏，等于在压缩一个已经有损的模型。

端侧运行时

• **运行时（runtime）**加载模型、分配内存，并在目标硬件上执行推理。每个平台都有自己的首选运行时：

• 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 部署的标准。

• Core ML：Apple 在 iOS/macOS 上的运行时。根据模型特征自动使用神经引擎、GPU 或 CPU。模型用 coremltools 从 PyTorch/TensorFlow 转换。与 Apple 硬件深度集成（统一内存、神经引擎）。

• ExecuTorch：Meta 新的端侧 PyTorch 运行时。专为端侧部署设计，支持提前编译（AOT）和算子级的硬件加速器委派。是 PyTorch Mobile 的继任者。

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

• llama.cpp：单文件 C++ 的 LLM 推理引擎。支持 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）
    ↓ lowering
硬件特定 IR
    ↓ 后端优化
        - 分块和循环序（cache 友好的访问模式）
        - 向量化（SIMD，第 16 章）
        - 内存规划（复用缓冲区以降低峰值内存）
        - kernel 选择（为每个算子选最优实现）
    ↓ 代码生成
机器码 / NPU 指令
​

• **算子融合（operator fusion）**是影响最大的优化。一个 transformer 块大约有 20 个算子（矩阵乘、加、layernorm、softmax 等）。不融合的话，每个算子都要把输出写回显存，下一个再读回来。融合后，多个算子合并成一个 kernel，数据一直留在寄存器/缓存里。这能快 2-5 倍（第 16 章 roofline 模型）。

• 内存规划（memory planning）：编译器分析模型图，判断哪些张量生命周期重叠、可以共用同一块内存缓冲。一个有 100 个中间张量的模型，可能只需要 10 个张量的内存，因为大多数张量在别的张量诞生前就被消费、释放了。这对内存有限的设备至关重要。

硬件目标

移动 GPU

• Qualcomm Adreno（Android）：支持 OpenCL、Vulkan compute（第 16 章）以及高通自家的 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 推理设计的固定功能加速器。对于标准 ML 运算（矩阵乘、卷积、激活），它比 GPU 能效高得多。

• Apple 神经引擎（Neural Engine）：16 核，约 38 TOPS（INT8）。通过 Core ML 访问。非常适合视觉模型和端侧扩散模型。不能跑任意代码——只支持 Core ML 支持的算子。

• Qualcomm Hexagon NPU：集成在 Snapdragon SoC 里。支持 INT8 和 INT4 推理。通过 SNPE 或带 QNN 后端的 ONNX Runtime 访问。支撑背景虚化、语音识别、实时翻译等端侧功能。

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

• 委派模式（delegation pattern）：运行时把模型图在 NPU（跑支持的算子）和 CPU（跑不支持的算子）之间切分。让尽量多的部分跑在 NPU 上，是性能和能效的关键。

端侧 LLM

• 凭借小模型和激进量化，在手机和笔记本上跑 LLM 已经可行：

• 挑战：

  • 内存：3B Q4 模型装进 2 GB 没问题，但长对话的 KV-cache 会显著增加。手机上上下文长度通常限制在 2-4K token。
  • 热降频（thermal throttling）：持续推理会让手机发热。连续生成 30 秒后，SoC 会降频以防过热，性能下降 30-50%。
  • 电池：以 15 tokens/s 跑 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）**在不集中数据的前提下，跨众多设备训练模型。每个设备在自己的本地数据上训练，算出梯度更新，只把更新（不是数据）发给一个聚合更新的中心服务器。

• 算法（FedAvg）：

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

• 隐私：原始数据从不离开设备。服务器只看到聚合后的模型更新。**差分隐私（differential privacy）**给更新加噪声，使单个数据点无法从梯度反推出来。

• 通信效率：模型更新很大（和模型一样大）。压缩技术能减小它：梯度量化（发 INT8 梯度而非 FP32）、稀疏化（sparsification）（只发最大的那些梯度）、梯度累积（多做本地步数、少发几次）。

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

延迟优化

• 除了压缩，还有几种技术能降低端到端推理延迟：

• 提前退出（early exit）：在中间层加分类头。如果模型在第 6 层（共 24 层）就 confident，就直接返回预测，不跑 7-24 层。简单输入提前退出，困难输入用完整模型。对难易混合的任务，平均延迟显著下降。

• 模型切分（model partitioning）：把模型在 NPU（矩阵乘高效）、GPU（不规则运算高效）和 CPU（其余都兜底）之间切分。编译器根据 profiling 决定每个算子去哪跑。

• 缓存（caching）：对有重复查询的应用（自动补全、代码补全），缓存最近的计算。如果用户输入"如何"，而模型最近为"如何"生成过补全，缓存的 KV-cache 就能复用，完全跳过预填充阶段。

• 投机预取（speculative prefetching）：预测用户接下来会做什么，在用户开口前就启动推理。一个聊天应用可以在用户读当前回答时，就开始为可能的后续问题生成回复。

编程练习（使用 CoLab 或 notebook）

1. 模拟模型压缩流水线。从一个 float32 模型出发，依次应用蒸馏（mock）、剪枝和量化，并跟踪每一步的大小。

def compression_pipeline(original_params_M, original_bits=32):
    size_mb = original_params_M * 1e6 * original_bits / 8 / 1e6

    print(f"Original: {original_params_M}M params, {original_bits}-bit → {size_mb:.0f} MB")

    # 第 1 步：知识蒸馏（减少参数）
    distilled_params = original_params_M * 0.15  # 70B → 约 10B 等效
    size_mb = distilled_params * 1e6 * original_bits / 8 / 1e6
    print(f"After distillation ({distilled_params:.0f}M params): {size_mb:.0f} MB")

    # 第 2 步：结构化剪枝（再移除 30%）
    pruned_params = distilled_params * 0.7
    size_mb = pruned_params * 1e6 * original_bits / 8 / 1e6
    print(f"After pruning ({pruned_params:.0f}M params): {size_mb:.0f} MB")

    # 第 3 步：INT4 量化
    size_mb = pruned_params * 1e6 * 4 / 8 / 1e6
    print(f"After INT4 quantisation: {size_mb:.0f} MB")

    print(f"Total compression: {original_params_M * 1e6 * original_bits / 8 / 1e6 / size_mb:.0f}x")

print("=== Starting from 70B model ===")
compression_pipeline(70000)

print("\n=== Starting from 7B model ===")
compression_pipeline(7000)
​

2. 估算端侧推理延迟。给定模型的算子量和硬件规格，计算是否满足延迟目标。

def estimate_latency(model_name, params_M, bits, compute_tops, mem_bw_gbs, seq_len=256):
    """估算访存带宽受限模型的 token 生成延迟。"""
    # 模型大小（字节）
    model_bytes = params_M * 1e6 * bits / 8

    # 解码是访存受限的：每个 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 params @ {bits}-bit = {model_bytes/1e9:.1f} GB")
    print(f"  Memory bandwidth: {mem_bw_gbs} GB/s")
    print(f"  Time per token: {time_per_token_ms:.1f} ms")
    print(f"  Tokens/sec: {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)
​