成立半年估值$120億，Thinking Machines Lab首次出擊，為何發了篇技術博客？

文｜四木洋子

在 AI 巨頭們爭相比拼聲量、瘋狂搶奪人才的浮躁時代，有一家公司選擇了截然不同的路徑 —— 它就是 Thinking Machines。

這家由 OpenAI 前 CTO 米拉·穆拉蒂（Mira Murati）創立的新興 AI 公司，成立至今還不到半年，卻已經匯聚了一眾來自 OpenAI、Meta 等頂級 AI 機構的精英人才。

更令人矚目的是，今年夏天，Thinking Machines 在尚未發布任何產品的情況下，就完成了20億美元的種子輪融資，公司估值飆升至120億美元。

在數月的低調后，這家備受矚目的AI新星終于有了新動作——上周，它悄然推出了自己的技術博客專欄，并發布了首篇技術文章 "Defeating Nondeterminism in LLM Inference"（《戰勝大模型推理中的不確定性》）。

*原文：https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference

這篇深度技術文章直指一個困擾業界已久的核心難題：大語言模型前向推理過程中的不確定性問題。

這個問題從 LLM 誕生之初就如影隨形，但業界一直苦于找不到理想的解決方案。

而 Thinking Machines 的這篇文章不僅從算法、系統和硬件三個維度深入剖析了問題的根源所在，更難能可貴的是，他們還提出了切實可行的解決方案——這或許正是這家"神秘"公司以此作為首秀的原因所在。

經過實驗驗證，Thinking Machines 提出的方案體現出雙重突破：

• 一方面，它讓大語言模型能夠提供高度可復現的前向推理結果，解決了長期困擾研究者的一致性問題；

• 另一方面，該方案從根本上改善了強化學習的訓練動態，成功實現了基于大語言模型的真正同策略探索與利用平衡——這一直是這個領域亟待攻克的核心難題。

這些突破性成果印證了 Thinking Machines 團隊在 AI 前沿研究中所具備的深厚洞察力和卓越的系統性研究能力。

我們把本文的閱讀心得，分為內容總結、思考點評和原文翻譯三部分，供大家按需了解：

內容總結

為什么問大模型相同的問題，會得到不同的答案？

這個問題雖然有時會有點惱人，但并不難理解——因為語言模型中的結果涉及“采樣”，這本身就是一個"隨機選詞&輸出"的過程。

如果說隨機采樣約等于抽簽，那么采樣溫度（temperature）則直接決定了抽簽的權重分布。

但現在的問題在于，即便將采樣溫度調整到 0，我們仍然不能保證模型會輸出確定的答案。

*temperature = 0 意味著大模型預測過程總是選擇最高概率的詞元，這被稱為貪婪采樣，其結果總是確定的。

這就有點令人費解了。

Thinking Machines 的這篇文章從一個我們都不陌生的現象切入——同樣的問題，大語言模型有時會給出不同的答案。這看似隨機的表現，實際上揭示了大模型推理過程中一個深層次的技術問題：推理的"不確定性"（nondeterminism）。

文章的核心發現令人意外：這種不確定性的根源，竟然來自于推理批次不變性（batch invariance）的缺失。

什么意思呢？讓我們從 GPU 的工作機制說起。為了榨干硬件的每一分算力，GPU 會根據當前的批次大?。╞atch size）——也就是服務器同時處理的請求數量——動態調整自己的并行策略和計算單元的調度方式。批次大了，就用這種策略；批次小了，就換那種策略。

問題就出在這里。

由于浮點數運算本身的特性——它并不嚴格滿足數學上的加法結合律，不同的計算順序會產生微小的數值偏差。這些看似微不足道的差異，在 Transformer 層層疊疊的前向傳播過程中被逐步放大，最終導致了一個令人困擾的結果：相同的輸入，在不同的批次環境下，竟然會產生不同的輸出。

簡而言之，GPU 為了優化效率而采用的靈活調度策略，意外地成為了大模型推理不確定性的幕后推手。一個追求極致性能的設計，卻在無意中為模型的"善變"埋下了伏筆。

那么，Thinking Machines 是怎么解決的呢？簡單概括，讓模型的計算對批次大小敏感因素“免疫”。

在 GPU 計算中，每個操作都通過專門的計算內核（kernel）來執行，這就要求所有內核都必須具備批次不變性——即無論輸入批次如何變化，計算邏輯都應保持一致。

然而，在 Transformer 模型的實際運行中，某些操作天然地會因批次規模的改變而產生計算行為上的差異，這類操作主要集中在各種歸約計算（reduction）環節。認識到這一挑戰，Thinking Machines針對性地對三個核心的歸約操作 —— RMSNorm、矩陣乘法以及注意力機制 —— 進行了專門的優化設計。

• RMSNorm：讓每一個樣本都單獨在一個內核上完成 RMSNorm 的計算。這樣一來，無論批次里有多少樣本，都不需要跨樣本通信，完全在自己的內核里完成，得到的結果與單獨算該樣本時是一致的 。

• 矩陣乘法：簡單起見，可以把矩陣乘看作“逐元素相乘再累加”的過程。實現批次不變性的原則同樣是讓每個輸出元素的累加都在單一內核中按固定順序完成。具體方法是將輸出矩陣劃分成很多小塊（tile），每個內核負責自己那個塊的所有計算，并把該塊內部需要累加的部分串行做完。這樣每個輸出塊的加法順序是固定的，不受別的核心影響。

• 注意力機制：注意力計算涉及更復雜的步驟（比如 KV 的點積 和 Softmax 歸一化等），而且還牽涉序列長度（token 數量）維度的變化。解決方案是固定歸約塊的大小。簡單理解：原先注意力在解碼時為了適應不同 token 數，可能動態地選擇需要多少塊來并行計算；現在改為預先固定每一塊的大小，這樣無論總共多少 token，就算需要的塊數變多，每個塊內部的計算順序和規模都是一致的。

