文｜四木洋子

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

這家由 OpenAI 前 CTO 米拉·穆拉蒂（Mira Murati）創(chuàng)立的新興 AI 公司，成立至今還不到半年，卻已經(jīng)匯聚了一眾來(lái)自 OpenAI、Meta 等頂級(jí) AI 機(jī)構(gòu)的精英人才。

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

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

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

這篇深度技術(shù)文章直指一個(gè)困擾業(yè)界已久的核心難題：大語(yǔ)言模型前向推理過(guò)程中的不確定性問(wèn)題。

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

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

經(jīng)過(guò)實(shí)驗(yàn)驗(yàn)證，Thinking Machines 提出的方案體現(xiàn)出雙重突破：

• 一方面，它讓大語(yǔ)言模型能夠提供高度可復(fù)現(xiàn)的前向推理結(jié)果，解決了長(zhǎng)期困擾研究者的一致性問(wèn)題；

• 另一方面，該方案從根本上改善了強(qiáng)化學(xué)習(xí)的訓(xùn)練動(dòng)態(tài)，成功實(shí)現(xiàn)了基于大語(yǔ)言模型的真正同策略探索與利用平衡——這一直是這個(gè)領(lǐng)域亟待攻克的核心難題。

這些突破性成果印證了 Thinking Machines 團(tuán)隊(duì)在 AI 前沿研究中所具備的深厚洞察力和卓越的系統(tǒng)性研究能力。

我們把本文的閱讀心得，分為內(nèi)容總結(jié)、思考點(diǎn)評(píng)和原文翻譯三部分，供大家按需了解：

內(nèi)容總結(jié)

為什么問(wèn)大模型相同的問(wèn)題，會(huì)得到不同的答案？

這個(gè)問(wèn)題雖然有時(shí)會(huì)有點(diǎn)惱人，但并不難理解——因?yàn)檎Z(yǔ)言模型中的結(jié)果涉及“采樣”，這本身就是一個(gè)"隨機(jī)選詞&輸出"的過(guò)程。

如果說(shuō)隨機(jī)采樣約等于抽簽，那么采樣溫度（temperature）則直接決定了抽簽的權(quán)重分布。

但現(xiàn)在的問(wèn)題在于，即便將采樣溫度調(diào)整到 0，我們?nèi)匀徊荒鼙ＷC模型會(huì)輸出確定的答案。

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

這就有點(diǎn)令人費(fèi)解了。

Thinking Machines 的這篇文章從一個(gè)我們都不陌生的現(xiàn)象切入——同樣的問(wèn)題，大語(yǔ)言模型有時(shí)會(huì)給出不同的答案。這看似隨機(jī)的表現(xiàn)，實(shí)際上揭示了大模型推理過(guò)程中一個(gè)深層次的技術(shù)問(wèn)題：推理的"不確定性"（nondeterminism）。

文章的核心發(fā)現(xiàn)令人意外：這種不確定性的根源，竟然來(lái)自于推理批次不變性（batch invariance）的缺失。

什么意思呢？讓我們從 GPU 的工作機(jī)制說(shuō)起。為了榨干硬件的每一分算力，GPU 會(huì)根據(jù)當(dāng)前的批次大�。╞atch size）——也就是服務(wù)器同時(shí)處理的請(qǐng)求數(shù)量——?jiǎng)討B(tài)調(diào)整自己的并行策略和計(jì)算單元的調(diào)度方式。批次大了，就用這種策略；批次小了，就換那種策略。

問(wèn)題就出在這里。

由于浮點(diǎn)數(shù)運(yùn)算本身的特性——它并不嚴(yán)格滿足數(shù)學(xué)上的加法結(jié)合律，不同的計(jì)算順序會(huì)產(chǎn)生微小的數(shù)值偏差。這些看似微不足道的差異，在 Transformer 層層疊疊的前向傳播過(guò)程中被逐步放大，最終導(dǎo)致了一個(gè)令人困擾的結(jié)果：相同的輸入，在不同的批次環(huán)境下，竟然會(huì)產(chǎn)生不同的輸出。

簡(jiǎn)而言之，GPU 為了優(yōu)化效率而采用的靈活調(diào)度策略，意外地成為了大模型推理不確定性的幕后推手。一個(gè)追求極致性能的設(shè)計(jì)，卻在無(wú)意中為模型的"善變"埋下了伏筆。

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

在 GPU 計(jì)算中，每個(gè)操作都通過(guò)專門(mén)的計(jì)算內(nèi)核（kernel）來(lái)執(zhí)行，這就要求所有內(nèi)核都必須具備批次不變性——即無(wú)論輸入批次如何變化，計(jì)算邏輯都應(yīng)保持一致。

然而，在 Transformer 模型的實(shí)際運(yùn)行中，某些操作天然地會(huì)因批次規(guī)模的改變而產(chǎn)生計(jì)算行為上的差異，這類操作主要集中在各種歸約計(jì)算（reduction）環(huán)節(jié)。認(rèn)識(shí)到這一挑戰(zhàn)，Thinking Machines針對(duì)性地對(duì)三個(gè)核心的歸約操作 —— RMSNorm、矩陣乘法以及注意力機(jī)制 —— 進(jìn)行了專門(mén)的優(yōu)化設(shè)計(jì)。

• RMSNorm：讓每一個(gè)樣本都單獨(dú)在一個(gè)內(nèi)核上完成 RMSNorm 的計(jì)算。這樣一來(lái)，無(wú)論批次里有多少樣本，都不需要跨樣本通信，完全在自己的內(nèi)核里完成，得到的結(jié)果與單獨(dú)算該樣本時(shí)是一致的 。

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

