TensorRT-LLM 在 NVIDIA Blackwell GPU 上創(chuàng)下了 DeepSeek-R1 推理性能的世界紀錄，Multi-Token Prediction (MTP) 實現(xiàn)了大幅提速。我們在之前的博客[1]中介紹了 DeepSeek-R1 模型實現(xiàn)超低推理延遲的關鍵優(yōu)化措施。本文將深入探討 TensorRT-LLM 中的 MTP 實現(xiàn)與優(yōu)化。

MTP 在推理中的應用

受先前研究工作的啟發(fā)，MTP 用于輔助 DeepSeek-V3 的訓練，在主模型末尾添加額外的 MTP 模塊，并使用這些模塊預測更多 token。這可以將 MTP 的預測范圍擴展到每個位置的多個候選 token，從而提高模型準確性。這些 MTP 模塊還可用于推理過程中的投機采樣，以此進一步降低生成延遲。本章將介紹適用于 LLM 推理的 MTP 投機采樣算法。

背景

投機采樣是一種提高 LLM 推理速度和成本效益的主流技術，其基本原理是生成多個候選 token，與處理單個 token 相比能更加高效地利用 GPU，特別是在計算需求較低的解碼階段。投機采樣技術通常將此過程分為低成本的 draft 階段和并行驗證階段。Draft 階段使用小模型或主模型的一部分層預測候選 token。驗證階段則使用主模型確定接受多少個候選 token，這比每次迭代生成一個 token 高效得多。

圖 1. 驗證示例

圖 1 是一個有關如何驗證并接受這些候選 token 的示例。假設共有 5 個候選 token “ABCDE”，我們將它們拼接到輸入 token “G”，并向主模型輸入總共 6 個 token。采樣后，我們可以得到 6 個不同的預期 token，然后將預期 token 與候選 token 進行比較，并接受最長的前綴匹配 token。此示例中匹配的 token 是 “ABC”。由于 “H” 由主模型預測，且對應的輸入 token “C” 已被接受，因此 “H” 也將被接受。這樣一來，我們可以在單次迭代中接受四個 token。MTP 也使用此方法驗證并接受候選 token。在 MTP 的 draft 階段有兩種 MTP 方法：MTP Vanilla 和 MTP Eagle。它們適用于不同的推理場景。

MTP Vanilla

圖2. MTP Vanilla，其中ti為輸入token，di為預測的候選token，K為MTP模塊數(shù)量，hin為第n個MTP模塊的隱藏狀態(tài)。需注意h0表示主模型的隱藏狀態(tài)。

（注：圖改編自DeepSeek V3技術報告https://arxiv.org/pdf/2412.19437）

MTP Vanilla 方法與 MTP 訓練更為相似，該方法按序調用不同 MTP 模塊來預測多個候選 token，支持包含多個不同 MTP 模塊權重的模型 checkpoint，且每個 MTP 模塊均擁有獨立的 KV 緩存。

圖2是MTP Vanilla的推理流程。在預填充階段，假設總共有四個輸入token，我們將在主模型前向傳播后獲得輸出tokent5和隱藏狀態(tài)。輸出token將附加到輸入token上，然后移除第一個token，得到t2到t5作為第一個MTP模塊的輸入token。主模型的隱藏狀態(tài)將直接作為第一個MTP模塊的輸入，用于預測第一個候選token。對于接下來的幾個MTP模塊，我們將新生成的候選token與對應最后一個輸入token的隱藏狀態(tài)附加到輸入token和隱藏狀態(tài)中。然后，我們將移除第一個token以準備下一個MTP模塊的輸入。通過這種方式，我們可以盡可能多地保留主模型中的信息，有助于MTP模塊做出更準確的預測。

