第 11 章 · Attention 入门

学习目标

• 把"缩放点积注意力"翻译成代码：QKᵀ → softmax → ·V
• 实现 fused QKV projection（1 个 GEMM 代替 3 个）
• 看清朴素实现的 O(T²) HBM 物化问题——为下一章 FlashAttention 埋伏笔
• 会处理 causal mask

前置知识

已完成 Ch09–10，能实现 GEMM、数值稳定 Softmax，并理解 row-major 张量布局。

核心概念

11.1 数学复习

Scaled Dot-Product Attention (single head, single batch):

S = Q @ K^T / sqrt(D)     # (T, T)
P = softmax_row(S + mask) # (T, T)
O = P @ V                 # (T, D)

• Q, K, V 形状 (T, D)，T = 序列长度，D = head_dim（GPT-2 small 是 64）
• scale = 1/√D 防止 softmax 输入太大
• causal mask：自回归生成时，token i 不能看到 j>i 的 token，所以 S[i,j>i] = -∞

关键代码

11.2 Fused QKV Projection

输入 X: (T, D)。原本要做三次 GEMM：

Q = X @ W_q   (T, D) @ (D, D)
K = X @ W_k
V = X @ W_v

把 W_qkv = [W_q; W_k; W_v] 水平拼成 (D, 3D)，一次 GEMM 出 (T, 3D) 再 split：

// 1) 1 个 GEMM
gemm(X, W_qkv, QKV, T, 3*D, D);

// 2) split (T, 3*D) → 三个 (T, D)
__global__ void split_qkv(...) {
    int t = ..., d = ...;
    q[t*D+d] = qkv[t*3*D + 0*D + d];
    k[t*D+d] = qkv[t*3*D + 1*D + d];
    v[t*D+d] = qkv[t*3*D + 2*D + d];
}

好处：

• 3 个 GEMM → 1 个 → kernel launch 开销少 2/3
• GPU 利用率高（单个大 GEMM 比三个中等 GEMM 容易跑满）
• 权重一次 H2D，cache 局部性更好

11.3 朴素三阶段 Attention

源码：attention_naive.cu

// Step 1: S = Q @ K^T * scale   (T, T)
__global__ void qkt_scale(const float* Q, const float* K, float* S,
                          int T, int D, float scale, bool causal) {
    int i = ...,  j = ...;
    if (causal && j > i) { S[i*T+j] = -1e30f; return; }
    float s = 0;
    for (int d = 0; d < D; ++d) s += Q[i*D+d] * K[j*D+d];
    S[i*T + j] = s * scale;
}

// Step 2: P = softmax_row(S)
softmax_rows<256><<<T, 256>>>(S, T, T);

// Step 3: O = P @ V             (T, D)
__global__ void pv_kernel(const float* P, const float* V, float* O, int T, int D) {
    int i = ..., d = ...;
    float acc = 0;
    for (int j = 0; j < T; ++j) acc += P[i*T+j] * V[j*D+d];
    O[i*D+d] = acc;
}

11.4 朴素实现的两个致命问题

问题 1：HBM 物化 T×T

S 和 P 都是 T × T，必须先全部写到 HBM 再读。T=2048 时 S 占 16 MiB，T=8192 时 S 占 256 MiB—— 不光慢（多 3 次全量 HBM 读写），显存还放不下长上下文。

问题 2：softmax 是 memory-bound

softmax 算每个元素 O(1) FLOP 但要扫两遍 T。kernel 间没法融合，每次都得回 HBM。

graph LR
    Q["Q in HBM"] --> K1["QKᵀ kernel"]
    K1 --> S1["S in HBM"]
    S1 --> K2["softmax kernel"]
    K2 --> P1["P in HBM"]
    P1 --> K3["PV kernel"]
    K3 --> O1["O in HBM"]
    style S1 fill:#f3f1e8,stroke:#a82c30
    style P1 fill:#f3f1e8,stroke:#a82c30

两次完整的 T×T 显存读写——这就是 FlashAttention 要彻底干掉的地方。

运行结果

11.5 性能与瓶颈

