Benchmarking Attention for Tabular Foundation Models

1 Introduction

• Tabular in-context learners commonly alternate attention across samples and across features, since full attention across both dimensions is prohibitively expensive.Some architectures instead use a multi-stage row- and column-wise approach.Source: §1 · p. 1
• Efficient-attention implementations have largely targeted long, one-dimensional language-model sequences, while tabular attention has asymmetric row and column lengths, small hidden dimensions, and bidirectional attention.Row attention can involve hundreds to tens of thousands of rows, whereas column attention typically spans tens to low hundreds of features.Source: §1 · pp. 1–2
• The standard tabular tensor layout makes column attention contiguous in memory, while row attention requires transposed views that can trigger costly copies.A backend supporting strided tensors can avoid some of those copies.Source: §1 · p. 2
• The paper benchmarks five training and two inference-only backends across realistic row and column shapes on A100, H100, and B200 GPUs.It also reports an AI agent-based optimization of the FA4 kernel on B200.Source: §1 · p. 2

2.1 General attention

• Optimized attention implementations reduce memory footprint to O(N), and available backends include PyTorch SDPA, cuDNN, FlashAttention, and inference-only vLLM and SageAttention.FlashAttention variants target different NVIDIA GPU generations; vLLM and SageAttention support forward passes only.Source: §2.1 · p. 2

2.2 Tabular attention

• Applying full attention across both table dimensions is infeasible because of quadratic scaling, and locality-based Vision Transformer approaches do not directly fit row- and column-permutation-invariant tabular data.Recent models therefore apply attention alternately along rows and columns.Source: §2.2 · pp. 2–3
• Column attention attends across features for each sample, while row attention attends across samples for each feature; their sequence lengths are C and R, respectively.For PyTorch-native backends, row attention requires copies of queries, keys, and values, whereas FlashAttention can use a zero-copy transpose view.Source: §2.2 · p. 3
• The row-first layout keeps column attention contiguous but row attention non-contiguous; it may still be preferred because rows typically outnumber columns and its compact stride can improve L1 cache hits.The row-first layout is also standard in many current tabular models.Source: §2.2 · p. 3
• Tabular models use smaller hidden dimensions than LLMs, and some backends restrict supported head dimensions.The passage gives examples ranging from 12 heads of dimension 64 in ConTextTab to eight heads of dimension 16 in TabICL.Source: §2.2 · p. 3

3 Benchmark and Experiments

• The benchmark fixes B=1, varies column sequence length C from 16 to 2048 at R=1024, and varies row sequence length R from 32 to 131072 at C=64.The main configuration uses 12 heads with head dimension 64; other configurations are included in the appendix.Source: §3 · p. 3
• Five backends support forward and backward passes, while vLLM and SageAttention are evaluated for forward throughput only.Wrappers are tested against a naive PyTorch implementation for correctness.Source: §3 · p. 4
• The benchmark focuses on standard full bidirectional attention, approximating the zero-query limit and single-query inference rather than modeling train-test splits in detail.TabPFN and Mitra use cross-attention formulations, whereas ConTextTab uses masking.Source: §3 · p. 4
• Timing uses CUDA events, cold-cache runs, stabilized GPU clocks, warmups, and repeated measurements; throughput and peak memory are reported.Backward time is estimated by subtracting mean forward time from combined forward-plus-backward time.Source: §3 · p. 4
• Measurements cover A100, H100, and B200 GPUs, using bfloat16 and non-causal attention, with backend availability and head-dimension constraints varying by hardware.Figure 1 reports forward throughput for column and row attention at 12 heads and head dimension 64; some backend measurements are unavailable because of support constraints or GPU memory limits.Source: §3 · p. 4, Figure 1 · p. 5
• The self-contained setup records configuration and timing statistics, flushes results after each completed shape, and releases benchmark code, results, and plotting scripts.The release is available at the paper's stated repository.Source: §3 · p. 4

4 Results

