編輯：Panda

OpenAI 發論文的頻率是越來越低了。

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

至於原因嘛，讓該公司自家的 ChatGPT 來說吧：“截至目前，OpenAI 在 2025 年在 arXiv 上公開發布的論文數量相對較少，可能反映了其對研究成果公開策略的謹慎態度，可能出於商業保密或安全考慮。”

不過近日，OpenAI 也確實發佈了一份完全由自己人蔘與的、實打實的研究論文，其中提出了一種用於高效張量映射的統一代數框架 Linear Layouts。這是一種使用二元線性代數而非比特表示（bit representation）的張量佈局的通用代數形式，解決了 Triton 等深度學習編譯器中長期存在的難題。

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

簡單來說：張量佈局 = 邏輯張量與硬件資源（例如內存、線程、向量單元）之間的映射關係。下圖給出了兩個佈局示例。

對於現代深度學習工作負載而言，所需要的張量佈局需要滿足幾個要求：

高效（爲了性能）。

靈活（以支持多種算子）。

可組合（爲了變換和優化）。

然而，當前的佈局系統卻難以充分滿足這些需求，而是往往：

需要根據實際需求設計，而且往往是硬編碼的（需要手動編寫規則）。

不可擴展（每一對佈局都需要二次組合）。

容易出錯，尤其是在像 Triton 這樣的低層級的後端中 —— 截至目前，Triton 的 GitHub 庫中提交的 12% 的 Bug 與佈局有關。

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

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

爲此，需要解決一些技術難題：

在將張量映射到硬件資源方面，需要一種通用且可組合的表示方法。

佈局轉換應該用統一的形式來表達，甚至需要包含諸如數據交換（data swizzling）等複雜變換。

這種表示必須與低級硬件優化無縫集成，以確保高效的數據訪問和計算。

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

相關背景知識

GPU 架構

在設計上，現代 GPU 的目標是通過包含多層硬件資源的分層執行模型來充分利用並行性。

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

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

請注意，這些指令要求數據分佈在線程和 Warp 之間，或者以特殊佈局駐留在共享內存或特殊內存單元（例如 Blackwell 上的 Tensor Memory）中，才能產生正確的結果。

然而，這些佈局通常不會爲加載 / 存儲等其他操作帶來最佳性能，而且並非總是可以使用特定指令將數據直接從全局內存複製到特殊內存單元。

因此，通常必須對數據進行重新排列，以便將用於內存訪問的佈局轉換爲計算單元偏好的佈局。

簡而言之，要實現峯值性能，不僅需要利用這些專用單元，還需要精心設計張量佈局和轉換。

Triton 語言和編譯器

Triton 是一種類似於 Python 的用於特定領域的語言，其設計目標是提供用於編寫高性能深度學習原語的靈活接口。Triton 的編譯器後端使用了 MLIR，支持多層次抽象表達。

究其核心，Triton 內核遵循單程序多數據 (SPMD) 模型，其中計算被劃分爲多個抽象的 Triton 程序實例。這種設計允許開發者主要關注 CTA 級別的並行性即可。在 Triton 中，“張量”一詞指的是從原始 PyTorch 張量中提取的塊，它們用作 GPU 核的輸入和輸出。

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

𝔽₂ 數學基礎

我們可將兩個元素 {0, 1} 的域表示爲 𝔽₂。在 𝔽₂ 中，所有算術運算均以 2 爲模執行。

例如，加法定義爲

，其對應於邏輯異或（XOR）。

而乘法定義爲

，對應於邏輯與（AND）。

在 𝔽₂ 上，一個基本運算是矩陣乘法。令

是元素在 𝔽₂ 中的兩個矩陣。乘積 𝐶 = 𝐴𝐵 的各個元素的定義爲

這類似於標準矩陣乘法，不同之處在於所有算術運算都在 𝔽₂ 中執行。

𝔽₂ 中的算術運算與二進制邏輯自然契合，使得該領域的運算在硬件實現中非常高效。因此，𝔽₂ 廣泛應用於密碼學和糾錯碼等領域。

傳統佈局

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

在最高層級，佈局分爲分佈式（Distributed）佈局和內存（（Memory）佈局。前者是指張量元素分佈在不同的執行單元中，而後者是指張量元素存儲在特定的特殊內存中。

分佈式佈局又可進一步分爲 Blocked、Sliced、MMA 和 MMA Input 佈局等類型，而內存佈局又可進一步分爲 Unswizzled 和 Swizzled 佈局。

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