长上下文场景下 FlashAttention 不是"快几倍"——是"能不能跑"的差别。

自检清单

练习题

1. 把 attention_naive 扩成 multi-head：B=1, H=4, T=128, D=64。
2. 给 fused_qkv 用 Ch9 的 gemm_reg_tile 替换 tiled gemm，看是否更快。
3. 测一下 T 从 64 到 4096 的耗时曲线，画出 O(T²) 增长曲线。
4. 实现一个 attention 输出与 PyTorch 的 F.scaled_dot_product_attention 对比（Colab 上）。

11.8 工业实战：MHA/MQA/GQA、prefill vs decode、causal mask 优化

11.8.1 MHA / MQA / GQA — KV cache 的演进

例：Llama 3 8B 用 n_q=32, n_kv=8 (GQA-4)，KV cache 立刻缩到 MHA 的 1/4。

实现差异

// MHA: 各 head 完全独立
attention(Q[b, h, :, :], K[b, h, :, :], V[b, h, :, :]) → O[b, h, :, :]

// MQA: 所有 Q head 共享同一对 K/V
attention(Q[b, h, :, :], K[b, 0, :, :], V[b, 0, :, :]) → O[b, h, :, :]

// GQA: head 分组共享
int g = h / group_size;
attention(Q[b, h, :, :], K[b, g, :, :], V[b, g, :, :]) → O[b, h, :, :]

kernel 层面只是 K/V 索引换一下，FlashAttention 默认支持三种。

11.8.2 Prefill vs Decode — 两个完全不同的世界

LLM 推理本质上是两个独立 phase：

为什么要分两个 kernel

同一个 attention kernel 在 T=2048 和 T=1 上的最优 tile 完全不同：

• Prefill 用 64×64 Q tile，T 个 K/V tile 流过
• Decode 用 1×N Q tile（实际是 vector），K/V 直接从 cache 拉

vLLM、TensorRT-LLM 都分别实现两个 kernel：flash_attn_varlen_func（prefill）和 flash_attn_with_kvcache（decode）。

11.8.3 Causal mask 的高效实现

Ch11 朴素实现把 -∞ 写到所有 j>i 的位置，但 FlashAttention 里这是浪费：

// ❌ 朴素：扫全 N×N, 一半位置算了又 mask 掉
if (j > i) S[i, j] = -inf;

// ✅ 优化 1: 跳过完全 mask 掉的 tile (block-level skip)
// FlashAttention 里, Q tile 在 [i_lo, i_hi), K/V tile 在 [j_lo, j_hi):
//   - 如果 j_lo > i_hi: 整个 K/V tile 都被 mask, 直接 skip 不计算
//   - 如果 j_hi <= i_lo: 完全 visible, 不用 mask
//   - 部分重叠: 逐元素 mask
if (j_lo > i_hi) continue;             // skip 这个 K/V tile
if (j_hi <= i_lo) compute_no_mask();   // 不用 mask
else compute_with_mask();              // 部分 mask

对长 context，causal mask 让大量上三角 tile 可以 skip；精确节省取决于 tile 大小和调度方式。FlashAttention v2 的 causal 实现会利用这个结构，但仍要处理负载均衡。

11.8.4 长 context 的工业方案

T 从 2K 增到 128K 的发展：

• FlashAttention: O(T²) FLOPs 不变, 但显存从 O(T²) 降到 O(T)，让 32K 可行
• Sliding Window Attention (Mistral 7B): 只 attend 最近 W=4K 个 token，把 attention 复杂度从 O(T²) 降到 O(T·W)
• Dilated / sparse attention (Longformer): 间隔抽取若干位置
• Linear / Mamba: 完全放弃 softmax attention，用 state-space 模型 O(T)
• YaRN / NTK-aware: RoPE 频率重 scaling，把 4K 训练的模型外推到 128K 推理

11.8.5 attention dtype 实战

FlashAttention 严格遵循此模式：fp16 I/O + fp32 累加。不要偷懒全 fp16，长 T 时 P @ V 会丢精度，模型输出乱。

11.8.6 attention 在 LLM 推理总耗时中的占比

