Benchmarking Attention for Tabular Foundation Models¶
会议: NeurIPS 2026
arXiv: 2609.31306
OpenReview: https://openreview.net/forum?id=1MbSbrblLi
代码: https://github.com/SAP-samples/tabular-attention-benchmark
领域: 自监督学习
关键词: 表格基座模型, 行列注意力, 内存连续性, FlashAttention, 硬件基准测试
一句话总结¶
针对表格基座模型中二维交替行列注意力的特殊形态，系统评测了横跨三代 GPU 的 7 类注意力后端，量化了内存步幅转置与张量连续性开销，提出了行注意力与列注意力解耦的混合后端调度策略并实现端到端加速。
研究背景与动机¶
表格数据构成了企业级机器学习应用的主体场景。随着预训练上下文学习范式的发展，以 TabPFN、TabICL、Mitra 和 ConTextTab 为代表的表格基座模型在公开基准上展现出了逼近甚至超越梯度提升决策树（GBDT）的强劲预测能力。这类模型在架构上普遍依赖对隐层嵌入执行二维交错的注意力机制：由于全表单元格联合注意力的计算复杂度随行列相乘呈四次方剧烈膨胀，现有模型通过交替执行跨样本的行注意力（Row Attention）与跨特征的列注意力（Column Attention）来建模二维上下文依赖。
然而，近年来高效注意力算子的演进与优化主要由大语言模型（LLM）的算力需求所驱动。无论是 FlashAttention 家族还是 vLLM、SageAttention 等推理引擎，其底层优化均针对一维长序列、因果掩码解码以及数千隐藏维度的场景定制，忽视了表格领域的独特性质。表格注意力天然具有高度非对称性：特征列数通常仅在数十至数百区间，而样本行数则可延伸至数万级别；当前表格模型的单头隐藏维度普遍偏小（例如 16 到 64，远低于 LLM 常见的 128）；且表格注意力属于非因果的全双向注意力。更为关键的是，标准五维表格张量布局为行优先形态 \((B, R, C, H, D)\)，列注意力天然沿着连续内存轴展开，而行注意力在转置至序列维度时不可避免地破坏连续性，导致底层频繁触发昂贵的数据搬运。
这种张量布局与算子底层内存访问模式之间的脱节，使得盲目将大语言模型的注意力库套用于表格模型时往往遭遇严峻的性能退化。本文的核心 idea 是：系统刻画表格二维交替注意力的计算密度与内存布局非对称性，严格基准测试 5 种训练与 2 种推理后端在跨三代 GPU 上的吞吐边界，揭示内存连续性开销在中小序列上的瓶颈地位，并提出依据行列维度与头维度解耦选择最优算子的混合后端策略。
方法详解¶
整体框架¶
在基于 Transformer 的表格基座模型中，输入表格嵌入被组织为行优先五维张量，其形状为 \((B, R, C, H, D)\)，其中 \(B\) 为批大小，\(R\) 为样本行数，\(C\) 为特征列数，\(H\) 为注意力头数，\(D\) 为单头维度。为了避免全表单元格注意力的灾难性复杂度，模型在每个表征块内交替执行列注意力与行注意力。在执行列注意力时，张量视角被展平为有效批大小 \(B_{\text{eff}} = B \cdot R\)、序列长度为 \(C\) 的算子输入；在执行行注意力时，张量必须经历维度置换，转换为有效批大小 \(B_{\text{eff}} = B \cdot C\)、序列长度为 \(R\) 的运算形态。
评测框架在独立无缓存污染的环境下量化了不同底层算子在这一交错流水线中的表现，覆盖了 PyTorch 原生 SDPA（基于 xFormers 的高效内核与 cuDNN 内核）、FlashAttention 演进序列（FA2、FA3、FA4）以及推理专用的量化与 Triton 算子（SageAttention 2++、vLLM）。整个系统的数据流向与核心瓶颈转换机制如下所示：
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入表格嵌入张量<br/>(B, R, C, H, D) 行优先布局"] --> B["1. 列注意力路径：跨特征建模<br/>零拷贝展平为有效批 (B·R, H, C, D)"]
A --> C["2. 行注意力路径：跨样本建模<br/>轴置换转置为 (B·C, H, R, D)"]
B --> D["3. 算子步幅适配与连续性开销<br/>原生连续布局，仅需单次输出连续化"]
C --> D
D -->|内存受限短序列| E["4. 列注意力与行注意力的非对称计算密度<br/>列算子受限于内存带宽，行算子迈向算力受限"]
D -->|跨步内存拷贝惩罚| E
E --> F["5. 跨硬件与维度的混合后端调度策略<br/>列选 cuDNN，行选中长序列 FA3/FA4 或 SageAttention"]关键设计¶
1. 列注意力与行注意力的非对称计算密度：揭示运算边界向 Roofline 顶点的迁移规律 在非因果注意力机制中，设有效批大小为 \(B_{\text{eff}}\)、头数为 \(H\)、序列长度为 \(N\)、单头维度为 \(D\)。前向计算的总浮点操作量为 \(4 \cdot B_{\text{eff}} \cdot H \cdot N^2 \cdot D\)，而在 bfloat16 精度下输入与输出的最优总显存传输量（读取 Q、K、V 并写入 O）为 \(4 \cdot B_{\text{eff}} \cdot H \cdot N \cdot D \cdot 2\) 字节。两者之比决定了注意力的操作强度（Operational Intensity, OI）： $\(\text{OI} = \frac{4 \cdot B_{\text{eff}} \cdot H \cdot N^2 \cdot D}{4 \cdot B_{\text{eff}} \cdot H \cdot N \cdot D \cdot 2} = \frac{N}{2}\)$ 在现代加速器如 NVIDIA H100 NVL（约 2000 TFLOPS BF16 张量核心算力，3.9 TB/s HBM3 显存带宽）上，Roofline 模型的拐点约为 500 FLOP/byte，对应临界序列长度 \(N \approx 1000\)。这一数学规律直接界定了表格注意力的根本分歧：列注意力的序列长度 \(C \in [16, 256]\)，对应的计算强度仅有 \(\text{OI} \in [8, 128]\)，完全深陷于显存带宽受限区（Memory-bound）；而行注意力在常见表格中样本数 \(R \in [512, 16384]\)，计算强度攀升至 \(\text{OI} \in [256, 8192]\)，迅速跨越拐点进入算力受限区（Compute-bound）。因此，针对大语言模型长序列算力优化优化的内核无法在列注意力上发挥优势，具有极低算子启动开销和优秀带宽利用率的算子（如 cuDNN）在列方向更为出色。
2. 算子步幅适配与连续性开销：量化非连续转置对中间张量的性能吞噬
表格模型中默认的行优先布局 \((B, R, C, H, D)\) 使列维度相邻，因此列注意力可以通过零拷贝的 .view() 变换直接提供给底层，仅在注意力计算输出后需要一次连续化以恢复布局。相反，行注意力要求沿着样本轴计算，置换操作 .permute(0, 2, 1, 3, 4) 仅改变了张量的步幅元数据（Strides），底层的物理存储仍呈现跳跃分布。诸如 PyTorch SDPA（Efficient 与 cuDNN）等后端强制要求输入张量具备标准的行主序连续存储，这迫使框架对行注意力的 Query、Key、Value 三者各执行一次显式的 .contiguous() 内存深拷贝。连续化耗时为：
$\(t_{\text{copy}} = \frac{3 \cdot B_{\text{eff}} \cdot H \cdot N \cdot D \cdot 2}{\beta_{\text{copy}}(N)}\)$
由于跨步非连续读取无法完全饱和 GPU 的突发传输机制，实测带宽 \(\beta_{\text{copy}}\) 在小尺寸时仅为 783 GB/s，在大尺寸下最高也只能饱和至 1.45 TB/s，远低于 H100 的 3.9 TB/s 理论峰值。在样本量处于 128 至 4096 的关键训练区间内，显存数据拷贝耗时占据了算子总执行时间的绝大部分。支持紧凑非连续跨步输入的算子（如 FlashAttention-3 与 FlashAttention-4）能够直接在非连续视图上展开计算，免除 Q、K、V 的深拷贝税，从而在行注意力上获得断层式的吞吐领先。
3. 跨硬件与维度的混合后端调度策略：针对表格超参特性的解耦派发 传统深度学习框架倾向于在整个模型内硬编码单一种类的注意力后端，但基准测试表明没有任何单一算子能在表格的所有切片下保持领先。一方面，表格模型的隐藏层维度远小于 LLM（例如 TabICL 使用 \(D=16\)，TabPFN 使用 \(D=32\)，ConTextTab 使用 \(D=64\)），而前沿算子（如 FA4）在特定硬件上对头维度设置了严苛的最低阈值（H100 上要求 \(D=128\)，B200 上要求 \(D \ge 64\)），导致其直接与大量主流轻量表格模型绝缘；同时，连续拷贝的开销与单头维度 \(D\) 线性正相关，使大头维度模型从跨步跳过中获益更多。另一方面，SageAttention 2++ 等推理端 INT4 量化算子由于启动开销大而在短列场景性能垫底，但在行长超越 16k 时凭借高压缩率在 H100 上实现断层式反超。因此，最优方案必须是动态分流的“混合派发架构”：列注意力统一派发至 cuDNN；行注意力在短序列时交由 SDPA Efficient，中长序列转交具备跨步支持的 FA3/FA4；超长序列推理时按需切换至低比特量化后端。
损失函数 / 训练策略¶
基准测试评测了基于 bfloat16 精度的双向全注意力前向与反向传播。在硬件与运行时环境层面，测试统一采用各后端所兼容的最新配置栈：针对 SDPA、FA3 与 FA4 统一调度 PyTorch 2.10、CUDA 13.0 与 cuDNN 9.15.1；针对 FA2 则回退至 PyTorch 2.8 与 CUDA 12.8。测试流程通过专用的 GPU 缓冲区写入在每次调用前主动刷新 50MB 级别的 L2 缓存，彻底杜绝缓存命中对测速造成的失真；在正式采样前执行 5 次预热迭代并预先执行密集矩阵乘法以锁死 GPU 核心时钟频率。反向传播耗时通过从联合计算总耗时中剔除前向均值精确剥离，浮点吞吐量分别按 \(4 B_{\text{eff}} H L^2 D / t\)（前向）与 \(2.5 \times\) 前向 FLOPs（反向）标准统一换算为 TFLOPS。
实验关键数据¶
主实验¶
主实验在 NVIDIA A100（80GB HBM2e）、H100（80GB HBM3）以及 Blackwell B200（192GB HBM3e）三代代表性加速器上全面测定。配置固定为经典配置 \(H=12, D=64\)（BF16），覆盖特征列数 \(C \in [16, 2048]\) 与样本行数 \(R \in [32, 131072]\)。下表汇总了 H100 上最具代表性区间内的前向计算吞吐核心数据（单位：TFLOPS）：
| 场景 | 序列长度 \(L\) | SDPA (Efficient) | SDPA (cuDNN) | FlashAttention-2 | FlashAttention-3 | 最佳方案与对比优势 |
|---|---|---|---|---|---|---|
| 列注意力 (\(B \cdot R=1024\)) | \(C = 16\) | 42.1 | 22.3 | 19.8 | 18.5 | SDPA Efficient 启动开销最低，约领先 2.1× |
| 列注意力 (\(B \cdot R=1024\)) | \(C = 128\) | 168.4 | 312.5 | 235.1 | 228.6 | cuDNN 逼近内存受限上限，领先 FA3 36.7% |
| 列注意力 (\(B \cdot R=1024\)) | \(C = 512\) | 315.2 | 524.8 | 468.2 | 471.0 | cuDNN 持续主导中等特征列吞吐 |
| 列注意力 (\(B \cdot R=1024\)) | \(C = 2048\) | 442.0 | 608.2 | 592.4 | 615.1 | FA3 在长序列列注意力上微弱超越 cuDNN (1.1%) |
| 行注意力 (\(B \cdot C=64\)) | \(R = 64\) | 24.6 | 38.2 | 26.1 | 31.4 | cuDNN 占优，此时数据量小、内存拷贝耗时短 |
| 行注意力 (\(B \cdot C=64\)) | \(R = 512\) | 142.1 | 215.4 | 338.2 | 485.6 | FA3 免除连续化拷贝，大幅领先 cuDNN 2.25× |
| 行注意力 (\(B \cdot C=64\)) | \(R = 2048\) | 305.8 | 451.2 | 682.0 | 842.1 | FA3 跨步支持拉开代差，绝对加速达到 1.87× |
| 行注意力 (\(B \cdot C=64\)) | \(R = 8192\) | 520.4 | 698.5 | 884.2 | 1085.3 | FA3 充分释放长序列算力，进入计算受限区 |
消融实验¶
针对内存连续性拷贝（.contiguous()）开销的消融验证了显存搬运在 H100 上的吞吐惩罚。下表展示了行注意力配置下（\(B_{\text{eff}}=64, H=12, D=64\)），实际拷贝耗时与理论带宽饱和度分析：
| 样本行数 \(N\) | 张量总内存 (Q+K+V) | 拷贝耗时 \(t_{\text{copy}}\) (\(\mu\)s) | 实测跨步拷贝带宽 \(\beta_{\text{copy}}\) | 理论峰值占比 | 对端到端算子影响说明 |
|---|---|---|---|---|---|
| \(N = 32\) | 9 MB | 24.1 | 783 GB/s | 20.1% | 跨步访问碎片化严重，内存总线利用率低 |
| \(N = 128\) | 36 MB | 56.8 | 1330 GB/s | 34.1% | 拷贝开销占据 cuDNN 总执行时间的近 60% |
| \(N = 1024\) | 288 MB | 422.7 | 1429 GB/s | 36.6% | 拷贝时间与 FA3 内核计算总时间相当 |
| \(N = 4096\) | 1152 MB | 1646.0 | 1468 GB/s | 37.6% | 达到非连续搬运带宽饱和上限（~1.47 TB/s） |
| \(N = 16384\) | 4608 MB | 6665.1 | 1450 GB/s | 37.2% | 计算量 \(\mathcal{O}(N^2)\) 渐进占优，拷贝惩罚比率开始衰减 |
在端到端系统验证中，将该混合调度方案集成至 TabPFN-v2.6 并在 TabArena 基准（单折 lite 评测集，共计 48 个真实数据集分类与回归任务）上运行：优化后的混合调度相较于 TabPFN 官方默认调度在 92% 的数据集上取得了稳定加速，端到端整体平均推理延迟降低 7.4%，在样本数超万行的大表上平均加速达到 14.0%，单任务推理时间最多缩短近 2.8 秒（最高达 29% 加速比），且全任务预测准确率与 \(R^2\) 指标差异严格被限制在 \(\pm 0.004\) 以内。
关键发现¶
- 列注意力是内存受限，cuDNN 近乎全域通吃：列注意力长度绝大多数分布在 16 至 256 之间，算子操作强度 \(\text{OI} \le 128\)，远低于硬件的 500 FLOP/byte 拐点。cuDNN 凭借精细的带宽饱和度调度与轻量启动，在绝大部分场景中显著超越 FlashAttention 系列。
- 行注意力受制于连续性拷贝，FA3 凭借跨步张量支持彻底逆转：在样本行数位于 128 至 4096 区间时，传统 SDPA 必须执行三次针对非连续内存的搬运，拷贝耗时直接吞噬了计算时间。FA3 支持紧凑非连续输入，从而在该区间相对 cuDNN 斩获 1.5× 至 2.4× 的加速，且该优势随单头维度 \(D\) 增大而线性扩大。
- 量化近似注意力具有双刃剑效应：SageAttention 2++ 能够支持高达 128k 样本的超长行推理，并在 TabPFN 上无损精度甚至微弱提升 Elo 分数；但在 ConTextTab 上会导致致命的性能崩塌（平均准确率下降约 3 个百分点，Elo 分数暴跌近 200 分），需严格按模型架构单独验证。
亮点与洞察¶
- 算子与数据维度的正交解耦思路：打破了全局采用单一底层后端的惯性思维，明确指出表格模型的横向特征流与纵向样本流处于完全不同的计算密度与内存连续性工况，应当在单次前向内混合调度不同厂商与特性的算子。
- 量化了被忽视的连续性税收：揭示了在行优先的五维张量体系下，看似无害的
.contiguous()调用在真实 GPU 上的显存访问模式会发生跨步退化，其实测带宽上限被锁死在理论物理带宽的 37% 左右，成为了限制中型表格吞吐的罪魁祸首。 - 闭环探索了 AI Agent 驱动的底层内核自进化：利用基准测试作为高保真反馈环境，调用 8 个基于 Claude Code 的自主代理在 NVIDIA B200 上针对 FA4 的 CuTe DSL 展开了长达 24 小时的全自动内核重构，成功挖掘出融合 Softmax 单遍扫描、双倍 TMEM 寄存器写入等 5 项内核级优化，在中长列上实现了 11%–14% 的前向吞吐提升。
局限与展望¶
- 跨注意力与掩码切片未纳入隔离基准：评测目前以理想化的全双向注意力展开，尚未覆盖 TabPFN 与 Mitra 所依赖的训练集/测试集 Query-Key 交叉注意力架构，以及部分模型采用的特定非对称掩码模式。
- 端到端多卡分布式拓扑尚未覆盖：当前基准集中于单卡合成与单卡推理验证，未涉及跨多 GPU 数据并行或张量切分下的通信开销与张量收集对连续性开销的进一步放大。
- 多阶段解耦模型扩展空间：针对类似 TabICLv2 与 TabPFN-3 这种完全拆分行/列注意力的多阶段架构，可以进一步定制零轴置换的专门优化流水线，从根本上在网络顶层规避转置开销。
相关工作与启发¶
- vs FlashAttention 演进系列 (FA1/2/3/4)：FA 系列主要针对自然语言 1D 长文本、因果掩码以及大头维度（128）设计；本文明确指出表格数据维度天然与 FA 预设工况冲突，在列注意力上 FA 甚至落后于 cuDNN，证明了通用 LLM 算子无法盲目迁移至表格领域。
- vs PyTorch 原生 SDPA (cuDNN / xFormers)：SDPA 在常规连续张量与极短序列上拥有最低的内核调度开销，但因缺乏非连续张量步幅算子支持，在行注意力的轴变换中被迫引入大量中间显存拷贝，成为吞吐洼地。
- vs 纯树模型与传统表格深度学习：树模型在推理时高度依赖 CPU 内存检索，而基于交替注意力的表格基座模型在算力上严重依赖 GPU。本文通过底层算子重排与混合架构优化，大幅消减了基座模型相对 GBDT 的推理延迟劣势，为企业级实时预测部署奠定了基础设施基础。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ 首次针对表格基座模型特有的 2D 交替注意力及五维张量步幅特性建立系统性算子基准，揭示了内存拷贝开销的定量本质。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 7 款训练与推理后端、三代主流企业级 GPU、多种头维度超参组合、Roofline 理论建模以及全套真实 TabArena 端到端落地验证。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨清晰，系统分析了操作强度与非连续内存带宽，论证深刻且洞见切实可行。
- 价值: ⭐⭐⭐⭐⭐ 为后续表格基座模型设计、算子原生定制以及推理部署选型提供了无可争议的实践指南，且已开源全套开箱即用套件。