現代 AI 應用越來越依賴于將龐大的參數數量與數百萬個令牌的上下文窗口相結合的模型。無論是經過數月對話的 AI 智能體、通過像整個百科全書一樣大的案例法進行推理的法律助理，還是在不斷擴展的資源庫中對 copilot 進行編碼，保持遠程上下文對于相關性和一致性至關重要。最重要的是，用戶期望快速的交互式響應。
對如此龐大的數據進行解碼以及讓多個 GPU 快速擴展并相互通信的需求不斷增長，這凸顯了 FP4 計算和 NVIDIA Blackwell 系統提供的高帶寬大型 NVLink 域的重要性。本博客中介紹的 Helix Parallelism 是與 Blackwell 共同設計的。與使用超長上下文進行實時解碼的最知名的先前并行方法相比，它在給定延遲下的并發用戶數量增加了 32 倍。
換言之，它讓 AI 智能體和虛擬助理能夠以前所未有的速度為更多人提供服務。
(注意：此博客中的上下文是指之前生成的令牌序列，其中間鍵和值表示形式存儲為 KV 緩存，并在每個解碼步驟中進行訪問。)
解碼瓶頸：KV 緩存和 FFN 權重讀取
為支持大規模實時解碼，系統必須克服解碼 (即生成) 階段的兩大瓶頸：
- 鍵值 (KV) 緩存流式傳輸：在處理數百萬個令牌的上下文時，每個 GPU 必須從每個樣本的 DRAM 中讀取大量過去令牌 ( KV 緩存) 的歷史記錄。這種恒定的流式傳輸反過來會使 DRAM 帶寬飽和，增加 token 到 token 的延遲 (TTL) ，并隨著上下文長度的增長迅速成為主要的瓶頸。
- 前饋網絡 (FFN) 權重加載：在自回歸解碼期間，生成每個新 token 都需要從 DRAM 加載大型前饋網絡 (FFN) 權重。在小批量的低延遲場景中，這種內存訪問成本無法很好地分攤，這使得 FFN 權重讀取成為主要的延遲來源。
使用傳統的并行策略，很難同時優化 KV 緩存流和 FFN 權重加載這兩個瓶頸。
我們以張量并行 (TP) 為例：增加 TP 可以通過在多個 GPU 之間分配權重負載和改進 TTL 來幫助減少 FFN 停頓，但只能達到一定程度。在 Llama 模型中使用的 Grouped Query Attention (GQA) 或 DeepSeek 模型中的 Multi-Latent Attention (MLA) 等注意力方案中，多個查詢頭共享有限數量的 KV 頭。如圖 2 (c) 所示，當 TP 超過 KV 頭的數量時，系統最終會跨 GPU 復制每個樣本的數百萬個令牌 KV 緩存，以實現自注意力。因此，即使 TP 增加，KV 讀取量仍然很高，再次使 DRAM 帶寬飽和，并限制了可擴展性。就 MLA 而言，TP 的上限僅為 1，以避免 KV 緩存重復。
那么，開發者如何在不犧牲實時交互性的情況下擴展模型大小和上下文長度呢？Helix Parallelism 提供了一條前進之路。
Helix 執行流程
Helix 是一種混合分片策略，可在時間工作流中分解注意力和 FFN 的并行策略，有效解決數百萬個令牌解碼期間的 KV 緩存和 FFN 權重讀取瓶頸。
圖 1 (下圖) 顯示了 Helix 如何在單個 Transformer 層中編排注意力和 FFN 的執行。受 DNA 螺旋結構的啟發，Helix 將多個并行維度 ( KV、張量和專家) 交織成一個統一的執行循環。通過將用于注意力和 FFN 的并行策略解，Helix 允許每個階段在針對其自身瓶頸進行調整的配置中運行，同時重復使用相同的 GPU 池。Helix 重用方法可保持 GPU 在各個階段的高效利用，消除計算流經模型時的空閑時間。
注意力階段
Helix 通過沿 KVP GPU 序列維度對數百萬個令牌的 KV 緩存進行分片來應用 KV 并行 (KVP) ，同時在注意力頭 (TPA) 中應用張量并行 (TPA) ，其中 TPA 是每個 QKV 投影分割的 GPU 數量，并保持小于等于 KV 頭數量以避免重復。
圖 2 (d) 中還通過簡化的玩具示例說明了這種分片策略。這將導致總共 N* 個 KVPxTPA GPU 協同進行注意力計算，而不會在 GPU 之間重復 KV 緩存。此處的 N 表示用于端到端推理執行的 GPU 總數。同一組 N 個 GPU 將在 FFN 階段重復使用。
為避免注意力前的全集，Helix 可確保每個 KVP GPU 擁有與其本地 KV 頭相關的所有查詢頭，并冗余計算 QKV 投影。這將在每個 KV 分片上啟用完全本地的 FlashAttention。
在本地 FlashAttention 之后，沿著 KVP GPU 中的查詢頭維度的單個多對多，交換部分注意力輸出和對數和表達式標量。重要的是，此通信成本隨批量大小和隱藏維度而擴展，并且獨立于 KV 緩存長度，即使上下文長度擴展為數百萬個令牌，也能高效運行。每個 GPU 在本地重建 SoftMax 標準化的精確輸出。
這種多對多還會觸發將 KVP GPU 重新分配為 TP 組 (TP = N+ KVP x TPA) ，用于注意力輸出線性計算。至關重要的是，這種多對多階段受益于 NVLink/ NVL72 的高帶寬互連，可實現大型 GPU 數量之間的快速集合通信。圖 2 概要介紹了關注的不同分片方案及其相應布局。
為進一步降低 TTL，我們引入 Helix 重疊工作流 – 批量 (HOP-B) ，這是一種細粒度工作流技術，可重疊整個批量的通信和計算，如圖 3 所示。計算一個 token 的注意力輸出后，Helix 會啟動該 token 的多對多交換，同時計算下一個 token 的注意力。這種緊密重疊隱藏了有用工作背后的通信延遲，保持了較高的 GPU 利用率，并進一步加速了實時解碼。
在圖 3 (上圖) 中，頂部顯示了在沒有 HOP-B 的情況下，八個請求以鎖步方式執行注意力計算。然后是順序多對多通信。圖表底部顯示，使用 HOP-B 時，一個請求的通信與下一個請求的計算重疊，通過細粒度流水線減少 TTL。
FFN 階段
注意力集中后，系統會重新調配相同的 N* KVPxTPA GPU 池，而不會在空閑時間內執行 FFN 塊。多對多步驟的輸出已按隱藏維度在 N 個 GPU 中進行劃分，因此注意力后線性投影可在 TP 模式 (TP = N) 下立即運行。每個 GPU 使用其權重分片執行局部矩陣乘法運算，并參與跨 TP = N GPU 的全局歸約，以構建正確的輸出。
在注意力后線性投影之后，Helix 使用密集模型中的 1D TP (N+ TPF) 布局或 MoE 模型中的 2D TP x 專家并行 (N+ TPFxEP) 網格重新配置相同的 N 個 GPU 以進行 FFN 計算。
分布式 KV 連接
在解碼期間，每個新 token 都會廣播到所有 KVP GPU 進行查詢計算。為了防止出現 DRAM 熱點，Helix 以循環方式在 KVP 等級中錯開 KV 緩存更新，例如，令牌 1 到 16 轉到 KVP 0，17 到 32 轉到 KVP 1，以此類推。這可確保統一的 KV 增長、平衡 GPU 之間的顯存使用情況，并保持一致的吞吐量，而不管序列長度或批量大小如何。
Blackwell 上的模擬結果
Helix 為長上下文 LLM 解碼設置了新的性能基準。圖 4 (下圖) 顯示了在使用 (假設) 100 萬個令牌的上下文進行解碼時，DeepSeek-R1 671B 的標準化吞吐量延遲帕累托邊界。通過對數千種配置進行全面模擬、系統化地改變模型分區策略 ( TP、EP、PP 和 KVP) 和批量大小，以找到吞吐量 – 延遲的最佳權衡，Pareto 前沿得到了推導。具體而言：
- 對于固定的延遲預算，Helix 可以將并發用戶的數量增加多達 32 倍 (即，使 token/s/ GPU 增加 32 倍) 。
- 對于低并發設置，Helix 可將用戶交互性提升高達 1.5 倍 (即，將可實現的最小 TTL 降低高達 1.5 倍) 。
通過在所有可用設備中對 KV 緩存和 FFN 權重進行分片，顯著降低 DRAM 壓力并提高計算效率，實現了這些提升。Helix 推動了吞吐量延遲方面的發展，即使在較低的延遲下也能實現更高的吞吐量。如需了解更多詳情，請參閱此處的論文。
敬請關注
與 Blackwell 的最新功能共同設計的 Helix Parallelism 為在不影響交互性的情況下大規模提供數百萬個令牌模型提供了藍圖。請持續關注我們將此優化引入推理框架。
?
