HiSparse:把长上下文 decode 的容量瓶颈从 GPU 显存搬到 CPU 内存

长上下文推理里,稀疏注意力解决了”算不动”的问题,却没解决”装不下”的问题。一个 1M 上下文的请求,即便每步 decode 只看 top-k 个 token,它的完整 KV cache 仍然得常驻 GPU HBM,否则没法被快速访问。结果是 batch size 很快被显存撑爆,throughput 在并发还没拉起来时就见顶了。

HiSparse(Hierarchical Sparse Attention)就是 SGLang 社区针对这个容量瓶颈做的 decode 优化:把完整但冷的 KV 放进 CPU pinned memory,GPU 上每个请求只留一个固定大小的”热”buffer,decode 时按需把 top-k 的 KV 从 host swap-in 回来。实测在 GLM-5.1-FP8 的长上下文场景下吞吐最高 5×,256 并发时约 3× 基线。

这篇文档面向想搞清楚 HiSparse 到底省了什么、代价是什么、能不能扩大系统 KV 容量的工程师。代码引用基于本仓库 release/v0.5.15。

它解决的是什么问题:capacity-bound,不是 compute-bound

这是理解 HiSparse 的前提,也是它和一般”加速注意力”工作最不一样的地方。

DeepSeek Sparse Attention(DSA)、NSA 这类稀疏注意力,在每个 decode step 本来就只挑选一小部分 KV 做注意力。算力上,它已经把 self-attention 的二次开销降下来了。但被选中的 top-k 可能落在序列的任意位置,所以完整的 KV 历史必须全部待在显存里,才能被随机访问。

这就出现了一个尴尬的局面:算力富余,显存先爆。

sparse attention is often capacity-bound rather than compute-bound. The KV cache for the full context must remain in GPU HBM … limiting the achievable batch size and overall throughput.

基线(不开 HiSparse 的稀疏注意力)的吞吐在并发数还很低时就早早见顶,因为 KV cache 占用很快撞上 HBM 容量上限。HiSparse 把这条曲线拉成了近似线性,并发越高、显存压力越大,收益越明显。

注意反过来也成立:低并发时 HiSparse 是亏的。额外的 host↔︎device IO 抵消了省显存的好处。它不是”开了就快”,而是”高并发长上下文才快”。这点后面会再讲。

为什么能省 KV cache:分层 + 按需 swap-in

前提:依赖模型自带的稀疏选择

HiSparse 只能用于 DSA 架构(DeepSeek-V3.2、GLM-5.1)和 DeepSeek-V4。原因很直接:这些模型每步 decode 本来就只看 top-k。既然模型本来就不需要全量 KV 参与,那 GPU 上就没必要常驻全量,只留 top-k 在显存、其余放 CPU,是无精度损失的。

It is possible to keep only the top-k KV on GPU while storing the full KV in host memory — without accuracy loss.

普通 dense 注意力不能这么干,因为它每步都要看全部 KV。HiSparse 的”省”是建立在模型已经稀疏这个事实之上的。

每个 decode step 的流程

来自本仓库 docs/advanced_features/hisparse_guide.md 的工作流:

1. Forward decode — 正常前向,生成下一个 token
2. Top-k selection — 用注意力分数选出最相关的 token 位置
3. Swap-in — 专用 CUDA kernel 把 top-k 的 KV 从 host 搬到 device buffer:
   • 短序列(seq_len ≤ device_buffer_size):快路径,KV 本来就在 buffer 里,不搬
   • 长序列:命中检测 → LRU 重排 → miss 时执行 host→device copy
4. Decode attention — 只用 device 上的 top-k 位置算注意力
5. Eager backup — 异步把”上一个 token”的 KV 从 device 拷回 host,给下一轮腾位置

第 3 步是整个设计的核心,下面单独讲。

核心:swap-in CUDA kernel

swap-in kernel 把三件事融合进一个 kernel:

1. 识别 top-k 在 device buffer 里哪些是 cache miss
2. 用 LRU 策略挑出要被驱逐的 slot
3. 更新 page table,把缺的条目从 host 搬到 device

