利用 NVIDIA 組件提升 GPU 推理的吞吐
本實踐中,唯品會 AI 平臺與 NVIDIA 團隊合作,結合 NVIDIA TensorRT 和 NVIDIA Merlin HierarchicalKV(HKV)將推理的稠密網絡和熱 Embedding 全置于 GPU 上進行加速,吞吐相比 CPU 推理服務提升高于 3 倍。
應對 GPU 推理上的難題
唯品會(NYSE: VIPS)成立于 2008 年 8 月,總部設在中國廣州,旗下網站于同年 12 月 8 日上線。唯品會主營業務為互聯網在線銷售品牌折扣商品,涵蓋名品服飾鞋包、美妝、母嬰、居家、生活等全品類。
唯品會 AI 平臺服務于公司搜索、推薦、廣告等業務團隊,提供公司級一站式服務平臺。搜索、推薦、廣告等業務旨在通過算法模型迭代,不斷優化用戶購買體驗,從而提升點擊率和轉化率等業務指標,最終實現公司銷售業績增長。
在使用 GPU 打開推理算力天花板過程中,遇到了如下問題:
- 稠密網絡,如何獲取更好的 GPU 推理性能;
- Embedding table 如何使用 GPU 加速查詢。
為了解決上面的問題,我們選擇使用了 NVIDIA TensorRT 和 Merlin HierarchicalKV。具體原因如下:
- 稠密網絡使用 TensorRT 推理,通過 TensorRT 和自研 Plugin 方式獲取更好的推理性能;
- HierarchicalKV 是一個高性能 GPU Table 實現,我們將熱 Embedding 緩存在 GPU 中,冷 Embedding 則通過內存和分布式 KV 存儲,加速查表過程。
GPU 推理服務設計方案
AI 平臺支持搜索、推薦、廣告等所有算法業務,提供大規模分布式訓練、推理、實時模型等基礎引擎平臺,打造屬于唯品會自己的 AI 基礎能力引擎。

圖1. GPU 推理服務工作流程圖
如上圖所示,支持 GPU 推理服務,可以分為如下幾步:
- TensorPS(自研訓練框架)
- 支持離線和實時訓練;
- 離線訓練:生成天級全量模型,完成后同步給 Odin;
- 實時訓練:生成小時級別的全量模型和分鐘級別的增量模型,完成后同步給 Odin;
- Odin(模型協調者)
- (離線/實時)單機模型的全量模型:觸發 TensorRT Converter;
- (離線/實時)分布式模型的全量模型:同時觸發 TensorRT Converter 和 Reshard;
- (離線/實時)單機/分布式模型的增量模型:觸發 TensorRT Converter;
- TensorRT Converter(模型轉換器)
- 將 Dense 網絡轉換成 TensorRT Engine;
- 轉化完成,如果是全量模型,向模型管理 API 匯報全量版本;如果是增量模型,向模型管理 API 匯報增量版本;
- Reshard(參數分片模塊)
- 對模型參數分片后,向模型管理 API 匯報版本;
- 分片后參數,同步到分布式在線參數服務 Atreus;
- Thor(自研推理服務)
- 單機模型:通過模型管理 API 獲取全量模型版本,拉取模型并啟動推理服務 Thor;
- 分布式模型:需要部署分布式參數服務 Atreus 和 推理服務 Thor;
- 如果開啟了實時模型特性,Thor 會定時通過模型管理 API 獲取增量版本,拉取并更新增量模型;
- Atreus(自研分布式在線參數服務)
- 僅用于分布式模型,可支持 TB 級參數;
- 如果開啟了實時模型特性,Atreus 會定時通過模型管理 API 獲取增量版本,拉取并更新增量參數。
GPU 模型推理

圖 2. 前向計算流程圖
如上圖所示,前向計算可以分為如下幾步:
- H2D 拷貝(CPU -> GPU);
- Embedding 層,使用 GPU Table lookup(GPU);
- Dense 層,使用 TensorRT + 自研 Plugin 推理(GPU);
- D2H 拷貝(GPU -> CPU)。
稠密網絡使用 TensorRT 在 GPU 上計算

圖 3. 稠密網絡 TensorRT 推理優化
如上圖所示:
- 稠密網絡使用 TensorRT 推理,結合自定義 Plugin 實現推理性能優化。
利用 HierarchicalKV 實現 GPU Table lookup

