feifeibear/long-context-attention

解決的問題

YunChang 實作了統一序列平行（USP），解決了兩種常見於長上下文 LLM 訓練與推論的分散式注意力方法——DeepSpeed-Ulysses 與 Ring-Attention 的限制。Ulysses 受限於注意力頭數（不適用於 GQA 或 MQA），而 Ring-Attention 在計算與通訊上通常效率較低，且在大規模部署時可能引發死結。

工作原理

USP 是一種混合序列平行方法，結合了 DeepSpeed-Ulysses 與 Ring-Attention 的優勢。使用者可設定 Ulysses 度與 Ring 度，將序列跨 GPU 分區。此專案整合多種注意力後端，包括 FlashAttention（v2 與 v3）、Torch Efficient，以及其他如 SageAttention 與 FlashInfer。同時提供負載平衡策略（例如 "zigzag" 與 "stripe"）用於因果遮罩，以優化效能。

適用對象

專為需要在多個 GPU 上擴展序列長度的研究人員與工程師設計，特別適用於使用 GQA/MQA 架構或異質網路硬體的長上下文生成式 AI 模型訓練或部署場景。

亮點

• 混合方法：融合 Ulysses 與 Ring 注意力，克服頭數限制與通訊低效問題。
• 廣泛硬體支援：相容 NVIDIA GPU（H100、A100 等）與 NPU。
• 靈活後端：支援 FlashAttention v2/v3、SageAttention 與 FlashInfer。
• 整合性：已應用於 NVIDIA 的 TransformerEngine，並與 Megatron-LM 框架相容。