傳統 Triton 佈局系統要求每個佈局定義自己的接口方法，例如每個線程的元素數量和連續元素的數量。此外，必須爲每個佈局顯式實現對張量元素的索引以及佈局之間的轉換。這種方法導致佈局構造和轉換常出現 bug。

Linear Layouts（線性佈局）

下面將簡單介紹線性佈局的定義、一些基本的線性佈局算子、創建各種 Triton 佈局以作爲線性佈局實例，以及應用於 Triton 的通用佈局引擎。

一個示例

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

此外，在 Triton 的編程模型中，張量的維度以及與每個張量相關的佈局子部分（例如每個線程的寄存器和線程數量）都被限制爲 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)。

此外，對於任何偶數線程 𝑡_𝑘，𝑘 的最後一位與 𝑟_0 中 𝑗 的倒數第二位匹配，𝑘 的倒數第二位與 𝑟_0 中 𝑗 的倒數第三位匹配。例如，𝑡_10 = 𝑡_0𝑏1010 的 𝑟_0 位於 (2, 4) = (0𝑏10, 0𝑏100)。這種系統性對齊持續存在，表明二次冪結構足以清晰地決定了每個線程元素的分佈。

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

帶標註的向量空間。如此，可爲佈局中的每一位（bit）分配標籤。輸入 𝑣 位於

空間中，建模了 Reg × Thr × Wrp 的空間。輸出 𝑤 遵循

結構，表示邏輯張量 (𝑖, 𝑗) 的兩個維度。

定義與構造

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

定義 2（組合 / Composition）。給定 𝔽₂ 上的向量空間 𝑈 、𝑉 、𝑊 以及線性佈局 𝐿₁ : 𝑈 → 𝑉 和 𝐿₂ : 𝑉 → 𝑊 ，它們的組合定義爲：

將 𝐿₁ 和 𝐿₂ 表示爲矩陣 𝑀₁ 和 𝑀₂ ，表示 𝐿₂ ◦ 𝐿₁ 的矩陣由 𝔽₂ 上的（逐標籤）矩陣乘法 𝑀₂𝑀₁ 給出。

定義 3（積 / Product）。給定 𝔽₂ 上的兩個向量空間 𝑈 和 𝑉，定義它們的積爲：

給定兩個線性佈局𝐿₁ : 𝑈₁ → 𝑉₁, 𝐿₂ : 𝑈₂ → 𝑉₂，且 𝑢1₁∈ 𝑈₁, 𝑢₂ ∈ 𝑈₂，則定義它們的積爲：

將 𝐿₁ 和 𝐿₂ 表示成矩陣 𝑀₁ 和 𝑀₂，則表示 𝐿₁ × 𝐿₂ 的矩陣爲：

定義 4（左除 / Left Division）。若矩陣 𝑀 具有如下結構，則矩陣 𝑀 左側可被矩陣 𝑀₁ 整除：

這裏將左側的除法記爲

。此運算可在線性佈局中逐標籤處理。

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

定義 5（右逆 / Right Inverse）。在 𝔽₂ 上的滿射線性佈局 𝐿 : 𝑈 → 𝑉 具有一個右逆。

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

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

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

不僅如此，OpenAI 表示，線性佈局爲在語言前端和編譯器後端開發算法提供了結構化的基礎。他們也在論文中給出了一些關鍵示例，這裏就不過多展開。接下來簡單看看新提出的線性佈局的實際表現。

評估

OpenAI 將優化版 Triton（集成了基於線性佈局的優化，即 Triton-Linear）與未集成這些優化的基準 Triton 進行了比較。Triton 和 TritonLinear 之間的主要區別如下：

Triton 使用傳統的數據佈局，不支持任意分佈式佈局的實用程序或它們之間的轉換，因此容易出現 bug。

Triton 未採用論文中描述的優化代碼生成。例如，佈局轉換始終通過共享內存進行，對高效硬件原語的使用有限。

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

爲了比較 Triton 和 Triton-Linear 的性能，該團隊構建了一些合成微基準來進行測試，這方面的結果請訪問原論文查看。這裏僅看看它們在實際基準測試中表現。

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

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

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

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

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

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

總體而言，由於缺乏 ldmatrix 等高效的硬件原語，Triton-Linear 在 AMD GPU 上實現的加速低於在英偉達 GPU 的。

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