CUDA黑魔法：假装CUTLASS库

摘要：很简单的一个黑魔法，只要在你的函数名前加上cutlass_，假装是CUTLASS库，有可能获得一定的性能提升。

我发现这个问题来源于Triton这笔commitade3d49

这不是在开玩笑，代码里真有。在FP8提升了 150 TFLOPS

那我们根据这个commit去build Triton，计算机没有magic，是可复现的。运行后看看最后的cubin吧，大小确实不一样。

cutlass_gluon_attention.ptx和gluon_attention.ptx 除了kernel name外，其他全部一样。

但是经过ptxas得到的cubin分别是321KB 和280 KB，函数名加上cutlass后编译器做了编译优化，展开了一部分东西，带来了代码的膨胀。

我没有B200，所以我没有复现性能影响，我在找其他可以带来优化的例子。发现这个的真是个人才，自古优化藏不住，唯有性能得人心。

更小复现案例

去装一下Triton实在太麻烦了，我们可以直接运行最后的ptxas命令，去复现下。

    def make_cubin(self, src, metadata, opt, capability):
        ptxas = get_ptxas().path
        with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc, \
            tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog:
            fsrc.write(src)
            fsrc.flush()
            fbin = fsrc.name + '.o'

            line_info = ["-lineinfo", "-suppress-debug-info"] if knobs.compilation.disable_line_info else ["-lineinfo"]
            fmad = [] if opt.enable_fp_fusion else ['--fmad=false']
            arch = sm_arch_from_capability(capability)

            # Disable ptxas optimizations if requested
            disable_opt = ['--opt-level', '0'] if knobs.nvidia.disable_ptxas_opt else []

            # Accept more ptxas options if provided
            ptx_extra_options = opt.ptx_options.split(" ") if opt.ptx_options else []

            ptxas_cmd = [
                ptxas, *line_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', fsrc.name, '-o',
                fbin
            ]
            try:
                subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog)

ptxas工具是来自于NVIDIA官网的，Triton的setup会下载

wget https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-12.8.61-archive.tar.xz
tar -xf cuda_nvcc-linux-x86_64-12.8.61-archive.tar.xz

得到ptxas工具后。你可以clone 我的项目，也可以把ptx下载下来做为输入，就可以得到cubin了。

git clone https://github.com/OpenMLIR/cuda-magic
cd cuda-magic
cuda_nvcc-linux-x86_64-12.8.61-archive/bin/ptxas -lineinfo -v --gpu-name=sm_100a triton_cache/gluon_attention/attention_kernel.ptx -o gluon_attention.cubin
cuda_nvcc-linux-x86_64-12.8.61-archive/bin/ptxas -lineinfo -v --gpu-name=sm_100a triton_cache/cutlass_gluon_attention/attention_kernel.ptx -o cutlass_gluon_attention.cubin

命令执行完成后你可以用ls -lh命令看到2个cubin的大小不一样。

做为编译器工程师，我的疑问是为什么不在PTX文件里加入提示，真是个耐人寻味的问题。