Kernel Backend、Triton Blocking 与 Persistent Kernel

它位于推理栈的哪一层

对 torch.softmax(x, dim=-1) 或 x @ weight，用户先声明数学语义；framework 再根据输入和运行时状态选择 backend，最终启动一个或多个具体 kernels：

API semantics
  softmax / matmul / scaled_dot_product_attention
        ↓
framework operator、graph lowering 与 dispatch policy
  ATen / Inductor / backend selection
        ↓
实现资源
  cuBLAS(Lt) library · CUTLASS templates · Triton JIT · FlashInfer primitives
        ↓
当前 shape、dtype、layout、GPU 对应的 concrete kernel(s)

本页讨论最下面两层怎样把 tensor 映射到 GPU 工作单元。它不重新定义 softmax、GEMM 或 attention，也不等同于 Backend Dispatch/Fallback：dispatch 回答“为什么选择这条路径”，kernel backend 回答“这条路径怎样分块、读取、归约和写回”。

从 stable row softmax 看 blocking

输入 $X\in\mathbb{R}^{R\times N}$，逐行 stable softmax 为：

教学 kernel 令一个 Triton program instance 拥有一整行，launch grid 为 (R,)。第 $r$ 个 program 生成 offsets = 0..BLOCK_SIZE-1，其中：

offsets < N 的 lanes 才能读写真实元素；其余 load 为 $-\infty$，因而不改变 max 和指数和。tl.max 与 tl.sum 在 FP32 中归约，最后 store 回输入 dtype。

BLOCK_SIZE 是编译期 tile 宽度，不是“启动了 1024 个 CUDA threads”的同义词；num_warps=4 才指定一个 program 使用多少个 warps 协作。Triton compiler 仍会决定指令、register/shared-memory 使用和最终 CUDA launch 细节。

小尺寸 worked example

取一行 X[0,:]=[0,1,2,3,4]，$N=5$，所以 BLOCK_SIZE=8：

offset       0       1       2       3       4       5       6       7
loaded       0       1       2       3       4      -∞      -∞      -∞
valid?       ✓       ✓       ✓       ✓       ✓       ✗       ✗       ✗
softmax   .0117   .0317   .0861   .2341   .6364       0       0       0
store        ✓       ✓       ✓       ✓       ✓       ✗       ✗       ✗

$m=4$ 避免直接计算较大指数；三个 masked lanes 既不参与有效归约，也不会越界访问。对 $N=1000$，tile 补到 1024，每个 program 有 24 个 masked lanes，即 2.4% 的 lane 宽度。若 $N$ 刚超过 1024，补到 2048 会明显增加无效工作和 register pressure，单行单 program 策略未必仍合适。

完整可运行的 Triton 实验

实验问题：一个 blocked Triton kernel 能否保持 PyTorch stable softmax 语义？本环境下它实际启动什么 kernel，速度是否一定超过 PyTorch backend？

运行条件：需要 CUDA 与 Triton；脚本固定 seed，以 [4096,1000] FP16 contiguous tensor 运行。A/B 先 assert_close，然后各自 warmup 25 次，用 CUDA Event 测 100 轮中位数；profiler 只用于 kernel 名取证，不混入计时。

import platform,statistics
import torch,triton
import triton.language as tl
from torch.profiler import ProfilerActivity,profile

torch.manual_seed(223)
assert torch.cuda.is_available(),"This Triton experiment requires a CUDA GPU"
DEVICE=torch.device("cuda"); DTYPE=torch.float16

@triton.jit
def row_softmax_kernel(output_ptr,input_ptr,input_row_stride,output_row_stride,
                       n_cols:tl.constexpr,BLOCK_SIZE:tl.constexpr):
    row=tl.program_id(axis=0)
    offsets=tl.arange(0,BLOCK_SIZE)
    mask=offsets<n_cols
    values=tl.load(input_ptr+row*input_row_stride+offsets,
                   mask=mask,other=-float("inf")).to(tl.float32)
    values=values-tl.max(values,axis=0)
    numerators=tl.exp(values)
    probabilities=numerators/tl.sum(numerators,axis=0)
    tl.store(output_ptr+row*output_row_stride+offsets,probabilities,mask=mask)

def triton_softmax(x):
    assert x.ndim==2 and x.is_cuda and x.is_contiguous()
    n_rows,n_cols=x.shape
    block_size=triton.next_power_of_2(n_cols)
    assert block_size<=65536,"The teaching kernel keeps one row in one program"
    num_warps=8 if block_size>=2048 else 4
    output=torch.empty_like(x)
    row_softmax_kernel[(n_rows,)](
        output,x,x.stride(0),output.stride(0),n_cols=n_cols,
        BLOCK_SIZE=block_size,num_warps=num_warps)
    return output