• On H100, SageAttention performs poorly for column attention below 2k columns but overtakes other backends for row attention beyond 16k rows, scaling to 128k rows for smaller head dimensions.SageAttention is approximate, so its effect on predictive performance must be evaluated for the model at hand.Source: Figure 2 · p. 6, §4 · pp. 6–7
• Across GPU families, cuDNN is state-of-the-art or close for column attention; on B200, FA4 leads through 256 columns, while cuDNN wins at longer sequences.On H100, cuDNN dominates from 128 to 1k columns and FA3 is slightly faster at longer sequences.Source: §4 · p. 6
• On A100 and H100, FA3 dominates row attention except below 128 rows, reaching up to 3.5× speedup over SDPA and around 1.5× over cuDNN.The passage attributes this advantage to avoiding copies through stride support and to better tuned-kernel performance on longer sequences.Source: §4 · p. 6
• On B200, FA4 leads row attention except beyond 8k rows, where cuDNN overtakes it; the relative copy overhead diminishes at longer sequences, and the crossover depends on batch size and hidden dimension.The backward-pass behavior is similar.Source: §4 · p. 6
• 2.4×: At very large head dimensions, FA3 reaches this speedup over cuDNN for row attention; column attention can reach 1.4× only at much larger dimensions.For row attention, FA3's advantage grows with head dimension as the avoided contiguous-copy overhead scales linearly with D.Source: Figure 3 · p. 7, §4 · p. 7
• For H100 with H=12 and D=64, the authors recommend cuDNN for column attention and FA3 for row attention, with possible SDPA (efficient) fallback at very short lengths.There is no one-fits-all choice; model dimensions, GPU family, and sequence range affect the optimal backend.Source: §4 · pp. 7–8

Operational intensity of tabular attention

• Column attention at C=16–256 is memory-bound (OI=8–128), while row attention at R=512–16,384 transitions toward compute-bound operation (OI=256–8192).For the H100 NVL, the ridge point is approximately 500 FLOP/byte at N≈1000; cuDNN approaches the memory-bound ceiling at C=128, while longer column sequences remain limited by insufficient parallelism.Source: §5 · p. 8

Quantifying the copy overhead for row attention

• At N=32, copy bandwidth is 783 GB/s; for N≥2048, it saturates at 1.45 TB/s, both below the 3.9 TB/s peak HBM bandwidth.Non-sequential reads in strided-to-contiguous copies prevent full bandwidth utilization.Source: §5 · pp. 8–9
• At N=256–1024, measured copy overhead explains the majority of the FA3–cuDNN throughput gap; at N>2048, FA3’s copy-burdened estimated throughput exceeds measured cuDNN throughput, indicating greater long-sequence kernel efficiency.Source: §5 · p. 9
• Because copy overhead grows linearly with head dimension D, FA3’s advantage over cuDNN grows with D.This explains the pattern reported for the head-dimension ablation in Figure 3.Source: §5 · p. 9

Implications

• The analysis supports a mixed-backend strategy: cuDNN for memory-bound column attention and copy-avoiding FA3 for longer row attention.The authors predict that higher compute-to-bandwidth ratios will extend the memory-bound regime for columns and increase row-attention copy overhead, further favoring strided-input backends.Source: §5 · p. 9

6 End-to-end Validation

• The full-inference evaluation is an initial, inference-only validation restricted to TabPFN-v2.6 on one H100 GPU.It uses a single query row and requires model-specific surgery, so the authors characterize it as supportive rather than comprehensive evidence across architectures.Source: §6 · p. 9
• Up to 15% speedup at fixed sample size across column lengths and more than 300% versus naive SDPA baselines in the row sweep; FA3 gains at longer sequences over TabPFN’s default dispatch.For narrow tables cuDNN performs best; at longer column sequences, cuDNN, FA3, and default dispatch perform similarly.Source: §6 · p. 9, Figure 5 · p. 10
• 7.4% mean speedup over TabPFN’s default dispatch, with gains up to 30% and latency reductions up to 1 s, across TabArena single-fold lite tasks.The combination uses cuDNN for column attention, memory-efficient SDPA for row attention below 2k rows, and FA3 for larger row counts; speedups vary widely with table size.Source: §6 · p. 9

7 Conclusion

• FlashAttention-3/4 dominate row attention, while efficient SDPA excels on short column sequences and cuDNN on longer ones; backend gaps can exceed 3× for a shape.The preferred FlashAttention implementation depends on GPU family.Source: §7 · p. 10
• Contiguity overhead can dominate medium-sequence runtime, while FA3/FA4 reduce it by supporting more compact input strides.The measured impact depends on batch size and hidden dimension.Source: §7 · p. 10

