0.1.0rc1 peaks 0.1 GB higher than 5ee2256 on fused attention at 32K tokens #119

import resource
import numpy as np, pyarrow as pa
from datafusion import SessionContext
from einfold.datafusion import EinFold

N, D = 32768, 64
rng = np.random.default_rng(0)
t, d = np.repeat(np.arange(N), D), np.tile(np.arange(D), N)
einfold = EinFold()
ctx = SessionContext().with_extensions(einfold)
for n in "qkv":
    ctx.from_arrow(pa.table({"t": t, "d": d, "x": rng.normal(size=N * D)}), n)
sql = """
WITH scores 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 WHERE q.t >= k.t GROUP BY q.t, k.t),
weights AS (SELECT i, j, exp(s - MAX(s) OVER (PARTITION BY i)) AS w FROM scores),
norm AS (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
FROM weights w JOIN v ON w.j = v.t JOIN norm n ON w.i = n.i
GROUP BY w.i, v.d"""
peak = lambda: resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6
before = peak()
for _ in range(3):
    ctx.sql(sql).collect()
print(f"pipelines={einfold.plans_with_pipelines()} peak RSS {before:.2f} -> {peak():.2f} GB (+{peak() - before:.2f})")