通過這三種歸約計算方法的優化，Thinking Machines 實現了 Transformer 模型前向計算的批次確定性。

有了確定性之后，除了能夠實現可復現的模型推理結果，還能顯著改善 RL 訓練效果。

RL 的基本原理是通過獎勵機制來強化模型的正確輸出。然而，如果每次答案都有細微差別，訓練數據就會變得很「嘈雜」，而更一致的響應能讓整個 RL 訓練過程變得「絲滑」很多。

在 RL 領域，On-policy（同策略）和 Off-policy（異策略）是兩種核心訓練范式，它們的區別在于學習時使用的"行為策略"與"目標策略"是否為同一個策略。目前，很多研究工作都在強調同策略（On-policy）優化的重要性。

但這里出現了一個有趣的問題：GPU 的不確定性既可能影響采樣時的策略表現（生成不同的數據），也可能影響訓練時的策略表現（產生不同的梯度計算結果）。

這就導致了一個核心矛盾——明明是同一個策略，但在采樣和訓練階段跑出來的結果卻不一致。結果，原本設計為同策略的訓練被迫變成了異策略（Off-policy）訓練，這種現象可以稱為"偽同策略 RL"（Faked On-Policy RL）。

解決方案是通過確定性推理（通過 batch-invariant kernel 等方式），確保推理和訓練逐位一致，保證真正的同策略訓練。

思考點評

大咖云集的 Thinking Machines 為啥選擇這項工作作為自己的首次亮相？

答案很簡單：這個問題實在太重要了，而能夠針對這一問題提出系統性解決方案的團隊卻寥寥無幾。

Faked On-policy RL for LLM，這是當前世界范圍內普遍存在的痛點問題。

令人意外的是，即便是那些在前沿的 AI 頂會上發表、明確宣稱"我們采用同策略強化學習"的研究工作，實際情況往往是：研究者們在算法層面精心打磨，但承載算法實現的底層框架和算子等系統組件，卻在無形中背離了 On-Policy RL 的初衷和愿景。

而這一根本性問題，恰恰是目前開源社區（如 veRL、MS-Swift 等 RL 訓練框架的提供方）難以攻克的 ——原因在于維護人員往往術業有專攻：擅長框架設計的未必精通算子開發，而算子專家又不一定能駕馭整體框架架構。這種技能壁壘的存在，使得端到端的解決方案變得格外珍貴。

目前我們看到的相似研究風格的博客包括但不限于以下幾個：

• 如 Thinking Machines 這篇文章中所引用的，微軟高劍峰組在前兩個月發現了目前 RL for LLM 訓練中普遍存在的這一問題。雖然他們在《Your Efficient RL Framework Secretly Brings You Off-Policy RL Training》中提出了緩解方案，但坦率地說，這種方式仍顯得不夠優雅且不夠徹底。

• 同樣在算法與系統（算子）層面尋求改進的，還有最近以兼職身份加入 Thinking Machines 的 Songlin Yang 提出的 DeltaNet 系列 也是類似的風格。

• 其他更多的博客文章則側重于單點優化和改進，如系統（算子）層面 Tri Dao 的 FlashAttention 系列，以及純算法層面的典型如 John Schulman 的《Approximating KL Divergence》。

而 Thinking Machines 的做法，實際上是拿軟件（算法）—— 系統—— 硬件這三者串聯所產生的推理不確定性問題，進行了一次由表及里的深度剖析，系統性地梳理推理過程中不確定性的根本來源。

他們的研究路徑頗有章法：

從「數值計算」的底層細節，到「系統（內核）原理」的中層機制，再到算法層面的優化改進，層層遞進，向外界展示了那些在一般研究人員眼中"平平無奇"的技術細節，是如何最終影響到終端模型訓練效果的。

更重要的是，在解決了高劍峰等人發現的「現在的 RL for LLM 都是偽同策略 RL」這一關鍵問題之后，才有可能在此基礎上構建出真正的滿血版 RL-ed LLMs（經過強化學習訓練的語言模型）。

而這樣的模型，又將為其他更宏大的目標提供堅實支撐——比如構建多模態通用 Agent，或是提供面向客戶的優質 RL 解決方案等。

在文章的結尾部分，Thinking Machines 的研究人員特別強調：

不能忽視推理過程中的不確定性，只有抓住并深入分析這些不確定性，找到切實可行的解決方案，才能最終實現真正的同策略強化學習，進而把 RL for LLM 這件事做得更加扎實。

對其他投身大模型研究的人來說，Thinking Machines 這次的工作傳遞出一個清晰的信號：

想要繼續提升模型表現，既需要自頂向下的宏觀思考，也需要自底向上的細致打磨，端到端解決問題的系統思維和扎實的動手能力缺一不可——「宏觀規劃 + 微觀操作」的組合拳必須打得漂亮。

以下為原文翻譯：

重復性是一切科學進步的基石。然而，從大語言模型中獲得可重復的結果卻異常困難。

例如，你可能會發現，多次向 ChatGPT 提出相同的問題會得到不同的結果。這本身并不令人驚訝，因為從語言模型中獲取結果涉及“采樣”，這是一個將語言模型輸出轉換為概率分布并概率性地選擇一個標記的過程。

更令人驚訝的可能是，即使我們將溫度調整到 0（這意味著 LLM 總是選擇最高概率的標記，這被稱為貪婪采樣），LLM API 在實踐中仍然不是確定性的。即使在你自己的硬件上使用 vLLM 或 SGLang 等開源推理庫運行推理，采樣仍然不是確定性的。

