Show HN: Automatic deduction of FlashAttention kernel from ~25 lines of SQL

A Python benchmark demonstrates how causal softmax attention can be expressed as SQL and optimized into a FlashAttention-style computation. It represents queries, keys, and values as tables, computes query-key scores through joins and aggregation, applies a causal mask by excluding future-token pairs, then forms a softmax-weighted average of value vectors. Two equivalent SQL formulations calculate the row-wise softmax maximum differently: one uses a window function, the other a grouped maximum joined back to the scores. The central claim is that einfold, a query optimizer that rewrites SQL plans as array algorithms, recognizes the shared structure and fuses either formulation into a streaming plan that avoids storing the full matrix of token-to-token scores. The optimizer uses a general rule for join-derived, softmax-weighted averages, rather than a rule written specifically for attention.

The script is designed to test correctness and differentiation across DataFusion, DataFusion with einfold, and DuckDB, comparing forward outputs with a JAX reference. Where a SQL differentiation backend is available, it also computes gradients of a scalar loss with respect to every element of Q, K, and V, and checks them against JAX. Further checks ensure masked pairs behave as absent rows - even when values contain infinity - and show that the same optimization applies to a structurally similar Gaussian-kernel weighted average, the Nadaraya-Watson estimator. The benchmark reports execution operators and timings for configurable sequence lengths; its emphasis is validating the rewrite and its semantics, not presenting performance results in the supplied content.