OpenAI久违发了篇「正经」论文：线性布局实现高效张量计算

OpenAI久违发了篇「正经」论文：线性布局实现高效张量计算

仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接

来源：机器之心
编辑：Panda
OpenAI发论文的频率是越来越低了。
如果你看到了一份来自OpenAI的新PDF文件，那多半也是新模型的系统卡或相关增补文件或基准测试，很少有新的研究论文。
至于原因嘛，让该公司自家的ChatGPT来说吧：「截至目前，OpenAI在2025年在arXiv上公开发布的论文数量相对较少，可能反映了其对研究成果公开策略的谨慎态度，可能出于商业保密或安全考虑。」
不过近日，OpenAI也确实发布了一份完全由自己人参与的、实打实的研究论文，其中提出了一种用于高效张量映射的统一代数框架LinearLayouts。这是一种使用二元线性代数而非比特表示（bitrepresentation）的张量布局的通用代数形式，解决了Triton等深度学习编译器中长期存在的难题。
论文标题：LinearLayouts:RobustCodeGenerationofEfficientTensorComputationUsing𝔽₂
论文地址：
https ://arxiv. org/pdf/2505. 23819.pdf
要理解这项研究的意义，首先需要先理解一下什么是张量布局（tensorlayouts）。
简单来说：张量布局=逻辑张量与硬件资源（例如内存、线程、向量单元）之间的映射关系。下图给出了两个布局示例。
对于现代深度学习工作负载而言，所需要的张量布局需要满足几个要求：
高效（为了性能）。
灵活（以支持多种算子）。
可组合（为了变换和优化）。
然而，当前的布局系统却难以充分满足这些需求，而是往往：
需要根据实际需求设计，而且往往是硬编码的（需要手动编写规则）。
不可扩展（每一对布局都需要二次组合）。
容易出错，尤其是在像Triton这样的低层级的后端中——截至目前，Triton的GitHub库中提交的12%的Bug与布局有关。
另外，深度学习硬件（如GPU）的日益复杂也导致张量布局日益复杂。
例如，为了实现高效的矩阵乘法，英伟达在Ampere、Hopper和Blackwell等不同代际的GPU上采用了不同的使用TensorCore的布局，并且每种布局在使用不同数据类型时都有不同的变体。AMD和英特尔等其它GPU供应商在利用其类似TensorCore的技术进行加速时，也使用了不同的布局。因此，硬件架构的快速发展和多样化的深度学习模型需要一种新的张量布局建模方法。
为此，需要解决一些技术难题：
在将张量映射到硬件资源方面，需要一种通用且可组合的表示方法。
布局转换应该用统一的形式来表达，甚至需要包含诸如数据交换（dataswizzling）等复杂变换。
这种表示必须与低级硬件优化无缝集成，以确保高效的数据访问和计算。
不过，在介绍OpenAI这篇论文的贡献之前，我们需要先了解一些基础概念。
相关背景知识
GPU架构
在设计上，现代GPU的目标是通过包含多层硬件资源的分层执行模型来充分利用并行性。
其关键执行单元包括协作线程阵列(CTA)、Warp和线程。每个GPU线程都可以访问私有寄存器——这些寄存器提供最低延迟的存储空间，但容量有限。常规指令可以由各个线程独立执行。然而，某些特殊功能单元必须在更高的粒度级别上执行。
例如，英伟达的mma（矩阵乘法累加）指令利用TensorCore的方式是并行执行由各个Warp发出的多个乘加运算。而wgmma（Warp组矩阵乘法累加）等高级变体则是通过在多个Warp上同时执行矩阵乘法而对这些功能进行了扩展。AMD也引入了类似的原语，例如mfma（矩阵融合乘加）指令。
请注意，这些指令要求数据分布在线程和Warp之间，或者以特殊布局驻留在共享内存或特殊内存单元（例如Blackwell上的TensorMemory）中，才能产生正确的结果。
然而，这些布局通常不会为加载/存储等其他操作带来最佳性能，而且并非总是可以使用特定指令将数据直接从全局内存复制到特殊内存单元。
因此，通常必须对数据进行重新排列，以便将用于内存访问的布局转换为计算单元偏好的布局。
简而言之，要实现峰值性能，不仅需要利用这些专用单元，还需要精心设计张量布局和转换。
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布局并使用TensorCore和类似的单元。
𝔽₂数学基础
我们可将两个元素{0,1}的域表示为𝔽₂。在𝔽₂中，所有算术运算均以2为模执行。
例如，加法定义为，其对应于逻辑异或（XOR）。
而乘法定义为，对应于逻辑与（AND）。
在𝔽₂上，一个基本运算是矩阵乘法。令是元素在𝔽₂中的两个矩阵。乘积𝐶=𝐴𝐵的各个元素的定义为
这类似于标准矩阵乘法，不同之处在于所有算术运算都在𝔽₂中执行。
𝔽₂中的算术运算与二进制逻辑自然契合，使得该领域的运算在硬件实现中非常高效。因此，𝔽₂广泛应用于密码学和纠错码等领域。
传统布局
图2列出了Triton中所有可用的布局。
在最高层级，布局分为分布式（Distributed）布局和内存（（Memory）布局。前者是指张量元素分布在不同的执行单元中，而后者是指张量元素存储在特定的特殊内存中。
分布式布局又可进一步分为Blocked、Sliced、MMA和MMAInput布局等类型，而内存布局又可进一步分为Unswizzled和Swizzled布局。
Blocked布局通常用于连续的内存访问。MMA和MMA输入布局用于矩阵乘法运算（例如tt. dot）的输出和输入。MMA布局可以根据其映射到的硬件指令进一步分类，例如英伟达GPU上的mma和wgmma，或AMDGPU上的mfma。Sliced布局是从其父布局中提取一个维度，用作广播或某个归约运算的输出。
传统Triton布局系统要求每个布局定义自己的接口方法，例如每个线程的元素数量和连续元素的数量。此外，必须为每个布局显式实现对张量元素的索引以及布局之间的转换。这种方法导致布局构造和转换常出现bug。
LinearLayouts（线性布局）
下面将简单介绍线性布局的定义、一些基本的线性布局算子、创建各种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（线性布局/LinearLayouts）。线性布局的定义是在𝔽₂上的向量空间之间的线性映射。
定义2（组合/Composition）。给定𝔽₂上的向量空间𝑈、𝑉、𝑊以及线性布局𝐿₁:𝑈→𝑉和𝐿₂:𝑉→𝑊，它们的组合定义为：
将𝐿₁和𝐿₂表示为矩阵𝑀₁和𝑀₂，表示𝐿₂◦𝐿₁的矩阵由𝔽₂上的（逐标签）矩阵乘法𝑀₂𝑀₁给出。
定义3（积/Product）。给定𝔽₂上的两个向量空间𝑈和𝑉，定义它们的积为：
给定两个线性布局𝐿₁:𝑈₁→𝑉₁,𝐿₂:𝑈₂→𝑉₂，且𝑢1₁∈𝑈₁,𝑢₂∈𝑈₂，则定义它们的积为：
将𝐿₁和𝐿₂表示成矩阵𝑀₁和𝑀₂，则表示𝐿₁×𝐿₂的矩阵为：
定义4（左除/LeftDivision）。若矩阵𝑀具有如下结构，则矩阵𝑀左侧可被矩阵𝑀₁整除：
这里将左侧的除法记为。此运算可在线性布局中逐标签处理。
左除法可用于确定布局是否可以分解为满足高效硬件原语（例如ldmatrix）的较小布局，
定义5（右逆/RightInverse）。在𝔽₂上的满射线性布局𝐿:𝑈→𝑉具有一个右逆。
如果𝑀是𝐿的一个矩阵表示，其形状为𝑚×𝑛，则可将𝑀−1定义为𝑀𝑋=𝐼_𝑚的𝑛×𝑚最小二乘解，其中𝐼_𝑚是𝑚×𝑚的单位矩阵。具体来说，它可以通过对𝔽₂进行高斯消元法计算得出。
当需要从逻辑张量的坐标中恢复硬件索引时，需要使用求逆运算。
对线性布局的更详细完备性说明请访问原论文，其中涉及到说明分块布局、mma和wgmma的输入和输出布局、线性布局的slice、每个分布式布局、MMAswizzled布局、内存布局都是线性布局。另外，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个案例，因此他们使用了误差线（errorbars）来表示每个基准测试的最小和最大加速。
需要注意的是，由于硬件限制，并非所有基准测试都适用于每个平台。例如，某些基准测试需要仅在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在AMDGPU上实现的加速低于在英伟达GPU的。
对于OpenAIOpen的这个研究，你有什么看法呢？
阅读最新前沿科技趋势报告，请访问欧米伽研究所的“未来知识库”
https ://wx. zsxq.com/group/454854145828
未来知识库是“欧米伽未来研究所”建立的在线知识库平台，收藏的资料范围包括人工智能、脑科学、互联网、超级智能，数智大脑、能源、军事、经济、人类风险等等领域的前沿进展与未来趋势。目前拥有超过8000篇重要资料。每周更新不少于100篇世界范围最新研究资料。欢迎扫描二维码或访问https ://wx. zsxq.com/group/454854145828进入。
截止到3月31日”未来知识库”精选的百部前沿科技趋势报告
（加入未来知识库，全部资料免费阅读和下载）
牛津未来研究院《将人工智能安全视为全球公共产品的影响、挑战与研究重点》
麦肯锡：超级智能机构：赋能人们释放人工智能的全部潜力
AAAI2025关于人工智能研究未来研究报告
斯坦福：2025斯坦福新兴技术评论：十项关键技术及其政策影响分析报告（191页）
壳牌：2025能源安全远景报告：能源与人工智能（57页）
盖洛普&牛津幸福研究中心：2025年世界幸福报告（260页）
Schwab：2025未来共生：以集体社会创新破解重大社会挑战研究报告（36页）
IMD：2024年全球数字竞争力排名报告：跨越数字鸿沟人才培养与数字法治是关键（214页）
DS系列专题：DeepSeek技术溯源及前沿探索，50页ppt
联合国人居署：2024全球城市负责任人工智能评估报告：利用AI构建以人为本的智慧城市（86页）
TechUK：2025全球复杂多变背景下的英国科技产业：战略韧性与增长路径研究报告（52页）
NAVEXGlobal：2024年十大风险与合规趋势报告（42页）
《具身物理交互在机器人-机器人及机器人-人协作中的应用》122页
2025-2035年人形机器人发展趋势报告53页
EvaluatePharma：2024年全球生物制药行业展望报告：增长驱动力分析（29页）
【AAAI2025教程】基础模型与具身智能体的交汇，350页ppt
Tracxn：2025全球飞行汽车行业市场研究报告（45页）
谷歌：2024人工智能短跑选手（AISprinters）：捕捉新兴市场AI经济机遇报告（39页）
【斯坦福博士论文】构建类人化具身智能体：从人类行为中学习
《基于传感器的机器学习车辆分类》最新170页
美国安全与新兴技术中心：2025CSET对美国人工智能行动计划的建议（18页）
罗兰贝格：2024人形机器人的崛起：从科幻到现实：如何参与潜在变革研究报告（11页）
兰德公司：2025从研究到现实：NHS的研究和创新是实现十年计划的关键报告（209页）
康桥汇世（CambridgeAssociates）：2025年全球经济展望报告（44页）
国际能源署：2025迈向核能新时代
麦肯锡：人工智能现状，组织如何重塑自身以获取价值
威立（Wiley）：2025全球科研人员人工智能研究报告（38页）
牛津经济研究院：2025TikTok对美国就业的量化影响研究报告：470万岗位（14页）
国际能源署（IEA）：能效2024研究报告（127页）
Workday：2025发挥人类潜能：人工智能（AI）技能革命研究报告（20页）
CertiK：Hack3D：2024年Web3. 0安全报告（28页）
世界经济论坛：工业制造中的前沿技术：人工智能代理的崛起》报告
迈向推理时代：大型语言模型的长链推理研究综述
波士顿咨询：2025亚太地区生成式AI的崛起研究报告：从技术追赶者到全球领导者的跨越（15页）
安联（Allianz）：2025新势力崛起：全球芯片战争与半导体产业格局重构研究报告（33页）
IMT：2025具身智能（EmbodiedAI）概念、核心要素及未来进展：趋势与挑战研究报告（25页）
IEEE：2025具身智能（EmbodiedAI）综述：从模拟器到研究任务的调查分析报告（15页）
CCAV：2025当AI接管方向盘：自动驾驶场景下的人机交互认知重构、变革及对策研究报告（124页）
《强化学习自我博弈方法在兵棋推演分析与开发中的应用》最新132页
《面向科学发现的智能体人工智能：进展、挑战与未来方向综述》
全国机器人标准化技术委员会：人形机器人标准化白皮书（2024版）（96页）
美国国家科学委员会（NSB）：2024年研究与发展-美国趋势及国际比较（51页）
艾昆纬（IQVIA）：2025骨科手术机器人技术的崛起白皮书：创新及未来方向（17页）
NPL&Beauhurst：2025英国量子产业洞察报告：私人和公共投资的作用（25页）
IEAPVPS：2024光伏系统经济与技术关键绩效指标（KPI）使用最佳实践指南（65页）
AGI智能时代：2025让DeepSeek更有趣更有深度的思考研究分析报告（24页）
2025军事领域人工智能应用场景、国内外军事人工智能发展现状及未来趋势分析报告（37页）
华为：2025鸿蒙生态应用开发白皮书（133页
《超级智能战略研究报告》
中美技术差距分析报告2025
欧洲量子产业联盟（QuIC）：2024年全球量子技术专利态势分析白皮书（34页）
美国能源部：2021超级高铁技术（Hyperloop）对电网和交通能源的影响研究报告（60页）
罗马大学：2025超级高铁（Hyperloop）：第五种新型交通方式-技术研发进展、优势及局限性研究报告（72页）
兰德公司：2025灾难性网络风险保险研究报告：市场趋势与政策选择（93页）
GTI：2024先进感知技术白皮书（36页）
AAAI：2025人工智能研究的未来报告：17大关键议题（88页）
安联Allianz2025新势力崛起全球芯片战争与半导体产业格局重构研究报告
威达信：2025全球洪水风险研究报告：现状、趋势及应对措施（22页）
兰德公司：迈向人工智能治理研究报告：2024EqualAI峰会洞察及建议（19页）
哈佛商业评论：2025人工智能时代下的现代软件开发实践报告（12页）
德安华：全球航空航天、国防及政府服务研究报告：2024年回顾及2025年展望（27页）
奥雅纳：2024塑造超级高铁（Hyperloop）的未来：监管如何推动发展与创新研究报告（28页）
HSOAC：2025美国新兴技术与风险评估报告：太空领域和关键基础设施（24页）
Dealroom：2025欧洲经济与科技创新发展态势、挑战及策略研究报告（76页）
《无人机辅助的天空地一体化网络：学习算法技术综述》
谷歌云（GoogleCloud）：2025年AI商业趋势白皮书（49页）
《新兴技术与风险分析：太空领域与关键基础设施》最新报告
150页！《DeepSeek大模型生态报告》
军事人工智能行业研究报告：技术奇点驱动应用加速智能化重塑现代战争形态-250309（40页）
真格基金：2024美国独角兽观察报告（56页）
璞跃（PlugandPlay）：2025未来商业研究报告：六大趋势分析（67页）
国际电工委员会（IEC）：2025智能水电技术与市场展望报告（90页）
RWS：2025智驭AI冲击波：人机协作的未来研究报告（39页）
国际电工委员会（IEC）：2025智能水电技术与市场展望报告（90页）
RWS：2025智驭AI冲击波：人机协作的未来研究报告（39页）
未来今日研究所2025年科技趋势报告第18版1000页
模拟真实世界：多模态生成模型的统一综述
中国信息协会低空经济分会：低空经济发展报告（2024-2025）（117页）
浙江大学：2025语言解码双生花：人类经验与AI算法的镜像之旅（42页）
人形机器人行业：由“外”到“内”智能革命-250306（51页）
大成：2025年全球人工智能趋势报告：关键法律问题（28页）
北京大学：2025年DeepSeek原理和落地应用报告（57页）
欧盟委员会人工智能与未来工作研究报告
加州大学伯克利分校：面向科学发现的多模态基础模型：在化学、材料和生物学中的应用
电子行业：从柔性传感到人形机器人触觉革命-250226（35页）
RT轨道交通：2024年中国城市轨道交通市场数据报告（188页）
FastMoss：2024年度TikTok生态发展白皮书（122页）
CheckPoint：2025年网络安全报告-主要威胁、新兴趋势和CISO建议（57页）
【AAAI2025教程】评估大型语言模型：挑战与方法，199页ppt
《21世纪美国的主导地位：核聚变》最新报告
沃尔特基金会（VoltaFoundation）：2024年全球电池行业年度报告（518页）
斯坦福：2025斯坦福新兴技术评论：十项关键技术及其政策影响分析报告（191页）
国际科学理事会：2025为人工智能做好国家研究生态系统的准备-2025年战略与进展报告（英文版）（118页）
光子盒：2025全球量子计算产业发展展望报告（184页）
奥纬论坛：2025塑造未来的城市研究报告：全球1500个城市的商业吸引力指数排名（124页）
FutureMatters：2024新兴技术与经济韧性：日本未来发展路径前瞻报告（17页）
《人类与人工智能协作的科学与艺术》284页博士论文
《论多智能体决策的复杂性：从博弈学习到部分监控》115页
《2025年技术展望》56页slides
大语言模型在多智能体自动驾驶系统中的应用：近期进展综述
【牛津大学博士论文】不确定性量化与因果考量在非策略决策制定中的应用
皮尤研究中心：2024美国民众对气候变化及应对政策的态度调研报告：气候政策对美国经济影响的多元观点审视（28页）
空间计算行业深度：发展趋势、关键技术、行业应用及相关公司深度梳理-250224（33页）
Gartner：2025网络安全中的AI：明确战略方向研究报告（16页）
北京大学：2025年DeepSeek系列报告-提示词工程和落地场景（86页）
北京大学：2025年DeepSeek系列报告-DeepSeek与AIGC应用（99页）
CIC工信安全：2024全球人工智能立法的主要模式、各国实践及发展趋势研究报告（42页）
中科闻歌：2025年人工智能技术发展与应用探索报告（61页）
AGI智能时代：2025年Grok-3大模型：技术突破与未来展望报告（28页）
上下滑动查看更多