跳转至

AdaFuse: Accelerating Dynamic Adapter Inference via Token-Level Pre-Gating and Fused Kernel Optimization

会议: AAAI 2026
arXiv: 2603.11873v1
代码: 无
领域: 模型压缩 / LLM高效推理
关键词: 动态适配器, MoE-LoRA, CUDA内核优化, 推理加速, Token级路由

一句话总结

针对动态MoE-LoRA适配器推理延迟暴增（250%-950%）的问题，提出了一种token级预门控架构，只在第一层做一次全局路由决策，配合自研的SGMM融合CUDA内核将所有激活的LoRA适配器一次性合并进骨干网络，在保持精度的同时将解码延迟降低2.4倍。

背景与动机

MoE与LoRA的结合（动态适配器）是当前增强LLM多任务能力的热门路线——通过条件激活不同的LoRA专家来提升模型容量。然而作者发现一个被普遍忽视的实际问题：虽然动态适配器只增加1%-5%参数量和不到1%的FLOPs，但推理延迟却暴增250%-950%。以MoRAL为例，解码延迟从2.4ms/token飙升到8.6ms/token。

根本原因不是计算量，而是CUDA内核的碎片化调用。每个LoRA适配器需要两次CUDA内核调用（down projection和up projection），多层动态路由导致内核启动次数与层数成正比。由于CUDA内核启动延迟不随矩阵大小线性缩放，LoRA的小矩阵反而导致"启动开销远大于实际计算"的窘境。

直接对已有方法做layer-by-layer合并也行不通——MoRAL加上简单合并策略后延迟仍然高达4.5ms/token（比纯骨干模型高88%），因为合并本身也需要额外的CUDA内核调用。

核心问题

如何在保留动态适配器表达能力的同时，消除由碎片化CUDA内核调用引起的推理延迟开销？ 这本质上是一个算法-系统协同设计问题：现有layer-wise或block-wise路由结构天然不利于内核级优化。

方法详解

整体框架

AdaFuse的核心理念是"decide once, apply everywhere"——输入token只在第一个扩展线性层经过一次Top-2路由器，得到的专家激活权重直接应用于所有后续层。这样，每个token的执行路径在推理前就完全确定，从"动态"变成了"按token静态"，为系统级优化打开了大门。

推理流程分三步：(1) 路由器确定当前token激活哪两个LoRA专家及权重；(2) SGMM内核将激活的LoRA参数一次性融合进骨干网络的所有层；(3) 融合后的骨干网络像原始模型一样做前向计算。

关键设计

  1. Token级预门控（Token-Level Pre-Gating）: 不同于现有方法在每层或每个block独立做路由决策，AdaFuse只在第一个线性层放一个Top-2路由器。这个设计基于一个观察：语义相似的token在不同层倾向于激活一致的专家模式。微调时公式为 \(y^l = f^l(x^l) + \sum_{i=1}^{N} G^1(x^1)_i \cdot E_i^l(x^l)\)，其中 \(G^1\) 只依赖第一层输入。这个"一次决策"的架构将路由计算量从 \(O(L)\) 降到 \(O(1)\)，更关键的是使得适配器可以被预合并。

  2. 融合适配器切换（Fused Adapter Switching）: 当相邻token激活的专家不同时，需要先"反合并"上一步的适配器，再合并新的适配器。AdaFuse将上一token的反合并（取负权重）和当前token的合并（正权重）拼接成一个统一操作：\(\text{Fused\_LoRA\_DOWN} = \text{concat}[-(\text{LoRA\_DOWN})_{t-1}, (\text{LoRA\_DOWN})_t]\)。这样只需一次CUDA内核调用就能完成所有层的适配器切换。

  3. SGMM CUDA内核（Segmented Gather Matrix Multiplication）: 改编自Punica的SGMV思路，专为融合适配器切换设计的定制CUDA内核。SGMM将所有层的LoRA合并操作打包成一个batched GEMM，利用分块策略让不同层的合并任务分配到不同thread block并行执行。采用预取缓冲机制隐藏显存加载延迟，原地加法(in-place addition)减少内存开销。

损失函数 / 训练策略

训练阶段与标准MoE-LoRA训练一致，使用标准的指令微调损失。关键区别仅在于路由器只在第一层，所有层共享同一组路由权重。预填充阶段未做特殊优化（LLM生成延迟主要来自解码阶段）。

实验关键数据

数据集/指标 指标 AdaFuse 最强基线 说明
通用能力(5项平均, Llama2-7B) Accuracy 60.12% 60.45% (PESC) 持平，MMLU和TruthfulQA上最优
领域任务(3项平均, Llama2-7B) Accuracy 83.60% 84.20% (MoLA) 接近，CommonsenseQA上最优(79.03%)
领域任务(3项平均, Mistral-7B) Accuracy 87.24% 87.06% (PESC) 略优
解码延迟(Llama2-7B) ms/token 3.1 (+29%) 8.5 (+254%, PESC) 比最快基线快2.7倍
显存占用(Llama2-7B) GB 13.8 (+7%) 13.1 (+2%, PESC) 略高但可接受

