vLLM × HPC-Ops:來自騰訊混元的高性能 Attention 與 MoE 後端
vLLM × HPC-Ops:來自騰訊混元的高性能 Attention 與 MoE 後端
為何這很重要
生產環境中的 LLM 推論現在需要處理動態、混合長度的 Batch 以及日益龐大的 MoE 模型,因此延遲主要取決於 Kernel 調度工作與數據移動的效率,而非單純的矩陣乘法(matmul)吞吐量。HPC-Ops 提供了負載平衡的 Attention 以及完全融合的 FP8 MoE Kernel,減少了閒置的 SM 週期並消除了每階段的開銷,直接解決了這些現實世界的瓶頸。
HPC-Ops:生產級算子庫,現已整合至 vLLM
HPC-Ops 是由騰訊混元 AI Infra 團隊開發的開源算子庫,專注於 Attention、MoE、GEMM、Sampling、Normalization 以及計算與通信融合,並針對 Hopper GPU 提供原生 BF16 與 FP8 支援。其 Attention 與 MoE Kernel 已透過 PR #46020 (Attention) 與 PR #45924 (MoE) 作為一等公民後端併入 vLLM 主分支,無需修改原始碼或使用長期分支(fork)。
Attention 後端:動態負載平衡調度
HPC-Ops 的 Attention 後端透過每步(per-step)負載平衡調度器,消除了混合長度解碼(mixed-length decode)中由靜態 split-KV 調度引起的閒置時間。該調度器將每個 KV 序列切分為統一的 64-token tiles,根據實際工作量將 tiles 分配給 CTA,並運行持久性 Kernel Grid 以保持所有 SM 處於飽和狀態。這在混合長度解碼的 H20 上,比靜態 split-KV 調度快 up to 2.95×,且比 FlashInfer 與 FlashAttention 平均快 2.25×。
一個融合的 Attention 前導算子 (HpcRopeNorm) 將 QK-Norm、RoPE、KV-cache 寫入,以及在 FP8 下的 Query 量化整合進單個 Kernel,消除了獨立的記憶體受限(memory-bound)啟動及其 HBM 往返開銷。
與 vLLM 的整合是透過 HpcAttentionBackend 完成的,它繼承自 vLLM 的 AttentionBackend 基類,並透過標準後端機制進行註冊。
MoE 後端:融合、低延遲的 FP8 MoE 流水線
HPC-Ops 的 MoE 後端透過將 Routing、Gate-Up GEMM、Activation/Quantization、Down GEMM 與 top-k 加權歸約(weighted reduction)融合進單個執行流水線,消除了傳統 MoE 路徑的每階段開銷。它使用共享記憶體計數(shared-memory counting)來構建連續的專家範圍,透過 Routing indices 直接讀取 tokens,將中間結果保留在暫存器或共享記憶體中,並採用以佔用率為優先的 Warp Group 與持久性 Grid,將不均勻的專家 tiles 均勻分配到各個 CTA。程式化依賴啟動(Programmatic Dependent Launch)則透過重疊 Kernel 啟動來消除空隙(bubbles)。
該流水線以 FP8 運行,並具備逐張量(per-tensor)與逐塊(block-wise)縮放,與基準輸出品質一致。在 H20 上,於 TP8/EP1 配置下比 Triton 與 CUTLASS 平均快 1.59×,於 TP1/EP8 配置下快 1.21×。
與 vLLM 的整合是透過 HPCExperts 完成的,它繼承自 FusedMoEExpertsModular 並透過標準後端註冊機制進行註冊。
在 vLLM 中使用 HPC-Ops 後端
要使用這些後端,請先從原始碼安裝 HPC-Ops:
git clone https://github.com/Tencent/hpc-ops.git
cd hpc-ops
make wheel
python3 -m pip install dist/*.whl
為 Hy3 (BF16) 模型啟用 Attention 後端:
vllm serve tencent/Hy3 \
--tensor-parallel-size 8 \
--attention-backend HPC_ATTN
針對 Hy3-FP8,需添加 KV-cache dtype 與 block size:
vllm serve tencent/Hy3-FP8 \
--tensor-parallel-size 8 \
--attention-backend HPC_ATTN \
--kv-cache-dtype fp8_e4m3 \
--block-size 64
啟用 MoE 後端(僅限 FP8):
vllm serve tencent/Hy3-FP8 \
--tensor-parallel-size 8 \
--moe-backend hpc
目前這些後端僅支援 NVIDIA Hopper 架構的 GPU,在 H20 上表現最佳。
H20 上的性能表現
Fused MoE vs Triton / CUTLASS
在 Batch size 從 4 到 16384 的範圍內,HPC-Ops MoE 的延遲均低於 Triton 與 CUTLASS。在 TP8/EP1 下平均加速 1.59×;在 TP1/EP8 下平均加速 1.21×,在低延遲解碼常見的中小 Batch size 下增益最明顯。
混合長度 Batch 下的 Attention 解碼
動態調度優於靜態 split-KV、FlashInfer 與 FlashAttention。加速比隨偏斜度(skew)增加:從均勻 64×0.5K batches 的 1.00×,提升至 1×128K + 31×4K 混合模式下的 2.95×。在所有測試的分佈中,相對於 FlashInfer/FlashAttention 中最優者的平均優勢為 2.25×。
跨 Prefill、Extend 與 Decode 形狀的 Attention 性能
對於每個測試的形狀,HPC-Ops 的 Attention 延遲均與 FlashAttention、Triton 或 FlashInfer 中最快的版本持平或更快(例如:q512 prefill 為 0.047 ms vs 0.069 ms;8q1s1k decode 為 0.019 ms vs 0.031 ms)。
Hy3 (8× H20) 端到端性能
同時使用 HPC-Ops 的 Attention 與 MoE 後端,與 vLLM 預設後端相比,TTFT 降低了約 24%,TPOT 降低了約 17%。隨著 Batch size 增加,改進效果愈發顯著,在測試的最大 Batch size 下,TTFT 與 TPOT 分別可降低約 30%。
下一步計畫
騰訊混元團隊將持續與 vLLM 社群合作,優化這些後端,並隨著技術成熟將更多工作併入主分支,同時歡迎用戶提供回饋、回報問題與提供 Benchmark。
致謝
感謝騰訊混元 AI Infra 團隊 (Sethran Liu, Chase Shao, Shengy Wei, Theo Cheng, Ryann Xue, Lando Jiang, Looper Zhao, Haank Lin, Aiden Ren, Lehua Ding, Chengv Jiang, Steven Kuang, Liqi He, Kipper Gong, Reedlau Liu, Raccoon Liu, Dick Zhu)、騰訊網絡平台部、vLLM/Inferact 貢獻者 (Kaichao You, Yongye Zhu, Yifan Qiao)、NVIDIA 工程師 (Yuanhang Sun, Perkz Zheng, Yuxi Chi, Jiang Shao, Jun Gu, Meng Wang, River Liu, Gary Ji, Chandler Zhou),以及為此提供基礎與對照基準的廣大開源 Kernel 社群。