FlashQLA – 通義實驗室開源的高性能線性注意力算子庫

FlashQLA是什麼

FlashQLA 是通義實驗室開源的基於 TileLang 實現的高性能線性注意力算子庫。FlashQLA 通過算子融合、Gate 驅動卡內序列並行及 Warp-Specialized 優化，在 Hopper 上較 FLA Triton 實現 2–3× 前向與 2× 反向加速，覆蓋 2B 至 397B 模型，提升預訓練與端側推理效率。FlashQLA 需 SM90、CUDA 12.8+、PyTorch 2.8+ 環境。

FlashQLA的主要功能

• 高性能線性注意力算子庫：面向 Qwen 全系列 Gated Delta Network（GDN）注意力層進行深度優化。
• 算子融合加速：將 GDN Chunked Prefill 的前向與反向流程進行合理的算子融合與性能優化。
• 全規格模型覆蓋：支持從 2B 到 397B 的多規格模型，覆蓋 TP1 至 TP8 場景。
• 雙層級 API 接口：提供對齊 FLA 簽名的 high-level API，以及底層 fwd / bwd 入口。
• 變長序列支持：內置 varlen 變長序列處理能力，適配真實訓練與推理數據分佈。

FlashQLA的技術原理

• TileLang Warp-Specialized Kernel：基於 TileLang 構建關鍵 fused kernel，採用 warpgroup specialization 實現數據搬運、Tensor Core 計算與 CUDA Core 計算的重疊。
• 自動化卡內序列並行（AutoCP）：利用 GDN gate 的指數衰減性質，在 TP、長序列、小頭數等場景下自動開啓卡內序列並行，提高 GPU SM 利用率。
• 滑動窗口 warmup 機制：針對具備衰減性質的線性注意力頭，僅用 6–8 個 chunk 的 warmup 即可精確獲得子序列初始狀態，捨棄修正量 M 矩陣的計算。
• 硬件友好的代數改寫：對 GDN Chunked Prefill 的前向和反向流程進行代數變換，在不影響數值精度的前提下有效降低 Tensor Core、CUDA Core 及 SFU 開銷。
• 兼顧訪存與並行的折中架構：將計算流程拆分爲兩個 fused kernel 並在中間插入 CP 預處理，避免 fully-fused kernel 在小 batch / TP 場景下 GPU 利用率低的問題。

如何使用FlashQLA

• 環境檢查：確認硬件爲 NVIDIA SM90（Hopper 架構）且軟件環境滿足 CUDA 12.8+、PyTorch 2.8+ 的要求。
• 安裝部署：從 GitHub 克隆 FlashQLA 倉庫並通過 pip 完成編譯安裝。
• 模塊導入：在 Python 中導入 chunk_gated_delta_rule 函數。
• 數據準備：準備好輸入張量 q、k、v 以及 gate 參數 g、beta，確保各張量形狀符合接口要求。
• 執行計算：調用 chunk_gated_delta_rule 並傳入對應參數，獲取輸出結果 O 和最終狀態。
• 高級配置：如需處理變長序列，可傳入 cu_seqlens 參數；如需狀態續傳，可傳入 initial_state。
• 自動優化：AutoCP 序列並行會根據 batch 大小和序列長度自動觸發，無需手動配置。

FlashQLA的關鍵信息和使用要求

• 發佈方：通義實驗室 / QwenTeam
• 開源地址：github.com/QwenLM/FlashQLA
• 硬件要求：NVIDIA SM90（Hopper 架構，如 H200）
• 軟件要求：CUDA 12.8+，PyTorch 2.8+
• 支持模型：Qwen3.5 / Qwen3.6 系列（head dim 覆蓋 64 至 8，對應 TP1 至 TP8）
• 加速效果：前向 2–3×，反向 2×（相較 FLA Triton Kernel）

FlashQLA的核心優勢

• 兼顧訪存與並行的折中架構：將計算拆分爲兩個 fused kernel 在中間插入 CP 預處理，避免 fully-fused kernel 在小 batch / TP 場景下 GPU 利用率低的問題，通過合理拆分減少 HBM 反覆讀寫中間變量的訪存開銷。
• AutoCP 自動開啓機制：僅在 batch_size × num_heads ≤ 40 或 batch_size × num_heads ≤ 56 且 seq_len ≥ 8192 時自動觸發卡內序列並行，避免不必要的冗餘計算，自適應平衡並行度與訪存代價。
• 滑動窗口 warmup 機制：用 GDN gate 的指數衰減性質，對 60–80% 的線性注意力頭僅需 6–8 個 chunk 的 warmup 可精確獲得子序列初始狀態，直接捨棄修正量 M 矩陣的計算，大幅降低 CP 預處理開銷。
• Warp-Specialized 計算重疊：基於 TileLang 的 warpgroup specialization 設計，在同一個 SM 內實現生產者與消費者 warpgroup 協同，通過 ping-pong 結構遮蓋數據搬運與 Tensor Core / CUDA Core 計算。
• 硬件友好的代數改寫：對前向和反向流程進行代數變換與化簡，在不影響數值精度的前提下有效降低 Tensor Core、CUDA Core 及 SFU 的硬件開銷。

FlashQLA的項目地址

• 項目官網：https://qwen.ai/blog?id=flashqla
• GitHub倉庫：https://github.com/QwenLM/FlashQLA

FlashQLA的同類競品對比

FlashQLA的應用場景

• 超大模型預訓練：覆蓋 397B / 122B / 35B / 27B 等全系列 Qwen 模型，支持 256K 長上下文訓練，顯著降低注意力層在端到端訓練中的算力與時間開銷
• 端側 agentic 推理：針對 batch_size=1、小尺寸模型（如 2B / 0.8B）的 chunked prefill 場景，通過 AutoCP 提升小頭數下的 GPU 利用率，加速端側 Agent 實時響應
• 大模型線上部署：在 TP（Tensor Parallelism）場景下處理 coding agent 等長序列輸入，解決 chunked prefill 開不出足夠大 batch 時的 GPU 利用率瓶頸，提升服務吞吐
• 通用 GDN / 線性注意力架構加速：適用任何基於 Gated Delta Network 或線性注意力架構的 LLM 訓練與推理，提供開箱即用的高性能算子替換方案