Triton黑魔法：多层级 runner 工具

摘要：我开发了一个 Triton 的多层级 runner 工具，支持 ttir、ttgir、llir、ptx、cubin 等多个编译阶段的运行和调试。相比直接维护 Triton 的 Python 源码，这个工具更关注稳定的中间产物（IR 和二进制），让 Triton 的kernel内核调试、部署和研究更加轻量、可靠、高性能。

项目地址：OpenMLIR/triton-ml-runner

triton-ml-runner(Triton multi-level runner,TMLR)是一个面向 OpenAI/Triton 的多层级 runner 工具，用于调试 Triton IR，支持在多个编译阶段直接运行 GPU kernel，包括 ttir、ttgir、llir、ptx、cubin。该工具旨在提升 Triton 用户对编译流程的可观测性与可控性，同时降低对 Triton 源码的编译Pass pipeline的限制，提升调试开发与部署效率。

一、项目介绍

本项目是一个 Triton 编译流程中的 多层级kernel运行工具，主要支持以下中间表示或输出格式：

ttir（Triton IR）
ttgir（Triton GPU IR）
llir（LLVM IR）
ptx（NVIDIA PTX 汇编）
cubin（NVIDIA GPU 可执行二进制）

与直接使用 @triton.jit 源码不同，该工具可以脱离 @triton.jit 修饰的Python kernel依赖，聚焦于更稳定、可复用的中间产物和二进制运行过程。这对kernel调试、cubin部署、自定义 Pass 开发等场景尤为重要。

run 示例

gluon/01-attention-forward.py 是Triton正在实现的高性能 attention kernel，使用了Gluon 模块来生成ttgir，相当于用户可以使用Python来写Triton GPU IR来获得更好的性能。由于直接面向IR编程了，layout、shared memory 以及fence均可以由用户控制，由Jeff Niu 负责此模块的开发。Jeff Niu 的IRDL: an IR definition language for SSA compilers 论文。我这个项目将为 Gluon 提供支持。

What is Gluon?

二、使用方式

所有runner最终执行的均是cubin_runner，IR输入均编译为cubin再运行。除了本身的文件依赖外，ttir_runner和ttgir_runner需要提供优化选项options，llir_runner、ptx_runner和cubin_runner需要提供响应的metadata.json文件，Python代码里仅需提供输入文件的文件夹路径即可。

triton-ml-runner 目前仅支持 Triton v3.3.x，其他版本还在开发。

使用前先安装，这里建议pip使用-e(editable)以可编辑模式安装一个本地包，因为项目并不稳定，你可能需要修改我的代码。pip install会将src文件夹下内容安装到你的Python环境中。

git clone https://github.com/OpenMLIR/triton-ml-runner
cd triton-ml-runner
pip install -e .

安装好后，可以按照README去运行相应arch的runner的example，目前提供了sm90(H20)、sm86(A10)、sm75(Tesla T4) 这3个compute capability 的，比如H20可以运行如下命令。

python examples/ttir_runner/matmul.py

python examples/ttgir_runner/sm90/matmul.py

python examples/llir_runner/sm90/matmul.py

python examples/ptx_runner/sm90/matmul.py

python examples/cubin_runner/sm90/matmul-with-tma-v2.py

几个arch target的Python代码是一致的，如果没有你的arch，你需要使用TRITON_CACHE_DIR=$PWD/.cache 得到对应的源文件之后再运行。我目前手上只有这些卡，有sm_100(B200)可以共享的欢迎私信我。

1、cubin_runner

Triton黑魔法：cubin runner 我结合Triton源码分析了这层级runner的设计。这里介绍下使用，我们以cubin_runner/sm90/matmul-with-tma-v2.py 为例。

之前的编译和 launch 是JIT(Just in time, 即时编译)执行，grid 是 grid launch 参数，指定线程块布局。后面是具体的函数参数，和我们的@triton.jit修饰的函数一一对应。这个写法会触发 JIT 并运行。

    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']), triton.cdiv(K, META['BLOCK_SIZE_K']), )
    matmul_kernel_make_tensor_desciptor[grid](
        a, b, c,
        M, N, K,
        BLOCK_SIZE_M=128,
        BLOCK_SIZE_K=64,
        BLOCK_SIZE_N=64,
    )

我们为了调用需要准备kernel_name和对应的.cubin和.json文件。为了方便文件管理，这里均在当前Python文件所在文件夹处理。我们需要将之前的@triton.jit修饰的kernel换成@triton_ml_runner.jit，并将对应的文件路径做为cubin_dir参数传给kernel。

import triton_ml_runner

# @triton.jit
@triton_ml_runner.jit
def matmul_kernel_make_tensor_desciptor(a_ptr, b_ptr, c_ptr,  #
                                        M, N, K,  #
                                        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
                                        BLOCK_SIZE_K: tl.constexpr,  #
                                        ):
...

def matmul(a, b):
...
    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']), triton.cdiv(K, META['BLOCK_SIZE_K']), )
    import os
    current_dir = os.path.dirname(os.path.abspath(__file__))
    matmul_kernel_make_tensor_desciptor[grid](
        a, b, c,
        M, N, K,
        BLOCK_SIZE_M=128,
        BLOCK_SIZE_K=64,
        BLOCK_SIZE_N=64,
        cubin_dir=current_dir
    )

代码编写结束后使用Python调用这个Python程序就可以了。

python examples/cubin_runner/sm90/matmul-with-tma-v2.py

2、ttir_runner

