加速 Roblox 上的 3D 創作人工智慧推論

• Roblox 導入 CUDA 圖形運算與 KV 快取技術，以加速 3D 網格生成，實現更流暢的迭代。
• 在推出時，Cube 3D 模型可在 7.8 毫秒內生成代幣（從 60.5 毫秒縮短），並在 4 秒內生成完整物件（從 31 秒縮短）。 

今年稍早，Roblox 分享了我們 Cube 3D 基礎模型的首項功能。透過 Cube 3D，創作者可直接根據文字提示生成 3D 模型與環境。 我們從一開始就將延遲優化列為優先事項，因為我們意識到，緩慢的生成時間會打斷這個本質上屬於迭代的過程。在 Cube 3D 於三月推出之前，我們已經將推論步驟的速度提升了 7.8 倍，並為開發者和使用者帶來更流暢的體驗。 

自推出以來，已有超過 578,000 個物件在多個知名體驗中被生成。開發者也表達了希望讓使用者能在體驗中透過「貓」、「漢堡」等文字提示生成 3D 物件的意願。最值得注意的是《Mic Up》——這款使用語音聊天的熱門聚會遊戲，便運用 Cube 3D 為玩家提供了一種有趣且互動的物件生成方式。 在該遊戲的實作中，玩家可開啟左側選單以存取額外功能，其中包括一個 AI 圖示。點擊該圖示後，玩家即可輸入文字提示來生成 3D 物件。對使用者而言，過長的生成時間會造成使用阻礙，使他們無法體驗到創意即時轉化為 3D 模型的魔力。 

我們希望將 3D 生成體驗從「停頓等待」的互動模式，轉變為反應靈敏且自然的體驗，從而實現快速實驗。能夠迅速將物件加入場景中，對開發者而言至關重要。為了加速 Cube 3D，我們首先對推論流程進行了效能分析，以找出效能瓶頸。儘管使用了強大的 GPU，我們仍發現操作之間存在顯著的閒置時間。

現代深度學習框架依賴 CPU 來調度並在 GPU 上執行運算（或稱核心）。CPU 會預先準備每項運算，將其傳送至 GPU，並在等待確認後才開始準備下一項運算。這種等待會造成調度瓶頸，導致 CPU 在準備下一批工作時，GPU 可能處於閒置狀態。理想情況下，我們希望 CPU 能領先於 GPU 運作，預先準備並排入運算任務，使 GPU 始終有工作可處理。

對於 Cube 3D 等變壓器型模型中的自迴歸解碼器而言，這尤其成問題，因為這些模型需要依序處理輸入並生成標記。這些模型單次生成便需執行數千個獨立運算，且隨著序列中的每個步驟推進，運算開銷會不斷累積。

「我們希望打造能實現四維互動的系統，」工程副總裁 Anupam Singh 解釋 Roblox 選擇自迴歸方法的原因時表示。「我們不僅想造出一輛車，還希望能夠打開車門並坐進車內。」

每項操作所耗費的：

• 準備每個核心所需的 CPU 時間
• 啟動核心所產生的開銷
• GPU 執行時間（實際運算）
• 檢查完成狀態時的同步開銷

對於在 GPU 上執行迅速的小型運算而言，此開銷可能佔據推論時間的大部分。GPU 實際進行運算的時間，可能僅佔總推論時間的一小部分。

為解決此瓶頸，我們運用了 CUDA Graphs——這項功能可讓使用者在不需 CPU 介入的情況下，記錄並重播一連串的 GPU 運算序列。Cube 3D 架構中的自迴歸解碼器元件會處理文字提示，並透過固定長度的向量產生形狀標記。 

雖然在功能上類似於傳統的大型語言模型（LLM），但我們的雙流解碼器架構有一個重要差異——它使用兩個並行注意力流。其中一個流專門處理條件標記，另一個則專門處理形狀標記。現成的 LLM 推論引擎無法滿足我們的需求，因此我們需要針對特定架構量身打造的自訂實作。

不妨將 CUDA Graphs 視為為 GPU 錄製的巨集。與其讓 CPU 逐一發出指令，它會先記錄整段 GPU 運算序列（即圖），並透過單一 CPU 指令啟動整個圖。此方法大幅減少了核心啟動的開銷，因為在推論過程中，CPU 無需再逐一調度每項操作。一旦圖被啟動，GPU 便會自主執行整段序列，無需等待進一步的指令。

CUDA 圖確實存在一些限制。由於圖結構需要預先確定，因此要求批次大小和輸入維度必須固定。這意味著必須針對每種批次大小或輸入形狀建立獨立的圖。對於我們 Cube 3D 的應用案例而言，這項限制尚可接受，因為我們可以針對常見的輸入維度來標準化推論流程。

為了在 Cube 3D 模型中實作 CUDA 圖，我們必須調整方法。在傳統的大型語言模型（LLMs）中，注意力運算總是針對相同的序列長度進行，因此提供了一個靜態的處理形狀。然而，在我們的自訂雙流架構中，部分注意力層僅針對序列長度運作，而其他層則同時針對序列長度與條件長度進行運算。

儘管面臨這些挑戰，在實作 CUDA Graphs 後，我們仍獲得了顯著成果。我們使用「每輸出代碼單元所需時間」（TPOT）來衡量推論過程中每個代碼單元的生成時間。實作 CUDA Graphs 後，我們的 TPOT 從 60.5 毫秒改善至 20.5 毫秒，提升了 2.9 倍。整體生成時間則從 31 秒縮短至 10.5 秒，減少了 66%。

為了進一步降低延遲，我們實作了 KV 快取，這是在大型語言模型（LLM）推論中的一項標準做法，並已在業界證明極為有效。 

在 Cube 3D 等基於 Transformer 的模型中，每次生成標記都需要根據先前生成的所有標記來計算鍵 (K) 和值 (V) 矩陣。隨著序列越來越長，為每個標記重新計算這些矩陣的效率會越來越低。

KV 快取透過以下方式解決此問題：

1. 儲存所有先前已生成字元所對應的 K 矩陣與 V 矩陣
2. 僅針對新標記計算 K 和 V 矩陣
3. 將這些新矩陣追加至已快取的值中

此方法可消除冗餘運算，減少每個新代碼片段所需的處理工作量。隨著生成的序列越來越長，此優化效果將更加顯著。

我們將 KV 快取與 CUDA 圖演算實作整合的方法，與傳統的大型語言模型（LLM）推論類似。加入 KV 快取後，我們的 TPOT 縮短至僅 7.8 毫秒。整體生成時間減少了 87%，從原本的 31 秒縮短至僅 4 秒。這顯著的時間縮減，讓創作者使用此工具時能獲得更佳的效率。

我們正在探索能進一步降低延遲並提升使用者體驗的技術，包括最佳化的核心模組、用於加速推論的模型量化、針對特定硬體的優化，以及平行標記生成。

當我們擴展至完整的場景生成與理解時，這項工作將變得更加複雜，因為許多 3D 元素需要在佈局中相互配合、在上下文中協同運作。我們也希望所創建的 3D 物件與世界能具備完整功能，例如門能開闔、輪子能轉動等。為此，我們需要快速的生成與迭代能力，以擴展至整個場景、具備完整功能的物件以及虛擬角色。 隨著我們持續擴展 Cube 3D 基礎模型，我們期待能分享進一步的改進與新功能，並見證創作者社群運用這些工具所建構的沉浸式世界。