Meta 发布 Triton 低级扩展 TLX，以 MIMW 模型优化生产级 GPU 内核

Meta 工程博客·2026-11-11 00:59· 1 小时前AI 评分54

我们如何构建 TLX：面向大规模生产 ML 的硬件原生 GPU 编译器

AI 导读

Meta 推出 Triton 低级扩展 TLX，通过显式 warp 组编排、同步和内存编排，在保留 Triton 编程模型的同时掌控现代 GPU 执行结构。

正文 · 原文

• Triton Low-Level Extensions (TLX) 是 Triton 的扩展，帮助工程师在保留 Triton 高效编程模型的同时优化 GPU 性能。 
• TLX 引入了显式的 warp 组和集群级编排，让开发者能够在不牺牲开发速度的情况下采用新的硬件能力。 
• TLX 已在 RecSys 训练、RecSys 推理和大规模 LLM 预训练中展现出性能提升。

现代 GPU 架构越来越依赖面向执行的特性（如 warp 特化、异步内存移动、张量内存和集群级执行）来优化性能。但这些特性也暴露出用于控制的低级编程模型与用于生产力的高级编程模型之间日益扩大的差距。为弥合这一差距，我们构建了 Triton Low-Level Extensions (TLX)，这一扩展 让工程师能够显式控制 GPU 执行，同时保留 Triton 高效的编程模型。 

如今，TLX 支撑着 Meta 的各项生产工作负载，从 AI 训练到大规模 LLM 预训练，使工程师能够在不牺牲开发速度的情况下快速采用新的硬件能力。在 NVIDIA 和 AMD GPU 上，TLX 正在推动 GPU 性能的前沿——将最先进的 kernel 与显式执行编排和模型感知融合相结合，带来超越独立库所能实现的端到端改进。 

从 AI 模型到 GPU Kernel

要理解 TLX 的定位，不妨梳理从 AI 模型到最终执行的 GPU 指令的路径。在 Meta，机器学习（ML）工作负载构建在以 PyTorch 为中心的通用软件栈之上。模型代码经过图捕获、优化、编译和 lowering 转换，最终成为在大型加速器集群上运行的 GPU kernel。过去几年间，Triton 已从实验性 kernel 语言演变为该软件栈的基础组件。相当一部分生产 GPU kernel——无论是编译器生成还是手工优化——现在都用 Triton 实现。因此，这一层的改进可以广泛传播到训练和推理工作负载中，使 kernel 性能成为提升端到端系统效率日益重要的杠杆。

这一转变也改变了内核的编写者。内核优化不再是 GPU 专家的专属领域。ML 工程师会例行分析工作负载、识别瓶颈，并通过自动调优、融合、内存布局改进或算法重构来优化最热点的内核。Triton 的高效性让许多此类迭代能在数小时内完成，而非数天，使工程师能够贴近模型开发流程。因此，GPU 内核开发已日益成为 ML 工程师标准优化工具箱的一部分。

Triton 已通过融合和编译器管理的内存访问优化，为许多内存受限的工作负载提供了高效的基础。在这些优势之上，TLX 解决了一个互补的挑战：在现代 GPU 上协调执行组、异步数据移动和共享资源。

问题所在：当硬件演进快于编程模型

尽管 Triton 大幅提升了开发者生产力，现代 GPU 架构的演进却越来越强调执行编排，而非单纯的计算。NVIDIA 的 Hopper 引入了 warp 特化和张量内存加速器（TMA）。NVIDIA 的 Blackwell GPU 通过张量内存（TMEM）、集群级执行、配对 CTA 执行以及日益异步的执行机制进一步扩展了这一趋势。这些能力并非可以隐藏在编译器优化背后的新指令，它们从根本上改变了高性能内核的组织方式。

例如，一个现代 FlashAttention 内核可能专门分配一个 warp 组来发起 TMA 传输，而另一个 warp 组执行张量核心计算。一个 Blackwell GEMM 内核可能同时协调多个 warp 组、TMEM 分配和集群资源。在这两种情况下，性能不仅取决于所执行的计算，还取决于多个执行代理如何协作。 

