StreamIndex：基于流式 Top-k 的内存受限压缩稀疏注意力
摘要
DeepSeek-V3.2 和 V4 引入压缩稀疏注意力 (CSA)：一个 lightning indexer（对压缩 key 的学习评分投影）对其进行打分，按查询顺序选择 top-k，然后稀疏注意力 kernel 只读取这些。公开的 CSA 实现会在 top-k 降低之前物化一个 [B, S, H_I, T] FP32 score 张量。当 H_I=64 indexer heads 且 V4-Flash 压缩比 m=4 时，该中间变量在序列长度为 S=65,536 时为 256 GB，超过任何单 GPU 的高带宽内存 (HBM) 预算。我们提出 StreamIndex，一种 Triton 实现的 CSA 流水线，其核心组件是分块 partition-merge top-k 驱动器，从不物化完整的中间结果。在单个 NVIDIA H200 上进行的 V4-shaped 输入在 indexer 步骤（层）级别的合成但真实的测试中，物化路径在 S=65,536 时因 V4-Flash 维度而耗尽内存 (OOM)，而 StreamIndex 能将同一 indexer 扩展到 S=1,048,576，仅使用 6.21 GB 峰值 HBM，实现 32 倍的规模扩展。分块驱动器与 TileLang 的流水线注意力 kernel 组合：在 S=262,144 时，物化 indexer 搭配 TileLang 注意力 OOM，而分块 indexer 搭配相同注意力仅需 1.97 秒，峰值 18.56 GB。我们的贡献针对 indexer 步骤；我们不声称更快的注意力 kernel 或真正的 checkpoint end-to-end 行为。代码：https://github.com/RightNow-AI/StreamIndex。
引用
@article{arxiv.2605.02568,
title = {StreamIndex: Memory-Bounded Compressed Sparse Attention via Streaming Top-k},
author = {Jaber Jaber and Osama Jaber},
journal= {arXiv preprint arXiv:2605.02568},
year = {2026}
}
备注
11 pages, 3 figures, 7 tables, 2 algorithms, 36 references. Memory-bounded indexer kernel for DeepSeek-V4 CSA via chunked partition-merge top-k. Code: https://github.com/RightNow-AI/StreamIndex