智源研究院提出 Triton-TLE 分层语言扩展，并通过 FlagTree 编译优化实现百倍级自动调优加速

关注

原创最新推荐文章于 2026-11-10 00:57:14 发布 198 阅读 AI 写同款· GEO 优化 ›

8 月 1 日，Meet AI Compiler 技术沙龙第 9 期在北京举办。本期活动聚焦 AI 编译技术的最新进展，多位来自产业界和科研机构的专家围绕编程语言、算子开发、编译优化与推理执行展开分享，呈现 AI 编译器从上层语言表达到硬件执行的协同演进。

其中，BAAI AI 编译器研究员郭晖与肖航以「FlagTree：Triton-TLE 语言扩展、Tile IR 后端与编译优化实践」为题，介绍了 FlagTree 团队围绕 Triton 开展的一系列探索。

BAAI - AI 编译器研究员 - 郭晖老师

BAAI - AI 编译器研究员 - 肖航老师

面对日益复杂的硬件架构和更加不规则的模型算子，团队提出 Triton Language Extension（TLE），尝试在高级 DSL 与底层硬件控制之间建立一条渐进式通道；同时在 FlagTree 编译器上围绕 TileIR 后端、自动调优、数据布局、指令调度等方面开展编译优化。

这套工作的核心目标，是在保留 Triton 易用性和社区生态的同时，让开发者可以根据实际需要逐步深入硬件细节，以相对统一的开发体系覆盖快速算子优化、架构感知调优和原生代码级优化。

两位老师与到场观众深度交流

HyperAI 在不违原意的前提下，对分享内容进行了整理汇总。

关注微信公众号「HyperAI超神经」，后台回复关键字「0801 AI 编译器」，即可获取确认授权的讲师演讲 PPT。

在 Triton 之上，重新平衡抽象与性能

近年来，AI 编译器面临一个越来越突出的矛盾：硬件架构和模型算子不断复杂化，但开发者仍希望使用较高层的 DSL，获得接近专家手写 Kernel 的性能。

抽象过高，编译器未必能获得足够的信息完成深度优化；抽象过低，开发者又会重新回到 CUDA 或厂商私有语言的复杂开发模式中。

Triton 的成功，在于将 GPU 算子开发从 Thread 级提升至 Tile 级。用户使用 Python DSL 描述数据块之间的计算关系，线程映射、寄存器分配、数据布局、流水线和同步等工作主要交给编译器处理。

这种方式降低了高性能算子的开发门槛，也形成了庞大的社区生态。但随着新一代 GPU、领域专用架构和国产 AI 芯片不断发展，Triton 原有抽象也开始遇到边界。

一方面，如果编译器后端尚未支持新硬件的存储结构、通信机制或计算单元，前端开发者很难自行利用这些能力；另一方面，越来越多的关键算子需要精确控制存储层级、并行粒度、CTA 间协作，以及通信与计算的重叠，原有 Triton 代码有时难以完整表达这些意图。

Gluon、TLX、TileLang 等新语言和 DSL 的出现，也反映了同一趋势：AI 算子开发已经不只是编写单个 Kernel，还需要表达数据布局、并行层次、流水线、通信拓扑和硬件特性。

TLE 并不试图替代 Triton，而是在其语法和生态之上进行分层扩展。它由 TLE-Lite、TLE-Struct 和 TLE-Raw 三个层次组成，分别对应轻量级语义提示、架构感知控制和原生代码级优化。

TLE-Lite 面向算法工程师和快速调优场景。开发者不必直接关心底层硬件，而是向编译器补充更明确的结构信息，例如某个 Tensor 需要按照子 Tile 访问，某项计算运行在分布式 Mesh 上，或者一个 CTA 内部采用 Producer-Consumer Pipeline。以子 Tile 操作为例，开发者可以从较大的 Tensor 中直接取出一个逻辑子块，完成激活、归一化或统计计算后再写回，而不必手动计算偏移、构造 Mask 和处理边界。

