FlashQLA是基于TileLang构建的高性能线性注意力内核,专为个人设备上的智能体AI设计。其核心创新包括门控驱动的自动片内计算并行、硬件友好的代数重构以及TileLang融合的Warp专用内核,通过提升流处理器利用率,在前向传播上实现2-3倍加速,反向传播实现2倍加速。该技术在小模型、长上下文工作负载和张量并行设置中效果显著,虽然在大批次处理时内存I/O开销略高,但在边缘设备和长上下文场景中实际性能更优。反向传播通过16级Warp专用流水线在严格片上内存限制下实现了核心级加速。相关资源已开源。
2 倍加速的背后是 Warp 特化流水线和自动 Copy 策略,像给手机 GPU 开了条专用跑道,做端侧 Agent 的可以直接拉代码试试。
🚀 Introducing FlashQLA: high-performance linear attention kernels built on TileLang.
⚡ 2–3× forward speedup. 2× backward speedup.
💻 Purpose-built for agentic AI on your personal devices.
💡Key insights:
1. Gate-driven automatic intra-card CP.
2. Hardware-friendly algebraic reformulation.
3. TileLang fused warp-specialized kernels.
FlashQLA boosts SM utilization via automatic intra-device CP. The gains are especially pronounced for TP setups, small models, and long-context workloads.
Instead of fusing the entire GDN flow into a single kernel, we split it into two kernels optimized for CP and backward efficiency. At large batch sizes this incurs extra memory I/O overhead vs. a fully fused approach, but it delivers better real-world performance on edge devices and long-context workloads.
The backward pass was the hardest part: we built a 16-stage warp-specialized pipeline under extremely tight on-chip memory constraints, ultimately achieving 2×+ kernel-level speedups.
We hope this is useful to the community!🫶🫶
Learn more:
📖 Blog: https://qwen.ai/blog?id=flashqla
💻 Code: https://github.com/QwenLM/FlashQLA
来源:Qwen · x.com