所以：短 prompt 优化 FFN，长 prompt 优化 attention（FlashAttention），decode 优化权重加载（量化），大 batch 优化 PagedAttention。

11.9 研究前沿（2025-2026）：MLA、线性 attention、SSM 混合、Reasoning

11.9.1 Multi-head Latent Attention (MLA) — DeepSeek 杀器

DeepSeek-V2/V3 引入，核心价值是把完整 K/V cache 压到 latent 表示里；显存节省比例由 latent 维度、head 数和 dtype 决定，应按模型配置计算。

// 标准 MHA: cache K, V shape (T, n_head, d_head), 共 2 × n_head × d_head
// MLA: 把 K/V 投影到低维潜空间 c_kv (T, d_c)
c_kv = X @ W_DKV    # (T, d_c)        d_c=512 远小于 n_head*d_head
K_h  = c_kv @ W_UK[h] + position      # 各 head 上投影回 d_head
V_h  = c_kv @ W_UV[h]

// Key point: cache c_kv instead of full K/V; compute bytes from model config.

工程关键："吸收"W_UK 进 W_Q，让 Q @ K^T 直接 = Q' @ c_kv^T，不必构造完整 K：

Q @ K^T = (Q @ W_UK^T) @ c_kv^T
          ─────────────
          推理时只算一次, Q' 跟着 token 走

DeepSeek 2025 开源 FlashMLA kernel，FA v2 的 MLA 变种。MLA 已成为 2025 中国大模型主流（Qwen 3 也用）。

11.9.2 Linear Attention 重生

O(N²) → O(N·D²)：用 kernel function φ + 结合律：

softmax(QK^T) V          O(N²·D)
       ↓ φ(Q) @ φ(K)^T 近似 + 结合律
φ(Q) @ (φ(K)^T @ V)      O(N·D²)
              ↑
        (D,D) 大小, 与 N 无关

2024-2026 主要变体：

• Lightning Attention-2（MiniMax 2024）：右乘累加，128K 实用
• RWKV-7（2025）：完全 RNN-like，零 KV cache
• RetNet（Microsoft）：训练并行、推理 O(1) per step
• Gated Linear Attention (GLA)：加门控提升精度
• Lightning Indexer：稀疏路由 + softmax/线性混合

评价：精度仍略输 softmax attention，但 1M+ context 唯一可行。

11.9.3 Mamba / SSM 混合架构

• Mamba（2023 末）：selective SSM，跟 Transformer 同精度 + 线性复杂度
• Mamba-2（2024）：state space ↔ attention 对偶, 跑 Tensor Core
• Jamba（AI21）：Transformer + Mamba + MoE 混合, 256K context
• Zamba / Hymba：各种 Mamba+Attn 混合
• Samba（Microsoft 2024）：Mamba + 滑动窗口 attention

共识：纯 SSM 仍不及 Transformer，但 SSM + 少量 Attention 块在长 context 推理 KV cache 大幅缩小，是产业新方向。Llama 4 据传用类似混合。

11.9.4 Sparse Attention 新进展

• Sliding Window（Mistral/Mixtral）：固定 4K-8K 窗口
• Native Sparse Attention (NSA)（DeepSeek 2025）：训练学习的稀疏 pattern
• MoBA（Kimi 2024）：K/V 分块，attention 选 top-k 块
• InfLLM / Quest：路由式 KV cache 检索

11.9.5 Reasoning 模型对 attention 的新挑战

o1 / o3 / R1 / Claude reasoning：单次回答生成 10K-100K thinking token。

• 极长 generation 的 KV cache 持续膨胀, 每 token 写入是 mem-bound
• 动态 batch: reasoning 长度方差极大 (10K vs 100K), scheduling 难度激增
• Test-time compute scaling: 推理算力成为训练同等阶, 推理优化 ROI ×10

这是 2025-2026 推理优化（FA v3/v4、FlashMLA、MoBA、Lookahead 解码）远比训练优化更受关注的原因。

11.9.6 一图总结 2026 attention 全景

           ┌─ 标准 MHA    GPT-2/3/4 老路, 已过时
