3. 功能支持

3.1. 基本功能

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

3.2. 扩展功能

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

3.2.1. triton.Config

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

3.2.1.1. pipeline

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

  • Triton 2.1

    triton.Config({...},...,pipeline="basic",...)
    
  • Triton 3.0

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

支持参数:

  • basic:基础N-buffer优化。其中 num-stages 大于 2 的buffer数据会暂存在寄存器当中。 basicpipeline 的默认配置。

  • 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 2.1

    triton.Config({...},...,scenario="flashattn-fwd",...)
    
  • Triton 3.0

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

支持参数:

scenario 可以搭配 pipeline 使用,当前版本曦云系列GPU支持以下参数设置:

  • flashattn-fwd:针对flashattn前向类型的Triton算子进行编译层面的重排和优化。

  • flashattn-bwd:针对flashattn反向类型的Triton算子进行编译层面的重排和优化。

  • mla:针对mla或者extendattn前向类型的Triton算子进行编译层面的重排和优化。

  • unrollpipelinecpasync 时,对Triton算子中的循环进行展开(默认为不展开); pipelinebasic 时,Triton算子中的循环默认是展开的; pipeline 为空时无效。

  • rollpipelinebasic 时,对Triton算子中的循环进行不展开(默认为展开); pipelinecpasync 时,Triton算子中的循环默认是不展开的; pipeline 为空时无效。

  • unprefetch:当 pipelinebasic 时,不会进行prefetch操作。

  • fullstage: 使用更激进的策略对更多的 Op 进行 N-buffer pipeline 来掩盖访存带来的延迟,同时会占用更多的设备资源。当 pipelinebasic 时,开启会占用更多的寄存器;当 pipelinecpasync 时,开启会占用更多的共享内存。

  • storeCoalesce:优先开启全局内存写入时的大位宽优化。

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

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

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

  • flashattn-fwd;flashattn-bwd;mla

  • unroll;roll

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

3.2.2. triton.autotune

triton.autotune 是Triton中用于自动调优的装饰器。 用户提供可选的 triton.Config 列表,在运行时,Triton会执行“tune”过程,从这些配置中选取最合适的配置来运行Kernel,并将结果缓存下来。

为了进一步提升性能,mcTriton提供将自动调优(autotune)结果持久化至硬盘的功能。 通过此功能,程序可以在每次重新运行时跳过“tune”阶段,直接使用之前运行时保存的自动调优结果。操作步骤如下所示。

  1. 开启自动调优结果持久化功能的环境变量。

    export TRITON_ENABLE_PERSISTENT_AUTOTUNE_CONFIGS=1
    
  2. (可选)配置持久化路径。未配置时默认保存至$HOME/.triton/cache/configs

    export TRITON_AUTOTUNE_CONFIG_PATH=$USER_DEFINED_PERSISTENT_PATH
    
  3. 运行Kernel。

    Kernel运行后会将自动调优结果保存到持久化路径下,后续程序再次运行时优先从该路径读取之前的缓存结果。