TLX以及开箱即用, 原文地址为博客, 作者包括、普扬·洛菲、伊恩、谢恩·奈、、、 池标点符号使用有些混乱, 可能需整理。注解: 此篇文章乃是针对公众号阅读所进行的中文翻译、译述文稿, 留存了原文的技术主要脉络、接口、性能方面的数据以及参考资料, 然而它并非逐句严格按照字面意思进行的翻译。外部链接统一被写作“文本链接”这样的形式, 正文当中的图片已经转存至图床。【博客翻译, 译述】, 插件扩展, 开箱即用的 TLX, 与自定义编译器 Pass。TL;DR- 3.7进行了引入, 这一引入使得其能够允许在运行期间加载编译器pass, 还能加载MLIR以及它的op, 甚至可以达到扩展层的DSL, 并且在这个过程中不用去维护fork, 也无需为了安装插件而重新开展编译。率先使用这套机制里, Meta的TLX是首个主要使用者。往昔运用TLX时, 往往得构建Meta的实验性fork如今它能够当作独立的-utlx包来安装, 并借助将其加载到上游。博客提供的测试表明, 这条动态加载路径跟原fork产生相同的PTX/, 于H100和MI350上未察觉到额外性能损失。此次改动所解决的, 是生态当中一个相当实际的问题, 即高性能, 常常需要定制编译流水线, 然而过去一旦跨越上游已有的能力范畴, 团队就极易趋向于长期维护分支。1. 为什么 需要插件系统设为默认的编译流水线涵盖了诸多常见情形 , 然而在生产级别上进行性能优化时往往还得进一步深入推进。比如说。过往常常出现的这种做法是进行fork, 将、pass以及还有是直接编译进去。在短时段之内它是十分灵活的, 然而后续的维护工作却特别繁重: 上游方面的升级会导致merge、API产生变化以及行为出现差异在进行fork之后固定于旧版本的状况下, 同样无法获得新的硬件予以的支持以及bug fix。一个边界被换了, core提供出稳定的加载以及hook机制, 扩展作为独立进行发布, 研究者能够持续开发自定义pass和op, 使用者依旧运行上游。2. 插件如何接入 编译流水线插件从本质上来说, 是在运行的时候进行加载的那类.so 文件, 在把插件包安装完成之后, 要将动态库路径给写入进去, 当使得启动编译的时候 就可以发现并且对它进行加载。es .()你提供的内容似乎不太完整且存在错误, 不太能按照正确需求进行改写, 请提供完整无误的内容。插件系统的关键并非仅仅是“再多注册几个op”, 在.py的各个阶段添加了hook, 扩展能够参与从高层IR直至目标代码的一整条路径。指标IR即TTIR呈现下降态势, 指标IR也就是TTGIR同样呈现下降态势, LLVM IR出现下降, PTX /。通过这些 hook插件可以和 AMD 都支持这套机制。2.1 三个扩展层级原文把 API 能力分成三层。第一层存在自定义 pass, 这个自定义 pass 并不需要引入新的内容, 它呢适合应用于对已有 IR做局部转换。第二层是自定义的MLIR以及pass, 插件能够将独立编译的内容加载进去, 接着把标准IR重写成自定义op, 而后交给专门的去处理。顶层 DSL op 在第三层之处。新的语法跟语义于侧能够被引入用以展开扩展, 如此这般就能让作者去采用新的编程这种可以成为抽象的方式, 而采用这种方式不必对本体进行修改。2.2 可以按 切换在插件方面, 并非规定整个进程只能运用一套编译流水线, 代码能够进行设置hook, 自设置hook之后, 后续调用的部分采用自定义方式, 在取消hook操作之后, 再回归到默认路径。插件的个数以及自定义的数量不存在硬性的限定数值。插件承担着缓存管理的职责具体指, 只有当配置实实在在产生变化的时候才理应引发重新编译的操作。关于 TLX 的 utlx 库已将这一部分进行了封装处理, 一般的用户无需自行去实现 cache hash。3. TLX 提供了哪些能力TLX 是面向那些, 有着需要进行显式控制内存, 以及数据搬运, 还有异步执行情况的。它会将某些, 以往只是能够在 fork 里面去使用的底层能力, 给暴露成为 DSL op。有一种办法能用一些软件流水线, 它能让它的制作者自身去对好多缓冲运用表达功能体现过程, 先做出预先取来接下来的小图像切块的动作, 接着进行那个当下的小图像切块所需要的计算, 还要凭借着“_wait”去控制存在的依赖状况。有一种编译器还是负有相应责任的 , 只是其开展调度的意图已经不全部依靠着靠着一般常识作出的推理判断作为依据去进行了。4. 同一套 TLX如何映射到 H100 和 MI350TLX的接口能够跨越多种情况和AMD来运用, 其底层的实现是依据硬件各自进行区分的。4.1 H100TMA WGMMA GEMM在于其上, TLX的异步负载能够被映射至TMA, 矩阵乘加可用WGMMA。GEMM会先期分派多个, 于其中预取前几级数据, 接着在K循环内交替执行:等待当前计算所需的 load group。对已经就绪的 发起异步 dot。把下一块 A/B tile 预取进即将复用的 。循环结束后等待所有 dot并写回结果。这般能够将, 搬运和, Core计算叠加一块儿。原文对比了, H100上, FP16 GEMM的, 以及, TLX。把这组数据解读作为“插件加载没有拖慢原来的 TLX ” , 会更合适一些。在不同 shape 上, TLX 和各有所得得失存在差异大多数是落在几个百分点以内。4.2 MI350通过寄存器组织流水线MI350所走的是另外一条数据路径, 先是运用普通load将下一块数据读进寄存器, 接着再借由写入, 在计算当前tile之际, 然后利用把已经准备妥当的数据取回寄存器并且执行tl.dot。管理方式依然是多级流水线, 仅是数据搬运不依赖那个的TMA/WGMMA。博客所给出的方阵GEMM数据如下:于这四个测试尺寸之处, 相较于TLX而言, 高出了百分之十一点八至百分之十五点二。5. 从 到 GPU MODE团队针对GPU MODE所对应的任务, 对插件路径展开了验证。此任务涵盖五个GEMM, 还有一个, 会有输出, 并且包含layer norm、gate以及, 较为贴近真实的复合。19.2毫秒是 pile 基线所耗时间, 其主要时间耗费在了 GEMM 上。加载 TLX 插件之后, 参赛实现运用了 warp - GEMM, 并且持续融合周围算子, 最终达成了12.0毫秒, 相较于 pile 基线而言加速了1.61倍。依旧是以平常的样子现身于编译链之中的 TLX, 所以优化能够持续拓展至围绕 GEMM 的算子融合。后来团队又对 B200 添加了 CLC。GPU MODE 任务说明6. 动态加载是否会改变专门针对插件版本以及Meta fork最终所输出的两者展开了比较, 得出的结果是:其原因相对直面, TLX pass以及op的加载形式产生了改变, 然而它们依旧归入同一套MLIR, 插件在进行执行期间并未增添中间层次。曾保留住了最初的硬件优化, 仅仅是将“扩展怎样进入 ”的途径从 fork转变为运行时插件, 这般情形极为实用。7. 如何开始使用原文所给出的安装方式, 是要先去构建能够启用扩展支持的, 然后呢, 要再去安装那种在PyPI上的TLX插件哦:译者注, 时间为2026年7月24日, 当前PyPI上的-utlx 3.7.1含有原生.so, 且与ABI绑定。项目页要求必须使用上游的v3.7.0 tag, 并且记录了该tag中一个会对插件op返回值产生影响的一行修复。在实际进行安装时, 应当以PyPI项目的说明作为准则, 千万不要随意去混用其他的东西。设置插件路径后TLX op 就可以在 中使用es .()os.将空字符串、.so通过os.path.join进行连接, 其作用等同于tlx。相关资料8. 接下来还会扩展到哪里原文列出的后续方向包括当下, 这些方向处于各异的开发阶段, 并非都能被看成是 3.7 已然交付的功能。已实现的是插件基础设施, 还有 TLX 作为独立包接入上游的途径。小结将, 与core的版本维护, 进行了拆分。开发者能够持续编写, 与硬件相关的、pass以及DSL op, 使用者无需, 为了一个扩展, 而长期绑定, 某个fork。TLX 给出了首个完整样例, 同一编程模型, 在 H100 上可调用 TMA/WGMMA, 于 MI350 上依循寄存器, 动态插件版本和原来的 fork 保持一致 , 对于那些要长期追踪上游, 且必定得保留定制编译能力的团队而言, 相较于维护一个愈发难以升级的 fork 要更简便。