在生成階段會有一點差異。預測的tokent5和候選token將作為主模型的輸入。在主模型前向傳播后，我們將進行驗證以獲取被接受的token。此示例假設j個候選tokend6~$d_{j+5}$被接受，然后準備MTP模塊的輸入與預填充階段不同，我們將此前所有接收token中的最后K個token及其對應的隱藏狀態(tài)作為第一個MTP模塊的輸入，該示例中最后一個被接受的token是t_{j+6}，則輸入為t_{j+7-K}~t_{j+6}，然后我們可以獲得第一次MTP模塊前向傳播后的第一個候選token。對于后續(xù)的MTP模塊，我們可通過與預填充階段中的MTP模塊類似的方式準備它們的輸入，所有這些MTP模塊都具有相同的輸入序列長度。在預測完所有新的候選token后，由于驗證階段很多token沒有被接收，需要將這些被舍棄的候選token的鍵/值從模型的KV緩存中清除，以確保后續(xù)計算的正確性。

MTP Eagle

圖 3. 使用與圖 2 相同標記的 MTP Eagle

（注：圖改編自 DeepSeek V3 技術報告 https://arxiv.org/pdf/2412.19437)

MTP Eagle 可以被視為 Eagle 投機采樣方法的變體，但目前僅支持鏈式解碼。該方法復用同一個 MTP 模塊并重復多次來預測候選 token，因此 MTP Eagle 用于支持僅包含 1 個 MTP 模塊的模型 checkpoint，官方 DeepSeek-V3 和 DeepSeek-R1 的checkpoint 中僅包含 1 個 MTP 模塊，適用于 MTP Eagle。其與 MTP Vanilla 的另一個差異在于 KV 緩存，MTP Eagle 方法中的 MTP 模塊在預測多個候選 Token 時復用相同的 KV 緩存。

圖 3 是一個 MTP Eagle 示例。在預填充階段，第一個 MTP 模塊的輸入與 MTP Vanilla 相同，但是在后續(xù) MTP 模塊前向傳播過程中有些不同，首先 MTP Eagle 使用相同的 MTP 模塊預測候選 token 并復用相同的 KV 緩存，其次我們只需輸入 1 個 token 的 token ID 和其隱藏狀態(tài)，該 token 是前一個 MTP 模塊預測得到的候選 token。通過這種方式，我們只使用 1 個 MTP 模塊就可以預測總共 K 個候選 token。

生成階段的驗證過程與 MTP Vanilla 類似。在獲得接受的 token 后，我們將所有被接收的 token 及其對應的隱藏狀態(tài)作為第一個 MTP 模塊前向傳播的輸入，與需要存儲過去 token 和隱藏狀態(tài)的 MTP Vanilla 不同，這種方法實現(xiàn)起來要簡單得多。后續(xù) MTP 模塊前向傳播采用與預填充階段相同的方法來準備輸入。在預測所有候選 token 后，與 MTP Vanilla 類似，需要清除模型 KV 緩存中所有被舍棄 token 的鍵 / 值。

TensorRT-LLM 中的 MTP 實現(xiàn)方式

基本實現(xiàn)方式

TensorRT-LLM 提供了兩條不同的 MTP 實現(xiàn)路徑，一條對應 MTP Vanilla，另一條對應 MTP Eagle。MTP Eagle 是 DeepSeek-V3 和 DeepSeek-R1 模型的默認路徑。

圖 4. TensorRT-LLM 中的 MTP 工作流

圖 4 是 TensorRT-LLM 中的完整 MTP 工作流。兩條路徑共用同一運行時工作流，差別則在于 MTP 模塊的前向傳播。在預填充階段，輸入中不包含候選 token，TensorRT-LLM 模型引擎從請求中獲取輸入 ID，并將其輸入到模型引擎進行前向傳播以獲取下一個 token 和隱藏狀態(tài)。隨后準備 MTP 模塊輸入，MTP 模塊通過前向傳播來預測候選 token。

生成工作流更為復雜，需要同時進行驗證和預測候選 token，預測的新 token 和候選 token 為主模型的輸入。在主模型前向傳播后，我們可以從輸出的 logits 中采樣并獲取新的 token，然后將它們與輸入的候選 token 進行比較，以獲取最終接受的 token 來完成對候選 token 的驗證。我們將使用接受的 token 和隱藏狀態(tài)啟動一個新的 draft 階段，該階段使用 MTP 層預測下一輪迭代的新候選 token。最后，我們需要回滾 KV 緩存以清除與被拒絕 token 對應的鍵值對。

