大家好，我是Ai學(xué)習(xí)的老章

推薦一篇迎合文章，F(xiàn)rom：Horace He 與 Thinking Machines 的其他成員

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

例如，你可能會(huì)發(fā)現(xiàn)多次向 ChatGPT 提出同一個(gè)問題會(huì)得到不同的結(jié)果。這本身并不奇怪，因?yàn)閺恼Z言模型獲取結(jié)果涉及“采樣”，這是一個(gè)將語言模型的輸出轉(zhuǎn)換為概率分布并依概率選擇 token 的過程。

更令人驚訝的可能是，即使我們將溫度調(diào)到 0這意味著 LLM 總是選擇概率最高的 token，這被稱為貪婪采樣。（理論上使采樣變得確定），LLM API 在實(shí)際中仍然不具備確定性。即便在你自己的硬件上使用 vLLM 或 SGLang 這類 OSS 推理庫運(yùn)行推理，采樣依舊不是確定性的。

但為什么 LLM 推理引擎不具備確定性？一種常見假設(shè)是，浮點(diǎn)非結(jié)合性與并發(fā)執(zhí)行共同作用，導(dǎo)致非確定性，具體取決于哪個(gè)并發(fā)核心先完成。我們將此稱為 LLM 推理非確定性的“并發(fā) + 浮點(diǎn)”假設(shè)。例如，最近的一篇 arXiv 預(yù)印本寫道：

GPU 中的浮點(diǎn)運(yùn)算表現(xiàn)出非結(jié)合性，意味著 (a+b)+c≠a+(b+c)(a+b)+c=a+(b+c) 由于有限精度和舍入誤差。這一特性直接影響 Transformer 架構(gòu)中注意力分?jǐn)?shù)和 logits 的計(jì)算，其中跨多個(gè)線程的并行操作可能因執(zhí)行順序不同而產(chǎn)生不同結(jié)果。

你也可以在其他地方看到“并發(fā) + 浮點(diǎn)”這一假設(shè)被反復(fù)提及，比如這里（“存在速度權(quán)衡，為了讓端點(diǎn)足夠快，我們使用 GPU，而 GPU 會(huì)進(jìn)行并行 [非確定性] 計(jì)算。任何現(xiàn)代 GPU 上的神經(jīng)網(wǎng)絡(luò)計(jì)算都會(huì)受到這些影響?！保蛘哌@里（“由于 GPU 高度并行化，每次執(zhí)行時(shí)加法或乘法的順序可能不同，這會(huì)導(dǎo)致輸出出現(xiàn)微小差異?！保?/p>

雖然這一假設(shè)并非完全錯(cuò)誤，但它并沒有揭示全貌。例如，即使在 GPU 上，對(duì)同一數(shù)據(jù)反復(fù)執(zhí)行相同的矩陣乘法，結(jié)果也會(huì)始終位級(jí)相等。我們確實(shí)在使用浮點(diǎn)數(shù)，我們的 GPU 也確實(shí)擁有大量并發(fā)。為什么在這個(gè)測(cè)試中我們看不到非確定性？

A = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16) B = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16) ref = torch.mm(A, B) for _ in range(1000):     assert (torch.mm(A, B) - ref).abs().max().item() == 0

要理解 LLM 推理非確定性的真正原因，我們必須看得更深。

不幸的是，就連定義“LLM 推理是確定性的”到底意味著什么都很困難。可能令人困惑的是，以下所有說法同時(shí)成立：

1. 某些 GPU 上的 kernel 是非確定性的。

2. 然而，語言模型前向傳播中使用的所有 kernel 都是確定性的。

3. 此外，LLM 推理服務(wù)器（如 vLLM）的前向傳播也可以被認(rèn)為是確定性的。

4. 然而，從任何使用推理服務(wù)器的人的角度來看，結(jié)果都是非確定性的。

在這篇文章中，我們將解釋“并發(fā) + 浮點(diǎn)”假設(shè)為何偏離靶心，揭開 LLM 推理非確定性的真正元兇，并說明如何擊敗非確定性，在 LLM 推理中獲得真正可復(fù)現(xiàn)的結(jié)果。

原罪：浮點(diǎn)運(yùn)算的非結(jié)合性

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

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

(a+b)+c≠a+(b+c)(a+b)+c=a+(b+c)

(0.1 + 1e20) - 1e20 >>> 0 0.1 + (1e20 - 1e20) >>> 0.1

諷刺的是，正是打破結(jié)合律才讓浮點(diǎn)數(shù)變得有用。

浮點(diǎn)數(shù)之所以有用，是因?yàn)樗鼈冊(cè)试S“動(dòng)態(tài)”的精度水平。為了便于解釋，我們將使用十進(jìn)制（而非二進(jìn)制），其中浮點(diǎn)數(shù)的格式為 mantissa?10exponentmantissa?10exponent 。我們還將使用 3 位尾數(shù)和 1 位指數(shù)。

例如，對(duì)于值 3450，我們可以精確地表示為 3.45?1033.45?103 。我們也可以表示更小的值，如 0.486，表示為 4.86?10?14.86?10?1 。通過這種方式，浮點(diǎn)數(shù)讓我們既能表示非常小的值，也能表示非常大的值。在科學(xué)領(lǐng)域，我們可能會(huì)說浮點(diǎn)數(shù)讓我們保持恒定的“有效數(shù)字”位數(shù)。

