Batched attention's backward still joins the softmax normalizer on one key (sequences × heads too many rows inside a step; 12.8 s of 30 s with #220) #222

"""Batched attention's backward on DuckDB: the largest relation inside each of
ddx's steps, and the joins that pair every sequence and head."""
import re
import duckdb, numpy as np, pyarrow as pa
import ddxdb.duckdb as dd

S, H, T, C = 3, 2, 8, 4  # sequences, heads, positions, channels per head
inner = []
orig = dd.DuckDBBackend.materialize
def materialize(self, name, plan):
    sql = dd.to_sql(plan, "duckdb")
    if sql.startswith("WITH "):
        head = sql[: sql.rindex(") SELECT ") + 1]
        for cte, body in re.findall(r"(__ddx_q\d+) AS \((.*?)\)(?=, __ddx_q\d+ AS \(|$)", head[5:]):
            n = self.con.execute(f"{head} SELECT count(*) FROM {cte}").fetchone()[0]
            on = re.search(r" ON (.*)$", body)
            inner.append((name.split("_", 3)[-1], cte, n, on.group(1) if on else ""))
    orig(self, name, plan)
dd.DuckDBBackend.materialize = materialize

con = duckdb.connect()
idx = np.indices((S, T, H, C)).reshape(4, -1)
for name, phase in (("q", 0.0), ("k", 1.0), ("v", 2.0)):
    con.register(name, pa.table({"seq": idx[0], "pos": idx[1], "h": idx[2], "c": idx[3],
                                 "val": np.sin(np.arange(idx.shape[1]) * 0.37 + phase)}))
LOSS = """WITH scores AS (
  SELECT q.seq, q.h, q.pos AS i, k.pos AS j, SUM(q.val * k.val) / 2.0 AS s
  FROM q JOIN k ON q.seq = k.seq AND q.h = k.h AND q.c = k.c WHERE q.pos >= k.pos
  GROUP BY q.seq, q.h, q.pos, k.pos),
weights AS (SELECT seq, h, i, j, exp(s - MAX(s) OVER (PARTITION BY seq, h, i)) AS p FROM scores),
norm AS (SELECT seq, h, i, SUM(p) AS z FROM weights GROUP BY seq, h, i),
heads AS (
  SELECT a.seq, a.i AS pos, a.h, v.c, SUM(a.p / n.z * v.val) AS val
  FROM weights a
  JOIN v ON a.seq = v.seq AND a.h = v.h AND a.j = v.pos
  JOIN norm n ON a.seq = n.seq AND a.h = n.h AND a.i = n.i
  GROUP BY a.seq, a.i, a.h, v.c),
loss AS (SELECT SUM(val * val) / 2 AS l FROM heads) """
dd.sql(con, LOSS + "SELECT seq, pos, h, c, grad(loss, val) AS g FROM q ORDER BY seq, pos, h, c").fetchall()
print(f"{S * H * T * (T + 1) // 2} weights rows, {S * H * T} norm rows")
for step, cte, n, on in sorted(inner, key=lambda r: -r[2])[:8]:
    print(f"  {step:>14} {cte:>9}: {n:5} rows" + (f"  ON {on}" if on else ""))