由于编译器能够识别这是一次规则的 Tile 访问，便可以围绕数据布局、向量化、Bank Conflict 和寄存器复用进一步优化。这种方式尤其适合稀疏 Attention、局部归一化、分块统计和路由类算子。在分布式场景中，TLE 使用 Device Mesh 描述节点、GPU、Block Cluster 和 Block 等不同层级，将其组织为统一的多维拓扑。

开发者面向 Mesh 维度编写程序，编译器和 Runtime 再将逻辑拓扑映射到真实硬件和通信机制。由此，Ring 通信、Barrier 同步和分片访问不再只是代码中散落的 Rank 和通信组，而成为可分析的结构化语义。一旦通信关系被显式表达，编译器便有机会开展拓扑感知调度、通信计算重叠、Barrier 合并和死锁检查。

TLE-Lite 还通过 Pipeline 原语，将 CTA 内部协作抽象为 Producer-Consumer 模型。开发者主要描述数据由谁生产、由谁消费，底层 Barrier、缓冲区复用和同步机制则交给编译器处理。这并不是屏蔽底层能力，而是将其转化为可分析、可验证和可优化的程序结构。

从轻量语义到原生透传，覆盖不同优化深度

如果说 TLE-Lite 主要解决跨平台语义表达，那么 TLE-Struct 更关注架构感知和精细调优。

不同 GPU、DSA 和 AI 加速器在存储层次、执行单元、同步机制和片上网络上存在明显差异。TLE-Struct 因此向开发者暴露层次化的并行和存储结构，允许其显式定义数据布局、计算映射和内存层级，但不必直接绑定某一家厂商的私有接口。

例如，同一块本地 Buffer 在 GPU 后端可以映射至 Shared Memory，在 DSA 后端则可以映射至 Scratchpad 或片上 SRAM。用户表达的是结构化的内存意图，编译器负责将其转换为适合目标硬件的地址空间和访存指令。

MoE 中的 Expert Counting 是一个典型场景。该算子本质上是统计 Token 被路由至不同 Expert 的数量，容易受到共享内存布局、并发更新、Bank Conflict 和跨 Block 汇总等因素影响。

借助 TLE-Struct，开发者可以显式组织计数器布局，将不同 Expert 或 Token 映射至本地存储的不同区域；编译器获得这些结构信息后，再生成合适的访问和同步方式。

TLE-Raw 则面向性能优化专家，为厂商原生代码保留接口。

部分性能路径必须直接使用 CUDA、汇编或专用 Intrinsic。如果强行将其重新封装进高层 DSL，反而可能限制性能或增加迁移成本。TLE-Raw 允许开发者在 Triton/TLE 体系中内联原生代码，或直接接入厂商编译管线。

以 All-Gather GEMM 为例，该算子同时涉及通信、矩阵计算、Buffer 和同步管理。开发者可以复用底层通信能力，同时继续使用结构化的 Tensor 和 Tile 表达计算，最后由统一编译管线将不同部分组织为可调用算子。

这种分层设计，使算法工程师、算子开发者和性能专家能够在同一套体系内选择不同的优化深度，而不必一开始就进入最底层编程。

性能测试也表明，TLE 的抽象开销整体可控。

在 Radix Select 测试中，TLE 复刻了 TensorRT-LLM 的算法，在多组 Shape 下达到其约 85%—97% 的性能。对于需要跨平台维护和快速迭代的团队而言，以更统一的代码获得接近专家实现的性能，具备较高工程价值。

在 128K 上下文的 SparseMLA 场景中，TLE 通过 Pipeline 原语表达不同执行角色之间的协作，性能达到 FlashMLA 基线的约 90%。

团队还在单节点 8 张 NVIDIA H100 上测试了 All-Gather，并进一步探索 GEMM 与 All-Gather 融合。其重点不是简单替代通信库，而是将通信纳入算子级表达和编译优化范围，使数据传输、局部计算、同步和后续消费可以共同形成流水。