如果將兩個(gè)具有相同指數(shù)的浮點(diǎn)數(shù)相加，看起來類似于整數(shù)加法。例如，123（ 1.23?1021.23?102 ）+ 456（ 4.56?1024.56?102 ）結(jié)果為 579（ 5.79?1025.79?102 ）。

但當(dāng)我們將兩個(gè)指數(shù)不同的浮點(diǎn)數(shù)相加時(shí)，比如 1230 和 23.4，會(huì)發(fā)生什么？此時(shí)精確結(jié)果為 1253.4。然而，我們一次只能保留 3 位有效數(shù)字。因此，浮點(diǎn)加法會(huì)舍去最后兩位，得到 1.25?1031.25?103 （即 1250）。

我們需要 3 位有效數(shù)字來表示 1230，也需要 3 位有效數(shù)字來表示 23.4。然而，將這兩個(gè)數(shù)相加后，結(jié)果需要 5 位有效數(shù)字才能精確表示（1253.4）。我們的浮點(diǎn)格式只能把末尾的 34 截掉。某種意義上，我們實(shí)際上先把原來的 23.4 四舍五入成了 20.0，然后再相加。

但此時(shí)，我們已經(jīng)破壞了信息。請(qǐng)注意，每當(dāng)我們把兩個(gè)“尺度”不同（即指數(shù)不同）的浮點(diǎn)數(shù)相加時(shí)，都可能發(fā)生這種情況。而指數(shù)不同的浮點(diǎn)數(shù)相加在實(shí)際中非常常見。事實(shí)上，如果我們能保證永遠(yuǎn)不需要不同的指數(shù)，那干脆用整數(shù)就好了！

換句話說，每當(dāng)我們以不同順序相加浮點(diǎn)數(shù)時(shí)，就可能得到完全不同的結(jié)果。舉個(gè)極端的例子，僅因求和順序不同，這個(gè)數(shù)組就可能產(chǎn)生 102 種不同的結(jié)果。

import random vals = [1e-10, 1e-5, 1e-2, 1] vals = vals + [-v for v in vals] results = [] random.seed(42) for _ in range(10000):     random.shuffle(vals)     results.append(sum(vals)) results = sorted(set(results)) print(f"There are {len(results)} unique results: {results}") # Output: # There are 102 unique results: [-8.326672684688674e-17, -7.45931094670027e-17, ..., 8.326672684688674e-17]

盡管這是輸出不一致的根本原因，但它并未直接解釋非確定性的來源。它無法幫助我們理解浮點(diǎn)值為何會(huì)以不同順序相加、何時(shí)發(fā)生這種情況，以及如何避免。

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

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

如上所述，對(duì)于內(nèi)核為何以不同順序相加數(shù)字，一種常見解釋是“并發(fā) + 浮點(diǎn)”假說。該假說指出，如果并發(fā)線程完成的順序是非確定性的，并且累加順序依賴于并發(fā)線程完成的順序（例如使用原子加法），那么我們的累加順序也將是非確定性的。

令人困惑的是，盡管這可能導(dǎo)致非確定性內(nèi)核，但在 LLM 推理的非確定性中，并發(fā)（以及原子加法）最終卻完全無關(guān)！為了解釋真正的罪魁禍?zhǔn)资鞘裁?，我們首先來理解為什么現(xiàn)代 GPU 內(nèi)核很少需要原子加法。

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

通常，GPU 會(huì)在許多“核心”（即 SM）上并發(fā)地啟動(dòng)一個(gè)程序。由于這些核心之間沒有內(nèi)在的同步機(jī)制，當(dāng)它們需要相互通信時(shí)就會(huì)帶來挑戰(zhàn)。例如，如果所有核心都必須累加到同一個(gè)元素，你可以使用“原子加”（有時(shí)稱為“fetch-and-add”）。原子加是“非確定性的”——結(jié)果累加的順序完全取決于哪個(gè)核心最先完成。

具體地說，假設(shè)你正在用 100 個(gè)核心歸約一個(gè) 100 元素的向量（例如torch.sum()）。雖然你可以并行加載全部 100 個(gè)元素，但最終必須歸約到單個(gè)元素。實(shí)現(xiàn)這一點(diǎn)的一種方法是使用某種“原子加”原語，硬件保證所有加法都會(huì)被處理，但不保證順序。

原子加確保每個(gè)核心的貢獻(xiàn)都會(huì)體現(xiàn)在最終和中。然而，它并不保證這些貢獻(xiàn)將以何種順序被累加。順序完全取決于哪個(gè)核心最先完成，這是一種非確定性屬性。因此，多次執(zhí)行同一個(gè)并行程序可能會(huì)產(chǎn)生非確定性的輸出。

這通常就是人們所說的“非確定性”——你用完全相同的輸入兩次執(zhí)行同一個(gè)內(nèi)核，卻得到了不同的結(jié)果。這被稱為“運(yùn)行間非確定性”，即你用完全相同的依賴項(xiàng)兩次運(yùn)行同一個(gè) Python 腳本，卻得到了不同的結(jié)果。

