StreamIndex让大序列注意力不爆显存，突破65K长度限制。
StreamIndex: Memory-Bounded Compressed Sparse Attention via Streaming Top-k
- 用分块合并的top-k算法避免存储完整得分矩阵
- 在单张H200上支持104万序列长度，峰值显存仅6.21GB
- 适合训练长序列模型的研究者和工程师
DeepSeek-V3.2和V4引入压缩稀疏注意力（CSA）：通过学习的评分投影对压缩后的键进行打分，每个查询选取top-k项，再由稀疏注意力核只读取这些项。现有CSA实现会在top-k前生成大小为[B, S, H_I, T]的FP32得分张量。当索引头数H_I=64、V4-Flash压缩比m=4时，序列长度S=65,536下中间张量达256GB，超出单卡高带宽内存（HBM）容量。本文提出StreamIndex，一个基于Triton的CSA流水线实现，其核心是分块分区-合并的top-k驱动器，全程不显式存储完整中间结果。在合成但真实感强的V4型输入下，单张NVIDIA H200上，传统路径在S=65,536时即显存溢出（OOM），而StreamIndex可运行至S=1,048,576，峰值显存仅6.21GB，扩展倍数达32倍。与传统路径的精确基线相比，小序列下位级召回完全一致；在三个5点设计空间扫描（分块大小、键块大小、top-k）中，平均召回率约1.0000，最小召回率不低于0.9980。该分块驱动器可与TileLang的流水线注意力核协同工作：在S=262,144、V4-Flash维度下，传统索引器+TileLang注意力因显存不足无法运行，而分块索引器+同一注意力可在1.97秒内完成，峰值显存18.56GB。本工作聚焦于索引步骤优化，不声称加速注意力核或真实检查点端到端表现。代码已开源。
原文摘要 · Abstract (English)
DeepSeek-V3.2 and V4 introduce Compressed Sparse Attention (CSA): a lightning indexer (a learned scoring projection over compressed keys) scores them, the top-k are selected per query, and a sparse attention kernel reads only those. Public CSA implementations materialize a [B, S, H_I, T] FP32 score tensor before the top-k reduction. With H_I=64 indexer heads and the V4-Flash compression ratio m=4, that intermediate is 256 GB at sequence length S=65,536, exceeding any single-GPU high-bandwidth-memory (HBM) budget. We present StreamIndex, a Triton implementation of the CSA pipeline whose central component is a chunked partition-merge top-k driver that never materializes the full intermediate. On synthetic-but-realistic V4-shaped inputs at the indexer-step (layer) level on a single NVIDIA H200, the materialize path runs out of memory (OOMs) at S=65,536 with V4-Flash dimensions; StreamIndex runs the same indexer to S=1,048,576 with 6.21 GB peak HBM, a 32x regime extension. Set-overlap recall against the materialize ground truth is bit-exact at small S where both fit; across three 5-point design-space sweeps (chunk size, key-tile size, top-k), mean recall rounds to 1.0000 with min recall at least 0.9980 in every cell. The chunked driver composes with TileLang's pipelined attention kernel: at S=262,144 with V4-Flash dimensions, the materialize indexer paired with TileLang attention OOMs while the chunked indexer paired with the same attention runs in 1.97 s at 18.56 GB peak. Our contribution targets the indexer step; we make no claim of a faster attention kernel or of real-checkpoint end-to-end behavior. Code: https://github.com/RightNow-AI/StreamIndex.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。