8 月 1 日Meet AI Compiler 技术沙龙第 9 期在北京举办。本期活动聚焦 AI 编译技术的最新进展多位来自产业界和科研机构的专家围绕编程语言、算子开发、编译优化与推理执行展开分享呈现 AI 编译器从上层语言表达到硬件执行的协同演进。其中BAAI AI 编译器研究员郭晖与肖航以「FlagTreeTriton-TLE 语言扩展、Tile IR 后端与编译优化实践」为题介绍了 FlagTree 团队围绕 Triton 开展的一系列探索。BAAI - AI 编译器研究员 - 郭晖老师BAAI - AI 编译器研究员 - 肖航老师面对日益复杂的硬件架构和更加不规则的模型算子团队提出 Triton Language ExtensionTLE尝试在高级 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 OperationTKO共同表达计算、数据视图和必要依赖。其中TensorView 描述全局指针、Shape 和 StridePartitionView 在此基础上增加分块映射能力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.217.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 端到端测试中TP4 时首 Token 时延降低 20.22%总吞吐提升 12.63%TP8 时首 Token 时延降低 12.87%总吞吐提升 7.65%。未来工作展望下一阶段FlagTree 将重点推进 NUMA 编程模型和 MegaKernel Compiler。面向多节点、多 GPU以及 GPU 内部的 Cluster 和本地 SRAM数据访问会呈现明显的远近差异。TLE 希望显式描述 Tensor 如何沿 Mesh 切分以及如何在不同分布方式之间转换使编译器能够判断计算应靠近哪些数据哪些通信可以被合并或与计算重叠。MegaKernel Compiler 则尝试将优化范围从单个 Kernel 扩展至一段模型执行过程。编译器将识别可以协同优化的计算拆解为具体任务再统一决定任务如何分配、并行和流水最终生成一个或少数几个 MegaKernel。TLE 在其中承担前端语义与硬件结构之间的桥梁。Tile、Mesh、Pipeline、Memory Layout 和原生能力接口使编译器看到的不再只是一串独立算子而是一张可以调度、放置、同步和融合的任务图。从 TLE-Lite 的轻量语义提示到 TLE-Struct 的架构感知控制再到 TLE-Raw 的原生能力透传结合各种编译优化技术FlagTree 正尝试在 Triton 易用性、跨平台迁移和极致性能之间建立一套更具弹性的分层体系并让新硬件能力更快转化为模型和算子的实际性能。