arXiv:2607.06601cs.LGcs.AI2026-07

统一调度三类计算资源,让大模型推理更省力还更准。

TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation

  • 用一个轻量控制器同时决定注意力、专家和缓存位宽
  • 在相同算力下比单独优化效果更好,且更抗罕见词干扰
  • 适合追求高效推理的开发者,尤其对长文本和代码友好

条件计算可将语言模型质量与每令牌推理成本解耦,但现有技术仅独立优化单一维度:专家混合(MoE)稀疏化前馈网络,深度混合(MoD)跳过整块变换器,KV缓存量化压缩注意力内存。我们提出三者决策强耦合,应联合处理:一个值得全注意力的词可能也需高精度缓存,无论由哪个专家处理。引入TriRoute,一个跨三轴共享的轻量控制器,为每层每个令牌生成协调策略:(i) 注意力模式(跳过/局部/完整),(ii) 稀疏专家集(空专家可恢复MoD),(iii) KV缓存位宽。控制器通过异构松弛(Gumbel-Softmax结合直通估计,专家使用负载均衡top-k门控)在拉格朗日预算约束下端到端训练,将平均计算与内存成本转为可控参数。识别出联合训练中的跨轴路由坍塌级联现象,并通过各轴归一化和耦合感知平衡损失解决。在160M至1.3B参数的解码器模型上,于计算最优令牌数下,TriRoute在匹配推理FLOPs与内存的前提下,优于最优独立组合,同时更好保持对罕见实体、代码和算术任务的鲁棒性。事后分析显示:控制器将全注意力与高精度缓存分配给句首位置、稀有子词和命名实体,而廉价路由功能词。

原文摘要 · Abstract (English)

Conditional computation can decouple language model quality from per-token inference cost, yet leading techniques act on a single axis in isolation: Mixture-of-Experts (MoE) sparsifies the FFN, Mixture-of-Depths (MoD) skips whole transformer blocks, and KV-cache quantization compresses attention memory. We argue these three decisions (attention resolution, expert selection, and cache bit-width) are strongly coupled and should be made jointly: a token rare enough to warrant full attention may also need high-precision caching regardless of which expert processes it. We introduce TriRoute, a single lightweight controller shared across all three axes that, for every token at every layer, emits a coordinated policy: (i) an attention mode (skip/local/full), (ii) a sparse set of FFN experts (with a null expert recovering MoD), and (iii) a KV-cache bit-width. The controller trains end-to-end via a heterogeneous relaxation (Gumbel-Softmax with straight-through estimation for categorical decisions and load-balanced top-k gating for experts) under a Lagrangian budget constraint that turns the average compute and memory cost into a controllable knob. We identify a cross-axis routing-collapse cascade in naive joint training, where collapse on one axis propagates to the others, and address it with per-axis normalization and a coupling-aware balancing loss. On decoder-only models from 160M to 1.3B parameters at compute-optimal token counts, TriRoute Pareto-dominates the best independent MoD+MoE+KV-quantization combination at matched inference FLOPs and memory, while better preserving tail-case robustness on rare entities, code, and arithmetic that pure perplexity optimization erodes. Post-hoc analysis reveals interpretable structure: the controller allocates full attention and high-precision cache to sentence-initial positions, rare subwords, and named entities, while cheaply routing function words.

条件计算推理优化MoEKV缓存

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。