arXiv'26 | APEX：端侧 MoE 的专家预取，别再 top-k 一刀切了

August 16, 2026

2026   ·   LLM   MoE   端侧推理   推理优化   论文解读  

arXiv’26 | APEX：端侧 MoE 的专家预取，别再 top-k 一刀切了

原文：APEX: Adaptive Expert Prefetching for Edge MoE Inference

1. 前言：先把背景交代清楚

你有没有想过这样一个问题：MoE 模型号称”稀疏激活、适合端侧”，为什么真放到手机/边缘盒子上，还是慢得像老牛拉车？

答案还是那个老矛盾：激活是稀疏的，加载不是。 端侧设备的片上内存装不下全部专家权重，绝大多数专家只能躺在 LPDDR5X 里。每来一个 token，router 挑出 top-k 个专家，再去片外内存把它们搬进来——搬的这段时间，计算阵列整块闲置，还在白白耗静态功耗。论文在 Granite-3.1-3B-A800M 上量了这笔账：专家加载占延迟的 43%、能耗的 29%。

端侧 MoE 的延迟与能耗拆解

预取（prefetching）是标准解法：趁当前层还在算，提前把下一层要用的专家搬进来。这条线我之前写过一篇相关的：《ST-MoE expert-prefetch》——它的核心洞察是 expert 激活不是随机的，有时空相关性，所以可以用历史激活模式预测下一层要哪些专家，提前搬上片。另一篇 《SMoE》 走的是另一条路：专家不在 GPU 上？那就别死等，找个”平替”专家直接算——本质上是用精度换不 stall。

今天这篇 APEX（arXiv 2608.11688）站在这两条路之间，它要指出的问题是：现有预取的”固定 top-k”策略，在端侧是一个两头堵的死局：

  • 预算给小了（比如只多取 δ=2 个），难 token 的专家预测 miss，照样 stall；
  • 预算给大了（δ=8 无脑超采），简单 token 白搬一堆用不上的专家，能耗爆炸。

而 ProMoE 这类固定 top-k 预取的代表工作，平均预测准确率 70–85% 看着不错——但平均数在这里毫无意义，任何一层的 miss 都会造成串行的纠错等待。

APEX 的答案：预算本身应该是每个 token 自适应的。给每个 token 实时估计”再多取几个专家能把覆盖率提到 90%”，按需分配。在模拟的高端端侧平台上，重叠准确率保住 97–98%（ProMoE 会掉到 79%），延迟比无预取低 42%，能耗比固定超采低 21.8%，EDP（延迟×能耗）最高改善 41%。

2. 平台：先把”端侧设备”说清楚

这篇是算法-硬件协同设计的工作，评估平台是 CHIPSIM——一个综合校准的联合仿真器，不是真机。目标设备定位在”近端旗舰边缘加速器”（Jetson Thor / Hailo-10H 这一档），chiplet 架构：

  • 计算 chiplet：4 个 vector processing array（每个 16×16 vector unit、vector length 32），BF16 峰值 24 TFLOPS，8MB SRAM；阵列 RTL 级实现、TSMC 28nm 综合、PrimeTime 功耗表征；
  • 封装内存：HBM3-class 819 GB/s（放共享的热数据）；
  • 片外：专家权重放 LPDDR5X，PCIe 6.0 ×16、256 GB/s；
  • 预取按真实 DMA 建模——受带宽竞争和排队影响，不是理想化的并行搬运。

这个设定比”在 A100 上模拟手机”要严谨得多，读这类论文时平台假设一定要先看清楚。

3. 方法：一个会看人下菜的预取路由器

3.1 Prefetch Router：预测本层专家

和常见”用上一层状态预测下一层专家”的做法不同，APEX 的 prefetch router 放在本层 attention 之前——用当前的 hidden state 预测本层 FFN 将要路由的专家，把 route→load→execute 的流水线改成 predict→prefetch→execute。

Router 本体是一个线性层 + softmax，从冻结的原 router 蒸馏（KL loss），基座模型完全不动。参数量开销 0.79M–34.11M，只占模型参数的 0.022%–0.060%，性能开销小于 0.051%。

3.2 自适应预算 δ̂(x)：这篇的灵魂

固定策略的问题在于不同 token 的”预测难度”差异极大。APEX 给每个 token 建了一个序数 logistic CDF 模型：

