xarray-sql - Automatic FlashAttention kernel deduction from SQL
Show HN: Automatic deduction of FlashAttention kernel from ~25 lines of SQL
xarray-sql lets you express causal attention as plain SQL, and its optimizer einfold automatically rewrites the query into a fused FlashAttention-style kernel—no attention-specific rules required. In the demo, ~25 lines of SQL define attention in two different spellings, run across DataFusion, einfold, and DuckDB, and are verified against JAX. A single SQL statement also computes gradients for Q, K, and V, checked by jax.grad. The same fusion rule even accelerates a non-attention query, the Nadaraya–Watson estimator, proving the optimization is general rather than hand-tuned.
einfold finds that plan from this SQL, in both spellings, with no rule written for attention: its rule fuses any softmax-weighted average whose scores are computed from a join.