戳破GPU泡沫：Moondream的流水線解碼技術

近日，Moondream團隊公佈了其推理引擎Photon的技術細節。Photon通過一種稱為“流水線解碼”的技術，顯著提升了視覺語言模型（VLM）的推理速度。在NVIDIA B200 GPU上，該引擎實現了約33毫秒的近實時推理，並將解碼吞吐量提高了35%。

傳統AI模型生成文本時，每個令牌（token）的生成都依賴前一個令牌，這種自迴歸特性導致瞭解碼循環中的CPU與GPU往返通信。GPU負責大部分數學運算，但CPU也需要完成請求調度、元數據設置、令牌採樣等任務。由於每個令牌的GPU計算量較小，而CPU的固定開銷在每個循環中都存在，因此GPU經常處於空閒等待狀態，這種現象被稱為“GPU泡沫”。

Photon通過流水線解碼解決了這一問題。其核心思想是讓CPU和GPU的工作重疊進行：在GPU處理當前令牌的前饋計算時，CPU可以同時處理上一個令牌的提交工作。為此，Photon引入了三種關鍵技術。

首先，乒乓槽位機制：為了支持兩個步驟的重疊，系統維護兩個解碼槽位，交替使用。每個槽位包含一組固定的緩衝區，包括輸入、輸出和採樣結果的內存區域。CPU在處理一個槽位的結果時，GPU可以同時運行另一個槽位的前饋計算。內存複製通過獨立的複製流進行，避免阻塞。

其次，先前饋後採樣機制：前饋計算不依賴約束解碼的掩碼，因此可以提前啓動。採樣步驟則依賴上一個令牌的提交結果，因此延遲到提交完成後進行。這種“提交-再最終化”的順序確保了掩碼的正確性，同時將提交工作從關鍵路徑中移除。

第三，殭屍機制：當序列提前結束時（如遇到結束令牌），該序列仍可能被包含在已啓動的前饋步驟中。Photon通過引用計數和“最終化早、釋放晚”的策略處理這些“殭屍”序列：序列被標記為最終化，但資源在引用計數歸零後才釋放，從而避免了複雜的取消邏輯。

此外，預填充（prefill）階段也被整合到同一流水線中。預填充是對新請求的提示和圖像的一次性處理，通常計算量大。Photon將其視為一種特殊類型的“前饋”步驟，與解碼步驟共享兩個槽位。這使得CPU和GPU的工作可以進一步重疊，特別是在輸出較短的場景中，顯著提高了整體效率。

總之，Photon通過精心設計的流水線解碼技術，有效消除了GPU泡沫，為實時AI推理提供了新的可能性。該引擎的設計思路對於優化GPU利用率和提升推理性能具有重要參考價值。