Attention Quantization for Tabular Foundation Models
This paper details a method to make tabular foundation models run faster by optimizing the attention calculation process, which is a key part of these models' performance.
This research paper proposes a quantization strategy for queries, keys, and values in tabular foundation models to FP8, using explicit FP8 matrix multiplication instructions to speed up attention calculation. The study finds that aligning quantization error between training and test rows is crucial to avoid accuracy drops. The Triton kernel achieves up to 1.7x speedup over 16-bit kernels, with no relevant accuracy loss on TabPFN-v3 and TabICLv2 across TabArena and BeyondArena.
With the recent rise and adoption of tabular foundation models, optimizing their inference performance becomes an emerging field for efficiency research. While the models are architecturally similar to transformer-based large language models (LLMs), the size and serving patterns differ significantly. We show that the focus should be on the attention calculation and less on weight or KV cache quantization, which are more popular in LLMs. We develop a quantization strategy for queries, keys, and values to FP8 and use explicit FP8 matrix multiplication instructions to speed up the attention calculation. We find that it is crucial to align the quantization error in the test rows with the quantization error in the training rows, as otherwise the accuracy drops drastically. Our Triton kernel achieves a speedup up to 1.7x over regular 16-bit kernels, and we show that on TabPFN-v3 and TabICLv2 there is no relevant accuracy loss across TabArena and BeyondArena.