ttir(Triton IR) 是架构无关(target-independence)的，所以你可以把的生成的ttir给任意arch去编译，所有arch共用了一份examples/ttir_runner/matmul_kernel.ttir。Python程序本身就是共享的。

同理这里需要你传入ttir_dir参数，不影响你原来的优化options设置。

# @triton.jit
@triton_ml_runner.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr
):
...
def matmul(a, b):
...
    import os
    current_dir = os.path.dirname(os.path.abspath(__file__))
    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=16,
        BLOCK_SIZE_N=16,
        ttir_dir=current_dir
    )

3、ttgir_runner

ttgir(Triton GPU IR)是架构有关的，会根据你的compute capability进入不同的Pass。这里不同的是有些编译选项已经在kernel里了，比如"ttg.num-warps" = 4和ttg.target = "cuda:90"等已经是moduleOp的attributes了。传入的xx_dir参数变化为了ttgir_dir

    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=16,
        BLOCK_SIZE_N=16,
        ttgir_dir=current_dir
    )

完整Python代码在examples/ttgir_runner/sm90/matmul.py:66

4、llir_runner

由于转换为llir将shared等信息都放进metadata了，所以还需要编译时的matmul_kernel.json。传入的xx_dir参数变化为了llir_dir

    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=16,
        BLOCK_SIZE_N=16,
        llir_dir=current_dir
    )

完整Python代码在examples/llir_runner/sm90/matmul.py:66

5、ptx_runner

ptx 同上，有变动的Python代码如下。

    matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_SIZE_M=16,
        BLOCK_SIZE_N=16,
        ptx_dir=current_dir
    )

完整Python代码在examples/ptx_runner/sm90/matmul.py:66

6、该选择哪个runner

示例中所有IR runner使用的均是Triton cache中的IR，如果使用MLIR_ENABLE_DUMP=1获得其中一个Pass执行后的代码呢？这个问题并不简单。

比如matmul-with-tma-v2.py 在B200上编译可以获得with_tma_v2/MLIR 文件夹下的80个IR文件。

你可以根据深度剖析 Triton编译器 MatMul优化（三）—— TMA 或 make_xx 看下它在哪个阶段。比如你修改了43-TritonGPUPipeline.mlir GPU 流水后的IR，由于还未执行TritonGPUFenceInsertion Pass, 这是个对ttng.tc_gen5_mma前插入了ttng.fence_async_shared来控制异步 shared-memory 操作都完成的Pass，会影响正确性。所以你应该执行的是ttir_runner。

三、更新计划（Roadmap）

[Roadmap] Plan of TMLR v0.2.0

1、易用性提升

参考 JITFunction 使用decorator来去掉signature_str，另外还需要简化现在的runner流程，现在略微复杂了些。

7.13更 已基本解决，易用性得到了非常大的提升。

2、bench解决

现有runner流程没有cache，对bench有影响，需要开发 or 复用 Triton 源码。

3、更多Triton版本支持

目前仅支持Triton v3.3.x，Triton v3.4.0发布在即，会进行跟进。另外由于Triton v3.3.x 部分编译Pass在旧卡上会获得负收益，比如更慢or更多的shared 占用, Triton v3.2.0 也会尽力支持。多版本支持后现在的代码可能也面临重构问题。

4、pass stage

提供自定义的编译stage接口，用户可以控制执行哪些Pass。Pass的添加本项目不会提供，要做的话会新开一个项目。

5、更多

项目初期难免被bug困扰，欢迎使用者来提issue，我会在空闲时尽量支持。有开发能力的也欢迎来提pull request。

四、相关文章

Triton黑魔法：cubin runner

深度剖析 Triton编译器 MatMul优化（三）—— TMA

深度剖析 Triton编译器 MatMul优化（二）—— MMA

深度剖析 Triton编译器 MatMul优化（一）—— FMA

浅析 Triton 执行流程

附录

1、优化选项

kernel JIT运行可以带的默认值

options = {"num_warps": 4, "num_ctas": 1, "num_stages": 3, "enable_fp_fusion": True, "launch_cooperative_grid": False}

CUDAOptions里有的，存储在metadata里

class CUDAOptions:
    num_warps: int = 4
    num_ctas: int = 1
    num_stages: int = 3
    warp_size: int = 32
    # maxnreg corresponds to the ptx parameter .maxnreg, which controls the
    # maximum number of 32-bit registers used by one thread.
    maxnreg: Optional[int] = None
    cluster_dims: tuple = (1, 1, 1)
    ptx_version: int = None
    ptx_options: str = None
    ir_override: Optional[str] = None  # filename of a user-defined IR (*.{ttir|ttgir|llir|ptx})
    enable_fp_fusion: bool = True
    launch_cooperative_grid: bool = False
    launch_pdl: bool = False
    supported_fp8_dtypes: Tuple[str] = ("fp8e5", "fp8e4b15")
    deprecated_fp8_dot_operand_dtypes: Tuple[str] = ()
    default_dot_input_precision: str = "tf32"
    allowed_dot_input_precisions: Tuple[str] = ("tf32", "tf32x3", "ieee")
    max_num_imprecise_acc_default: bool = None
    extern_libs: dict = None
    debug: bool = False
    backend_name: str = 'cuda'
    sanitize_overflow: bool = True
    arch: str = None

posted @ 2025-07-12 17:46 暴力都不会的蒟蒻阅读(164) 评论(0) 收藏举报

刷新页面返回顶部

BobHuang