跳到正文
aihot(aihot.news)· noreply@aihot.news (vLLM 官方部落格)·· 2 天前精選品質評分83

vLLM Day‑0 支援 Mistral Large 4,立即部署高效多模態模型

vLLM 釋出 Mistral Large 4 Day 0 支援

重點導讀

vLLM 已於 Day‑0 支援 Mistral Large 4,讓使用者可直接在 Docker 或 GPU 上即時部署這款 1 T 引數的多模態模型。

推薦理由

vLLM Day‑0 支援 Mistral Large 4 讓開發者可即時部署高效多模態模型,並提供實驗資料供最佳化參考。

正文

vLLM 官方部落格· Mistral and the vLLM Team·2026-10-29 08:00· 2 天前精選AI 評分71

Day 0 support for Mistral Large 4 on vLLM

AI 導讀

vLLM 團隊宣佈對 Mistral 新旗艦模型 Mistral Large 4 提供 Day 0 支援,該模型為 1T 引數、每 token 啟用 49B 的原生多模態 MoE 模型,採用 MLA 稀疏注意力加 DSA top-k 索引器、跨層 KV 共享、NoPE、SSMax 和 LeRoPE 等新架構,並內建 DSpark 推測解碼草稿器與 Pixtral 視覺編碼器。

推薦理由

vLLM 團隊詳解支援 Mistral Large 4 所需的架構適配與量化、長上下文、推測解碼等實測資料,可作部署參考。

我們很高興地宣佈,vLLM 已實現 Mistral Large 4 的 day-0 支援。Mistral Large 4 是 Mistral 推出的新一代前沿混合專家(MoE)原生多模態模型。它把 DeepSeek 風格的稀疏注意力棧(MLA 加 DSA top-k 索引器)與多項新架構特性結合在一起:跨層 KV 共享、NoPE 層、SSMax 和 LeRoPE。模型自帶 DSpark 投機解碼草稿器和 Pixtral 視覺編碼器。

支援它遠不止加一個模型檔案那麼簡單。ML4 執行在 vLLM 的模型執行器 V2(MRV2)上,並支援可斷開的 CUDA 圖。它的各個構建模組放在新的 vllm/models/mistral/ 包中,這項工作還順帶改進了 decode 上下文並行(DCP)、prefill 上下文並行(PCP)、DSpark、基於 NIXL 的分離式服務以及 HiSparse 主機記憶體 KV 解除安裝。

亮點

  • 架構:72 層 MoE,128 個路由專家(top-4)加 1 個共享專家,隱藏維度 6144。注意力採用 MLA(KV LoRA 秩 512,RoPE 維度 64),並配有 DSA 稀疏索引器,每個查詢選取 top 2048 個 key。總引數 1 萬億,每 token 啟用 490 億(含嵌入層和輸出層為 520 億)。
  • 檢查點:FP8(塊縮放權重)和 NVFP4(專家用 NVFP4,注意力/共享專家用 FP8)。兩者預設都使用 FP8 潛在 KV 快取(fp8_ds_mla)。
  • 多模態:Pixtral 視覺編碼器處理影像輸入,支援每個提示詞包含多張圖片。
  • 投機解碼:內建 DSpark 草稿器(3 層滑動視窗 MLA,深度 3),每步接受約 2.9 個 token,在 TP8 上輸出吞吐約提升 2 倍。
  • 長上下文:DCP 在 TP8 上帶來約 7.9 倍的 KV 容量。PCP+DCP 預填充把首 token 時間(TTFT)最多縮短 1.6 倍。DSpark 在兩種模式下都能工作。
  • 分離式服務:NIXL 預填充/解碼(P/D)分離,支援 DP/EP、TP 和 PCP×DCP 預填充節點,兩側都啟用 DSpark。
  • HiSparse:主機記憶體稀疏 KV 解除安裝,現已支援跨層 KV 共享。
  • AMD:在 AMD Instinct MI350 系列 GPU 上通過 ROCm 實現 day-0 支援,包括 DSpark、影像輸入和 MXFP4 專家的專家並行。

用法

Mistral Large 4 的支援包含在 vllm/vllm-openai:mistral-large-4 Docker 映象中,該映象還包含 mistral-common 1.13。下面的命令在映象內執行 vllm serve:

docker pull vllm/vllm-openai:mistral-large-4

8 塊 GPU 上的 FP8:

vllm serve mistralai/Mistral-Large-4-1T-A52B \
  --tensor-parallel-size 8 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice \
  --reasoning-parser mistral

4 塊 GPU 上的 NVFP4(Blackwell;在 8 塊 GPU 上使用 --tensor-parallel-size 8):

vllm serve mistralai/Mistral-Large-4-1T-A52B-NVFP4 \
  --tensor-parallel-size 4 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice \
  --reasoning-parser mistral