盡管并發(fā)原子加法確實(shí)會(huì)讓內(nèi)核變得非確定性，但絕大多數(shù)內(nèi)核并不需要原子加法。事實(shí)上，在 LLM 的典型前向傳播中，通常連一個(gè)原子加法都不會(huì)出現(xiàn)。

考慮到并行化歸約操作可以從原子加法中受益，這一點(diǎn)可能會(huì)令人驚訝。原子加法最終不被需要主要有兩個(gè)原因。

1. 通常沿著“批次”維度已經(jīng)有足夠的并行度，因此我們無需沿著歸約維度進(jìn)行并行化。例如，假設(shè)我們不是歸約單個(gè) 100 維向量，而是并行歸約 500 個(gè)向量。在這種情況下，我們可以在每個(gè)核心上歸約一個(gè)完整的向量，并讓每個(gè)核心處理不同的向量。

2. 隨著時(shí)間的推移，大多數(shù)神經(jīng)網(wǎng)絡(luò)庫都采用了多種策略，在不影響性能的前提下實(shí)現(xiàn)確定性。例如，我們可以進(jìn)行“拆分”（或樹形）歸約，將 100 個(gè)元素的歸約拆分為五個(gè) 20 個(gè)元素的歸約（從而實(shí)現(xiàn)五路并行）。然后，為了合并剩下的五個(gè)元素，我們可以執(zhí)行一次單獨(dú)的“清理”歸約（這部分不再并行，但元素?cái)?shù)量很少，開銷極低），或者使用信號(hào)量（確保每個(gè)并發(fā)線程塊按確定順序累加）。信號(hào)量策略的描述可在此處找到。

由于這兩個(gè)因素，在絕大多數(shù)神經(jīng)網(wǎng)絡(luò)運(yùn)算中，避免原子加操作帶來的性能損失可以忽略不計(jì)。

仍有少數(shù)常見運(yùn)算在避免原子加時(shí)會(huì)帶來顯著的性能損失。例如，PyTorch 中的scatter_add（a[b] += c）。然而，在 LLMs 中唯一常用的就是 FlashAttention 的反向傳播。有趣的事實(shí)：你知道嗎？廣泛使用的 Triton 版 FlashAttention 反向?qū)崿F(xiàn)，在算法上與 Tri Dao 的 FlashAttention-2 論文并不相同？標(biāo)準(zhǔn)的 Triton 實(shí)現(xiàn)會(huì)在反向傳播中額外重新計(jì)算，從而避免原子操作，但代價(jià)是 FLOPs 增加 40%！

然而，LLM 的前向傳播中沒有任何需要原子加法的操作。因此，LLM 的前向傳播實(shí)際上是“運(yùn)行到運(yùn)行確定性的”。

從推理服務(wù)器的角度來看，它是確定性的。給定完全相同的用戶請(qǐng)求，它總會(huì)給出相同的確定性輸出。

維基百科寫道：“確定性算法是指給定特定輸入時(shí)，總會(huì)產(chǎn)生相同輸出的算法。”而在這個(gè)場(chǎng)景下，給定完全相同的輸入（即推理服務(wù)器正在處理的完全相同的請(qǐng)求），前向傳播總會(huì)產(chǎn)生完全相同的輸出。

然而，前向傳播本身“確定”并不足以保證包含它的整個(gè)系統(tǒng)也是確定的。例如，如果我們的請(qǐng)求輸出依賴于并行的用戶請(qǐng)求（例如 batch-norm）呢？由于每個(gè)單獨(dú)請(qǐng)求都無法預(yù)知并行請(qǐng)求會(huì)是什么，從它們的角度看，我們的整體 LLM 推理也是非確定的！

事實(shí)證明，我們的請(qǐng)求輸出確實(shí)依賴于并行的用戶請(qǐng)求。并不是因?yàn)槲覀円阅撤N方式在批次之間泄露信息——而是我們的前向傳播缺乏“批次不變性”，導(dǎo)致我們的請(qǐng)求輸出依賴于前向傳播的批次大小。

批次不變性與“確定性”

為了解釋 batch invariance，讓我們簡(jiǎn)化系統(tǒng)，只看 matmul。你可以假設(shè)所有 matmul 實(shí)現(xiàn)都是“運(yùn)行間確定”的。這并不完全正確，但大多數(shù)常見的 matmul 實(shí)現(xiàn)確實(shí)具有這一特性。然而，它們并不是“batch 不變”的。換句話說，當(dāng) batch size 改變時(shí)，batch 中的每個(gè)元素都可能得到不同的結(jié)果。

從數(shù)學(xué)角度來看，這是一個(gè)相當(dāng)不尋常的特性。矩陣乘法在 batch 的每個(gè)元素上應(yīng)該是“獨(dú)立”的——batch 中的其他元素或 batch 的大小都不應(yīng)影響 batch 中某個(gè)特定元素的計(jì)算結(jié)果。

然而，正如我們憑經(jīng)驗(yàn)觀察到的那樣，事實(shí)并非如此。

import torch torch.set_default_device('cuda')  B = 2048 D = 4096 a = torch.linspace(-1000, 1000, B*D).reshape(B, D) b = torch.linspace(-1000, 1000, D*D).reshape(D, D) # Doing a matrix vector multiplication by taking # the first element of the batch out1 = torch.mm(a[:1], b) # Doing a matrix matrix multiplication and then taking # the first element of the batch out2 = torch.mm(a, b)[:1] print((out1 - out2).abs().max()) # tensor(1669.2500, device='cuda:0')

