【工程】kv_offload 迁移至社区 HiCache 初稿

1. 结论

<aside>

当前 kv_offload是一种主动的、per-layer 、Decode 侧重点优化的KV 管理策略（每层 Indexer → miss 判定 → H2D 回填 → Attention），与社区 HiCache 的被动的，按需加载、Prefill 重点优化的路径（prefix 匹配 → 全层预取 start_loading）在设计目标和粒度上完全不同。因此，初步将 kv_offload 特性作为 Decode 侧的一个优化开关引入，且保持社区的代码风格

</aside>

Decode 侧实现细节对照表

实现细节	当前	社区	WIP
Host pool 分配	`NSATokenToKVPoolOffload`：同时分配好 CPU+GPU，KV 在 CPU pinned，index 在 GPU	`NSATokenToKVPoolHost`仅 HiRadixCache prefix 场景使用，搬 kv+indexer，page 对齐，位于 CPU	`NSAKVOnlyTokenToKVPoolHost` ，和 `SparseCacheController` 一一对应，社区layer-first 布局，只保存 KV cache，位于 CPU pinned
GPU Pool 分配	`sparse_kv_buffer`：独立连续 GPU buffer，代替 dense 语义的 `kv_buffer`	无，KV 全在 `token_to_kv_pool` ，按需申请 slots	无独立 `sparse_kv_buffer`；在 `NSATokenToKVPool.kv_buffer` 内按 request 申请连续 sparse slots（`alloc_contiguous`），请求结束后释放
`req_to_token` 语义	存 CPU slot, Attn 不读取	存 GPU slot， Attn 读取	保持 GPU device slot；新增 `req_seq_to_host_slot` 维护 `seq position -> host slot` 映射
`topk_indices` 语义	值即 CPU host slot，直接查 `cpu_slot_to_gpu_slot`	值为序列位置（`0..seq_len-1`），需翻译为 host slot	先 `map_seq_positions_to_host_slots`（seq pos -> host slot），再 `build_gpu_page_table`（host slot -> GPU slot）
controller	自定义 CUDA kernel 直接操作，无 controller 语义	`HiCacheController`：page 粒度，全层循环	`SparseCacheController`：token 粒度，单层 H2D + per-layer D2H；LRU 更新逻辑在 `LRUCacheState`
LRU 状态	GPU 侧 per-layer LRU（`cpu_slot_to_gpu_slot`、`cache_ids` 等）	`RadixCache` 的 CPU 侧 LRU（prefix 复用）	GPU 侧 per-layer 操作
D2H 触发	Decode 阶段每 step per-layer 频率，以 token 为粒度写回 L2	Decode 阶段，整个请求结束后为触发点，按 page 粒度写回 L2/L3	1. 首次 decode 先备份全量 cache 到 host

Decode 阶段每 step 增量 D2H，按 token 粒度、per-layer 执行 | | H2D 触发 | Decode 阶段每层 Indexer topk → miss → per-token H2D | Prefill 阶段 prefix 匹配后全量 load_back | Decode 阶段每层执行 Indexer topk -> miss -> prepare_layer_swap -> load_tokens_h2d_one_layer ，接入了当前的 H2D CUDA OP | | NSA backend 集成 | OffloadArgs、event 同步、per-layer 回填完整闭环 | 无 offload 逻辑 | 接入 lru_state.prepare_layer_swap + build_gpu_page_table，并通过 decode_cache_controller + CUDA event 保证层间同步 | | decode 增量 KV | 每 step per-layer set_kv_cpu_op 写回 CPU | 请求结束时 page 对齐 offload | 每 step 在 scheduler_output_processor 按层调用 transfer_kv_per_layer_mla 方法，不新增OP | | 后处理 | CompactCacheIdsBlock + ReorderOutCacheLocSparse | 无 | 接入 lru_post_process，内部执行 compact + reorder | | server_args | enable_kv_offload | disaggregation_decode_enable_offload_kvcache | enable_decode_nsa_kv_offload |

取舍

新增 SparseCacheController：HiCacheController 粒度是 page 级别（start_loading 对同一批 indices 循环所有 layers），效率低且语义不匹配 per-layer token miss 回填。且阿里最新的 Roadmap 也提到要新增该 Controller，说明是有必要的。

[Roadmap-HiCache]: Unified Storage Framework for SGLang Across Diverse Scenarios · Issue #18239 · sgl-project/sglang

新增 KV-only host pool：社区 NSATokenToKVPoolHost（仅用于 HiRadixCache prefix 场景）同时搬 kv_cache 和 indexer_buffer，且 indexer 搬运要求 indices 按 page_size 对齐——与 per-token sparse 回填冲突；DecodeKVCacheOffloadManager 不感知 NSA，无对应 host pool
新增 LRUCacheState：社区 RadixCache 的 LRU 是 CPU 侧 prefix 复用，无法解决 NSA decode 每层 topk 的 L2 命中问题，引入独立文件，做增量开发
GPU 上需要独立的 sparse buffer：对齐 sparse_kv_buffer，H2D 直接写入连续 buffer，attention 从中读取
不引入 L3 storage，不做 dual attention 优化代入

2. 当前的实现逻辑

数据结构

NSATokenToKVPoolOffload (memory_pool.py)

继承自 NSATokenToKVPool
kv_buffer: List[Tensor]，per-layer，shape [size+page_size, 1, 656]，CPU pinned memory
sparse_kv_buffer: List[Tensor]，per-layer，shape [int(size*cache_ratio)//page_size*page_size, 1, 656]，GPU 连续
index_k_with_scale_buffer: per-layer，始终在 GPU，Indexer 不需等待 host load

ReqToTokenPool 的 LRU 状态（enable_kv_offload=True 时额外分配）