25 行 SQL 写出注意力：einfold 从查询计划里推出 FlashAttention 式融合

xarray-sql 团队 10 月 30 日在 Hacker News 展示了一个实验：把因果注意力写成两段 SQL（窗口函数版和 GROUP BY 版），再由查询优化器 einfold 自动改写成不存储 N×N 分数矩阵的融合计划，也就是 FlashAttention 的思路。脚本用 JAX 校验前向结果和 Q/K/V 梯度，还加了一个「非注意力」的 Nadaraya–Watson 对照，确认优化规则不是只对注意力开小灶。

这个实验做了什么

xarray-sql 把 Xarray 数据集当成表来查。团队在 benchmarks/ml/01_attention.py 里，用两段 SQL 写出因果注意力：一段用窗口函数算每行的最大值，另一段用 GROUP BY 加自连接。Q、K、V 都是 (t, d, x) 三列的表，注意力就是两次 join 上的求和。脚本在 DataFusion、带 einfold 的 DataFusion 和 DuckDB 上各跑一遍，再用 JAX 的同一段注意力做数值对照。

关键不在 SQL，在优化器

朴素计划会把 N×N 的分数矩阵整个存下来。FlashAttention 的思路是不存，分块流式算，维护一个滚动 softmax。einfold 是一个把 SQL 计划改写成数组算法的优化器：它没有为注意力写专门规则，而是识别「带 softmax 权重的平均值，且分数来自 join」这个模式，自动融合。作者还特意加了一个不是注意力的对照——Nadaraya–Watson 估计，确认同一条规则也会加速它。

为什么值得看

它把「模型里的算子」和「数据库里的查询」放进了同一个优化框架。如果这类改写可靠，未来做推理或科学计算的人可以用声明式查询描述计算，让优化器决定用哪套 kernel，而不是手写 CUDA 或绑定某个框架。对数据平台来说，这也意味着窗口聚合、相似度检索这类查询有機會直接吃到融合执行计划。

现在还不能当生产方案

这是一个 benchmark 演示，不是经过生产验证的推理引擎；einfold 还是早期项目，规则覆盖和数值稳定性都需要更多测试。脚本本身值得读：它同时校验前向、梯度、掩码语义和一个非注意力对照，这种「连反例一起测」的做法，比只看一个加速比更有参考价值。

如果优化器能认出注意力，写 kernel 的人就该开始写规则了。