OpenAI 發(fā)論文的頻率是越來越低了。

如果你看到了一份來自 OpenAI 的新 PDF 文件，那多半也是新模型的系統(tǒng)卡或相關(guān)增補文件或基準(zhǔn)測試，很少有新的研究論文。

至于原因嘛，讓該公司自家的 ChatGPT 來說吧：「截至目前，OpenAI 在 2025 年在 arXiv 上公開發(fā)布的論文數(shù)量相對較少，可能反映了其對研究成果公開策略的謹(jǐn)慎態(tài)度，可能出于商業(yè)保密或安全考慮?！?/span>

不過近日，OpenAI 也確實發(fā)布了一份完全由自己人參與的、實打?qū)嵉难芯空撐?，其中提出了一種用于高效張量映射的統(tǒng)一代數(shù)框架 Linear Layouts。這是一種使用二元線性代數(shù)而非比特表示（bit representation）的張量布局的通用代數(shù)形式，解決了 Triton 等深度學(xué)習(xí)編譯器中長期存在的難題。

• 論文標(biāo)題：Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using ???
• 論文地址：https://arxiv.org/pdf/2505.23819.pdf

要理解這項研究的意義，首先需要先理解一下什么是張量布局（tensor layouts）。

簡單來說：張量布局 = 邏輯張量與硬件資源（例如內(nèi)存、線程、向量單元）之間的映射關(guān)系。下圖給出了兩個布局示例。

對于現(xiàn)代深度學(xué)習(xí)工作負(fù)載而言，所需要的張量布局需要滿足幾個要求：

• 高效（為了性能）。
• 靈活（以支持多種算子）。
• 可組合（為了變換和優(yōu)化）。

然而，當(dāng)前的布局系統(tǒng)卻難以充分滿足這些需求，而是往往：

• 需要根據(jù)實際需求設(shè)計，而且往往是硬編碼的（需要手動編寫規(guī)則）。
• 不可擴展（每一對布局都需要二次組合）。
• 容易出錯，尤其是在像 Triton 這樣的低層級的后端中 —— 截至目前，Triton 的 GitHub 庫中提交的 12% 的 Bug 與布局有關(guān)。

另外，深度學(xué)習(xí)硬件（如 GPU）的日益復(fù)雜也導(dǎo)致張量布局日益復(fù)雜。

例如，為了實現(xiàn)高效的矩陣乘法，英偉達(dá)在 Ampere、Hopper 和 Blackwell 等不同代際的 GPU 上采用了不同的使用 Tensor Core 的布局，并且每種布局在使用不同數(shù)據(jù)類型時都有不同的變體。AMD 和英特爾等其它 GPU 供應(yīng)商在利用其類似 Tensor Core 的技術(shù)進(jìn)行加速時，也使用了不同的布局。因此，硬件架構(gòu)的快速發(fā)展和多樣化的深度學(xué)習(xí)模型需要一種新的張量布局建模方法。

為此，需要解決一些技術(shù)難題：

1. 在將張量映射到硬件資源方面，需要一種通用且可組合的表示方法。
2. 布局轉(zhuǎn)換應(yīng)該用統(tǒng)一的形式來表達(dá)，甚至需要包含諸如數(shù)據(jù)交換（data swizzling）等復(fù)雜變換。
3. 這種表示必須與低級硬件優(yōu)化無縫集成，以確保高效的數(shù)據(jù)訪問和計算。

不過，在介紹 OpenAI 這篇論文的貢獻(xiàn)之前，我們需要先了解一些基礎(chǔ)概念。

相關(guān)背景知識

GPU 架構(gòu)

在設(shè)計上，現(xiàn)代 GPU 的目標(biāo)是通過包含多層硬件資源的分層執(zhí)行模型來充分利用并行性。

其關(guān)鍵執(zhí)行單元包括協(xié)作線程陣列 (CTA)、Warp 和線程。每個 GPU 線程都可以訪問私有寄存器 —— 這些寄存器提供最低延遲的存儲空間，但容量有限。常規(guī)指令可以由各個線程獨立執(zhí)行。然而，某些特殊功能單元必須在更高的粒度級別上執(zhí)行。

例如，英偉達(dá)的 mma（矩陣乘法累加）指令利用 Tensor Core 的方式是并行執(zhí)行由各個 Warp 發(fā)出的多個乘加運算。而 wgmma（Warp 組矩陣乘法累加）等高級變體則是通過在多個 Warp 上同時執(zhí)行矩陣乘法而對這些功能進(jìn)行了擴展。AMD 也引入了類似的原語，例如 mfma（矩陣融合乘加）指令。

