🤖 AI Summary
This work addresses the memory bottleneck of Compressed Sparse Attention (CSA), which is constrained by the excessive size of intermediate score tensors and thus limited to single-GPU memory, hindering its application to long sequences. The paper introduces the first CSA indexer that avoids fully materializing intermediate tensors, leveraging a streaming partition-merge top-k algorithm to efficiently generate attention indices without instantiating the full score matrix. Integrated with a Triton-based chunked top-k dispatcher, a TileLang-pipelined attention kernel, and key sharding, the proposed method achieves remarkable scalability: on an NVIDIA H200 GPU, it processes sequences up to 1,048,576 tokens with a peak memory footprint of only 6.21 GB—32× longer than prior approaches—while maintaining high fidelity, with an average recall of 1.0000 (minimum 0.9980).
📝 Abstract
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.