softmax    │
attention ─┼─ GQA         Llama 2/3, Mistral — 当前实用主流
 O(N²)     │
           └─ MLA         DeepSeek V2/V3 — KV cache 结构性压缩

           ┌─ Sliding W.  Mistral 7B, 长 context 推理
sparse    ─┼─ NSA / MoBA  DeepSeek / Kimi 2025
attention  │
           └─ InfLLM      检索式 KV 取舍

           ┌─ Lightning   MiniMax — 线性 + softmax 混合
linear /  ─┼─ RWKV-7      纯 RNN-like, 零 KV
SSM        │
           ├─ Mamba-2     selective SSM, Tensor Core 友好
           │
           └─ Jamba/Samba SSM+Attn 混合 (产业方向)

对工程师的影响：单一 attention kernel 不够。生产推理引擎需要 FlashInfer / xFormers / FlexAttention 这种"attention dispatch 层"，根据模型自动选 backend。

常见坑

• causal mask 写错方向（j<i 还是 j>i）→ 模型胡言乱语
• scale 写成 D 而不是 sqrt(D) → softmax 几乎是 one-hot，模型急于自我重复
• fused QKV split 时下标写错 → Q/K/V 张冠李戴，输出 nan
• 长 T 直接 OOM → 必须分块 / FlashAttention
• MLA 没"吸收"W_UK 进 Q → 每步显式构造 K，KV 显存优化白做
• 用线性 attention 训练，但推理时切回 softmax → 数值分布对不上，效果暴跌
• Mamba + Attn 混合模型推理时，把不同层 KV cache 当一样大小处理 → 跨层显存计算错
• Reasoning 模型 max_tokens 设小（如 4K），实际 thinking 要 30K → 模型被截断输出乱码

11.11 CUDA 官方手册精讲（CUDA Programming Guide 13.2（核验：2026-07-20））

Memory Traffic 分析、MHA/MQA/GQA layout、Causal Mask 优化

朴素 attention 的 HBM 流量公式

11.5 节给了"T=4096, D=64, S 占 64 MiB"的结果，但没拆账。这里用 CUDA Programming Guide 13.2 Asynchronous Data Copies 的 transaction-bytes 视角分析朴素三阶段 attention：

设单头 (single-head): Q, K, V 形状均为 (T, D)；fp16 = 2 字节。

关键观察：

