Mission 1
TODO 1 + 两层分块循环
先读懂函数：分块不改答案，只改计算顺序
🎯 先猜一猜（下注解锁）
代码一开始还没有读过任何 score，却要先给每行的最大值 m 一个初值。
m 为什么应该从负无穷开始，而不是从 0 开始？
先押一个假设，下面的讲解就会解锁。猜错完全没关系——带着疑问读，记得最牢。
先补的知识
- 本 Notebook 暂时没有 Batch 和 Head 维度，q、k、v 都是 [seq_len, dim]。先把二维版本学明白，再迁移到真实多头实现。
- Python 切片 start:start+block_size 即使越过末尾也不会报错，它会自然返回较短的最后一块。
- FlashAttention 仍然计算精确 Attention；这里的核心是逐块读取 K/V，并且不把完整 [seq_len, seq_len] 分数矩阵长期保存下来。
- 对每一行 query，算法只需要持续记住三个状态：当前输出 out、见过的最大值 m、指数和 l。
图解原理
先把代码当成两层翻页：外层固定一小块 Q，内层把 K/V 一块块翻完。每翻一页，就把这一页的信息并入 m、l 和 out；翻完全部 K/V 后，这块 Q 的答案就完成了。
1. 固定 Q 块一次只处理几行 query
2. 扫描 K/V 块内层循环逐块读取
3. 保存滚动状态每行只保留 out、m、l
普通写法：先造出完整分数表
Q [S,D]K^T [D,S]Scores [S,S]
序列变长时，S×S 的中间矩阵按平方增长。这里的模拟作业，就是为了理解如何避免长期保存它。
分块写法：固定 Q 块，逐块读 K/V
Q_i [Bq,D]K_j^T [D,Bk]S_ij [Bq,Bk]
每次只出现一个小分数块。数学结果没有近似，变化的是计算顺序和中间数据的生存时间。
out[seq_len, dim]，每个 query 当前已经归一化的输出
m[seq_len, 1]，每行目前见过的最大 score；初始值是 -inf
l[seq_len, 1]，以当前 m 为基准的指数和；初始值是 0
为什么是列向量[seq_len, 1] 可以沿最后一维广播到 score 块的每一列
初始化的含义比写法更重要：还没看任何 K 块，所以输出贡献和指数和都是 0；最大值设为 -inf，保证第一块中任何正常分数都能成为新最大值。
语法热身：用另一个张量练初始化、切片和最后一块
rows, width = 5, 3
x = torch.randn(rows, width)
buffer = torch.zeros((rows, width), device=x.device, dtype=x.dtype)
running_max = torch.full((rows, 1), -float('inf'), device=x.device, dtype=x.dtype)
running_sum = torch.zeros((rows, 1), device=x.device, dtype=x.dtype)
for start in range(0, rows, 3):
chunk = x[start:start + 3]
print(chunk.shape) # [3, 3]，然后是 [2, 3]从例子迁移到 TODO 1
buffer对应需要返回的输出状态，shape 要同时包含序列行和特征列。running_max与running_sum都是每个 query 一份统计量，所以第二维保留为 1。device和dtype跟随输入，可以避免 CPU/GPU 或精度不一致。- 不要手工补齐最后一块；Notebook 的第二个测试专门覆盖 seq_len 不能整除 block_size 的情况。
巩固一下
seq_len=5、block_size=3 时，第二个切片 q[3:6] 的 shape 是什么？
回到 notebook 的作业
这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。