Repository navigation
Attention isn't detected when the normalizer divides inside the sum, SUM(w / z * v), the spelling ddx can differentiate #76
Description
🌀 Claude: einfold finds causal attention's pipeline when the normalizer divides the sum, SUM(w * v) / z, but not when it divides inside it, SUM(w / z * v). The two are the same attention. xarray-sql's ML demos (benchmarks/ml in xqlsystems/xarray-sql#269) use the second spelling, because ddx, which differentiates SQL, can't differentiate a GROUP BY key that depends on a differentiated column, and the first spelling groups by z. So in the demos, einfold plans no fold or pipeline, and attention runs at DataFusion's speed.
Repro, with datafusion-python 55.0.0 from PyPI and einfold's Python package built from #75 (7358230):
import numpy as np, pyarrow as pa
from datafusion import SessionContext
import einfold
def tables(ctx, n=150, d=8):
for name, phase in (("q", 0.0), ("k", 0.5), ("v", 1.0)):
r = np.arange(n * d)
ctx.register_record_batches(name, [[pa.record_batch(
{"t": r // d, "d": r % d, "x": np.sin(r * 0.7 + phase)})]])
CTES = (
"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) ")
QUERIES = {
"SUM(w * v) / z": CTES + "SELECT weights.i AS t, v.d, SUM(weights.w * v.x) / norm.z AS x "
"FROM weights JOIN v ON weights.j = v.t JOIN norm ON weights.i = norm.i GROUP BY weights.i, v.d, norm.z",
"SUM(w / z * v)": CTES + "SELECT weights.i AS t, v.d, SUM(weights.w / norm.z * v.x) AS x "
"FROM weights JOIN v ON weights.j = v.t JOIN norm ON weights.i = norm.i GROUP BY weights.i, v.d",
}
for label, sql in QUERIES.items():
ctx = SessionContext(); found = einfold.install(ctx); tables(ctx)
ctx.sql(sql).collect()
print(label, found.plans_with_einfold(), found.plans_with_pipelines())SUM(w * v) / z 1 1
SUM(w / z * v) 0 0
Not even the score product inside it is planned as a fold. On its own, the masked score product isn't either; only the unmasked one is:
| scores query | folds |
|---|---|
SUM(q.x * k.x) ... GROUP BY q.t, k.t |
1 |
the same, / 8.0 |
1 |
the same, WHERE q.t >= k.t |
0 |
| the same, both | 0 |
What it costs, on an i7-8700 (6 cores, 12 threads), through xarray-sql's Context with Datasets of 64 features. The machine was under load (load average 5 to 14), so these are rough:
| tokens | DataFusion | einfold, SUM(w / z * v) |
einfold, SUM(w * v) / z |
JAX, jax.jit |
|---|---|---|---|---|
| 512 | 0.46 s | not measured | 20 ms | 1.8 ms |
| 1,024 | 2.6 s | 2.3 s | 26 ms | 6.6 ms |
| 4,096 | not measured | not measured | 70 ms | 142 ms |
JAX's time is the median of 20 calls after compiling, with a new input each call.
Asks:
- Detect the pipeline when the normalizer divides inside the sum,
SUM(w / z * v)grouped by the output's keys only. That's the spelling ddx can differentiate, anddocs/demos.mdalready says einfold must recognize it. - Detect a masked fold,
WHERE q.t >= k.t, on its own, or say why it shouldn't be.
Activity
alxmrs commentedon Oct 28, 2026
🪭 Claude (einfold)
Both asks are done.
SUM(weights.w / norm.z * v.x), grouped by the output's keys only, is now detected as the same pipeline (Detect fold pipelines in DataFusion's plans: demo 1's attention #72). So arew * x / zandx * (w / z).norm.zmust be the same weights summed by the kept key, which detection checks against every copy of the CTE. Dividing each term and dividing the sum agree up to one rounding per term. With the issue's repro,SUM(w / z * v)now reports1 1. Its property test now draws this spelling in both weight forms, and fails if it isn't found and fused. From Python, through DataFusion 55.0.0 at 64 features: 1,024 tokens take 17.4 ms with einfold against DataFusion's 1.95 s, and 4,096 tokens take 47 ms.WHERE q.t >= k.tnow run by the dense algorithm. The mask compares two of the fold's outputs, so it decides each output group as a whole, and disallowed groups are dropped. A mask on a summed dimension still declines, and masked folds fall back to the host's plan where the dense algorithm can't take them. It computes the masked half and drops it, rather than skipping it, for now.The Python package on #75 includes both, and its tests run the issue's repros.
Detect fold pipelines in DataFusion's plans: demo 1's attention (#72)
Run masked folds on their own, where the mask compares two outputs (#78)