Re-encoding the support set dominates `tabfm_classify`: 71-80% of a call, paid again on every call

[maxdemarzi]

Measured while costing a two-model routing system built on this extension. Posting the numbers
because they make a fairly strong case for a reusable support-set context, and because the shape of
the cost is not obvious from the API.

What was measured

One idle 16-core AMD EPYC 4564P, 124 GB, nothing else running. tabicl-v2, CPU EP, e=1,
--test-chunk 128, 500 features per call. Two UCR datasets, each run at two query-batch sizes in
the same process invocation, so the two points differ only in the number of query rows:

dataset	support rows	query rows	seconds per call
Herring	64	14	1.525
Herring	64	64	1.978
ScreenType	375	22	5.265
ScreenType	375	128	6.289

Solving seconds = a + b * query_rows on each pair:

dataset	fixed per call	marginal per query row	fixed share of a 128-row call
Herring	1.398 s	9.0 ms	71%
ScreenType	5.053 s	9.7 ms	80%

The fixed term scales with the support set: 1.40 s at 64 support rows and 5.05 s at 375, about
11.8 ms per support row across the two. The marginal term is stable at ~9 ms per query row and
agrees with a third estimate — a regression of per-call seconds on (support rows, query rows) across
51 archived runs gives 9.3 ms, independently of the two fits above.

These are read from the per-call timings rather than wall clock, so process startup (0.5 s of a
63 s call) and feature computation are excluded and not attributed to the model.

Why this is worth a feature

The support set does not change between calls in any deployment I can think of. It is the training
data. But tabfm_classify(train, y, test := ...) re-encodes it on every call, so:

• Chunking is punished twice. --test-chunk is the documented way to bound memory, and each
  chunk re-sends the whole support set. Measured on ItalyPowerDemand: chunk-128 took 1074 s against
  493 s unchunked, 2.18x, for byte-identical predictions (1029/1029 rows). The fit above,
  which never saw that run, predicts 2.04x for it from the support/query split alone — so most of
  that factor is the repeated support pass.
• Serving a small batch is nearly as expensive as serving everything. Routing — run a cheap
  model everywhere, escalate only the rows it is unsure of — should be the natural way to use an
  expensive in-context model. Measured, escalating 22% of a Herring batch cost 77% of running
  the teacher on every row, and 17% of a ScreenType batch cost 83%. The row count falls and the
  cost does not, because the fixed pass is most of it.
• It gets worse as the training set grows, which is the opposite of the usual expectation, and
  it is the fixed term that grows.

An ensemble over feature groups multiplies all of this: 40 groups against one dataset is 40 support
passes over the same rows in different feature subsets, plus another factor per chunk.

Why it cannot be done downstream

The tensor contract is a single forward pass over support and query concatenated:

inputs:  x [1, T, H]    support and query rows in one tensor
         y [1, T]
         train_size     the index that splits support from query
outputs: logits [1, T, C]

Nothing intermediate is exposed, so an extension-level or caller-level cache has nothing to hold.
Reusing the encoded support would need the exported graph split at the support/query boundary, with
the encoded support as an output of the first stage and an input of the second — which is export
work in tools/export_*, per architecture, and not something to send as an outside patch. Hence an
issue rather than a PR.

What a user-facing shape might look like

Not a design proposal, just the smallest thing that would capture the win:

-- encode once
CREATE TABLE ctx AS SELECT tabfm_prepare('train_tbl', 'y', model := 'tabicl-v2', features := [...]);

-- reuse across calls; support rows never re-encoded
FROM tabfm_classify(context := ctx, test := 'batch_1');
FROM tabfm_classify(context := ctx, test := 'batch_2');

Invalidation is easy to state: the handle is keyed by (model, feature list, support contents), and
any change to those is a new handle.

Where the numbers come from

All reproducible from https://github.com/maxdemarzi/duckdb-rocket — scripts/route_serve.py serve --compare produces the two-batch-size fit and prints the decomposition; reference/RESULTS.md
section "What routing actually costs" has the full working, and the chunking measurement is under
"Where the wall clock goes".

Happy to test a branch against the 28-dataset harness here if that is useful — it covers CPU and
CUDA and would give you an accuracy-identity check as well as timings.

[maxdemarzi]

Two follow-ups that make this look considerably more tractable than the issue text assumed, at
least for tabicl-v2.

1. TabICL already implements the cache; the export routes around it

From tools/export_tabicl/src/export_tabicl/tabicl_patches.py:

TabICL's own forward dispatches on self.training — the training-mode code path is the clean,
single-forward one (no KV cache / InferenceManager / feature-shuffle Python loops), so the export
wrapper puts the model in train() and calls the three stage modules directly.

