GEMM、GEMV、dtype 与算术强度

同一个线性层为何在 prefill 与 decode 很不一样

把 batch 和 token 维展平，Transformer 线性层可写成：

$M=B\times T$。Prefill 中 $T$ 较大，通常是标准 GEMM；单请求单 token decode 中 $M=1$，数学上接近 GEMV。库 API 仍可能把它提交给 GEMM kernel，所以 GEMM/GEMV 是 shape 与复用特征，不一定是用户可见 API 名称的二分。

Prefill/Decode 页面给出了阶段 shape；本页解释为何同一组权重在两个阶段的瓶颈不同。

FLOPs、最低字节数与算术强度

矩阵乘的乘加约为：

若理想地只从显存各读取一次 $X,W$ 并写一次 $Y$，元素大小为 $s$ bytes，最低数据量为：

算术强度（Arithmetic Intensity）为：

Roofline 上界：

这里的 $Q_{min}$ 是算法级下界，不是 profiler 实测 DRAM bytes。cache miss、workspace、layout 转换、低效 tile、epilogue 和重复读取都可能增加真实流量；计算与访存也会重叠。

具体代入：$K=N=4096$，FP16

当 $M=1$ 时，几乎每个 FP16 权重元素只服务一个乘加，约 2 FLOPs / 2 bytes，权重复用极低。$M$ 增大后，同一权重 tile 被更多 rows 复用，AI 大幅提高，更可能接近 Tensor Core 算力上限。

dtype 改变什么

• $s$ 变小会降低权重/激活 payload，理论上提高按 byte 计算的 AI；但只有 backend 有对应低精度 kernel、packing 和累加路径时才能兑现。
• FP16/BF16 常用更高精度 accumulator；FP8/INT8/INT4 还涉及 scale、zero point、反量化与校准误差。
• Weight-only quantization 主要减小 decode 时反复读取的权重；prefill 是否加速取决于反量化能否融合、Tensor Core 支持和大 GEMM 的算力瓶颈。
• torch.float32 在支持 TF32 的 CUDA backend 上可能实际使用 TF32 乘法语义；API dtype 与执行单元精度要分开说明。

可运行 CUDA 实验

实验问题：固定同一 $4096\times4096$ FP16 权重，只改变 $M$，理论 AI 与同步后的 kernel 调用延迟呈现什么趋势？A/B 输出先用 FP32 reference 校验。

import platform, statistics, time
import torch

torch.manual_seed(61)
assert torch.cuda.is_available(), "本 benchmark 需要 CUDA；公式本身不依赖 GPU"
DEVICE=torch.device("cuda"); DTYPE=torch.float16

def model(m,k,n,s):
    flops=2*m*k*n; minimum_bytes=(m*k+k*n+m*n)*s
    return flops,minimum_bytes,flops/minimum_bytes

def benchmark(x,w,warmup=20,rounds=50):
    for _ in range(warmup): x@w
    torch.cuda.synchronize(); samples=[]
    for _ in range(rounds):
        start=time.perf_counter(); y=x@w; torch.cuda.synchronize()
        samples.append((time.perf_counter()-start)*1e6)
    return y,statistics.median(samples)

def main():
    k=n=4096; w=torch.randn(k,n,device=DEVICE,dtype=DTYPE); s=w.element_size()
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={torch.cuda.get_device_name(0)} dtype={DTYPE} W={(k,n)}")
    for label,m in (("decode",1),("small_batch_decode",16),("prefill",512)):
        x=torch.randn(m,k,device=DEVICE,dtype=DTYPE)
        y,latency=benchmark(x,w)
        reference=x.float()@w.float()
        torch.testing.assert_close(y.float(),reference,rtol=2e-3,atol=2e-2)
        flops,minimum_bytes,intensity=model(m,k,n,s)
        print(f"{label}: X=({m},{k}) Y=({m},{n}) latency_us={latency:.3f} "
              f"flops={flops/1e9:.3f}G minimum_bytes={minimum_bytes/2**20:.3f}MiB "
              f"arithmetic_intensity={intensity:.3f}FLOP/byte")

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

完整文件位于 examples/gemm_gemv_roofline.py。2026-08-13 在 RTX 3080、PyTorch 2.13.0+cu130 上实测中位数：$M=1$ 为 59.957 µs，$M=16$ 为 60.934 µs，$M=512$ 为 329.905 µs。前两者延迟接近，但完成的 FLOPs 相差 16 倍，展示了小矩阵的带宽/启动与利用率问题；prefill 绝对延迟更高，却在单位时间完成远多得多的工作。

这个结果不是 cuBLAS 的普遍性能表：设备、clock、PyTorch/cublasLt 版本、矩阵对齐、transpose、epilogue 和并发都会改变 kernel。计时范围仅是同步后的 x @ w，不含模型层其他算子、queue 或 streaming。

常见错误与生产映射

• 用 nvidia-smi 的利用率百分比直接判断 compute-bound；应结合 profiler 的 FLOPs、DRAM throughput、occupancy 和 kernel 时间。
• 只看总 FLOPs 推断延迟，忽略每步重新读取模型权重与 KV。
• 用 CPU wall-clock 测 CUDA 异步提交而不 synchronize，得到 launch 时间而非执行时间。
• 把 PyTorch matmul API 当成固定 kernel；生产中可能 dispatch 到 cuBLAS/cuBLASLt、CUTLASS、Triton 或量化 fused kernel。
• 认为量化 payload 按位宽缩小就必然同比加速；metadata、反量化、fallback 和小 batch kernel 效率可能抵消收益。

上下游关系

• Prefill/Decode 提供 $M=B\times T$ 的阶段来源；本页解释 compute-bound 与 memory-bound 倾向。
• MHA/MQA/GQA 改变 Q/K/V projection 的 $N$ 和 KV 读取量；不会让 Query projection 消失。
• 张量并行 把 $K$ 或 $N$ 切到多个 rank，减少局部 GEMM 但加入 collective；是否加速要比较被省的 matmul 时间与通信。
• Backend Dispatch/Fallback展示 x @ W、ATen operator 与具体 CUDA kernel 是不同证据层，并解释 shape/dtype/layout 如何改变选择。
• Kernel Backend、Triton Blocking 与 Persistent Kernel说明 tile、warp、workspace 与 autotuning 怎样决定理论算术强度能否在具体实现中兑现。
• 推理量化基础进一步把 weight-only、W8A8 的 payload、scale 和整数累加语义展开，避免把位宽缩减直接等同于 GEMM 加速。
• 后续 CUDA Graph 页面应把 kernel 本身耗时与 CPU launch/framework overhead 分开；Graph 不会让同一个大 GEMM 获得同比例算力提升。

参考资料

• NVIDIA — Matrix Multiplication Background User’s Guide
• NVIDIA cuBLAS documentation
• PyTorch — torch.matmul