請(qǐng)注意，這是“運(yùn)行間確定性”。如果你多次運(yùn)行該腳本，它會(huì)確定性地返回相同的結(jié)果。它并非“硬件/軟件版本不變”——你的 GPU/PyTorch 版本可能會(huì)返回不同的值，但它應(yīng)該確定性地返回相同的值。

然而，當(dāng)一個(gè)非批次不變的內(nèi)核被用作更大推理系統(tǒng)的一部分時(shí)，系統(tǒng)就可能變得非確定性。當(dāng)你向推理端點(diǎn)發(fā)出查詢時(shí)，服務(wù)器當(dāng)前的負(fù)載量從用戶角度來看實(shí)際上是“非確定性”的。負(fù)載決定了內(nèi)核運(yùn)行的批次大小，從而改變了每個(gè)單獨(dú)請(qǐng)求的最終結(jié)果！

盡管推理服務(wù)器本身可以被認(rèn)為是“確定性的”，但對(duì)單個(gè)用戶而言情況卻不同。從單個(gè)用戶的角度來看，其他并發(fā)用戶并不是系統(tǒng)的“輸入”，而是系統(tǒng)的一種非確定性屬性。這使得 LLM 推理在每個(gè)用戶看來都是“非確定性”的。

如果你將某個(gè)內(nèi)核不具備不變性的屬性（例如 batch-size）與該屬性的非確定性（例如服務(wù)器當(dāng)前負(fù)載）組合在一起，就會(huì)得到一個(gè)非確定性系統(tǒng)。

換句話說，幾乎所有 LLM 推理端點(diǎn)之所以非確定，根本原因就是負(fù)載（進(jìn)而導(dǎo)致 batch-size）本身在不可預(yù)測(cè)地變化！這種非確定性并非 GPU 獨(dú)有——無論是 CPU 還是 TPU 提供的 LLM 推理端點(diǎn)，同樣會(huì)受這一非確定性來源的影響。

因此，若想在我們的推理服務(wù)器中避免非確定性，就必須在內(nèi)核層面實(shí)現(xiàn) batch 不變性。為了弄清如何做到這一點(diǎn)，我們先來看看為什么內(nèi)核一開始就不具備 batch 不變性。

我們?nèi)绾巫寖?nèi)核具備 batch 不變性？

為了讓 transformer 實(shí)現(xiàn)對(duì) batch 不敏感，我們必須讓每個(gè) kernel 都對(duì) batch 不敏感。幸運(yùn)的是，我們可以假設(shè)所有逐點(diǎn)運(yùn)算都是對(duì) batch 不敏感的。盡管對(duì)于 PyTorch 中的所有 kernel 來說確實(shí)如此，但這并非必然成立。例如，CPU 上的一些 kernel 實(shí)現(xiàn)會(huì)在數(shù)組的某些部分使用向量化 intrinsic，而在其他部分使用非向量化 intrinsic，而這些 intrinsic 的數(shù)值結(jié)果并不總是逐位一致。因此，我們只需關(guān)注涉及歸約的 3 種操作——RMSNorm、矩陣乘法和注意力。與并行相關(guān)的歸約不在本文討論范圍內(nèi)，但同樣的原則適用。一個(gè)可能有用的信息是：在 Blackwell 以及使用 CUDA 12.8+ 的 Hopper 上，NVLink-Sharp 的 in-switch 歸約是確定性的。和許多事情一樣，這些信息可以在 NCCL 的 GitHub issues 中找到。

方便的是，這些也按難度遞增的順序排列。每一項(xiàng)都需要額外考慮，才能在合理性能下實(shí)現(xiàn)批次不變性。我們先從 RMSNorm 說起。

批次不變的 RMSNorm
數(shù)據(jù)并行的 RMSNorm 理想情況下，我們希望并行策略中核心之間無需通信。一種實(shí)現(xiàn)方法是將每個(gè)批次元素分配給單獨(dú)的核心，從而保證所有歸約操作完全在一個(gè)核心內(nèi)完成。這就是所謂的“數(shù)據(jù)并行”策略，因?yàn)槲覀冎皇茄刂鵁o需通信的維度進(jìn)行并行。在此示例中，我們有四行和四個(gè)核心，正好占滿所有核心。

RMSNorm 的實(shí)現(xiàn)如下：

# x: [batch_size, hidden_dim] # weight: [hidden_dim] def rms_norm(x, weight):     return x * torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True)) * weight

批不變性的要求是：無論內(nèi)核的批大小如何，每個(gè)元素的歸約順序都必須固定。請(qǐng)注意，這并不意味著我們必須始終使用相同的歸約策略。例如，如果我們改變要?dú)w約的元素?cái)?shù)量，即使歸約策略發(fā)生變化，我們?nèi)匀豢梢员３峙蛔冃?。Quack 博客文章中有一些很好的示例，展示了可以使用的各種歸約策略的層次結(jié)構(gòu)（例如線程歸約、warp 歸約、block 歸約、cluster 歸約）。

因此，只有當(dāng)我們的批大小影響歸約策略時(shí)，我們才會(huì)破壞批不變性。

