add KV cache dtype as explicit argument [PRI-347]

[jmkuebler]

Issue

Fixes #1119 and #631.
Also this is an alternative to #1099

Motivation and Context

kv_cache_dtypeis a tradeoff decision and should be possible to control explicitly

---

Public API Changes

• No Public API changes
• Yes, Public API changes (Details below)

---

How Has This Been Tested?

I fixed the unit tests to not rely on Monkeypatching anymore.
I also ran the script shared in #631

import random
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
import numpy as np
import torch

from tabpfn import TabPFNRegressor

X, y = load_diabetes(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.33, random_state=42
)


def _set_seeds() -> None:
    torch.manual_seed(0)
    np.random.seed(0)
    random.seed(0)


_set_seeds()
reg = TabPFNRegressor(fit_mode="low_memory", inference_precision=torch.float64)
reg.fit(X_train, y_train)
preds_no_cache = reg.predict(X_test)

reg = TabPFNRegressor(fit_mode="low_memory", inference_precision=torch.float64)
reg.fit(X_train, y_train)
preds_no_cache_repeat = reg.predict(X_test)

_set_seeds()
reg = TabPFNRegressor(fit_mode="fit_preprocessors", inference_precision=torch.float64)
reg.fit(X_train, y_train)
preds_cache_preproc = reg.predict(X_test)

_set_seeds()
reg = TabPFNRegressor(fit_mode="fit_with_cache", inference_precision=torch.float64, kv_cache_dtype="auto")
reg.fit(X_train, y_train)
preds_kv_cache_auto = reg.predict(X_test)

_set_seeds()
reg = TabPFNRegressor(fit_mode="fit_with_cache", inference_precision=torch.float64, kv_cache_dtype="int8")
reg.fit(X_train, y_train)
preds_kv_cache_int8 = reg.predict(X_test)


def _max_diff(a: np.ndarray, b: np.ndarray) -> float:
    return np.max(np.abs(a - b) / np.abs(a))


print("max relative diffs")
print("no_cache vs no_cache_repeat:", _max_diff(preds_no_cache, preds_no_cache_repeat))
print("no_cache vs cache_preproc:", _max_diff(preds_no_cache, preds_cache_preproc))
print("no_cache vs kv_cache (auto):", _max_diff(preds_no_cache, preds_kv_cache_auto))
print("no_cache vs kv_cache int8:", _max_diff(preds_no_cache, preds_kv_cache_int8))

Using kv_cache_dtype="auto" fixes the issue. If int8is used we now throw a warning, when someone forces the inference dtype to sth else.

/usr/lib/python3.12/contextlib.py:81: UserWarning: inference_precision forces a dtype but kv_cache_dtype is 'int8', so the KV cache is still stored as int8 regardless of the forced dtype. Pass kv_cache_dtype='auto' to keep the KV cache in the forced dtype.
  return func(*args, **kwds)
max relative diffs
no_cache vs no_cache_repeat: 0.0
no_cache vs cache_preproc: 0.0
no_cache vs kv_cache (auto): 0.0

---

Checklist

• The changes have been tested locally.
• Documentation has been updated (if the public API or usage changes).
• A changelog entry has been added (see changelog/README.md), or "no changelog needed" label requested.
• The code follows the project's style guidelines.
• I have considered the impact of these changes on the public API.

---

[gemini-code-assist]

Code Review

This pull request replaces the boolean flags quantize_kv_cache and maybe_quantize_kv_cache with a new kv_cache_dtype parameter (accepting 'auto' or 'int8') across the TabPFN classifier, regressor, inference engines, and architecture configurations. The code reviewer recommends using getattr(self, "kv_cache_dtype", "int8") when accessing this attribute on estimators to ensure backward compatibility with older unpickled models. Additionally, the reviewer suggests adding validation to ensure kv_cache_dtype is restricted to valid values.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[jmkuebler]

/gemini review

[gemini-code-assist]

Code Review

This pull request replaces the boolean quantize_kv_cache parameter with a more explicit kv_cache_dtype argument (accepting 'auto' or 'int8') across the TabPFN classifier, regressor, inference engines, and cache size calculations. It also introduces a warning when a forced inference precision is overridden by the int8 KV cache. The review feedback highlights a lack of runtime validation for this new parameter in public APIs, specifically in warn_if_kv_cache_dtype_overrides_inference_precision and get_cache_size, which could lead to silent fallbacks on invalid inputs.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[bejaeger]

Looks great thanks a lot!

I forgot to mention that not all our models currently supported kv cache quantization. So not only is the new kv_cache_dtype argument only relevant when fit_mode is fit_with_cache but also only if the model is v3 (which is the only one that currently supports int8). This is not ideal and might be confusing if we have "int8" as the default. Should we have another "magic" option like "auto" that is "smallest" or "optimized" which will always pick the smallest supported kv cache dtype?

[oscarkey]

Just a few thoughts! I'm on holiday tomorrow and monday, so happy with whatever's decided if you want to merge :)

[bejaeger]

Very nice! thanks for all the changes, looks very clean now.
Approving to avoid blocking you but please consider my other comments before merging.

[bejaeger]

Can we make this a list instead for each model in models and below in build_cache use the model_index? Could otherwise bite us in the future. I'd say we then want to adjust the warning in the resolve function to only be thrown once?

[jmkuebler]

I moved the resolution now within each cache build, so it is per estimator.