消融实验要点

  • SGMM内核至关重要: 把SGMM换成简单逐层合并，延迟从3.1ms升至4.2ms（+35%），验证了内核级优化的必要性
  • 预门控 vs 逐层路由: 改为逐层路由（类似MoRAL），延迟回升到8.6ms——预门控是加速的关键
  • Top-k 选择: Top-2 是最优平衡（Top-1 精度下降，Top-4 延迟上升）
  • LoRA秩影响: 秩从64增到128，精度微升但延迟几乎不变——SGMM内核对秩不敏感
  • 专家数量: N=8 在精度和开销间取得最优平衡
  • 预门控+简单合并 vs MoRAL+简单合并: 4.2ms vs 4.5ms，说明token级预门控架构本身也有贡献，但更大收益来自SGMM
  • MOLA延迟极高(25.3ms/token): 因为MOLA的block-wise路由导致每层都做MoE计算，内核调用次数最多

亮点

  • 问题洞察犀利: 指出动态适配器的瓶颈不在于FLOPs而在于CUDA内核启动开销，这个诊断非常精准，对领域有启发价值
  • 算法-系统协同设计: 不是简单地在系统层面打补丁，而是重新设计算法结构使其对系统友好。"先做算法可优化，再做系统高效"的思路值得借鉴
  • 融合切换的正负拼接trick: 把"反合并旧适配器+合并新适配器"统一成一个操作，简洁有效
  • 2.4倍加速几乎无精度损失: 将动态适配器的延迟开销从250%+降到仅29%，实用价值很高

局限性 / 可改进方向

  • 预门控的假设需验证: "不同层的路由模式一致"这个假设在更大模型或更复杂任务上是否依然成立？论文未做分析
  • 仅优化解码阶段: 预填充阶段未优化，对长输入场景（如RAG、长文档）的加速效果未评估
  • 模型规模受限: 仅在7B模型上验证，更大模型（13B、70B）上SGMM的扩展性未知
  • 显存略增: 需要同时维护原始权重和LoRA参数用于动态合并/反合并，内存效率仍有优化空间
  • 无代码开源: 难以复现，尤其是SGMM内核这种系统级实现

与相关工作的对比

  • vs MoRAL/MOLA（layer-wise路由）: 这些方法在每层独立做路由决策，CUDA调用次数与层数成正比。AdaFuse通过全局一次路由从根本上消除了这个问题。精度相当但延迟大幅优于这些方法
  • vs PESC（block-wise路由）: PESC比layer-wise方法更高效但仍需block间的路由切换。AdaFuse在延迟上比PESC快2.7倍，精度略低0.3%（通用任务）但在特定任务上持平甚至更优
  • vs 静态LoRA: 静态LoRA可以训后直接合并进骨干网络，零推理开销。AdaFuse的理想是逼近这个上限，目前延迟开销仅29%，代价是多一次路由计算和SGMM切换

启发与关联

  • 可迁移的设计原则: "算法设计要对系统友好"这个原则可以应用到视觉Transformer的MoE适配器（如VLM的domain-specific adapter切换），以及更广泛的PEFT serving场景
  • 预门控的思想可以推广到多模态模型中不同模态使用不同适配器的场景——例如VLM中视觉token和语言token可能需要不同的LoRA专家
  • "decide-once, apply-everywhere"是一种通用的延迟优化范式——任何需要逐层/逐块动态决策的架构都可以考虑提前做全局决策

局限性 / 可改进方向

  • 仅在7B模型测试：更大规模模型（13B/70B）的效果和加速比未验证
  • Prefilling阶段未优化：仅优化解码阶段，prefilling阶段延迟与现有方法持平
  • 全局路由假设的局限：假设不同层应激活相同的专家——对某些任务可能不成立（如底层需要语法专家而高层需要语义专家）
  • 专家数量有限：N=4-16个专家，更大规模MoE的适配性未探索
  • 训练成本未充分讨论：需要微调LoRA专家+路由器的训练成本

评分

  • 新颖性: ⭐⭐⭐⭐ 算法-系统协同设计思路清晰，token级预门控+融合内核有创新性
  • 实验充分度: ⭐⭐⭐ 只在7B模型上测试，缺少更大规模验证
  • 写作质量: ⭐⭐⭐⭐ 问题分析透彻，延迟剖析图表丰富
  • 价值: ⭐⭐⭐⭐ 解决了动态适配器落地的关键瓶颈，对LLM serving有实际意义