• 注意力機(jī)制：注意力計(jì)算涉及更復(fù)雜的步驟（比如 KV 的點(diǎn)積 和 Softmax 歸一化等），而且還牽涉序列長(zhǎng)度（token 數(shù)量）維度的變化。解決方案是固定歸約塊的大小。簡(jiǎn)單理解：原先注意力在解碼時(shí)為了適應(yīng)不同 token 數(shù)，可能動(dòng)態(tài)地選擇需要多少塊來(lái)并行計(jì)算；現(xiàn)在改為預(yù)先固定每一塊的大小，這樣無(wú)論總共多少 token，就算需要的塊數(shù)變多，每個(gè)塊內(nèi)部的計(jì)算順序和規(guī)模都是一致的。

通過(guò)這三種歸約計(jì)算方法的優(yōu)化，Thinking Machines 實(shí)現(xiàn)了 Transformer 模型前向計(jì)算的批次確定性。

有了確定性之后，除了能夠?qū)崿F(xiàn)可復(fù)現(xiàn)的模型推理結(jié)果，還能顯著改善 RL 訓(xùn)練效果。

RL 的基本原理是通過(guò)獎(jiǎng)勵(lì)機(jī)制來(lái)強(qiáng)化模型的正確輸出。然而，如果每次答案都有細(xì)微差別，訓(xùn)練數(shù)據(jù)就會(huì)變得很「嘈雜」，而更一致的響應(yīng)能讓整個(gè) RL 訓(xùn)練過(guò)程變得「絲滑」很多。

在 RL 領(lǐng)域，On-policy（同策略）和 Off-policy（異策略）是兩種核心訓(xùn)練范式，它們的區(qū)別在于學(xué)習(xí)時(shí)使用的"行為策略"與"目標(biāo)策略"是否為同一個(gè)策略。目前，很多研究工作都在強(qiáng)調(diào)同策略（On-policy）優(yōu)化的重要性。

但這里出現(xiàn)了一個(gè)有趣的問(wèn)題：GPU 的不確定性既可能影響采樣時(shí)的策略表現(xiàn)（生成不同的數(shù)據(jù)），也可能影響訓(xùn)練時(shí)的策略表現(xiàn)（產(chǎn)生不同的梯度計(jì)算結(jié)果）。

這就導(dǎo)致了一個(gè)核心矛盾——明明是同一個(gè)策略，但在采樣和訓(xùn)練階段跑出來(lái)的結(jié)果卻不一致。結(jié)果，原本設(shè)計(jì)為同策略的訓(xùn)練被迫變成了異策略（Off-policy）訓(xùn)練，這種現(xiàn)象可以稱為"偽同策略 RL"（Faked On-Policy RL）。

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

思考點(diǎn)評(píng)

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

答案很簡(jiǎn)單：這個(gè)問(wèn)題實(shí)在太重要了，而能夠針對(duì)這一問(wèn)題提出系統(tǒng)性解決方案的團(tuán)隊(duì)卻寥寥無(wú)幾。

Faked On-policy RL for LLM，這是當(dāng)前世界范圍內(nèi)普遍存在的痛點(diǎn)問(wèn)題。

令人意外的是，即便是那些在前沿的 AI 頂會(huì)上發(fā)表、明確宣稱"我們采用同策略強(qiáng)化學(xué)習(xí)"的研究工作，實(shí)際情況往往是：研究者們?cè)谒惴▽用婢拇蚰�，但承載算法實(shí)現(xiàn)的底層框架和算子等系統(tǒng)組件，卻在無(wú)形中背離了 On-Policy RL 的初衷和愿景。

而這一根本性問(wèn)題，恰恰是目前開(kāi)源社區(qū)（如 veRL、MS-Swift 等 RL 訓(xùn)練框架的提供方）難以攻克的 ——原因在于維護(hù)人員往往術(shù)業(yè)有專攻：擅長(zhǎng)框架設(shè)計(jì)的未必精通算子開(kāi)發(fā)，而算子專家又不一定能駕馭整體框架架構(gòu)。這種技能壁壘的存在，使得端到端的解決方案變得格外珍貴。

目前我們看到的相似研究風(fēng)格的博客包括但不限于以下幾個(gè)：

• 如 Thinking Machines 這篇文章中所引用的，微軟高劍峰組在前兩個(gè)月發(fā)現(xiàn)了目前 RL for LLM 訓(xùn)練中普遍存在的這一問(wèn)題。雖然他們?cè)凇禮our Efficient RL Framework Secretly Brings You Off-Policy RL Training》中提出了緩解方案，但坦率地說(shuō)，這種方式仍顯得不夠優(yōu)雅且不夠徹底。

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

• 其他更多的博客文章則側(cè)重于單點(diǎn)優(yōu)化和改進(jìn)，如系統(tǒng)（算子）層面 Tri Dao 的 FlashAttention 系列，以及純算法層面的典型如 John Schulman 的《Approximating KL Divergence》。

而 Thinking Machines 的做法，實(shí)際上是拿軟件（算法）—— 系統(tǒng)—— 硬件這三者串聯(lián)所產(chǎn)生的推理不確定性問(wèn)題，進(jìn)行了一次由表及里的深度剖析，系統(tǒng)性地梳理推理過(guò)程中不確定性的根本來(lái)源。

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

從「數(shù)值計(jì)算」的底層細(xì)節(jié)，到「系統(tǒng)（內(nèi)核）原理」的中層機(jī)制，再到算法層面的優(yōu)化改進(jìn)，層層遞進(jìn)，向外界展示了那些在一般研究人員眼中"平平無(wú)奇"的技術(shù)細(xì)節(jié)，是如何最終影響到終端模型訓(xùn)練效果的。

更重要的是，在解決了高劍峰等人發(fā)現(xiàn)的「現(xiàn)在的 RL for LLM 都是偽同策略 RL」這一關(guān)鍵問(wèn)題之后，才有可能在此基礎(chǔ)上構(gòu)建出真正的滿血版 RL-ed LLMs（經(jīng)過(guò)強(qiáng)化學(xué)習(xí)訓(xùn)練的語(yǔ)言模型）。

