Triton本文指 OpenAI Triton先看官方介绍Triton is a language and compiler for parallel programming. It aims to provide a Python-based programming environment for productively writing custom DNN compute kernels capable of running at maximal throughput on modern GPU hardware.Triton 是一个用于并行编程的 Python DSL也是一个将计算映射到并行硬件的编译器框架。两个视角算法研究使用 Python DSL 快速实现高性能算子无缝接入 PyTorch 生态编译器开发为自研 AI 芯片提供编译基础设施可基于 MLIR 构建完整编译栈它是一座连接高级算法思想与底层硬件并行性的桥梁兼顾了生产力、性能与可移植性。本文聚焦于 AI 编译器开发想的是怎么把 Triton MLIR 上手玩起来。不过在此之前先聊聊我写这篇文章的初衷。为什么写 Triton我也是初次接触 Triton。那为什么想到写它了呢因为目前 AI 编译器的主流方向是 Triton MLIR正好打算深入学习一下。但真要写 Triton现在又有问题了怎么写呢如今问一问 AI 不就好了。所以文风与内容打算稍作变化。啥意思会多点‘废话’多唠叨几句不再简要成手册这就由 AI 来做吧。我呢多谈谈自己的想法一是让大家感觉更像是在和一个人交流二是分享个人实践的经验希望能对大家有所裨益。GoCoding 定位让大家了解技术上新的东西或能引导大家进入新的领域。技术深入跟随会议或社区进展就好。怎么来学 Triton官方渠道不用多说当然从官方渠道开始。不建议看我这样的二手信息。那我写文的价值呢一对我个人最有用二让关注我的人能了解新的技术、发现感兴趣的领域一起学习。AI 辅助与 AI 聊聊我的目标看看它建议的学习路径。现在可真方便 ✌️那我们开始吧首先了解一样东西我不仅仅会了解它是什么、用做什么。就如文章开头对 Triton 的描述。我一般还会去了解背景它的由来、发展历程。挖掘得细一点有时还会发现一些有意思的小故事。其次就是对比。为什么是它它的优势在哪它在往什么方向发展。也就是了解一下它“到哪里去”的问题。更远的未来的可能性与不确定性学习并深入时可能会想一想。之后才是学习。笨办法是我用得比较多的。就是把一本书或官方教程全部写一遍、跑一遍。当然现在 AI 开道免费导师可以试试聪明办法了。又想起我自己说过的那句AI不懂的人用则是引导者现在就按上述思考开始吧Triton 的发展阶段发展阶段大致时间核心特征与关键事件技术栈与生态主要目标用户初创与原型2018 - 2021在OpenAI内部诞生核心是为NVIDIA GPU提供比CUDA更友好的编程抽象。发表了开创性的学术论文。独立的编译器深度集成于PyTorch。后端主要为CUDA/NVVM。AI研究人员、需要编写定制CUDA算子的开发者。成熟与生态扩张2021 - 2023关键转折基于MLIR重构。从“编译器”变为“编译器框架”奠定跨硬件基础。前端DSL稳定社区采用度激增如FlashAttention。MLIR成为核心中间层。后端扩展至AMD GPUROCM和Intel GPUOpenCL/SPIR-V。高性能计算工程师、框架开发者、其他硬件厂商的编译团队。平台化与工业部署2023 - 至今成为自研AI芯片的“上游编译前端”。其Triton IRMLIR Dialect成为连接算法与多种硬件的关键抽象层。基于MLIR的完整编译栈。出现针对不同自研硬件的定制后端。生态与Torch-MLIR等融合。AI芯片公司的编译器团队、构建端到端编译基础设施的架构师。Triton 的核心优势Triton 核心优势用 Python 写 GPU 算子基于 MLIR 轻松跨硬件已是 PyTorch 高性能算子的事实标准前端。Triton 与 TVM, XLA 的核心差异维度TritonTVMXLA设计哲学算子优先、Python 原生端到端编译、跨平台静态图优化、框架集成编程体验Python DSL算法思维多层 IR编译思维HLO IR框架约束硬件适配MLIR 方言栈灵活定制Relay/TOPI 多层抽象LLVM 专用后端生态定位PyTorch 生态的算子加速器多框架通用编译器TensorFlow/JAX 专属编译器Triton MLIR 跨硬件编译架构Triton MLIR 的学习路径学习路径可以遵循“从用户到开发者从抽象到具体”的原则分为四个层次层次一作为算子开发者用户体验层目标学会使用Triton Python DSL快速实现高性能内核。基础入门阅读Triton官方教程理解其编程模型tl.program_id,tl.arange,tl.load/store。亲手实现几个经典算子向量加法、矩阵乘法、Softmax。关键概念块Block、指针运算、掩码Mask、原子操作。进阶掌握学习内存优化技巧利用tl.cache,tl.make_block_ptr进行向量化加载/存储理解不同内存空间DRAM, L2 Cache, Shared Memory的差异。学习高级调度如流水线Pipeline和异步复制async_copy。学习如何将Triton内核封装为PyTorch的torch.autograd.Function并集成到模型中使用。实践尝试复现或优化一个经典论文中的融合算子如FlashAttention V1/V2的核心部分。层次二理解编译流程系统认知层目标理解从Triton Python代码到目标代码的完整编译过程。追踪编译链使用torch.compile或Triton的JIT功能运行一个简单内核并学习如何输出其关键中间表示。查看Triton IR了解高级Python操作如何被降低为Triton IR操作。查看MLIR理解Triton IR是如何被转换为MLIR中的Triton,LLVM,NVVM等Dialect的。这是理解其跨硬件能力的关键。学习MLIR基础概念无需深入建立认知即可什么是Dialect、Operation、Attribute、Type。MLIR的核心优势多层中间表示、可重用的转换和优化基础设施。实践使用TRITON_DUMP_GRAPH1等环境变量导出一个内核的编译流程图对照官方架构图理解每一步。层次三探索编译器开发贡献者/定制者层目标能够为Triton编译器添加新特性或为其定制新硬件后端。深入Triton编译器代码定位关键目录python/(DSL),lib/(核心编译器逻辑C)。学习从Python AST到Triton IR的生成过程。学习基于MLIR的Pass是如何对Triton Dialect进行优化的。学习为MLIR添加Pass编写一个简单的MLIR Pass对Triton Dialect的IR进行模式匹配和重写。后端开发入门以连接自研芯片为例路径一推荐将Triton Dialect lowering到已有的、更通用的MLIR Dialect如Vector,GPU然后利用社区或自定义的后端从这些Dialect生成代码。路径二更深入定义自己的硬件Dialect并实现从Triton Dialect到该硬件Dialect的转换规则。这需要对MLIR的方言转换框架Pattern Rewrite, Dialect Conversion有扎实理解。层次四架构与集成系统架构师层目标设计基于TritonMLIR的完整编译栈用于产品化部署。性能调优与自动化研究Triton的自动调优器Auto-Tuner原理学习如何为新的硬件特性添加调优空间。探索基于机器学习的内核自动生成与优化。端到端集成研究如何将Triton与模型级编译器如Torch-MLIR, IREE结合实现从PyTorch模型到混合内核部分由Triton生成的无缝编译。设计运行时Runtime接口管理自定义硬件上Triton内核的加载、启动和资源分配。紧跟社区关注Triton和MLIR的RFC、会议如MLIR Open Design Meetings、论文把握技术演进方向。Triton MLIR 动手开始学习路径有了那就动手开始吧。我根据个人情况想法是对四个层次竖着切一刀每一部分都实践一个简单用例先对整体建立概念再继续深入了解我关心的部分。每个人都可以依据自己的兴趣与实际情况去开始不必非得按部就班地学。结语有兴趣的小伙伴一起来学吧。
Triton 开始
Triton本文指 OpenAI Triton先看官方介绍Triton is a language and compiler for parallel programming. It aims to provide a Python-based programming environment for productively writing custom DNN compute kernels capable of running at maximal throughput on modern GPU hardware.Triton 是一个用于并行编程的 Python DSL也是一个将计算映射到并行硬件的编译器框架。两个视角算法研究使用 Python DSL 快速实现高性能算子无缝接入 PyTorch 生态编译器开发为自研 AI 芯片提供编译基础设施可基于 MLIR 构建完整编译栈它是一座连接高级算法思想与底层硬件并行性的桥梁兼顾了生产力、性能与可移植性。本文聚焦于 AI 编译器开发想的是怎么把 Triton MLIR 上手玩起来。不过在此之前先聊聊我写这篇文章的初衷。为什么写 Triton我也是初次接触 Triton。那为什么想到写它了呢因为目前 AI 编译器的主流方向是 Triton MLIR正好打算深入学习一下。但真要写 Triton现在又有问题了怎么写呢如今问一问 AI 不就好了。所以文风与内容打算稍作变化。啥意思会多点‘废话’多唠叨几句不再简要成手册这就由 AI 来做吧。我呢多谈谈自己的想法一是让大家感觉更像是在和一个人交流二是分享个人实践的经验希望能对大家有所裨益。GoCoding 定位让大家了解技术上新的东西或能引导大家进入新的领域。技术深入跟随会议或社区进展就好。怎么来学 Triton官方渠道不用多说当然从官方渠道开始。不建议看我这样的二手信息。那我写文的价值呢一对我个人最有用二让关注我的人能了解新的技术、发现感兴趣的领域一起学习。AI 辅助与 AI 聊聊我的目标看看它建议的学习路径。现在可真方便 ✌️那我们开始吧首先了解一样东西我不仅仅会了解它是什么、用做什么。就如文章开头对 Triton 的描述。我一般还会去了解背景它的由来、发展历程。挖掘得细一点有时还会发现一些有意思的小故事。其次就是对比。为什么是它它的优势在哪它在往什么方向发展。也就是了解一下它“到哪里去”的问题。更远的未来的可能性与不确定性学习并深入时可能会想一想。之后才是学习。笨办法是我用得比较多的。就是把一本书或官方教程全部写一遍、跑一遍。当然现在 AI 开道免费导师可以试试聪明办法了。又想起我自己说过的那句AI不懂的人用则是引导者现在就按上述思考开始吧Triton 的发展阶段发展阶段大致时间核心特征与关键事件技术栈与生态主要目标用户初创与原型2018 - 2021在OpenAI内部诞生核心是为NVIDIA GPU提供比CUDA更友好的编程抽象。发表了开创性的学术论文。独立的编译器深度集成于PyTorch。后端主要为CUDA/NVVM。AI研究人员、需要编写定制CUDA算子的开发者。成熟与生态扩张2021 - 2023关键转折基于MLIR重构。从“编译器”变为“编译器框架”奠定跨硬件基础。前端DSL稳定社区采用度激增如FlashAttention。MLIR成为核心中间层。后端扩展至AMD GPUROCM和Intel GPUOpenCL/SPIR-V。高性能计算工程师、框架开发者、其他硬件厂商的编译团队。平台化与工业部署2023 - 至今成为自研AI芯片的“上游编译前端”。其Triton IRMLIR Dialect成为连接算法与多种硬件的关键抽象层。基于MLIR的完整编译栈。出现针对不同自研硬件的定制后端。生态与Torch-MLIR等融合。AI芯片公司的编译器团队、构建端到端编译基础设施的架构师。Triton 的核心优势Triton 核心优势用 Python 写 GPU 算子基于 MLIR 轻松跨硬件已是 PyTorch 高性能算子的事实标准前端。Triton 与 TVM, XLA 的核心差异维度TritonTVMXLA设计哲学算子优先、Python 原生端到端编译、跨平台静态图优化、框架集成编程体验Python DSL算法思维多层 IR编译思维HLO IR框架约束硬件适配MLIR 方言栈灵活定制Relay/TOPI 多层抽象LLVM 专用后端生态定位PyTorch 生态的算子加速器多框架通用编译器TensorFlow/JAX 专属编译器Triton MLIR 跨硬件编译架构Triton MLIR 的学习路径学习路径可以遵循“从用户到开发者从抽象到具体”的原则分为四个层次层次一作为算子开发者用户体验层目标学会使用Triton Python DSL快速实现高性能内核。基础入门阅读Triton官方教程理解其编程模型tl.program_id,tl.arange,tl.load/store。亲手实现几个经典算子向量加法、矩阵乘法、Softmax。关键概念块Block、指针运算、掩码Mask、原子操作。进阶掌握学习内存优化技巧利用tl.cache,tl.make_block_ptr进行向量化加载/存储理解不同内存空间DRAM, L2 Cache, Shared Memory的差异。学习高级调度如流水线Pipeline和异步复制async_copy。学习如何将Triton内核封装为PyTorch的torch.autograd.Function并集成到模型中使用。实践尝试复现或优化一个经典论文中的融合算子如FlashAttention V1/V2的核心部分。层次二理解编译流程系统认知层目标理解从Triton Python代码到目标代码的完整编译过程。追踪编译链使用torch.compile或Triton的JIT功能运行一个简单内核并学习如何输出其关键中间表示。查看Triton IR了解高级Python操作如何被降低为Triton IR操作。查看MLIR理解Triton IR是如何被转换为MLIR中的Triton,LLVM,NVVM等Dialect的。这是理解其跨硬件能力的关键。学习MLIR基础概念无需深入建立认知即可什么是Dialect、Operation、Attribute、Type。MLIR的核心优势多层中间表示、可重用的转换和优化基础设施。实践使用TRITON_DUMP_GRAPH1等环境变量导出一个内核的编译流程图对照官方架构图理解每一步。层次三探索编译器开发贡献者/定制者层目标能够为Triton编译器添加新特性或为其定制新硬件后端。深入Triton编译器代码定位关键目录python/(DSL),lib/(核心编译器逻辑C)。学习从Python AST到Triton IR的生成过程。学习基于MLIR的Pass是如何对Triton Dialect进行优化的。学习为MLIR添加Pass编写一个简单的MLIR Pass对Triton Dialect的IR进行模式匹配和重写。后端开发入门以连接自研芯片为例路径一推荐将Triton Dialect lowering到已有的、更通用的MLIR Dialect如Vector,GPU然后利用社区或自定义的后端从这些Dialect生成代码。路径二更深入定义自己的硬件Dialect并实现从Triton Dialect到该硬件Dialect的转换规则。这需要对MLIR的方言转换框架Pattern Rewrite, Dialect Conversion有扎实理解。层次四架构与集成系统架构师层目标设计基于TritonMLIR的完整编译栈用于产品化部署。性能调优与自动化研究Triton的自动调优器Auto-Tuner原理学习如何为新的硬件特性添加调优空间。探索基于机器学习的内核自动生成与优化。端到端集成研究如何将Triton与模型级编译器如Torch-MLIR, IREE结合实现从PyTorch模型到混合内核部分由Triton生成的无缝编译。设计运行时Runtime接口管理自定义硬件上Triton内核的加载、启动和资源分配。紧跟社区关注Triton和MLIR的RFC、会议如MLIR Open Design Meetings、论文把握技术演进方向。Triton MLIR 动手开始学习路径有了那就动手开始吧。我根据个人情况想法是对四个层次竖着切一刀每一部分都实践一个简单用例先对整体建立概念再继续深入了解我关心的部分。每个人都可以依据自己的兴趣与实际情况去开始不必非得按部就班地学。结语有兴趣的小伙伴一起来学吧。