請注意，這些指令要求數(shù)據(jù)分布在線程和 Warp 之間，或者以特殊布局駐留在共享內(nèi)存或特殊內(nèi)存單元（例如 Blackwell 上的 Tensor Memory）中，才能產(chǎn)生正確的結(jié)果。

然而，這些布局通常不會為加載 / 存儲等其他操作帶來最佳性能，而且并非總是可以使用特定指令將數(shù)據(jù)直接從全局內(nèi)存復(fù)制到特殊內(nèi)存單元。

因此，通常必須對數(shù)據(jù)進(jìn)行重新排列，以便將用于內(nèi)存訪問的布局轉(zhuǎn)換為計算單元偏好的布局。

簡而言之，要實現(xiàn)峰值性能，不僅需要利用這些專用單元，還需要精心設(shè)計張量布局和轉(zhuǎn)換。

Triton 語言和編譯器

Triton 是一種類似于 Python 的用于特定領(lǐng)域的語言，其設(shè)計目標(biāo)是提供用于編寫高性能深度學(xué)習(xí)原語的靈活接口。Triton 的編譯器后端使用了 MLIR，支持多層次抽象表達(dá)。

究其核心，Triton 內(nèi)核遵循單程序多數(shù)據(jù) (SPMD) 模型，其中計算被劃分為多個抽象的 Triton 程序?qū)嵗＿@種設(shè)計允許開發(fā)者主要關(guān)注 CTA 級別的并行性即可。在 Triton 中，「張量」一詞指的是從原始 PyTorch 張量中提取的塊，它們用作 GPU 核的輸入和輸出。

在編譯過程中，Triton 的 Python 代碼首先被翻譯成 Triton 方言 (tt)，然后進(jìn)一步翻譯成 TritonGPU 方言 (ttg)。在此過程中，每個張量都與特定的布局相關(guān)聯(lián)，以充分利用現(xiàn)代 GPU 上可用的硬件功能單元。例如，當(dāng)遇到 dot 類算子（例如 tt.dot 和 tt.dot_scaled）時，會采用 mma 布局并使用 Tensor Core 和類似的單元。

??? 數(shù)學(xué)基礎(chǔ)

我們可將兩個元素 {0, 1} 的域表示為 ???。在 ??? 中，所有算術(shù)運算均以 2 為模執(zhí)行。

例如，加法定義為 ，其對應(yīng)于邏輯異或（XOR）。

而乘法定義為，對應(yīng)于邏輯與（AND）。

在 ??? 上，一個基本運算是矩陣乘法。令是元素在 ??? 中的兩個矩陣。乘積 ?? = ???? 的各個元素的定義為

這類似于標(biāo)準(zhǔn)矩陣乘法，不同之處在于所有算術(shù)運算都在 ??? 中執(zhí)行。

??? 中的算術(shù)運算與二進(jìn)制邏輯自然契合，使得該領(lǐng)域的運算在硬件實現(xiàn)中非常高效。因此，??? 廣泛應(yīng)用于密碼學(xué)和糾錯碼等領(lǐng)域。

傳統(tǒng)布局

圖 2 列出了 Triton 中所有可用的布局。

在最高層級，布局分為分布式（Distributed）布局和內(nèi)存（（Memory）布局。前者是指張量元素分布在不同的執(zhí)行單元中，而后者是指張量元素存儲在特定的特殊內(nèi)存中。

分布式布局又可進(jìn)一步分為 Blocked、Sliced、MMA 和 MMA Input 布局等類型，而內(nèi)存布局又可進(jìn)一步分為 Unswizzled 和 Swizzled 布局。

Blocked 布局通常用于連續(xù)的內(nèi)存訪問。MMA 和 MMA 輸入布局用于矩陣乘法運算（例如 tt.dot）的輸出和輸入。MMA 布局可以根據(jù)其映射到的硬件指令進(jìn)一步分類，例如英偉達(dá) GPU 上的 mma 和 wgmma，或 AMD GPU 上的 mfma。Sliced 布局是從其父布局中提取一個維度，用作廣播或某個歸約運算的輸出。

傳統(tǒng) Triton 布局系統(tǒng)要求每個布局定義自己的接口方法，例如每個線程的元素數(shù)量和連續(xù)元素的數(shù)量。此外，必須為每個布局顯式實現(xiàn)對張量元素的索引以及布局之間的轉(zhuǎn)換。這種方法導(dǎo)致布局構(gòu)造和轉(zhuǎn)換常出現(xiàn) bug。

Linear Layouts（線性布局）

下面將簡單介紹線性布局的定義、一些基本的線性布局算子、創(chuàng)建各種 Triton 布局以作為線性布局實例，以及應(yīng)用于 Triton 的通用布局引擎。