而這樣的模型，又將為其他更宏大的目標(biāo)提供堅(jiān)實(shí)支撐——比如構(gòu)建多模態(tài)通用 Agent，或是提供面向客戶的優(yōu)質(zhì) RL 解決方案等。

在文章的結(jié)尾部分，Thinking Machines 的研究人員特別強(qiáng)調(diào)：

不能忽視推理過(guò)程中的不確定性，只有抓住并深入分析這些不確定性，找到切實(shí)可行的解決方案，才能最終實(shí)現(xiàn)真正的同策略強(qiáng)化學(xué)習(xí)，進(jìn)而把 RL for LLM 這件事做得更加扎實(shí)。

對(duì)其他投身大模型研究的人來(lái)說(shuō)，Thinking Machines 這次的工作傳遞出一個(gè)清晰的信號(hào)：

想要繼續(xù)提升模型表現(xiàn)，既需要自頂向下的宏觀思考，也需要自底向上的細(xì)致打磨，端到端解決問(wèn)題的系統(tǒng)思維和扎實(shí)的動(dòng)手能力缺一不可——「宏觀規(guī)劃 + 微觀操作」的組合拳必須打得漂亮。

以下為原文翻譯：

重復(fù)性是一切科學(xué)進(jìn)步的基石。然而，從大語(yǔ)言模型中獲得可重復(fù)的結(jié)果卻異常困難。

例如，你可能會(huì)發(fā)現(xiàn)，多次向 ChatGPT 提出相同的問(wèn)題會(huì)得到不同的結(jié)果。這本身并不令人驚訝，因?yàn)閺恼Z(yǔ)言模型中獲取結(jié)果涉及“采樣”，這是一個(gè)將語(yǔ)言模型輸出轉(zhuǎn)換為概率分布并概率性地選擇一個(gè)標(biāo)記的過(guò)程。

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

但為什么 LLM 推理引擎不是100%確定性的呢？一個(gè)常見(jiàn)的假設(shè)是，浮點(diǎn)非關(guān)聯(lián)性和并發(fā)執(zhí)行的某種組合導(dǎo)致了基于哪個(gè)并發(fā)核心首先完成的不確定性。我們將此稱為 LLM 推理不確定性的“并發(fā) + 浮點(diǎn)”假設(shè)。例如，最近的一篇 arXiv 文章寫(xiě)道：

“GPU 中的浮點(diǎn)運(yùn)算表現(xiàn)出非關(guān)聯(lián)性，這意味著 (a + b) + c ≠ a + (b + c)，這是由于有限精度和舍入誤差造成的。此屬性直接影響 Transformer 架構(gòu)中注意力分?jǐn)?shù)和詞表輸出分?jǐn)?shù)（logits）的計(jì)算，其中跨多個(gè)線程的并行操作可能會(huì)根據(jù)執(zhí)行順序產(chǎn)生不同的結(jié)果�！�

你還可以在其他地方找到“并發(fā) + 浮點(diǎn)”假設(shè)的重復(fù)，例如：“存在速度權(quán)衡，為了使端點(diǎn)快速，使用了 GPU，它進(jìn)行并行不確定性計(jì)算。任何現(xiàn)代 GPU 神經(jīng)網(wǎng)絡(luò)計(jì)算都將受制于這些”，或“因?yàn)?GPU 是高度并行化的，所以每次執(zhí)行加法或乘法的順序可能不同，這可能會(huì)導(dǎo)致輸出中的微小差異”。

雖然這個(gè)假設(shè)并非完全錯(cuò)誤，但它并未揭示全貌。例如，即使在 GPU 上，反復(fù)對(duì)相同的數(shù)據(jù)執(zhí)行相同的矩陣乘法運(yùn)算，每次得到的結(jié)果在二進(jìn)制層面都是完全相同的。我們確實(shí)在用浮點(diǎn)數(shù)運(yùn)算，而且 GPU 也確實(shí)有很多并發(fā)操作。那為什么在這個(gè)測(cè)試?yán)餂](méi)有看到不確定性呢？

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

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

然而，語(yǔ)言模型正向傳遞中使用的所有內(nèi)核都是確定性的。

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

盡管如此，從使用推理服務(wù)器的任何人的角度來(lái)看，結(jié)果都是不確定性的。

在這篇文章中，我們將解釋為什么“并發(fā) + 浮點(diǎn)數(shù)”假設(shè)未能達(dá)到目標(biāo)，揭示導(dǎo)致 LLM 推理不確定性背后的真正原因，并解釋如何克服不確定性并在 LLM 推理中獲得真正可重現(xiàn)的結(jié)果。

最初的“原罪”：浮點(diǎn)數(shù)的非結(jié)合性

在討論不確定性之前，有必要解釋一下為什么會(huì)存在數(shù)值差異。畢竟，我們通常認(rèn)為機(jī)器學(xué)習(xí)模型是遵循交換律或結(jié)合律等結(jié)構(gòu)規(guī)則的數(shù)學(xué)函數(shù)。我們的機(jī)器學(xué)習(xí)庫(kù)難道不應(yīng)該為我們提供一個(gè)“數(shù)學(xué)上正確”的結(jié)果嗎？

罪魁禍?zhǔn)拙褪歉↑c(diǎn)數(shù)的非結(jié)合性。也就是說(shuō)，對(duì)于浮點(diǎn)數(shù)：

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

具有諷刺意味的是，正是這種對(duì)結(jié)合律的破壞使得浮點(diǎn)數(shù)變得有用。

浮點(diǎn)數(shù)之所以有用，是因?yàn)樗鼈冊(cè)试S“動(dòng)態(tài)”的精度級(jí)別。

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

例如，對(duì)于數(shù)值 3450，我們可以精確地表示為 3.45 * 10^3。我們也可以表示小得多的值，如 0.486，表示為 4.86 * 10^-1。通過(guò)這種方式，浮點(diǎn)數(shù)允許我們表示非常小和非常大的值。在科學(xué)領(lǐng)域，我們可能會(huì)說(shuō)浮點(diǎn)數(shù)允許我們保持恒定數(shù)量的“有效數(shù)字”。