但為什么 LLM 推理引擎不是100%確定性的呢？一個常見的假設是，浮點非關聯性和并發執行的某種組合導致了基于哪個并發核心首先完成的不確定性。我們將此稱為 LLM 推理不確定性的“并發 + 浮點”假設。例如，最近的一篇 arXiv 文章寫道：

“GPU 中的浮點運算表現出非關聯性，這意味著 (a + b) + c ≠ a + (b + c)，這是由于有限精度和舍入誤差造成的。此屬性直接影響 Transformer 架構中注意力分數和詞表輸出分數（logits）的計算，其中跨多個線程的并行操作可能會根據執行順序產生不同的結果?！?/p>

你還可以在其他地方找到“并發 + 浮點”假設的重復，例如：“存在速度權衡，為了使端點快速，使用了 GPU，它進行并行不確定性計算。任何現代 GPU 神經網絡計算都將受制于這些”，或“因為 GPU 是高度并行化的，所以每次執行加法或乘法的順序可能不同，這可能會導致輸出中的微小差異”。

雖然這個假設并非完全錯誤，但它并未揭示全貌。例如，即使在 GPU 上，反復對相同的數據執行相同的矩陣乘法運算，每次得到的結果在二進制層面都是完全相同的。我們確實在用浮點數運算，而且 GPU 也確實有很多并發操作。那為什么在這個測試里沒有看到不確定性呢？

為了理解 LLM 推理不確定性的真正原因，我們必須深入研究。

不幸的是，即使定義 LLM 推理的確定性意味著什么也很困難。一些令人困惑的事情是，以下所有陳述同時都是正確的：GPU 上的一些內核是不確定性的。

然而，語言模型正向傳遞中使用的所有內核都是確定性的。

此外，LLM 推理引擎（如 vLLM）的前向計算也可以聲稱是確定性的。

盡管如此，從使用推理服務器的任何人的角度來看，結果都是不確定性的。

在這篇文章中，我們將解釋為什么“并發 + 浮點數”假設未能達到目標，揭示導致 LLM 推理不確定性背后的真正原因，并解釋如何克服不確定性并在 LLM 推理中獲得真正可重現的結果。

最初的“原罪”：浮點數的非結合性

在討論不確定性之前，有必要解釋一下為什么會存在數值差異。畢竟，我們通常認為機器學習模型是遵循交換律或結合律等結構規則的數學函數。我們的機器學習庫難道不應該為我們提供一個“數學上正確”的結果嗎？

罪魁禍首就是浮點數的非結合性。也就是說，對于浮點數：

(a + b) + c ≠ a + (b + c)

具有諷刺意味的是，正是這種對結合律的破壞使得浮點數變得有用。

浮點數之所以有用，是因為它們允許“動態”的精度級別。

為了便于解釋，我們將使用十進制（而非二進制），其中浮點數采用 尾數 * 10^指數 的格式。我們還將使用 3 位尾數和 1 位指數。

例如，對于數值 3450，我們可以精確地表示為 3.45 * 10^3。我們也可以表示小得多的值，如 0.486，表示為 4.86 * 10^-1。通過這種方式，浮點數允許我們表示非常小和非常大的值。在科學領域，我們可能會說浮點數允許我們保持恒定數量的“有效數字”。

如果你將兩個具有相同指數的浮點數相加，它看起來類似于整數加法。例如，123 (1.23 * 10^2) + 456 (4.56 * 10^2) 的結果是 579 (5.79 * 10^2)。

但是，當我們相加兩個具有不同指數的浮點數時會發生什么，比如 1230 和 23.4？在這種情況下，精確結果是 1253.4。然而，我們一次只能保持 3 位精度。因此，浮點加法將丟棄最后兩位，得到值 1.25 * 10^3（或 1250）。

Figure 1：我們需要 3 位精度來表示 1230，也需要 3 位精度來表示 23.4。然而，將這兩個數相加會得到一個需要 5 位精度才能表示的數（1253.4）。我們的浮點格式必須丟棄末尾的 34。從某種意義上說，我們在相加之前，實際上已經將原始的 23.4 四舍五入到 20.0 了。

然而，在這一點上，我們已經丟失了信息。請注意，每當我們相加兩個具有不同“尺度”（即不同指數）的浮點數時，都可能發生這種情況。而相加具有不同指數的浮點數是經常發生的。事實上，如果能保證從不需要不同的指數，我們就可以直接使用整數了！

換句話說，每次我們以不同的順序將浮點數相加時，都可能得到一個完全不同的結果。舉一個極端的例子，對這個數組求和，根據相加順序的不同，可能會有 102 種不同的結果。

盡管這是導致輸出不一致的根本原因，但它并沒有直接回答不確定性從何而來。它沒有幫助我們理解為什么浮點值會以不同的順序相加、何時發生這種情況以及如何避免。

答案在于內核（kernels）是如何實現的。

為什么內核不總是以相同的順序相加數字？

既然運算順序會影響結果，那么為什么 GPU 內核在執行加法時，順序不是固定的呢？

這就回到我們之前提到的一個流行假設：GPU 上的多個線程（并發執行）完成的順序不確定，如果加法（累加）操作的最終結果依賴于這些不確定的完成順序（例如，如果多個線程同時嘗試向同一個內存位置累加，需要用到“原子加法”），那么最終的累加順序就會不確定，從而導致結果不確定。

令人困惑的是，盡管這（并發與浮點非關聯性）可能會導致內核出現不確定性，但在 LLM 推理的不確定性中，并發（以及原子加法 atomic adds）實際上完全沒有參與其中！ 要解釋真正的罪魁禍首是什么，我們首先需要理解為什么現代 GPU 內核很少需要使用原子加法。

*atomic add = 一種讓很多線程同時往同一個變量里“加數”而不丟結果的機制，但加法順序不固定，所以浮點數計算可能有細微差異。