一個示例

在 GPU 編程中，大多數(shù)參數(shù)都是 2 的冪：一個 Warp 由 32 或 64 個線程組成，一個 Warp 組包含 4 個 Warp，矩陣乘法內(nèi)聯(lián)函數(shù)（例如 mma 和 wgmma）要求 Tile 尺寸為 16 × ??，其中 ?? ≥ 1。

此外，在 Triton 的編程模型中，張量的維度以及與每個張量相關(guān)的布局子部分（例如每個線程的寄存器和線程數(shù)量）都被限制為 2 的冪。在圖 1 中，布局 A 有一個 16 × 16 的張量，其使用了多個 2 × 2 的寄存器、4 × 8 的線程和 2 × 1 的 Warp。

由于這些量都是 2 的冪，因此使用其坐標(biāo)的比特表示，可以直觀地可視化布局 A 中元素的分布（如圖 1 所示）。所有線程的寄存器 0 (??_0) 都位于坐標(biāo) (??, ??)，其中 ?? 和 ?? 的最后幾位（bit）均為 0。例如，線程 ??_1 的 ??_0 位于 (0, 2) = (0??00, 0??10)。作為對比，??_1 元素的坐標(biāo)中，?? 的最后一位始終為 0，而 ?? 的最后一位始終為 1。例如，??_9 的 ??_1 位于 (2, 3) = (0??10, 0??11)。

此外，對于任何偶數(shù)線程 ??_??，?? 的最后一位與 ??_0 中 ?? 的倒數(shù)第二位匹配，?? 的倒數(shù)第二位與 ??_0 中 ?? 的倒數(shù)第三位匹配。例如，??_10 = ??_0??1010 的 ??_0 位于 (2, 4) = (0??10, 0??100)。這種系統(tǒng)性對齊持續(xù)存在，表明二次冪結(jié)構(gòu)足以清晰地決定了每個線程元素的分布。

綜上所述，假設(shè)一個大小為 8 的向量 ?? 表示一個 Warp 中線程的一個元素，其中前 2 位表示寄存器 (Reg)，接下來的 5 位表示線程 (Thr)，最后一位則表示 Warp (Wrp)，則可以如此定義布局 ??：

帶標(biāo)注的向量空間。如此，可為布局中的每一位（bit）分配標(biāo)簽。輸入 ?? 位于   空間中，建模了 Reg × Thr × Wrp 的空間。輸出 ?? 遵循 結(jié)構(gòu)，表示邏輯張量 (??, ??) 的兩個維度。

定義與構(gòu)造

定義 1（線性布局 / Linear Layouts）。線性布局的定義是在 ??? 上的向量空間之間的線性映射。

定義 2（組合 / Composition）。給定 ??? 上的向量空間 ?? 、?? 、?? 以及線性布局 ??? : ?? → ?? 和 ??? : ?? → ?? ，它們的組合定義為：

將 ??? 和 ??? 表示為矩陣 ??? 和 ??? ，表示 ??? ? ??? 的矩陣由 ??? 上的（逐標(biāo)簽）矩陣乘法 ?????? 給出。

定義 3（積 / Product）。給定 ??? 上的兩個向量空間 ?? 和 ??，定義它們的積為：

給定兩個線性布局??? : ??? → ???, ??? : ??? → ???，且 ??1?∈ ???, ??? ∈ ???，則定義它們的積為：

將 ??? 和 ??? 表示成矩陣 ??? 和 ???，則表示 ??? × ??? 的矩陣為：

定義 4（左除 / Left Division）。若矩陣 ?? 具有如下結(jié)構(gòu)，則矩陣 ?? 左側(cè)可被矩陣 ??? 整除：

這里將左側(cè)的除法記為 。此運算可在線性布局中逐標(biāo)簽處理。

左除法可用于確定布局是否可以分解為滿足高效硬件原語（例如 ldmatrix）的較小布局，

定義 5（右逆 / Right Inverse）。在 ??? 上的滿射線性布局 ?? : ?? → ?? 具有一個右逆。

如果 ?? 是 ?? 的一個矩陣表示，其形狀為?? × ??，則可將???1 定義為 ???? = ??_?? 的 ?? × ?? 最小二乘解，其中 ??_?? 是 ?? × ?? 的單位矩陣。具體來說，它可以通過對???進(jìn)行高斯消元法計算得出。

當(dāng)需要從邏輯張量的坐標(biāo)中恢復(fù)硬件索引時，需要使用求逆運算。