讓我們來看看 RMSNorm 的標(biāo)準(zhǔn)并行策略。一般而言，并行算法通過最小化跨核心的通信來獲得收益。為了本次討論的目的，你可以假設(shè)當(dāng)我們提到“cores”時(shí)，我們指的是 SMs。更具體地說，這里重要的屬性是：我們內(nèi)核啟動(dòng)的 threadblock 數(shù)量大于 SMs 的數(shù)量。因此，我們可以從一個(gè)簡(jiǎn)單的策略開始：將每個(gè)批次元素分配給單個(gè)核心，如上圖所示。

增大批次大小不會(huì)影響我們的歸約策略；如果批次大小為 200 就能為內(nèi)核提供足夠的并行度，那么批次大小為 2000 時(shí)肯定也能提供足夠的并行度。

更大批次的數(shù)據(jù)并行 RMSNorm 將數(shù)據(jù)并行策略擴(kuò)展到更大批次非常簡(jiǎn)單——不再讓每個(gè)核心處理一行，而是讓每個(gè)核心按順序處理不同的行。這保持了批次不變性，因?yàn)槊總€(gè)批次元素的歸約策略保持不變。

另一方面，減小批大小也會(huì)帶來挑戰(zhàn)。由于我們將每個(gè)批元素分配給一個(gè)核心，當(dāng)批大小減小時(shí)，最終會(huì)出現(xiàn)核心數(shù)量多于批元素的情況，導(dǎo)致部分核心閑置。

遇到這種情況，優(yōu)秀的內(nèi)核工程師會(huì)采用上一節(jié)提到的解決方案（原子加法或拆分歸約），以保持高并行度，從而維持良好性能。然而，這會(huì)改變歸約策略，導(dǎo)致該內(nèi)核不再具備批不變性。

拆分歸約的 RMSNorm 如果批大小較小，我們的數(shù)據(jù)并行策略可能無法提供足夠的并行度來充分利用所有核心。此時(shí)，將歸約操作“拆分”到多個(gè)核心上執(zhí)行可能更高效，從而充分利用 GPU。但這會(huì)失去批不變性，因?yàn)槲覀儾辉僖韵嗤樞驓w約每個(gè)元素。

最簡(jiǎn)單的解決方案是直接忽略這些情況。這并非完全不合理——小批大小意味著內(nèi)核本身執(zhí)行速度較快，因此性能下降可能不會(huì)造成災(zāi)難性后果。

如果我們被迫優(yōu)化這一用例，一種方法是始終采用一種即便在極小批量下也具備足夠并行度的歸約策略。這種策略在較大批量時(shí)會(huì)產(chǎn)生過量并行，但能在整個(gè)尺寸范圍內(nèi)都獲得尚可（而非峰值）的性能。

批不變矩陣乘法
數(shù)據(jù)并行 Matmul 與 RMSNorm 類似，matmul 的標(biāo)準(zhǔn)并行策略是“數(shù)據(jù)并行”，將整個(gè)規(guī)約操作保留在一個(gè)核心內(nèi)。最直觀的做法是把輸出張量拆分成 2D 瓦片，并將每塊瓦片分配給不同的核心。每個(gè)核心隨后計(jì)算屬于該瓦片的點(diǎn)積，再次在單個(gè)核心內(nèi)完成全部規(guī)約。

與 RMSNorm 不同的是，圍繞算術(shù)強(qiáng)度以及充分利用 Tensor Core 的額外約束，迫使我們?cè)诟咝?matmul 內(nèi)核中拆分 2D 瓦片，而不是單個(gè)輸出元素。

本質(zhì)上，你可以把矩陣乘法看作一個(gè)逐點(diǎn)操作后再進(jìn)行規(guī)約。于是，如果我們通過將輸出分塊來并行化矩陣乘法，就得到了一種類似的“數(shù)據(jù)并行”內(nèi)核策略，使每次規(guī)約都保留在單個(gè)核心內(nèi)。

與 RMSNorm 類似，我們的“批”維度（M 和 N）也可能變得過小，從而被迫沿著歸約維度（K）切分。盡管有兩個(gè)“批”維度，矩陣乘法仍需要每個(gè)核心承擔(dān)更多“工作量”，才能有效利用 Tensor Core。例如，對(duì)于 [1024, K] × [K, 1024] 的矩陣乘法，若采用標(biāo)準(zhǔn)的 2D 瓦片大小 [128, 128]，數(shù)據(jù)并行策略只能將其拆分到 64 個(gè)核心，不足以讓 GPU 飽和。

在矩陣乘法中沿著歸約維度切分被稱為 Split-K Matmul。與 RMSNorm 一樣，這種策略會(huì)破壞批不變性。

另一種有趣的矩陣乘法并行策略是 stream-k。stream-k 的有趣之處在于，它比典型的矩陣乘法具有更少的“不變性”。如前所述，大多數(shù)矩陣乘法庫并非 batch-invariant，但至少可以稱為 batch-position-invariant（即改變 batch 中元素的位置不會(huì)影響數(shù)值結(jié)果）。然而，stream-k 連 batch-position-invariant 都不是！其核心洞見是：通過為不同的輸出 tile 以不同方式沿 k 維度切分，可以獲得更均衡的負(fù)載，但利用這一點(diǎn)會(huì)使我們的 kernel 也不再具備 batch-position-invariant 特性。

