Skip to content

Disable KV cache quantization when inference dtype specified.#1099

Draft
oscarkey wants to merge 2 commits into
mainfrom
ok/fit-modes-diff
Draft

Disable KV cache quantization when inference dtype specified.#1099
oscarkey wants to merge 2 commits into
mainfrom
ok/fit-modes-diff

Conversation

@oscarkey

@oscarkey oscarkey commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

A limitation of this approach is that you can't have the dtype automatically determined and also use this for the KV cache, but I'd guess anyone who is in deep enough to disable cache quantization is happy specifying a dtype.

Closes #631
Fixes PRI-154

A limitation of this approach is that you can't have the dtype
automatically determined and also use this for the KV cache, but I'd
guess anyone who is in deep enough to disable cache quantization is
happy specifying a dtype.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request updates the key-value cache quantization logic in TabPFN to only apply when no specific inference dtype is forced (i.e., when inference_precision is set to "auto" or "autocast"). It updates the docstrings across the classifier, regressor, and inference modules, and refactors the test suite to verify this behavior. The review feedback correctly identifies an inaccuracy in the classifier and regressor docstrings regarding which inference_precision values disable quantization, and suggests correcting a spelling inconsistency ("quantisation" vs "quantization").

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.

Comment thread src/tabpfn/classifier.py
Comment on lines +393 to +394
Setting `inference_precision` to a value other than "auto" disables
quantisation.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The docstring states that setting inference_precision to a value other than 'auto' disables quantization. However, if inference_precision is set to 'autocast', the internal forced_inference_dtype_ remains None, meaning quantization is still enabled. Only specifying a concrete torch.dtype actually disables quantization. Additionally, 'quantisation' is spelled with an 's' here, which is inconsistent with the American spelling ('quantization') used throughout the rest of the codebase.

Suggested change
Setting `inference_precision` to a value other than "auto" disables
quantisation.
Setting inference_precision to a specific torch.dtype disables
quantization.

Comment thread src/tabpfn/regressor.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PRI-154 Fit modes give different results

1 participant