NVFP4 不以質量換速度。在多模態、長上下文、智慧體程式設計、業務自動化和網路安全評測中,它在每項任務上與 FP8 的差距都在 2.1 分以內,平均差異為 −0.3 分。token 消耗也相當:平均補全 token 數與 FP8 相差約 4% 以內。它還更快、更省:在 TP8 下,不使用投機解碼時 NVFP4 每 token 解碼速度比 FP8 快 1.08–1.19 倍(啟用 DSpark 時為 1.18–1.22 倍),KV 快取容量約為 1.5 倍。

*FP8 與 NVFP4 在核心評測上的得分和 token 消耗對比。誤差條為通過自助法估計的 95% 置信區間。在所有領域,包括多模態推理(GDP.pdf、MathVision)、長上下文(Graphwalks 512k)、程式設計智慧體(DeepSWE、SWE-bench Pro、Terminal-Bench 4)、業務自動化(AutomationBench)和網路安全(CyberGym),NVFP4 檢查點的表現都與 FP8 一致。

長上下文(DCP)。新增 --decode-context-parallel-size 8 可將潛在 KV 快取分片到 8 個 rank 上。這要求使用預設的 FP8 KV 快取。

多模態。傳入 --limit-mm-per-prompt '{"image": N}' 可允許每個提示詞最多 N 張圖片。

多節點、分離式 P/D 以及其他已測試的配置,參見 Mistral Large 4 配方。

實現細節

模型架構

元件Mistral Large 4
層數 / 隱藏維度 / 注意力頭數72 / 6144 / 64
注意力MLA,無 query LoRA,KV LoRA 秩 512,RoPE 維度 64
稀疏注意力DSA indexer,top-k 2048,32 個 index 頭;44 層各自擁有一個 indexer,28 層複用最近的前一個 indexer 層的 top-k
位置編碼YaRN + LeRoPE(按頻率學習縮放);每第 4 層為 NoPE
query 縮放每個注意力層和 indexer 都使用 SSMax(隨長度變化的 query 縮放)
KV 共享9 層(39、43、…、71)讀取第 35 層的 latent cache(每 4 層 1 層)
MoE第 0 層為稠密層;第 1–71 層:128 個專家,top-4,1 個共享專家
輸出logit soft cap 30
視覺Pixtral 編碼器(30 層,隱藏維度 2048,patch 14)

*72 個解碼器層:每第 4 層為 NoPE,28 層複用最近的前一個 indexer 層的 top-k,NoPE 層 39–71 讀取第 35 層的 latent KV,且只有第 0 層是稠密層。*

在 vLLM 中支援各部分所需的工作:

稀疏 MLA + DSA indexer。 ML4 複用了 vLLM 的稀疏 MLA 棧(帶 fp8_ds_mla latent cache 的 FlashMLA 稀疏實現,外加帶 FP8 indexer cache 的稀疏 indexer)。ML4 特有的差異疊加在其上:共享 indexer 層複用最近的前一個 indexer 層的 top-k;indexer 擁有自己的 SSMax;NoPE 層上使用 NoPE indexer。每個 token 的 KV cache 儲存 47,136 位元組:41,328 位元組用於 MLA latent,5,808 位元組用於 indexer。

跨層 KV 共享。 九個 NoPE 層讀取第 35 層寫入的 latent KV,而不是寫入自己的 KV。我們在 kv_sharing_target_layer_name 到 MLAAttention 以及 AttentionLayerBase 中加入了支援。消費層跳過 cache 寫入,V2 runner 將它們解析為目標層的 KV cache 規格。消費層不佔用 KV 記憶體:72 層中有 63 層擁有自己的 cache。 該特性與模型無關,正在 vllm#64095 中向上遊合入。

NoPE、SSMax 與 LeRoPE。 NoPE 層以 RoPE 維度為零的 MLA 形式表達,因此使用相同的 kernel 和 cache 佈局。LeRoPE 被摺疊進 YaRN 頻率。SSMax 被融合進 MLA RoPE kernel(見下文融合部分)。

Mistral 模型包。 移植沒有在 deepseek_v2.py 中增加開關,而是基於 vllm/models/mistral/ 下可複用的 Mistral 元件構建:注意力、indexer、SSMax、MLP、MoE、融合運算元以及逐層規格(common/layer_spec.py)。ML4 模型本身是 vllm/models/mistral_large4/nvidia/model.py,DSpark drafter 緊挨著它 (dspark.py)。逐層特性(RoPE/NoPE、indexer 擁有者/共享者、KV 擁有者/消費者)來自基於 params.json 構建的層規格,而不是共享程式碼中的模型特定分支。

模型 runner V2。 ML4 僅在 MRV2 上執行。V2 注意力工具會發現 KV 共享層並查詢其目標規格。DSpark 接入 V2 的 DFlash/DSpark speculator。