Split-K 矩陣乘法 如果我們的 batch 維度非常小，可能無法提供足夠的并行度，此時(shí)就需要使用 split-k 矩陣乘法。在這個(gè)例子中，我們將每個(gè)規(guī)約操作拆分到兩個(gè)核心上，這兩個(gè)核心分別累加，最后再合并結(jié)果。然而，把每個(gè)規(guī)約拆分到兩個(gè)核心，仍讓我們能夠充分利用八個(gè)核心。

矩陣乘法還有一個(gè)額外的復(fù)雜性——張量核心指令。對(duì)于歸約操作，我們可以一次只處理一行，而高效的矩陣乘法內(nèi)核必須一次處理整個(gè)“瓦片”。

每條張量核心指令（例如wgmma.mma_async.sync.aligned.m64n128k16）內(nèi)部可能采用不同的歸約順序。選擇不同張量核心指令的一個(gè)原因可能是批次非常小。例如，如果我們使用一條對(duì)長(zhǎng)度為 256 的瓦片進(jìn)行運(yùn)算的張量核心 PTX 指令，而批次大小只有 32，那么幾乎浪費(fèi)了所有算力！當(dāng)批次大小為 1 時(shí)，最快的內(nèi)核通常完全不使用張量核心。

填充的 Tensor-Core 指令 如果 batch size 太小，我們可能會(huì)遇到連一個(gè) 2D tile 都無法放入輸出的情況。此時(shí)，最有效的方法是切換到更小的 tensor-core 指令，或者干脆不用 tensor-core！然而，這兩種選擇都會(huì)使我們的 kernel 無法保持 batch 不變性。

因此，確保 matmul 的 batch 不變性最簡(jiǎn)單的方法是：編譯一個(gè) kernel 配置，并在所有形狀下都使用它。雖然會(huì)損失一些性能，但在 LLM 推理中這通常不會(huì)帶來災(zāi)難性后果。特別是，split-k 在 M 和 N 都很小時(shí)才最需要，而幸運(yùn)的是，在我們的場(chǎng)景里 N（即模型維度）通常非常大！

盡管實(shí)現(xiàn)了批次不變性，與 cuBLAS 相比我們只損失了約 20% 的性能。請(qǐng)注意，這也不是一個(gè)經(jīng)過優(yōu)化的 Triton 內(nèi)核（例如沒有使用 TMA）。然而，性能中的一些模式可以說明我們的批次不變需求在何處導(dǎo)致性能下降。首先，在極小的批次規(guī)模下，由于指令過大且并行度不足，我們損失了大量性能。其次，隨著批次規(guī)模增加，會(huì)出現(xiàn)一種“拼圖”模式，這是由量化效應(yīng)（包括 tile 和 wave）引起的，通常通過改變 tile 大小可以緩解。你可以在這里了解更多關(guān)于這些量化效應(yīng)的信息。

批次不變注意力

FlashAttention2 策略 我們沿著 Q 并行化，同時(shí)沿著 K/V 進(jìn)行歸約。這意味著我們的整個(gè)歸約可以保持在單個(gè)核心內(nèi)，使其成為另一種數(shù)據(jù)并行策略。

在為矩陣乘法實(shí)現(xiàn)批次不變性之后，注意力機(jī)制又引入了兩個(gè)額外的難題——恰如其分，因?yàn)樗瑑蓚€(gè)矩陣乘法。

1. 與 RMSNorm 和 matmul 僅沿特征維度進(jìn)行歸約不同，我們現(xiàn)在同時(shí)沿特征維度和序列維度進(jìn)行歸約。

2. 由于上述原因，注意力機(jī)制必須處理各種影響序列處理方式的推理優(yōu)化（分塊預(yù)填充、前綴緩存等）。

因此，為了在 LLM 推理中實(shí)現(xiàn)確定性，我們的數(shù)值計(jì)算必須不受以下兩個(gè)因素影響：一次處理多少請(qǐng)求，以及推理引擎如何對(duì)每個(gè)請(qǐng)求進(jìn)行切片。

讓我們首先回顧 FlashAttention2 首次引入的標(biāo)準(zhǔn)注意力并行策略。與 RMSNorm 和 Matmul 類似，默認(rèn)策略是“數(shù)據(jù)并行”策略。由于我們沿著 key/value 張量進(jìn)行歸約，數(shù)據(jù)并行策略只能沿著 query 張量進(jìn)行并行化。

例如，根據(jù)推理引擎的選擇，一個(gè)序列可能會(huì)被分塊處理（如分塊預(yù)填充），也可能一次性處理（如果預(yù)填充未被拆分）。為了實(shí)現(xiàn)“批處理不變性”，必須確保某個(gè) token 的規(guī)約順序不依賴于其序列中同時(shí)被處理的其他 token 數(shù)量。如果你將 KV 緩存中的 K/V 值與當(dāng)前正在處理的 token 的 K/V 值分開規(guī)約（如 vLLM 的 Triton attention kernel 所做的那樣），就無法實(shí)現(xiàn)這一點(diǎn)。例如，在處理序列中的第 1000 個(gè)查詢 token 時(shí)，無論 KV 緩存中有 0 個(gè) token（預(yù)填充）還是 999 個(gè) token（解碼），其規(guī)約順序都必須完全一致。

