专题演讲嘉宾：章程瑞东

自 Qwen-3.5 采用混合注意力结构以来，Qwen 系列模型的线性注意力层 Gated Delta Network（GDN）在训练与推理中的开销日益显著。现有实现面临两大难题：一是多步算子需要反复读写 HBM，访存瓶颈严重；二是状态递推的串行性使算子并行度受限于注意力头数，在长序列、小 batch 或 TP 场景下 GPU 利用率低下。为此，我们提出 FlashQLA —— 基于 TileLang 的高性能线性注意力算子库。核心方案包括：将 GDN Chunked Prefill 的前向与反向流程各拆分为两个融合算子，通过数据复用降低访存开销，并在算子内通过 warp-specialization 实现不同计算的重叠；基于门控衰减的自动化卡内序列并行，并建立数学模型动态调节并行粒度；针对具有指数衰减特性的注意力头，舍弃昂贵的修正矩阵计算，改用轻量级预热获得精确状态。在 Hopper 显卡上，FlashQLA 相比 FLA Triton kernel 实现前向 2–3 倍、反向 2 倍加速，且在 TP 增大时加速比持续提升，显著优化了预训练与端侧推理效率。

演讲提纲

1. 背景与挑战

• Qwen 模型训练与推理 Prefill 过程中 GDN 模块的耗时
• 性能瓶颈剖析：
  • 访存密集型：分步 kernel 大量读写 HBM，中间变量开销大
  • 并行度受限：状态递推依赖导致 thread block 数量不足，SM 利用率低（尤其在 TP、小 batch 场景）

2. 优化方案

• 设计目标：兼顾访存优化与序列并行，避免极端方案（Fully-fused / 多步算子）的缺陷
• 核心策略：
  • 拆分为两个 fused kernel，在 kernel 间插入 CP 预处理，平衡数据重用与并行度
  • 门控驱动的动态并行：根据 batch_size × num_heads 和序列长度自动开启卡内序列并行
  • 利用门控衰减性质：采用 warmup 替代昂贵的修正矩阵计算，大幅降低预处理开销

3. 系统设计思路

• 对前向与反向流程做等价的代数改写，降低计算复杂度
• 从门控衰减到滑动窗口及背后的代数原理
• 自动卡内序列并行的参数设计

4. 算子实现细节

• 计算与访存负载分析
• 片上资源分析
• 如何用 Tilelang 手动实现 warp-specialization

5. 性能评估

6. 开放问题

• GDN 模块剩余的优化空间
• KDA 及其他线性注意力结构的优化方案
• 芯片硬件方面对线性注意力算子的限制

实践痛点

• 不能覆盖所有场景
• 算子复用性较差

听众收益

• 线性注意力模块在系统上的一些性质
• 与算法创新相结合的系统优化流程
• 复习线性代数