当我们在 Meta 的 ML 基础设施中采用 Hopper 和 Blackwell 时，我们反复遇到同一个挑战：CUDA 可以表达这些执行结构，但这样做通常需要大量架构特定的编排代码。Triton 提供了生产力显著更高的编程模型，但现代硬件所需的许多执行关系仍隐藏在编译器抽象之后。

结果是一个日益严重的编译器追赶问题。每一代新硬件都会引入新的执行能力，而程序员要么等待编译器支持到位，要么退回到日益专用的实现中。我们意识到，挑战不仅仅是缺少编译器特性，而是缺少合适的编程抽象。

缺失的中间层：为什么 SIMT 和 SIMB 还不够

近期 GPU 上最关键的能力既不在线程粒度上运行，也不在块粒度上运行。Warp 特化将工作划分到相互协作的 warp 组之间。TMA 引入了专门负责内存移动的代理。TMEM 使执行组之间能够采用新的通信模式。集群执行让多个 CTA 在共享的执行上下文中协作。围绕这些特性来组织内核，意味着要推理多个执行代理如何协作、通信和同步——这是一种介于 CUDA 和 Triton 之间的结构层次。

这暴露了我们后来视为缺失抽象的问题。CUDA 的 SIMT 模型运行在过于精细的粒度上：它提供完全的控制，但要求开发者手动表达每一种执行关系。Triton 的分块编程模型运行在过于粗放的粒度上：它提升了生产力，却把许多执行结构变成了编译器的隐式决策。现代硬件越来越需要介于两者之间的东西。

我们将这一模型称为 多指令多 Warp（MIMW）。MIMW 将 warp 组提升为编排的基本单元。程序员不再需要推理单个线程或整个块，而是表达多个执行代理如何协作、通信、同步和共享资源。这正是许多 Hopper 和 Blackwell 优化所表达的自然层次，也是 TLX 所实现的抽象层。

图 2 展示了 GPU 编程模型的演进。CUDA 的 SIMT 模型直接在线程粒度上暴露执行，让开发者对调度、同步和内存移动拥有完全控制。Triton 走向 SIMB 模型，将抽象提升到分块张量计算并提高了生产力。MIMW 占据这两个极端之间的空间。它不再将线程或张量块视为主要执行抽象，而是将 warp 组提升为一等执行代理。

MIMW 的出现不是软件设计偏好，而是硬件演进方式的结果。Volta 和 Ampere 通过线程级和 warp 级编程暴露张量核心。Hopper 使异步数据移动（由 TMA 实现）和 warp 组执行（由 WGMMA 实现）成为高性能内核的核心。Blackwell 通过 TMEM、寄存器重分配、集群执行和成对 CTA MMA 进一步延续了这一趋势。每一代都将更多跨 warp 组——有时是跨 CTA——的协作推入必须由编程模型、而不仅仅是编译器来处理的领域。

要让 MIMW 实用化，需要将协作的 warp 组或 wave 组表达为一等执行任务，并对其间的共享内存缓冲区访问、同步和交接进行显式控制。TLX 将这些能力引入 Triton 现有的张量编程模型，让开发者能够将显式执行控制与熟悉的高层张量操作结合起来。

推出 TLX：缺失的中间层，已实现

我们希望避免用 TLX 再造一个独立的底层语言。相反，TLX 通过第二层面向执行的构造来扩展 Triton。程序员继续使用熟悉的 Triton 张量操作来表达计算，而 TLX 仅在需要对执行、内存移动或同步进行显式控制时使用。换句话说，Triton 回答的是应该执行什么计算，而 TLX 回答的是如何编排该计算。通过一个简单的分块矩阵乘法最容易看出这种差异。在标准 Triton 中，执行遵循同步的加载-计算循环：每个分块被加载、使用，然后在下一个分块开始前被丢弃。

# Standard Triton: synchronous load-compute loop semantics, 
# with compiler-managed software pipelining  
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptr + ...)   # Load A tile from global memory
        b = tl.load(b_ptr + ...)   # Load B tile from global memory
        acc += tl.dot(a, b)        # Compute matmul 
    tl.store(c_ptr + ..., acc)

