Repository navigation
program: choose step boundaries by what they save (fewer, larger steps) #186
Description
🌀 Claude (xarray-sql ML demos)
Summary
A program's steps are where ddx's work is now. With ddxdb 0.3.3, a cached training step of demo 2's GPT (xarray-sql's ML benchmarks, xqlsystems/xarray-sql#269) takes 1.0 s on DataFusion, and 0.70 s of that is running the program's 68 steps (1,335 operators in all, at most 16,384 rows in any step). This issue tracks choosing better step boundaries: fewer steps where a boundary buys nothing, larger steps that share their inputs, and occasionally one more step where many steps recompute the same relation.
The measurements below change what "fewer, larger steps" should mean. A step's fixed cost is small. What costs is work repeated across steps.
What a step costs
I timed DataFusionBackend.materialize (consume the Substrait plan, run it, register the result) on DataFusion 54.0.0, on an idle Intel i7-8700:
| step | per step |
|---|---|
SELECT i, x FROM t (1,024 rows) |
0.2 ms |
a GROUP BY over it |
0.6–1.0 ms |
a join and a GROUP BY |
1.4–2.2 ms |
| the GPT's steps, on average | 10 ms |
Merging steps only to have fewer of them would save little: 68 steps at 0.2–2 ms is 15–135 ms. Over the whole GPT program, consuming the Substrait plans takes 39 ms and registering the results 10 ms. The rest is planning and running each step's operators, so the lever is the operators.
Where the GPT's steps repeat work
-
27 of the 68 steps rebuild the token embedding,
batch ⋈ wte ⋈ wpe. Every step that reaches the residual stream reads it. It's a join with no aggregate, so ddx doesn't save it, and each step rebuilds it from the base tables:batchis read by 30 steps,wteby 28,wpeby 27. On DataFusion, each of those reads is also a slow scan when the table is an xarray-sql Dataset (Each scan of a registered Dataset rebuilds its batches in Python: 73% of the GPT demo's forward pass xarray-sql#272). -
Q, K and V are three copies of one computation. These step triples have the same shape and read the same inputs (
saved_0,saved_1, the embedding andln_1):- the forward projections
saved_2,saved_3andsaved_5; - the contractions
contraction_12,_14and_16; - the regions
region_13,_15and_17; - the gradients
grad_3_wq,grad_4_wkandgrad_5_wv.
Each triple reads its inputs three times. nanoGPT computes the same thing as one product with a stacked weight (
c_attn). - the forward projections
-
13 steps are read by exactly one later step:
contraction_1,_5,_12,_14and_16;window_0,_10and_11;cotangent_1,_9,_10,_13and_14.
Storing them buys nothing: their reader could compute them inline, with no recomputation.
Proposal: choose step boundaries by what they save
- Inline a step that one step reads. This is the inverse of
opt::splitand has to respect itsMAX_DEPTH. It is always safe and saves the step's fixed cost and its write: small, about 13 × 1–2 ms on the GPT, but free. - Stack independent steps that share their inputs. Steps of the same shape over the same inputs (Q, K and V above, or the per-table gradients of a weight read by every layer) become one step with a tag column. Each reader filters on its tag, and the gradient outputs become filters or views of the stacked table. That reads the shared inputs once instead of three times, which is what a fused QKV projection does. This needs a rule for recognizing "same shape, same inputs" in the plans ddx emits; the contraction rule's output is regular enough that it may only need to compare plans with the weight table abstracted.
- Store a relation that many steps rebuild, even if it isn't an aggregate. The embedding is one more step, and 26 fewer rebuilds. An embedding join costs 1.8 ms over in-memory tables and 5.6 ms through xarray-sql's scan today, so this is roughly 50–150 ms of the GPT's 0.70 s. In general, step boundaries are a materialization choice: store a relation that more than one step reads and that costs more to recompute than to write; inline one read once.
- On an engine that computes a CTE once, use CTEs. DuckDB computes a CTE that is read twice only once, and PostgreSQL has
AS MATERIALIZED. There, a chain of steps, or a whole program, can be one statement, with each would-be step a CTE, and nothing goes to the client between steps. This bears on the DuckDB backend (ddxdb: one backend for SQL engines, DuckDB-Python first #148, Exit: xarray-sql's ML demos differentiate on DuckDB #150), which is starting now: itsmaterializeshould beCREATE TEMP TABLE … AS, keeping the data in the engine, and its target profile could say that CTEs are computed once. DataFusion recomputes a CTE at every read (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth xarray-sql#267), so it needs steps, as Plan queries for engines with no Substrait producer (Polars, ADBC) #159 notes.
Items 1–3 change the program and help every engine. Item 4 is a per-engine way of running it.
A cheaper knob, separately
The steps are small, and DataFusion plans each one for every core. Setting datafusion.execution.target_partitions = 1 takes the GPT's 68 steps from 0.70 s to 0.57 s (−19%). With 2 or 4 partitions it's 0.63–0.64 s. A backend could run steps below some row count with one partition, using DataFusion's statistics or the program's own knowledge of each step's size. This is independent of the boundaries above and stacks with #157 (running independent steps concurrently): fewer partitions per step leaves cores free for other steps.
How to measure
The profile example (#158) gives each step's time. The gate for this issue is the GPT's cached step, comparing the step count, the total operators and the base-table reads before and after. Attention (demo 1) is the other check. In #153's profile, made before #162 changed its program, window_0 rebuilt the (i, j, d) join that region_1 stores, which would be case 3 again. I haven't checked whether 0.3.3's program still does.
Part of #158.
Activity
alxmrs commentedon Oct 21, 2026
A related lever for item 4, from the SQL writer in #189 (
ddx_ad::to_sql, which writes each step as SQL for engines whose Substrait consumer can't run it, DuckDB first):to_sqlwrites each relation as a CTE, and a relation the plan already shares (aReferenceRel) as one CTE read wherever it is shared. DuckDB computes a CTE that is read twice only once. So ifto_sqlalso wrote structurally identical subtrees as one CTE, DuckDB would compute them once without ddx adding a step. This is hash-consing: key each subtree by its structure and write the first copy only. In the GPT's steps that covers thebatch ⋈ wte ⋈ wpeembedding join, which 27 steps rebuild. Within one step it would also cover the inputs that the Q, K and V branches read in common, though not the branches themselves, which differ in their weight table.It doesn't help DataFusion, which inlines a CTE at every read (xqlsystems/xarray-sql#267). So it belongs with this issue's per-engine item 4, not items 1–3. It needs no new step boundary, so it's cheap to try once the DuckDB backend (#148) can be measured.
alxmrs commentedon Oct 25, 2026
A profile of a cached GPT SGD step with ddx at 3c95aa7 (after ddxdb 0.4.1), DataFusion 55.0.0, Intel i7-8700. The step takes 0.55 s after the first, almost all inside
rewrite; collecting the 11 statements' results takes 13 ms.materializecalls take 0.41 s, 5.4 ms each, for steps whose median plan has 16 operators and whose largest result is 16,384 rows (74,627 rows in all).optimized_logical_plan0.10 s,execution_plan0.20 s, and collecting (which plans again, then runs) 0.34 s. About three quarters of a step is planning; running it is about 1–1.5 ms.view_sql, 13 calls, 0.07 s, with_same_plans0.05 s of it),select_all(0.015 s),plan_sqlandplan_key.Session settings move it a little:
target_partitions = 1target_partitions = 4optimizer.max_passes = 1For comparison, JAX (
jax.jit, same CPU) takes 0.37 ms for the whole SGD step. Since planning dominates, keeping each step's physical plan between runs looks like the biggest win, ahead of choosing step boundaries. That needs the step's inputs to be tables that can be swapped under a live plan, plusExecutionPlan::reset_statebetween runs (aHashJoinExeckeeps its build side in the plan). A smaller, independent win: run small steps with one partition.alxmrs commentedon Oct 29, 2026
🌀 Claude: A re-measure for #186, and a sketch. It's on ddx
mainplus #229 (ddxdb 0.4.5), DuckDB 1.5.6, DataFusion 55.0.0, Intel i7-8700. It uses xarray-sql's GPT step (demo 2) at two sizes: small (16 channels, context 16, batch 4) and base (1 layer, 4 heads, 128 channels, context 64, batch 12).Where the time goes now
On DuckDB, at base size, it's execution, mostly the matrix products. A cached
sql_allstep takes 13.7 s, of which 13.4 s is its 77 steps'CREATE TABLE AS.cotangent_111.7 s,grad_9_w_proj1.4 s,saved_111.2 s, each joining 768 × 512 × 128 rows. Next are Q, K and V's projections, gradients and contractions, about 0.3 s each.CREATE_TABLE_AS)On DataFusion, at small size, it's per-step overhead. A prepared step (
ddxdb.ad.prepare, without the fork'sin_place) takes about 300 ms for 77 steps, about 4 ms a step, mostly planning each step's physical plan and writing it.What each of #186's items can buy now
Inline a step that one later step reads, once. 13 of the 77 steps qualify:
contraction_0/4/11/13/15,cotangent_1/9/10/13/14/16andwindow_9/10. Two kinds of step are materialized on purpose and must not be inlined:tthat its reader reads more than once;So the rule is: inline a step only when its one reader reads it exactly once, and the result stays under
opt::split'sMAX_DEPTH. Expected: about a sixth of the steps on DataFusion, roughly 10–15% of a small prepared step. On DuckDB it saves only their writes, about 1–2%.Stack Q, K and V. What stacking shares is each triple's input rebuild. All rebuild joins together are under 4% of DuckDB's time, so stacking can save at most a few percent, at the cost of a new rule. Defer.
Store the embedding that 27 steps rebuild. Also within that 4% on DuckDB, and it needs a new kind of saved relation (one with no aggregate, whose cotangent accumulates by its inputs' dims). Defer until a model makes rebuilds matter.
Run a whole program as one DuckDB statement with CTEs. At most the writes' 4% plus statement overhead. Defer.
The products themselves (70%) are einfold's to speed up, not step boundaries'.
The plan
Build item 1 as
opt::inline, run on the final steps beforeopt::split:MAX_DEPTHor its size passesinlined_size's budget (perf(ddx-ad, ddxdb): run a step whose names would grow large as SQL #221).I'll measure step counts and the GPT's small (DataFusion) and base (DuckDB) steps before and after, and keep it only if it's faster.
perf(ddx-ad): compute a step that one later step reads once inside it (…
alxmrs commentedon Oct 31, 2026
Closing as done.
opt::inline), released in ddxdb 0.4.6.If a larger model makes the embedding rebuilds matter, item 3 gets its own issue.