OpenAI久違發(fā)了篇「正經(jīng)」論文：線性布局實現(xiàn)高效張量計算

2025-06-05 11:49:59

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

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

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

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

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

論文標題：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)代深度學習工作負載而言，所需要的張量布局需要滿足幾個要求：

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

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

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

另外，深度學習硬件（如 GPU）的日益復雜也導致張量布局日益復雜。

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

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

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

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

傳統(tǒng)布局

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

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

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

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

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

Linear Layouts（線性布局）

下面將簡單介紹線性布局的定義、一些基本的線性布局算子、創(chuàng)建各種 Triton 布局以作為線性布局實例，以及應用于 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 的冪，因此使用其坐標的比特表示，可以直觀地可視化布局 A 中元素的分布（如圖 1 所示）。所有線程的寄存器 0 (??_0) 都位于坐標 (??, ??)，其中 ?? 和 ?? 的最后幾位（bit）均為 0。例如，線程 ??_1 的 ??_0 位于 (0, 2) = (0??00, 0??10)。作為對比，??_1 元素的坐標中，?? 的最后一位始終為 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)足以清晰地決定了每個線程元素的分布。

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

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

定義與構(gòu)造

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

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

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

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

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

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

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

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

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

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

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

當需要從邏輯張量的坐標中恢復硬件索引時，需要使用求逆運算。

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

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

評估

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

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

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

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

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

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

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

可以看到，在 GH200 上，他們實現(xiàn)了 0.92 倍到 1.57 倍不等的加速，所有基準測試的平均加速均超過 1.0 倍。加速最顯著的基準測試是 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)外部定義，會重復從同一地址加載數(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)的加速低于在英偉達 GPU 的。

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

責任編輯：張燕妮來源：機器之心

OpenAI 深度學習模型

OpenAI久違發(fā)了篇「正經(jīng)」論文：線性布局實現(xiàn)高效張量計算

相關(guān)背景知識

傳統(tǒng)布局

Linear Layouts（線性布局）

評估