arXiv26 | ST-MoE：expert 激活不是随机的——用时空相关性把 expert 提前搬上片

arXiv'26 | ST-MoE：expert 激活不是随机的——用时空相关性把 expert 提前搬上片

原文：A Spatio-Temporal Expert Prefetching Framework for Efficient MoE-based LLM Inference

1. 前言

MoE 的老矛盾我在专栏里念叨过很多次：稀疏激活省的是计算，省不了内存——每个 token 只激活 top-k 个 expert，但你事先不知道是哪 k 个，router 出结果才知道，于是 expert 权重只能按需现搬。搬运在关键路径上，decode 阶段一层一停顿，90% 以上的参数全程闲着还得占着位置或者反复过总线。

这篇来自 UNC Charlotte 和 George Washington 大学的工作（作者里有 Ahmed Louri、Avinash Karanth 这些做片上网络的老牌体系结构学者）问了一个很自然的问题：expert 的选择真的是不可预测的吗？ 答案是否定的，而且规律有两个维度。

2. 观察：expert 激活的时空规律

作者在 DeepSeek、Qwen、Mixtral 等多个 MoE 模型、10 万 token 的语料上做了统计：

跨层与跨 token 的 expert 共激活热力图

两条规律：

  1. 空间相关（跨层）：同一个 token 在第 l 层选了 expert e，那它在第 l+1 层的 expert 选择有明显的统计偏好——左图热力图上那些深色条纹就是高频"expert 组合"。直觉解释：transformer 逐层精炼的是同一个 token 的表征，语义特征是连续的，路由决策自然层间相关。卡方检验 p 值全部 < 0.01
  2. 时间相关（跨 token）：同一层里，相邻 decode token 的 expert 选择重合度是随机基线（E(N)=K²/N）的接近 2 倍。因为自回归解码的每个 token 都基于前文生成，邻近 token 语义依赖强，触发的 expert 自然也像

既然可预测，就可以 prefetch——在第 l 层还在算的时候，提前把第 l+1 层大概率要用的 expert 搬上片，加载延迟就被计算掩盖了。

3. 方案：ST-MoE

整体框架分 profiling 和 runtime 两阶段：

ST-MoE 框架总览：profiling 建表，runtime 查表预测

3.1 两张表长什么样

预测机制轻得出奇，就两张查找表。先把表结构说清楚。

CCT（Cross-layer Correlation Table，跨层相关表）：每对相邻层一张，回答"这一层选了 expert e，下一层谁最可能被选"。构建过程（对应原文 Algorithm 1）：

  1. profiling 阶段跑 P 个 token，维护一个 E×E 的共现矩阵 Co-matrix（E 是每层 expert 数）：token 在第 i 层选了 e、在第 i+1 层选了 f，就给 Co-matrix[e][f] 加 1
  2. 统计完后，对每个 e 取共现次数最高的 Top-K 个 f 存进 CCT[e]（K 与模型的 top-k 路由值一致）
  3. 每个存进来的候选 expert 挂一个 2-bit 置信状态，初始化为 10

这个 2-bit 状态是从 CPU 分支预测器抄来的经典设计——四个状态：11 强偏好、10 弱偏好、01 弱不偏好、00 强不偏好。所以 CCT 的一个 entry 大概长这样（假设 K=4）：

CCT[e=7] = [(f=12, conf=11), (f=3, conf=10), (f=45, conf=10), (f=28, conf=01)]
含义：第 i 层选了 expert 7 时，第 i+1 层最常跟着被选的是 12、3、45、28

HT（History Table，历史表）：每层一张，只存一样东西——上一个 decode token 在这层实际选中的 Top-K expert。每个 token 解码完就整表覆写，永远反映"最近的过去"。

两张表分工明确：CCT 编码的是长期、跨层的统计规律（离线学 + 在线微调），HT 编码的是瞬时、跨 token的局部性（纯在线）。

3.2 预测怎么打分：走一个具体例子

预测发生在每层 gating 结果出来的瞬间（对应原文 Algorithm 2）。打分规则：

  1. 当前层实际选中的每个 expert e，去查 CCT[e]，其中每个候选 f 的置信分累加进 ConfScore[f]（置信状态数值化：10 记 2 分等）
  2. HT 里的每个 expert h，给 ConfScore[h] 加分
  3. 总分 ≥ 2（对应 2-bit 状态里"弱偏好"这一档）的 expert 进 prefetch 集合

拿一个 top-2 路由的玩具例子走一遍。假设第 i 层 router 选中了 expert {7, 21}，要预测第 i+1 层：

查 CCT[7]  → {12: 2分, 3: 2分}       （7 的高共现后继）
查 CCT[21] → {12: 2分, 45: 1分}      （21 的高共现后继）
查 HT[i+1] → {12, 3}                 （上个 token 在 i+1 层用的，各 +1）汇总：ConfScore[12] = 2+2+1 = 5      ← 两条空间线索 + 一条时间线索都指向它ConfScore[3]  = 2+1   = 3ConfScore[45] = 1
阈值 2 过滤 → prefetch {12, 3}

注意两个设计细节：其一，expert 12 被三路信源同时指认，分数最高——这正是"融合"的意义，单一信源都可能误报，交叉验证过的候选才值得花带宽去搬；其二，prefetch 集合的大小是动态的（分数过线的都要），不是死板的 top-k，路由行为混乱时集合自动变大（多搬点保命），规整时自动收窄（省带宽）。

