Skip to content

边缘推理

边缘推理在用户设备(手机、笔记本、物联网传感器)上运行模型,无需将数据发送到云端。本文件涵盖边缘约束、模型压缩流水线、设备端运行时、编译器栈、硬件目标(NPU、Neural Engine)、设备端LLM、联邦学习和延迟优化

  • 云端推理需要网络连接,增加延迟(50-200ms往返),按请求计费,并将用户数据发送到第三方服务器。边缘推理消除以上所有问题:模型本地运行、即时响应、每次推理零成本、数据保持隐私。

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

  • Cactusgithub.com/cactus-compute/cactus)是一个专为移动和可穿戴设备打造的低延迟AI引擎。它在生产中展示了本文件涵盖的许多技术:用于注意力和矩阵操作的自定义ARM SIMD内核(第16章)、KV缓存量化(第17章文件01)、分块预填充、Apple和Qualcomm芯片上的NPU加速推理、零拷贝内存映射以实现10倍更低的RAM使用,以及当设备端计算不足时自动回退到云端。Cactus支持iOS、Android、macOS和嵌入式Linux上的多模态推理(LLM、视觉、语音),提供Swift、Kotlin、Python、Flutter、React Native和Rust的SDK。其基准测试显示在INT4下,1.2B模型在M4 Pro上达到100 tokens/s的解码速度,在iPhone 17 Pro上达到48 tokens/s,这是优化边缘推理效果的具体实例。

边缘约束

