火韦先生
26-06-23 21:33

attention 层做 MoE 这条线,三年走下来居然变温和了。

2023 年 SwitchHead 出来时是激进派:262M 模型,44% 算力、27% 显存就能打平 baseline,KV 和 Q 一起压。听着很爽,但 KV cache 一动,长上下文就废了一半——这也是为啥后来没人真把它端上线。

最近这篇 GQE(Grouped Query Experts)就老实多了:在 GQA group 内只对 query head 做 expert 路由,KV head 一根不动,专门保 KV cache 这条命脉。代价是要激活一半 Q heads 才能勉强 match baseline,节省幅度比 SwitchHead 小了一大截。

🤔 但我觉得这才是 long context 时代的正解。

做推理优化的都懂:模型上到 70B、context 128k 之后,瓶颈早不是 attention 的 FLOPs,是 KV cache 的显存和带宽。你激进压 KV 省下来的那点算力,换不来部署侧任何收益,反而把 prefix cache、PD 分离这些工程红利全堵死。

⚠️ 当然 GQE 也有几个问号没回答。

一个是 ROI——Q heads 在 attention 里占比本来就不大,FFN MoE 才是 transformer 算力的真大头,省一半 Q 计算端到端能省多少?大概率个位数。

另一个是规模——250M 参数加 30B token 在今天只能算 toy。MoE 的负载均衡、路由稳定性这些坑要 scale 到 7B、70B 才暴露,小规模看着甜的设计 scale up 经常翻车。

💡 不过方向是对的:attention 优化要先承认 KV cache 不能碰,再去琢磨能动哪儿。从激进派到妥协派,这条路线总算开始懂工程了。 arXiv: 2606.20945

发布于 日本