1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
# /// script
# requires-python = ">=3.11"
# dependencies = [
# "xarray-sql[ml,duckdb]",
# "ddxdb[duckdb]",
# "einfold[datafusion]>=0.1.0rc2",
# "jax",
# ]
#
# [tool.uv.sources]
# xarray-sql = { path = "../../", editable = true }
# ///
"""Demo 1: causal attention in SQL, on several engines, with its gradient.
Attention, the core of a transformer, is ``softmax(Q·Kᵀ/√d)·V``: each
token's output is an average of the value vectors ``V``, weighted by how
well the token's query in ``Q`` matches each earlier token's key in ``K``.
Here ``Q``, ``K`` and ``V`` are tables of ``(t, d, x)`` (token, feature,
value), registered from xarray Datasets, and attention is two sums over
joins: the scores, a sum of ``Q`` times ``K`` over features, and the
average of ``V`` weighted by ``exp(score)``.
As a SQL engine runs it, this is the naive plan: it stores every one of
the N×N scores. FlashAttention (Dao et al., 2022,
https://arxiv.org/abs/2205.14135) computes the same result without storing
them, by streaming blocks of keys past blocks of queries and keeping a
running softmax. einfold (https://github.com/xqlsystems/einfold), a query
optimizer that rewrites SQL plans into array algorithms, finds that plan
from this SQL, in both of the spellings below, with no rule written for
attention: its rule fuses any softmax-weighted average whose scores are
computed from a join, and the same rule fuses the control below, which is
not attention. This is the first of einfold's three target demos
(``docs/demos.md`` in its repository). The script checks:
* **forward**: each spelling of the model on each *host* (einfold's word
for an engine that runs a plan; here DataFusion, DataFusion with
einfold, and DuckDB), checked against the same attention in JAX,
Google's library for differentiable array programs;
* **gradient**: a loss that reads the model, and its gradient with
respect to every value of Q, K and V, as one SQL statement. ddx (https://github.com/xqlsystems/ddx),
which differentiates SQL queries, computes it, and ``jax.grad`` checks
it;
* **masks are absent rows**: an infinity in V reaches only the queries
allowed to see it, as SQL's semantics require of any faster plan;
* **a control that is not attention**: the Nadaraya–Watson estimator, a
Gaussian-kernel weighted average used in regridding. Any rule that
speeds up attention must speed up this query too;
* **einfold finds the fused plan**: each spelling, and the control, runs
as one of einfold's fused pipelines.
Run ``uv run benchmarks/ml/01_attention.py --help`` for the options.
"""
from __future__ import annotations
import jax
import jax.numpy as jnp
import numpy as np
import xarray as xr
from _bench import Report, assert_close, cli, hosts
from _hosts import EinFold
import xarray_sql as xql
jax.config.update("jax_enable_x64", True)
D = 64 # features per head; the scores are scaled by 1/√64 = 1/8.
# The model and its loss are separate, as in a neural network written in
# Python (compare attention() and loss() below). The model is a view, a
# named query: each spelling of attention defines the view ``attention``,
# and the forward pass and the loss read it as if it were a table.
# Attention in einfold's window spelling. The causal mask is a join
# predicate: a masked (i, j) pair is a row that does not exist. Subtracting
# each row's max only keeps exp() from overflowing: it cancels out of the
# ratio w / z, so its gradient sums to zero, and ddx differentiates it
# like any other window.
ATTENTION_WINDOW = """
CREATE OR REPLACE VIEW attention AS
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
"""
# The same attention as ddx spells it: each row's max as a GROUP BY,
# joined back to the scores.
ATTENTION_GROUP_BY = """
CREATE OR REPLACE VIEW attention AS
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),
row_max AS (
SELECT i, MAX(s) AS m FROM scores GROUP BY i),
weights AS (
SELECT s.i, s.j, exp(s.s - r.m) AS w
FROM scores s JOIN row_max r ON s.i = r.i),
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
"""
SPELLINGS = {"window": ATTENTION_WINDOW, "group by": ATTENTION_GROUP_BY}
# The forward pass: run the model.
FORWARD = "SELECT t, d, x FROM attention ORDER BY t, d"
# The loss, half the output's squared norm, calls the model by reading the
# view. The SELECT asks for its gradient with respect to every value of Q,
# K and V. The three grad() calls differentiate one loss, so ddx runs one
# backward pass for all three, as jax.grad(loss, argnums=(0, 1, 2)) does.
GRADIENT = """
WITH loss AS (
SELECT SUM(x * x) / 2 AS l FROM attention)
SELECT q.t, q.d,
grad(loss, q.x) AS dq,
grad(loss, k.x) AS dk,
grad(loss, v.x) AS dv
FROM q
JOIN k ON q.t = k.t AND q.d = k.d
JOIN v ON q.t = v.t AND q.d = v.d
ORDER BY q.t, q.d
"""
# Nadaraya–Watson: at each target point, the average of the sources' y
# weighted by a Gaussian kernel of their distance (bandwidth h = 0.5).
# Structurally attention: a weighted average with weights exp(score).
NADARAYA_WATSON = """
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
ORDER BY k.p
"""
def attention(q, k, v):
"""The reference: causal softmax attention, as JAX writes it."""
n = q.shape[0]
scores = q @ k.T / np.sqrt(D)
causal = jnp.tril(jnp.ones((n, n), dtype=bool))
weights = jax.nn.softmax(jnp.where(causal, scores, -jnp.inf), axis=-1)
return weights @ v
def loss(q, k, v):
return jnp.sum(attention(q, k, v) ** 2) / 2
def table(values: np.ndarray) -> xr.Dataset:
"""A (t, d) matrix as the XQL table ``(t, d, x)``."""
t, d = values.shape
return xr.Dataset(
{"x": (("t", "d"), values)},
coords={"t": np.arange(t), "d": np.arange(d)},
)
def matrix(rows, name: str = "x") -> np.ndarray:
"""The rows of a ``(t, d, …)`` result as a (t, d) matrix."""
return xql.to_dataset(rows, dims=["t", "d"])[name].values
def main():
parser = cli(__doc__.splitlines()[0], sizes=[64, 256], size_name="tokens")
args = parser.parse_args()
report = Report("attention", "tokens", args.reps)
rng = np.random.default_rng(0)
for n in args.sizes:
q, k, v = (rng.normal(size=(n, D)) for _ in range(3))
want = attention(q, k, v)
want_grads = jax.grad(loss, argnums=(0, 1, 2))(q, k, v)
for host in hosts(args.hosts):
for name, data in zip("qkv", (q, k, v)):
host.load(name, table(data))
for spelling, model in SPELLINGS.items():
host.execute(model)
rows = report.time(
f"forward ({spelling})", host, n, host.query, FORWARD
)
report.check(assert_close("forward", matrix(rows), want))
report.note("operators", host.operators(FORWARD))
step = f"gradient ({spelling})"
if not host.differentiates:
report.skip(step, host, n, "no ddx backend")
continue
rows = report.attempt(step, host, n, host.query, GRADIENT)
if rows is None:
continue
errors = [
assert_close(f"d/d{name}", matrix(rows, f"d{name}"), ref)
for name, ref in zip("qkv", want_grads)
]
report.check(max(errors))
check_masks_are_absent_rows(args.hosts)
check_nadaraya_watson(args.hosts, report)
check_einfold_fuses(args.hosts)
report.finish(args.csv)
def check_masks_are_absent_rows(host_names):
"""An infinity in V at token j reaches the queries at or after j, and
no others: a masked pair is skipped, not multiplied by a zero weight,
which would turn 0·∞ into NaN. (xarray-sql reads NaN as missing data,
a NULL that SUM skips, so an infinity is the value that tells.)"""
n, j = 16, 9
rng = np.random.default_rng(1)
q, k, v = (rng.normal(size=(n, D)) for _ in range(3))
v[j, 0] = np.inf
print(f"\n masks: +inf in V at token {j}, feature 0")
for host in hosts(host_names):
for name, data in zip("qkv", (q, k, v)):
host.load(name, table(data))
for spelling, model in SPELLINGS.items():
host.execute(model)
out = matrix(host.query(FORWARD))
assert np.isfinite(out[:j]).all(), (host.name, spelling)
assert np.isposinf(out[j:, 0]).all(), (host.name, spelling)
assert np.isfinite(out[:, 1:]).all(), (host.name, spelling)
print(f" {spelling:<9} {host.name:<11} ok")
def check_nadaraya_watson(host_names, report):
"""The control: a kernel-weighted average that is not attention."""
rng = np.random.default_rng(2)
targets, sources = rng.uniform(size=(64, 2)), rng.uniform(size=(256, 2))
y = np.sin(4 * sources[:, 0]) + np.cos(3 * sources[:, 1])
dist2 = ((targets[:, None, :] - sources[None, :, :]) ** 2).sum(-1)
w = np.exp(-dist2 / 0.25)
want = (w @ y) / w.sum(axis=1)
target = xr.Dataset(
{"x": ("p", targets[:, 0]), "y": ("p", targets[:, 1])},
coords={"p": np.arange(len(targets))},
)
source = xr.Dataset(
{"x": ("q", sources[:, 0]), "y": ("q", sources[:, 1]), "v": ("q", y)},
coords={"q": np.arange(len(sources))},
)
for host in hosts(host_names):
host.load("target", target)
host.load("source", source)
rows = report.time(
f"nadaraya-watson ({len(sources)} sources)",
host,
len(targets),
host.query,
NADARAYA_WATSON,
)
got = xql.to_dataset(rows, dims=["p"])["v"].values
report.check(assert_close("nadaraya-watson", got, want))
def check_einfold_fuses(host_names):
"""einfold runs each spelling of attention, and the control, as one
fused pipeline, computed tile by tile with a running softmax, rather
than as the joins and aggregates the SQL describes."""
if EinFold.name not in host_names:
return
print("\n einfold's plans")
host = EinFold()
rng = np.random.default_rng(3)
for name in "qkv":
host.load(name, table(rng.normal(size=(128, D))))
for spelling, model in SPELLINGS.items():
host.execute(model)
before = host.einfold.plans_with_pipelines()
host.query(FORWARD)
assert host.einfold.plans_with_pipelines() > before, spelling
print(f" {spelling:<16} fused")
targets, sources = rng.uniform(size=(16, 2)), rng.uniform(size=(64, 2))
host.load(
"target",
xr.Dataset(
{"x": ("p", targets[:, 0]), "y": ("p", targets[:, 1])},
coords={"p": np.arange(len(targets))},
),
)
host.load(
"source",
xr.Dataset(
{
"x": ("q", sources[:, 0]),
"y": ("q", sources[:, 1]),
"v": ("q", sources.sum(axis=1)),
},
coords={"q": np.arange(len(sources))},
),
)
before = host.einfold.plans_with_pipelines()
host.query(NADARAYA_WATSON)
assert host.einfold.plans_with_pipelines() > before, "nadaraya-watson"
print(f" {'nadaraya-watson':<16} fused")
if __name__ == "__main__":
main()