什么時候需要用到原子加法？

通常情況下，GPU 會在許多“核心”（即 SM，Streaming Multiprocessors）上并發地啟動一個程序。由于這些核心之間沒有天然的同步機制，如果它們需要彼此通信，就會產生挑戰。比如說，如果所有核心都必須把結果累加到同一個元素上，就可以使用“atomic add”（有時也叫“fetch-and-add”）。

atomic add 是“不確定性的” —— 結果的累加順序完全取決于哪個核心先完成。

具體來說，假設你要用 100 個核心對一個 100 元素的向量做歸約（reduction）（例如 torch.sum()）。雖然你可以并行加載這 100 個元素，但最終必須把它們歸約成一個結果。一種實現方式就是使用某種 atomic add 原語，在這種情況下，硬件能保證所有加法都會被處理，但不能保證執行的順序。

Figure2：atomic add（原子加法） 確保了每個核心的貢獻都會體現在最終的總和中。然而，它并不保證這些貢獻是按照什么順序被加進去的。順序完全取決于哪個核心先完成，這是一個不確定性特性。因此，執行同一個并行程序多次，可能會產生不確定性輸出。

這通常就是人們所說的"不確定性"——用完全相同的輸入執行同一個內核兩次，卻得到截然不同的結果。

這種現象被稱為運行間不確定性（run-to-run nondeterminism）：你用相同的依賴環境運行同一個 Python 腳本兩次，結果卻不一樣。

雖然并發的 atomic add 操作確實會讓內核變得不確定，但對于絕大多數內核來說，atomic add 其實并非必需品。 事實上，在 LLM 的典型前向傳播過程中，通常連一個 atomic add 操作都不會出現。

這個現象可能會讓人感到驚訝——畢竟在對歸約（reduction）操作進行并行化時，atomic add 確實能帶來性能收益。 但 atomic add 最終之所以并非必需，主要有兩個原因：

首先，在“batch（批次）”維度上通常已經有足夠的并行性，因此我們并不需要在歸約維度上再做并行化。舉個例子，假設我們面對的不是對單個 100 維向量做歸約，而是需要同時對 500 個向量做歸約。在這種情況下，我們完全可以讓每個核心負責處理一個完整的向量，讓不同核心操作不同的向量，這樣就自然地避免了競爭條件。

其次，隨著時間的推移，大多數神經網絡庫都逐漸采用了多種巧妙的策略，在不犧牲性能的前提下實現了計算的確定性。比如，我們可以采用"分塊歸約（split reduction）"或"樹狀歸約（tree reduction）"的方式：將一個 100 元素的歸約任務拆分成 5 個獨立的 20 元素歸約（從而實現五路并行處理）。

接下來，為了將剩余的這 5 個結果合并起來，我們有兩種選擇：

執行一次單獨的“清理歸約（clean-up reduction）” —— 雖然這一步不再并行，但由于處理的元素數量很少，計算代價相當低廉；

或者采用信號量（semaphore）機制，它能夠確保每個并發的線程塊按照確定的順序進行累加操作。

這樣一來，我們既保持了高效的并行計算，又避免了不確定性帶來的困擾。

由于以上兩個因素，對于絕大多數神經網絡運算來說，避免使用 atomic add 帶來的性能損失幾乎可以忽略不計。

當然，仍有一些常見操作在避免 atomic add 時會造成顯著的性能損失，比如 PyTorch 中的 scatter_add 操作（a[b] += c）。不過在 LLM 領域，唯一常用到這類操作的場景就是 FlashAttention 的反向傳播了。

這里有個有趣的事實：你知道嗎？被廣泛使用的 Triton 實現版本，在算法層面其實與 Tri Dao 的 FlashAttention-2 原版論文并不相同。

標準的 Triton 實現會在反向傳播中進行額外的重計算，以此來規避 atomic add 的使用，但代價是增加了 40% 的 FLOPs！這是一個典型的以計算換確定性的權衡。

然而，LLM 的前向傳播并不涉及需要 atomic add 的操作。因此，LLM 的前向傳播實際上具備了運行間確定性（run-to-run deterministic）的特性。

從推理服務器的角度來看，這意味著什么呢？它是完全確定性的——只要用戶的請求完全相同，服務器總是會返回相同的輸出結果。

正如維基百科所定義的："一個確定性算法，是指在給定某個輸入時，總是會產生相同輸出的算法。"而 LLM 的前向推理恰好滿足了這個定義。

在這里，給定完全相同的輸入（即推理服務器正在處理的完全相同的請求），前向傳播總是會產生完全相同的輸出。

然而，僅僅前向傳播本身具有"確定性"，并不足以保證包含它的整個系統也是確定性的。

舉個例子，如果我們的請求輸出會受到并行用戶請求的影響（比如在 batch normalization 的情況下），那情況就大不相同了。

由于每個獨立請求都無法預知其他并行請求的內容，從單個請求的視角來看，整個 LLM 推理系統仍然表現出不確定性！

事實上，我們的請求輸出確實會受到并行用戶請求的影響。

這種現象并非源于批次間的信息泄露，而是因為我們的前向傳播缺乏批次不變性（batch invariance）——換句話說，請求的輸出結果會依賴于前向傳播時的具體批次大小。

批次不變性與確定性

為了說明批次不變性的問題，我們將系統簡化，專門考察矩陣乘法（matmul）操作。

在討論中，我們可以假設所有的矩陣乘法實現都具有"運行間確定性（run-to-run deterministic）"——盡管這個假設在實際情況中并非完全準確，但大多數常見的矩陣乘法實現確實表現出這種特性。