如果你將兩個(gè)具有相同指數(shù)的浮點(diǎn)數(shù)相加，它看起來(lái)類似于整數(shù)加法。例如，123 (1.23 * 10^2) + 456 (4.56 * 10^2) 的結(jié)果是 579 (5.79 * 10^2)。

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

Figure 1：我們需要 3 位精度來(lái)表示 1230，也需要 3 位精度來(lái)表示 23.4。然而，將這兩個(gè)數(shù)相加會(huì)得到一個(gè)需要 5 位精度才能表示的數(shù)（1253.4）。我們的浮點(diǎn)格式必須丟棄末尾的 34。從某種意義上說(shuō)，我們?cè)谙嗉又�，�?shí)際上已經(jīng)將原始的 23.4 四舍五入到 20.0 了。

然而，在這一點(diǎn)上，我們已經(jīng)丟失了信息。請(qǐng)注意，每當(dāng)我們相加兩個(gè)具有不同“尺度”（即不同指數(shù)）的浮點(diǎn)數(shù)時(shí)，都可能發(fā)生這種情況。而相加具有不同指數(shù)的浮點(diǎn)數(shù)是經(jīng)常發(fā)生的。事實(shí)上，如果能保證從不需要不同的指數(shù)，我們就可以直接使用整數(shù)了！

換句話說(shuō)，每次我們以不同的順序?qū)⒏↑c(diǎn)數(shù)相加時(shí)，都可能得到一個(gè)完全不同的結(jié)果。舉一個(gè)極端的例子，對(duì)這個(gè)數(shù)組求和，根據(jù)相加順序的不同，可能會(huì)有 102 種不同的結(jié)果。

盡管這是導(dǎo)致輸出不一致的根本原因，但它并沒(méi)有直接回答不確定性從何而來(lái)。它沒(méi)有幫助我們理解為什么浮點(diǎn)值會(huì)以不同的順序相加、何時(shí)發(fā)生這種情況以及如何避免。

答案在于內(nèi)核（kernels）是如何實(shí)現(xiàn)的。

為什么內(nèi)核不總是以相同的順序相加數(shù)字？

既然運(yùn)算順序會(huì)影響結(jié)果，那么為什么 GPU 內(nèi)核在執(zhí)行加法時(shí)，順序不是固定的呢？

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

令人困惑的是，盡管這（并發(fā)與浮點(diǎn)非關(guān)聯(lián)性）可能會(huì)導(dǎo)致內(nèi)核出現(xiàn)不確定性，但在 LLM 推理的不確定性中，并發(fā)（以及原子加法 atomic adds）實(shí)際上完全沒(méi)有參與其中！ 要解釋真正的罪魁禍?zhǔn)资鞘裁�，我們首先需要理解為什么現(xiàn)代 GPU 內(nèi)核很少需要使用原子加法。

*atomic add = 一種讓很多線程同時(shí)往同一個(gè)變量里“加數(shù)”而不丟結(jié)果的機(jī)制，但加法順序不固定，所以浮點(diǎn)數(shù)計(jì)算可能有細(xì)微差異。

什么時(shí)候需要用到原子加法？

通常情況下，GPU 會(huì)在許多“核心”（即 SM，Streaming Multiprocessors）上并發(fā)地啟動(dòng)一個(gè)程序。由于這些核心之間沒(méi)有天然的同步機(jī)制，如果它們需要彼此通信，就會(huì)產(chǎn)生挑戰(zhàn)。比如說(shuō)，如果所有核心都必須把結(jié)果累加到同一個(gè)元素上，就可以使用“atomic add”（有時(shí)也叫“fetch-and-add”）。

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

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

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

這通常就是人們所說(shuō)的"不確定性"——用完全相同的輸入執(zhí)行同一個(gè)內(nèi)核兩次，卻得到截然不同的結(jié)果。

這種現(xiàn)象被稱為運(yùn)行間不確定性（run-to-run nondeterminism）：你用相同的依賴環(huán)境運(yùn)行同一個(gè) Python 腳本兩次，結(jié)果卻不一樣。

雖然并發(fā)的 atomic add 操作確實(shí)會(huì)讓內(nèi)核變得不確定，但對(duì)于絕大多數(shù)內(nèi)核來(lái)說(shuō)，atomic add 其實(shí)并非必需品。 事實(shí)上，在 LLM 的典型前向傳播過(guò)程中，通常連一個(gè) atomic add 操作都不會(huì)出現(xiàn)。

這個(gè)現(xiàn)象可能會(huì)讓人感到驚訝——畢竟在對(duì)歸約（reduction）操作進(jìn)行并行化時(shí)，atomic add 確實(shí)能帶來(lái)性能收益。 但 atomic add 最終之所以并非必需，主要有兩個(gè)原因：

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

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

接下來(lái)，為了將剩余的這 5 個(gè)結(jié)果合并起來(lái)，我們有兩種選擇：

執(zhí)行一次單獨(dú)的“清理歸約（clean-up reduction）” —— 雖然這一步不再并行，但由于處理的元素?cái)?shù)量很少，計(jì)算代價(jià)相當(dāng)?shù)土?/p>

或者采用信號(hào)量（semaphore）機(jī)制，它能夠確保每個(gè)并發(fā)的線程塊按照確定的順序進(jìn)行累加操作。

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

由于以上兩個(gè)因素，對(duì)于絕大多數(shù)神經(jīng)網(wǎng)絡(luò)運(yùn)算來(lái)說(shuō)，避免使用 atomic add 帶來(lái)的性能損失幾乎可以忽略不計(jì)。

當(dāng)然，仍有一些常見(jiàn)操作在避免 atomic add 時(shí)會(huì)造成顯著的性能損失，比如 PyTorch 中的 scatter_add 操作（a[b] += c）。不過(guò)在 LLM 領(lǐng)域，唯一常用到這類操作的場(chǎng)景就是 FlashAttention 的反向傳播了。

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