TLX 保留相同的计算，但使执行结构显式化。生产者 warp 组异步移动数据，消费者 warp 组执行张量核心计算，同步通过轻量级屏障直接表达。计算本身仍然是 Triton；只有编排层发生了变化。结果是一个深度流水线化的内核，在不离开 Python 或不牺牲 Triton 自动调优基础设施的情况下，重叠内存移动、计算和输出处理。更重要的是，执行结构在源代码中可见，而不是隐藏在编译器启发式规则中。

# TLX on Blackwell: explicit warp-group orchestration and pipelining 
@triton.jit
def matmul_kernel(desc_a, desc_b, desc_c, M, N, K, ...):
    buffers_a = tlx.local_alloc((BLOCK_M, BLOCK_K), tl.float16, NUM_STAGES)
    buffers_b = tlx.local_alloc((BLOCK_K, BLOCK_N), tl.float16, NUM_STAGES)
    acc = tlx.local_alloc((BLOCK_M, BLOCK_N), tl.float32, 1, tlx.storage_kind.tmem)
    load_bars = tlx.alloc_barriers(num_barriers=NUM_STAGES, arrive_count=1)
    dot_bars  = tlx.alloc_barriers(num_barriers=NUM_STAGES, arrive_count=1)
    with tlx.async_tasks():
        # Warp group 0: PRODUCER — async TMA loads into shared memory
        with tlx.async_task("default"):
             # Wait for last matmul
            tlx.barrier_wait(dot_bars[num_k_tiles - 1], phase)  
            result = tlx.local_load(acc[0])
            desc_c.store(offsets, result.to(tl.float16))
            
        # Warp group 1: CONSUMER — tensor-core matmul from shared memory
        with tlx.async_task(num_warps=1):
            for k in range(num_k_tiles):
                buf = k % NUM_STAGES
                 # Wait for TMA load to complete
                tlx.barrier_wait(load_bars[buf], phase)      
                tlx.async_dot(buffers_a[buf], buffers_b[buf], acc[0],
                              use_acc=(k > 0), mBarriers=[dot_bars[buf]])
         
        # Warp group 2: EPILOGUE — store result to global memory
        with tlx.async_task(num_warps=1):
            for k in range(num_k_tiles):
                buf = k % NUM_STAGES
                 # Wait for consumer to release buffer
                tlx.barrier_wait(dot_bars[buf], phase)       
                tlx.barrier_expect_bytes(load_bars[buf], ...)
                tlx.async_descriptor_load(desc_a, buffers_a[buf], offsets, load_bars[buf])
                tlx.async_descriptor_load(desc_b, buffers_b[buf], offsets, load_bars[buf])

TLX 在五个关键领域扩展了 Triton： 

Warp 组编排： TLX 将 warp 组视为一等执行代理。程序员可以为协作的 warp 组分配显式角色——例如，移动数据的数据生产者 warp 组、执行张量核心计算的消费者 warp 组，以及处理结果的收尾 warp 组。 

同步： TLX 在 warp 组级别而非线程级别暴露同步。程序员使用 mbarrier 等硬件原语来表达协作执行组之间的依赖关系——例如软件流水线中的生产者-消费者关系。细粒度的线程级同步仍由底层执行模型和编译器处理，让开发者专注于协调越来越决定现代 GPU 性能的更大执行结构。

内存编排： TLX 提供对 TMA、异步张量核心操作和 TMEM 等硬件能力的直接访问。这些原语允许内存移动和计算在程序定义的执行结构内自然重叠，同时避免 CUDA 中通常需要的大量底层设置和簿记工作。

集群执行： TLX 将集群级执行作为一等编程构造暴露出来。多个 CTA 可以通过分布式共享内存和集群级同步原语进行协作，从而实现原本需要经过全局内存往返的片上通信模式。

