Disable KV cache quantization when inference dtype specified.#1099
Disable KV cache quantization when inference dtype specified.#1099oscarkey wants to merge 2 commits into
Conversation
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.
There was a problem hiding this comment.
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.
| Setting `inference_precision` to a value other than "auto" disables | ||
| quantisation. |
There was a problem hiding this comment.
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.
| Setting `inference_precision` to a value other than "auto" disables | |
| quantisation. | |
| Setting inference_precision to a specific torch.dtype disables | |
| quantization. |
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