\[p_\delta(x) = \Pr(\delta \geq \delta^* \mid x) = \sigma(\theta_\delta - w^\top x)\]

其中 $\delta^*$ 是”恰好覆盖真实路由专家所需的额外预算”。推理时选最小的 δ 使 $p_\delta(x) \geq \tau$（置信阈值 τ=0.90 默认）——预测有把握的 token 只多取一两个，没把握的 token 才大手笔超采。

训练分两步：先 KL 蒸馏 prefetch router，再用累积二元交叉熵拟合这个 CDF（标签来自 oracle δ*）。整个过程离线完成，不改基座。

重叠准确率随预算的变化

3.3 两种执行模式

  • Correctness-preserving（默认）：预测漏了的专家异步补取，已就绪的专家先算，聚合前补齐——语义与原 MoE 完全一致，代价只是没藏住的纠错延迟；
  • Stall-free（可选）：绝不等待，缺失的专家直接用”已预取的、router 分数最高的候选”顶替——这就是 SMoE 那篇”找平替”思想的端侧版，简单任务上再赚 2%–3% 延迟，但 k 很小的模型（如 Phi-mini-MoE，k=2）会明显掉点，论文建议这类模型关掉。

4. 效果

4.1 重叠准确率：自适应的碾压局

τ=0.90 下，APEX 在所有层保持 >97–98% 的专家重叠率；ProMoE 在 Granite-1B 的第 5 层会掉到 79%。平均每个 token 额外预取的专家数：Granite-1B 4.17 个、Granite-3B 2.86 个、Phi-7B 0.67 个、DeepSeek-V2-Lite 1.98 个——模型越”好预测”，预算自动越小，这正是自适应的意义。

逐层重叠准确率

4.2 延迟与能耗（Granite-3B，512 token）

  • 延迟 11.41ms：比无预取（19.77ms）低 42%，比 ProMoE（15.39ms）低 26%，比固定超采 (k+4)/(k+8) 分别低 20%/40%；
  • 能耗 287.3mJ：比无预取低 9.5%，比 ProMoE 低 5.8%；固定 k+8 超采最差（比 APEX 多 21.8%）——固定预算省下的延迟全在能耗上还回去了；
  • 能耗拆解很有意思：开预取后片外 I/O 从 48mJ 涨到 56mJ，但计算阵列的闲置漏电从 57mJ 崩到 8mJ——用 8mJ 的搬运换 49mJ 的干等，这就是预取生意的本质。

延迟与能耗对比

4.3 鲁棒性

EDP 相对无预取改善 36–49%、相对 ProMoE 改善 16–30%（各模型不同）；收益在 32–1024 GB/s 带宽区间、4/8/16/32-bit 专家精度下都成立——和量化是正交可叠加的。Stall-free 模式的精度代价：Granite 系列 PPL 7.88→7.99 这种量级，可接受；Phi-7B 掉得明显（64.6→60.4），印证了”k 小则每个专家都是命根子”。

5. 我的 Take

端侧 MoE 这两年被”稀疏激活”的话术罩着一层滤镜，真实瓶颈一直是专家权重的搬运。APEX 这个工作我最欣赏两点：

一是它抓住了预取问题的正确度量——不是平均准确率，而是尾部覆盖率。固定 top-k 的 80% 平均准确率听起来能打，但那一层 miss 带来的串行等待会把前面九层的完美预取全部白干。用 CDF 建模”这个 token 需要多少预算”是个很干净的解法，本质上就是把置信度校准用到了内存搬运调度上。

二是诚实：correctness-preserving 是默认模式，stall-free 的精度代价明码标价，连”哪些模型不适合”都写清楚了。和 SMoE 的”平替”路线、ST-MoE 的”时空相关性预测”路线对照着看，端侧 MoE 的解题空间基本就铺开了：预测得准（ST-MoE）、取得巧（APEX）、算得动（SMoE），三件事各占一角。

如果这篇文章涉及的 MoE 推理优化你想系统深入，可以看看我之前出版的《动手学AutoML：从 NAS 到大语言模型优化实战》，书里讲 NAS 在大模型上的延伸时正好覆盖了稀疏模型结构搜索这条线，和 MoE 的效率优化是同一套思路。

动手学AutoML书籍封面

Flag Counter