生成式推薦（GR）
系統正成為大規模個性化推薦的強大新方向。GR不再將推薦視為一系列孤立的檢索、排序和預測階段，而是將推薦重新定義為用戶行為的序列建模問題。用戶的交互、上下文、候選物品和行為都被轉化為高基數事件流中的Token，模型學習從該序列中生成或預測下一個相關物品。
這種方法對現代推薦工作負載尤為適用，因為用戶歷史記錄往往很長，物品目錄持續變化，個性化質量依賴於對豐富序列行為的建模。但這也帶來了服務層面的挑戰：儘管存在長歷史記錄、大型嵌入表和序列密集型模型架構，GR模型仍需要實現低延遲推理。
NVIDIA Dynamo-Triton（原NVIDIA Triton Inference Server）現已通過NVIDIA recsys-examples代碼庫支持端到端的分層序列轉導單元（HSTU）GR推理工作流。該工作流結合了HSTU、PyTorch預先編譯（Ahead-of-Time Inductor）、基於FlexKV的KV緩存、原生C++驗證、NV嵌入緩存以及Dynamo-Triton部署。
最終結果是一條具有出色延遲性能的HSTU排序模型服務路徑。在NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU上，動態批次大小為8時，相比未使用KV緩存的同款AOTI配置，Dynamo-Triton搭配PyTorch AOTI在GPU KV緩存命中率達到100%的情況下，三層HSTU模型最高提速4.47倍，八層模型最高提速5.93倍。
本文將介紹如何藉助NVIDIA Dynamo-Triton、PyTorch AOTI和FlexKV，將HSTU生成式推薦系統從PyTorch開發階段推進至生產推理階段。您將學習如何導出並預先編譯模型，在Python和原生C++環境中驗證最終的部署產物，並通過Dynamo-Triton進行服務部署，而無需為單獨的運行時重寫模型。
本文還探討了基於GPU的KV緩存如何減少重複計算，並給出了基準測試結果，證明該部署工作流可將延遲降低最高達5.93倍，突顯了該方案的實際性能優勢。
為何使用HSTU進行生成式推薦
HSTU是為處理高基數、非平穩事件流的GR工作負載而提出的。在傳統推薦系統中，檢索和排序通常由一系列專門的模型和特徵流水線構建而成。而GR將推薦建模為序列預測問題，使模型能夠在一個具備序列感知能力的架構中，統一推理用戶上下文、物品歷史、行為歷史和候選物品。
在NVIDIA HSTU排序示例中，模型輸入由分類Token構建而成。上下文Token代表用戶側資訊，物品Token代表具體物品，可選的行為Token則代表用戶與這些物品的交互行為。
HSTU預處理路徑會檢索嵌入向量，在存在行為Token時將物品嵌入和行為嵌入交錯排列，附加上下文資訊，並應用位置編碼。隨後HSTU模組處理該序列，預測頭輸出多任務排序結果。
這種結構非常適合需要考慮時近性、順序和重複交互模式的推薦系統。但這也意味著，當每個請求都要反覆處理長歷史序列時，推理成本會變得高昂。生產系統需要在保留HSTU建模優勢的同時，減少服務過程中的冗餘計算。
為何大型序列推薦系統的服務具有挑戰性
服務大型序列推薦系統與服務小型稠密排序模型截然不同。服務架構必須處理不規則序列輸入、龐大的分類嵌入狀態、長歷史記錄，以及同一用戶可能反覆請求、每次只攜帶少量新資訊的請求模式。每次請求都重新計算用戶歷史的完整鍵值狀態會浪費計算資源並增加延遲。
這正是KV緩存變得重要的原因。KV緩儲存存先前序列計算中可復用的鍵值數據，使模型避免重新計算用戶歷史中已緩存的部分。對於推薦系統推理而言，當用戶的長期歷史大體保持穩定、僅有新的候選物品或近期行為到來時，這種方式尤為有用。
NVIDIA HSTU推理工作流包含一個KVCacheManager，利用GPU內存和主機儲存來緩存KV數據。GPU緩存以分頁KV數據表的形式組織，支持查找、分配、追加和淘汰操作。當GPU緩存空間受限時，系統會按照類似LRU的策略淘汰較舊的用戶數據。主機側儲存提供了另一層緩存KV數據的空間，該工作流還包含基於FlexKV的KV緩存運行時後端。
HSTU注意力核心可以從分頁緩存中讀取KV數據，導出的推理路徑包含支持緩存感知的自定義操作，用於查找、分配、加載、追加和卸載。這使得服務路徑能夠在減少冗餘計算的同時，保留模型的序列語義。
用於原生推理的PyTorch AOTI
PyTorch AOTI（預先編譯Inductor）工作流從一個PyTorch模型開始，使用torch.export和PyTorch AOTI對其進行導出。AOTI會將模型提前編譯為一個可由原生C++運行時加載的軟體包。這降低了Python運行時開銷，為Dynamo-Triton PyTorch AOTI後端提供了便於部署的產物。
導出的模型包包含AOTI模型存檔以及元數據和嵌入表文件。在NVIDIA示例中，嵌入實現結合了DynamicEmb推理嵌入表和NV嵌入緩存，後者僅將常用嵌入儲存在GPU內存中，同時將整個表保留在CPU內存中，從而降低GPU內存占用。導出路徑會在編譯後的.pt2存檔旁寫入層元數據和嵌入表數據，這樣模型加載時就無需產生不必要的重複嵌入表副本。
該工作流通過多種方式驗證同一導出產物。Python導出腳本生成軟體包並重放張量數據。原生C++可執行文件加載並重放導出模型，以驗證正確性和性能。Dynamo-Triton部署隨後使用同一個AOTI包和重放路徑，這有助於保持開發驗證與生產服務的一致性。
Dynamo-Triton部署路徑
Dynamo-Triton為導出的HSTU模型提供了生產級服務層。AOTI部署使用Dynamo-Triton PyTorch後端，平台設置為"torch_aoti"。這使得Dynamo-Triton能夠加載並服務預先編譯的PyTorch模型包。
完整工作流包含以下五個階段：
構建所需的自定義算子和運行時庫
使用PyTorch AOTI導出HSTU排序模型
啟動基於FlexKV的KV緩存服務
通過原生C++重放驗證導出的產物
使用Dynamo-Triton服務導出的KV緩存AOTI模型
這種方法之所以重要，是因為推薦系統的服務不僅僅需要一個快速的模型核心。Dynamo-Triton帶來了模型倉庫管理、請求處理、後端集成、指標監控和部署結構。AOTI帶來了低開銷的編譯模型產物。NV嵌入緩存通過僅在GPU內存中保留嵌入表的熱點部分，降低了GPU內存需求。FlexKV通過緩存注意力模組，減少了長用戶歷史的重複計算。這些組件共同構成了一個面向實際生產環境的生成式推薦推理服務架構。
HSTU服務延遲基準測試
本基準測試比較了不同Dynamo-Triton後端、模型規模、批次大小和KV緩存狀態下的HSTU服務延遲。
測試目標是量化生產級HSTU服務架構的性能優勢。測試比較了Dynamo-Triton PyTorch AOTI後端與Python後端，衡量了GPU KV緩存帶來的額外延遲降低幅度，並評估了這些優勢在不同模型深度和批次大小下的擴展表現。最終，這些結果向開發者展示了從未緩存的基於Python的推理，遷移到編譯型、具備緩存感知能力的HSTU部署方案（基於Dynamo-Triton）後可獲得的性能提升。
recsys-examples中的基準測試結果使用了單GPU上的KuaiRand-1K排序配置。模型結構包括三層和八層HSTU變體，隱藏層大小為512，4個注意力頭，BF16模型權重，BF16 KV緩存，最大歷史序列長度共8,192個Token（4,096組物品與行為對）的歷史流，最大候選序列長度為100，以及6個上下文特徵。對齊前的有效序列長度為8,298個Token，導出時的最大對齊序列長度為8,320。
基準測試協議按每個邏輯請求報告延遲。對於Dynamo-Triton AOTI基準測試，每次Dynamo-Triton調用包含一個邏輯批次，每個邏輯請求的延遲通過將端到端的處理時間除以Dynamo-Triton調用次數再除以邏輯批次大小計算得出。數據集加載、驗證、重新分批、用戶ID生成、伺服器啟動、預熱以及預熱後的休眠時間均未計入測量時間。
用於Dynamo-Triton後端對比和批次大小結果測試的硬體為NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU。
Dynamo-Triton後端對比
在Dynamo-Triton批次大小為2時，即便沒有KV緩存命中，PyTorch AOTI相比Dynamo-Triton Python後端也能改善延遲表現。而在20GB GPU KV緩存命中的情況下，延遲改善幅度明顯更大。
這些結果展示了兩種不同的增益來源。首先，AOTI相比Python後端降低了服務開銷。其次，KV緩存命中通過復用已緩存的序列狀態減少了模型的計算量。在更深的八層模型上，這種緩存帶來的收益尤為明顯，因為避免重複計算所帶來的影響更大。
AOTI與KV緩存下的批次大小擴展性
按批次大小劃分的PyTorch AOTI後端結果顯示，隨著邏輯批次大小的增加，KV緩存的效果愈發顯著。
在批次大小為8時，三層HSTU模型在GPU KV緩存命中的情況下，每個邏輯請求的延遲達到0.423毫秒。八層HSTU模型的延遲則達到0.678毫秒。對於長序列排序推理而言，這是相當出色的結果，展示了編譯模型執行與緩存感知服務相結合所帶來的價值。
加速HSTU GR推理有何好處
推薦系統在嚴格的延遲預算下運行。額外的排序延遲會影響頁面加載時間、資訊流響應速度和廣告投放的時效要求。與此同時，日益具備序列感知能力和個性化能力的模型往往需要更多的推理算力，尤其是隨著用戶歷史記錄的增長。
HSTU服務工作流結合了多項互補技術來應對這一挑戰。HSTU提供了生成式推薦架構，PyTorch AOTInductor生成預先編譯的部署產物，基於FlexKV的KV緩存使之前計算的注意力狀態得以復用，NVIDIA Dynamo-Triton則提供了生產級服務環境。
當連續請求共享用戶交互歷史中未發生變化的前綴部分時，這種方法尤為有價值。模型無需重新計算該部分序列的注意力，而是可以復用已緩存的鍵值狀態，只需針對新追加的Token進行計算。隨著序列變長、模型加深，這種節省效果會進一步放大，因為跨多個HSTU層的重複計算原本會帶來顯著的延遲增加。
開始加速HSTU GR推理
您可以通過NVIDIA/recsys-examples GitHub代碼庫復現並擴展此工作流。HSTU概述介紹了GR模型結構，包括上下文Token、物品Token、行為Token、嵌入表、HSTU模組和預測頭。
AOTI推理指南詳細介紹了構建所需鏡像和庫、準備KuaiRand-1K數據、訓練檢查點、導出KV緩存AOTI模型、使用C++重放進行驗證、打包Dynamo-Triton運行時鏡像，以及通過Dynamo-Triton伺服器重放請求的完整流程。
要了解更多資訊，請查閱以下相關資源：
HSTU生成式推薦系統概述
Dynamo-Triton、FlexKV與PyTorch AOTI集成方案
HSTU推理基準測試
Dynamo-Triton HSTU文檔
NV嵌入緩存
致謝
本文是NVIDIA多個團隊跨職能協作的成果。我們要感謝J、Runchu Zhao、Yulu Liu、Lin Hu、Zhuofan Li、Jacob Subag和Tomer Bar-On的貢獻。
Q&A
Q1：什麼是HSTU生成式推薦系統？
A：HSTU（分層序列轉導單元）是一種用於生成式推薦的模型架構，它將推薦問題重新定義為序列建模任務，能夠統一處理用戶上下文、物品歷史和行為歷史，而不是依賴傳統的檢索、排序等孤立階段。
Q2：使用Dynamo-Triton部署HSTU模型能帶來多大的性能提升？
A：在NVIDIA RTX PRO 6000 Blackwell GPU上，動態批次大小為8時，相比未使用KV緩存的配置，三層HSTU模型最高提速4.47倍，八層HSTU模型最高提速5.93倍，GPU KV緩存命中率達到100%。
Q3：KV緩存在HSTU推理中起什麼作用？
A：KV緩儲存存先前序列計算中可復用的鍵值數據，使模型避免對用戶歷史中未變化的部分重新計算，僅需處理新追加的Token，從而顯著降低長序列推理的延遲和計算成本。