可选布局：与 CUDA 和 CUTLASS 不同，TLX 不要求程序员在表达计算之前显式指定布局。编译器会自动推断合理的默认布局，程序员仅在出于性能原因想要覆盖默认布局时才指定布局。因此，即使完全不写任何布局注解，也能编写出正确的 kernel，布局调优可以在优化过程中逐步引入。

一个重要结果是，kernel 融合变得显著更容易。融合是最有效的 GPU 优化之一，因为它消除了中间内存流量和 kernel 启动开销。然而，在传统的 CUDA 风格工作流中，融合往往被布局管理所主导。每个阶段可能假定不同的分块、swizzling 或共享内存映射，迫使程序员在操作边界之间手动协调布局。 

由于 TLX 会自动推断布局，操作通常可以在不手动协调中间布局的情况下进行融合。编译器会尽可能解决布局兼容性问题，程序员仅在特定布局选择能提升性能时才介入。这显著降低了构建复杂融合 kernel 所需的工程资源，并使得原本需要多个 kernel 或大量布局专用基础设施的多阶段流水线成为可能。

综合来看，这些能力使 TLX 占据了 Triton 的生产力与 CUDA 的控制力之间的空间。开发者既保留了 Python、自动调优和分块张量编程的优势，又能直接访问那些在现代 GPU 上日益决定性能的执行结构。

TLX 在生产环境中：真实世界影响

过去几个月，我们一直在 Meta 的 ML 基础设施中生产运行 TLX。三个部署——涵盖 RecSys 训练、RecSys 推理和大规模 LLM 预训练——凸显了 TLX 所赋能的不同方面。

案例研究 1：GEM 模型训练

训练 Meta 的 生成式广告模型（GEM） 比 LLM 训练更难优化，因为更小的输入形状、每层 10 多种 norm、异构的模块类型以及架构多样性都限制了 MFU。 通过系统性的模型-内核协同设计方法，涵盖模型架构、自定义内核、编译器优化和硬件特定调优，最新版本的 GEM 训练在数千块 NVIDIA Blackwell GPU 上达到了 O(10^24) 的总训练 flops，端到端 MFU 为 20-25%。

TLX 在多 DSL 策略中的角色 

我们在 Blackwell 上评估了三种编程模型——CUTLASS、CuTeDSL 和 TLX——发现 TLX 在大多数内核上提供了生产力与性能之间的最佳平衡。它结合了类似 Triton 的简洁性、自动调优和可组合性，并可直接访问 Blackwell 特有的能力，例如 async_dot、张量内存操作和 mbarrier 同步。

对于 GEM 中的内核广义点积注意力（GDPA），TLX 以极少的实现复杂度实现了接近最优的性能，通过将模型中众多类注意力操作统一到单个 GPU 内核来加速训练，前向速度提升最高达 2 倍。尽管不同内核有不同需求，但 TLX 被证明是大多数 GEM 工作负载的正确抽象，在这些场景中开发速度和性能都很重要。

TLX 使我们能够超越对单个内核的优化，转向对模型更大范围进行内核化和融合。对于 GEM 工作负载，减少内存移动通常比单纯优化计算带来更大收益，因此可组合性和融合成为首要关注点。为加速这一过程，我们构建了计算效率智能体（ACE），一个 AI 辅助的优化框架，可自动化性能分析、分析和调优。ACE 缩短了内核的优化周期，使工程师能够专注于更高层次的优化机会，而非底层的性能调试。

案例研究 2：RecSys 的内核内广播优化

在 RecSys 推理中，用户特征在所有被评分的候选项之间是相同的，但它们通常在每个批次中被复制数十次以匹配候选项维度。在广泛的候选项与用户广播比例范围内验证后，这种广播成为内存带宽和计算浪费的重要来源，其开销随候选项数量线性增长。

内核内广播优化（IKBO）通过将广播移入内核本身来消除这种低效。内核不再物化复制的用户特征，而是以各自自然的、不匹配的批次大小消费用户和候选项输入，并在计算过程中隐式执行广播。这种方法通过模型架构、推理运行时和内核实现的协同设计得以实现。

TLX 为这一优化的最后阶段提供了构建模块

