Show HN：从约 25 行 SQL 中自动推导 FlashAttention 内核

# require-python = ">=3.11" # "xarray-sql[ml,duckdb]", # "einfold[datafusion]>=0.1.0rc2", # xarray-sql = { path = "../../", editable = true } """演示 1：SQL 中的因果注意力，在多个引擎上，具有梯度。注意力是 Transformer 的核心，是“softmax(Q·Kᵀ/√d)·V”：每个标记的输出是值向量“V”的平均值，根据“Q”中标记的查询与“K”中每个早期标记的键的匹配程度进行加权。这里“Q”、“K”和“V”是从 xarray 数据集注册的“(t, d, x)”表（标记、特征、值），注意力是两个总和连接：分数、特征上的“Q”乘以“K”的总和，以及由“exp(score)”加权的“V”的平均值，当 SQL 引擎运行它时，这是一个简单的计划：它存储每一个 N×N 分数（Dao 等人，2022，https://arxiv.org/abs/2205.14135）通过流式传输键块来计算相同的结果。 einfold (https://github.com/xqlsystems/einfold) 是一个将 SQL 计划重写为数组算法的查询优化器，在下面的两种拼写中找到该计划，没有为注意编写规则：它的规则融合了任何其分数是通过连接计算的 softmax 加权平均值，并且相同的规则融合了下面的控制，这是 einfold 的三个目标演示中的第一个。 （其存储库中的“docs/demos.md”）： * **forward**：每个 *host* 上模型的每个拼写（einfold 表示运行计划的引擎；这里是 DataFusion、带有 einfold 的 DataFusion 和 DuckDB），根据 JAX（Google 的可微数组程序库）中的相同注意进行检查； * **gradient**：读取模型的损失及其相对于 Q、K 和 V 的每个值的梯度，如下所示一个 SQL 语句（https://github.com/xqlsystems/ddx），它区分 SQL 查询，计算它，并检查 * **掩码是否存在**：V 中的无穷大仅到达允许看到它的查询，因为 SQL 的语义需要任何更快的计划； * **不注意的控制**：Nadaraya-Watson 估计器，用于重新网格化的高斯内核加权平均值。任何加速注意力的规则也必须加速这个查询； * **einfold 找到融合计划**：每个拼写和控件都作为 einfold 的融合管道之一运行。 ( "jax_enable_x64" , True ) D = 64 # 每个头的特征；分数按 1/√64 = 1/8 缩放。 # 模型和它的损失是分开的，就像用 Python 编写的神经网络一样（比较下面的 focus() 和 loss()）。 # einfold 的窗口拼写中的注意力是一个连接 # 谓词：屏蔽的 (i, j) 对是不存在的行 # 减去每行的最大值只会使 exp() 免于溢出：它取消了 # 比率 w / z，因此其梯度总和为零，并且 ddx 将其区分 CREATE OR REPLACE VIEW Attention AS SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s FROM q JOIN k ON q.d = k.d SELECT i, j, exp(s - MAX(s) OVER (PARTITION BY i)) AS w SELECT i, SUM(w) AS z FROM Weights GROUP BY i) SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x # 同样注意正如 ddx 的拼写：每行的最大值作为 GROUP BY，# CREATE OR REPLACE VIEW Attention AS SELECT q.t AS i, k.t AS j, SUM(q.x * k.x) / 8.0 AS s FROM q JOIN k ON q.d = k.d SELECT i, MAX(s) AS m FROM Scores GROUP BY i), SELECT s.i, s.j, exp(s.s - r.m) AS w FROM 分数 s JOIN row_max r ON s.i = r.i), SELECT i, SUM(w) AS z FROM 权重 GROUP BY i) SELECT w.i AS t, v.d, SUM(w.w / n.z * v.x) AS x SPELLINGS = { "window" : ATTENTION_WINDOW , "group by" : ATTENTION_GROUP_BY } # 前向传递：运行模型。 FORWARD = "SELECT t, d, x FROM Attention ORDER BY t, d" # 损失是输出平方范数的一半，通过读取 # 视图来调用模型。 SELECT 询问其相对于 Q、# K 和 V 的每个值的梯度。三个 grad() 调用区分一个损失，因此 ddx 对这三个损失运行一次反向传递，如下所示jax.grad(loss, argnums=(0, 1, 2)) 这样做。 SELECT SUM(x * x) / 2 AS l FROM Attention) JOIN k ON q.t = k.t AND q.d = k.d JOIN v ON q.t = v.t AND q.d = v.d # Nadaraya–Watson：在每个目标点，源 y 的平均值由其距离的高斯核加权(带宽 h = 0.5). # 结构注意力：权重 exp(score) 的加权平均值 SELECT t.p, sq, exp(-(power(t.x - s.x, 2) + power(t.y - s.y, 2)) / 0.25) AS w FROM target t CROSS JOIN source s) SELECT k.p, SUM(k.w * s.v) / SUM(k.w) AS v FROM kernel k JOIN source s ON k.q = s.q """参考：因果softmax注意力，如JAX