필사 모드: 编写核函数的三个层次 — 用同一个问题比较 CUDA C++、Triton 与 CUTLASS

// softmax_cuda.cu
// 构建: nvcc -O3 -arch=sm_80 softmax_cuda.cu -o softmax_cuda
#include <cstdio>
#include <cfloat>
#include <cuda_runtime.h>

static const int WARP = 32;
static const int BLOCK = 256;             // 每个 block 有 8 个 warp

__inline__ __device__ float warpMax(float v) {
  for (int off = WARP / 2; off > 0; off >>= 1)
    v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, off));
  return v;
}
__inline__ __device__ float warpSum(float v) {
  for (int off = WARP / 2; off > 0; off >>= 1)
    v += __shfl_xor_sync(0xffffffffu, v, off);
  return v;
}

__global__ void softmaxKernel(float *out, const float *in, int nCols) {
  __shared__ float part[BLOCK / WARP];
  __shared__ float bcast;                 // 专门用于广播的槽位

  const long row = blockIdx.x;
  const float *src = in + row * nCols;
  float *dst = out + row * nCols;

  const int tid = threadIdx.x;
  const int lane = tid % WARP, wid = tid / WARP;
  const int nWarps = BLOCK / WARP;

  // 第 1 遍: 该行的最大值
  float m = -FLT_MAX;
  for (int i = tid; i < nCols; i += BLOCK) m = fmaxf(m, src[i]);
  m = warpMax(m);
  if (lane == 0) part[wid] = m;
  __syncthreads();
  if (tid == 0) {
    float t = part[0];
    for (int w = 1; w < nWarps; w++) t = fmaxf(t, part[w]);
    bcast = t;
  }
  __syncthreads();
  m = bcast;
  __syncthreads();                        // 在重用 part[] 之前必不可少

  // 第 2 遍: exp 的总和
  float s = 0.f;
  for (int i = tid; i < nCols; i += BLOCK) s += __expf(src[i] - m);
  s = warpSum(s);
  if (lane == 0) part[wid] = s;
  __syncthreads();
  if (tid == 0) {
    float t = 0.f;
    for (int w = 0; w < nWarps; w++) t += part[w];
    bcast = t;
  }
  __syncthreads();
  const float inv = 1.f / bcast;

  // 第 3 遍: 归一化并写出
  for (int i = tid; i < nCols; i += BLOCK) dst[i] = __expf(src[i] - m) * inv;
}

int main() {
  const int R = 8192, C = 4096;
  size_t bytes = (size_t)R * C * sizeof(float);
  float *h = (float *)malloc(bytes);
  for (long i = 0; i < (long)R * C; i++) h[i] = (float)((i * 37) % 100) * 0.01f;

  float *d_in, *d_out;
  cudaMalloc(&d_in, bytes); cudaMalloc(&d_out, bytes);
  cudaMemcpy(d_in, h, bytes, cudaMemcpyHostToDevice);

  for (int i = 0; i < 5; i++) softmaxKernel<<<R, BLOCK>>>(d_out, d_in, C);
  cudaDeviceSynchronize();

  cudaEvent_t a, b; cudaEventCreate(&a); cudaEventCreate(&b);
  cudaEventRecord(a);
  for (int i = 0; i < 50; i++) softmaxKernel<<<R, BLOCK>>>(d_out, d_in, C);
  cudaEventRecord(b); cudaEventSynchronize(b);

  float ms; cudaEventElapsedTime(&ms, a, b);
  double per = ms / 50.0;
  // 最小流量: 读一次、写一次。(虽然分 3 遍，但缓存会吸收掉)
  printf("cuda   %7.3f ms  %7.1f GB/s\n", per,
         2.0 * bytes / (per * 1e-3) / 1e9);
  return 0;
}

# softmax_triton.py
# 运行: python softmax_triton.py   (基于 triton 3.7.x、torch 2.13)
import torch
import triton
import triton.language as tl


@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_stride, out_stride, n_cols,
                   BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols

    # 把一整行一次性载入 SRAM。填充位置是 -inf，不影响 max/sum。
    x = tl.load(in_ptr + row * in_stride + cols, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    y = num / tl.sum(num, axis=0)
    tl.store(out_ptr + row * out_stride + cols, y, mask=mask)


def softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    num_warps = 4 if BLOCK_SIZE < 2048 else (8 if BLOCK_SIZE < 8192 else 16)
    out = torch.empty_like(x)
    softmax_kernel[(n_rows,)](
        out, x, x.stride(0), out.stride(0), n_cols,
        BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps,
    )
    return out


if __name__ == "__main__":
    torch.manual_seed(0)
    x = torch.randn(8192, 4096, device="cuda", dtype=torch.float32)

    ours, ref = softmax(x), torch.softmax(x, axis=1)
    assert torch.allclose(ours, ref, atol=1e-5), "正确性验证失败"

    gb = 2 * x.numel() * x.element_size() / 1e9
    for name, fn in [
        ("triton", lambda: softmax(x)),
        ("torch ", lambda: torch.softmax(x, axis=1)),
    ]:
        ms = triton.testing.do_bench(fn, warmup=25, rep=100)
        print(f"{name}  {ms:7.3f} ms  {gb / (ms * 1e-3):7.1f} GB/s")

Python 函数(@triton.jit)
   │  遍历 Python AST 生成 IR
   ▼
Triton IR(TTIR)           瓦片运算，尚无硬件概念
   │  布局分配、合并访存(coalescing)、流水线化 pass
   ▼
TritonGPU IR(TTGIR)       warp/线程布局与共享内存已确定
   │
   ▼
LLVM IR
   │  NVPTX 后端            │  AMDGPU 后端
   ▼                        ▼
PTX → (ptxas) → SASS       AMDGCN → 目标文件

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64},  num_warps=4, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64},  num_warps=8, num_stages=3),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=4),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128}, num_warps=4, num_stages=4),
    ],
    key=["M", "N", "K"],          # 这些值一变，就会重新调优
)
@triton.jit
def matmul_kernel(...):
    ...

import triton, torch

# 编译一次核函数，拿到句柄
k = softmax_kernel.warmup(
    torch.empty(1, 8, device="cuda"), torch.empty(1, 8, device="cuda"),
    8, 8, 8, BLOCK_SIZE=8, grid=(1,),
)
k._init_handles()

print(k.asm.keys())        # ttir, ttgir, llir, ptx, cubin
print(k.asm["ttgir"][:800])  # 用肉眼确认布局分配的结果
print(k.n_regs, k.n_spills)  # 一旦发生寄存器溢出，这里就不会是 0

// CUTLASS 3.x/4.x 风格 GEMM 配置的骨架
// 完整、可实际构建的示例请参考代码仓库中的 examples 目录。
using ElementA = cutlass::half_t;
using ElementB = cutlass::half_t;
using ElementC = float;

// 用多大的瓦片、分几个阶段的流水线来处理
using TileShape   = cute::Shape<cute::_128, cute::_128, cute::_64>;
using ClusterShape = cute::Shape<cute::_1, cute::_1, cute::_1>;

// 把主循环(数据搬运 + Tensor Core 累加)
// 和 epilogue(附加在结果上的后处理)分开组装
using CollectiveMainloop = /* CollectiveBuilder<...> */;
using CollectiveEpilogue = /* CollectiveBuilder<...> */;

using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
    cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;