本仓库 python/sglang/jit_kernel/hisparse.py 里这个 kernel 的 Python 接口是 load_cache_to_device_buffer_mla / load_cache_to_device_buffer_dsv4_mla,关键参数一眼就能看出设计:lru_slots(LRU 槽)、host_cache_locs、device_buffer_locs、top_k_tokens、seq_lens。命中检测 + LRU 驱逐 + 搬运,全在 kernel 里完成,而不是拆成几个来回的 Python 调用。

热 buffer 的大小和驱逐策略直接影响 miss 率。博客给的数据:把 hot buffer 从 2048 提到 4096 slot,配合 LRU,miss 数大幅下降,直接换来关键路径上更低的 swap-in 延迟。

1
2
3
4
5
6
7
8
9

                  ┌─────────────── GPU HBM ───────────────┐
每个请求  ───────▶ │  device buffer (固定, 如 6144 tokens)   │  ← 只放热集
                  │  top-k 在这里做 attention               │
                  └──────────────────┬─────────────────────┘
                                     │ swap-in (on-demand, top-k miss)
                  ┌──────────────────▼─────────────────────┐
                  │  CPU pinned memory (host pool)          │  ← 放完整 KV
                  │  host_to_device_ratio 决定大小 (5~10×)   │     便宜、大、慢
                  └─────────────────────────────────────────┘

容量怎么算:per-request 显存从 O(seq_len) 降到 O(固定)

这点直接看 python/sglang/srt/managers/hisparse_coordinator.py 最清楚:

• 每个 request 在 GPU 上只占固定大小的 device buffer:device_buffer_size(例如 6144 tokens),和序列长度无关。
• host pool 的大小由 host_to_device_ratio 决定。HiSparseCoordinator.__init__ 里对 DSA 走 MLATokenToKVPoolHost(host_to_device_ratio=...),对 DSV4 走按页分配的 DeepSeekV4PagedHostPool。
• req_to_device_buffer 形状是 (max_num_req_slots, padded_buffer_size),padded_buffer_size = device_buffer_size + page_size,给新 token 预留一页。

对比一下开销结构:

“省 KV cache”的本质,就是把不活跃的 KV 从贵的 HBM 挪到便宜的 DRAM,HBM 里每个请求只留一个固定热集。同样的 HBM 能塞下多得多的并发请求。

和 PD 分离部署怎么配合:direct-to-host

HiSparse 目前只在 PD disaggregation 模式下、且只在 decode 侧开启,prefill 侧完全无感。

有个很自然的问题:prefill 算完的 KV 要搬到 decode,本来 decode 侧就要收 KV,现在 decode 侧的 KV 又分了 host/device 两层,这一搬不就更麻烦了吗?

HiSparse 的做法是 direct-to-host:prefill 侧通过 RDMA 把 KV 直接写进 decode 侧的 host pool,绕过 decode 的 GPU。

1
2
3
4
5
6
7

Prefill GPU  ──RDMA──▶  Decode Host Pool (CPU pinned memory)
                              │
                              ▼
                     alloc device buffer (如 4KB tokens)
                              │
                              ▼
                     swap-in kernel (on-demand top-k)

好处有两个:一是省掉了 KV 迁移时 decode 侧 GPU 上的瞬时显存尖峰(不用先落到 GPU 再转存);二是去掉了中间的 staging DMA 步骤。

DeepSeek-V4 这里有个细节:dsv4 的 KV 结构是分层的(C4 / C128 / c4_indexer),direct-to-host 路径只把 C4 KV 写进 host pool,c4_indexer 和 C128 KV 仍然是 device-to-device 传输。这部分由 PR #24880 引入,本仓库里 python/sglang/srt/disaggregation/decode.py 和 disaggregation/nixl/conn.py 是相关路径。

怎么用

两个参数,都在 decode 实例上:

配置项含义:

host_to_device_ratio 按机器主存配:约 1 TB 主存 → ratio 5;约 2 TB 主存 → ratio 10。

Decode 实例(开 HiSparse):

1
2
3
4
5
6
7
8
9
10
11
12
13