圖 4. 基于 HierarchicalKV 的 GPU Table
如上圖所示,查表過程可以分為如下幾步:
- 將 keys 拷貝到 GPU;
- 將 keys concat 成一個大的 merged keys,減少后續查表次數;
- merged keys 查 GPU Table,輸出 merged values,并輸出未命中 missed keys 和 missed indices;
- 拷貝 missed keys 到 CPU;
- 查詢 Atreus(分布式參數服務器),獲取 missed values;
- missed values 拷貝到 GPU;
- 將 missed values 更新到 merged values;
- 將 merged values 輸出 Split 成多個 Tensor(和 keys 一一對應);
- 對 missed keys 進行去重;
- 去重之后,異步更新 GPU Table。
TensorRT Converter 實現 GPU 模型轉換

圖 5. TensorRT Converter 轉換流程
如上圖所示,TensorRT Converter 可以分為如下幾步:
- Freeze CPU 模型;
- 切分模型 Graph 成 Sparse 和 Dense 兩個子圖,Sparse 圖在 GPU 上執行,Dense 圖經過圖優化后使用 TensorRT 推理;
- Dense 圖轉化成 ONNX 模型;
- 優化 ONNX 模型,把圖中 OP 替換成自定義的高性能 TensorRT Plugin;
- 轉換 ONNX 模型成 TensorRT Engine;
- 合并 Sparse 圖和 TensorRT Engine 生成 GPU 模型。
自研 CUDA Kernel,提高性能
- GPU Table 加速查表
基于 HierarchicalKV 增強了 find 接口,支持獲取未命中 keys indices 等信息,在高命中率情況下有更好的性能,并貢獻給社區:
void find(const size_type n,
const key_type* keys, // (n)
value_type* values, // (n, DIM)
key_type* missed_keys, // (n)
int* missed_indices, // (n)
int* missed_size, // scalar
score_type* scores = nullptr, // (n)
cudaStream_t stream = 0) const
- GPU 支持 CSR(Compressed sparse row) 格式的序列特征
- 根據統計,序列特征有 85%+ 的數據都是填充值,使用 CSR 格式壓縮序列特征可以大幅度減小序列特征大小。考慮到搜推序列數據的特殊性(填充值都在序列尾部),這里僅使用 value 和 offset 兩個序列表示原始稀疏矩陣,如下圖:

圖 6. CSR 的稀疏矩陣
- 通過 Fusion 的方式,減少 Lookup 過程 CUDA Kernel 數量,提升推理性能。

圖 7. Lookup 過程優化對比
- 優化前:N 個輸入對應 N 個 Lookup CUDA Kernel;
- 優化后:通過提前合并,將 CUDA Kernel 數量減少為 3 個(Concat、Lookup 和 Split)。
- 通過 Fusion 的方式,減少 CSR 處理過程 CUDA Kernel 數量,提升推理性能,下圖以 ReduceSum 舉例。

圖 8. CSR 處理優化流程
- 優化前:N 對輸入對應 N 個 ReduceSum CUDA Kernel;
- 優化后:通過提前合并,將 CUDA Kernel 數量減少為 4 個(2 個 Concat,1 個 ReduceSum 和 1 個 Split)。
- H2D,合并 CPU -> GPU 內存拷貝
搜推模型中有較多的特征輸入,GPU 推理中需要將這些 Tensor 從 CPU 拷貝到 GPU,頻繁小內存的 cudaMemcpy 會導致性能下降,最佳實踐是將這些 Tensor 打包在一塊連續內存中,將整個大內存 H2D 拷貝到 GPU。
- Tile 算子融合
搜推模型中有超過 200 個 Tile,大量的 Kernel Launch 會帶來 GPU 推理性能惡化,最佳實踐是進行 Kernel Fusion,在一個大的算子中執行多個小 Kernel,從而充分發揮 GPU 的并發優勢。
持續在搜廣推場景中 GPU 加速
唯品會 AI 平臺,一直追求性能上的極致,未來將會持續與 NVIDIA 技術團隊合作,繼續探索使用 HierarchicalKV 在訓練超大型模型上的 GPU 性能優化,在提升 GPU 性能方面進行不斷地探索和實踐,也會對 Generative Recommenders 進行探索和實踐。