設計工具

Invalid input. Special characters are not supported.

洞察

一百萬個詞元的上下文:優點、缺點與挑戰

Wes Vaske、Katya Giannios | 2025 年 5 月

最近,出現了一系列有關「記憶體」和上下文長度的 AI 公告。

在這篇部落格中,我將回顧這些公告的內容,以及它們對設計 AI 工廠的系統架構師意味著什麼。

長上下文是一種優勢

我們大多數人與大型語言模型(LLM)互動時,會輸入查詢資訊,然後得到回覆。輸入和輸出的總大小就是查詢的「上下文視窗」。在聊天用例中,輸入和輸出通常在 100 到 1,000 個詞元範圍內。

隨著硬體效能的提高和支持更大上下文視窗的模型的開發,我們發現在一些用例中,上下文視窗非常重要。

我每天都會與我的編碼助手互動。當我需要生成程式碼時,我希望模型能理解我現有的程式碼庫。我可以在生成請求的上下文中加入現有的程式碼,以滿足這項需求。透過使用更多程式碼作為上下文,編碼助手能以更高的層次運作,並對我的程式碼庫進行更大規模的架構調整。

實際上,更長的上下文視窗能夠使 AI 編碼助手變得更聰明。

其他用例還包括存取電子郵件、SharePoint 網站、醫療記錄、法律程序等。由此不難看出,擁有更多可用資料能夠獲得更好的輸出。

什麼是預填充和 KV 快取?

我們已經知道為什麼需要更大的上下文視窗,現在,我們需要介紹現代 LLM 的一個基礎概念——鍵值(key-value,簡稱 KV)快取。

目前使用的所有主要 LLM 均採用 Transformer 架構,每個推論查詢都遵循相同的步驟:

嵌入 -> 預填充 -> 解碼 -> 後處理

嵌入階段包括將文字查詢轉化為 AI 模型可以操作的向量表示。

預填充階段包括透過模型執行查詢(以及任何提供的上下文,例如輔助文字或程式碼行),並透過評估每個子詞來進行處理,進而生成內部表徵。這些表徵的一部分是一組稱為鍵(K)與值(V)的向量

解碼階段是最有趣的地方。它從預填充階段獲取資料,並根據查詢開始生成新詞元(單詞)。預填充階段的 K 和 V 向量可以針對每個新生成的詞元重新計算,也可以快取並重複使用。在解碼過程中,我們會建立新的 KV 向量,並同時利用先前的向量來計算每個新詞元。

最後,後處理階段會將解碼階段生成的詞元轉換為人類可讀的文字。

從這一流程拆解中可以看出,在預填充階段完成之前,LLM 無法開始生成新的詞元。這項限制代表,若預填充耗時過長,將會直接影響使用者體驗。

預填充需要記憶體和時間

我的同事 Katya Giannios 是一名應用數學博士,她一直在為推論的系統架構建模,以便我們可以估算出某些情況下的預填充時間。

記憶體和時間表

免責聲明:數字均為估計值;我們仍在實驗室進行測試,以更加瞭解系統。

下表顯示以下模型的推論效能:
  • Meta 的 Llama 4 Maverick 模型
    • 4,000 億個參數
    • FP8 權重
    • 最大上下文視窗為 1,000,000 個詞元
  • 運行於 NVIDIA DGX B200 伺服器
    • 8 個 NVIDIA B200 GPU
    • 1.4TB HBM3E
  • 單使用者
  • 1,000 個詞元輸出

我們使用單個使用者以簡化操作,因為多個並發使用者會大大增加負載和 KV 快取大小。

我們可以看到,模型可以輕鬆處理 10,000 個詞元(甚至更多)以內的小型上下文,預填充時間不到一秒。但是,當上下文長度達到最大值時,LLM 開始生成任何輸出之前的預填充時間(出現第一個詞元的時間)超過 2 分鐘。

如果我們乘以 10 個使用者,每個使用者的上下文長度都是 250,000,則預填充時間將超過 30 秒,KV 快取總大小將達到 39GB。

