ggaaooppeenngg

FlashAttention 没有发明新的数学。它是 online softmax 递推与两个 GEMM 在同一个 tile 循环里的交织–分数块 $S$ 算出来就地做 softmax 得到 $\tilde{P}$，$\tilde{P}$ 就地乘上 $V$，中间矩阵全程驻留在 SRAM，HBM 流量从 $O(N^2)$ 降到 $O(N \cdot d)$。本文用 TileLang（v0.1.13）把 FA-2 论文的 Algorithm 1 写成可运行的 kernel：先补齐数学地基，再讲清楚"怎么切"，然后逐行拆解主循环，最后做数值验证、测速与调参。只覆盖前向；TileLang 五原语与 T.Pipelined 流水线机制见上一篇《TileLang 编程基本知识点》，本文直接使用其结论。

TileLang 的定位可以用一句话概括：把 CUDA kernel 里「怎么切 tile、数据放哪级存储、什么时候预取」这几件事变成显式的 Python 语句，其余的寄存器分配、指令选择、同步插入交给编译器。它不是又一个 Triton–Triton 隐藏 shared memory，TileLang 让你直接写 T.alloc_shared。这个差别决定了两者能碰到的性能天花板不同。

本文整理 TileLang 的编程基本知识点：它是什么、编程模型的 5 个原语、循环原语与 T.Pipelined 的工作原理、同一份代码在不同 GPU 架构上如何 lowering、四个进阶主题（Split-K / warp specialization / autotune / Blackwell 两条路径），以及写完 kernel 之后的三个标准动作（验证 / dump 源码 / bench）。

0. 为什么要固定大小的记忆状态

标准 softmax 注意力解码时必须缓存全部历史的 Key/Value，显存和单步延迟都是 $O(n \cdot d)$，上下文翻倍就一起翻倍。把历史压缩成一个固定大小的状态 $S \in \mathbb{R}^{d \times d}$，显存和单步开销变成 $O(d^2)$，与 $n$ 无关。

代价是这个状态必须会写、会改、会忘。两条路线各给出一半答案：线性注意力造出固定状态，SSM 教它怎么遗忘。