可斷開的 CUDA graph 與融合。 ML4 執行時無需 torch.compile。它使用可斷開的 CUDA graph,該特性會針對該架構自動啟用(當文本模型位於多模態封裝內時同樣如此)。相比 eager 模式,它們帶來了約 6 倍加速。 手寫融合替代了編譯器原本會做的工作,同時支援進一步最佳化:

  • 第一輪:SSMax、MLA RoPE + SSMax、indexer Q 路徑和 indexer K 路徑。這些將 decode TPOT 降低了 8–10%。
  • 解碼環節:FP32 router GEMM;indexer query 的 LayerNorm 用一次 Triton 啟動完成;all-reduce、殘差和 RMSNorm 合併進一個 FlashInfer kernel;indexer K 直接寫入 FP8 cache。在 TP8 FP8 下,c1 TPOT 從 10.22 ms 降到 9.00 ms(-12%),c8 從 16.01 ms 降到 15.08 ms。

長上下文:解碼與預填充的上下文並行

DCP。 MLA 只保留一個 latent KV head,因此在 TP 下 cache 會在每個 rank 上覆制一份。DCP 改為按 token 分片。在 TP8 上,DCP8 把 KV 容量從 2.38M 提升到 18.88M token(FP8),從 3.60M 提升到 28.76M token(NVFP4)。

Long-context capacity on 8x GB300: TP8 vs TP8 DCP8
8× GB300 上的長上下文容量:TP8 與 TP8 DCP8 對比

併發 128k / 256k / 1M token 請求,KV cache 在 8× GB300 上以 ML4 NVFP4 執行,TP8 與 TP8 DCP8 對比:floor(KV pool / 輸入長度) - 3.60M 與 28.76M token,gpu-memory-utilization 0.85。

開箱即用時,搭配稀疏 MLA 的 DCP 在每個解碼步都要付一筆不小的固定開銷。在 60k 上下文下做 profile 顯示,每步約有 7 ms 花在 DCP 專屬工作上:

  • 對 LSE 和 indexer top-k 候選做 NCCL all-gather(共 44 個 indexer 層);
  • reduce-scatter;
  • 穩定的 top-k 合併;
  • query gather;
  • 小型 mask kernel 和啟動間隙。

