Meta把GPU执行控制搬进Triton：TLX让生产训练与推理直接编排warp组
核心概要
Meta 推出 Triton 低层扩展 TLX，在保留 Triton 编程模型的前提下，让工程师显式编排 warp 组、异步数据搬运与集群执行。现代 GPU 的 warp 特化、张量内存和集群执行能力此前要么藏在编译器抽象之后，要么需要大量架构专用 CUDA 代码。TLX 已在 Meta 的推荐系统训练、推荐系统推理和大规模 LLM 预训练中投入生产，其中 GEM 训练在数千块 NVIDIA Blackwell GPU 上达到 20-25% 端到端 MFU。
深度剖析
GPU 编程多了一层可显式书写的执行结构：工程师可以在 Triton 里指定哪些 warp 组负责搬运数据、哪些负责张量核心计算、哪些处理结果，而不必退回架构专用的 CUDA 编排代码。 此前 CUDA 的 SIMT 模型把控制放在线程粒度，Triton 的分块模型把执行关系交给编译器隐式决定，两者之间缺少一个以 warp 组为单位的抽象层。 Meta 在 NVIDIA Blackwell 上比较了 CUTLASS、CuTeDSL 和 TLX 三种编程模型，认为 TLX 在多数内核上兼顾了生产率与性能，并已将其用于 GEM 训练、推荐推理和大规模 LLM 预训练三类生产负载。
TLX 以 Triton 扩展的形式提供第二层面向执行的构造，计算仍用 Triton 张量操作表达，只有需要显式控制执行、内存搬运或同步的地方才调用 TLX。 这一层把 warp 组当作一等执行任务，用 mbarrier 等硬件原语表达生产者与消费者之间的依赖，并直接开放 TMA、异步张量核心操作和 TMEM 等能力。 在 Blackwell 上的矩阵乘示例中，生产者、消费者和收尾三个 warp 组通过轻量屏障交接共享内存缓冲区，形成深度流水线，同时保留 Python 与 Triton 的自动调优设施。
在推荐系统推理中，TLX 支撑的内核内广播优化把用户特征的复制搬进内核内部，让内核直接消费批大小不匹配的用户与候选输入。 此前用户特征在每个候选批次中被复制数十次以对齐候选维度，这部分广播随候选数量线性增长，浪费内存带宽和计算。 经过矩阵乘分解、内存对齐和广播融合后，剩余瓶颈由一个 warp 特化、多阶段的 TLX 内核处理，在 NVIDIA H100 上取得 4 倍加速；对协同设计的模型，计算密集场景的净延迟最多降低三分之二。
启示与展望
已展示的范围是 Meta 生产环境中的三类负载：GEM 推荐模型训练、推荐系统推理的 IKBO 以及 LLM 预训练的分布式分组 GEMM；硬件覆盖 NVIDIA Hopper、Blackwell 和 AMD Mi300、Mi350x。TLX 目前作为 Triton 的扩展使用，面向需要显式执行控制的性能关键内核。
TLX 对更多架构（包括 NVIDIA 与 AMD 新一代以及 MTIA 等定制加速器）的支持仍在开发中，MIMW 模型本身与硬件无关，难点在于高效映射到不同指令集。自动推断 warp 组编排的更高层抽象，以及作为 torch.compile 后端，也仍在推进。