3. 功能支持

3.1. 基本功能

mcTriton以兼容Triton的原生使用方式为设计目标。大部分情况下，用户可以参考官方文档获得mcTriton的使用方式。

3.1.1. tl.dot的数据类型支持

tl.dot 是Triton中用于矩阵乘的核心算子。C50X系列与C60X系列产品对 tl.dot 的数据类型支持情况对比如下：

表 3.1 C50X系列与C60X系列的tl.dot数据类型支持对比
数据类型	C50X系列	C60X系列
`fp32`	支持（mma）	支持（fma）
`tf32`	支持（mma）	支持（mma）
`fp16`	支持（mma）	支持（mma）
`bf16`	支持（mma）	支持（mma）
`int8`	支持（mma）	支持（mma）
`fp8`	不支持	支持（mma）

3.2. 扩展功能

本章介绍mcTriton相较官方Triton提供的额外功能。

3.2.1. triton.Config

出于性能调优的目的，mcTriton扩展了 triton.Config 中的配置选项，提供给用户做不同场景下的性能调优。

3.2.1.1. pipeline

pipeline 选项用于指定在for循环中 triton.language.dot 场景下使用的不同优化方式。具体使用方法：

triton.Config({...,'pipeline': "basic",...},...)

支持参数：

basic：基础N-buffer优化。其中 num-stages 大于 2 的buffer数据会暂存在寄存器当中。 basic 是 pipeline 的默认配置。
cpasync：使用cp.async功能在pipeline中将数据直接从全局内存拷贝到共享内存进行N-buffer优化，该配置下 num-stages 增大的同时，共享内存用量也会同步变大。
当 pipeline 设置为空时，关闭所有N-buffer优化。

对于不同的Triton kernel，可能获得最优性能的 pipeline 参数是不同的，可以配合 triton.autotune 搜索得到当前kernel的最佳pipeline。

3.2.1.2. scenario

scenario 选项用于指定在 for 循环中 triton.language.dot 场景下一些特定的编译选项和指令重排方式。具体使用方法：

triton.Config({...,'scenario': "flashattn-fwd",...},...)

支持参数：

scenario 可以搭配 pipeline 使用，当前版本沐曦通用GPU支持以下参数设置：

flashattn-fwd：针对flashattn前向类型的Triton算子进行编译层面的重排和优化。
flashattn-bwd：针对flashattn反向类型的Triton算子进行编译层面的重排和优化。
mla：针对mla或者extendattn前向类型的Triton算子进行编译层面的重排和优化。
unroll： pipeline 为 cpasync 时，对Triton算子中的循环进行展开（默认为不展开）； pipeline 为 basic 时，Triton算子中的循环默认是展开的； pipeline 为空时无效。
roll： pipeline 为 basic 时，对Triton算子中的循环进行不展开（默认为展开）； pipeline 为 cpasync 时，Triton算子中的循环默认是不展开的； pipeline 为空时无效。
unprefetch：当 pipeline 为 basic 时，不会进行prefetch操作。
fullstage: 使用更激进的策略对更多的 Op 进行 N-buffer pipeline 来掩盖访存带来的延迟，同时会占用更多的设备资源。当 pipeline 为 basic 时，开启会占用更多的寄存器；当 pipeline 为 cpasync 时，开启会占用更多的共享内存。
storeCoalesce：优先开启全局内存写入时的大位宽优化。

备注

Triton3.6分支不再支持flash-attn以及相关的 scenario 参数（包括 flashattn-fwd、flashattn-bwd、mla）。Triton3.0与Triton3.6对 scenario 参数的支持情况对比如下：

表 3.2 Triton3.0与Triton3.6的scenario参数支持对比
scenario参数	Triton3.0	Triton3.6
`unroll`	支持	支持
`roll`	支持	支持
`unprefetch`	支持	支持
`fullstage`	支持	支持
`storeCoalesce`	支持	支持
`flashattn-fwd`	支持	不支持
`flashattn-bwd`	支持	不支持
`mla`	支持	不支持

多个scenario的参数设置可以进行叠加成为scenario组合，组合中每个参数之间使用 “;” 进行间隔，例如：

triton.Config({...,'scenario': "unprefetch;roll;fullstage",...},...)

以下scenario之间存在冲突，避开组合当中同时出现：

flashattn-fwd;flashattn-bwd;mla
unroll;roll

对于不同的Triton kernel，可能获得最优性能的scenario组合是不同的，可以配合 triton.autotune 搜索得到当前kernel的最佳scenario组合。

小技巧

对于类似矩阵乘的单 dot 计算，使用大tile（如 128*128*128）时，建议设置 pipeline="cpasync"，优化效果更好。

autotune用例

triton.autotune 配合 triton.Config 的使用示例如下：

@triton.autotune(
    configs=[
        triton.Config({...,'pipeline': "basic",'scenario': "unprefetch",...}, num_warps=4, num_stages=2),
        triton.Config({...,'pipeline': "cpasync",...}, num_warps=8, num_stages=3),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def matmul_kernel(...):
    ...

3.2.2. MCTRITON_LOG_LEVEL

Triton3.6支持通过 MCTRITON_LOG_LEVEL 环境变量开启Triton代码中的日志打印功能，日志级别由高到低依次为 TRACE、DEBUG、INFO、WARN、ERROR、CRITICAL，级别越高，打印的日志内容越多。

表 3.3 MCTRITON_LOG_LEVEL支持的日志级别
日志级别	说明
`TRACE`	打印全部跟踪信息，如layout、Tensor的shape等。
`DEBUG`	打印用于确认优化和函数执行的调试信息及关键数值。
`INFO`	打印基础信息，如config设置、Kernel的寄存器/共享内存使用情况等。
`WARN`	打印编译器行为可能导致性能不达预期的消息。
`ERROR`	打印编译器行为将导致错误的消息。
`CRITICAL`	仅打印用户自定义的消息。

使用示例

设置 MCTRITON_LOG_LEVEL 环境变量后运行Kernel：

export MCTRITON_LOG_LEVEL=INFO

每条打印信息包含时间、打印位置、进程号、logger名称、日志级别和消息内容。