然而，一個令人意外的現象是：這些實現往往不具備"批次不變性（batch-invariant）"。也就是說，當批次大小發生變化時，批次中每個元素的計算結果可能會隨之改變。

從數學的直覺來看，這確實是一個頗為反常的現象。按理說，矩陣乘法應該對批次中的每個元素"獨立"進行計算——無論批次中包含其他什么元素，也無論批次的總體大小如何，都不應該影響到批次中任意特定元素的最終結果。

但令人困惑的是，實際的觀察結果卻與這種數學直覺相悖。

請注意，這里指的是"運行間確定性（run-to-run deterministic）"。如果你在同一環境中多次運行這個腳本，它會確定性地返回完全相同的結果。

不過，這種確定性是有條件的 —— 它并非"跨硬件/軟件版本不變"。當你的GPU型號或PyTorch版本發生變化時，可能會得到不同的數值結果。但關鍵在于：在同一個運行環境下，結果始終是可重現的。

然而，問題出現在更復雜的場景中。當這種缺乏批次不變性的內核被集成到大型推理系統時，整個系統的行為就變得不可預測了。

想象一下這樣的情況：當你向推理服務端點發送請求時，服務器當前的負載情況對你來說是"黑盒"的 —— 你無法預知也無法控制。而這個負載直接決定了內核運行時的批次大小，進而影響每個看似獨立的請求的最終輸出結果！這就是為什么同樣的輸入在不同時刻可能產生不同輸出的根本原因。

Figure4：盡管從整體上可以說推理服務器是“確定性的”，但對單個用戶來說情況不同。從單個用戶的角度來看，其他并發用戶并不是系統的“輸入”，而是系統的一個不確定性屬性。這使得從每個用戶的角度來看，LLM 推理是“不確定性的”。

當某種內核缺乏批次大小不變性時，一旦遇到不確定的外部因素（比如服務器負載的波動），整個系統就會變得不確定起來。

實際上，幾乎所有 LLM 推理端點的不確定性都可以追溯到同一個根源：服務器負載的動態變化會導致批次大小的不可預測變動！

這種不確定性并非 GPU 所獨有——無論是基于 CPU 還是 TPU 的 LLM 推理端點，都會面臨同樣的挑戰。

那么，如果我們希望在推理服務器中消除這種不確定性，就必須在內核層面實現批次不變性。要理解具體該如何實現，我們首先需要弄清楚：為什么內核最初就不具備批次不變性呢？

我們該如何讓內核實現批次不變性？

為了讓 Transformer 的實現具有批次不變性，我們需要確保其中的每一個內核都滿足這個要求。

好消息是，我們可以安全地假設所有逐點操作（pointwise operation）都是批次不變的。

雖然這在 PyTorch 的所有內核中確實成立，但從理論上講，這并非絕對保證。舉個例子，在某些 CPU 內核實現中，可能會對數組的部分區域使用向量化指令（vectorized intrinsics），而對其他區域使用標量指令，這兩種指令的數值結果未必總能做到逐位（bitwise）完全一致。

因此，我們真正需要關注的是三類涉及歸約（reduction）的操作：RMSNorm、矩陣乘法（matrix multiplication）和注意力機制（attention）。

雖然與并行性相關的歸約問題不在本次討論范圍之內，但相同的處理原則同樣適用。

順便分享一個實用的小知識：在 Blackwell 架構以及搭配 CUDA 12.8+ 的 Hopper 上，NVLink-Sharp 的交換機內歸約（in-switch reductions）是確定性的。和很多技術細節一樣，這個信息藏在 NCCL 的 GitHub issue 討論中。

巧合的是，這三類操作的實現難度也恰好呈遞增趨勢。想要在保持合理性能的前提下實現批次不變性，每一個操作都需要我們投入額外的思考和設計。

那么，讓我們先從 RMSNorm 開始講起。

批次不變的 RMSNorm

Figure 5：理想情況下，我們希望在并行化策略中避免核心間的通信。實現這一目標的一種方法是將一個批次元素分配給一個核心，從而保證每次歸約都完全在一個核心內完成。這就是所謂的“數據并行”策略，因為我們只是沿著一個不需要通信的維度進行并行化。在這個例子中，我們有四行和四個核心，充分利用了我們的核心。

對批量不變性的要求是：無論內核的批量大小是多少，每個元素的歸約順序（reduction order）都必須保持一致。

需要注意的是，這并不意味著我們必須始終使用完全相同的歸約策略。舉個例子，當要歸約的元素數量發生變化時，即使歸約策略相應調整，我們依然可以維持批量不變性。Quack 博客文章中提供了一些精彩的示例，清晰地展示了各種歸約策略的層級關系 —— 從線程歸約、warp 歸約到 block 歸約、cluster 歸約。

換句話說，只有當批量大小的變化直接影響到歸約策略的選擇時，批量不變性才會被打破。

現在讓我們來看看 RMSNorm 的標準并行策略。

眾所周知，并行算法的性能優勢主要來源于最大化減少核心之間的通信開銷。為了便于理解，這里我們可以將"核心（cores）"簡單理解為 SM（Streaming Multiprocessors，流式多處理器）。在實際應用中，有一個關鍵特性需要把握：我們啟動的線程塊（threadblocks）數量通常要超過可用的 SM 數量?；谶@個前提，一個自然的策略就是為每個批量元素分配一個獨立的核心，正如上圖所展示的那樣。

這種分配策略的美妙之處在于：增加批量大小并不會改變我們的歸約策略。如果批量大小為 200 時已經能為內核提供充足的并行性，那么當批量大小增加到 2000 時，這種并行性只會變得更加充沛。

Figure 6：數據并行 RMSNorm（適用于更大的批次）將數據并行策略擴展到更大的批次相當直接 —— 與其讓每個核心只處理一行，不如讓每個核心依次處理不同的行。這樣可以保持批次不變性，因為每個批次元素的歸約策略仍然是相同的。

