Qwen 把线性注意力的推理效率压到了新台阶,2-3 倍加速对想做本地 Agent 的开发者是实打实的,不是论文灌水,是能跑在设备上的代码。
FlashQLA是基于TileLang开发的高性能线性注意力内核,专为提升个人设备上智能体AI性能而设计。它实现了2-3倍的前向传播加速和2倍的反向传播加速。其核心技术包括门控驱动的片上自动计算与通信重叠、硬件友好的代数重构,以及TileLang融合的Warp专用内核。该设计通过自动片上通信重叠显著提升了流处理器利用率,在张量并行、小模型和长上下文任务中效果突出。尽管在大批量处理时,其将GDN流程拆分为两个内核的策略会带来额外内存开销,但在边缘设备和长上下文实际场景中性能更优。反向传播部分通过构建16级、严格片上内存限制下的Warp专用流水线,实现了超过2倍的内核级加速。
🚀 推出 FlashQLA:基于 TileLang 构建的高性能线性注意力内核。
⚡ 前向 2–3 倍加速。后向 2 倍加速。 💻 专为个人设备上的智能体 AI 打造。
💡关键洞察: 1. 门控驱动的自动卡内通信原语。 2. 对硬件友好的代数重构。 3. TileLang 融合的 warp 专用内核。
FlashQLA 通过自动设备内通信原语提升了 SM 利用率。对于 TP 设置、小模型和长上下文工作负载,加速效果尤为明显。
我们没有将整个 GDN 流程融合到单个内核中,而是将其拆分为两个内核,分别针对通信原语和后向效率进行了优化。在大批量下,这会比完全融合的方法带来额外的内存 I/O 开销,但在边缘设备和长上下文工作负载上却实现了更好的实际性能。
后向传播是最难的部分:我们在极其紧张的片上内存约束下构建了一个 16 阶段 warp 专用流水线,最终实现了 2 倍以上的内核级加速。
我们希望这对社区有用!🫶🫶 了解更多: 📖 博客:https://qwen.ai/blog?id=flashqla 💻 代码:https://github.com/QwenLM/FlashQLA