除了回滾 KV 緩存外，所有這些過程均在模型引擎前向傳播中完成。這樣我們可以使用一個模型引擎支持 MTP 推理，并且 MTP 更容易與其他功能兼容，如 CUDA graph 和重疊調度器等。啟用 CUDA graph 后，驗證和 draft 階段均可在同一個 graph 中運行，大幅降低了 CPU 開銷。

MTP 模塊

圖 5. MTP 模型架構

圖 5 介紹了 MTP Vanilla、MTP Eagle 的基本模型架構以及基本 MTP 模塊設計。由于 MTP Vanilla 需要 K 個輸入 token，若接受的 token 數(shù)量少于輸入 token 數(shù)量（即 j

MTP 模塊遵循了 DeepSeek-V3 中的設計，其中嵌入層和輸出頭與主模型共享，這可以減少 GPU 顯存消耗。

支持 MTP 的注意力機制

注意力機制也是支持 MTP 推理的重要組成部分，變化主要體現(xiàn)在生成階段的注意力內核上。在普通請求情況下，生成階段僅有一個輸入 token，但在 MTP 情況下，將有 K+1 個輸入 token。由于 MTP 按順序預測額外的 token，預測的候選 token 會被鏈式連接，盡管我們有 MTP Eagle 路徑，但目前只支持基于鏈式的 MTP Eagle。因此，一個因果掩碼即可滿足注意力內核對 MTP 的支持需求。在我們的實現(xiàn)方式中，TensorRT-LLM 將在 Hopper GPU 上使用FP8 FlashMLA生成內核，在 Blackwell 上使用 TensorRT-LLM 定制注意力內核來實現(xiàn)更高的性能。

如何運行帶有 MTP 功能的 DeepSeek 模型

如要運行帶有 MTP 功能的 DeepSeek-V3 / R1 模型，請使用 examples/llm-api/quickstart_advanced.py 并添加以下選項：

cd examples/llm-api
python quickstart_advanced.py--model_dir--spec_decode_algoMTP--spec_decode_max_draft_lenN

如要在有 MTP 功能的情況下進行最低延遲性能基準測試，請按照本文檔準備數(shù)據集，然后按照以下步驟操作：

YOUR_DATA_PATH=


cat>./extra-llm-api-config.yml<

	

	MTP 優(yōu)化：寬松接受

	DeepSeek-R1 是一種先輸出一些思考 token，然后為用戶提供實際輸出的推理模型。其思考過程通常會消耗大量 token，而思考過程的輸出質量對最終答案的影響有限，因此我們希望采用一種更積極的接受策略——寬松接受以加快思考解碼階段的速度，該策略將權衡加速幅度與輸出質量。從實驗結果來看，寬松接受對輸出質量的影響不大。

	寬松接受

	

	圖 6. 使用 MTP nextn=4 和 top-3 的寬松接受示例。

	如圖 1 所示，在之前的驗證和接受過程中，我們使用 top-1 從主模型的 logits 中采樣以獲取“預期” token，此時僅有一個選項與候選 token 進行比較，我們稱之為“嚴格接受”。

	至于寬松接受，我們首先從 logits 中采樣 top-N 個 token，將有更多候選 token 與輸入的候選 token 進行比較。為確保接受的 token 盡可能準確，我們還引入了一個概率閾值 delta，我們可通過對 logits 應用 softmax 函數(shù)獲得 token 概率。在獲得前 N 個候選 token 后，我們會移除概率小于（前 1 個概率 - delta）的 token。如此一來，我們可獲得多個候選 token，且所有候選 token 的概率均較高。隨后，我們可以將輸入候選 token 與這些候選 token 進行比較，若其中一個匹配，即可接受該候選 token，從而提高接受率。圖 6 比較了嚴格接受與寬松接受。

	需注意，寬松接受僅在思考階段使用，而嚴格接受仍在非思考階段使用。目前寬松接受僅支持 DeepSeek-R1 模型。

	如何運行使用寬松接受策略的

	DeepSeek-R1 模型

	如要運行使用 MTP 寬松接受的 DeepSeek-R1 模型，請使用 examples/llm-api/quickstart_advanced.py 并添加以下選項：

	
