缩放点积注意力 · LLM Glossary

数学语义与所在层次

对一个 attention 调用，设：

Q\in\mathbb{R}^{B\times H_q\times T\times D},\quad K,V\in\mathbb{R}^{B\times H_{kv}\times S\times D}

缩放点积注意力（Scaled Dot-Product Attention, SDPA）计算：

A=\frac{QK^{\mathsf T}}{\sqrt D}+M, \qquad P=\operatorname{softmax}(A), \qquad O=PV

$M$ 是 causal、padding 或业务约束产生的 additive mask；不可见位置通常加 $-\infty$ 。输出 $O\in\mathbb{R}^{B\times H_q\times T\times D}$ 。除非使用 GQA/MQA 共享，公式中的 head 乘法是同编号 Query head 与其负责的 KV head 配对。

它位于 Transformer attention 子层的核心数学位置：上游线性投影产生 Q/K/V，下游把各 head 拼回 hidden dimension 并做 output projection。它不是完整的 Multi-Head Attention block，也不负责 KV Cache 生命周期、物理分页或请求调度。

为什么除以 $\sqrt D$

若 $q_i,k_i$ 独立、均值为 0、方差约为 1，则未缩放点积 $q\cdot k=\sum_{i=1}^Dq_ik_i$ 的方差随 $D$ 增长到约 $D$ 。较大的 logits 会把 softmax 推入饱和区，概率接近 one-hot，梯度与低精度数值更敏感。除以 $\sqrt D$ 后，logit 方差回到约 1 的量级。

这不是在保证数值永远稳定。实现仍要用减去行最大值的稳定 softmax；低精度累加、极端 mask 和 backend 运算顺序仍可能带来差异。

一个 $T=S=3$ 的 causal worked example

假设单 head、 $D=2$ ，缩放前点积分数为：

QK^{\mathsf T}= \begin{bmatrix} 2&1&0\\ 1&2&1\\ 0&1&2 \end{bmatrix}

causal mask 为：

M= \begin{bmatrix} 0&-\infty&-\infty\\ 0&0&-\infty\\ 0&0&0 \end{bmatrix}

第一行只能看到位置 0，softmax 后必为 $[1,0,0]$ ；第二行能看到 0、1；第三行能看到全部历史。无论 backend 是否物化矩阵，数学上的每一行概率都非负且和为 1，不可见位置概率为 0。

中间量	shape	layout / 生命周期
Q	$[B,H_q,T,D]$	当前调用输入；prefill 常有 $T>1$ ，decode 常有 $T=1$
K/V	$[B,H_{kv},S,D]$	prefill 可为当前序列，decode 可来自历史 KV Cache
scores / probabilities	$[B,H_q,T,S]$	reference 实现会物化；FlashAttention 类 kernel 分 tile 在线维护状态
output	$[B,H_q,T,D]$	每个 Query 位置一个加权 Value 向量

可运行 reference 与 PyTorch SDPA 对齐

实验问题：手写 causal GQA attention 是否与 PyTorch SDPA 数值一致？causal mask 和 softmax 行和是否满足不变量？

import math, platform
import torch
import torch.nn.functional as F

torch.manual_seed(17)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float32

def reference_attention(q, k, v, causal=True):
    # q:[B,Hq,T,D], k/v:[B,Hkv,S,D]
    q_heads, kv_heads = q.shape[1], k.shape[1]
    assert q_heads % kv_heads == 0 and k.shape == v.shape
    k = k.repeat_interleave(q_heads // kv_heads, dim=1)
    v = v.repeat_interleave(q_heads // kv_heads, dim=1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if causal:
        target, source = q.shape[-2], k.shape[-2]
        assert target == source, "本例只实现同长度 causal self-attention"
        mask = torch.triu(torch.full((target, source), -torch.inf, device=q.device), 1)
        scores = scores + mask
    probs = torch.softmax(scores, dim=-1)
    return probs @ v, probs

def main():
    b, t, hq, hkv, d = 2, 5, 4, 2, 8
    q = torch.randn(b, hq, t, d, device=DEVICE, dtype=DTYPE)
    k = torch.randn(b, hkv, t, d, device=DEVICE, dtype=DTYPE)
    v = torch.randn_like(k)
    reference, probs = reference_attention(q, k, v)
    sdpa = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
    torch.testing.assert_close(sdpa, reference, rtol=1e-5, atol=1e-6)
    torch.testing.assert_close(probs.sum(-1), torch.ones_like(probs[..., 0]))
    forbidden = torch.triu(torch.ones(t, t, dtype=torch.bool, device=DEVICE), 1)
    assert torch.count_nonzero(probs.masked_select(forbidden)).item() == 0

    expanded = F.scaled_dot_product_attention(
        q, k.repeat_interleave(hq // hkv, 1),
        v.repeat_interleave(hq // hkv, 1), is_causal=True)
    torch.testing.assert_close(sdpa, expanded, rtol=1e-5, atol=1e-6)
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE}")
    print(f"Q={tuple(q.shape)} K=V={tuple(k.shape)} scores={tuple(probs.shape)} O={tuple(sdpa.shape)}")
    print(f"row_sum_error={(probs.sum(-1)-1).abs().max():.3e}")
    print(f"sdpa_max_abs_diff={(sdpa-reference).abs().max():.3e}")

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

同版完整文件位于 examples/attention_reference.py。2026-08-13 在 PyTorch 2.13.0+cu130、RTX 3080 上实际运行得到 Q=[2,4,5,8]、K=V=[2,2,5,8]、scores [2,4,5,5]，softmax 最大行和误差 1.192e-07，SDPA 与 reference 最大绝对差 2.384e-07。

代码显式复制 GQA heads 只为证明语义；不能用它衡量生产 GQA 的内存或速度。PyTorch backend 还可能改变浮点累加顺序，因此 BF16/FP16 应使用与误差来源匹配的容差。

成本与何时失效

reference 实现的两个矩阵乘主项约为 $O(BH_qTSD)$ ，而完整 score/probability 中间量需要 $O(BH_qTS)$ 元素。prefill self-attention 中 $T=S$ ，逻辑计算量随序列长度二次增长；使用 KV Cache 的单步 decode 中 $T=1,S$ 为历史长度，单步 attention 工作随 $S$ 线性增长。

FlashAttention 通过 tiling 与 online softmax 避免把完整 $T\times S$ 中间矩阵反复写回显存，但不改变上面的数学结果或 dense attention 的依赖范围。PagedAttention 则管理历史 KV 的物理 block；它也不等价于 SDPA 或 FlashAttention。

常见失败点包括：

把 boolean mask 的 True 语义在不同 API 间套反；
decode 时把 absolute position 与 causal mask 对齐错位；
同时传 is_causal=True 和不兼容的 explicit mask；
$H_q$ 不能被 $H_{kv}$ 整除却强行启用 GQA；
以为 fused backend 一定可用，忽略 dtype、head dimension、mask、dropout、硬件和软件版本约束。

上下游关系

MHA / MQA / GQA 决定 $H_q$ 与 $H_{kv}$ 以及 head 配对方式；SDPA 本身不定义这些投影参数如何共享。
KV Cache 为 decode 提供历史 K/V，使 SDPA 的 $T=1$ Query 能读取 $S$ 个历史位置；KV 页面反向解释缓存为何没有消除 attention 读取。
FlashAttention 从本页的 reference 语义出发，证明 online softmax 与 dense softmax 对齐，并区分 API backend 与具体 kernel。