帶 KV 緩存的 FlashAttention 之所以把 KV 緩存與當(dāng)前 KV 值分開顯式處理會(huì)破壞批不變性，原因有些微妙，與“邊界條件”有關(guān)。具體來說，假設(shè)塊大小為 32，而當(dāng)前 KV 緩存中有 80 個(gè)元素。我們?cè)儆?jì)算 48 個(gè)尚未緩存的元素。此時(shí)，需要 3 個(gè)塊（2 個(gè)完整塊 + 1 個(gè)掩碼塊）來計(jì)算 “P cache”，再需要 2 個(gè)塊（1 個(gè)完整塊 + 1 個(gè)掩碼塊）來計(jì)算 “P”。因此總共需要 5 個(gè)塊來完成歸約，而我們總共只有 4 個(gè)塊（即 128 個(gè)元素）需要計(jì)算，這必然會(huì)改變歸約順序。

例如，如果 KV 緩存為空，我們一次性處理 128 個(gè)元素，那么這兩種情況必須得到完全相同的數(shù)值，才能保證 attention 的“批不變性”。

為解決此問題，我們只需在 attention kernel 之前更新 KV 緩存和頁表，確保無論處理多少 token，鍵和值的布局始終一致。

有了這些額外細(xì)節(jié)（以及上一節(jié)提到的所有內(nèi)容，如一致的 tile 大?。?，我們就能實(shí)現(xiàn)一個(gè)不受 batch 影響的 attention 實(shí)現(xiàn)！

然而，這里有一個(gè)顯著的問題。與矩陣乘法不同，我們?cè)?LLM 推理中看到的 attention 形狀通常確實(shí)需要一個(gè) split-reduction 內(nèi)核，通常稱為 Split-KV 或 FlashDecoding。這是因?yàn)槿绻覀儾谎刂?reduction 維度并行化，就只能沿著 batch 維度、head 維度和“query 長(zhǎng)度”維度并行化。在 attention 的 decode 階段，query 長(zhǎng)度非常小，因此除非 batch size 非常大，否則我們通常無法充分利用 GPU。

不幸的是，這次不能像對(duì)待 RMSNorm 和 Matmul 那樣輕易忽略這種情況。例如，如果你有一個(gè)非常長(zhǎng)的 KV cache，即使只處理一個(gè)請(qǐng)求，attention 內(nèi)核也可能需要很長(zhǎng)時(shí)間。

固定 # Split-KV 策略（即 FlashDecode） 如果查詢長(zhǎng)度變得非常小（如在解碼階段），我們可能會(huì)遇到內(nèi)核中幾乎沒有任何并行性的情況。此時(shí)，我們需要再次沿歸約維度——這次是 KV 維度——進(jìn)行切分。沿 KV 維度切分的典型策略是：先確定需要多少并行度，然后將 KV 維度均勻劃分。例如，如果 KV 長(zhǎng)度為 1000 且需要 4 個(gè)分片，每個(gè)核心將處理 250 個(gè)元素。

不幸的是，這也破壞了批不變性，因?yàn)槲覀兊木_歸約策略取決于在任何給定請(qǐng)求中我們要處理序列中的多少個(gè)查詢 token。

此外，常用于注意力的 split-reduction 策略也對(duì)批不變性提出了挑戰(zhàn)。例如，F(xiàn)lashInfer 的“平衡調(diào)度算法”會(huì)選擇仍能飽和所有 GPU 核心的最大 split-size，從而使歸約策略不再是“批不變”的。然而，與 RMSNorm/Matmul 不同，僅固定一個(gè)與 batch size 無關(guān)的 split 數(shù)量是不夠的。

相反，為了實(shí)現(xiàn)批不變性，我們必須采用“固定 split-size”策略。換句話說，我們不再固定 split 的數(shù)量，而是固定每個(gè) split 的大小，從而得到可變的 split 數(shù)量。這樣，無論處理多少 token，我們都能保證始終執(zhí)行完全相同的歸約順序。這需要對(duì) FlexAttention 內(nèi)部做一些修改，這些修改尚未包含在我們的代碼發(fā)布中。我們將在不久的將來將其上游！

固定大小 Split-KV 策略 此策略與前一種策略的唯一區(qū)別在于，我們的拆分現(xiàn)在是“固定大小”的。例如，如果 KV 長(zhǎng)度為 1000，我們不再將其拆分為四個(gè)等長(zhǎng) 250 的片段，而是拆分為三個(gè)固定大小 256 的片段和一個(gè) 232 的片段。

這使得我們能夠保持批次不變性，因?yàn)槲覀兊臍w約策略不再依賴于我們一次性處理的查詢 token 數(shù)量！

實(shí)現(xiàn)

我們通過利用 vLLM 的 FlexAttention 后端以及 torch.Library，在 vLLM 之上實(shí)現(xiàn)了確定性推理的演示。借助 torch.Library，我們能夠以非侵入式的方式替換掉大多數(shù)相關(guān)的 PyTorch 算子。你可以在 thinking-machines-lab/batch-invariant-ops 找到“批不變”內(nèi)核庫，以及以“確定性”模式運(yùn)行 vLLM 的示例。

實(shí)驗(yàn) 補(bǔ)全結(jié)果有多不確定？