我們把這些替換為:

  • a2a combine,用於 FlashMLA sparse(已合入 vllm#61285):用對稱記憶體的 all-to-all LSE combine 取代 all-gather + reduce-scatter。ML4 預設啟用。
  • 直接 top-k 交換:每個 rank 把打包好的 (score, id) 候選推送到對端的對稱記憶體 inbox,由 selector 就地讀取。這取代了每個 indexer 層的一次 NCCL all-gather 和一次 6-pass radix select。
  • Top-k 壓縮:對 32 個 token 以上的 batch,每個 rank 壓縮自己有效的 top-k 槽位,並用 topk_length 限定 kernel 範圍,而不是掃描 2048 個塞滿 -1 的槽位。在 5k token 的預填充 profile 中,attention 從 165 ms 降到 38 ms。
  • 一個融合 kernel 負責處理空行(即某個 rank 不擁有某條 query 的任何 top-k 槽位)。

DCP 專屬開銷降到每步 3.2 ms。再加上解碼融合,在 TP8 DCP8 FP8 下(60k 輸入 / 1k 輸出):

之前之後
c1 TTFT p506.89 s4.73 s
c1 TPOT p5017.82 ms12.50 ms
c64 輸出 tok/s139.1204.0 (+47%)
DCP8 before and after the collectives port
DCP8 在集合通訊移植前後的對比

TP8 DCP8 下 TTFT 和 TPOT p50 隨併發變化(60k 輸入 / 1k 輸出,無投機解碼),對比 DCP 集合通訊移植與解碼融合之前(a2db3f0b9d)和之後(bbe5926511)。

這些 DCP 最佳化隨 Mistral Large 4 支援在 day 0 一同上線。

何時使用 DCP。DCP 用每步延遲換 KV 容量。在同樣的 8x GB300 NVFP4 部署和工作負載下(無投機解碼),DCP8 每 token 的解碼速度比 TP8 慢約 16–28%:8k 輸入時 c1 TPOT 為 11.1 vs 8.7 ms,60k 時為 11.4 vs 9.1 ms。在最多 128 個 8k–60k token 的併發請求下,TP8 的 3.6M token 池從來不是瓶頸,因此 TP8 在那裡吞吐更高。一旦 KV 成為限制,DCP8 才划算:超長上下文,或大量併發的長生命週期會話。在我們的 agentic P/D 執行中(約 100k token 會話),一個 TP4 DCP4 解碼節點在 32 個會話下把 TTFT 保持在 0.7 s,而普通 TP4 解碼節點 KV 耗盡,TTFT 攀升到 15 s。

PCP + DCP 預填充。對長 prompt,PCP 把一個 prompt 的 token 拆分到各 rank。PCP+DCP 現在在 ML4 上支援 CUDA graphs,包括 DCP > PCP 和 DSpark。兩項修復讓這成為可能:

  • 在 PCP 下,MLA KV-cache 寫入會先跨 PCP rank 對輸入做 all-gather,而這次 gather 取決於該批次的 prefill/decode 劃分。在分段圖(piecewise graph)中捕獲時,這個劃分在捕獲時就被固定了。現在只有 gather 作為 eager 斷點執行,寫入固定大小的緩衝區,而 cache 寫入本身仍留在圖裡。
  • DCP 組的 rank 順序與集合通訊返回資料的順序保持一致。

在 TP1×PCP4×DCP4 上啟用圖後,GSM8K-200 從 0.315 提升到 0.975。配合 DSpark 可達 0.980。作為 P/D 的 prefill 節點,PCP4×DCP4 給出我們測到的最低 TTFT(見 P/D 一節)。

DSpark 投機解碼

ML4 在 checkpoint 中內建了一個 DSpark drafter(mtp_dspark.* 個張量,以及 params.json 中的一個 mtp 塊)。DSpark 是一種並行 drafter:一個小型非因果 transformer 在一次前向中提出一個塊內所有 draft 位置。它以從目標模型第 1、18、36、54、71 層引出的隱藏狀態為條件。ML4 的 drafter 還增加了:

  • RNN Markov 頭:一個門控迴圈狀態(rank 512)在各 draft 位置間執行,因此每個位置的 logits 都依賴於它之前 draft 出的 token;
  • 位置專屬的 mask 輸入:每個 draft 槽位都有自己的學習輸入嵌入,而不是共用一個 mask token;
  • 滑動視窗 MLA draft 層:3 個 MLA 層,視窗 2048 token,BF16 KV cache,位於各自的 KV cache 組。Triton MLA 後端為它們新增了滑動視窗解碼。

*一個 DSpark 步驟:目標模型的引出狀態餵給 draft 的上下文 K/V,[anchor, mask_1, mask_2] 塊一次通過 3 個滑動視窗 MLA 層,RNN 頭依次輸出 3 個 draft token,再由目標模型驗證。*

用 --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' 啟用。

結果。TP8,200 條真實提示詞,貪心解碼(temperature 0),最多輸出 512 token:

FP8 基線 TPOTFP8 DSpark TPOT(AL)NVFP4 基線 TPOTNVFP4 DSpark TPOT(AL)
c18.84 ms3.50 ms (2.90)8.21 ms2.95 ms (2.93)
c815.28 ms7.08 ms (2.85)12.85 ms5.80 ms (2.86)
c3229.92 ms11.57 ms (2.85)25.38 ms9.80 ms (2.86)

該表使用貪心解碼。取樣時接受率也保持接近:在 temperature 1.0 / top-p 0.95 下,併發 8 時 AL 約為 2.8。溫度更高時,在投機配置中設定 "draft_sample_method": "probabilistic":T=1.5 時它把 AL 從 2.04 提升到 2.36–2.53(吞吐量 +12–22%)。得益於標準的拒絕取樣,兩種模式都是無損的。

在相同併發下,DSpark 在 FP8 上把輸出吞吐量提升 1.9–2.1 倍,在 NVFP4 上提升 1.9–2.2 倍。它不改變模型生成的內容:每個 draft token 都會對照目標模型驗證,我們也做了實測確認。當基線模型重新給 DSpark 的輸出打分時,這些輸出與基線模型自身貪心選擇的一致程度,和基線模型自己的輸出一樣。在 GSM8K(1,319 題)上,FP8 不用 DSpark 得 0.909,用 DSpark 得 0.918。

*每 GPU 輸出 tok/s 對每使用者 tok/s(1000 / TPOT p50),虛線為不用 DSpark,實線為用 DSpark(深度 3)。Short:上表中的 200 條提示詞 TP8 執行。8k / 32k / 60k:ML4 NVFP4 共置於 8× GB300,TP8 和 TP8 DCP8,`vllm bench serve` 使用真實文本提示詞和 1k 輸出 token;60k 時只有 DSpark 執行。*

DSpark + DCP。在此之前,滑動視窗 MLA 草稿模型無法在 DCP 下執行。現在,草稿模型的小視窗快取(3 層 × 約 2k token,每請求約 7 MB)會在每個 DCP rank 上覆制一份,而目標快取仍按 DCP 分片。這把已有支援(vllm#56723)擴充套件到了帶視窗的 MLA 草稿模型。在 TP8 DCP8、60k 上下文下,c1 TPOT 從 11.82 ms 降到 5.22 ms(2.26 倍,AL 2.85),GSM8K-200 保持在 0.965。 代價是 KV 容量:在 DCP8 下,每個 64 token 的草稿塊持有一個塊 ID,其目標頁覆蓋 512 個 token。在 64k 上下文下,容量降到純 DCP8 的約 0.69 倍。在 max_num_seqs 128 下,平均上下文低於約 99k 時這一限制不會觸發。

P/D 中的 DSpark。DCP prefill 節點可以向 TP decode 節點交接,兩側都啟用 DSpark。NIXL 從一個 prefill rank 拉取複製的草稿視窗,並以 DCP 分片形式拉取目標 KV。無論採用哪種 P/D 佈局,接受率都相同:8k / 32k / 60k 輸入下約為 2.98 / 2.86 / 2.78。

滑動視窗 MLA 草稿層以及 DCP 下的 DSpark 正在 vllm#63795 和 vllm#63796 中合入上游。

基於 NIXL 的分離式服務

我們在 4+4 塊 GB300 GPU(一個 NVL72 機架)上用 NIXL 對 1P1D 做了基準測試。配置如下:

  • 兩側均啟用 NVFP4 與 DSpark,全面使用 CUDA graphs,不啟用字首快取;
  • vllm bench serve 使用固定長度提示詞:每個輸入長度 320 條唯一的單輪聊天提示詞,每條恰好 8,192 / 32,768 / 60,000 個 token。每條提示詞都是一段自然文本(WikiText-2 和 GSM8K 題目)後接一個問題,提示詞之間沒有共享字首,以 token ID 形式傳送;
  • 每請求 1,024 個輸出 token(貪心解碼,忽略 EOS);
  • 併發掃描和固定 QPS 掃描。

所有配置經代理在 GSM8K-200 上得分 0.965–0.980,NIXL 傳輸零失敗,請求零失敗。

*每 GPU 輸出 tok/s 對比每使用者 tok/s,1P1D,4+4 GB300。實線為併發掃描;空心/虛線標記為固定 QPS 掃描。*

在給定最低每使用者速度下,每 GPU(8 塊 GPU)的最佳輸出 tok/s:

P → D8k,≥40 tok/s/使用者8k,≥10032k,≥4032k,≥10060k,≥40 / ≥100c1 TTFT 8k / 32k / 60k
P DP4/EP → D DP4/EP718.2176.7202.6139.2108.80.49 / 1.90 / 3.65 s
P DP4/EP → D TP4459.2194.2205.2150.0107.20.46 / 1.88 / 3.66 s
P PCP4×DCP4 → D TP4350.6192.691.591.549.80.35 / 1.37 / 2.57 s
P TP4 → D TP4259.2185.362.862.832.30.52 / 2.07 / 4.02 s

AMD

Mistral Large 4 從釋出首日起就通過 vLLM 的 ROCm 構建在 AMD Instinct MI350 系列 GPU 上執行。稀疏 MLA 執行在 AITER 的 MLA decode 核心上,DSA 索引器執行在 AITER 的 gfx950 索引器上。

MXFP4 專家。在 ROCm 上,路由專家以 OCP MXFP4 形式執行在 AITER 的融合 MoE 上。有兩種方式可以獲得它們:

  • AMD Quark MXFP4 檢查點(Mistral-Large-4-1T-A52B-Quark-MXFP4)。Quark 只重新量化路由專家,權重使用靜態縮放、啟用使用動態縮放。注意力、共享專家、稠密前饋層、路由器、輸出頭、視覺編碼器和 DSpark 草稿模型都保持原始精度。檢查點從約 1,061 GB 縮小到 578 GB,因此可以在四塊 MI350 系列 GPU(每塊 288 GB)上服務。服務命令、評測結果和復現步驟見模型卡。
  • 載入時重新量化 FP8 檢查點(VLLM_FP8_MOE_REQUANT=mxfp4)。專家在載入過程中逐層轉換,因此峰值記憶體也足夠低,TP4 同樣可用。TP8 下每個 GPU 的專家權重從約 125 GiB 降到 69 GiB,在 1–64 併發請求下,輸出吞吐量比直接服務 FP8 專家提升 22–71%。VLLM_FP8_MOE_REQUANT_ACT 設定 MXFP4 專家計算時使用的啟用格式:a8w4(FP8 啟用,預設)或 a4w4(MXFP4 啟用)。

無論哪種方式,一個 8-GPU 節點都能跑兩個 TP4 副本,每個副本留給 KV 快取和啟用的 HBM 約為 FP8 專家在同樣四塊 GPU 上所能留下的 6 倍。

功能。 ROCm 上支援 TP8 和 TP4、專家並行、通過 Pixtral 編碼器輸入影像,以及使用檢查點內建草稿模型的 DSpark 投機解碼。在 MI350X 上以 TP4 配合 FP8 KV 快取執行時,DSpark 在真實提示詞上每步接受 3.1 個 token,單請求下每 token 解碼延遲降低 2.7 倍,8 請求下降低 1.9 倍,32 請求下降低 1.8 倍(輸出吞吐量分別為 2.1 倍、1.6 倍和 1.4 倍)。ROCm 上潛在 KV 快取預設是 FP8(E4M3),KV 容量是 bf16 的 1.86 倍;--kv-cache-dtype bfloat16 可將其關閉。

AITER。VLLM_ROCM_USE_AITER=1 在 AMD 的核心庫 AITER 上執行 vLLM 的熱點路徑:融合 MoE(MXFP4 專家跑在兩階段 FlyDSL 核心上)、塊 FP8 和 bf16 GEMM、MLA 解碼核心、RMSNorm 以及張量並行自定義 all-reduce。ROCm 映象自帶針對該模型 MoE 和 GEMM 形狀的 AITER 調優表,AITER 在啟動時會將其與自己的調優表合併;在 TP4 下,1–8 併發請求時輸出吞吐量提升 9–16%。還有兩項設定能進一步加分:VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 讓稀疏 MLA 跑在 AITER 的 Triton 核心上,AITER_CUSTOM_AR_MAX_SIZE=134217728 讓大型預填充步驟(最大 128 MiB)的 all-reduce 留在 AITER 的自定義 all-reduce 上。兩者合計在 TP4 下再提升 2–5% 的輸出吞吐量,並將首 token 時間降低 5–7%。下面兩條命令都會啟用這三項設定。

專用解碼。對於延遲敏感的服務,--additional-config '{"decode_backend": "latency_specialized"}' 會把 1、2 或 4 請求解碼步驟中的每個 MoE 層作為一次常駐 FlyDSL 核心啟動來執行。該啟動融合了注意力、索引器、路由器、MXFP4 專家以及兩次張量並行 all-reduce。在 MI350X 上以 TP4 執行時,與同樣 GPU 上 vLLM 的標準解碼相比,單請求下每 token 解碼延遲降低 1.5 倍,四請求下降低 1.1 倍。它可以與 DSpark 組合:單請求的驗證塊跑在融合核心上,每 token 延遲比不做投機的標準解碼低 2.9 倍。更大的批次和預填充走標準路徑。它執行在預設 FP8 KV 快取上,要求 TP4 和 VLLM_FP8_MOE_REQUANT_ACT=a4w4,否則會拒絕啟動並給出原因。該核心改編自 ROCm ATOM 的 GLM-5.2 融合解碼。

在 MI350 系列上服務。兩條命令都會啟用 AITER,並在載入時將 FP8 專家重新量化為 MXFP4:

# TP8, one replica per node
VLLM_ROCM_USE_AITER=1 VLLM_FP8_MOE_REQUANT=mxfp4 VLLM_FP8_MOE_REQUANT_ACT=a8w4 \
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 AITER_CUSTOM_AR_MAX_SIZE=134217728 \
vllm serve mistralai/Mistral-Large-4-1T-A52B --tensor-parallel-size 8 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice --reasoning-parser mistral
 
# TP4 with the specialized decode, for few concurrent users; two replicas per node (second on GPUs 4-7)
VLLM_ROCM_USE_AITER=1 VLLM_FP8_MOE_REQUANT=mxfp4 VLLM_FP8_MOE_REQUANT_ACT=a4w4 \
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 AITER_CUSTOM_AR_MAX_SIZE=134217728 \
vllm serve mistralai/Mistral-Large-4-1T-A52B --tensor-parallel-size 4 \
  --additional-config '{"decode_backend": "latency_specialized"}' --max-num-seqs 64 \
  --speculative-config '{"method": "dspark", "num_speculative_tokens": 3}' \
  --tool-call-parser mistral --enable-auto-tool-choice --reasoning-parser mistral

ROCm 實現細節

AITER 上的稀疏 MLA。在 ROCm 上,ML4 的稀疏 MLA 執行在 ROCM_AITER_MLA_SPARSE 後端。索引器選出的 top-2048 位置會被轉換成全域性快取槽位,每個查詢 token 在槽位流中各自佔一條記錄。隨後 AITER 的 MLA 解碼核心以常駐模式讀取這條流,工作切分後設資料提前構建好,以便在 CU 之間均衡負載。AITER 的 MLA 核心要求頭數為 16 的倍數:TP4 下每塊 GPU 正好分到 ML4 的 64 個頭中的 16 個,而 TP8 下的 8 個頭會被補齊到 16。VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 則改用 AITER 的 Triton 稀疏 MLA 核心,它不需要補齊頭數,也不需要單獨一份 FP8 的查詢副本。bf16 快取上的 NoPE 層沒有 RoPE 部分可供 AITER 核心讀取,因此它們執行在 vLLM 的 ragged Triton 稀疏注意力核心上。

AITER 上的 DSA 索引器。一個融合 Triton 核心把每個索引器鍵量化為 FP8(使用單一 scale)並寫入索引器快取,每個索引器層每 token 佔 132 位元組。在一個快取塊內,鍵按 16×16 的瓦片儲存,這正是 AITER 分頁讀取器期望的佈局。寫入方和所有讀取方(插入、prefill gather 和分頁 logits)共用同一條佈局規則,因此不會出現不一致。Prefill 打分執行在 AITER 的 fp8_mqa_logits 上,解碼打分執行在 AITER 的分頁 MQA-logits 核心上,兩者都讀取 FP8 鍵。top-2048 在形狀受支援時由 AITER 的索引器 top-k 給出,否則由 vLLM 的 top_k_per_row 給出。共享索引器層複用最近的前一個索引器層的 top-k,與其他平臺上的做法一致。

FP8 潛在快取。ROCm 沒有 fp8_ds_mla 佈局,因此當檢查點要求 FP8 KV 快取時,ROCm 改用普通的 E4M3 快取。全部 576 個潛在維度(512 個 NoPE 和 64 個 RoPE)都以 FP8 儲存,使用一個逐張量 scale,每層每 token 佔 576 位元組,而 fp8_ds_mla 為 656 位元組,bf16 為 1,152 位元組。查詢也被量化為 FP8,因此 AITER 的核心執行的是 FP8 × FP8。算上 63 個擁有快取的層加上索引器,一個 token 佔 42,096 位元組,而 bf16 下為 78,384 位元組:這就是 AMD 部分所說的 1.86× KV 容量。

載入時逐層將專家量化為 MXFP4。使用 VLLM_FP8_MOE_REQUANT=mxfp4 時,檢查點載入器保持不變:它仍然載入帶塊 scale 的 block-FP8 專家權重。這些引數起初位於 meta 裝置上,逐層載入器一次只物化一層。某一層的專家一載入完成,每個專家就會被反量化為 bf16,再重新量化為 MXFP4(E2M1 數值,每 32 個數值一個 E8M0 scale),隨後釋放該層的 FP8 副本。峰值記憶體是 MXFP4 模型加上一層的 FP8 專家,這正是 1T 模型能裝進四塊 GPU 的原因。權重誤差與直接把 bf16 權重量化為 MXFP4 大致相當。之後專家執行在 AITER 的融合 MoE 上:a8w4 使用帶 FP8 啟用的兩階段 FlyDSL 核心,a4w4 則使用 MXFP4 × MXFP4。

專用解碼:每層一個常駐核心。使用 decode_backend=latency_specialized 時,1、2 或 4 個請求的純解碼步驟(或單個請求的 DSpark 驗證塊)會把 71 個 MoE 層中的每一層作為一次 FlyDSL 核心啟動來執行。該啟動使用 256 個 CTA,全部同時常駐在 GPU 的 256 個 CU 上。其內部依次執行:

  • 輸入歸一化,以及帶 SSMax 和 LeRoPE 的查詢投影與潛在 KV 投影;
  • 整個 DSA 索引器:投影、鍵快取寫入、打分以及精確的 top-2048;
  • 在潛在快取上做稀疏 MLA,以 bf16 或 FP8 原地讀取;
  • 輸出投影、softmax top-4 路由,以及共享專家和路由 MXFP4 專家。

嵌入、稠密層 0 和最終歸一化仍留在 vLLM 的模組上。

兩次張量並行 all-reduce 都在核內完成,走的是四塊 GPU 之間的對等記憶體,所以一層只需一次啟動,不用回主機往返。這些核按 1、2、4 行構建,正好覆蓋 vLLM 的 CUDA 圖捕獲尺寸,並在這些圖內重放。其他步驟仍走 vLLM 的層:預填充、混合批次、更寬的批次。該核改編自 ROCm ATOM 的 GLM-5.2 融合解碼,並針對 ML4 的層型別做了擴充套件:

  • 查詢不使用 LoRA;
  • 查詢和索引查詢上使用 SSMax;
  • softmax 路由;
  • NoPE 層;
  • 讀取第 35 層潛在快取的 KV 共享層;
  • 索引共享層。

故障即停。 等待對等 GPU 的常駐核絕不能把伺服器掛死。核內每次等待超過 10 秒就放棄,放棄時會在每個 rank 上置一個粘性故障字,其他 rank 的等待隨即結束。每次前向之後,裝置端斷言會在計算 logits 之前檢查該字。若該字已置位,本步剩餘工作不再執行,worker 中止、引擎停止,失敗步驟產生的任何 token 都不會到達客戶端。這項檢查每次啟動約耗時 1 µs。

這條路徑同樣預設拒絕。如果部署用到了核不支援的任何東西,啟動就會失敗,並給出一個錯誤列出全部原因;所有 TP rank 在任何集合通訊開始前都會對每項檢查達成一致。不支援的特性包括其他 TP 尺寸、專家並行或上下文並行、LoRA、KV 聯結器、其他投機方法,以及未以 a4w4 形式儲存的專家。

TPU

Mistral Large 4 從第 0 天起就執行在 TPU v7x 上,使用 vllm-torchtpu,並支援 Pixtral 視覺編碼器。我們驗證過的配置把主機內張量並行、跨主機資料並行與專家並行結合起來,共使用 16 塊 TPU v7x 晶片。

在 TPU v7x 上服務: 以下命令在一臺主機上啟動單個數據並行 worker:

vllm serve mistralai/Mistral-Large-4-1T-A52B \
  --tensor-parallel-size 8 --data-parallel-size 4 --data-parallel-size-local 1 \
  --data-parallel-start-rank "$TPU_WORKER_ID" \
  --data-parallel-address "$HEAD_IP" --data-parallel-rpc-port 29696 \
  --enable-expert-parallel --enable-ep-weight-filter \
  --dtype bfloat16 --kv-cache-dtype fp8_e4m3

TPU 實現依賴 vllm-torchtpu 的稀疏 MLA 和 DSA 索引器 Pallas 核,並把 MLA 擴充套件為跨層 KV 快取共享。MLA 潛在 KV 快取以 FP8 儲存,稠密線性投影支援 FP8 權重和啟用量化。路由專家使用僅權重量化 FP4(W4A16)。

HiSparse:主機記憶體稀疏 KV 解除安裝

稀疏注意力每個查詢只讀取 2048 個鍵,所以長請求的大部分 KV 不必留在 GPU 記憶體裡。HiSparse 把完整的潛在 KV 放在鎖頁主機記憶體中,保留一份常駐 GPU 工作集加每請求熱緩衝區,未命中時從主機複製選中的行。用 HiSparseConnector 啟用:

--kv-transfer-config '{"kv_connector": "HiSparseConnector", "kv_role": "kv_both",
                       "kv_connector_extra_config": {"host_pool_gib": 440}}'

已知問題。 當張量並行組位於單節點時,還要設定 VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm。一旦 HiSparse 把 KV 溢位到主機,預設的單節點 FlashInfer all-reduce(MNNVL)可能以 CUDA error: unspecified launch failure 讓引擎崩潰(vllm#64327)。修復在 flashinfer#7782 中。

針對 ML4,我們讓 HiSparse 支援了跨層 KV 共享。九個消費層繫結到第 35 層的常駐檢視和主機副本,只保留自己的熱緩衝區和 LRU(每層有自己的索引器)。沒有資料重複,也不需要環境變數。與把共享層實體化相比,同樣的主機池能多容納 14.3% 的 token。我們還修復了:

  • 從多模態模型的文本配置中讀取 index_topk;
  • CUDA 圖捕獲期間出現雙重主機映象。

容量。在單臺 GB300 節點上(TP4、NVFP4),四個 TP rank 共享同一個去重後的主機記憶體池,因此 host_pool_gib 是總量。在 440 GiB 的池下,HiSparse 啟動時可接納 877 萬 KV token,穩態下最高可容納 1125 萬,而僅靠 GPU 只有 197 萬(4.5 倍到 5.7 倍)。GSM8K-200 結果不變(兩者均為 0.975),大海撈針測試在約 9.8 萬(10/10)和約 24.8 萬(5/5)token 下均通過。

HiSparse KV capacity
HiSparse KV 容量

單臺 GB300 節點上一個 TP4 伺服器的 KV token 容量(NVFP4,max_model_len 262k,關閉字首快取):僅 GPU 為 197 萬 token;啟用 HiSparse 並由四個 TP rank 共享 440 GiB 主機記憶體池後,啟動時可接納 877 萬(4.5 倍),穩態下最高 1125 萬(5.7 倍)。複用第 35 層 KV 的層在主機上只存一份。

下一步

  • 上游合入:MLA 跨層 KV 共享、DSpark 滑動視窗 MLA 草稿,以及 DCP 下的 DSpark 和 HiSparse KV 共享支援。
  • 一個在 P/D 中執行 DSpark 的 DCP 解碼節點,以及獨立的草稿塊池,使 DCP 下的 DSpark 保持完整 KV 容量。
  • HiSparse 支援 NIXL 匯入的 GPU 落地,以及面向多節點 TP 的節點共享主機記憶體池。
  • 更多面向低併發延遲的解碼融合。

致謝

vLLM 對 Mistral Large 4 的支援由 Nicolò Lucchesi 和 Mistral Inference 團隊牽頭合作完成。我們還要感謝審閱、設計併合入大量底層工作的 vLLM 維護者:Lucas Wilkinson、Flora Feng、Matt Bonanni 以及 Red Hat 團隊的其他成員;整個 Inferact 團隊;Cyrus Leung 以及整個 vLLM 社群。

在 AMD 方面,我們還要感謝 Hongxia Yang、Felix Marty、Andy Luo、Indira Vats 在 AMD Instinct GPU 上的 Day 0 支援與最佳化,以及 Bowen Bao、Chuan Li、Peng Sun 的總體支援。

來源:vLLM 官方部落格 · vllm.ai

來源: aihot (aihot.news)
開啟原文