A.1 Quick overview of attention backends

• All evaluated backends compute exact scaled dot-product attention but differ in kernel design, hardware specialization, and tensor-layout support.Source: §A.1 · p. 13

A.2 A Note on Tabular Attention in different Model Families

• The benchmark uses full bidirectional attention rather than each architecture’s exact formulation; TabICL-like multi-stage attention is outside its scope.The authors say key row/column asymmetries still transfer to multi-stage models, while recurring axis permutations and associated copies do not.Source: §A.2 · pp. 13–14

A.3 Tensor layout considerations

• The row-first layout has identical theoretical copy cost to column-first, but may improve cache locality for feature-wise MLP layers.Table 1 reports layout copy-cost ratios within 0.3% of unity.Source: §A.3 · p. 14, Table 1 · p. 14

• The .contiguous() call allocates memory and copies tensor data, an O(N) operation that can become a non-negligible repeated overhead.Source: §A.4 · p. 14
• Row attention transposes the row-first layout into a non-contiguous tensor, requiring copies for queries, keys, and values; column attention can use zero-copy views.Stride-compatible backends such as FlashAttention-3 can operate directly on transposed views and reduce intermediate copies.Source: §A.4 · p. 15
• The H100 measurement reports bandwidth for three consecutive Q, K, and V copies at row-attention shapes, with bandwidth increasing with tensor size but remaining below peak HBM bandwidth.The table specifies Beff=64, H=12, D=64, and bfloat16.Source: §A.4 · p. 15, Table 2 · p. 15

B Additional Results

• The appendix presents a selective subset of configurations, while raw measurements can be used to generate plots for all benchmarked GPU families.Source: B Additional Results · p. 16

B.1 Batch Size Sweep

• For H100 with 12 heads and D=64, row-attention throughput changes little across batch sizes 1, 2, 4, and 8; column throughput is slightly lower at larger batches only in the long-sequence regime.At smaller column counts, batch size produces no observed throughput difference.Source: §B.1 · p. 16

B.2 Peak Memory

• H100 peak-memory scaling is similar across backends, with only minor absolute differences and FA3 showing a slightly smaller footprint; no backend has a clear memory advantage.The reported configuration is 12 heads with head dimension 64.Source: §B.2 · p. 16, Figure 7 · p. 17

B.3 Throughput

• Supplementary throughput plots cover hidden dimensions 128, 192, and 768; the 128-dimension, 8-head, D=16 setting was not evaluated on B200 because FA4 lacks support for D=16.Plots are listed for A100 and H100 at D=16, and for A100, H100, and B200 at D=32 and D=64.Source: §B.3 · pp. 16–17

• The mixed dispatch is faster on 92% of TabArena-lite tasks, with a mean speedup of 7.4% and a maximum of 29%.On a handful of small tasks it is marginally slower; maximum latency reduction is roughly 2.8 seconds on the largest tasks.Source: §B.4 · p. 17
• Mean speedup is 14% above 10,000 training rows versus 4.5% below 10,000, while predictive-score differences reach at most 0.004.The comparison uses accuracy for classification and R2 for regression tasks.Source: §B.4 · p. 17

B.5 End-to-end validation of SageAttention

• The evaluation compared SageAttention with each model's base implementation on 48 CARTE and TabArena-lite classification tasks spanning 500–100,000 rows and 4–1,700 columns.Accuracy was the target metric, and statistical significance was assessed with Friedman and post-hoc Nemenyi tests.Source: §B.5 · p. 17
• With ConTextTab, SageAttention caused a statistically significant average accuracy loss of about 3 percentage points (close to 200 Elo points), confirmed by the critical-difference diagram.In a separate SageAttention comparison, accuracy rose from 84.4% to 84.5% and Elo by 39 points, but the change was not statistically significant.Source: §B.5 · p. 18

C.1 Optimization setup

