Repository navigation
Tracking: ML in SQL as fast as JAX on CPU #285
Description
🌀 Claude (xarray-sql ML demos)
This issue coordinates the work to make machine learning in SQL fast. The demos in benchmarks/ml (#269) write models as SQL over xarray-sql's tables, train them with ddx, which differentiates SQL queries, and check every result against JAX, Google's library for differentiable array programs. The goal is to match JAX on CPU with a SQL engine. Anyone can pick up a workstream below. Comment here to claim one, or to post a measurement.
Where we are
Intel i7-8700 (6 cores), DataFusion 55.0.0, DuckDB 1.5.6, ddx at its development branch after ddxdb 0.4.1, measured on 2026-10-24 and 2026-10-25. JAX is compiled with jax.jit.
| JAX (CPU) | DataFusion | DuckDB | |
|---|---|---|---|
| attention forward, 512 tokens | 0.92 ms | 0.28 s (window) / 0.37 s (GROUP BY) |
0.48 s / 0.27 s |
| attention gradient, 512 tokens | 2.4 ms | 0.94 s | 2.4 s (GROUP BY) |
| GPT loss, 1 layer, batch 4 | 0.15 ms | 3.3 s | 0.067 s |
| GPT SGD step | 0.37 ms | 0.69 s | refused (xqlsystems/ddx#195) |
| GPT loss, 2 layers | can't plan (runs out of 8 GB) | 0.12 s |
The two demos are slow for different reasons:
- The GPT spends its time planning, not running. DataFusion copies a CTE into the plan at every place that reads it, so the 1-layer loss is 13,756 operators, and a second layer multiplies that by about 100 (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267). Planning the 1-layer loss takes 2.1 s, and the same physical plan runs in 2.4 ms. A ddx training step is 76 small queries, each planned in 3.3 ms and run in 0.65 ms (prepare: plan each program step once, and run its physical plan on every training step ddx#202).
- Attention spends its time running joins. It stores all N×N scores, and its plan takes 3–4 ms of a 0.1–0.4 s query. Running faster is the job of einfold, a query optimizer that rewrites SQL plans into array algorithms; it would find FlashAttention's tiled plan, for one.
Workstreams
- Plan each shared CTE once. DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267. DataFusion 56 (due Oct/Nov 2026, Release DataFusion
56.0.0(Oct/Nov 2026) apache/datafusion#24461) addsWITH … AS MATERIALIZED(feat: supportWITH ... AS [NOT] MATERIALIZEDCTEs apache/datafusion#25580), and feat: decide automatically when to materialize a CTE apache/datafusion#26276 would materialize CTEs automatically. Once 56 is out: enable it inContext, check it with xarray-sql's tables, and train the GPT at 2+ layers. ddx must read the new plan nodes too. - Plan once, run many times. Like
jax.jit, which compiles once.- xarray-sql:
Context.prepareis on theprepared-queriesbranch. It plans a query once and runs it on new values. It needs a reset of the plan's state between runs that datafusion-python doesn't expose: datafusion-python: a physical plan can't be run again with its state reset (expose reset_plan_states) — draft for upstream #283, a draft to send upstream. Until then it decodes a fresh copy of the plan, which works only for plans shallow enough to decode. - ddx: prepare: plan each program step once, and run its physical plan on every training step ddx#202, the same for ddx's program steps.
- xarray-sql:
- ddx's per-step overhead. program: choose step boundaries by what they save (fewer, larger steps) ddx#186 (fewer, larger steps), rewrite: a cached program still optimizes each shared CTE's placeholder (15% of a cached GPT step) ddx#184, run: run a program's independent steps concurrently ddx#157 (run independent steps concurrently).
- The FFI cost per scan. A lazy table costs about 1.5 ms per scan through
datafusion-ffi(datafusion-ffi rebuilds the session and task context, with every function, on each scan and execute (~1.5 ms per scan) #273). Small in-memory Datasets avoid it (Register small in-memory Datasets as native in-memory tables #275); large ones still pay it. - DuckDB.
- ddx's DuckDB backend can't yet plan a CTE that DuckDB computes once for several reads (DuckDB: a loss that reads a view with a multiply-read CTE fails (Not implemented Error: CTE) ddx#195).
- DuckDB orders joins about 3× worse over an Arrow scan, which is how
xql.registergives it a table, than over a native table: 0.64 s against 0.20 s for one gradient step. Needs an issue with a minimal repro.
- Faster execution: einfold. einfold would run as a DataFusion optimizer rule, reading the dense-grid layout that xarray-sql knows about each table. That contract still needs writing down.
- Bigger models.
02_nanogpt.pytakes--layers,--embed,--headsand--context. nanoGPT's Shakespeare model is 6, 384, 6 and 256. Every workstream above moves this one.
Good places to start
- Run the demos on your machine, and post the tables here (see below). Other CPUs and OSes are welcome.
- Peak memory per host, from a fresh process per measurement.
- The DuckDB join-order repro in workstream 5.
- A new model as a demo. See
benchmarks/ml/CONTRIBUTING.md. A new demo is also a new benchmark.
How to measure
uv run benchmarks/ml/01_attention.py --sizes 64 256 512 --reps 3 --csv attention.csv
uv run benchmarks/ml/02_nanogpt.py --steps 5 --reps 3 --csv nanogpt.csvWith a post, include:
- the CPU;
python -c "import datafusion, duckdb, ddxdb; print(datafusion.__version__, duckdb.__version__, ddxdb.__version__)";- the ddx commit, if you built ddx from a checkout.
To see whether a query's time goes to planning or running, time ctx.sql(q).execution_plan() (planning) against ctx.sql(q).to_arrow_table() (planning and running).
Activity
Run the ML demos in CI, and say how to contribute one
alxmrs commentedon Oct 25, 2026
🌀 Claude: The demos, re-recorded on ddxdb 0.4.2 (DataFusion 55.0.0, DuckDB 1.5.5, i7-8700). Every check against JAX passes.
GROUP BYThese match the numbers from ddx's development branch: 0.4.2 is correctness fixes, not speed. On DuckDB, a gradient with respect to a table that a view doesn't read now works. A view whose CTE computes a window and is read twice (a softmax) is still refused, filed as xqlsystems/ddx#204 with a minimal repro. With ddx#203's
prepare(not in 0.4.2), a GPT SGD step drops to 0.53 s (details).alxmrs commentedon Oct 25, 2026
🌀 Claude: Workstream 6 (einfold) now has an issue: xqlsystems/einfold#39. It's xarray-sql's facts (what xarray-sql guarantees about each table it registers, checked on #275) and a DataFusion host with einfold on in
_hosts.py, so every demo can report einfold on against off, as einfold's design measures it.alxmrs commentedon Oct 25, 2026
🪭 Claude (einfold): workstream 6, where einfold stands.
These are Rust-side numbers. None of them is from the demos yet. Intel i7-8700 (6 cores, 12 threads), DataFusion 55.2.0 (Rust), ddx
main(921aef1). No DuckDB and no ddxdb were involved.einfold's fold operator against DataFusion, on ddx's matrix-product gradients (xqlsystems/einfold#38, best of 5, 12 threads, from rows to rows):
FoldDataFusion 55 made the host 2.7× faster on Ā than 54.1 was (233 ms then), so the room shrank from 25×.
ddx's whole
matmulgradient at n = 50,000 takes 379 ms on a multi-threaded runtime. ddx's ownad_perfreports 868–904 ms because it runs on Tokio's single-threaded runtime. Per-step profile, and what einfold's first milestone needs from ddx: xqlsystems/ddx#208.For the demos: ddxdb runs ddx's steps on the Python
datafusionwheel, which can't decode an extension relation from Substrait. So einfold's speedups won't show inbenchmarks/mluntil there's a bridge. The options, and the plan, are on xqlsystems/einfold#39.Run the ML demos in CI, and say how to contribute one
alxmrs commentedon Oct 26, 2026
🌀 Claude: per-query overhead, workstream 1. Where a GPT training step's time goes now, and what's been done about it.
Measured 2026-10-25 on a quiet machine: Intel i7-8700 (6 cores, 12 threads), datafusion-python 55.0.0, ddxdb at 3dc0e62 (ddx#203's branch, with
ad.prepare). The 1-layer GPT isbenchmarks/ml/02_nanogpt.py's, at batch 4, trained with SGD.A prepared gradient step (about 430 ms)
ad.prepare's function runs 78 plans per call: the 67 steps of ddx's program, each stored as a table, and the 11 statements that read them. Three quarters of the call's time is DataFusion's, insideto_arrow_table. Decoding Substrait, registering tables and Python glue take about 60 ms together. For the 67 steps:use_row_number_estimates_to_optimize_partitioningon aContextmade with no config. A whole SGD step in the demo takes 0.415 s instead of 0.655 s, and attention at 512 tokens is unchanged.Around the step, in the demo's loop (about 74 ms per step)
CREATE VIEW modelagainto_dataset)from_dataset)The demo creates the view again on every step because a DataFusion view keeps reading the tables it was created over (#271). With an in-memory table (#275), writing the new weights and batch into the registered arrays makes re-registering and re-creating the view unnecessary. Five steps that way give the same losses and weights as the demo's way, to 1e-12. The demo can switch to it once #287 (
Context.prepare) lands.The forward loss (3.1–3.6 s)
The loss query still takes 3 s, nearly all of it planning, because DataFusion inlines a CTE at every reference (#267): 13,756 operators, or 9,882 with #289.
Context.prepare(#287) would plan it once, but it can't keep this plan. The plan encodes as 35 MB and is nested deeper than protobuf decodes ("recursion limit"), soContext.prepareplans it again on every call, as its docstring says it will. A planned loss runs in a few milliseconds, so the fix is #283's binding ofreset_plan_states, which resets a plan's state without decoding it. That makes #283 the largest single win for the GPT loop on DataFusion.Next
Split a query across partitions only where enough rows reach it (#289)
Run the ML demos in CI, and say how to contribute one
alxmrs commentedon Oct 26, 2026
🪭 Claude (einfold): first end-to-end number with einfold on, from Rust. Intel i7-8700, DataFusion 55.2.0, ddx
main921aef1, no DuckDB, no ddxdb.ddx's
matmulgradient at n = 50,000, multi-threaded runtime: 368 ms with einfold off, 224 ms on (1.6×), with the same gradients (xqlsystems/einfold#49). einfold now finds folds in DataFusion's plans by itself, with no change to ddx. Next for the demos: the FFI query-planner bridge into datafusion-python (xqlsystems/einfold#39), which uses the same detection code.alxmrs commentedon Oct 26, 2026
🌀 Claude: how fast the local stack is now, 2026-10-26. That's
mainwith #289, #287'sContext.prepare, #269's demos, and ddx#203'sad.prepare(ddxdb built from 3dc0e62), on datafusion-python 55.0.0 and an Intel i7-8700. Other agents were building on the machine (load average 4–7), so times are best-of and somewhat inflated. The configurations ran interleaved, and every result matches JAX.GPT, 1 layer, SGD, a whole training step
Per later step, in ms; best of 2 rounds of 6 steps:
sql_allContext.preparefor the loss,ad.preparefor the step, writes in placeLosses and weights agree between the two to 1e-12. Preparing removes the setup and a quarter of the gradient step. The loss is the wall: its plan is too deep to copy, so
Context.prepareplans it again from the optimized logical plan on every call, and physically planning 9,882 operators takes about 3 s. A reused physical plan runs it in a few milliseconds, but that needs #283'sreset_plan_statesbinding. JAX takes 0.15 ms for the loss and 0.37 ms for an SGD step.Attention, called with new Q, K and V each time
Best of 4 later calls, in ms; every value checked against JAX to 1e-14:
At 512 tokens, planning the forward pass takes about 4 ms and running it about 400 ms, so preparing doesn't help: the time is in running the joins. JAX takes 0.92 ms forward and 2.4 ms with gradients. That gap is einfold's M2, FlashAttention from SQL.
einfold M1
einfold's M1 gate (the
m1/gatebranch, measured by the einfold session at 03:36 today) has ddx's matmul gradient at n = 50,000 running 1.6× faster with einfold: 224 ms against 368, short of M1's 4× gate. The rest is ddx's join back around each contraction (xqlsystems/ddx#208). einfold is Rust-only, on DataFusion 55.2.0, so it can't run in these Python demos yet.What would move these numbers
alxmrs commentedon Oct 26, 2026
🌀 Claude: Reset plans make a GPT training step 4.1× faster.
Context.preparenow runs a reset copy of its plan on each call, using a binding of DataFusion'sreset_plan_statesadded on the xqlsystems/datafusion-python fork (#283).One SGD step of
02_nanogpt.py: the median of later steps, best of 3 interleaved rounds, in ms.prepare, copying plans (#287)prepare, resetting plans, no dynamic filter pushdownprepareresets only when the session turns all fourenable_*dynamic_filter_pushdownsettings off, and otherwise copies the plan as before. Turning pushdown off made no measurable difference to today's loop.Setup:
ExecutionPlan.reset_state(the DataFusion 55.2.0 crates), release build.3dc0e62(ddx#203'sad.preparefor the gradient step).alxmrs commentedon Oct 26, 2026
🌀 Claude: Running each shared CTE once makes a GPT training step 10× faster, with no fork needed.
The loss plan was 9,882 operators, but only 87 distinct subplans. DataFusion inlines a CTE or view at every place it's read, and the GPT's layers each read the layer before several times.
Context.preparenow splits such a query (branchshared-ctes, on top of #287): each CTE or view read more than once runs once per call, and its rows go into a table, written in place, that the rest of the query reads. The rewrite usessqlparserin xarray-sql's native module.One SGD step of
02_nanogpt.py: the median of later steps, best of 3 interleaved rounds, in ms.prepare, sharing CTEsreset_stateprepare, sharing CTEs and resetting plans, no dynamic filter pushdownSetup: Intel i7-8700 (12 threads), load average about 1. ddx at
3dc0e62(ddx#203). xarray-sql is #287 plus #269 plusshared-ctes, on a local branch.alxmrs commentedon Oct 26, 2026
🌀 Claude: The GPT training step is now about 30× faster than the demo loop. Each of ddx's gradient steps keeps its physical plan and runs it again.
The setup: an Intel i7-8700, DataFusion 55.2.0 (datafusion-python 55.0.0 from the xqlsystems fork, with
ExecutionPlan.reset_state), ddx-local 3dc0e62, pyarrow 23.0.0 and Python 3.12. Batch 4, dynamic filter pushdown off. Timings are the median of later steps of 6, best of 3 rounds.preparewith shared CTEs (#290) +ad.prepare, 12 partitionsDataFusionBackendkeeps each step's plan, and writes its rows in place. Details and code are on perf(ddx-ad, ddx-datafusion, ddxdb): run each prepared step's physical plan again ddx#214.sql_all's to within 4e-16.alxmrs commentedon Oct 26, 2026
When we compare against JAX, are we including its jit?
alxmrs commentedon Oct 26, 2026
🌀 Claude: The JAX number we compare against leaves out compiling, and so does ours.
In both, the setup and the first call are left out, and later steps are timed. I measured JAX again, as the SQL loop is timed: batch and weights passed as arguments, a new batch each step, the median of 51 later steps. Intel i7-8700, JAX 0.12, CPU only, load about 5.
jax.jit(tracing and compiling)Context.preparefor the loss,ad.preparewith plans kept, 1 partitionalxmrs commentedon Oct 26, 2026
🌀 Claude: Where the time goes in the SQL step's first call, our "compile": about 1.04 s in all, against JAX's 370 ms to trace and compile. One partition, i7-8700, load about 1.
view_sql: checking that themodelview's SQL plans as the view does, which optimizes the deep model viewplan_sql,select_all,returns_rows(My timers planned each step twice, so I subtracted one planning, about 200 ms, from their total of 1.24 s.)
The view check is about two-thirds of the first call. Without it, our first call would be about 360 ms, which is about JAX's compile.
alxmrs commentedon Oct 26, 2026
🪭 Claude (einfold): einfold's first milestone (M1) passes its gate. einfold is an XQL Systems project that runs tensor contractions written in SQL with fused kernels inside the engine. The gate is ddx's
matmulgradient program, end to end, at n = 50,000; ddx is XQL Systems' automatic differentiation of SQL queries. With einfold installed, it's 4.0–4.1× faster than the same program on the same DataFusion, over five runs: 369 ms against 90, with the same gradients.mainat 921aef1. DuckDB and ddxdb aren't involved: this runs in Rust, with einfold as DataFusion's query planner (einfold_datafusion::install).A Python path for xarray-sql's demos (an FFI query planner for datafusion-python) is next, starting as a spike.