资源 云端GPU(H100) 笔记本(M4) 手机(骁龙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
  • 云端GPU和手机NPU之间的计算差距约为20倍。GPU和微控制器之间约为1,000,000倍。不同设备需要不同级别的压缩和不同的模型架构。

模型压缩流水线

  • 对于边缘部署,压缩不是单一技术——它是一条流水线,按顺序应用互补技术:
完整模型(FP32, 70B参数)
    ↓ 知识蒸馏 → 更小的模型(7B参数)
    ↓ 结构化剪枝 → 删除冗余的头/层(4B有效)
    ↓ 量化(INT4)→ 4倍更小(2 GB)
    ↓ 编译器优化 → 融合内核、优化内存布局
    ↓ 运行时 → 设备端执行
  • 每一步减小大小和延迟。顺序很重要:先蒸馏(缩减架构),再剪枝(移除结构),再量化(降低精度),再编译(针对目标硬件优化)。量化后蒸馏会试图压缩已经是有损的模型。

设备端运行时

  • 运行时加载模型、分配内存并在目标硬件上执行推理。每个平台有其偏好的运行时:

  • 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运行时。根据模型特性自动使用Neural Engine、GPU或CPU。模型通过coremltools从PyTorch/TensorFlow转换。与Apple硬件紧密集成(统一内存、Neural Engine)。

  • ExecuTorch:Meta的新设备端PyTorch运行时。专为边缘部署设计,具有提前编译和算子级委托到硬件加速器。PyTorch Mobile的后继者。

  • TensorRT:NVIDIA的GPU推理优化运行时(第15章)。融合层、选择最优内核并自动量化。在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
    ↓ 后端优化
        - Tiling和循环排序(缓存友好的访问模式)
        - 向量化(SIMD,第16章)
        - 内存规划(重用缓冲区以最小化峰值内存)
        - 内核选择(为每个算子选择最佳实现)
    ↓ 代码生成
机器码 / NPU指令
  • 算子融合是影响最大的优化。一个Transformer块有约20个操作(矩阵乘法、加法、LayerNorm、Softmax等)。不融合时,每个操作写入其输出到内存,下一个再读取回来。融合后,多个操作被组合到单个内核中,数据保持在寄存器/缓存中。这可以快2-5倍(第16章,roofline模型)。

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

硬件目标

移动GPU

  • Qualcomm Adreno(Android):支持OpenCL、Vulkan计算(第16章)和Qualcomm专有的SNPE(骁龙神经处理引擎)。Adreno GPU有256-1024个ALU,支持FP16和INT8。

  • ARM Mali(Android):支持OpenCL和Vulkan。Mali GPU使用基于Tile的架构(与桌面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:集成在骁龙SoC中。支持INT8和INT4推理。通过SNPE或ONNX Runtime配合QNN后端访问。驱动设备端功能如背景虚化、语音识别和实时翻译。

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

  • 委托模式(Delegation):运行时在NPU(用于支持的操作)和CPU(用于不支持的操作)之间拆分模型图。最大化在NPU上运行的分数是性能和能效的关键。

设备端LLM

  • 在手机和笔记本上运行LLM已通过小型模型和激进量化变得可行:
模型 参数 量化后大小 目标设备 性能
Phi-3 Mini 3.8B ~2 GB(Q4) 手机/笔记本 iPhone 15上~15 tokens/s
Gemma 2B 2B ~1.5 GB(Q4) 手机 Pixel 8上~20 tokens/s
Llama 3.2 1B 1B ~700 MB(Q4) 手机 ~30 tokens/s
Llama 3.2 3B 3B ~2 GB(Q4) 手机/笔记本 ~15 tokens/s
Llama 3.1 8B 8B ~4.5 GB(Q4) 笔记本 M2上~20 tokens/s
  • 挑战

    • 内存:3B Q4模型可装入2 GB,但长对话的KV缓存额外增加显著。上下文长度在手机上通常限制在2-4K tokens。
    • 热节流:持续推理加热手机。连续生成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)上。

联邦学习

  • 联邦学习跨多个设备训练模型,而不集中数据。每个设备在其本地数据上训练,计算梯度更新,只发送更新(而非数据)到中心服务器,服务器聚合更新。

  • 算法(FedAvg):

    1. 服务器将当前模型发送给\(K\)个选定设备。
    2. 每个设备在其本地数据上微调模型若干步。
    3. 每个设备将其更新后的模型(或差异)发回服务器。
    4. 服务器平均更新:\(W_{\text{new}} = \frac{1}{K} \sum_{k=1}^{K} W_k\)
    5. 重复。
  • 隐私:原始数据从不离开设备。服务器只看到聚合后的模型更新。差分隐私向更新添加噪声,使得无法从梯度逆推个别数据点。

  • 通信效率:模型更新很大(与模型相同大小)。压缩技术减少此开销:梯度量化(发送INT8梯度而非FP32)、稀疏化(仅发送最大梯度)和梯度累积(多做本地步骤,少发送)。

  • 应用:Google的键盘预测(Gboard)、Apple的语音识别、健康监测(在敏感健康数据上训练而不集中数据)。

延迟优化

  • 除了压缩之外,还有几种技术减少端到端推理延迟:

  • 提前退出:在中间层添加分类头。如果模型在第6层(共24层)就置信,则返回预测而不运行第7-24层。简单输入提前退出,困难输入使用完整模型。对于混合简单和困难输入的任务,平均延迟显著下降。

  • 模型分区:在NPU(高效矩阵乘法)、GPU(高效不规则操作)和CPU(处理其余)之间拆分模型。编译器基于性能分析决定每个操作去往何处。

  • 缓存:对于有重复查询的应用(自动补全、代码补全),缓存最近的运算。如果用户输入"How do I",且模型最近已为"How do I"生成补全,则可重用缓存的KV缓存,完全跳过预填充阶段。

  • 推测预取:预测用户下一步会做什么,在用户提问之前开始推理。聊天应用可以在用户阅读当前回复时开始为可能的追问生成回复。

编程任务(使用CoLab或notebook)

  1. 模拟模型压缩流水线。从float32模型开始,应用蒸馏(模拟)、剪枝和量化,跟踪每一步的大小。

    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步:知识蒸馏(减少参数)
        distilled_params = original_params_M * 0.15  # 70B → ~10B等效
        size_mb = distilled_params * 1e6 * original_bits / 8 / 1e6
        print(f"蒸馏后 ({distilled_params:.0f}M 参数): {size_mb:.0f} MB")
    
        # 第2步:结构化剪枝(移除剩余30%)
        pruned_params = distilled_params * 0.7
        size_mb = pruned_params * 1e6 * original_bits / 8 / 1e6
        print(f"剪枝后 ({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}x")
    
    print("=== 起始于70B模型 ===")
    compression_pipeline(70000)
    
    print("\n=== 起始于7B模型 ===")
    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 参数 @ {bits}-位 = {model_bytes/1e9:.1f} GB")
        print(f"  内存带宽: {mem_bw_gbs} GB/s")
        print(f"  每token时间: {time_per_token_ms:.1f} ms")
        print(f"  Tokens/秒: {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)
    
    # 手机 (骁龙8 Gen 3): ~50 GB/s LPDDR5
    print("=== 骁龙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)