def benchmark_us(fn,warmup=25,rounds=100):
    for _ in range(warmup): fn()
    torch.cuda.synchronize(); samples=[]
    for _ in range(rounds):
        start=torch.cuda.Event(enable_timing=True)
        end=torch.cuda.Event(enable_timing=True)
        start.record(); fn(); end.record(); end.synchronize()
        samples.append(start.elapsed_time(end)*1000)
    return statistics.median(samples)

def cuda_kernel_names(fn):
    with profile(activities=[ProfilerActivity.CPU,ProfilerActivity.CUDA]) as prof: fn()
    torch.cuda.synchronize()
    return sorted({e.name for e in prof.events()
                   if str(e.device_type)=="DeviceType.CUDA"})

def main():
    n_rows,n_cols=4096,1000
    x=torch.randn(n_rows,n_cols,device=DEVICE,dtype=DTYPE)
    reference=torch.softmax(x,dim=-1); actual=triton_softmax(x)
    torch.testing.assert_close(actual,reference,rtol=2e-3,atol=2e-3)
    torch.testing.assert_close(actual.float().sum(-1),torch.ones(n_rows,device=DEVICE),
                               rtol=2e-3,atol=2e-3)
    torch_us=benchmark_us(lambda:torch.softmax(x,dim=-1))
    triton_us=benchmark_us(lambda:triton_softmax(x))
    kernels=cuda_kernel_names(lambda:triton_softmax(x))
    block_size=triton.next_power_of_2(n_cols)
    print(f"python={platform.python_version()} torch={torch.__version__} triton={triton.__version__}")
    print(f"device={torch.cuda.get_device_name(0)} capability={torch.cuda.get_device_capability(0)}")
    print(f"X={tuple(x.shape)} dtype={x.dtype} layout=contiguous row_stride={x.stride(0)}")
    print(f"grid=({n_rows},) BLOCK_SIZE={block_size} masked_lanes_per_program={block_size-n_cols}")
    print(f"max_abs_diff={(actual-reference).abs().max():.3e}")
    print(f"row_sum_max_error={(actual.float().sum(-1)-1).abs().max():.3e}")
    print(f"torch_softmax_median_us={torch_us:.3f} triton_softmax_median_us={triton_us:.3f}")
    print(f"triton_profiler_cuda_kernels={kernels}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/kernel_backend_triton.py。2026-08-13 在 Python 3.11.12、PyTorch 2.13.0+cu130、Triton 3.7.1、RTX 3080 上，最大绝对差为 3.815e-06，最大行和误差为 4.423e-05；profiler 观察到 row_softmax_kernel。PyTorch / 教学 Triton 中位数分别为 31.744 / 35.840 µs：本 shape 上教学实现更慢。

这不是跨 GPU 的 softmax 排名。计时发生在 warmup/JIT 完成之后，两个函数都返回新输出；caching allocator 已被预热。tl.exp 的近似、归约顺序和 FP16 store 会造成低位差异。实验没有证明 attention/GEMM kernel 性能，也没有证明换一个 $R,N$、dtype、stride 或 Triton/CUDA 版本仍保持同一排序。

Backend 怎样选择 blocking

Tile、program、warp 与硬件 block

• Tile 是算法的数据分块，例如 [BLOCK_M,BLOCK_N,BLOCK_K] GEMM 子矩阵；它决定数据复用和边界 mask。
• Triton program 是一份 SPMD 实例，处理一个或多个 tiles；launch grid 决定有多少实例。
• Warp 是 NVIDIA GPU 的 32-lane 执行组；num_warps影响一个 program 的并行度和资源占用。
• CUDA thread block / CTA 是具体 CUDA 执行与协作单元。Triton program 通常会 lowering 到 CTA，但不能把 DSL 中的每个 vector lane 逐一等同为 CUDA thread。

增大 tile 可能提高复用、减少重复 load，却也增加 registers/shared memory，降低同一 SM 能同时驻留的 programs；减小 tile 可增加并行工作数，却可能增加边界开销和 HBM traffic。好的 blocking 依赖 shape、dtype、GPU、alignment、warp 数、流水级数和相邻 epilogue，不是从公式唯一推导出来的常数。

Layout、alignment、workspace 与 autotuning

同一逻辑 shape 仍可能因 stride、transpose、packing、pointer alignment 而选到不同实现。以 GEMM 为例，backend 还需决定 tile shape、Tensor Core instruction、split-K、epilogue fusion 与 workspace：

• cuBLAS 以成熟 BLAS API 承载常见路径；cuBLASLt 暴露更灵活的 operation/layout、heuristic algorithm 和 fused epilogue，某些算法需要 workspace。
• CUTLASS 让实现者用 C++ 模板组合 collective mainloop、tile scheduler 与 epilogue；它是构造 kernels 的工具，不是 torch.matmul 必然经过的公开证据。
• Triton 以 constexpr/meta-parameters 编译专门变体，可用 autotune 在候选 BLOCK_*、warps、stages 中实测选择；首次 JIT/autotune 成本必须与 steady state 分开。
• FlashInfer 提供面向 serving 的 prefill/decode、paged/ragged KV、sampling 等 primitives；其内部实现与支持矩阵会变化，“用了 FlashInfer”仍不是某个固定 kernel 的名称。

因此，“API 能调用”只证明入口接受参数。命中 fast path 还可能要求特定 dtype、head dimension、contiguous/strided layout、alignment、group size、GPU capability、workspace budget 与 graph-capture 条件。推理量化尤其依赖 packing 与 scale layout；不匹配时可能先 unpack/dequantize、换通用 kernel 或直接失败。

普通 launch 与 persistent kernel 的准确边界

本实验 grid 为 (R,)：有 $R$ 个 programs，每个处理一行后终止。它是 blocked/fused kernel，不是 persistent kernel。

Persistent kernel 通常只启动与可用 SM/occupancy 同量级的有限 programs，让它们在 device 端循环取得多个 rows、tiles、tokens 或任务。可能的收益包括：

• 减少大量细碎工作单元的 launch/scheduling 开销；
• 让复用的数据或调度状态更久留在 registers/shared memory/cache；
• 在融合 decode、MoE 或动态任务队列中跨多个工作单元保持设备侧控制流。

代价是常驻程序长期占用 SM、register/shared memory，可能阻碍其他 kernels；工作量不均会产生尾部，device-side queue/synchronization 更复杂；大 GEMM 已能高效占满 GPU 时，persistent scheduling 也未必改善性能。它与 CUDA Graph不同：Graph 减少 host 重放一串 launches 的开销，并不要求单个 kernel 常驻；两者可以组合。

从教学 kernel 到生产实现还差什么

• 一行必须装进单个 program；很宽的行会产生高 register pressure、spill，甚至超过编译/硬件可行范围。生产 softmax 可能分块、多阶段归约或使用 persistent row loop。
• 只支持 contiguous 2D input，没有 arbitrary stride、空维度、非 CUDA、autograd、in-place/alias、ragged rows 或多 dtype policy。
• 固定启发式 num_warps，没有针对 shape/GPU autotune，也没有编译缓存、部署包与冷启动治理。
• 没分析 occupancy、register count、shared memory、DRAM throughput 或 numerical ULP；一个 kernel 名不能替代这些 profiler metrics。
• Softmax 已是单算子。生产优化更常把 mask、scale、softmax 和后续操作融合，或像 FlashAttention那样改变整个 attention 的 tile/IO 计划，而不是孤立替换 softmax。
• Custom kernel 必须继续满足 stream dependency、allocator lifetime 和 graph-capture 约束；这些由 GPU Runtime展开。

与推理主链路的连接

• GEMM/GEMV 与算术强度判断算法更可能受 compute、memory 还是 launch 限制；本页解释 backend 如何用 tile 实际兑现或浪费这些机会。
• FlashAttention把 blocking 和 online softmax扩展到 Q/K/V tiles；教学 row softmax只展示单行归约，不实现 attention。
• Backend Dispatch/Fallback提供从 API 到 concrete kernel 的取证流程；本页再解释 concrete kernel 内部的 program/tile/warp 映射。
• GPU Runtime说明 fusion、streams、torch.compile 与 allocator；generated Triton kernel 仍在这些运行时约束中执行。
• 推理量化改变 payload、scale 和 packing；只有相应 kernel backend 原生消费这些布局时，容量优势才可能转为延迟收益。

参考资料

• Triton — Introduction
• Triton — Fused Softmax Tutorial
• NVIDIA — cuBLAS / cuBLASLt Documentation
• NVIDIA CUTLASS — Overview
• FlashInfer Documentation