標(biāo)準(zhǔn)的 Triton 實(shí)現(xiàn)會(huì)在反向傳播中進(jìn)行額外的重計(jì)算，以此來(lái)規(guī)避 atomic add 的使用，但代價(jià)是增加了 40% 的 FLOPs！這是一個(gè)典型的以計(jì)算換確定性的權(quán)衡。

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

從推理服務(wù)器的角度來(lái)看，這意味著什么呢？它是完全確定性的——只要用戶的請(qǐng)求完全相同，服務(wù)器總是會(huì)返回相同的輸出結(jié)果。

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

在這里，給定完全相同的輸入（即推理服務(wù)器正在處理的完全相同的請(qǐng)求），前向傳播總是會(huì)產(chǎn)生完全相同的輸出。

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

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

由于每個(gè)獨(dú)立請(qǐng)求都無(wú)法預(yù)知其他并行請(qǐng)求的內(nèi)容，從單個(gè)請(qǐng)求的視角來(lái)看，整個(gè) LLM 推理系統(tǒng)仍然表現(xiàn)出不確定性！

事實(shí)上，我們的請(qǐng)求輸出確實(shí)會(huì)受到并行用戶請(qǐng)求的影響。

這種現(xiàn)象并非源于批次間的信息泄露，而是因?yàn)槲覀兊那跋騻鞑ト狈ε尾蛔冃裕╞atch invariance）——換句話說(shuō)，請(qǐng)求的輸出結(jié)果會(huì)依賴于前向傳播時(shí)的具體批次大小。

批次不變性與確定性

為了說(shuō)明批次不變性的問(wèn)題，我們將系統(tǒng)簡(jiǎn)化，專門(mén)考察矩陣乘法（matmul）操作。

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

然而，一個(gè)令人意外的現(xiàn)象是：這些實(shí)現(xiàn)往往不具備"批次不變性（batch-invariant）"。也就是說(shuō)，當(dāng)批次大小發(fā)生變化時(shí)，批次中每個(gè)元素的計(jì)算結(jié)果可能會(huì)隨之改變。

從數(shù)學(xué)的直覺(jué)來(lái)看，這確實(shí)是一個(gè)頗為反常的現(xiàn)象。按理說(shuō)，矩陣乘法應(yīng)該對(duì)批次中的每個(gè)元素"獨(dú)立"進(jìn)行計(jì)算——無(wú)論批次中包含其他什么元素，也無(wú)論批次的總體大小如何，都不應(yīng)該影響到批次中任意特定元素的最終結(jié)果。

但令人困惑的是，實(shí)際的觀察結(jié)果卻與這種數(shù)學(xué)直覺(jué)相悖。

請(qǐng)注意，這里指的是"運(yùn)行間確定性（run-to-run deterministic）"。如果你在同一環(huán)境中多次運(yùn)行這個(gè)腳本，它會(huì)確定性地返回完全相同的結(jié)果。

不過(guò)，這種確定性是有條件的 —— 它并非"跨硬件/軟件版本不變"。當(dāng)你的GPU型號(hào)或PyTorch版本發(fā)生變化時(shí)，可能會(huì)得到不同的數(shù)值結(jié)果。但關(guān)鍵在于：在同一個(gè)運(yùn)行環(huán)境下，結(jié)果始終是可重現(xiàn)的。

然而，問(wèn)題出現(xiàn)在更復(fù)雜的場(chǎng)景中。當(dāng)這種缺乏批次不變性的內(nèi)核被集成到大型推理系統(tǒng)時(shí)，整個(gè)系統(tǒng)的行為就變得不可預(yù)測(cè)了。

想象一下這樣的情況：當(dāng)你向推理服務(wù)端點(diǎn)發(fā)送請(qǐng)求時(shí)，服務(wù)器當(dāng)前的負(fù)載情況對(duì)你來(lái)說(shuō)是"黑盒"的 —— 你無(wú)法預(yù)知也無(wú)法控制。而這個(gè)負(fù)載直接決定了內(nèi)核運(yùn)行時(shí)的批次大小，進(jìn)而影響每個(gè)看似獨(dú)立的請(qǐng)求的最終輸出結(jié)果！這就是為什么同樣的輸入在不同時(shí)刻可能產(chǎn)生不同輸出的根本原因。

Figure4：盡管從整體上可以說(shuō)推理服務(wù)器是“確定性的”，但對(duì)單個(gè)用戶來(lái)說(shuō)情況不同。從單個(gè)用戶的角度來(lái)看，其他并發(fā)用戶并不是系統(tǒng)的“輸入”，而是系統(tǒng)的一個(gè)不確定性屬性。這使得從每個(gè)用戶的角度來(lái)看，LLM 推理是“不確定性的”。

當(dāng)某種內(nèi)核缺乏批次大小不變性時(shí)，一旦遇到不確定的外部因素（比如服務(wù)器負(fù)載的波動(dòng)），整個(gè)系統(tǒng)就會(huì)變得不確定起來(lái)。

實(shí)際上，幾乎所有 LLM 推理端點(diǎn)的不確定性都可以追溯到同一個(gè)根源：服務(wù)器負(fù)載的動(dòng)態(tài)變化會(huì)導(dǎo)致批次大小的不可預(yù)測(cè)變動(dòng)！

這種不確定性并非 GPU 所獨(dú)有——無(wú)論是基于 CPU 還是 TPU 的 LLM 推理端點(diǎn)，都會(huì)面臨同樣的挑戰(zhàn)。

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

我們?cè)撊绾巫寖?nèi)核實(shí)現(xiàn)批次不變性？

為了讓 Transformer 的實(shí)現(xiàn)具有批次不變性，我們需要確保其中的每一個(gè)內(nèi)核都滿足這個(gè)要求。

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

