Summary
The JAX model already implements context caching: TabFM.prefill() and
TabFM.decode() in tabfm/src/jax/model.py, with ICLearningCache and the
cache_icl_input_only option, covered by
model_test.py::test_prefill_decode_consistency. The public TabFMClassifier
and TabFMRegressor never call them. Every predict_proba call re-encodes
the full training context, once per ensemble member.
What this costs
Measured on an RTX 5090 (PyTorch backend, bf16, n_estimators=4, 40k-row
context, creditcard dataset from OpenML), steady state:
| queries per predict call |
wall time |
| 1 |
25.7 s |
| 100 |
26.1 s |
| 1,000 |
27.1 s |
| 30,000 |
37.5 s |
Predict cost is a fixed cost of about 25 s to re-encode the context, plus
about 1.4 ms per query row. A single prediction costs the same as a thousand.
prefill/decode exists to move that fixed cost to fit time.
Proposal
Add an option to the sklearn wrappers, using TabPFN's naming so users can
find it:
clf = TabFMClassifier(model=model, fit_mode="fit_with_cache")
clf.fit(X_train, y_train) # runs prefill once per ensemble member view
clf.predict_proba(X_test) # runs decode against the cached context
The default fit_mode="fit" keeps current behavior.
We have a working branch with tests: 60 passing, parity against the default
path at about 1e-7 in fp32 including the ensemble() preset, and a mean
probability difference of about 1e-4 in bf16 on the released checkpoint.
Wrapper-level measurements on the JAX backend at a 10k-row context show 1.3x
to 2.4x faster predict calls after a one-time 28.9 s prefill. The gain grows
with context length: the default path's fixed cost grows from about 2.5 s at
10k rows to about 35 s at 40k rows per member, while the cached path does not
depend on context length. We are happy to send the PR if this direction is
welcome.
Two findings from the implementation
- Wrapping
decode in nnx.jit does not work at real scale. The jit copies
the multi-GB caches into the executable arena and fuses a transpose across
all 24 layers that runs out of memory at a 20k context and fails XLA
autotuning at 40k. Our branch calls decode eagerly, which matches how
model_test.py uses it.
decode builds an internal per-head key/value tensor that spans all 24
layers at once (shape [heads, layers, B, T, head_dim], about 1 GB per 10k
context rows per member). On a 32 GB GPU this caps cached prediction at
roughly a 10k context with one member. Chunking that tensor per layer would
unlock the larger contexts where caching helps most. We can file this
separately.
A PyTorch-backend equivalent would be a natural follow-up. prefill/decode
currently exist only in the JAX model.
Summary
The JAX model already implements context caching:
TabFM.prefill()andTabFM.decode()intabfm/src/jax/model.py, withICLearningCacheand thecache_icl_input_onlyoption, covered bymodel_test.py::test_prefill_decode_consistency. The publicTabFMClassifierand
TabFMRegressornever call them. Everypredict_probacall re-encodesthe full training context, once per ensemble member.
What this costs
Measured on an RTX 5090 (PyTorch backend, bf16,
n_estimators=4, 40k-rowcontext, creditcard dataset from OpenML), steady state:
Predict cost is a fixed cost of about 25 s to re-encode the context, plus
about 1.4 ms per query row. A single prediction costs the same as a thousand.
prefill/decodeexists to move that fixed cost to fit time.Proposal
Add an option to the sklearn wrappers, using TabPFN's naming so users can
find it:
The default
fit_mode="fit"keeps current behavior.We have a working branch with tests: 60 passing, parity against the default
path at about 1e-7 in fp32 including the
ensemble()preset, and a meanprobability difference of about 1e-4 in bf16 on the released checkpoint.
Wrapper-level measurements on the JAX backend at a 10k-row context show 1.3x
to 2.4x faster predict calls after a one-time 28.9 s prefill. The gain grows
with context length: the default path's fixed cost grows from about 2.5 s at
10k rows to about 35 s at 40k rows per member, while the cached path does not
depend on context length. We are happy to send the PR if this direction is
welcome.
Two findings from the implementation
decodeinnnx.jitdoes not work at real scale. The jit copiesthe multi-GB caches into the executable arena and fuses a transpose across
all 24 layers that runs out of memory at a 20k context and fails XLA
autotuning at 40k. Our branch calls
decodeeagerly, which matches howmodel_test.pyuses it.decodebuilds an internal per-head key/value tensor that spans all 24layers at once (shape
[heads, layers, B, T, head_dim], about 1 GB per 10kcontext rows per member). On a 32 GB GPU this caps cached prediction at
roughly a 10k context with one member. Chunking that tensor per layer would
unlock the larger contexts where caching helps most. We can file this
separately.
A PyTorch-backend equivalent would be a natural follow-up.
prefill/decodecurrently exist only in the JAX model.