• 当 T >> D 时，HBM 流量被 T² 项主导（S 和 P 反复读写）。
• FLOPs 也是 T² 但只系数 ~4·D。算术强度 ~0.7·D：D=64 时 ~45 FLOP/byte——远低于 A100 fp16 算力/带宽比 = 312 TF / 2 TB/s = 156 FLOP/byte → memory-bound。
• FlashAttention 把 S, P 完全消除在 HBM 上（只走 shared），HBM 流量直降到 ~4·T·D（只剩 Q/K/V/O 各一遍），算术强度涨到 ~T·D / (T·D) = ~T——T=2048 时 ~2048 FLOP/byte，compute-bound。

    graph LR
        Naive["朴素三阶段
HBM: 12·T² + 6·T·D B
算术强度 ~0.7·D
memory-bound"]
        FA["FlashAttention
HBM: ~8·T·D B
算术强度 ~T
compute-bound (大 T)"]
        Naive -.->|"省去 S, P 物化"| FA
        style Naive fill:#f3f1e8,stroke:#a82c30
        style FA fill:#f3f1e8,stroke:#2f5d3a
    

这就是为什么下一章的"省 HBM"常常直接体现在端到端速度上：不是 FLOPs 省了，而是把中间矩阵从 HBM 移回 shared/register，让 roofline 落点从 memory-bound 往 compute-bound 移动。具体加速必须按序列长度、head_dim、dtype 和实现测量。

Causal mask 的 tile-level skip 收益量化

11.8.3 给了 block-level skip 代码，"causal 让一半 tile 可以 skip → 复杂度 O(T²/2)"——这只是 FLOPs 账。完整收益还有 负载均衡，分两个维度看：

KV cache 的 VMM 视角：为什么 PagedAttention 选 16 token / 页

vLLM 的 PagedAttention block size 是软件层设计选择；CUDA Programming Guide 13.2 的 Virtual Memory Management 章节则给出底层虚拟内存映射约束，两者不能直接等同：

// 查询当前 GPU 的最小可分配粒度
    size_t granularity = 0;
    CUmemAllocationProp prop = {};
    prop.type           = CU_MEM_ALLOCATION_TYPE_PINNED;
    prop.location.type  = CU_MEM_LOCATION_TYPE_DEVICE;
    prop.location.id    = device;
    prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;

    cuMemGetAllocationGranularity(&granularity, &prop,
                                  CU_MEM_ALLOC_GRANULARITY_MINIMUM);
    // Print this value; do not assume it is identical across GPUs/drivers.
    

先用 cuMemGetAllocationGranularity 查询目标 GPU/driver 的最小粒度，再设计 KV cache 块大小。一个 KV cache 块要：

• 装下 N 个 token 的 K + V： N × 2 × n_kv_head × d_head × 2 (fp16) 字节
• 典型 Llama-3 8B: n_kv_head=8, d_head=128 → 每 token 4 KiB
• N = 16 token → 64 KiB / 块 (32 个块 = 2 MiB, 正好填满一个 VMM 最小分配)

    graph TB
        VMM["cuMemCreate
granularity = 2 MiB"]
        VMM --> Pool["KV pool:
N 个 2 MiB 物理块"]
        Pool --> Map["block_table[req][t/16]
→ 物理 block id"]
        Map --> Attn["attention kernel:
indirection load K, V"]
        style VMM fill:#f3f1e8,stroke:#a86420
        style Attn fill:#f3f1e8,stroke:#2f5d3a
    

下一章 §12.9.4 给出 attention kernel 怎么消费 block_table。

MHA vs MQA vs GQA 的 kernel 代码差异（v13.1 API 完整版）

11.8.1 给了一行伪代码，看着像三种变体要写三个 kernel；事实上生产 attention 库（FlashInfer / vLLM / FA v3）一份代码搞定，只在 K/V 加载 那一步换 indexing：

template <int N_Q_HEAD, int N_KV_HEAD, int D>
    __global__ void unified_attention_kernel(
        const __half* Q,          // (T, N_Q_HEAD, D)
        const __half* K,          // (T, N_KV_HEAD, D)  注意 N_KV_HEAD 维变
        const __half* V,          // (T, N_KV_HEAD, D)
        __half* O,
        int T, float scale)
    {
        static_assert(N_Q_HEAD % N_KV_HEAD == 0, "group_size must divide");
        constexpr int GROUP_SIZE = N_Q_HEAD / N_KV_HEAD;

        int h_q  = blockIdx.x;                        // 0..N_Q_HEAD-1
        int h_kv = h_q / GROUP_SIZE;                  // ★ 唯一区别 ★

        // ---- 1) load Q, K, V tile, swizzle 128B ----
        // Q: 永远用 h_q
        // K/V: 用 h_kv (MHA 时 h_kv == h_q, MQA 时 h_kv = 0, GQA 时 h_q/GROUP_SIZE)

        /* ... 后面 online softmax + GEMM 跟 single-head 完全一样 ... */
    }

    // 启动配置:
    // MHA: N_Q_HEAD = N_KV_HEAD = 32   → GROUP_SIZE = 1
    // MQA: N_Q_HEAD = 32, N_KV_HEAD = 1 → GROUP_SIZE = 32
    // GQA: N_Q_HEAD = 32, N_KV_HEAD = 8 → GROUP_SIZE = 4
    

kernel 层面统一写法的隐藏好处：用 PG §4.11 的 TMA descriptor 共享（一个 CUtensorMap 描述整个 K tensor），h_kv 仅参与 box_dim 起始坐标计算，硬件 cache 命中率不变——这就是为什么 FlashInfer 11.9.6 提到的"backend 自动 dispatch"可以做到 zero overhead。

下一章导览

第 12 章避免物化完整 attention score 矩阵，用 online softmax 与 tile 复用降低 HBM 流量。