雖然這在 PyTorch 的所有內(nèi)核中確實(shí)成立，但從理論上講，這并非絕對(duì)保證。舉個(gè)例子，在某些 CPU 內(nèi)核實(shí)現(xiàn)中，可能會(huì)對(duì)數(shù)組的部分區(qū)域使用向量化指令（vectorized intrinsics），而對(duì)其他區(qū)域使用標(biāo)量指令，這兩種指令的數(shù)值結(jié)果未必總能做到逐位（bitwise）完全一致。

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

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

順便分享一個(gè)實(shí)用的小知識(shí)：在 Blackwell 架構(gòu)以及搭配 CUDA 12.8+ 的 Hopper 上，NVLink-Sharp 的交換機(jī)內(nèi)歸約（in-switch reductions）是確定性的。和很多技術(shù)細(xì)節(jié)一樣，這個(gè)信息藏在 NCCL 的 GitHub issue 討論中。

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

那么，讓我們先從 RMSNorm 開(kāi)始講起。

批次不變的 RMSNorm

Figure 5：理想情況下，我們希望在并行化策略中避免核心間的通信。實(shí)現(xiàn)這一目標(biāo)的一種方法是將一個(gè)批次元素分配給一個(gè)核心，從而保證每次歸約都完全在一個(gè)核心內(nèi)完成。這就是所謂的“數(shù)據(jù)并行”策略，因?yàn)槲覀冎皇茄刂粋�(gè)不需要通信的維度進(jìn)行并行化。在這個(gè)例子中，我們有四行和四個(gè)核心，充分利用了我們的核心。

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

需要注意的是，這并不意味著我們必須始終使用完全相同的歸約策略。舉個(gè)例子，當(dāng)要?dú)w約的元素?cái)?shù)量發(fā)生變化時(shí)，即使歸約策略相應(yīng)調(diào)整，我們依然可以維持批量不變性。Quack 博客文章中提供了一些精彩的示例，清晰地展示了各種歸約策略的層級(jí)關(guān)系 —— 從線程歸約、warp 歸約到 block 歸約、cluster 歸約。

換句話說(shuō)，只有當(dāng)批量大小的變化直接影響到歸約策略的選擇時(shí)，批量不變性才會(huì)被打破。

現(xiàn)在讓我們來(lái)看看 RMSNorm 的標(biāo)準(zhǔn)并行策略。

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

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

Figure 6：數(shù)據(jù)并行 RMSNorm（適用于更大的批次）將數(shù)據(jù)并行策略擴(kuò)展到更大的批次相當(dāng)直接 —— 與其讓每個(gè)核心只處理一行，不如讓每個(gè)核心依次處理不同的行。這樣可以保持批次不變性，因?yàn)槊總�(gè)批次元素的歸約策略仍然是相同的。

另一方面，減小批次大小可能會(huì)帶來(lái)挑戰(zhàn)。

因?yàn)槲覀優(yōu)槊總�(gè)批次元素分配一個(gè)核心，當(dāng)批次大小逐漸減小時(shí)，勢(shì)必會(huì)遇到核心數(shù)量超過(guò)批次元素?cái)?shù)量的情況，這時(shí)部分核心就會(huì)處于空閑狀態(tài)。

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

Figure 7：分裂歸約 RMSNorm。如果我們的推理批次較小，那么數(shù)據(jù)并行策略可能無(wú)法提供足夠的并行度來(lái)充分利用所有核心。在這種情況下，一個(gè)更高效的做法是將單個(gè)歸約任務(wù)分解到多個(gè)核心上并行處理，從而讓 GPU 資源得到更好的利用。不過(guò)，這種方法會(huì)帶來(lái)一個(gè)權(quán)衡：我們可能會(huì)失去批次不變性。這是因?yàn)楫?dāng)多個(gè)核心并行處理同一個(gè)歸約任務(wù)時(shí)，元素的處理順序可能不再保持一致，進(jìn)而影響最終結(jié)果的確定性。

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

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

如果確實(shí)需要優(yōu)化這類場(chǎng)景，一個(gè)可行的思路是始終采用某種歸約策略，確保即便在很小的批次下也能提供充足的并行度。雖然這種策略在大批次時(shí)會(huì)產(chǎn)生一定程度的過(guò)度并行，但它能讓我們?cè)诟鞣N批次規(guī)模下都獲得相對(duì)穩(wěn)定（盡管不是最優(yōu)）的性能表現(xiàn)。

批次不變的矩陣乘法

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

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

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

然而，就像我們?cè)?RMSNorm 中遇到的問(wèn)題一樣，矩陣乘法的"批次"維度（M 和 N）有時(shí)會(huì)顯得太小，迫使我們不得不沿著歸約維度 (K) 進(jìn)行切分。雖然矩陣乘法擁有兩個(gè)"批次"維度的優(yōu)勢(shì)，但要想充分發(fā)揮 Tensor Cores 的威力，我們?nèi)匀恍枰诿總�(gè)核心上分配足夠多的"工作量"。

舉個(gè)具體例子：假設(shè)你有一個(gè)形狀為 [1024, K] × [K, 1024] 的矩陣乘法，如果采用標(biāo)準(zhǔn)的二維 tile 大小 [128, 128]，那么數(shù)據(jù)并行策略最多只能將這次計(jì)算分配給 64 個(gè)核心。對(duì)于現(xiàn)代 GPU 來(lái)說(shuō)，這樣的并行度顯然還不足以讓硬件"吃飽"。

在矩陣乘法領(lǐng)域，沿著歸約維度進(jìn)行切分的策略有一個(gè)專門(mén)的名字：Split-K Matmul。不過(guò)需要注意的是，和 RMSNorm 的情況類似，采用 Split-K 策略同樣會(huì)破壞批次不變性 (batch invariance)——這是我們?cè)谶x擇并行策略時(shí)需要權(quán)衡的一個(gè)重要考量。