另一方面，減小批次大小可能會帶來挑戰。

因為我們為每個批次元素分配一個核心，當批次大小逐漸減小時，勢必會遇到核心數量超過批次元素數量的情況，這時部分核心就會處于空閑狀態。

面對這種情況，經驗豐富的內核工程師通常會采用前一節介紹的某種解決方案（比如原子加法或分裂歸約）來充分利用這些空閑核心，以保持良好的并行效率。然而，這樣做的代價是改變了原有的歸約策略，使得內核無法繼續維持批次不變性這一重要特性。

Figure 7：分裂歸約 RMSNorm。如果我們的推理批次較小，那么數據并行策略可能無法提供足夠的并行度來充分利用所有核心。在這種情況下，一個更高效的做法是將單個歸約任務分解到多個核心上并行處理，從而讓 GPU 資源得到更好的利用。不過，這種方法會帶來一個權衡：我們可能會失去批次不變性。這是因為當多個核心并行處理同一個歸約任務時，元素的處理順序可能不再保持一致，進而影響最終結果的確定性。

最簡單的解決辦法就是干脆完全忽略這些情況。

這其實也算合理 —— 畢竟小批次本身執行速度就很快，即便性能有所下降，通常也不至于造成嚴重問題。

如果確實需要優化這類場景，一個可行的思路是始終采用某種歸約策略，確保即便在很小的批次下也能提供充足的并行度。雖然這種策略在大批次時會產生一定程度的過度并行，但它能讓我們在各種批次規模下都獲得相對穩定（盡管不是最優）的性能表現。

批次不變的矩陣乘法

Figure 8：數據并行矩陣乘法：與 RMSNorm 類似，矩陣乘法的標準并行策略采用"數據并行"的思路，即將整個歸約過程集中在單個核心內完成。最直接的實現方式是：將輸出張量劃分為二維小塊（2D tiles），每個小塊分配給不同的核心處理。各核心負責計算所分配小塊內的點積運算，歸約操作仍在核心內部完成。然而，矩陣乘法相比 RMSNorm 面臨著更多約束條件——算術強度（arithmetic intensity）的要求以及對 Tensor Cores 的高效利用等因素，都對內核設計提出了更嚴格的要求。正是這些約束，使得高效的矩陣乘法內核必須采用二維小塊的切分策略，而非按單個輸出元素進行切分的簡單方式。

從本質上看，你可以把矩陣乘法理解為：先執行逐點運算 (pointwise operation)，再執行一次歸約 (reduction)。

基于這個理解，當我們通過把輸出切分成小塊 (tiles) 來并行化矩陣乘法時，實際上得到了類似"數據并行"的內核策略——每次歸約運算仍然保持在單個核心內部完成，這樣就避免了跨核心的通信開銷。

然而，就像我們在 RMSNorm 中遇到的問題一樣，矩陣乘法的"批次"維度（M 和 N）有時會顯得太小，迫使我們不得不沿著歸約維度 (K) 進行切分。雖然矩陣乘法擁有兩個"批次"維度的優勢，但要想充分發揮 Tensor Cores 的威力，我們仍然需要在每個核心上分配足夠多的"工作量"。

舉個具體例子：假設你有一個形狀為 [1024, K] × [K, 1024] 的矩陣乘法，如果采用標準的二維 tile 大小 [128, 128]，那么數據并行策略最多只能將這次計算分配給 64 個核心。對于現代 GPU 來說，這樣的并行度顯然還不足以讓硬件"吃飽"。

在矩陣乘法領域，沿著歸約維度進行切分的策略有一個專門的名字：Split-K Matmul。不過需要注意的是，和 RMSNorm 的情況類似，采用 Split-K 策略同樣會破壞批次不變性 (batch invariance)——這是我們在選擇并行策略時需要權衡的一個重要考量。

Figure 9：Split-K 矩陣乘法。如果我們的批次維度非常小，就可能沒有足夠的并行度，這時就需要使用 Split-K 矩陣乘法。在這個例子中，我們把每一次歸約拆分給兩個核心來執行，這兩個核心會各自累加，最后再把結果合并。這樣一來，雖然每次歸約被分到兩個核心上，但我們仍然能夠利用八個核心。

矩陣乘法還有一個額外的復雜點 —— Tensor Core 指令。

在歸約運算里，我們可以一次只處理一行；但在高效的矩陣乘法內核中，必須一次處理整個“tile”（小矩陣塊）。

每條 Tensor Core 指令（比如

wgmma.mma_async.sync.aligned.m64n128k16）在內部可能會有不同的歸約順序。選擇不同的 Tensor Core 指令的一個原因，可能就是 批次大小非常小。

舉個例子：如果我們使用的 Tensor Core PTX 指令需要處理長度為 256 的 tile，但實際 batch size 只有 32，那幾乎 90% 的算力都被浪費掉了！

當 batch size = 1 時，最快的內核通常甚至完全不會使用 Tensor Cores。

Figure10：填充 (Padded) Tensor Core 指令。如果批次大小太小，就可能出現一種情況：我們甚至無法在輸出中放下一個完整的二維 tile。在這種情況下，最有效的做法是切換到更小的 Tensor Cores 指令，或者干脆完全不用 Tensor Cores！然而，這兩種選擇都會使我們的內核無法保持批次不變性 (batch-invariance)。

因此，確保矩陣乘法批次不變性最簡單的方法，就是只編譯一種內核配置，并把它應用到所有形狀上。

雖然這樣會損失一部分性能，但在大語言模型（LLM）的推理過程中，這通常并不是災難性的。尤其是，Split-K 最需要的情況是 M 和 N 都很小，而幸運的是，在我們的場景中，N（也就是模型維度）通常都相當大！

