Source-linked AI summary
Attention Quantization for Tabular Foundation Models
Jonas M. Kübler, Benjamin Jäger, Klemens Flöge, Noah Hollmann, Frank Hutter
TL;DR
Tabular foundation models are relatively small, making weight quantization less beneficial, while attention's effect on predictive performance under quantization is uncertain. The paper quantizes attention computations to FP8 and aligns train- and test-row quantization. It reports up to 1.7x end-to-end inference speedup without relevant quality degradation on TabArena and BeyondArena.
Problem
Tabular foundation models are fairly small, limiting the benefits of weight quantization, while the effect of quantizing attention on predictive performance is unknown.
Source: §1 · p. 2Method
The method quantizes Q, K, and V in fused attention to FP8 and aligns train- and test-row quantization by reusing the train-query scale for test queries.
Source: §2 · p. 2Results
Up to 1.7x end-to-end inference speedup was achieved without relevant quality degradation on TabArena and BeyondArena.
Source: §4 · p. 5Takeaways & Limitations
Attention quantization offers a post-training optimization direction for tabular foundation models, with coordination between train and test rows required.
Source: §4 · p. 5
Abstract
from arXiv · showhide
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.
1 Introduction
Tabular foundation models create a need to reduce inference cost, but their smaller size limits the benefits of weight quantization and the effect of attention quantization on predictive performance is unknown. The paper studies FP8 attention, finding it can be nearly lossless when train–train and test–train attention are both quantized, with kernel speedups up to 1.7x.
1 Introduction
- Transformer adoption in tabular models makes reducing inference cost and latency an important efficiency objective.Quantization has been a successful approach to reducing costs and latency for LLMs.Source: §1 · p. 1
- Tabular foundation models are fairly small, limiting weight-quantization benefits, while attention quantization's effect on predictive performance is unknown.Attention computation is relevant because its cost scales quadratically with training-row count, which can reach hundreds of thousands or millions.Source: §1 · p. 2
- FP8 matrix instructions offer twice the throughput of 16-bit counterparts, but fused-attention speedups depend on hardware and other operations can become bottlenecks.Exponentiation is one example of an operation that can limit the benefit.Source: §1 · p. 2
- Both train–train and test–train attention can be quantized nearly losslessly, with kernel speedups up to 1.7x on an NVIDIA RTX Pro 6000.The authors also report similar accuracy and speed results on TabICLv2.Source: §1 · p. 2
2 Method
The method targets the two attention matrix multiplications that dominate fused-attention runtime and the quadratic row-scaling cost of tabular in-context learning. It quantizes Q, K, and V to FP8, reuses train-query scales for test queries, and implements the approach in a Triton kernel based on FA2.
2 Method
- Fused attention avoids materializing P in global memory, so QK⊤ and PV dominate runtime and are the matrix multiplications this method quantizes.The softmax is row-wise; α is a scalar.Source: §2 · p. 2
- Tabular in-context learning runs train–train and test–train attention, with forward-pass FLOPs growing as O(N^2 + MN), compared with O(N + M) for weight matrix multiplications.Keys and values come from training rows alone, while queries come from both training and test rows.Source: §2 · p. 2
- Q, K, and V are quantized to e4m3fn FP8 using a dynamic per-head absmax scale, while P uses a fixed in-kernel scale and is not materialized globally.The V scale is applied in the epilogue; Q and K scales are applied in the kernel with the softmax scale.Source: §2 · p. 2, Figure 2 · p. 3
- Test queries reuse the scale computed on train queries, preserving batch-invariant predictions without a relevant accuracy change.A dynamically computed test-query scale would couple test rows and make predictions depend on the exact test batch.Source: §2 · p. 3
- The Triton kernel follows FA2's design, and quantization overhead is amortized for N > 8192, so the final proposal gates on that threshold.The baseline is FA2 via torch scaled dot product attention; the quantization approach is similar to FA3.Source: §2 · p. 3
3 Experiments
Experiments use TabPFN-v3 on RTX Pro 6000 Blackwell, with three preprocessing seeds for the TabArena quality evaluation. The study compares attention-quantization strategies, coordination between train and test attention, and latency under the final gated configuration.
3.1 Quality impact on TabArena
- Quantizing both attention calls minimizes degradation, though the ungated configuration has a tiny but statistically detectable Elo regression.Sharing the train-query scale with test queries does not materially alter results and makes predictions batch invariant.Source: §3.1 · p. 3, Table 1 · p. 4
3.2 Coordination between self-attention and cross-attention is needed
- Gaussian-noise experiments show that attention inputs tolerate errors when both calls receive the same K,V version.Independent K,V perturbations across calls do not provide the same coordination.Source: §3.2 · p. 4, Figure 3 · p. 4
3.3 Latency measurements and final configuration experiments
- FP8 speedups reach 1.72x for the kernel and up to 1.67x end-to-end; with the kernel gated above 8,192 training rows, results are essentially on par with baseline on BeyondArena.Source: §3.3 · p. 4
4 Conclusion and Outlook
Attention quantization speeds inference on large tabular datasets without relevant quality degradation, while requiring coordination between train and test rows.
4 Conclusion and Outlook
- For datasets with hundreds of thousands of training rows, attention quantization speeds inference end-to-end by up to 1.7x without relevant quality degradation on TabArena and BeyondArena.The authors report that quantization also cuts cost by a similar factor and that coordination between train and test rows is required.Source: §4 · p. 5
A.1 Details on the Gaussian noise experiment
The Gaussian-noise experiment probes how perturbation type and magnitude affect TabArena performance, including whether train–train and test–train attention calls share their K,V perturbations.
A.1 Details on the Gaussian noise experiment
- Test rows attend to the first K,V head, and the calls should receive the same K,V version.This coordination is the critical condition identified for the TabPFN-v3 architecture.Source: §A.1 · p. 6
- Random perturbations are applied individually to Q, K, and V, with each head scaled by its own RMS and each layer drawing a seeded field.This makes the run reproducible from the seed alone.Source: §A.1 · p. 6
- The perturbation-strength sweep uses multiples of 0.0265, the error usually observed from quantization.The tested multipliers are 1 through 5.Source: §A.1 · p. 6
- The ablations compare perturbing only train–train, only test–train, both independently, or both with coordinated K,V perturbations.In the coordinated case, K,V are perturbed once before the head slice, while Q remains independently perturbed.Source: §A.1 · p. 6
A.2 Per-dataset behavior of the gated final configuration
The gated final configuration is essentially on par with the baseline, with a slight mean Elo drop on BeyondArena that the authors attribute to noise. Per-dataset results show where FP8 is activated and how its relative-error reporting can magnify tiny changes on near-optimal datasets.
A.2 Per-dataset behavior of the gated final configuration
- FP8 activates on 18 of 51 TabArena datasets and 56 of 142 BeyondArena datasets, with no trend toward greater degradation on larger datasets.The gate applies per run; datasets with differently sized splits may therefore be active on some splits but not others.Source: §A.2 · p. 7
- Relative error is reported because the Arenas mix roc_auc, log_loss, and rmse, but it inflates tiny movements on near-optimal datasets.For APSFailure, around 50k training rows, a +2% relative error on one split corresponds to an absolute AUC change of 8 · 10^-5.Source: §A.2 · p. 7
- TabPFN preprocessing variation across seeds is around five times larger than the within-seed FP8 deviations.The authors use this comparison when interpreting the BeyondArena result as noise.Source: §A.2 · p. 7
- BeyondArena shows a slight mean Elo drop, while FP8 wins on average in relative error; the authors attribute this pattern to noise rather than claiming improvement.The comparison uses the final gated configuration against the unquantized baseline at the same seed.Source: §A.2 · p. 7
A.3 Experiments on TabICLv2
On TabICLv2, the authors quantize a single attention call with per-head dynamic scales and evaluate the accuracy impact without gating. Results are nearly lossless and reach 1.56x end-to-end speedup on 512k training rows.
A.3 Experiments on TabICLv2
- TabICLv2 uses the same multi-head attention for train and test rows, so the authors quantize one attention call with per-head dynamic scales.Unlike TabPFN-v3, TabICLv2 does not have separate train–train and test–train attention calls.Source: §A.3 · p. 7
- The ungated TabICLv2 experiment excludes small datasets run in FP32 to avoid confounding FP8 effects with a change from 32-bit to 16-bit computation.The ungated setup was chosen to investigate accuracy impact even where it did not provide speedup.Source: §A.3 · pp. 7–8
- On TabICLv2, attention quantization is nearly lossless with a non-significant p-value, and end-to-end speedup reaches 1.56x on 512k training rows.Other operations take a larger share of overall latency on TabICLv2 than on TabPFN-v3.Source: §A.3 · p. 8
A.4 Generalization to different head dimensions and hardware
The kernel was tested across head dimensions and hardware platforms, with speedups varying by configuration. Its benefits are bounded on some GPUs by kernel design and softmax bottlenecks.
A.4 Generalization to different head dimensions and hardware
- On B200 or H100 GPUs, the current FA2 design lacks TMA support, and B200 softmax bottlenecks can erase gains from FP8 matrix multiplications.The paper says the kernel would need TMA to reach 16-bit performance on hardware like B200 or H100.Source: §A.4 · p. 8
- Up to 1.82x speedups were observed at head dimension 128 on an NVIDIA RTX PRO 6000.The Table 6 caption identifies the single-attention-call setting.Source: §A.4 · p. 8, Table 6 · p. 10
- Up to 1.89x speedups over torch’s SDPA were observed on an NVIDIA L4 at head dimension 64, though the baseline kernel is not perfectly optimized.The authors note their own 16-bit variant is faster than torch’s SDPA.Source: §A.4 · p. 8, Table 7 · p. 11