docs(cp): design NSA prefill CP shared KV phase2
This commit is contained in:
719
docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md
Normal file
719
docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md
Normal file
@@ -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 还没接好时会出现难定位的错误。
|
||||
Reference in New Issue
Block a user