python3 -m sglang.launch_server \
    --model-path /path/to/model --trust-remote-code \
    --port 8000 --host 0.0.0.0 \
    --context-length 81920 \
    --tp-size 8 --dp-size 8 --enable-dp-attention \
    --mem-fraction-static 0.85 \
    --disable-radix-cache \
    --disaggregation-mode decode \
    --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \
    --dist-init-addr 127.0.0.1:5757 \
    --nnodes 1 --node-rank 0 \
    --enable-hisparse \
    --hisparse-config='{"top_k": 2048, "device_buffer_size": 6144, "host_to_device_ratio": 10}'

几个要点:

• prefill 实例不需要 --enable-hisparse,它对 HiSparse 无感。
• DSA 模型的 --kv-cache-dtype 默认 auto,Blackwell(SM100+)上解析成 fp8_e4m3,老架构是 bfloat16。decode 后端按 KV dtype 自动选:bfloat16 → flashmla_sparse,fp8_e4m3 → flashmla_kv。
• DeepSeek-V4 不走 DSA 后端,它用自己的 dsv4 后端,默认 fp8_e4m3 KV cache。
• 实测有个坑(来自 PR #24880 的讨论):prefill + PD 时不要给 prefill 加 --enable-hisparse。当前实现里 prefill 会根据 decode 发来的 metadata 决定走不走 hisparse 的发送路径,如果 prefill 自己也开了 hisparse 用的 page size 是 64,而逻辑 page size 是 256,会出错。

能扩大系统”本来”的 KV 容量吗

这个问题要拆成三层回答,别混在一起。

总的可寻址 KV 历史:能,而且大幅扩大。 GPU HBM 是固定的,HiSparse 不能凭空加显存。但它把 CPU DRAM 当第二层,系统能承载的”完整 KV 历史”从受限于 HBM(百 GB 级)跃升到受限于主存(TB 级)。从 host_to_device_ratio 和主存配置看,host pool 可以比 device pool 大一个数量级。能容纳的上下文长度 / 并发请求数,被显著扩大了。

GPU 上的热工作集:反而更小,但这是好事。 每个请求在 GPU 上只占 device_buffer_size(6144)而非全长度。同样 HBM 能容纳的并发请求数大幅上升,这是 throughput 近线性增长的直接来源。

代价:不是免费午餐。 top-k miss 时要从 host 搬数据,有 PCIe/IO 开销。低并发时这个开销会抵消省显存的好处,反而更慢。它需要 PD 模式、只开在 decode 侧。DSV4 的 direct-to-host 走 RDMA 直写 host,绕过 GPU,消除了迁移时的瞬时显存尖峰。

一句话:HiSparse 没有增加 GPU 物理显存,但通过”CPU DRAM 存全量 + GPU 存热集”的分层结构,把系统有效可服务的 KV 容量和 decode 并发度提升了一个数量级,代价是引入 host↔︎device 的 IO,适合高并发长上下文场景。

已知限制和待办

HiSparse 还是实验特性,roadmap(Issue #28874)里列了一堆没干完的活,挑几个重要的:

• 和 MTP 的兼容:目前在推进,DSV4 + MTP 的组合需要小心。
• kernel 效率:swap-in kernel 的 shared memory 用得重,并行策略还有优化空间。
• IO 隐藏:在做 TwoBatch Overlap 和基于 IndexShared 的 IO/compute 重叠,降低长上下文下 IO 的影响。博客也提到,未来 Grace Blackwell 这类 CPU-GPU 带宽更高的平台会进一步缓解这个开销。
• 和 HiCache / decode L2 RadixTree 集成:还没完全打通,host kvcache 管理和 KV transfer 的 allocator 还需要对齐。
• 跨 TP rank 共享 CPU 内存:在做(PR #27370)。
• 启动 OOM:max_total_num_tokens 估计偏高会导致启动 OOM(PR #28752 在修)。

历史上还出过精度 bug(Issue #21545,NSA + hisparse 在 DeepSeek-V3.2 上输出错误)。所以别把它当稳定生产特性无脑开,尤其是低并发场景。

相关文件

参考

• HiSparse: Turbocharging Sparse Attention with Hierarchical Memory — LMSYS Blog
• HiSparse 官方文档 — docs.sglang.io
• PR #20343: HiSparse for Sparse Attention
• PR #24880: DeepSeek V4 support for HiSparse direct Prefill-to-Decode DRAM
• Issue #28874: HiSparse Roadmap for Long-Context Sparse Serving