對線性布局的更詳細(xì)完備性說明請訪問原論文，其中涉及到說明分塊布局、mma 和 wgmma 的輸入和輸出布局、線性布局的 slice、每個分布式布局、MMA swizzled 布局、內(nèi)存布局都是線性布局。另外，OpenAI 也在 Triton 說明了如何實現(xiàn)布局轉(zhuǎn)換以及形狀操作。

不僅如此，OpenAI 表示，線性布局為在語言前端和編譯器后端開發(fā)算法提供了結(jié)構(gòu)化的基礎(chǔ)。他們也在論文中給出了一些關(guān)鍵示例，這里就不過多展開。接下來簡單看看新提出的線性布局的實際表現(xiàn)。

評估

OpenAI 將優(yōu)化版 Triton（集成了基于線性布局的優(yōu)化，即 Triton-Linear）與未集成這些優(yōu)化的基準(zhǔn) Triton 進(jìn)行了比較。Triton 和 TritonLinear 之間的主要區(qū)別如下：

• Triton 使用傳統(tǒng)的數(shù)據(jù)布局，不支持任意分布式布局的實用程序或它們之間的轉(zhuǎn)換，因此容易出現(xiàn) bug。
• Triton 未采用論文中描述的優(yōu)化代碼生成。例如，布局轉(zhuǎn)換始終通過共享內(nèi)存進(jìn)行，對高效硬件原語的使用有限。

參與評估的硬件平臺見表 1。

為了比較 Triton 和 Triton-Linear 的性能，該團隊構(gòu)建了一些合成微基準(zhǔn)來進(jìn)行測試，這方面的結(jié)果請訪問原論文查看。這里僅看看它們在實際基準(zhǔn)測試中表現(xiàn)。

在三個不同的平臺上，OpenAI 運行了 TritonBench 中的 18 個基準(zhǔn)測試。圖 7、圖 8 和圖 9 中展示了 Triton-Linear 在三個平臺上的性能提升。

由于每個基準(zhǔn)測試包含多個輸入，總計 420 個案例，因此他們使用了誤差線（error bars）來表示每個基準(zhǔn)測試的最小和最大加速。

需要注意的是，由于硬件限制，并非所有基準(zhǔn)測試都適用于每個平臺。例如，某些基準(zhǔn)測試需要僅在 GH200 上才有的大型共享內(nèi)存，而一些核使用的張量描述符依賴于 TMA 引擎，而 RTX4090 和 MI250 上均不支持 TMA 引擎。

可以看到，在 GH200 上，他們實現(xiàn)了 0.92 倍到 1.57 倍不等的加速，所有基準(zhǔn)測試的平均加速均超過 1.0 倍。加速最顯著的基準(zhǔn)測試是 int4_gemm、ops_gemm 和 streamk_gemm。

可以觀察到，高效的硬件原語（例如 ldmatrix 和 stmatrix）在這些核中被廣泛用于布局轉(zhuǎn)換以及共享內(nèi)存的加載和存儲操作。值得注意的是，layer_norm 實現(xiàn)了從 0.99 倍到 1.57 倍的加速 —— 在不同形狀之間表現(xiàn)出了顯著差異。對于某些輸入形狀，Triton-Linear 能夠檢測「等效」布局之間的轉(zhuǎn)換，從而將轉(zhuǎn)換過程降低為 no-op（無操作）。這種優(yōu)化在舊版布局系統(tǒng)中無法實現(xiàn)，因為它無法直接比較不同類型的布局（例如，Blocked 布局和 Sliced 布局）。

在 RTX4090 上，新方法實現(xiàn)了 1.00 倍到 1.51 倍的加速。由于 mma (RTX4090) 和 wgmma (GH200) 指令之間的差異，他們在 template_attention 上實現(xiàn)了更高的加速。在本例中，tt.dot 運算的左操作數(shù)在循環(huán)外部定義，會重復(fù)從同一地址加載數(shù)據(jù)，因此 ldmatrix 和常規(guī)共享內(nèi)存指令均可實現(xiàn)高吞吐量。雖然右操作數(shù)在每次迭代中都會更新，但 wgmma 會直接在共享內(nèi)存中訪問它，只有在 RTX4090 上，經(jīng)過優(yōu)化后，它才會被降級到 ldmatrix 中。因此，在 GH200 上實現(xiàn)的加速相對較低。在 MI250 上，新方法實現(xiàn)了 0.98 倍到 1.18 倍的加速。

總體而言，由于缺乏 ldmatrix 等高效的硬件原語，Triton-Linear 在 AMD GPU 上實現(xiàn)的加速低于在英偉達(dá) GPU 的。

對于 OpenAI Open 的這個研究，你有什么看法呢？