tabicl._model.kv_cache is a real module in the installed package. So the machinery exists upstream
and the ONNX export deliberately takes the path without it — for good tracing reasons, but it means
the cost is an artifact of the export rather than of the architecture.

2. Query rows are independent given the support — measured

The property a cached support set needs is that a query row's prediction does not depend on which
other query rows share its call. For TabICL that follows from embed_with_test=False (the export
wrapper's own setting: the column embedding is computed from training rows only) plus a row-wise
interaction stage, but it is worth measuring rather than reading off the source.

Trained tabicl-classifier-v2 weights via the sklearn API, 64 support rows, 12 query rows,
8 features:

controls  predictions vary across query rows : 9.813e-01
          a different support set changes them: 7.135e-01
          flipped support labels change them  : 9.988e-01
claim     12 queries in one call vs split 7+5 : 1.416e-07
          query #3 in a batch of 12 vs alone  : 2.561e-09

The controls are load-bearing. A first attempt used a randomly initialised TabICL and every
difference came out at exactly 0.0 including the controls, because a random TabICL emits a constant
vector regardless of input — a degenerate model passes an independence test trivially.

So on this architecture the support half genuinely is reusable: splitting the export at the
support/query boundary is a refactor of a computation that is already separable, not a change to
what the model computes. It should also mean the existing --test-chunk is exact by construction,
which matches the 1029/1029 identical predictions reported for it.

Script: https://github.com/maxdemarzi/duckdb-rocket/blob/main/scripts/query_independence.py —
happy to send it as a test in a PR here instead if that is more useful.

Caveat, stated because it limits the above: this is TabICL only. TabPFN's exported contract has
no train_size input at all and infers the split from len(y), and Mitra's note describes
in-graph rank normalization over the feature axis; neither has been checked for the same property,
and a per-architecture answer is likely.

[maxdemarzi]

One implementation detail I got far enough to check, since it decides where the split can go.

The cacheable state is not just the ICL stage's keys and values. Query-row representations
already depend on the support set before the ICL stage: running the export wrapper's own
col_embedder(..., embed_with_test=False) then row_interactor over the same 12 query rows against
two different support sets gives

QUERY-row reps, two different support sets : max|diff| = 5.65
SUPPORT-row reps, same comparison (sanity) : max|diff| = 7.09

on trained v2 weights. embed_with_test=False stops training rows attending to test rows; it does
not make the test rows' embedding support-independent, which is expected for a target-aware column
embedding.

So a split has to carry whatever ColEmbedding derives from the support (inducing-point state /
column statistics) across the boundary as well as the ICL keys and values. That is a restructuring
of ColEmbedding into a support-derive half and an apply half, which is more than the monkeypatch
surface tabicl_patches.py currently keeps to — and is why I have stopped at reporting rather than
sending a patch.

None of this weakens the request. Query independence was measured end to end on the public API, so
whatever that state turns out to be, computing it once per support set and reusing it is exact.
It just means the handle holds more than a KV cache.

[maxdemarzi]

Went ahead and prototyped it in PyTorch. It works, predictions are identical, and it is worth
1.8-5.9x depending on shape.

Prototype (monkeypatch on the installed tabicl, no edits, in the spirit of tabicl_patches.py):
https://github.com/maxdemarzi/duckdb-rocket/blob/main/scripts/context_cache_proto.py

Where the split goes

InducedSelfAttentionBlock.induced_attention, unmodified:

hidden = self.multihead_attn1(ind_vectors, src[..., :train_size, :], src[..., :train_size, :])
out    = self.multihead_attn2(src, hidden, hidden)

hidden depends only on the support rows; every row's output is an independent attention against
it. It stacks — block k's hidden derives from block k-1's support outputs — so the cacheable
state is one hidden per column-ISAB block, plus the support representations the ICL stage reads.

The prepare pass records those; the query pass replays them and runs only the query rows through
multihead_attn2 and the row interactor. That is where the win is: at S=500/Q=128 those two stages
are 90% of the forward (col_embedder 48%, row_interactor 42%, icl_predictor 10%) and both were
reprocessing 500 support rows on every call.

Measured, trained tabicl-classifier-v2 weights, CPU

support	query	per-batch	amortised over 10 batches	max abs logit diff	argmax agreement
500	128	2.91x	2.42x	7.4e-05	100%
500	32	5.86x	3.91x	7.3e-05	100%
375	128	2.23x	1.92x	8.7e-05	100%
100	128	1.79x	1.64x	1.8e-05	100%
500	500	2.10x	1.90x	7.7e-05	100%

Predictions identical throughout. The logits differ by 4.5e-06 relative on a scale of 12.6, which is
fp32 reduction-order noise over ~600 accumulated terms rather than a different computation. So
numerically equivalent, not bit-exact — float64 would settle it but the model hardcodes .float()
in layers.py.

Note the shape of the win: largest at a small query batch against a large support set, which is
the escalation case from the original post that currently costs 77-83% of running the model on
everything.

One trap, because it cost me an hour and would cost the same again

ISAB.forward skips a column group whose values all equal skip_value, via
out[~skip_mask] = self.induced_attention(src[~skip_mask], train_size). That boolean index flattens
the batch dims, so a skipped group changes the shape the cached hidden has to match. The
reserved cls-token columns are padded with exactly skip_value and survive in the combined pass only
because the target-aware step adds y_emb to the support rows and perturbs them. A query-only batch
adds no y_emb, so those 4 columns look skippable and 104 groups silently become 100 — against a
cache recorded at 104.

The split has to inherit the combined pass's skip decision rather than re-derive it. Any real
implementation should carry that mask across the boundary with the rest of the state.

What I have not done

The ONNX side. This is a PyTorch-level proof that the decomposition is exact and profitable; turning
it into two exported graphs plus a context handle in the extension is your call on API and is where
the remaining risk sits (the dynamo export of this model is, per your own README, the hard case).

Happy to take a run at the two-graph export if you would want that as a PR, or to leave it here as
evidence — whichever is more use. Either way the 28-dataset harness on our side can check accuracy
identity and timings against the current single-graph path on CPU and CUDA.

[maxdemarzi]

Thanks for merging #38 — and for #36, which turns out to be load-bearing for the case below.

Served the routing path end to end this morning to check a config change on our side, and it
produced a sharper version of this issue's measurement than anything in the thread so far. Same
protocol as before: all arms on one idle 16-vCPU box in one run, per-group timings rather than wall
clock.

The fixed pass reaching 100% of a call

SemgHandMovementCh2 — 450 support rows, 6 classes, 500 features per call — fitted over 40 groups:

6.111 s fixed per group per call + 0.1 ms per query row

The context pass is 100% of the full-batch cost to the printed precision. Query rows are free;
the support set is the entire bill. For contrast, on the same box in the same hour:

dataset	support rows	fixed per group per call	per query row
Herring	64	1.411 s	9.0 ms
ScreenType	375	5.033 s	9.5 ms
SemgHandMovementCh2	450	6.111 s	0.1 ms

Those three fixed terms are an affine function of the support row count to within 2%:

0.62 s + 12.0 ms per support row,  per group, per call

Note it does not track series length (512 / 720 / 1500 timepoints) — the caller has already
reduced each series to 500 features, so what the model re-encodes is rows, and that is exactly the
quantity a cached support set removes.

For calibration on the harness: the G=40 arms reproduced our 2026-08-15 measurements on different
rented hardware within 1% (Herring 82.4 → 82.8 s, ScreenType 267.9 → 266.1 s, fits 1.398+9.0 →
1.411+9.0 and 5.053+9.7 → 5.033+9.5).

A second argument I had not appreciated: chunking and the memory ceiling

SemgHandMovementCh2 cannot run a 128-row batch at all here. It is OOM-killed by the kernel
(exit -9, after scoring 1 group of 40) against a 32 GB container. --test-chunk 32 is what makes
it complete — and that is the part relevant to this issue:

--test-chunk is the only lever that bounds a call's peak, and today it is priced per chunk.
Dropping to 32 turns one call per group into four, each re-encoding all 450 support rows. That is
why the full-batch arm costs 1013.5 s against the escalated arm's 273.4 s — a ratio of 3.7, which
tracks the call count (4:1), not the row count (128:24 = 5.3).

So the two features interact in a way worth naming: with #36 a user can now discover their memory
ceiling cleanly instead of being SIGKILLed, and --test-chunk is how they respond to it — but the
response currently costs a full support re-encode per chunk. With a context handle, chunk size
becomes an almost free memory knob rather than a costly one. That seems like the strongest practical
case for the extension-side half: it is not only a speedup, it removes the penalty on the only
mechanism that makes large-support datasets runnable.

(Incidentally, this is also the one dataset in our set where routing looks good on cost — 25-29%
of running the model on everything — and the reason is not that it escalated fewer rows. It is that
the full batch needed four calls where the escalated batch needed one. Routing saves calls, which
is the same observation from the other direction.)

Status on our side

Nothing needed from you. We are pinned to the community build, so --split-context and
anofox_tabfm_max_memory are not reachable from our pipeline until the next release cuts — happy to
re-run the 28-dataset harness against both once they are published, on CPU and CUDA, and report
accuracy identity plus timings against the current single-graph path.

Full measurements and the eight run records: https://github.com/maxdemarzi/duckdb-rocket/blob/main/reference/serve_compare.json