Figure 11：盡管實現了批次不變性，我們相比 cuBLAS 只損失了大約 20% 的性能。需要注意的是，這里用的也不是經過優化的 Triton 內核（例如沒有使用 Thinking MachinesA）。不過，一些性能表現的模式還是能說明問題，揭示出批次不變性要求導致性能損失的地方。首先，請注意在 非常小的批次大小 下，由于指令過大以及并行度不足，我們損失了相當多的性能。其次，當批次大小逐漸增加時，會出現一種 “拼圖”式的性能波動模式，這是由量化效應（tile 和 wave 兩方面）造成的，而這些效應通?？梢酝ㄟ^改變 tile 大小來緩解。

批次不變的注意力機制

Figure 12：Flash Attention 2 策略。我們在 Q 維度 上進行并行，同時在 K/V 維度 上做歸約。這意味著整個歸約過程都能保持在單個核心內完成，使其成為另一種數據并行策略。

在矩陣乘法實現了批次不變性之后，注意力機制又引入了兩個額外的復雜點 —— 這也合乎邏輯，因為它本身包含了兩次矩陣乘法。

1. 與 RMSNorm 和矩陣乘法只在 特征維度 上做歸約不同，注意力機制需要同時在 特征維度 和 序列維度 上做歸約。

2. 因此，注意力機制必須應對多種推理優化方式，這些方式會影響序列的處理方式（例如 分塊預填充 chunked prefill、前綴緩存 prefix caching 等）。

因此，要在大語言模型 (LLM) 的推理中實現確定性，我們的數值必須對兩方面保持不變：既要與一次同時處理多少請求無關，也要與推理引擎如何切分每個請求無關。

接下來，我們先來看注意力機制的標準并行化策略，這個策略最早由 Flash Attention 2 提出。與 RMSNorm 和矩陣乘法類似，默認策略是一種“數據并行”策略。由于我們在 Key/Value 張量 上進行歸約，所以數據并行策略只能在 Query 張量 上實現并行。

舉個例子：根據推理引擎的不同選擇，一個序列可能被分成多個部分來處理（比如分塊預填充），也可能一次性處理全部（如果預填充沒有被拆分）。為了實現“批次不變性”，關鍵在于：某個 token 的歸約順序不能依賴于該序列中同時處理了多少其他 token。

如果你把 KV 緩存里的 K/V 值與當前處理的 token 的 K/V 值分開做歸約（就像 vLLM 的 Triton 注意力內核那樣），那就無法做到批次不變性。

舉例來說：當處理序列中的第 1000 個 Query token 時，歸約順序必須完全一致——不論 KV 緩存中有 0 個 token（預填充階段）還是有 999 個 token（解碼階段）。

Figure 13：帶 KV 緩存的 Flash Attention將 KV 緩存與當前 KV 值分開處理之所以會破壞批次不變性，這一點有些微妙，主要與 “邊界條件 (boundary conditions)” 有關。

舉個例子：假設塊大小是 32，而我們當前在 KV 緩存中有 80 個元素。接著我們又計算了額外 48 個不在緩存中的元素。

在這種情況下：

計算 P cache 需要 3 個塊（兩個完整塊 + 一個帶掩碼的塊）。

計算 P 需要 2 個塊（一個完整塊 + 一個帶掩碼的塊）。

這樣一來，總共需要 5 個塊來完成歸約，但實際上我們只有 128 個元素，即 4 個塊。這就必然會改變歸約的順序。

對比：如果我們一開始沒有 KV 緩存，而是一次性處理 128 個元素，那么只需要 4 個塊來完成歸約。因此，為了保證注意力機制的 批次不變性 (batch invariance)，我們必須確保這兩種情況下的數值結果完全一致。

為了解決這個問題，我們可以在進入注意力內核之前，先更新 KV 緩存 和 頁表 (page table)，以確保無論當前處理多少個 token，我們的 Key 和 Value 都始終以一致的方式排列。

結合這一點（以及上一節提到的所有細節，比如保持一致的 tile 大?。覀兙湍軌驅崿F一個 批次不變的注意力機制 (batch-invariant attention)！

然而，這里存在一個嚴重的問題。不同于矩陣乘法，在大語言模型 (LLM) 推理中遇到的注意力形狀往往確實需要 分裂歸約內核 (split-reduction kernel)，通常稱為 Split-KV 或 Flash Decoding。

原因是：如果我們不在歸約維度上并行，就只能在 批次維度 (batch dimension)、頭部維度 (head dimension) 和 Query 長度維度 (query length dimension) 上進行并行。

但在注意力的 解碼階段 (decode stage)，Query 長度非常小，所以除非批次很大，否則我們經常無法讓 GPU 得到充分利用。

遺憾的是，這種情況不像 RMSNorm 和矩陣乘法那樣可以輕松忽略。舉個例子：如果 KV 緩存非常長，即便我們只在處理一個請求，注意力內核也可能會耗費非常長的時間。

Figure 14：固定的 # Split-KV 策略（也叫 Flash Decode）。如果我們的 query 長度變得非常小（就像在解碼階段那樣），我們可能會遇到這樣一種情況：內核中幾乎沒有多少并行度可利用。在這些情況下，我們需要再次沿著歸約維度進行拆分 —— 這一次是 KV 維度。

典型的做法是先確定我們需要多少并行度，然后將 KV 維度均勻地劃分。比如，如果 KV 長度是 1000，而我們需要 4 個拆分，那么每個核心就會處理 250 個元素。

但不幸的是，這種方法也會破壞 批量不變性（batch invariance），因為我們精確的歸約策略取決于當前請求中要處理多少個序列的 query token。

