Mission 1
概念练习
先建状态账本：分块改变计算顺序，不改变 Attention 答案
🎯 先猜一猜
seq_len=5、block_size=3，外层循环第二次执行 i=3。
q[3:6] 的实际 shape 是什么？
先猜一个答案，再对照下面的讲解。可以随时修改选择，猜错会得到解释。
先补的知识
- 本题刻意省略 Batch 和 Head 维度，q、k、v 都是 [seq_len, dim]；返回 out 也必须是 [seq_len, dim]。
- 外层固定一块 Q，内层遍历所有 K/V 块。Python 切片允许终点超过长度，所以最后一块可以比 block_size 短。
- 每个 query 行都维护自己的最大值 m 和指数和 l，因此它们是 [seq_len, 1] 列向量，便于向 score 的列方向广播。
- q_block 已在外层乘 scale = 1/sqrt(dim)，内层计算分数时不能再次缩放。
图解原理
把内层循环看成逐页读取 K/V：每读一页，就把这页对当前 Q 块的贡献合并进三个滚动状态。out 记当前归一化输出，m 记目前见过的最大 score，l 记以 m 为基准的指数和。翻完所有页后，这块 Q 的精确 Attention 才完成。
外层 i固定 q[i:i+block_size]
内层 j逐块读取 k 和 v
滚动合并更新 out_i、m_i、l_i
写回保存当前 Q 块的最终状态
标准 Attention 的大中间量
Q [S,D]Kᵀ [D,S]Scores [S,S]
完整分数矩阵会随序列长度按平方增长。
当前模拟只保留一个小分数块
Q_i [Bq,D]K_jᵀ [D,Bk]S_ij [Bq,Bk]
结果仍是精确 Attention，改变的是中间量的生存时间。
out[seq_len, dim]；初始为 0，逐块形成最终输出m[seq_len, 1]；初始为 -inf，表示尚未看过任何 scorel[seq_len, 1]；初始为 0，表示尚无指数贡献scalePython float，值为 1 / sqrt(dim)device新张量至少跟随 q.device；可见测试使用 CPU float32
m 不能从 0 开始。若第一块 score 全是负数，0 会变成一个数据中从未出现的假最大值；-inf 才表示“空状态”。
语法热身：用分批处理成绩表练初始化和尾块切片
import torch
rows, width = 5, 3
scores = torch.randn(rows, width)
summary = torch.zeros((rows, width), device=scores.device)
running_max = torch.full((rows, 1), -float('inf'), device=scores.device)
running_sum = torch.zeros((rows, 1), device=scores.device)
for start in range(0, rows, 3):
batch = scores[start:start + 3]
print(batch.shape) # [3,3]，然后 [2,3]从语法例子迁移到 TODO 1
summary：对应 out，行数与序列长度一致，列数与 value 特征维一致running_max：对应 m；每一行只保存一个最大值，所以末维是 1running_sum：对应 l；每一行只保存一个指数和，所以末维是 1scores.device：对应 q.device，避免创建在错误设备上的状态张量
巩固一下
为什么 m 和 l 使用 [seq_len, 1]，而不是 [seq_len]？
学完这一段，试着做
用一个小动作确认自己理解了；最后再进入官方题目。