← 全部作业
← → 翻页 · O 总览 · ? 帮助
Assignment 4
A4

Attention Module

Offline Sliding-Window · Online（可选）· Mask 的艺术

全课最难之一 a4/src/modeling/attention.py
定位

为什么 Attention 难

  • 同一公式，多种 mask / layout / pack
  • 数值稳定：scale · temp · cap · clip · dropout
  • 工程变体：GQA · varlen · 矩形 Q/K
  • 可选 Online：Flash 风格分块 + lse
建议工期 4–6 天；赶进度只保 Offline 满分。
依赖

动手前你需要什么

  • A2 GroupRMSNorm → QK-Norm 会用到
  • A1 多头形状直觉：batch / head 当 batch-like 维
  • Blog 公式 + FlashAttention mask 约定
A2 没过，A4 的 QK-Norm 会卡住；先同步 norm 模块。
任务地图

两大部分

Task1 · Offline

完整滑动窗口注意力
多 pack / layout / mask

Task2 · Online（可选）

分块计算 + 在线合并 lse
对齐 Offline 数值

核心公式

每个 head 在算什么

Q: [b, sq, hq, hd] K: [b, skv, hkv, hd] V: [b, skv, hkv, hd] score = scale · QKᵀ (+ mask 置 -inf) A = softmax(score) (+ 稳定技巧) O = A V # O 的 layout / dtype / device 与 Q 一致

语义：用 query 从 value「知识库」里按 key 相似度取加权和。

Mask

三种常见掩码

Full

token 看全部

Causal

只能看自己及过去
j ≤ i

Sliding Window

只看窗口内
可与 causal 叠加

Mask 值

无关位置：-∞
相关位置：0

矩形 Q/K

bottom-right 对齐

  • sq ≠ skv 时，mask 是 [sq, skv] 矩形
  • 作业约定：从大方阵裁 右下角（Flash 风格）
  • 解码时：当前 query 在历史 KV 之后 → 对应 mask 最后一行
最容易写反：写成 top-left 对齐会静默错分。
稳定技巧

softmax 别爆

  • scale：常 1/√hd
  • temperature：P / temp（cap 为 None 时才用）
  • capping：cap · tanh(P/cap)
  • clipping：抑制 A 的 outlier
  • dropout：训练鲁棒性
  • QK-Norm：GroupRMSNorm 作用在 Q/K
Pack

QKV 怎么打包进来

Q_K_V : 三个独立张量 Q_KV : K,V 在 head 维拼；Q 单独 QKV : Q,K,V 都在 head 维拼 （此时序列长度等须一致）

枚举：AttnQKVPackFormat · 先解包再算，最后按 Q 的 layout 输出。

Layout

张量排法三种

BSHD

最常见
[b, s, h, d]

SBHD

偏分布式
[s, b, h, d]

THD

varlen 拼接
[total_s, h, d]

THD 辅助

cu_seqlens_q
cu_seqlens_kv

GQA / MQA

头数不一致怎么办

  • hq ≠ hkv，且 hq % hkv == 0
  • 策略：把 KV heads repeat_interleave 对齐到 Q
  • 参考：Llama Attention · PyTorch repeat_interleave
先让 head 维对齐，再做 matmul；别在错维上 softmax。
Offline 小结

实现检查单

  • 解 pack → 统一到可计算形式
  • 可选 QK-Norm（独立子层，Q/K 各一套）
  • GQA 重复 KV → 算 score → 加 mask
  • 稳定技巧 → softmax → dropout/clip
  • 乘 V → 还原 layout → 与 Q 同 meta
Online（可选）

分块 + 在线 softmax

  • 只支持 BSHD + Q_K_V
  • 每次 forward 一个 (bq, bkv) block
  • 更新全局 O 与 global_lse（float32）
  • 禁用 clip / dropout（无法直接作用全局）
  • 遍历所有 block 后 ≈ Offline 结果
思想来自 Online Softmax / FlashAttention；先懂 Offline 再碰。
Online · 关键

局部如何并到全局

每个 block 得到： O_local, lse_local 用 lse 重标定后累加到： global_o, global_lse # 就地更新，不显式 return # 末块可能 zero-pad 到 block_size
踩坑

A4 高频扣分点

  • mask 对齐成 top-left
  • scale / temp / cap 互斥条件错
  • GQA 没 repeat
  • 输出 layout 跟 Q 不一致
  • THD 忘 cu_seqlens
  • QK-Norm shape 展平错误
  • Online lse dtype 不是 fp32
  • 数值和 Offline 对不齐
节奏

建议拆天做

  1. D1：SDPA + scale + causal/full
  2. D2：window + 矩形 bottom-right
  3. D3：GQA + pack 格式
  4. D4：layout / THD
  5. D5：稳定技巧 + QK-Norm
  6. D6：可选 Online
策略

时间不够怎么砍

保 Task1 Offline 满分 → A5 才能拼
Online / 花活后置

Offline 优先 Online = Bonus
链接

A4 资料

  • Blog：A4 Attention Module
  • 本地：~/Documents/nju-llm/a4
  • FlashAttention interface · Llama Attention
  • 下一份：A5 Transformer Block
1 / 1

快捷键