Add ICL context caching for the PyTorch backend + expose in sklearn API

[astonishedrobo]

Scope

This PR adds in-context caching to the PyTorch backend, opt-in via
cache_context=True on TabFMClassifier / TabFMRegressor (default: False).
When it's on, two options control the cache's memory footprint:

• maybe_quantize_kv_cache (default: True): int8-quantize the ICL KV cache to
  reduce memory; set False to keep the cache full precision.
• keep_cache_on_device (default: True): caches stay on GPU by default; with
  False, they are stored on CPU between prediction calls.

This PR implements the JAX backend's prefill/decode cache path for PyTorch and
exposes it through cache_context=True. The unquantized cached path matches
uncached predictions within fp32 tolerance; the default int8 KV cache reduces
memory, with numerical impact measured in the benchmark below.

I also ran a benchmark on Adult (OpenML 1590) for throughput and accuracy; results
are in the Benchmark section.

Usage

The following block shows the usage example.

## Classifier
...
clf = TabFMClassifier(
    model=model,
    cache_context=True,
    maybe_quantize_kv_cache=True,   # int8-quantize the cached KV
    keep_cache_on_device=True,      # False stores cache on CPU between predict calls
)
clf.fit(X_train, y_train)

predictions = clf.predict(X_test)
probabilities = clf.predict_proba(X_test)

## Regressor
...
reg = TabFMRegressor(
    model=model,
    cache_context=True,
    maybe_quantize_kv_cache=True,
    keep_cache_on_device=True,
)
reg.fit(X_train, y_train)
predictions = reg.predict(X_test)

Benchmark

I benchmarked context caching on Adult (OpenML 1590), a binary classification
task: predict whether annual income is > 50K or <= 50K. I used the following
setup:

• Dataset split: 10,000 (stratified) for context, 38,842 for test.
• Inference batch: 2,000 test samples per prediction call (unrelated to the
  batch_size parameter, which sets the number of ensemble members per forward).
• Model dtype: bf16 for the throughput and accuracy runs. Cached and uncached
  run at the same precision, so any difference is the cache alone.
• Cache config: keep_cache_on_device=True, so all members' caches are kept on
  the GPU (not offloaded to CPU).
• Unquantized fp32 check: a separate run at fp32 precision (2,000 context,
  2,000 test, n_estimators=2), where rounding is negligible, to check the cached
  and uncached paths produce the same result.
• Hardware: NVIDIA L4 (24 GB VRAM).
• Metrics measured:
  • throughput (rows/s)
  • ROC-AUC
  • max|p_cached - p_uncached|: the largest change between the cached and
    uncached runs in the predicted probability of the > 50K class.
  • peak GPU memory

The following table summarizes the bf16 runs.

n_estimators	path	rows/s	speedup	ROC-AUC	max|p_cached - p_uncached|	peak GPU
4	uncached	118	1.0x	0.9289	N/A	4.5 GB
4	cached + int8	912	7.8x	0.9289	1.9e-2	11.4 GB
8	uncached	59	1.0x	0.9289	N/A	4.5 GB
8	cached + int8	457	7.8x	0.9289	1.3e-2	15.4 GB

Observations:

• Throughput: cached is ~7.8x over uncached at both n_estimators (4 and 8).
• Accuracy and numerical differences:
  • ROC-AUC is unchanged across all runs.
  • default int8 cache: max|p_cached - p_uncached| is ~1.9e-2 in the bf16
    model run, with no observed effect on ROC-AUC.
  • fp32 unquantized cache: max|p_cached - p_uncached| = 3.7e-6, within the
    1e-5 tolerance used for comparing cached vs uncached predictions.
• Memory: cached peak grows with n_estimators (11.4 GB at 4, 15.4 GB at 8) because
  all members' caches are kept on the GPU at once; uncached stays at 4.5 GB.

Changelog

• tabfm/src/pytorch/model.py
  • Added prefill() / decode() to TabFM: encode the context once and reuse
    it during prediction.
  • Added optional cache arguments to the attention modules (MultiheadAttention,
    MultiheadAttentionBlock, Encoder, InducedSelfAttentionBlock,
    SetTransformer, ColEmbedding, ICLearning) so decode() reuses the
    cached KV instead of recomputing it.
  • Added the int8 KV-cache path: QuantizedTensor, _quantize_tensor,
    ICLearningCache.quantize(), move_cache_to_device. Only the ICL KV is
    quantized; column representations and training row counts stay full
    precision.
• tabfm/src/classifier_and_regressor.py
  • Added cache_context, maybe_quantize_kv_cache, keep_cache_on_device to
    both estimators. fit() builds the per-member cache; prediction uses
    decode() with the cached context. JAX raises NotImplementedError.
  • Added train-only / test-only tensor builders for the decode path:
    transform_context_only, transform_test_only, prepare_test_tensors.
• tabfm/src/pytorch/__init__.py
  • Added the package init.
• tabfm/src/pytorch/model_test.py
  • Added coverage that the cached path (prefill + decode) gives the same
    test-row output logits as the uncached full forward pass: max absolute
    difference <= 1e-5 in fp32; bf16 difference is reported for reference.
  • Added an int8-cache vs fp32-cache test that checks quantization error stays
    within tolerance (2% relative, observed 0.34%).
  • Added tests that an unsupported quantization dtype is rejected and that the
    cache is preserved after a device swap.
• tabfm/src/classifier_and_regressor_pytorch_test.py
  • Added cached-vs-uncached classifier/regressor tests with cache quantization
    disabled.

Disclaimer: Claude was used to understand/read the repo and author changes. However, the changes were supervised and were reviewed afterwards. The benchmarks were also run to evaluate the impact on accuracy.

[astonishedrobo]

Hi everyone, I wanted to kindly follow up on this PR. Is there something you will want me to add or modify something?