Figure 9：Split-K 矩陣乘法。如果我們的批次維度非常小，就可能沒(méi)有足夠的并行度，這時(shí)就需要使用 Split-K 矩陣乘法。在這個(gè)例子中，我們把每一次歸約拆分給兩個(gè)核心來(lái)執(zhí)行，這兩個(gè)核心會(huì)各自累加，最后再把結(jié)果合并。這樣一來(lái)，雖然每次歸約被分到兩個(gè)核心上，但我們?nèi)匀荒軌蚶冒藗�(gè)核心。

矩陣乘法還有一個(gè)額外的復(fù)雜點(diǎn) —— Tensor Core 指令。

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

每條 Tensor Core 指令（比如

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

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

當(dāng) batch size = 1 時(shí)，最快的內(nèi)核通常甚至完全不會(huì)使用 Tensor Cores。

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

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

雖然這樣會(huì)損失一部分性能，但在大語(yǔ)言模型（LLM）的推理過(guò)程中，這通常并不是災(zāi)難性的。尤其是，Split-K 最需要的情況是 M 和 N 都很小，而幸運(yùn)的是，在我們的場(chǎng)景中，N（也就是模型維度）通常都相當(dāng)大！

Figure 11：盡管實(shí)現(xiàn)了批次不變性，我們相比 cuBLAS 只損失了大約 20% 的性能。需要注意的是，這里用的也不是經(jīng)過(guò)優(yōu)化的 Triton 內(nèi)核（例如沒(méi)有使用 Thinking MachinesA）。不過(guò)，一些性能表現(xiàn)的模式還是能說(shuō)明問(wèn)題，揭示出批次不變性要求導(dǎo)致性能損失的地方。首先，請(qǐng)注意在 非常小的批次大小 下，由于指令過(guò)大以及并行度不足，我們損失了相當(dāng)多的性能。其次，當(dāng)批次大小逐漸增加時(shí)，會(huì)出現(xiàn)一種 “拼圖”式的性能波動(dòng)模式，這是由量化效應(yīng)（tile 和 wave 兩方面）造成的，而這些效應(yīng)通�？梢酝ㄟ^(guò)改變 tile 大小來(lái)緩解。

批次不變的注意力機(jī)制

Figure 12：Flash Attention 2 策略。我們?cè)?Q 維度 上進(jìn)行并行，同時(shí)在 K/V 維度 上做歸約。這意味著整個(gè)歸約過(guò)程都能保持在單個(gè)核心內(nèi)完成，使其成為另一種數(shù)據(jù)并行策略。

在矩陣乘法實(shí)現(xiàn)了批次不變性之后，注意力機(jī)制又引入了兩個(gè)額外的復(fù)雜點(diǎn) —— 這也合乎邏輯，因?yàn)樗旧戆藘纱尉仃嚦朔ā?/p>

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

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

因此，要在大語(yǔ)言模型 (LLM) 的推理中實(shí)現(xiàn)確定性，我們的數(shù)值必須對(duì)兩方面保持不變：既要與一次同時(shí)處理多少請(qǐng)求無(wú)關(guān)，也要與推理引擎如何切分每個(gè)請(qǐng)求無(wú)關(guān)。

接下來(lái)，我們先來(lái)看注意力機(jī)制的標(biāo)準(zhǔn)并行化策略，這個(gè)策略最早由 Flash Attention 2 提出。與 RMSNorm 和矩陣乘法類似，默認(rèn)策略是一種“數(shù)據(jù)并行”策略。由于我們?cè)?Key/Value 張量 上進(jìn)行歸約，所以數(shù)據(jù)并行策略只能在 Query 張量 上實(shí)現(xiàn)并行。

舉個(gè)例子：根據(jù)推理引擎的不同選擇，一個(gè)序列可能被分成多個(gè)部分來(lái)處理（比如分塊預(yù)填充），也可能一次性處理全部（如果預(yù)填充沒(méi)有被拆分）。為了實(shí)現(xiàn)“批次不變性”，關(guān)鍵在于：某個(gè) token 的歸約順序不能依賴于該序列中同時(shí)處理了多少其他 token。

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

舉例來(lái)說(shuō)：當(dāng)處理序列中的第 1000 個(gè) Query token 時(shí)，歸約順序必須完全一致——不論 KV 緩存中有 0 個(gè) token（預(yù)填充階段）還是有 999 個(gè) token（解碼階段）。

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

舉個(gè)例子：假設(shè)塊大小是 32，而我們當(dāng)前在 KV 緩存中有 80 個(gè)元素。接著我們又計(jì)算了額外 48 個(gè)不在緩存中的元素。

在這種情況下：

計(jì)算 P cache 需要 3 個(gè)塊（兩個(gè)完整塊 + 一個(gè)帶掩碼的塊）。

計(jì)算 P 需要 2 個(gè)塊（一個(gè)完整塊 + 一個(gè)帶掩碼的塊）。

這樣一來(lái)，總共需要 5 個(gè)塊來(lái)完成歸約，但實(shí)際上我們只有 128 個(gè)元素，即 4 個(gè)塊。這就必然會(huì)改變歸約的順序。

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

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

結(jié)合這一點(diǎn)（以及上一節(jié)提到的所有細(xì)節(jié)，比如保持一致的 tile 大�。覀兙湍軌�?qū)崿F(xiàn)一個(gè) 批次不變的注意力機(jī)制 (batch-invariant attention)！

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

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

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

遺憾的是，這種情況不像 RMSNorm 和矩陣乘法那樣可以輕松忽略。舉個(gè)例子：如果 KV 緩存非常長(zhǎng)，即便我們只在處理一個(gè)請(qǐng)求，注意力內(nèi)核也可能會(huì)耗費(fèi)非常長(zhǎng)的時(shí)間。