考虑我们的线性压缩内核。经过一系列优化——包括矩阵乘法分解、内存对齐和广播融合——剩余的瓶颈通过一个 warp 专用、多阶段 TLX 内核得到解决，该内核在 NVIDIA H100 上实现了 4 倍加速。最终实现将整个计算折叠为一个单一的持久巨型内核。生产者和消费者 warp 组使用乒乓调度并发执行，而用户和候选 GEMM 共享循环缓冲区，并通过 mbarrier 同步，而无需在阶段之间排空流水线。瓦片级释放-获取同步进一步实现了跨 CTA 的协调。

对于 FlashAttention，IKBO 协同设计通过将 K/V 读取分摊到共享同一用户上下文的候选上，将内核从 I/O 密集型转变为计算密集型。随后，TLX 叠加了具有持久性的 FA3 风格 warp 专用化，消除了在 RecSys 短查询长度下占主导地位的跨 CTA 启动停顿，并为协同设计的模型将计算密集型网络延迟最多降低三分之二。直接在 CUDA 中表达这种执行结构需要大量手动编排；TLX 允许在 Triton 的 Python 编程模型和自动调优框架内简洁地表达它。

如今，IKBO 帮助优化 Meta 的自适应排序模型，覆盖 GPU 和 Meta 训练和推理加速器（MTIA）。更广泛地说，它说明了我们处理 RecSys 推理方式的转变。我们不再围绕广播开销进行优化，而是在计算原语本身消除它，从而实现无论候选数量如何的密集交互质量。

案例研究 3：用于 LLM 预训练的分布式分组 GEMM

TLX 的生产影响超越了推荐工作负载。基于 TLX 的分布式分组 GEMM 内核目前已部署在 Meta 的大规模 LLM 预训练基础设施中，并在大型 GPU 集群上连续运行了较长的训练期。该内核加速了混合专家（MoE）模型所需的分组矩阵乘法，同时紧密集成计算和通信以降低端到端延迟。

在这种规模下，峰值吞吐量只是等式的一部分。可靠性、可调试性和长期性能稳定性同样重要。生产内核必须在数周的连续运行中保持性能，同时在模型、硬件和系统软件演进时保持可维护性。

TLX 被证明特别适合这种环境。它对 warp 专用执行、显式内存编排和异步流水线的支持，实现了大规模分布式训练所需的性能特征。同时，它与 Triton 的 Python 编程模型的集成保留了以生产速度迭代、验证和部署复杂内核所需的开发者生产力。

这次部署凸显了 TLX 设计理念的一个重要方面：它不仅仅是为了实现峰值内核性能，而是为了使先进的硬件能力能够在真实生产系统中实际采用。

深入了解 TLX 内核

如需深入了解 TLX 内核的技术细节，请参阅我们配套的 PyTorch 博客文章：

• 内核内广播优化：为推荐系统推理协同设计内核。
• TLX Block Attention：面向固定块稀疏自注意力的 Warp 专用 Blackwell 内核。
• 快速 2-单纯形注意力：TLX 中的硬件高效内核。
• 使用 TLX 启用集群启动控制。
• 实现免费归一化：通过 TLX 将归一化融合进 GEMM 和注意力。
• 使用 TLX 优化 Jagged Flash Attention：通往 Blackwell 上 SOTA FA4 之路。
• 使用 TLX 在 CDNA4 上优化 GEMM + 激活。

我们的经验：构建 GPU 编译器的建议

从真实工作负载开始，而非基准测试

我们构建 TLX 是为了解决 Meta 的实际生产问题，而不是为了赢得基准测试竞赛。GEM MFU 项目和 IKBO 都体现了这一原则。每个项目都始于一个具体的痛点（GEM 中 MFU 偏低，推荐系统中的广播开销），工具随后跟进。这种专注使我们始终面向真实世界的性能，而非合成分数。

生产力与性能是乘数关系，而非权衡取舍

