1. 问题与动机

多向量检索（ColBERT 系）用 token 级 MaxSim 打分——对每个 query token 取它与所有 doc token 的最大 cosine 相似度，再对 query token 求和——在 MS MARCO、BEIR 上拿到 SOTA 精度，代价是 O(Nq·Nd·d) 的计算量。论文做了 roofline 分析点出真正的瓶颈：现有 GPU 实现只跑到 HBM 峰值带宽的 5–18%。以典型 ColBERT 参数（Nq=32、Nd=128、d=128、batch=10⁴）算，naive 实现的算术强度只有 16.1 FLOP/byte，而 H100 的 compute/memory 交叉点在 591 FLOP/byte——这是个深度 memory-bound 的算子。根因是 naive 实现会把 Nq×Nd 的相似度矩阵整个物化到 HBM，而这块中间数据"算一次、求完 max 就丢"，白白消耗内存流量。正因为打分慢，现有系统（如 PLAID）只能把 MaxSim 限制在 CPU 或对候选集做激进剪枝，压低了召回上限。TileMaxSim 的目标就是把这个 memory-bound 算子打到接近带宽上界。

2. 方法详解

TileMaxSim 是一族 IO-aware 的 Triton kernel，核心是三招：

(1) Multi-Query SRAM Tiling（V2-MQ）：同时对 query token 和 doc token 分块。每个 block 在寄存器里放 BQ 个 query token、在 SRAM 里放 BN 个 doc token，文档 embedding 流过共享内存的同时，per-query-token 的最大值直接在寄存器里累加。关键是把每个文档 tile 从 HBM 读的次数从 Nq 次降到 ⌈Nq/BQ⌉ 次；当 BQ=Nq=32 时达到 single-pass 最优——每条文档 embedding 只从 HBM 读一次。

(2) Dimension Tiling：把 embedding 维度切成 128 宽的块，支持 d=64~768。对 d>128（会撑爆共享内存）的 embedding 先累加各块的部分点积、再求 max。带宽利用率随维度上升：d=64 时 44%、d=768 时升到 86%。

(3) Fused PQ Decompression+Scoring（TileMaxSim-PQ）：对 product quantization 压缩的库，避免物化解压后的向量。预计算 per-query 距离表 T[i,m,k]=qi[第 m 段]ᵀ·C[m,k]，打分时直接查共享内存里的表而不读 HBM。在 B=100K、Nq=32、Nd=128 下 IO 从 6.75B 字节降到 218M 字节（31× reduction），距离表约 50MB 正好落在 L2 cache。

实验环境：2×H100 SXM 80GB（HBM3，3.35 TB/s、132 SM、每 SM 228KB 共享内存），软件栈 PyTorch 2.6、Triton 3.2、CUDA 12.4。负载 d=128、Nq=32、Nd∈{64,128,256}，batch 100~500K，模型为 ColBERTv2（128 维）。

3. 核心技术贡献

第一，把 MaxSim 重新表述成"读一次、寄存器累加最大值"的 fused kernel，从根上消掉 Nq×Nd 相似度矩阵的物化——这是相对 naive einsum/loop 实现的本质改进，理论 IO 降 2.0×，实测带宽从 18.4% 拉到 80.2%（2687 GB/s / 3350 GB/s），逼近 FlashAttention 这类 memory-bound kernel 的利用率。

第二，BQ=Nq 的 single-pass tiling 是吞吐的关键开关：ablation 显示 BQ 从 16 调到 32，带宽提升 4.5×（872→1302 GB/s，对应 26.6M→39.7M docs/s），原因正是消除了文档的重复 re-read。

第三，fused PQ 把解压和打分合并成查表，让压缩库也能高吞吐打分（B=100K 时 PyTorch PQ 直接 OOM，TileMaxSim-PQ 仍 5.8M docs/s）。这三点都不改变排序结果，是纯工程层面的吞吐/延迟收益。

4. 实验结果分析

主吞吐（Table 1，d=128、Nq=32）：B=100K 时 81.2M docs/s，对比 PyTorch loop 的 372K docs/s 提速 218×；合成均长数据上 82M docs/s，真实 MS MARCO passage 上 71.6M/s（70% 带宽利用）。对 torch.compile（Table 2），B=100K 时 1.29ms vs 10.47ms（8.1×），因为编译器无法把 matmul→max→sum 这条 reduction 链 fuse 起来。对 WARP CPU 引擎，500K 候选下 82.0M vs 175K docs/s（469×），且 WARP 的延迟 96% 花在 query 编码上。

端到端（Table 15）：作为 PLAID 打分阶段的替换，100K 候选打分从 267.4ms 降到 1.21ms（220×），把整条 pipeline 端到端延迟砍掉 98%。

精度保真：MS MARCO 200 query、187565 对上，与参考 PyTorch 的最大打分差 9×10⁻⁶，MRR@10=0.3399、nDCG@10=0.3866、Recall@1000=0.7833 与参考六位小数完全一致；BEIR 的 SciFact/NFCorpus/TREC-COVID 上 nDCG@10 也逐一相同——结论是"任何 MaxSim 实现的 drop-in 替换"。

扩展性：batch 从 1K（17.9M/s、17.5% 带宽）到 500K（83.2M/s、81.4% 带宽）吞吐趋于常数；Nd=32 时 103.1M docs/s、Nd=512 降到 22.0M/s（但达到 2881 GB/s）；FP16/BF16 吞吐几乎相同、FP32 慢 1.6×。

5. 局限与未来方向

作者明确列出两条局限：一是 fixed-Nd padding——对变长语料按固定 Nd 补齐会浪费 38% 的 token 计算，用 length-sorted batching 可回收到 70M/s；二是假设数据全部 GPU-resident（embedding 已在显存里），没处理 doc embedding 超出显存、需从 host/SSD 流式加载的场景。多 GPU 目前只做了 embarrassingly-parallel 的 data-parallel 分片，没涉及跨卡通信优化。隐含但未展开的点：评测主要在 d=128 的 ColBERTv2，更高维（如 d=768 的多向量模型）虽支持但吞吐明显下降（14.7M/s）；query 编码本身的开销（在 WARP 对比里占 96% 延迟）不在本文优化范围，端到端收益要看编码是否也被加速。

6. 工业实践意义

对做语义检索/重排的团队，最直接的结论是：多向量 late-interaction（ColBERT/PLAID）过去因打分延迟难以上线，TileMaxSim 把 100K 候选的打分压到 1.2ms，让它重新进入可用区间。三点可操作：(1) 若已在用 ColBERTv2/PLAID，TileMaxSim 是 colbert_score_packed 的 drop-in 替换，不动候选生成/解压等其他环节、排序结果逐位一致，迁移风险低；(2) 打分变快后可以放宽剪枝、直接对更大候选集精排，换取更高 recall——这对召回-精排漏斗的设计是个新自由度；(3) fused PQ 查表 + 共享内存的思路可迁移到自家 ANN/重排打分 kernel。部署上要注意：方案绑定 H100/HBM3 与 Triton，需 GPU-resident embedding；变长文档要配 length-sorted batching，否则 padding 浪费近四成算力。roofline 给出的 80.2% 带宽上界，也是评估自研 kernel 还剩多少优化空间的现成标尺。代码已开源（github.com/ashutoshuiuc/tilemaxsim）。