Figure 14：固定的 # Split-KV 策略（也叫 Flash Decode）。如果我們的 query 長(zhǎng)度變得非常�。ň拖裨诮獯a階段那樣），我們可能會(huì)遇到這樣一種情況：內(nèi)核中幾乎沒(méi)有多少并行度可利用。在這些情況下，我們需要再次沿著歸約維度進(jìn)行拆分 —— 這一次是 KV 維度。

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

但不幸的是，這種方法也會(huì)破壞 批量不變性（batch invariance），因?yàn)槲覀兙_的歸約策略取決于當(dāng)前請(qǐng)求中要處理多少個(gè)序列的 query token。

此外，注意力中常用的 拆分歸約（split-reduction）策略 也會(huì)給批量不變性帶來(lái)挑戰(zhàn)。舉例來(lái)說(shuō)，F(xiàn)lashInfer 的“平衡調(diào)度算法（balanced scheduling algorithm）”會(huì)選擇一種最大化的拆分規(guī)模，以便仍能讓 GPU 的所有核心飽和，因此這種歸約策略并不是“批量不變的（batch-invariant）”。不過(guò)，與 RMSNorm/矩陣乘法不同，僅僅選擇一個(gè)固定的拆分?jǐn)?shù)量，而不管批量大小是多少，是不夠的。

為了實(shí)現(xiàn)批量不變性，我們必須采用“固定拆分大�。╢ixed split-size）”的策略。換句話說(shuō)，我們不是固定拆分的數(shù)量，而是固定每個(gè)拆分的大小，最終得到的拆分?jǐn)?shù)量則會(huì)隨之變化。這樣一來(lái)，我們就能保證無(wú)論處理多少 token，始終執(zhí)行相同的歸約順序。

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

這樣，我們就能夠保持 批量不變性（batch invariance），因?yàn)槲覀兊臍w約策略不再依賴于一次要處理多少個(gè) query token！

實(shí)現(xiàn)

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

實(shí)驗(yàn)

補(bǔ)全（completions）的不確定性有多大？

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

觀察補(bǔ)全結(jié)果的分歧點(diǎn)，我們發(fā)現(xiàn)前 102 個(gè) token 的補(bǔ)全完全一致！首次出現(xiàn)差異是在 第 103 個(gè) token。所有補(bǔ)全都會(huì)生成序列 “Feynman was born on May 11, 1918, in”。但其中 992 個(gè)補(bǔ)全繼續(xù)生成了 “Queens, New York”，而 8 個(gè)補(bǔ)全生成了 “New York City”。

另一方面，當(dāng)我們啟用 批量不變內(nèi)核 時(shí)，所有 1000 個(gè)補(bǔ)全結(jié)果完全一致。這正是從數(shù)學(xué)角度來(lái)看我們對(duì)采樣器的預(yù)期，但如果不使用批量不變內(nèi)核，我們無(wú)法實(shí)現(xiàn)這種確定性結(jié)果。

性能

在這里我們并沒(méi)有對(duì)批量不變內(nèi)核進(jìn)行大量性能優(yōu)化。不過(guò)，我們?nèi)匀贿M(jìn)行了一些實(shí)驗(yàn)來(lái)驗(yàn)證性能是否可用。

我們搭建了一個(gè) API 服務(wù)器，使用一張 GPU 運(yùn)行 Qwen-3-8B，并請(qǐng)求生成 1000 條序列，每條的輸出長(zhǎng)度在 90 到 110 之間。

大部分的性能下降，來(lái)自于 vLLM 中 FlexAttention 的集成尚未經(jīng)過(guò)深度優(yōu)化。盡管如此，我們可以看到性能并沒(méi)有糟糕到不可接受的程度。

滿血版 On-Policy 強(qiáng)化學(xué)習(xí)（RL）

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

當(dāng)然，如果連兩次完全相同的推理請(qǐng)求都無(wú)法得到逐位（bitwise）一致的結(jié)果，那么訓(xùn)練和推理之間就更不可能實(shí)現(xiàn)逐位一致。而一旦實(shí)現(xiàn)了 確定性推理，我們也就能夠修改訓(xùn)練棧，使得采樣和訓(xùn)練之間的結(jié)果逐位一致，從而實(shí)現(xiàn)真正的 on-policy RL。

我們?cè)?BigMath 上運(yùn)行了一個(gè) RLVR 設(shè)置實(shí)驗(yàn)，RL 策略從 Qwen 2.5-VL instruct 8B 初始化，最大采樣長(zhǎng)度為 4096。

如果我們?cè)谟?xùn)練中不做 off-policy 修正（即：重要性加權(quán)，importance weighting），獎(jiǎng)勵(lì)會(huì)在訓(xùn)練過(guò)程中途崩潰；而在訓(xùn)練中加入 off-policy 修正項(xiàng)，訓(xùn)練則能順利進(jìn)行。但如果我們能夠讓采樣器和訓(xùn)練器之間的結(jié)果逐位一致，那么我們就是完全的 on-policy（即 KL 散度為 0），同樣也能順利完成訓(xùn)練。

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

使用 重要性加權(quán) 時(shí)，KL 散度大約維持在 0.001，偶爾會(huì)有尖峰；

不使用 重要性加權(quán) 時(shí)，KL 散度最終會(huì)出現(xiàn)尖峰，并與獎(jiǎng)勵(lì)崩潰的時(shí)間點(diǎn)吻合；

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

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

結(jié)論

現(xiàn)代軟件系統(tǒng)往往由多層抽象構(gòu)成。在機(jī)器學(xué)習(xí)中，當(dāng)我們遇到不確定性和一些微妙的數(shù)值差異時(shí)，人們往往會(huì)選擇視而不見(jiàn)。

畢竟，我們的系統(tǒng)本來(lái)就是「概率性的」，再多一點(diǎn)不確定性又有何妨呢？單元測(cè)試掛掉時(shí)，把 atol/rtol 調(diào)大點(diǎn)又有什么問(wèn)題？訓(xùn)練器和采樣器之間的對(duì)數(shù)概率差異，應(yīng)該不算是真正的 bug 吧？

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

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