TLX 通常能提供与手工调优 CUDA 相当的性能，但工程师的迭代速度显著更快。一个仅略快但编写时间多出三倍的内核，在集群规模下是净损失。在 GEM 中，多 DSL 方法让我们选择 TLX 进行快速迭代；在 IKBO 中，基于 TLX 的实现耗时数周而非数月。

硬件特性只有在可访问时才有用 

WGMMA、TMA 和 TMEM 是强大的能力，但由于编程模型过于复杂，它们很少被使用。TLX 通过简洁的注解使它们变得可访问。IKBO 中的 warp 专用融合和 GEM 中的多 DSL 内核若用原生 CUDA 实现，将会复杂到令人望而却步。

编译器是用户界面的一部分 

错误消息、调试支持和性能分析集成与代码生成质量同样重要。我们在开发者体验上的投入不亚于优化器。当 GEM 工程师调试 MFU 问题或 IKBO 工程师遇到 warp 组同步问题时，他们需要清晰的反馈，而非晦涩的编译器错误。

演进优于革命

通过将 TLX 构建为 Triton 的扩展，我们利用了现有的生态系统、工具和专业知识。GEM 团队可以将 TLX 与 CUTLASS 和 CuTeDSL 一起使用；IKBO 团队已熟悉 Triton，只需学习性能关键部分的 TLX 注解。全新设计可能更优雅，但可能永远无法发布。

TLX 的下一步

TLX 正在更进一步。以下是我们正在开展的工作：

• 更多硬件：目前，TLX 面向 NVIDIA Hopper 和 Blackwell，以及 AMD Mi300 和 Mi350x。我们正在努力支持更多架构，包括 NVIDIA 和 AMD 的新一代产品，以及 MTIA 等定制加速器。MIMW 模型与硬件无关——挑战在于将其高效映射到不同的 ISA。GEM 已在 GPU 和 MTIA 上运行，IKBO 则横跨 GPU 和 MTIA。
• 更高层抽象：TLX 让开发者能够控制 warp group。接下来，我们希望再次提升抽象层级——从更高层的模式中自动推断 warp-group 编排。可以想象一个“auto-TLX”，为你自动编写注解。我们从 GEM（多 DSL 选择）和 IKBO（三阶段 warp 特化）中学到的模式，都是自动化的候选对象。
• 与 PyTorch 2.x 集成：我们正在努力让 TLX 成为 torch.compile 的后端。愿景是编写 PyTorch 代码，就能自动获得 TLX 优化的内核，常见模式无需手动注解。GEM 展示了系统性协同设计所能实现的能力，IKBO 则展示了针对特定模式进行优化的价值。我们希望将两者都自动化。

缺失的中间层确实存在，而且不仅适用于 GPU。随着加速器变得越来越复杂，我们需要能够在不牺牲生产力的前提下暴露恰当控制层级的编程模型。TLX 是我们朝这个方向迈出的第一步。

致谢 

TLX 是 Meta 内部 Triton、PyTorch、MSL、Ads Training Infra、Ads Ranking、RecSys 和 AI Compilers 团队的合作成果。我们感谢以下人员的贡献：

Triton： Shane Nay、Manman Ren、Oleksandr Stashuk、Karthik Manivannan、Zhijing Li、Yuanwei Fang、Shunting Zhang、Arda Unal、Yueming Hao、Taylor Robie、Puyan Lotfi、Corbin Robeck、Bert Maher、Alexey Loginov、Ying Liu、Bill Yoshimi。

PG： Jackie Xu、Jacky Zhou、Han Xu、Hao Yan、Jake Siso、Devashish Shankar、Shikai Li、Jianyu Huang、Andrew Tulloch、Boda Li、Jiao Jian、Zhuoran Zhao、Ao Cai、Shuyao Bi、Yuanjuan Yao、Ruizhen Guo、Doris Xiang、Jason Lu、Marvin Fan、Hai Zheng、Zhaokang Ke、Albert Chen、Mingyu Guan

UCSD： Yufei Ding 和 Yue Guan。

文章 How We Built TLX: A Hardware-Native GPU Compiler for Production ML at Scale 首次发表于 Engineering at Meta。

来源：Meta Engineering Blog · engineering.fb.com