diff --git a/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md b/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md new file mode 100644 index 000000000..8f42641a2 --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md @@ -0,0 +1,719 @@ +# NSA Prefill CP Phase 2 设计:Shared/Sharded Persistent KV Pool + +本文档设计 **Phase 2**:在 GLM/DeepSeek-NSA prefill CP 场景下,把当前每个 CP rank 都保存完整 KV cache 的实现,改造成 **CP group 共享逻辑 KV pool、每 rank 只保存自己 shard 的 persistent KV**。 + +Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它暂不解决 attention 真实计算时每个 rank 可能仍需要 materialize full/maxlen KV workspace 的问题;该问题进入 Phase 3。 + +--- + +## 1. 适用范围 + +Phase 2 首版只覆盖以下组合: + +- 模型路径: + - `python/sglang/srt/models/glm4_moe.py` + - `GlmMoeDsaForCausalLM` + - 继承 `python/sglang/srt/models/deepseek_v2.py::DeepseekV2ForCausalLM` +- Attention/KV 类型: + - NSA + MLA + - `NSATokenToKVPool` + - MLA latent KV + NSA indexer K cache 都必须 shard +- 并行: + - prefill 开启 CP + - decode 不开启 CP + - `nsa_prefill_cp_mode=in-seq-split` 是首版验证目标 + - 设计上必须给 `round-robin` CP mode 留扩展接口 +- PD transfer: + - 首版主路径:Mooncake + - decode 侧接收 full KV,decode KV pool 仍是非 CP/full layout +- Page: + - `page_size=64` + - persistent KV shard 以 page 为基本单位 + +明确不在 Phase 2 首版解决: + +- decode 侧 CP +- 非 NSA/MLA 模型 +- 非 Mooncake PD transfer backend 的完整实现 +- Phase 3 的 shard-aware NSA attention/topk/softmax +- 彻底消除运行时 full/maxlen KV materialization + +--- + +## 2. 当前行为与问题 + +### 2.1 `KV size` 的含义 + +当前日志: + +```text +KV Cache is allocated. #tokens: 237312, KV size: 22.14 GB +``` + +表示每个 rank 启动时预分配的 **persistent KV pool 物理空间**,包括: + +- MLA latent KV buffer +- NSA indexer K cache buffer +- page/dummy padding + +它不包含运行时 CP all-gather 临时 tensor、NSA topk/page table、attention workspace 或 hidden gather buffer。 + +### 2.2 当前 CP 下 persistent KV 是 replicated + +当前 prefill CP 路径中,attention 内部会把当前 chunk 的 KV/index key gather/rerange 成全局逻辑顺序,然后写入本 rank 的 KV pool。 + +结果是: + +```text +每个 CP rank 都分配 237312 tokens persistent KV pool +每个 CP rank 都保存完整 logical KV +CP8 总物理 KV 占用约 8 * 22.14GB +但逻辑 KV capacity 仍约 237312 tokens +``` + +也就是说,CP group 里的 KV memory 没有变成更大的逻辑容量,而是被复制消耗掉了。 + +### 2.3 当前 `max_total_num_tokens` 语义混合 + +当前代码中 `max_total_num_tokens` 同时承担多个含义: + +```text +1. 每 rank 物理 KV pool size +2. scheduler/radix 看到的逻辑 token capacity +3. allocator 的 loc 空间大小 +4. req_to_token/radix 存储的 KV loc 范围 +5. PD transfer 中 page index 的空间 +``` + +Phase 2 必须拆开这些语义,否则无法做到: + +```text +每 rank 物理 KV pool 不变 +CP group 逻辑 KV capacity 按 CP size 扩大 +``` + +--- + +## 3. Phase 2 目标 + +### 3.1 核心目标 + +当前: + +```text +physical_max_total_num_tokens_per_rank = 237312 +logical_max_total_num_tokens_cp_group = 237312 +persistent KV layout = replicated +``` + +Phase 2 后: + +```text +physical_max_total_num_tokens_per_rank = 237312 +logical_max_total_num_tokens_cp_group ≈ 237312 * attn_cp_size +persistent KV layout = sharded/shared across CP ranks +``` + +### 3.2 非目标 + +Phase 2 不保证每个 rank 的 runtime workspace 也缩小。Phase 2 可以使用兼容路径: + +```text +persistent KV at rest: sharded +attention runtime view: gather/materialize full logical KV if needed +``` + +这意味着 Phase 2 后仍可能存在运行时 OOM 风险。该风险来自 attention/topk/full-view materialization,不来自 persistent KV pool。Phase 3 再处理。 + +--- + +## 4. 核心设计:逻辑 KV loc 与物理 KV loc 分离 + +### 4.1 定义 + +Phase 2 引入两个 loc 空间: + +```text +logical KV loc/page: + CP group 统一可见,用于 scheduler、radix、req_to_token、PD ordering。 + +physical KV loc/page: + 当前 CP rank 本地 KV buffer 的真实 index,仅 owner rank 可读写。 +``` + +`req_to_token_pool.req_to_token` 和 radix cache 存 **logical loc**。只有在访问本地 KV buffer 时,才把 logical loc 转成 physical loc。 + +### 4.2 page-level shard mapping + +首版采用 page-level interleaved/round-robin owner mapping。注意当前 allocator page 0 是 dummy page,可用 page 从 1 开始,因此 owner 计算必须跳过 dummy page。 + +```python +def owner_cp_rank(logical_page_id: int, cp_size: int) -> int: + # logical_page_id == 0 是 dummy page,不参与真实 ownership + return (logical_page_id - 1) % cp_size + + +def logical_to_physical_page(logical_page_id: int, cp_size: int) -> int: + # 保留本地 physical page 0 作为 dummy page + return (logical_page_id - 1) // cp_size + 1 + + +def logical_to_physical_loc(logical_loc: int, cp_size: int) -> int: + logical_page = logical_loc // page_size + offset = logical_loc % page_size + physical_page = logical_to_physical_page(logical_page, cp_size) + return physical_page * page_size + offset +``` + +这样当每 rank 物理 page 数为 `P` 时,CP group 逻辑 page 数为: + +```text +logical_num_pages = P * cp_size +logical_tokens = logical_num_pages * page_size +``` + +### 4.3 为什么不用 request-contiguous split + +当前 `filter_kv_indices_for_cp_rank()` 按 request total pages 做 contiguous slice: + +```text +rank0: [0, N/8) +rank1: [N/8, 2N/8) +... +``` + +这不适合作为 Phase 2 persistent KV owner mapping,原因: + +1. chunked prefill 中 request 长度逐步增长,按 total pages 切会让 owner 依赖最终长度,不稳定。 +2. `in-seq-split` 的 attention 计算分块方式不应绑定 persistent KV 存储方式。 +3. PD transfer 对非连续 page owner 更自然,应该显式传 `logical_page_positions`。 + +Phase 2 的 persistent KV shard mapping 应独立于 attention CP split mode。 + +--- + +## 5. 关键模块改造 + +### 5.1 memory pool 初始化 + +相关文件: + +- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py` +- `python/sglang/srt/model_executor/model_runner.py` +- `python/sglang/srt/mem_cache/memory_pool.py` +- `python/sglang/srt/mem_cache/allocator.py` + +当前逻辑: + +```python +profiled_tokens = profile_max_num_token(...) +token_capacity = _resolve_token_capacity(profiled_tokens) +self.max_total_num_tokens = token_capacity +NSATokenToKVPool(size=self.max_total_num_tokens) +PagedTokenToKVPoolAllocator(self.max_total_num_tokens) +``` + +Phase 2 需要改成: + +```python +physical_max_total_num_tokens = profiled_tokens_aligned +logical_max_total_num_tokens = physical_max_total_num_tokens * attn_cp_size + +self.physical_max_total_num_tokens = physical_max_total_num_tokens +self.max_total_num_tokens = logical_max_total_num_tokens + +NSATokenToKVPool(size=physical_max_total_num_tokens) +CPSharedPagedTokenToKVPoolAllocator( + logical_size=logical_max_total_num_tokens, + physical_size=physical_max_total_num_tokens, + page_size=page_size, + cp_size=attn_cp_size, + cp_rank=attn_cp_rank, +) +``` + +日志必须区分物理和逻辑: + +```text +KV Cache is allocated. physical #tokens: 237312, KV size: 22.14 GB +CP shared KV enabled. logical #tokens: 1898496, cp_size=8, shard_policy=page_interleaved +``` + +### 5.2 allocator 与 radix + +相关文件: + +- `python/sglang/srt/mem_cache/allocator.py` +- `python/sglang/srt/mem_cache/common.py` +- `python/sglang/srt/mem_cache/radix_cache.py` +- `python/sglang/srt/managers/schedule_batch.py` +- `python/sglang/srt/managers/schedule_policy.py` +- `python/sglang/srt/managers/scheduler.py` + +Phase 2 要求 allocator 对 scheduler 暴露的是 **logical capacity**: + +```python +allocator.available_size() -> logical available tokens +``` + +`alloc_extend()` 返回 logical loc: + +```python +out_cache_loc = logical token locs +req_to_token[req, pos] = logical loc +radix values = logical loc +``` + +本地 KV buffer 不通过 logical loc 直接访问。任何 `set_*_buffer()` / `get_*_buffer()` 前必须经过 layout helper: + +```python +mask = layout.owned_by_this_rank(logical_locs) +physical_locs = layout.logical_to_physical(logical_locs[mask]) +``` + +free/evict 也基于 logical loc: + +```python +allocator.free(logical_locs) +``` + +由于 physical page 是 logical page 的确定性映射,free logical page 后对应 owner rank 的 physical page 自动可复用。 + +### 5.3 KV pool 写入路径 + +相关文件: + +- `python/sglang/srt/models/deepseek_v2.py` +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` +- `python/sglang/srt/layers/attention/nsa_backend.py` +- `python/sglang/srt/mem_cache/memory_pool.py` + +当前 CP prefill 写入语义: + +```text +local KV -> cp_all_gather_rerange_output -> full current chunk KV +每个 rank 都把 full current chunk KV 写入本地 KV pool +``` + +Phase 2 写入语义: + +```text +local KV -> cp_all_gather_rerange_output -> full current chunk KV +每个 rank 只筛选自己 owner 的 logical pages/tokens +只写入本地 physical KV pool +``` + +伪代码: + +```python +logical_locs = forward_batch.out_cache_loc +owned_mask = cp_kv_layout.owned_by_this_rank(logical_locs) +physical_locs = cp_kv_layout.logical_to_physical(logical_locs[owned_mask]) + +# MLA latent KV +forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, + physical_locs, + k_nope[owned_mask], + k_pe[owned_mask], +) + +# NSA indexer K cache +forward_batch.token_to_kv_pool.set_index_k_scale_buffer( + layer_id, + physical_locs, + index_k[owned_mask], + index_k_scale[owned_mask], +) +``` + +必须同时 shard: + +- MLA latent KV +- NSA indexer K cache + +否则 decode 或后续 NSA topk 会出现 latent KV 与 index K 不一致。 + +### 5.4 attention runtime full-view compatibility + +相关文件: + +- `python/sglang/srt/layers/attention/nsa_backend.py` +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` +- `python/sglang/srt/models/deepseek_v2.py` +- `python/sglang/srt/model_executor/forward_batch_info.py` + +Phase 2 暂时允许 attention 计算时 materialize full logical KV view。设计要求是把该行为显式隔离成 compatibility layer,而不是隐式假设 `req_to_token` 是 physical loc。 + +新增 metadata/interface: + +```python +forward_batch.cp_kv_layout +forward_batch.logical_out_cache_loc +forward_batch.runtime_attention_loc +forward_batch.uses_cp_shared_kv +``` + +Phase 2 compatibility 路径: + +```text +1. req_to_token/radix 提供 logical locs +2. runtime helper 按 logical loc 从 owner rank 收集 shard KV/index K +3. 每个 rank 得到当前 attention backend 可消费的 full logical view +4. 现有 NSA/MLA backend 尽量少改,继续消费 full-view runtime loc/table +``` + +这条路径可能需要 maxlen 级别 workspace。Phase 2 接受该风险,并要求日志显式打印: + +```text +CP shared KV runtime full-view materialization enabled. bytes=..., seq_len=..., layer=... +``` + +Phase 3 将把这条 compatibility layer 替换为 shard-aware NSA attention: + +```text +local shard index search -> global topk merge -> owner-aware sparse attention -> distributed reduce +``` + +### 5.5 `in-seq-split` 与 `round-robin` 扩展 + +Phase 2 persistent KV shard mapping 不绑定 attention CP mode。 + +首版验证: + +```text +nsa_prefill_cp_mode=in-seq-split +``` + +原因是当前生产命令使用该模式,且 Phase 1/当前 CP metadata 主要围绕它验证。 + +但 layout helper 必须提供 mode-neutral API: + +```python +class CpSharedKVLayout: + def owned_by_this_rank(self, logical_locs: torch.Tensor) -> torch.Tensor: ... + def logical_to_physical(self, logical_locs: torch.Tensor) -> torch.Tensor: ... + def logical_pages_owned_by_rank(self, logical_pages: np.ndarray, rank: int) -> np.ndarray: ... + def logical_page_positions_for_rank(self, logical_pages: np.ndarray, rank: int) -> np.ndarray: ... +``` + +这样未来 `nsa_prefill_cp_mode=round-robin` 只需要改 attention split/rerange 侧,不需要重写 persistent KV owner mapping 和 PD transfer semantics。 + +--- + +## 6. PD transfer 设计:prefill CP shards -> decode full KV + +### 6.1 目标语义 + +decode 不开 CP,因此 decode 侧仍是 full KV layout。 + +Phase 2 中,一次 PD transfer 的语义是: + +```text +prefill CP rank0 发送自己 owner 的 pages +prefill CP rank1 发送自己 owner 的 pages +... +prefill CP rank7 发送自己 owner 的 pages + +decode rank 接收所有 shards 后,拼成完整 decode KV pool +``` + +因此 shared KV 开启时必须强制: + +```text +SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1 +``` + +如果未开启,启动应报错,而不是静默只从 CP rank0 传 1/8 KV。 + +### 6.2 当前 Mooncake transfer 问题 + +相关文件: + +- `python/sglang/srt/disaggregation/common/conn.py` +- `python/sglang/srt/disaggregation/mooncake/conn.py` +- `python/sglang/srt/disaggregation/prefill.py` +- `python/sglang/srt/disaggregation/decode.py` +- `python/sglang/srt/disaggregation/utils.py` + +当前 Mooncake prefill chunk: + +```python +@dataclasses.dataclass +class TransferKVChunk: + room: int + prefill_kv_indices: np.ndarray + index_slice: slice + is_last_chunk: bool + prefill_aux_index: Optional[int] + state_indices: Optional[List[int]] +``` + +发送时通过: + +```python +chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice] +``` + +这要求当前 chunk 对应 decode dst indices 是连续 slice。 + +Phase 2 page-interleaved owner 下,rank 拥有的 logical pages 是非连续的: + +```text +rank0: logical pages 1, 9, 17, ... +rank1: logical pages 2, 10, 18, ... +``` + +因此不能继续用 `index_slice` 表达 mapping。 + +### 6.3 新 transfer chunk 语义 + +Phase 2 Mooncake chunk 应显式携带 logical page positions: + +```python +@dataclasses.dataclass +class TransferKVChunk: + room: int + prefill_kv_indices: np.ndarray # prefill source physical page ids + logical_page_positions: np.ndarray # absolute positions in the request-level page stream + is_last_chunk: bool + prefill_aux_index: Optional[int] + state_indices: Optional[List[int]] +``` + +`logical_page_positions` 是相对于 request 完整 page 序列的绝对位置,语义等价于旧 `index_slice` 覆盖的坐标系。chunked prefill 中,非首 chunk 必须加上当前 chunk 在 request page stream 中的起始 offset,不能只传 chunk-local positions。 + +发送时: + +```python +dst_page_indices = req.dst_kv_indices[kv_chunk.logical_page_positions] +``` + +而不是: + +```python +dst_page_indices = req.dst_kv_indices[kv_chunk.index_slice] +``` + +### 6.4 prefill 侧过滤 + +`send_kv_chunk()` 当前从 `req_to_token` 取出 logical loc: + +```python +kv_indices = req_to_token[req, start_idx:end_idx] +page_indices = kv_to_page_indices(kv_indices, page_size) +``` + +Phase 2 中 `page_indices` 是 logical page ids。发送前需要转换为本 rank physical page ids: + +```python +logical_pages = kv_to_page_indices(logical_locs, page_size) +chunk_page_start = start_idx // page_size +chunk_local_positions = np.arange(len(logical_pages), dtype=np.int32) +request_page_positions = chunk_page_start + chunk_local_positions +owned = layout.owned_pages_np(logical_pages) + +logical_page_positions = request_page_positions[owned] +src_physical_pages = layout.logical_pages_to_physical_np(logical_pages[owned]) +``` + +Mooncake 发送使用 `src_physical_pages`,decode dst 使用 `logical_page_positions` 选择。 + +### 6.5 NSA state/index cache transfer + +当前 NSA 在 last chunk 会设置: + +```python +state_indices = kv_to_page_indices(kv_indices_full, page_size) +``` + +Phase 2 下 `state_indices` 也必须与 shard ownership 一致。 + +对 NSA: + +```text +MLA latent KV pages 和 NSA index K pages 使用同一 logical page owner +同一个 logical page 的 latent KV 与 index K 必须由同一个 prefill CP rank 发送 +``` + +因此 state transfer 也要使用: + +```python +state_logical_pages -> owned filter -> state_physical_pages +state_logical_page_positions -> decode dst_state_indices selection +``` + +如果 Mooncake extra/state path 暂时仍只支持 slice,则 Phase 2 需要把 state path 同样升级为 explicit positions。 + +### 6.6 decode 侧 preallocation + +decode 不开 CP,decode `req_to_token` 仍可保存 physical/full loc。 + +Decode transfer queue 当前预分配: + +```python +kv_indices = req_to_token[req, :origin_input_len] +page_indices = kv_to_page_indices(kv_indices, page_size) +kv_receiver.init(page_indices, metadata_buffer_index, state_indices) +``` + +Phase 2 保持 decode full KV layout 不变。变化只在 prefill 多 CP rank 发送时,dst page 由 `logical_page_positions` 选择。 + +--- + +## 7. Scheduler / admission 设计 + +### 7.1 scheduler 使用 logical capacity + +Phase 2 中 scheduler 看到的 capacity 应是: + +```text +logical_max_total_num_tokens = physical_max_total_num_tokens_per_rank * attn_cp_size +``` + +涉及: + +- `PrefillAdder.rem_total_tokens` +- `ScheduleBatch.check_decode_mem()` +- metrics `max_total_num_tokens` +- disagg prefill request length check + +这些逻辑应基于 logical available tokens。 + +### 7.2 本地 physical capacity 保护 + +由于 owner mapping 是 deterministic page interleaving,只要 logical allocator 不超过: + +```text +physical_pages_per_rank * cp_size +``` + +每个 rank 的 physical pool 就不会超分。 + +启动时必须把 logical page 数设置为: + +```python +logical_num_pages = physical_num_pages * cp_size +``` + +不要使用无法被 CP size 均匀映射的 logical size。 + +### 7.3 decode 侧仍是瓶颈 + +Phase 2 只扩大 prefill CP group persistent KV capacity。decode 不开 CP,因此 decode KV pool 仍可能成为瓶颈。 + +PD disaggregation 下,如果 prompt KV 长度超过 decode side physical KV pool 能力,decode preallocation 会失败。这是 Phase 2 的已知限制,不应伪装成 prefill capacity 问题。 + +--- + +## 8. 配置与保护 + +建议新增显式开关: + +```text +--enable-nsa-prefill-cp-shared-kv +``` + +启用时必须校验: + +```text +1. enable_nsa_prefill_context_parallel == True +2. attn_cp_size > 1 +3. disaggregation_mode == prefill 时,decode CP disabled +4. nsa_prefill_cp_mode in {in-seq-split, round-robin},首版只允许 in-seq-split 通过运行校验 +5. page_size == 64 +6. token_to_kv_pool 是 NSATokenToKVPool +7. SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER == 1 when PD transfer is enabled +8. Mooncake backend for first implementation +``` + +如果任一条件不满足,启动时报明确错误。 + +建议日志: + +```text +CP shared KV enabled: physical_tokens_per_rank=237312, logical_tokens=1898496, cp_size=8, shard_policy=page_interleaved +CP shared KV PD transfer: all CP ranks transfer enabled, backend=mooncake +CP shared KV attention compatibility: full-view materialization enabled; Phase3 required for shard-aware runtime +``` + +--- + +## 9. 验证计划 + +### 9.1 单元测试 + +需要覆盖: + +1. logical page -> owner rank +2. logical page -> physical page +3. logical loc -> physical loc +4. page 0 dummy 不参与 ownership +5. logical allocator 返回 logical loc +6. free logical loc 后同一 logical page 可复用 +7. PD filtering 输出非连续 `logical_page_positions` +8. NSA state/index pages 与 MLA KV pages owner 一致 + +### 9.2 小规模集成测试 + +建议用 CP2 开始: + +```text +physical_pages_per_rank = small P +logical_pages = 2P +请求跨多个 pages +确认 rank0/rank1 各只写 owner pages +确认 decode 接收后 full KV page 顺序正确 +``` + +### 9.3 生产场景验证指标 + +必须打印并确认: + +```text +per-rank physical KV size 不增加 +logical max_total_num_tokens ≈ physical * cp_size +每 rank written KV pages ≈ total logical pages / cp_size +Mooncake 每个 CP rank 都发送 pages +decode 收到总 pages == request pages +attention compatibility full-view workspace bytes 可观测 +``` + +--- + +## 10. Phase 2 与 Phase 3 边界 + +Phase 2 完成后,系统应具备: + +```text +persistent KV pool: CP-sharded +scheduler/radix: logical capacity expanded +PD transfer: all prefill CP ranks send shards to non-CP decode +attention runtime: may use full-view materialization compatibility path +``` + +Phase 3 继续解决: + +```text +attention runtime 不再 materialize full/maxlen KV +NSA index/topk shard-aware +owner-aware sparse attention +distributed softmax/output reduce +``` + +Phase 2 不应把 Phase 3 的 runtime workspace 问题隐藏起来。相反,Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来,方便 Phase 3 定位和优化。 + +--- + +## 11. 推荐实施顺序 + +1. 新增 `CpSharedKVLayout`,只做映射和测试,不接入主流程。 +2. 新增 logical/physical token capacity 字段和启动日志,默认不开启 shared KV。 +3. 新增 CP shared allocator,返回 logical loc,persistent pool 仍按 physical size 分配。 +4. 改 KV 写入路径:MLA latent KV 和 NSA index K 只写 owner pages。 +5. 接入 Mooncake PD transfer explicit `logical_page_positions`。 +6. 接入 attention runtime full-view compatibility layer。 +7. 开启 scheduler logical capacity。 +8. 加端到端验证和 metrics。 + +这个顺序的原则是:先建立可验证的 loc 语义,再扩大 scheduler capacity。不要先把 admission 放大,否则 attention/PD/KV write 还没接好时会出现难定位的错误。