這就是大上下文讓人頭痛的一面——計算長上下文的 KV 值會嚴重影響使用者體驗。除非我們能改善預填充時間,否則長上下文將不適合互動式用例。

為什麼要重複使用 KV 快取?

回到我最初舉的使用長上下文例子——AI 編碼助手,我們預期連續查詢的上下文會有大量重複使用。當編碼助手在類別中生成方法時,每次查詢都會存取基類。

與其每次查詢都重新生成 KV 快取,不如只生成一次,然後在後續查詢中重複使用。然而,這種方法面臨的挑戰是 KV 快取的體積過大。即使在最新的 LLM 中進行最佳化,KV 快取仍會耗盡 AI 系統的全部記憶體。

卸載拯救了我們!

NVIDIA Dynamo 庫具有 KV 快取管理器,可將 KV 快取從 GPU 記憶體遷移到其他可用裝置。它還實施 KV 重用技術,將 KV 快取轉化為多個會話和使用者共享的 KV 池。我們討論的從 HBM 遷移過來的 1,000,000 個詞元的 KV 快取(約 15GB)只是更大 KV 池的一小部分。

第一站是系統記憶體。DGX 範例支持 4TB 系統記憶體,到達 GPU 的總頻寬為 1 TB/s。憑藉如此大的頻寬,將 100 萬個詞元的 KV 快取從 CPU 記憶體載入 GPU 記憶體(在理想情況下)只需要 15 毫秒,而重新運算則需要 2 分鐘以上,這一點我們在前面提到過。

時間上的縮短對使用者來說意義重大!對上下文進行初始運算後,從 CPU 記憶體重新載入的速度非常快,可以實現互動式使用者體驗。

CPU 系統記憶體仍然面臨與 GPU 記憶體相同的問題——容量有限。

下一站是 DGX 伺服器中的本地 NVMe 硬碟。使用八個 PCIe Gen5 NVMe 硬碟(如美光 9550 高效能 Gen5 硬碟),下一層 KV 快取卸載容量可從 30TB 到 245TB 不等,總頻寬為 112 GB/s。

雖然儲存層的速度明顯比系統記憶體慢,但載入 1,000,000 個詞元的 KV 快取仍然只需 140 毫秒。

使用這些卸載和遷移技術可以縮短生成第一個詞元的時間,進而改善使用者體驗。由於 GPU 將更多的時間用於生成輸出,而不是重做以前做過的工作,因此對於改進總體擁有成本(TCO)很有幫助。

高效能儲存裝置實現長上下文推論

過去幾年,生成式 AI 的格局發生了很大變化。以前,儲存裝置是配角,現在看來,為 AI 工廠提供高效能儲存裝置將是實現具有良好使用者體驗的智慧型 AI 的關鍵。

在去年的 FMS 和今年春季的 GTC 上,美光展示了我們即將推出的 PCIe Gen6 NVMe 硬碟的效能。圍繞長上下文與 KV 快取管理的這些發展顯示,將最快速的 NVMe 快閃記憶體連接至最新 GPU,將成為成功部署 AI 的關鍵。

 

SMTS 系統效能工程師

Wes Vaske

Wes Vaske 是 Micron Technology 的資深技術人員(SMTS),以及系統效能工程師。憑藉著在儲存解決方案和 AI 基礎架構方面的深厚經驗,Wes 在推動美光的資料人工智慧和機器學習功能中扮演了關鍵角色。他以在進行 AI 訓練系統基準測試和最佳化儲存裝置效能,以滿足次世代 GPU 需求方面的專業知識而聞名。加入美光之前,Wes 曾擔任 Dell 的系統工程師。他擁有愛荷華州立大學的學士學位。

機器學習主任工程師

Katya Giannios

Katya Giannios 是美光科技(Micron Technology)技術與產品團隊(TPG)可擴充記憶體系統路徑探索與策略小組的機器學習工程師。她的工作內容包括商業工作負載的分析與特性剖析,並著重在 AI 軟硬體相輔設計的應用。Katya 曾於德國慕尼黑工業大學(TUM)與德國慕尼黑馬克斯普朗克等離子物理研究所合作的計畫攻讀學位,並取得應用數學博士學位。