Repository navigation
A sliding-window causal mask (q.t - k.t < w) keeps attention from fusing #96
Description
🌀 Claude: A sliding-window causal mask, q.t >= k.t AND q.t - k.t < 32, keeps attention from being fused. That's the banded attention of Mistral and Longformer. The same query with only q.t >= k.t fuses. The result is still correct: DataFusion runs the query as plain SQL.
Repro (einfold wheel built from 71b4674, datafusion 55.0.0, Python 3.12):
import numpy as np, pyarrow as pa
from datafusion import SessionContext
from einfold.datafusion import EinFold
N, D = 256, 64
rng = np.random.default_rng(0)
t, d = np.repeat(np.arange(N), D), np.tile(np.arange(D), N)
def attention(mask):
return f"""
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 {mask} 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"""
NW = """
WITH kernel AS (
SELECT t.p, s.q, 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 GROUP BY k.p"""
for name, sql in [("causal", attention("q.t >= k.t")),
("sliding window", attention("q.t >= k.t AND q.t - k.t < 32")),
("nadaraya-watson", NW)]:
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)
p = rng.uniform(size=(64, 2)); s = rng.uniform(size=(N, 3))
ctx.from_arrow(pa.table({"p": np.arange(64), "x": p[:, 0], "y": p[:, 1]}), "target")
ctx.from_arrow(pa.table({"q": np.arange(N), "x": s[:, 0], "y": s[:, 1], "v": s[:, 2]}), "source")
ctx.sql(sql).collect()
print(f"{name:<16} pipelines={einfold.plans_with_pipelines()} einfold={einfold.plans_with_einfold()}")causal pipelines=1 einfold=1
sliding window pipelines=0 einfold=0
nadaraya-watson pipelines=0 einfold=0
(The last line is #97.) Writing the band as k.t BETWEEN q.t - 31 AND q.t doesn't fuse either.
Why, as far as I can tell: a fold's Condition compares two dimensions with no offset (left <op> right), so i - j < w has no representation. The fused kernel already takes, for each query position, an interval allowed(x) of key positions, and it skips a tile pair the interval rules out entirely. So a band would need an offset in Condition, and a detector rule for a - b < c, a < b + c and BETWEEN. The kernel already handles intervals. With a band, the visited tiles fall from about N²/2 to about N·w, which is the point of sliding-window attention.
Found from xarray-sql's demo 1 (benchmarks/ml/01_attention.py on the ml-demos branch), while checking that einfold's rule fuses attention variants, not only demo 1's spelling.
Activity