斯坦福 CS336 從頭開始的語言建模：推理引擎與全棧創新

斯坦福 CS336 從頭開始的語言建模：推理引擎與全棧創新

推理：將電力轉化為智慧的引擎

推理是關鍵階段，將訓練好的語言模型從靜態的數學對象（操作的有向無環圖）轉變為可實際使用的工具。了解底層的推理引擎與 GPU 核心程式碼對於「全棧創新」至關重要，因為這些組件決定了模型在生產環境中的實際效能、延遲與功能。

Token 的生命週期

當向推理系統發出請求時，會遵循以下特定流程：

  1. 排程 (Scheduling)： 請求被導向特定的 GPU，可能會將 prefill 與 decode 操作分別安排在不同機器上。
  2. KV 快取查詢 (KV Cache Lookup)： 系統檢查請求或其變體是否曾出現過，以節省計算資源。
  3. 執行 (Execution)： 核心機器學習程式碼被執行，通常會根據模型大小在節點或 GPU 之間平行化。
  4. Token 生成 (Token Generation)： 系統輸出 token，隨後會進行停止序列與安全性檢查。

Prefill 與 Decode 工作負載

推理由兩個本質上計算特性截然不同的階段組成：

  • Prefill： 此階段處理初始輸入提示（例如 10,000 個 token），計算初始激活。它是 計算受限 (compute‑bound)，類似於訓練（但不包含反向傳播），需要大量 FLOPs。
  • Decode： 此階段逐一生成 token。它是 記憶體頻寬受限 (memory‑bandwidth‑bound)，因為必須從記憶體載入整個模型才能產生單個 token，儘管所需 FLOPs 相對較少。

由於這些差異，現代堆疊常將 prefill 與 decode 分離到不同的工作者或專用硬體上（例如使用 NVIDIA GPU 處理 prefill，使用 LPU/Grok 晶片處理 decode）。

生產挑戰與系統最佳化

每日服務數兆 token 會顯露出在小規模測試時不會出現的細微錯誤與瓶頸。這些包括因核心程式錯誤導致模型無限重複 token 的「毀滅迴圈」，或因 kernel 讀取未初始化的 GPU 記憶體而產生的語言突變（例如從英文突變為中文）。

KV 快取管理

為了最大化吞吐量，系統必須有效管理鍵值（KV）快取。由於 GPU 記憶體有限，通常採用分層儲存方式：

  • GPU 記憶體： 存取速度最快，但空間最受限。
  • CPU DRAM： 較慢，用於卸載較少使用的激活。
  • SSD/磁碟： 最慢，用於長期保存會話資料。

這類似於傳統作業系統的記憶體管理，常使用最近最少使用 (LRU) 策略決定哪些激活要被驅逐至較慢的儲存層。

快取感知的分散化

一個簡單但有效的最佳化是根據快取命中率路由請求。將「新鮮」請求（快取命中率低、prefill 成本高）送至一組 GPU，將「溫熱」請求（快取命中率高、prefill 成本低）送至另一組 GPU，服務速度可提升最高約 40%。

Megakernel：實現「光速」推理

傳統推理引擎會逐一執行操作（例如先執行 Norm kernel，再執行 MatMul kernel）。這會因 kernel 啟動開銷與「尾部效應」而產生大量空閒時間，GPU 必須等批次中最長的序列完成後才能開始下一個操作。

Megakernel 方法

Megakernel 將多個操作融合成單一 kernel，將 GPU 視為大型分散式系統而非順序運算子。這允許操作重疊，例如：

  • 在 QKV 投影與 RoPE 縮放仍在執行時，同時載入 KV 快取。
  • 在注意力運算尚未完全結束前，就預先載入 O 投影的權重。

使用 ThunderKittens 函式庫進行低階 CUDA 控制，Megakernel 可達到接近「光速」的效能，在 H100 GPU 上實現高達 72% 的頻寬利用率。

Parcae：穩定的遞迴架構

雖然大多數當前 LLM 透過增加參數與資料規模來擴展，Parcae 研究則探索透過 遞迴（在 transformer 中循環區塊）來擴展模型。

穩定性問題

簡單的循環 transformer 往往極不穩定，若學習率等超參數稍有調整就會出現損失尖峰與 NaN。這種不穩定源於變換矩陣的譜半徑；若矩陣範數大於 1，激活在循環中會指數級爆炸。

Parcae 解決方案

Parcae 透過重新參數化遞迴系統的 A 與 B 矩陣來穩定訓練：

  • 矩陣 A： 限制為負對角矩陣，確保系統穩定，激活最終會衰減而非爆炸。
  • 矩陣 B： 以簡單的線性範數加以控制。

此數學約束保證了平滑的損失曲線，使模型在每個參數上能達到比傳統 transformer 更高的品質。

遞迴的擴展律

初步發現顯示遞迴應與資料與參數一起擴展。具體而言，隨著訓練資料量增加，遞迴次數也應相應提升，以維持計算最佳化。這暗示在固定參數預算下，透過循環預訓練可能產生更高品質的模型。

摘要：客座講師 Dan Fu 討論了推理引擎與 GPU 核心程式碼在將大型語言模型從數學對象轉化為可用智慧方面的關鍵角色，並介紹了 Megakernel 與 Parcae 遞迴架構等最佳化手段。

標題：斯坦福 CS336 從頭開始的語言建模：推理引擎與全棧創新

Sources