我們使用Qwen/Qwen3-235B-A22B-Instruct-2507，在溫度 0 下對(duì)提示“告訴我關(guān)于理查德·費(fèi)曼的事”（非思考模式）采樣 1000 次補(bǔ)全，每次生成 1000 個(gè) token。令人驚訝的是，我們得到了 80 種不同的補(bǔ)全結(jié)果，其中最常見的一種出現(xiàn)了 78 次。

觀察這些補(bǔ)全結(jié)果出現(xiàn)差異的位置，我們發(fā)現(xiàn)前 102 個(gè) token 實(shí)際上完全相同！第一次出現(xiàn)分歧是在第 103 個(gè) token。所有補(bǔ)全都生成了序列“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é)上期望的結(jié)果，但如果沒有批不變內(nèi)核，我們就無法獲得確定性的輸出。

性能

我們尚未對(duì)批不變內(nèi)核的性能進(jìn)行顯著優(yōu)化。不過，讓我們運(yùn)行一些實(shí)驗(yàn)來驗(yàn)證性能是否仍然可用。

我們將使用一塊 GPU 啟動(dòng)一個(gè) API 服務(wù)器，運(yùn)行 Qwen-3-8B，并請(qǐng)求 1000 條序列，輸出長(zhǎng)度在 90 到 110 之間。

配置

時(shí)間（秒）

vLLM 默認(rèn)

26

未優(yōu)化的確定性 vLLM

55

+ 改進(jìn)的注意力內(nèi)核

42

大部分性能下降源于 vLLM 中的 FlexAttention 集成尚未經(jīng)過深度優(yōu)化。盡管如此，性能表現(xiàn)并不算災(zāi)難級(jí)。

真正的 on-policy RL

正如研究人員所指出的，訓(xùn)練與推理之間數(shù)值上的差異，會(huì)隱式地將我們的 on-policy RL 變成 off-policy RL。

當(dāng)然，如果連兩次完全相同的推理請(qǐng)求都無法得到按位一致的結(jié)果，就更不可能在訓(xùn)練與推理之間實(shí)現(xiàn)按位一致。而確定性推理使我們能夠進(jìn)一步改造訓(xùn)練棧，從而在采樣與訓(xùn)練之間獲得按位一致的結(jié)果，最終實(shí)現(xiàn)真正的 on-policy RL。

我們?cè)?Bigmath 上使用 RLVR 設(shè)置進(jìn)行實(shí)驗(yàn)，RL 策略從 Qwen 2.5-VL instruct 8B 初始化，最大 rollout 長(zhǎng)度為 4096。

如果我們訓(xùn)練時(shí)不進(jìn)行 off-policy 校正（即不使用重要性加權(quán)），獎(jiǎng)勵(lì)會(huì)在訓(xùn)練中途崩潰；而加入 off-policy 校正項(xiàng)后，訓(xùn)練可以順利進(jìn)行。然而，如果我們的采樣器和訓(xùn)練器在比特級(jí)別完全一致，我們就完全處于 on-policy（即 KL 散度為 0），同樣可以順利訓(xùn)練。

我們還可以繪制采樣器與訓(xùn)練器之間 logprobs 的 KL 散度，三條曲線表現(xiàn)出明顯不同的行為。使用重要性加權(quán)時(shí)，KL 散度保持在約 0.001，偶爾出現(xiàn)峰值。然而，不使用重要性加權(quán)時(shí)，KL 散度最終會(huì)在獎(jiǎng)勵(lì)崩潰的同一時(shí)間點(diǎn)出現(xiàn)飆升。當(dāng)然，在運(yùn)行“True On-Policy RL”時(shí)，KL 散度始終為 0，表明訓(xùn)練策略與采樣策略之間沒有差異。

請(qǐng)注意，未使用重要性加權(quán)的運(yùn)行在第 318 步左右出現(xiàn)了顯著的損失尖峰，同時(shí) logprobs 的 KL 散度也相應(yīng)飆升。相比之下，無論是采用 off-policy 修正還是“True On-Policy”運(yùn)行，RL 都能平穩(wěn)繼續(xù)。藍(lán)色線顯示的“True On-Policy”并非 bug——它只是 0 處的一條平直線。

結(jié)論

現(xiàn)代軟件系統(tǒng)包含多層抽象。在機(jī)器學(xué)習(xí)中，當(dāng)我們遇到非確定性和細(xì)微數(shù)值差異時(shí)，往往會(huì)傾向于掩蓋它們。畢竟，我們的系統(tǒng)已經(jīng)是“概率性”的，再多一點(diǎn)非確定性又何妨？把失敗單元測(cè)試的 atol/rtol 調(diào)高一點(diǎn)又有什么問題？訓(xùn)練器和采樣器之間 logprobs 的差異大概不是真正的 bug，對(duì)吧？

我們拒絕這種失敗主義。只需稍加努力，我們就能理解非確定性的根本原因，甚至解決它們！我們希望這篇博客文章能為社區(qū)提供如何消除推理系統(tǒng)中非確定性的扎實(shí)理解，并激勵(lì)更多人全面掌握自己的系統(tǒng)。

引用

@article{he2025nondeterminism,   author = {Horace He and Thinking Machines Lab},   title = {Defeating Nondeterminism in LLM Inference},   journal = {Thinking Machines Lab: Connectionism},   year = {2025},   note = {https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/},   doi = {10.64434/tml.20250910} }