此外，注意力中常用的 拆分歸約（split-reduction）策略 也會給批量不變性帶來挑戰。舉例來說，FlashInfer 的“平衡調度算法（balanced scheduling algorithm）”會選擇一種最大化的拆分規模，以便仍能讓 GPU 的所有核心飽和，因此這種歸約策略并不是“批量不變的（batch-invariant）”。不過，與 RMSNorm/矩陣乘法不同，僅僅選擇一個固定的拆分數量，而不管批量大小是多少，是不夠的。

為了實現批量不變性，我們必須采用“固定拆分大小（fixed split-size）”的策略。換句話說，我們不是固定拆分的數量，而是固定每個拆分的大小，最終得到的拆分數量則會隨之變化。這樣一來，我們就能保證無論處理多少 token，始終執行相同的歸約順序。

Figure 15：固定大小的 Split-KV 策略這種策略與前一種策略的唯一不同在于，我們的拆分現在是“固定大小”的。比如，如果 KV 長度是 1000，那么我們不會再把它均分成四段、每段長度 250，而是拆分成三段固定長度 256 的部分，以及一段長度 232 的部分。

這樣，我們就能夠保持 批量不變性（batch invariance），因為我們的歸約策略不再依賴于一次要處理多少個 query token！

實現

我們基于 vLLM，利用其 Flex Attention 后端以及 torch.Library，提供了一個確定性推理的演示。通過 torch.Library，我們能夠以一種非侵入的方式替換掉大部分相關的 PyTorch 算子。你可以在 thinking-machines-lab/batch-invariant-ops 中找到“批量不變（batch-invariant）”內核的庫，以及在 vLLM 中運行“確定性（deterministic）”模式的示例。

實驗

補全（completions）的不確定性有多大？

我們使用 Qwen/Qwen3-235B-A22B-Instruct-2507，在溫度設為 0 的情況下采樣 1000 次補全，提示詞為 “Tell me about Richard Feynman”（非思考模式），每次生成 1000 個 token。令人驚訝的是，我們一共得到了 80 種不同的補全，其中最常見的一種出現了 78 次。

觀察補全結果的分歧點，我們發現前 102 個 token 的補全完全一致！首次出現差異是在 第 103 個 token。所有補全都會生成序列 “Feynman was born on May 11, 1918, in”。但其中 992 個補全繼續生成了 “Queens, New York”，而 8 個補全生成了 “New York City”。

另一方面，當我們啟用 批量不變內核 時，所有 1000 個補全結果完全一致。這正是從數學角度來看我們對采樣器的預期，但如果不使用批量不變內核，我們無法實現這種確定性結果。

性能

在這里我們并沒有對批量不變內核進行大量性能優化。不過，我們仍然進行了一些實驗來驗證性能是否可用。

我們搭建了一個 API 服務器，使用一張 GPU 運行 Qwen-3-8B，并請求生成 1000 條序列，每條的輸出長度在 90 到 110 之間。

大部分的性能下降，來自于 vLLM 中 FlexAttention 的集成尚未經過深度優化。盡管如此，我們可以看到性能并沒有糟糕到不可接受的程度。

滿血版 On-Policy 強化學習（RL）

正如研究者指出的那樣，訓練與推理階段數值的不一致，實際上會把我們的 on-policy RL 變成 off-policy RL。

當然，如果連兩次完全相同的推理請求都無法得到逐位（bitwise）一致的結果，那么訓練和推理之間就更不可能實現逐位一致。而一旦實現了 確定性推理，我們也就能夠修改訓練棧，使得采樣和訓練之間的結果逐位一致，從而實現真正的 on-policy RL。

我們在 BigMath 上運行了一個 RLVR 設置實驗，RL 策略從 Qwen 2.5-VL instruct 8B 初始化，最大采樣長度為 4096。

如果我們在訓練中不做 off-policy 修正（即：重要性加權，importance weighting），獎勵會在訓練過程中途崩潰；而在訓練中加入 off-policy 修正項，訓練則能順利進行。但如果我們能夠讓采樣器和訓練器之間的結果逐位一致，那么我們就是完全的 on-policy（即 KL 散度為 0），同樣也能順利完成訓練。

我們還可以繪制 采樣器與訓練器之間 logprob 的 KL 散度曲線，結果顯示三種運行方式的行為有明顯不同：

使用 重要性加權 時，KL 散度大約維持在 0.001，偶爾會有尖峰；

不使用 重要性加權 時，KL 散度最終會出現尖峰，并與獎勵崩潰的時間點吻合；

而在運行 “真正的 On-Policy RL” 時，KL 散度始終保持在 0，表明訓練策略與采樣策略之間不存在任何差異。

Figure 16：請注意，在沒有使用 重要性加權 的運行中，大約在 第 318 步時出現了顯著的損失尖峰，同時伴隨著 logprob 的 KL 散度的尖峰。與此同時，不論是使用 off-policy 修正，還是運行 “真正的 On-Policy”，都能讓 RL 順利地繼續進行。圖中那條代表 “True On-Policy” 的藍色線并不是 bug —— 它只是一直平坦地保持在 0。

結論

現代軟件系統往往由多層抽象構成。在機器學習中，當我們遇到不確定性和一些微妙的數值差異時，人們往往會選擇視而不見。

畢竟，我們的系統本來就是「概率性的」，再多一點不確定性又有何妨呢？單元測試掛掉時，把 atol/rtol 調大點又有什么問題？訓練器和采樣器之間的對數概率差異，應該不算是真正的 bug 吧？

但我們拒絕這種消極心態！只要稍微多付出一些努力，我們就能理解不確定性的真正根源，甚至徹底解決它們。

我們希望這篇博文能為社區提供一套可靠的思路，幫助大家在推理系統中更好地應對不確定性，并激勵更多人深入理解自己的系統。