技术教程｜深入详解众智 FlagOS 社区 Triton-TLE-Raw，解锁算力极致潜能

Triton 语言凭借简洁高效的编程范式与强大的自动优化能力已成为当前 AI 算子开发领域的主流核心工具。但在追求硬件极限性能的极致调优场景中Triton 通用的 DSL 抽象层往往无法充分表达硬件原生特性。无论是精细的指令级并行控制、内存屏障同步策略还是芯片专属的寄存器分配规则都难以通过通用语法完整传递最终导致编译器无法生成贴合硬件底层的最优代码让核心算子遭遇难以突破的性能瓶颈。针对这一行业核心痛点众智 FlagOS 社区基于 FlagOS 体系核心的统一 AI 编译器 FlagTree推出了全新的 Triton 语言扩展体系Triton-TLETriton Language Extensions从三个层级全面扩展 Triton 的原生能力精准匹配算子开发者不同层级的开发需求。作为面向多元 AI 芯片打造的统一开源系统软件栈众智 FlagOS 以「一次开发、跨芯迁移」为核心目标致力于打破不同硬件生态间的技术壁垒。而作为 FlagOS 技术体系核心的统一 AI 编译器 FlagTree基于 Triton 完成深度定制优化原生支持 31 种计算原语可在英伟达、华为昇腾、清微智能等主流 AI 芯片上实现性能大幅跃升调优效率与最终表现接近甚至超越芯片原厂的 C 语言算子实现。Triton-TLE 通过TLE-Lite、TLE-Struct、TLE-Raw三层架构构建了兼顾易用性、灵活性与极致性能的统一语言扩展体系。其中Triton-TLE-Raw以「原生透传、极致掌控」为核心理念专为需要突破硬件性能极限的极致调优场景打造是体系中触达硬件底层能力的核心抓手。本文将全面解析 Triton-TLE-Raw 的核心理念、环境部署、实操流程、核心特性与落地价值帮助 AI 算子开发者在保留 Triton 生态易用性的基础上精准掌控硬件底层能力实现算子性能的跨越式提升。01 什么是 Triton-TLE Triton-TLE 以「分层解耦、全场景覆盖」为核心构建了从初级到专家的三阶算子开发体系TLE-Lite实现轻量快速提效TLE-Struct提供硬件架构级进阶优化TLE-Raw作为体系的性能天花板面向专家级开发者提供硬件原生级极致操控能力TLE-Raw 支持直接调用硬件厂商原生编程语言实现对硬件的底层直控解锁极致性能。在完整编译流程中TLE-Lite 与 TLE-Struct 通过 FLIR 最终 Lowering 至 LLVM IRTLE-Raw 则通过硬件专属编译管线完成转换最终三者统一 Link 为完整 kernel供 Runtime 加载执行。这种设计打造了高效的分层开发模式算子通用逻辑、非热点路径可通过前两层快速实现兼顾开发效率与可移植性决定性能上限的核心热点路径可通过 TLE-Raw 完成原生级深度调优无需重构现有项目即可在一套代码中兼顾开发效率与硬件极限性能。作为 FlagTree 面向极致性能场景的底层扩展TLE-Raw 的核心价值在于打破 Triton DSL 的抽象边界。Triton-TLE-Raw 坚持原生透传、极致掌控的核心原则让开发者在标准的 Triton 开发体系内即可获得与手写原生代码一致的硬件操控能力。同时该组件复用 Triton 现有的编译流程与工程体系无需对现有项目进行重构即可快速接入原生优化能力并且支持多种硬件原生代码格式具备优秀的跨硬件适配能力。02 安装与环境部署Triton-TLE-Raw 作为 FlagTreeFlagTree 是 FlagOS 的重要组成部分的核心内置模块需要通过部署指定版本的 FlagTree 完成环境搭建同时配置必要的编译与运行依赖才能正常使用全部功能。2.1 安装 FlagTreeTriton-TLE-Raw 作为 FlagTree 的核心模块需通过 FlagTree 仓库部署指定稳定版本。gitclone https://github.com/flagos-ai/FlagTreecdFlagTreegitcheckout triton_v3.6.x# 切换到支持TLE-Raw的版本# 参考Wiki进行安装克隆并切换分支完成后按照项目 Wiki 文档*的指引完成编译与安装流程即可完成基础环境部署。*Wiki链接https://github.com/flagos-ai/FlagTree/wiki/2.2 安装 TLE-Raw 依赖在核心依赖配置中Clang 编译器是实现原生代码内联的关键工具主要用于将 CUDA、MLIR 等原生代码编译为 LLVM IR因此请确保 Clang 已安装并将其配置到系统环境变量中。除此之外开发者还可按照项目具体文档如 EDSL 设计文档*的要求安装所需的 Python 包与系统底层库保障 Triton-TLE-Raw 与 FlagTree 生态的协同稳定运行。*EDSL 设计文档https://github.com/flagos-ai/FlagTree/wiki/03 快速上手TLE-Raw 使用全流程详解TLE-Raw 的使用主要分为Triton 侧的声明调用和硬件域如CUDA侧的原生代码编写两部分。下面以 Vector-Add 为例展示完整流程。3.1 Triton侧——声明与调用硬件域函数在 Triton kernel 中需要声明硬件域函数并在适当位置调用它。# triton侧代码示例 (01-vector-add.py)frompathlibimportPathimporttorchimporttritonimporttriton.languageastlfromtriton.experimental.tle.rawimportdialectimporttriton.experimental.tle.language.rawastle_raw# 声明CUDA硬件域函数dialect(namecuda,filePath(__file__).parent/01-vector-add.cu)defedsl(*args,**kwargs):...# Triton kernel实现triton.jitdefadd_kernel(x_ptr,y_ptr,output_ptr,n_elements,BLOCK_SIZE:tl.constexpr,):tle_raw.call(edsl,[],[output_ptr,x_ptr,y_ptr,n_elements])3.2 CUDA侧——编写原生内核代码在独立的.cu文件中编写具体的 CUDA 内核实现。这部分代码将被 Clang 编译为 LLVM IR并内联到 Triton 生成的代码中。// CUDA侧代码示例 (01-vector-add.cu) // CUDA内核实现 __device__ void VectorAdd(__attribute__((address_space(1))) float *C, __attribute__((address_space(1))) const float *A, __attribute__((address_space(1))) const float *B, const int N) { const int idx blockIdx.x * blockDim.x threadIdx.x; for (int i idx; i N; i blockDim.x * gridDim.x) { C[i] A[i] B[i]; } }3.3 执行与验证代码编写完成后开发者可按照标准 Triton 程序的运行方式执行验证Triton-TLE-Raw 会自动完成原生代码的内联、编译与执行流程无需对原有执行逻辑进行改造即可快速验证计算结果的正确性。04 关键特性与高级用法Triton-TLE-Raw 提供了多项面向极致性能的高级特性能够帮助开发者更精细地控制硬件资源进一步挖掘算力潜力。4.1 支持共享内存Shared Memory共享内存是 GPU 的高速缓存合理使用可大幅减少全局内存访问开销。TLE-Raw 支持通过 Triton-TLE-Struct 在 Triton kernel 中分配共享内存并将其地址信息传递给 CUDA 侧函数。这允许在 CUDA 代码中直接操作共享内存实现更高效的数据复用。核心逻辑Triton 侧通过tle.alloc分配共享内存tle.copy将全局内存数据搬运至共享内存tle.local_ptr获取共享内存地址CUDA 侧接收地址偏移量结合共享内存基地址计算具体数据地址。4.2 数据类型与参数映射TLE-Raw 需要处理 Triton 侧的语义对象如tensor、ptr与 LLVM 侧的底层类型之间的映射。这是确保代码正确内联的关键。05 应用场景与性能案例Triton-TLE-Raw 的应用场景十分明确主要覆盖三类核心开发需求适用于 Triton 自动优化不足的算子对于一些计算密集、内存访问模式复杂的算子Triton 的自动优化可能无法达到手写 CUDA 的性能。此时可以使用 Triton-TLE-Raw 将核心计算部分替换为优化过的 CUDA 代码突破瓶颈。适用于需要精细控制硬件资源的场景支持手动管理共享内存、控制线程同步、或利用特定硬件指令集如 Tensor Core的场景充分释放硬件底层能力。适用于现有高性能 CUDA 算子的移植无需重写性能优异的 CUDA 算子Triton-TLE-Raw 可将其便捷地集成到 Triton 生态中兼顾生态兼容性与原生执行效率。在实际性能测试中针对长序列场景下的 TopK 算子*原生 Triton 实现因自动优化限制无法充分发挥硬件并行特性而采用 Triton-TLE-Raw 重写核心计算逻辑后单算子性能可达到原生 Triton 实现的 1.4 倍在长序列推理场景中能够显著降低延迟充分验证了 Triton-TLE-Raw 在性能优化上的核心价值。*TopK算子https://github.com/flagos-ai/FlagTree/blob/triton_v3.6.x/python/tutorials/tle/raw/mlir/05-topk.py结语TLE-Raw 通过“原生透传”的设计理念为 AI 算子开发者提供了兼顾 Triton 生态易用性与硬件极限性能的全新方案。作为 FlagOS 与 FlagTree 生态的关键组成它巧妙地平衡了 Triton 的易用性与硬件原生代码的高性能使得专家级的性能优化能够在一个统一、安全的框架内进行。未来FlagOS 生态将持续迭代升级Triton-TLE-Raw 也将不断拓展硬件原生支持、深化多架构适配推动 AI 算子开发迈向更高效、更极致的新阶段为多元 AI 算力生态筑牢坚实底层支撑。关于众智 FlagOS 社区为解决不同 AI 芯片大规模落地应用北京智源研究院联合众多科研机构、芯片企业、系统厂商、算法和软件相关单位等国内外机构共同发起并创立了众智 FlagOS 社区。成员单位包括北京智源研究院、中科院计算所、中科加禾、安谋科技、北京大学、北京师范大学、百度飞桨、硅基流动、寒武纪、海光信息、华为、基流科技、摩尔线程、沐曦科技、澎峰科技、清微智能、天数智芯、先进编译实验室、移动研究院、中国矿业大学(北京)等多家在 FlagOS 软件栈研发中做出卓越贡献的单位。FlagOS 是一款专为异构 AI 芯片打造的开源、统一系统软件栈支持 AI 模型一次开发即可无缝移植至各类硬件平台大幅降低迁移与适配成本。它包括大型算子库、统一AI编译器、并行训推框架、统一通信库等核心开源项目致力于构建「模型-系统-芯片」三层贯通的开放技术生态通过“一次开发跨芯迁移”释放硬件计算潜力打破不同芯片软件栈之间生态隔离。官网https://flagos.ioGitHub 项目地址https://github.com/flagos-aiGitCode 项目地址https://gitcode.com/flagos-ai