3.3 表怎么在线更新

真实 gating 结果出来后做核对（对应原文 Algorithm 3）：

  • CCT 里猜中的候选置信 +1（封顶 11），猜错的 −1；某候选降到 0 以下，直接换血——被真实结果中出现但表里没有的新 expert 顶替，新人置信重置为 10
  • HT 整表覆写为本 token 的真实选择

这套机制让表能跟着 workload 漂移走：任务从摘要切到代码，路由模式变了，老的相关性 entry 会被逐步扣分、换成新 pattern。所以说它是"应用感知"的——不同任务（摘要/数学/代码）学到的表内容完全不同。

预测准确率相当能打：

各模型各应用上的 expert 预测准确率

普遍超过 80%，DeepSeek 系平均 85-89%。有个很有意思的细节：MATH 任务的准确率稳定高于其他任务——数学推理的解码模式更结构化、token 间依赖更规整，路由行为也就更可预测。CNN/DM 摘要任务 token 转移更发散，预测最难。"任务的语义规整度决定路由可预测度"，这个观察值得记住。

3.4 硬件平台与流水线

配套的加速器由四部分组成：EPU（预测单元，存 CCT/HT + 算置信分）、Router（gating 网络的 MAC 阵列）、EMU（expert 到 PE 的动态映射 + permutation network）、可重构 PE 阵列（dataflow 可按数据复用机会切换）：

ST-MoE 硬件平台总览

执行流水线是精髓，看下图：

三个连续 MoE 层的流水线工作流

对着图走一遍：Layer 1 的 gating 结果一出来，Predictor（P）立刻预测 Layer 2 的 expert 并启动 prefetch（EP）——这段搬运和 Layer 1 的 expert 计算、Layer 2 的 attention 完全重叠。Layer 2 的 router 出真实结果后做验证（V）：预测对的 expert 已经在片上，直接开算；预测漏的走 Missed Loading（ML）现补。miss 的代价有界——最坏退化成按需加载，不会比 baseline 更差，模型精度则完全不受影响（预测只影响搬运时机，不改变路由结果，这是和 Adap-G 那类"少激活 expert"的有损方法的本质区别）。

4. 效果

四个 MoE 模型（Qwen1.5/2.0-MoE、DeepSeek-V2、DeepSeek-MoE）× 三个任务（CNN/DM、MATH、HumanEval），对比 GPU 基线、Adap-G、Pre-gated MoE：

归一化执行时间对比及时间分解

  • 执行时间比 GPU 基线平均降 60%，比 Adap-G 降 56%，比 Pre-gated 降 33%
  • 右边的时间分解图说明收益来源：GPU 基线里 Data Comm（灰色）占七成，ST-MoE 把这块几乎压没了——预测 + 重叠正是打在数据搬运上
  • 对比 Pre-gated MoE 的优势也讲得清楚：Pre-gated 要额外训练一个 pre-gating 函数，ST-MoE 纯查表零训练，还能利用跨 token 相关性

能耗方面比较诚实：比 GPU 基线多 10%（miss 补搬 expert 的代价），但 EDP（能耗延迟积）平均只有 GPU 的 0.4 倍——时间省得多，能耗多一点，综合效率大赚。

消融实验把贡献拆开（五个配置递进）：动态 dataflow 贡献 1.1→1.4 倍，加时间相关（HT）到 2.0 倍，再加空间相关（CCT）到 2.33 倍。空间和时间两路信源确实互补，谁也替代不了谁。

5. 一点个人 take

Expert prefetching 这个方向我一直在跟（自己做 ExpertFlow 时也深挖过 expert 激活的可预测性），说几点看法：

  1. 这篇的价值在于把"预测"做得极轻。之前的方案要么训练额外的预测网络（Pre-gated），要么用 LSTM/attention 做序列预测——预测器本身的开销和部署复杂度就把收益吃掉一块。ST-MoE 用两张查找表 + 阈值打分做到 80%+ 准确率，说明 expert 路由的可预测性是浅层统计规律，杀鸡不需要牛刀。
  2. 跨 token 相关性被严重低估了。大部分 prefetch 工作只用跨层信号（毕竟 gating 结果天然按层出），但消融显示 HT 这路信号贡献了 1.4→2.0 倍的跳变，比 CCT 的增量还大。上个 token 用过的 expert 大概率还热——这其实暗示 expert cache 的 LRU 策略天然就有效，两者是一体两面。
  3. 保留意见：实验是自研加速器 + 模拟器上的结果，profiling 阶段的领域依赖性也值得注意——表是按应用域建的，跨域 serving（一个服务同时接代码和闲聊请求）时 CCT 的质量会打折扣，这在真实多租户场景是个实际问题。另外 batch 变大后不同请求的 expert 需求会天然"并集化"，预测的价值会被稀释——这也是所有 expert 预测方法共同的适用边界：低 batch、边缘侧最香。

欢迎评论区交流，做 MoE 推理的朋友应该会有共鸣。

顺带扯一句题外话：ST-MoE 的 profiling-预测-验证闭环，本质是在"为每个负载自动适配系统行为"，这和 AutoML "为每个任务自动适配模型结构"是同一种思维方式。我们把 AutoML 方向的积累整理成了《动手学 AutoML：从 NAS 到大语言模型优化实战》，应用篇里 LLM 架构自动化、MoE 相关的模型融合都有覆盖，对"让系统自己学会调优"感兴趣的读者可以翻翻。

动手学AutoML书籍封面