• The B200 optimization targeted FA4 at 12 heads and 64 dimensions, using eight parallel agents split between row and column attention.Each agent had a separate GPU; the head configuration matched the largest hidden dimension used in a recent tabular foundation model.Source: §C.1 · p. 18, C Agent-based Kernel Optimization · p. 18
• Agents worked on short or long sequences using either structured priorities or open-ended exploration, with shared JSONL logs recording results and code changes.Four agents focused on each attention type, with two prescriptive and two creative agents in each group.Source: §C.1 · p. 18

C.2 Results

• The kernel fused softmax passes, doubled TMEM store repetition and correction-tile size, capped KV pipeline stages, and enabled dynamic CLC scheduling.These changes reduced repeated traversal, store instructions, pipeline prefetch, correction-loop iterations, or per-tile scheduling arithmetic.Source: §C.2 · p. 19
• 11–14% higher throughput was achieved for column attention at sequence lengths larger than 256, while shorter lengths regressed, likely from fixed overhead.The optimized kernel approximately maintained FA4's top performance at short lengths and narrowed its gap to cuDNN at longer lengths.Source: §C.2 · p. 19

C.3 Shortcomings and negative results

• No tested configuration yielded reproducible row-attention gains over baseline; further gains may require deeper architectural changes beyond current agent capabilities.The result may indicate that the FA4 row kernel is already near-optimal for this regime.Source: §C.3 · p. 19
• Merging agents' changes erased individual gains, and the best result came from one agent's checkpoint rather than a combined kernel.The authors suggest exploring sequential integration strategies.Source: §C.3 · p. 19
• Agents gravitated toward easy parameter tuning despite instructions, while prescriptive agents produced most optimizations and creative agents independently confirmed some discoveries.The findings indicate that open-ended exploration was less effective than structured search for this kernel task.Source: §C.3 · p. 19
• Proposed improvements include benchmark-validated external control, profiler-guided prompts, sequence-length-specific kernels, and incremental multi-agent coordination.The coordination protocol is proposed to address difficulties sharing and integrating discoveries.Source: §C.3 · p. 20

D A note on LLM usage

• Claude Code generated some plotting, visualization, Makefile, and documentation boilerplate; Claude also assisted with drafting and reviewing the paper under human feedback.The authors say this review included technical arguments, grammar, and typography.Source: D A note on LLM usage · p. 20
• The shared prompt directs agents to structurally optimize the FA4 forward kernel for B200 tabular workloads, using conditional specialization when changes help only some configurations.The workload uses bf16, head dimension 64, 12 heads, and non-causal attention; conditional variants preserve baseline logic on other configurations.Source: D A note on LLM usage · pp. 21, 23
• A/B tests use five interleaved pairs and paired t-tests; a per-configuration change counts only when p < 0.05 and effect size exceeds 0.5%.The overall verdict is improved only if at least one configuration improves and none regress; any regression makes the overall verdict regressed.Source: D A note on LLM usage · p. 23
• Profiling found all configurations memory-latency bound: compute throttling stayed below 1.5%, while memory and barrier stalls dominated.The guidance recommends focusing on memory access patterns and synchronization rather than compute throughput.Source: D A note on LLM usage · p. 23
• Profiling short column-attention sequences found 32% memory-latency stalls, 30% barrier stalls, 17% occupancy, 40–50K SMEM bank conflicts, and <0.2% math throttle.The profiled cases were sequence lengths 32, 128, and 512 with batch_eff=1024; the compute pipe was barely used.Source: D A note on LLM usage · p. 24
• Recommendations target SMEM bank conflicts, unnecessary single-N-block softmax rescaling, and an empty warp that still participates in barriers and register allocation.Suggested remedies include padding or swizzling SMEM accesses, compile-time removal of rescaling, and letting warp 15 exit for hdim=64.Source: D A note on LLM usage · p. 25
• For short sequences, proposed specializations reduce tile_n and SMEM allocation, simplify the single-N-block pipeline, and fuse the epilogue store into the MMA warp.At seqlen=32, tile_n=128 wastes 75% of loaded K/V data; the epilogue warp’s TMA store can remain on the critical path after MMA finishes.Source: D A note on LLM usage · p. 25
• Scheduler recommendations are to profile persistent-scheduler overhead for grids with ~12K+ CTAs and bypass CLC setup when short-sequence setup costs may exceed its benefit.The cited large-grid example uses batch_eff=1024 and nheads=12.Source: D A note on LLM usage · p. 25