cd examples/llm-api
python quickstart_advanced.py--model_dir--spec_decode_algoMTP--spec_decode_max_draft_lenN--use_relaxed_acceptance_for_thinking--relaxed_topk10--relaxed_delta0.6

	

	如要在使用 MTP 寬松接受策略的情況下進行最低延遲性能基準測試，請按照本文檔準備數(shù)據集，然后按照以下步驟操作：

	
YOUR_DATA_PATH=


cat>./extra-llm-api-config.yml<

	

	評估

	通過 MTP 投機采樣實現(xiàn)加速

	

	圖 7. DeepSeek-R1-FP4 671B 在不同 MTP next-n 設置下的最低延遲性能

	我們在 GPU 上測試了 DeepSeek-R1-FP4 模型在不同 MTP next-n 設置下的最低延遲 (batch size=1) 性能。測試中設置 MLA TP=8，MoE EP=2，總共使用 10 個輸入請求進行測試，輸入序列長度和輸出序列長度分別為 1K 和 2K。從圖 7 可以看出，在 8 個 GPU 上 MTP=3 可幫助實現(xiàn)最佳的最低延遲性能，與基線 nextn=0 相比提速 2.16 倍。而借助寬松接受策略可進一步減少最低延遲，提速 2.33 倍。我們還評估了 CUDA graph 和重疊調度器的優(yōu)勢。在此類最低延遲場景中，CUDA graph 可實現(xiàn)平均提速 7.22 倍，而重疊調度器則使延遲平均提速 1.03 倍。

	使用寬松接受策略時的準確性研究

	

	圖 8. 使用寬松接受策略時的消融結果，MTP nextn=3、top-10 和 delta=0.6。

	我們在不同數(shù)據集上驗證了寬松接受。圖 8 是 DeepSeek-R1-FP4 模型使用寬松接受策略時的消融結果，與嚴格接受相比，寬松接受對輸出質量的影響不大，精度略微下降，個別數(shù)據集上還達到了更高的精度。

	工作展望

	樹式投機采樣支持

	

	圖 9. 鏈式與樹式投機采樣技術的比較

	TensorRT-LLM PyTorch 后端目前僅支持 MTP Vanilla 和 MTP Eagle 等基于鏈式的投機采樣。為了提高接受率，Eagle2 和 Eagle3 等此前的先進方法廣泛采用樹式投機采樣技術，TensorRT-LLM 中的 MTP 也可擴展以支持樹式投機采樣技術。圖 9 比較了鏈式方法與樹式方法，無論是全樹式還是動態(tài)樹式方法均有助于擴展候選組合，從而提供更多候選 token 供選擇。

	Eagle3 支持

	Eagle3 同樣是一項重要的技術。Eagle3 論文的研究結果表明，使用不同層級的隱藏狀態(tài)預測候選 token 可顯著提高接受率。由于 TensorRT-LLM 已支持 Eagle-3，未來我們還計劃訓練一個 Eagle3 head，結合 DeepSeek-V3 / R1+Eagle3 以實現(xiàn)更好的加速。

	致謝

	支持和優(yōu)化 TensorRT-LLM 中的 MTP 是一項了不起的跨團隊合作成果。我們謹向所有為此做出貢獻的人士致以誠摯感謝，該項目涉及多個技術層面的系統(tǒng)與算法協(xié)同設計方法，包括內核優(yōu)化、運行時增強、算法改進以及性能測量與分析等。特別感謝 DeepSeek 團隊開發(fā)了這套 MTP 方法，為本文奠定了基礎。

	作者

	張國銘

	NVIDIA 性能架構師，目前主要從事大模型推理優(yōu)化。

	李繁榮

	NVIDIA Compute Arch 部門高級架構師，目前主要從事大模型推理優(yōu)化。