对于推理系统而言，这类通信计算融合往往比孤立提升某个 GEMM 的峰值性能更有意义，因为用户最终感知的是端到端等待时间。

Flagtree 编译优化实践

除 TLE 语言扩展外，FlagTree 团队还开展了两方面工作：一是接入 CUDA Tile IR 后端，二是面向真实模型工作负载实施多项 Triton 编译优化。

CUDA Tile IR 的核心理念，是让程序表达 Tile，而不是提前固化底层线程映射。其输入由 Tile Program 和三维 Tile Grid 组成，Dialect 内部则通过 Tile Compute、View 和 Token-Ordered Operation（TKO）共同表达计算、数据视图和必要依赖。

其中，TensorView 描述全局指针、Shape 和 Stride；PartitionView 在此基础上增加分块映射能力；Token 用于约束相关 TKO 之间的依赖关系；Memory Model 则单独规定内存语义与作用域。

FlagTree 并未整体替换 Triton NVIDIA 原生后端，而是在原有体系中增加 TileIR 路径，并通过 TLE 原语覆盖 TileIR的 View 和 Token 接口。满足条件的 Kernel 可以进入 TileIR 后端；暂不支持的 Kernel，则回退至 Native CUDA，两条路径可以在编译器中共存。

自动调优方面，团队提出 FlagOSTune，以解决 Triton Autotune 在真实模型中的覆盖与成本矛盾。

实际模型推理中，同一种算子可能对应大量不同 Shape。团队统计发现，六个模型、四类推理场景中包含 1994 个唯一 MM Shape，远高于日常 Benchmark 的覆盖范围。如果直接扩展候选配置，理论搜索规模又会迅速增长至数百万组。

FlagOSTune 通过扩展搜索空间提高性能，通过「模型预测+少量实测」的方式压缩搜索成本。系统先使用 XGBoost 对候选配置排序，只将最有潜力的少量配置送入真实编译和 GPU 测试，再结合遗传算法继续搜索。

在 NVIDIA, 摩尔线程和沐曦多款算力上提升了多个算子的性能，加速 1.21～7.35x；在 NVIDIA H20 MM 算子的多个 shape 实验中，搜索空间配置从 62 万余个压缩至 4070 个，调优效率加速 120 倍，性能几乎无损。

在数据布局方面，团队重点优化了数据布局转换的性能开销。

Triton 中的 convert_layout 通常涉及跨线程数据重排，需要写入 Shared Memory、同步后再读回，并不是零成本操作。对于小算子或访存密集型算子，少量布局转换就可能成为性能瓶颈。

FlagTree 对 Remove Layout Conversions 机制进行增强，通过代价模型、反向传播和局部联合求解等技术，减少不必要的数据布局转换。在 100 多个算子的测试中，净转换消除率约为 68%—79%，性能提升最高 71%。

另一项优化是指令重排。循环展开只是复制循环体，并不会自动把多次 Load 提前发射。开启编译器重排后，相互独立的 Load 可以提前执行，再穿插后续计算，从而利用硬件流水线隐藏访存延迟。在三类典型算子中，这项优化平均带来约 1.19—1.61 倍加速，峰值提升达到 2 倍。

团队还针对 DeepSeek-V4-Flash 的真实 Shape 优化 Fused Marlin MoE 算子的性能，通过 Fragment 合并减少数据加载时间。结合算法层面的优化，与 vLLM CUDA 相比，FlagGems 在测试的 53 个 Shape 上均取得加速，按shape出现频次加权计算，平均提升约 1.208 倍。

在 NVIDIA H20 的 DeepSeek-V4-Flash 端到端测试中，TP=4 时首 Token 时延降低 20.22%，总吞吐提升 12.63%；TP=8 时首 Token 时延降低 12.87%，总吞吐提升 7.65%。