diff --git a/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md b/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md index f0c16b52e..412e4e66a 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md +++ b/docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md @@ -644,6 +644,15 @@ Phase 1 **不改变** 当前 KV cache 传输格式与 allocator 语义。 ## 12. 后续阶段(本期不做) +注意:以下是 Phase 1 当时的初始拆分。后续 shared KV 工作已经细化为: + +```text +Phase 2: shared/sharded persistent KV pool +Phase 3: current chunk reuse +Phase 4: page-aligned CP token split +Phase 5: compute-owner shared KV layout +``` + ### Phase 2 - prefill shard-KV 保存 - 不再在每层 gather 完整 KV 再写本地 pool 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 index f1bf528f2..3f825d7a6 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md +++ b/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md @@ -26,7 +26,7 @@ Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它 3. **attention runtime 仍走 full-view compatibility**: - 现有 NSA attention/indexer kernel 仍假设可读 full KV/index view。 - shared KV 下运行时会 materialize dense full-view buffer:每 rank 拷贝本地 owner pages,其余填 0,然后 CP group all-reduce 得到完整 runtime view。 - - 这是 Phase 2 的主要性能成本;Phase 3 先去掉 current chunk 的重复 materialize,Phase 4 再做 shard-aware attention/topk。 + - 这是 Phase 2 的主要性能成本;Phase 3 先去掉 current chunk 的重复 materialize,Phase 4 先把 CP token split 对齐到 KV page,Phase 5 再改 compute-owner KV layout。 4. **Mooncake PD transfer 已支持 prefill CP shards -> decode full KV**: - prefill CP rank 发送本 rank owner physical pages。 - transfer chunk 携带 `logical_page_positions`,decode 侧用这些位置选择 full-layout dst pages。 @@ -38,7 +38,8 @@ Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它 当前仍未完成、属于 Phase 3 或后续工作: - Phase 3:复用已经 CP all-gather + rerange 的 current chunk KV/index,避免 shared KV compatibility path 重复 materialize current chunk。详见 `docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md`。 -- Phase 4:history/shared KV 的 selected-page materialize、NSA index/topk/attention 的 shard-aware 计算与 global topk merge。 +- Phase 4:page-aligned CP token split,为 compute-owner KV layout 建立“一个 page 只由一个 CP compute rank 计算”的不变量。详见 `docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md`。 +- Phase 5:compute-owner shared KV layout,让 current KV/index 可以按 CP-local loc 直接写本 rank physical pool。详见 `docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md`。 - `nsa_prefill_cp_mode=round-robin` 的 runtime wiring。 - decode 侧 CP/shared KV。 - Mooncake 之外的 PD transfer backend 完整支持。 @@ -774,13 +775,20 @@ mixed current/history 场景 current 直用,history 暂保留 Phase 2 compatib Phase 4 继续解决: ```text -history/shared KV selected-page materialize -NSA index/topk shard-aware -owner-aware sparse attention -distributed softmax/output reduce +page-aligned CP token split +确保 current KV page 不被多个 CP compute rank 同时切分 +为 compute-owner KV layout 提供 page ownership 不变量 ``` -Phase 2 不应把后续 runtime workspace 问题隐藏起来。相反,Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来,方便 Phase 3/Phase 4 定位和优化。 +Phase 5 继续解决: + +```text +compute-owner logical page allocation +local MLA KV / NSA index K 直接写本 rank physical pool +保持 scheduler/radix/PD transfer 的 logical loc 语义 +``` + +更深层的 history selected-page materialize、shard-aware NSA topk、owner-aware sparse attention 和 distributed softmax/output reduce 在 Phase 5 之后继续拆分。Phase 2 不应把后续 runtime workspace 问题隐藏起来。相反,Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来,方便 Phase 3/4/5 定位和优化。 --- diff --git a/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md b/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md index 091c93aa8..239edab05 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md @@ -2,7 +2,7 @@ 本文档定义 **Phase 3**:在 Phase 2 已经实现 CP shared/sharded persistent KV 的基础上,减少 shared KV compatibility path 引入的 **可避免重复 KV/index materialize**。 -Phase 3 的边界是:**复用当前 forward/chunk 中已经 CP all-gather + rerange 过的 current KV/index**,避免它们再次从 sharded persistent KV/index pool 里 materialize。更深层的 history KV shard-aware topk、selected KV exchange、distributed sparse attention 进入 Phase 4。 +Phase 3 的边界是:**复用当前 forward/chunk 中已经 CP all-gather + rerange 过的 current KV/index**,避免它们再次从 sharded persistent KV/index pool 里 materialize。后续拆分为:Phase 4 做 page-aligned CP token split,Phase 5 做 compute-owner shared KV layout;更深层的 history KV shard-aware topk、selected KV exchange、distributed sparse attention 继续后移。 --- @@ -91,7 +91,7 @@ history 仍走 Phase 2 compatibility materialize ### 1.3 非目标 -以下不属于 Phase 3,进入 Phase 4 或后续: +以下不属于 Phase 3,进入 Phase 4/5 或后续: ```text - shard-aware NSA topk @@ -372,9 +372,12 @@ metadata.page_table_1 ```text SGLANG_CP_SHARED_KV_CURRENT_REUSE=0/1 +SGLANG_CP_SHARED_KV_NVTX=0/1 ``` 建议初期默认关闭,验证稳定后再默认开启。 +`SGLANG_CP_SHARED_KV_NVTX=1` 只打开 shared-KV materialize 相关 NVTX range, +用于 Nsight Systems 定位 Phase2/Phase3 materialize、local copy、all-reduce 开销。 统计项: @@ -483,7 +486,7 @@ history index materialize + current compact index -> topk result maintains logical loc semantics ``` -如果 topk loc semantic 不容易保证,允许保留 Phase 2 fallback,并把 mixed indexer 推迟到 Phase 4 前置任务。 +如果 topk loc semantic 不容易保证,允许保留 Phase 2 fallback,并把 mixed indexer 推迟到 Phase 5 之后的 shard-aware runtime 任务。 --- @@ -578,12 +581,14 @@ SGLANG_CP_SHARED_KV_CURRENT_REUSE=1 ```text SGLANG_DEBUG_CP_SHARED_KV=1 SGLANG_CP_SHARED_KV_CURRENT_REUSE=1 +SGLANG_CP_SHARED_KV_NVTX=1 ``` 性能压测必须关闭: ```text SGLANG_DEBUG_CP_SHARED_KV=0 +SGLANG_CP_SHARED_KV_NVTX=0 ``` --- @@ -610,17 +615,23 @@ MLA current-only + MLA mixed + NSA indexer current-only --- -## 9. Phase 4 边界 +## 9. Phase 4 / Phase 5 边界 -Phase 4 继续解决 history/shared KV runtime 成本: +Phase 4 先解决 CP split 与 KV page 不对齐的问题: ```text -- history selected-page materialize -- shard-aware NSA topk -- local topk + global merge -- selected KV owner exchange -- distributed sparse attention -- owner-routed current write,消除写入前 current chunk all-gather 后再过滤 owner 的浪费 +- page-aligned in-seq-split +- 保证一个 current KV page 只由一个 CP compute rank 计算 +- 不改变 shared KV persistent layout ``` -Phase 3 不隐藏 Phase 4 问题。Phase 3 只保证:**已经 gather 过的 current chunk 不再被 shared KV compatibility path 重复 gather/materialize。** +Phase 5 再解决 persistent KV layout 与 compute owner 不一致的问题: + +```text +- compute-owner logical page allocation +- CP-local out_cache_loc +- local MLA KV / NSA index K direct persistent write +- 保持 materialize compatibility path 与 Mooncake PD transfer 正确 +``` + +history selected-page materialize、shard-aware NSA topk、local topk + global merge、selected KV owner exchange、distributed sparse attention 继续后移。Phase 3 不隐藏这些问题。Phase 3 只保证:**已经 gather 过的 current chunk 不再被 shared KV compatibility path 重复 gather/materialize。** diff --git a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md new file mode 100644 index 000000000..b4492e3d6 --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md @@ -0,0 +1,529 @@ +# NSA Prefill CP Phase 4 计划:Page-Aligned CP Token Split + +本文档定义 **Phase 4**:在不改变 shared KV persistent layout 的前提下,先把 NSA prefill CP 的 `in-seq-split` token 切分改造成 **page-aligned**。Phase 4 的目标是给后续 Phase 5 的 compute-owner KV layout 建立必要不变量: + +```text +同一个 KV page 内的 current-token 只由一个 CP compute rank 负责计算。 +``` + +当前 Phase 2/3 的 shared KV 仍使用 page-interleaved owner mapping。Phase 4 不改变 `CpSharedKVLayout`、allocator、PD transfer、materialize 策略;它只改变 CP compute token split 的边界。 + +--- + +## 0. 当前代码事实 + +### 0.1 NSA CP split 当前是 token 平均切分 + +主要路径: + +- `python/sglang/srt/models/deepseek_v2.py` + - `DeepseekV2ForCausalLM.forward()` 在 `can_cp_split(...)` 后设置 `forward_batch.nsa_cp_metadata`。 + - `DeepseekV2Model.forward()` 在 `nsa_use_prefill_cp(...)` 时对 `hidden_states` 和 `positions` 做 CP split。 +- `python/sglang/srt/models/deepseek_nextn.py` + - nextn 路径也复用同一套 NSA CP metadata/split helper。 +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `prepare_input_dp_with_cp_dsa(...)` + - `cp_split_and_rebuild_data(...)` + - `cp_split_and_rebuild_position(...)` + - `cp_all_gather_rerange_output(...)` + +当前 `in-seq-split` metadata 的核心逻辑: + +```python +cp_segment_num = cp_size * 2 +seq_per_batch = kv_len // cp_segment_num +split_list = seq_per_batch.repeat_interleave(cp_segment_num).int().tolist() +remainder = kv_len % cp_segment_num +if remainder > 0: + split_list[:remainder] = [x + 1 for x in split_list[:remainder]] +``` + +它只保证 token 数量大致均衡,不保证 segment 边界落在 KV page boundary。 + +### 0.2 `in-seq-split` 的 zigzag compute owner + +当前 `in-seq-split` 把序列切成 `2 * cp_size` 个 segment,然后按 zigzag 分配: + +```text +rank r owns: + segment r + segment (2 * cp_size - r - 1) +``` + +例如 `cp_size = 4`: + +```text +original: block0 block1 block2 block3 block4 block5 block6 block7 +rank0 : block0 + block7 +rank1 : block1 + block6 +rank2 : block2 + block5 +rank3 : block3 + block4 +``` + +Phase 4 保留这个 compute owner 规则,只把 `split_list` 从 token-level 改为 page-aware。 + +### 0.3 shared KV persistent layout 当前仍是 page-interleaved + +当前 layout: + +- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py` + +```python +owner = (logical_page - 1) % cp_size +physical_page = (logical_page - 1) // cp_size + 1 +``` + +该 owner 与 CP compute rank 没有关系。Phase 4 不改这里;Phase 5 再让 logical page 分配/owner 与 compute rank 对齐。 + +--- + +## 1. 为什么 Phase 4 要先做 page-aligned split + +Phase 5 想实现: + +```text +local hidden -> local MLA KV / local NSA index K -> 直接写本 rank physical KV/index pool +``` + +这要求一个 page 的 current-token 不能同时由多个 CP rank 计算。否则 page-level shared KV layout 无法判断该 page 应该由哪个 rank 持久化。 + +当前 token-average split 可能产生: + +```text +page N: + token offset 0..31 -> rank A + token offset 32..63 -> rank B +``` + +这会导致后续 compute-owner KV layout 无法只用 page owner 表达真实写入关系。 + +因此 Phase 4 的核心收益不是直接加速,而是建立后续布局改变的前置条件。 + +--- + +## 2. Phase 4 目标 + +### 2.1 核心目标 + +在 `--enable-nsa-prefill-cp-shared-kv` + `--nsa-prefill-cp-mode in-seq-split` 下,构造 page-aligned 的 `split_list`: + +```text +sum(split_list) == extend token count + +对所有内部边界 boundary: + (request_relative_start + boundary) % page_size == 0 +``` + +这里 `request_relative_start` 是当前 extend chunk 在 request KV 序列中的起点,即 `extend_prefix_lens_cpu[0]`。 + +### 2.2 保留的行为 + +Phase 4 不改变: + +- CP all-gather/rerange 语义。 +- `cp_reverse_index`、`reverse_split_len`、`per_rank_actual_token` 的用途。 +- attention/topk kernel 看到的 full-order runtime view。 +- persistent KV owner mapping。 +- PD transfer 的 page filtering。 +- materialize compatibility path。 + +### 2.3 非目标 + +Phase 4 不处理: + +- compute-owner KV layout。 +- local KV direct persistent write。 +- 删除 pre-write allgather。 +- shard-aware attention/topk。 +- `round-robin-split` 的 page-level layout 支持。 +- decode CP。 +- 非 NSA/MLA 路径。 + +--- + +## 3. Page-aligned split 规则 + +### 3.1 MVP 约束 + +Phase 4 第一版建议只支持: + +```text +batch size == 1 +nsa_prefill_cp_mode == in-seq-split +page_size > 1 +extend_prefix_lens_cpu[0] % page_size == 0 +``` + +原因: + +1. 当前 `prepare_input_dp_with_cp_dsa(...)` 注释和实现本身就是 batch=1 语义。 +2. 如果 `extend_prefix_lens_cpu[0]` 不对齐 page,当前 chunk 的第一个 page 是历史 partial page。这个 page 的 owner 需要与前序 chunk 已经写入的 owner 保持一致。该约束更适合 Phase 5 处理。 +3. shared KV 当前主配置 `page_size=64`,chunked prefill 的主 chunk 通常是 page-aligned;非对齐场景可以先 fallback 到旧 token split。 + +### 3.2 Page units + +在 MVP 约束成立时,当前 extend chunk 可以按 page unit 切: + +```python +num_full_pages = extend_len // page_size +tail_tokens = extend_len % page_size +num_units = num_full_pages + (1 if tail_tokens > 0 else 0) +``` + +每个 unit 是不可切分的: + +```text +full page unit: page_size tokens +tail unit : tail_tokens tokens +``` + +segment 边界只能落在 unit 边界。 + +### 3.3 Segment 分配 + +仍然构造 `cp_segment_num = cp_size * 2` 个 segment。 + +推荐算法: + +```python +base = num_units // cp_segment_num +rem = num_units % cp_segment_num + +unit_counts[i] = base + (1 if i < rem else 0) +``` + +再把 unit count 转成 token count: + +```python +split_list = [] +unit_cursor = 0 +for count in unit_counts: + token_count = 0 + for _ in range(count): + if unit_cursor < num_full_pages: + token_count += page_size + else: + token_count += tail_tokens + unit_cursor += 1 + split_list.append(token_count) +``` + +性质: + +```text +sum(split_list) == extend_len +所有非尾部内部边界 page aligned +尾部 partial page 只属于最后一个拿到 tail unit 的 segment +``` + +### 3.4 Too-short fallback + +如果 page 数太少,很多 rank 会拿到 0 token。当前 CP path 对 zero-token rank 的通信/attention/kernel 行为需要谨慎验证。 + +Phase 4 MVP 建议: + +```text +如果 num_units < cp_size: + fallback 到旧 token-average split +``` + +更保守的 gate: + +```text +如果 num_units < 2 * cp_size: + fallback 到旧 token-average split +``` + +推荐先使用保守 gate,跑通后再放宽。 + +--- + +## 4. Metadata 影响面 + +`NSAContextParallelMetadata` 当前字段: + +```python +split_list +max_rank_len +zigzag_index +per_rank_actual_token +reverse_split_len +cp_reverse_index +kv_len_prev / kv_len_next +actual_seq_q_prev / actual_seq_q_next +total_seq_lens +``` + +Phase 4 可以先复用这些字段,不强制加新字段。 + +但为了 Phase 5 衔接,建议增加只读/debug 字段: + +```python +page_aligned: bool +page_size: int +extend_prefix_len: int +segment_page_starts: list[int] +segment_page_ends: list[int] +``` + +这些字段的用途: + +- Debug/assert:确认 split boundary 没切 page。 +- Phase 5:从 segment/page position 推导 compute owner。 +- local `out_cache_loc` split:保证 local loc 对应完整 page run。 + +--- + +## 5. 需要改动的代码位置 + +### 5.1 CP split metadata 构造 + +文件: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + +新增 helper: + +```python +def build_page_aligned_in_seq_split_list( + *, + extend_len: int, + extend_prefix_len: int, + page_size: int, + cp_size: int, +) -> tuple[list[int], PageAlignedSplitInfo | None]: + ... +``` + +`prepare_input_dp_with_cp_dsa(...)` 需要拿到: + +- `page_size` +- `extend_prefix_len` +- 是否启用 shared KV / page-aligned split + +当前函数签名只有: + +```python +prepare_input_dp_with_cp_dsa(kv_len, cp_rank, cp_size, seqs_len) +``` + +因此 Phase 4 需要扩展签名,或直接传入 `forward_batch`。 + +推荐签名: + +```python +prepare_input_dp_with_cp_dsa( + kv_len, + cp_rank, + cp_size, + seqs_len, + *, + forward_batch=None, + page_size=None, +) +``` + +这样对调用点侵入较小。 + +### 5.2 调用点 + +文件: + +- `python/sglang/srt/models/deepseek_v2.py` +- `python/sglang/srt/models/deepseek_nextn.py` + +当前调用: + +```python +forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa( + len(input_ids), + self.cp_rank, + self.cp_size, + forward_batch.seq_lens_cpu.tolist(), +) +``` + +Phase 4 后需要传 `forward_batch` 或 `extend_prefix_lens_cpu`: + +```python +forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa( + len(input_ids), + self.cp_rank, + self.cp_size, + forward_batch.seq_lens_cpu.tolist(), + forward_batch=forward_batch, + page_size=forward_batch.token_to_kv_pool.page_size + if hasattr(forward_batch, "token_to_kv_pool") + else None, +) +``` + +实际代码中 `forward_batch.token_to_kv_pool` 在 `ForwardBatch` 上可用,或从 model runner/server args 获取 `page_size`。 + +### 5.3 Split application + +文件: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + +`cp_split_and_rebuild_data(...)` 和 `cp_split_and_rebuild_position(...)` 可以不变,因为它们只消费 `split_list` 和 `zigzag_index`。 + +### 5.4 Debug/assert + +新增 debug helper: + +```python +def assert_page_aligned_split( + split_list: list[int], + *, + extend_prefix_len: int, + page_size: int, + extend_len: int, +) -> None: + cursor = 0 + for split in split_list[:-1]: + cursor += split + assert (extend_prefix_len + cursor) % page_size == 0 + assert cursor + split_list[-1] == extend_len +``` + +该 assert 可以挂到已有 debug env: + +```text +SGLANG_DEBUG_CP_SHARED_KV=1 +``` + +正常推理路径不要引入 CUDA sync;这里全是 CPU metadata,成本可控。 + +--- + +## 6. Phase 4 数据流 + +### 当前 + +```text +input_ids / hidden_states + | + |-- prepare_input_dp_with_cp_dsa() + | split by token count + | + |-- cp_split_and_rebuild_data() + rank r gets segment r + segment mirror(r) + segment boundary may cut a KV page +``` + +### Phase 4 + +```text +input_ids / hidden_states + | + |-- prepare_input_dp_with_cp_dsa() + | split by request-relative KV pages + | + |-- cp_split_and_rebuild_data() + rank r gets segment r + segment mirror(r) + segment boundary does not cut a KV page +``` + +Persistent KV at rest 仍然是 Phase 2/3 的 shared layout。 + +--- + +## 7. 测试计划 + +### 7.1 Unit tests + +建议新增: + +```text +test/registered/unit/attention/test_nsa_cp_page_aligned_split.py +``` + +覆盖: + +1. `extend_len=32768, page_size=64, cp_size=8` + - `sum(split_list) == 32768` + - 所有内部 cumulative boundary 是 64 的倍数。 + - 每个 rank 两段 token 数接近。 +2. `extend_len=16384, page_size=64, cp_size=8` + - 同上。 +3. `extend_len=1000, page_size=64, cp_size=8` + - tail partial page 不被切分。 +4. `extend_prefix_len=128` + - boundary 检查使用 `prefix + boundary`。 +5. `extend_prefix_len=1` + - MVP fallback 到旧 split,或显式返回 `page_aligned=False`。 +6. too-short case + - `num_units < 2 * cp_size` fallback 或 `page_aligned=False`。 + +### 7.2 Invariant tests + +对生成的 `split_list` 验证: + +```python +cursor = 0 +for segment_len in split_list[:-1]: + cursor += segment_len + assert (extend_prefix_len + cursor) % page_size == 0 +``` + +对 rank ownership 验证: + +```text +rank r local token ranges: + segment r + segment 2*cp_size-r-1 + +每个 local range 是 page-aligned range,除 request tail partial 外不切 page。 +``` + +### 7.3 Runtime smoke + +使用当前 prefill 命令: + +```bash +--enable-nsa-prefill-context-parallel +--nsa-prefill-cp-mode in-seq-split +--enable-nsa-prefill-cp-shared-kv +--page-size 64 +--chunked-prefill-size 16384 +``` + +验证: + +- server 正常启动; +- prefill health 正常; +- router + decode 下短请求可返回; +- debug 日志显示 `page_aligned=True`; +- 与旧 split 对比输出没有明显异常。 + +--- + +## 8. Phase 4 完成标准 + +Phase 4 完成时应满足: + +```text +1. shared-KV + in-seq-split 下可以生成 page-aligned split_list。 +2. split_list 仍兼容现有 cp_split / cp_all_gather_rerange_output。 +3. 非 page-aligned prefix / too-short case 有明确 fallback 或 assert。 +4. 不改变 shared KV persistent layout。 +5. 不改变 PD transfer。 +6. 单测覆盖 split invariant。 +7. 启动测试不回归。 +``` + +--- + +## 9. Phase 4 到 Phase 5 的交付物 + +Phase 4 需要为 Phase 5 提供以下不变量/metadata: + +```text +page_aligned == True 时: + 每个 current logical page 只有一个 CP compute owner。 + owner 可由 segment index -> rank 规则推导。 + local out_cache_loc 可以用同一 split_list 切出来,不会切 page。 +``` + +Phase 5 将基于这个不变量,把 persistent KV page owner 改成 CP compute owner,并让 local KV/index 直接写本 rank physical pool。 diff --git a/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md b/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md new file mode 100644 index 000000000..48b4cd64f --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md @@ -0,0 +1,623 @@ +# NSA Prefill CP Phase 5 计划:Compute-Owner Shared KV Layout + +本文档定义 **Phase 5**:在 Phase 4 已经把 `in-seq-split` 改成 page-aligned split 后,把 shared KV persistent layout 从当前的 page-interleaved owner 改造成 **CP compute-owner layout**。 + +Phase 5 的核心目标: + +```text +local hidden -> local MLA KV / local NSA index K -> 直接写本 rank physical KV/index pool +``` + +避免当前 shared KV 路径中为了 persistent write 先把 current KV/index all-gather 成 full view,再过滤 owner pages 写入。 + +--- + +## 0. 与前序 Phase 的边界 + +### Phase 2/3 当前状态 + +当前 shared KV 已经实现: + +```text +logical KV loc/page: + scheduler / radix / req_to_token / PD transfer 统一可见 + +physical KV loc/page: + 每个 CP rank 本地实际 KV pool + +owner mapping: + owner = (logical_page - 1) % cp_size +``` + +当前 runtime compatibility path 仍可能 materialize full-view KV/index: + +```text +每 rank copy 本地 owner pages,其余填 0 +CP group all-reduce +得到完整 dense view +``` + +### Phase 4 前置条件 + +Phase 5 依赖 Phase 4 的不变量: + +```text +page_aligned == True 时,同一个 current KV page 只由一个 CP compute rank 计算。 +``` + +如果 Phase 4 fallback 到旧 token-average split,Phase 5 的 direct-write path 必须禁用。 + +--- + +## 1. 当前代码事实 + +### 1.1 当前 layout 是 logical page modulo owner + +文件: + +- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py` + +当前规则: + +```python +owner = (logical_page - 1) % cp_size +physical_page = (logical_page - 1) // cp_size + 1 +``` + +这与 CP compute rank 无关。 + +### 1.2 当前 allocator 只扩大 logical capacity + +文件: + +- `python/sglang/srt/mem_cache/allocator.py` +- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py` + +`CPSharedPagedTokenToKVPoolAllocator` 只要求: + +```python +logical_size == physical_size * cp_size +``` + +然后复用 `PagedTokenToKVPoolAllocator` 的 free page 顺序。它不会根据 CP split owner 分配 logical pages。 + +### 1.3 当前 MLA persistent write 依赖 full KV + +文件: + +- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` +- `python/sglang/srt/models/deepseek_v2.py` +- `python/sglang/srt/layers/attention/nsa_backend.py` + +当前 MLA path: + +```text +local hidden + -> local latent_cache / k_nope / k_pe + -> rebuild_cp_kv_cache() + CP all-gather + rerange into full current KV + -> nsa_backend.forward_extend() + _maybe_filter_shared_mla_kv_write() + filter logical pages by current layout owner + write owned physical pages +``` + +`_maybe_filter_shared_mla_kv_write(...)` 当前要求: + +```python +k.shape[0] == cache_loc.numel() +k_rope.shape[0] == cache_loc.numel() +``` + +也就是它只接受 full-order KV,不接受 CP-local KV。 + +### 1.4 当前 NSA index persistent write 也依赖 full key + +文件: + +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + +当前 indexer path: + +```text +local hidden + -> local index key + -> _get_q_k_bf16() + CP all-gather + rerange index key + -> _store_index_k_cache() + _filter_shared_index_write() + write owned physical index K/scale pages +``` + +`_filter_shared_index_write(...)` 当前同样要求: + +```python +key.shape[0] == forward_batch.out_cache_loc.numel() +``` + +--- + +## 2. Phase 5 目标 + +### 2.1 核心目标 + +在 Phase 4 `page_aligned=True` 的 batch 上: + +```text +1. logical page 的 owner 与 CP compute rank 一致。 +2. 本 rank local KV/index key 对应的 local out_cache_loc 全部 owned_by_this_rank。 +3. MLA KV 直接用 local KV + local physical loc 写入 persistent KV pool。 +4. NSA index K/scale 直接用 local key + local physical loc 写入 index pool。 +5. scheduler/radix/req_to_token 仍保存 CP group 统一 logical loc/page。 +6. decode 侧仍接收 full layout。 +``` + +### 2.2 性能目标 + +Phase 5 要去掉的是 **persistent write 前不必要的 current KV/index full gather 依赖**。 + +注意:如果 attention/topk 计算仍需要 full current KV/index,计算路径上的 all-gather 可以暂时保留。Phase 5 的第一步不承诺彻底消除所有 attention runtime communication。 + +更深层目标,例如 shard-aware topk、selected KV exchange、distributed sparse attention,可以放到后续 Phase。 + +### 2.3 非目标 + +Phase 5 不强制实现: + +- `round-robin-split` compute-owner page layout。 +- decode CP。 +- 完全取消 materialize compatibility path。 +- shard-aware distributed attention softmax/output reduce。 +- 非 Mooncake transfer backend 完整优化。 + +--- + +## 3. 推荐设计:allocation-aware modulo owner + +### 3.1 为什么不优先做 explicit owner table + +可选方案 A:新增 owner table。 + +```text +logical_page -> owner_rank +logical_page -> physical_page +``` + +优点: + +- 最灵活。 +- 可表达任意 owner 策略。 + +缺点: + +- materialize、PD transfer、debug、allocator 都要读 owner table。 +- table 生命周期要与 radix/cache eviction/free 同步。 +- 多 rank 必须一致维护 table,错误面更大。 + +### 3.2 推荐方案:保留 modulo layout 公式,但改变 logical page 分配 + +当前 `CpSharedKVLayout` 的公式其实可以继续使用: + +```python +owner = (logical_page - 1) % cp_size +physical_page = (logical_page - 1) // cp_size + 1 +``` + +Phase 5 不一定要把公式改掉,而是让 allocator 在分配 logical pages 时保证: + +```text +需要由 compute rank r 写入的 request page + -> 分配一个满足 (logical_page - 1) % cp_size == r 的 logical_page +``` + +这样: + +```text +logical page id 本身编码 owner +existing logical_to_physical mapping 继续成立 +materialize_shared_* 基本可以继续复用 +PD transfer filter 基本可以继续复用 +``` + +这仍然是 KV layout 改变,因为 logical page 的分配策略从“顺序 page-interleaved”变成“按 CP compute owner lane 分配”。 + +--- + +## 4. 关键不变量 + +Phase 5 direct-write path 只能在以下条件成立时启用: + +```text +forward_batch.uses_cp_shared_kv == True +forward_batch.nsa_cp_metadata.page_aligned == True +nsa_prefill_cp_mode == in-seq-split +batch size == 1 +page_size > 1 +``` + +并且对本 rank: + +```python +local_out_cache_loc = cp_split_out_cache_loc_1d( + forward_batch.out_cache_loc, + forward_batch.nsa_cp_metadata, +) + +assert all(layout.owned_by_this_rank(local_out_cache_loc[local_out_cache_loc > 0])) +``` + +如果 assert 不成立,必须 fallback 到 Phase 2/3 兼容路径,不能静默写错。 + +--- + +## 5. 需要改动的模块 + +### 5.1 CP-aware logical page allocation + +文件: + +- `python/sglang/srt/mem_cache/allocator.py` +- `python/sglang/srt/mem_cache/common.py` +- `python/sglang/srt/managers/schedule_batch.py` + +当前 `alloc_extend(...)` 只从统一 `free_pages` 顺序取新 page。 + +Phase 5 需要新增一条 shared-KV CP-aware allocation path: + +```text +输入: + prefix_len + extend_len + page_size + nsa_cp_metadata.page_aligned segment/page ownership + +输出: + full-order logical out_cache_loc + 其中每个 newly allocated page id 满足 owner(logical_page) == compute_owner(page) +``` + +推荐实现方式: + +```text +allocator 维护 cp_size 条 logical page free lane: + lane r: pages where (page - 1) % cp_size == r + +分配 page 时: + 根据 request page position 找 compute owner rank r + 从 lane r 取下一个 logical page + 写入 full-order out_cache_loc 对应 page 的 token locs +``` + +注意:所有 CP ranks 必须运行相同 deterministic allocation,得到相同 logical `out_cache_loc`。 + +### 5.2 local out_cache_loc split + +文件: + +- `python/sglang/srt/layers/attention/nsa/utils.py` +- `python/sglang/srt/model_executor/forward_batch_info.py` + +当前 `cp_split_and_rebuild_data(...)` 针对 2D hidden,不适合 1D loc tensor。 + +新增 helper: + +```python +def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor) -> torch.Tensor: + if is_nsa_prefill_cp_round_robin_split(): + return nsa_cp_round_robin_split_data(input_) + + input_list = list(torch.split(input_, forward_batch.nsa_cp_metadata.split_list, dim=0)) + return torch.cat( + [input_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index], + dim=0, + ) +``` + +并在 `ForwardBatch` 上缓存: + +```python +forward_batch.cp_local_out_cache_loc +``` + +避免每层重复 split。 + +### 5.3 MLA KV direct write + +文件: + +- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` +- `python/sglang/srt/models/deepseek_v2.py` +- `python/sglang/srt/layers/attention/nsa_backend.py` + +当前 `forward_absorb_prepare(...)` 在 `rebuild_cp_kv_cache(...)` 前有 local `k_nope/k_pe`。 + +Phase 5 需要保留 local KV: + +```text +local_k_nope +local_k_pe +``` + +然后: + +```text +local_k_nope/local_k_pe + cp_local_out_cache_loc + -> layout.logical_locs_to_physical(...) + -> token_to_kv_pool.set_mla_kv_buffer(...) +``` + +如果 attention 计算仍需要 full current KV: + +```text +local KV + -> direct persistent write + -> cp_all_gather_rerange_output for attention runtime +``` + +这样 Phase 5 第一版先解耦 persistent write 和 attention compute。 + +### 5.4 NSA index K/scale direct write + +文件: + +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + +当前 `_get_q_k_bf16(...)` 在 allgather 前已经有 local `key`。 + +Phase 5 需要: + +```text +local key + cp_local_out_cache_loc + -> act_quant / fused store + -> set_index_k_scale_buffer(...) +``` + +如果 topk 计算仍需要 gathered full key,保留计算用 allgather,但不再用 gathered key 做 persistent write。 + +### 5.5 materialize compatibility path + +文件: + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` + +如果采用 allocation-aware modulo owner,现有 materialize path 可以继续使用: + +```python +layout.owned_pages_mask(...) +layout.logical_pages_to_physical(...) +``` + +因为 logical page id 已经按 compute owner lane 分配。 + +Phase 5 只需要增加 debug assertion: + +```text +page_aligned direct-write path 写入后, +materialize_shared_token_kv_buffer(...) 读出的 current KV +与 no-direct-write baseline 对齐。 +``` + +### 5.6 Mooncake PD transfer + +文件: + +- `python/sglang/srt/disaggregation/utils.py` +- `python/sglang/srt/disaggregation/mooncake/conn.py` + +现有 shared KV transfer: + +```text +filter_kv_pages_for_cp_shared_kv(...) + -> layout.filter_owned_pages_np(...) + -> returns physical source pages + request logical positions +``` + +如果 Phase 5 采用 allocation-aware modulo owner,这条路径可以继续工作: + +```text +compute-owner page 被分配到对应 modulo owner lane +layout.filter_owned_pages_np 会选中正确 prefill CP rank 的 pages +logical_page_positions 仍用于 decode dst page selection +``` + +需要补充测试确认: + +```text +prefill source physical pages 与 decode dst pages 按 request positions 对齐。 +``` + +--- + +## 6. Phase 5 数据流 + +### 当前 Phase 2/3 + +```text +local hidden + | + +-- local MLA KV / local index K + | + +-- CP all-gather + rerange current KV/index + | + +-- filter by page-interleaved owner + | + +-- write owned physical pages + | + +-- attention/topk may materialize full-view shared KV/index +``` + +### Phase 5 direct-write + +```text +local hidden + | + +-- local MLA KV / local index K + | + +-- cp_local_out_cache_loc + | pages already allocated from this compute rank owner lane + | + +-- logical_to_physical + | + +-- direct write local physical KV/index pool + | + +-- attention/topk compute path keeps existing gather/materialize as needed +``` + +--- + +## 7. Correctness risks + +### 7.1 Prefix partial page + +如果 current chunk 从 page 中间开始: + +```text +extend_prefix_len % page_size != 0 +``` + +则第一个 current page 已经包含历史 tokens。这个 page 的 owner 必须与历史写入 owner 一致。 + +Phase 5 MVP 建议: + +```text +只有 extend_prefix_len % page_size == 0 时启用 direct-write layout。 +否则 fallback 到 Phase 2/3 path。 +``` + +后续可以支持 partial page continuation,但需要 allocator 保存 page owner continuation 语义。 + +### 7.2 Radix/cache hit + +radix prefix 命中会让 current extend 从已有 logical pages 之后开始。只要 `extend_prefix_len` page-aligned,new pages 可以按 compute owner lane 分配。 + +如果命中到 partial page,MVP fallback。 + +### 7.3 Page owner lane free/evict + +allocator 需要保证: + +```text +free(logical_page) 回到其 modulo owner lane +evict/release 不破坏 lane 去重 +所有 CP ranks 对 logical free state 一致 +``` + +否则不同 rank 可能得到不同 logical loc,导致 req_to_token/radix/PD 全部错位。 + +### 7.4 Current allgather 仍可能存在 + +Phase 5 第一版即使 direct-write 成功,也可能仍保留计算用 allgather: + +```text +local KV -> direct write +local KV -> allgather for attention/topk compute +``` + +这不是 Phase 5 失败。Phase 5 主要先消除 persistent write 对 full gathered KV 的依赖,并建立 compute-owner layout。 + +--- + +## 8. 测试计划 + +### 8.1 Unit tests + +建议新增/修改: + +```text +test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +test/registered/unit/attention/test_nsa_cp_compute_owner_layout.py +``` + +覆盖: + +1. owner-lane page allocation: + - compute owner 为 `r` 的 request pages 分配到 `(logical_page - 1) % cp_size == r`。 +2. `cp_split_and_rebuild_1d(...)`: + - 对 `out_cache_loc` 的 split 顺序与 hidden split 一致。 +3. local loc ownership: + - `layout.owned_by_this_rank(cp_local_out_cache_loc)` 全 True。 +4. MLA direct-write shape: + - local KV token 数等于 local loc token 数。 +5. NSA index direct-write shape: + - local key token 数等于 local loc token 数。 +6. materialize compatibility: + - direct-write 后 `materialize_shared_token_kv_buffer(...)` 读出的 dense KV 与 baseline 一致。 +7. PD filtering: + - `filter_kv_pages_for_cp_shared_kv(...)` 返回的 physical pages 与 request positions 对齐。 + +### 8.2 Runtime smoke + +使用 prefill CP + decode + router: + +```bash +--enable-nsa-prefill-context-parallel +--nsa-prefill-cp-mode in-seq-split +--enable-nsa-prefill-cp-shared-kv +--page-size 64 +--chunked-prefill-size 16384 +``` + +检查: + +- prefill 正常启动; +- decode 正常启动; +- router health 正常; +- 短 prompt 与长 prompt 均能返回正常内容; +- debug 下确认: + - `page_aligned=True` + - `cp_local_out_cache_loc` owned by current rank + - direct-write KV checksum 与 compatibility materialize checksum 对齐。 + +### 8.3 性能验证 + +NVTX 建议区分: + +```text +cp_shared_kv.write.mla.direct +cp_shared_kv.write.index.direct +cp_shared_kv.materialize.token.local_copy +cp_shared_kv.materialize.token.all_reduce +cp_shared_kv.materialize.paged.local_copy +cp_shared_kv.materialize.paged.all_reduce +``` + +Phase 5 预期: + +```text +persistent write 前的 full-KV filter/copy 开销下降 +materialize compatibility 开销不一定消失 +``` + +--- + +## 9. Phase 5 完成标准 + +Phase 5 完成时应满足: + +```text +1. Phase 4 page_aligned batch 可以启用 compute-owner logical page allocation。 +2. CP-local out_cache_loc 与 local KV/index key token 数一致。 +3. MLA KV 可以直接写本 rank physical KV pool。 +4. NSA index K/scale 可以直接写本 rank physical index pool。 +5. req_to_token/radix 仍保存 full logical loc。 +6. materialize compatibility path 仍能正确读出 full-view KV/index。 +7. Mooncake PD transfer 仍能把 prefill CP shards 传给 decode full KV。 +8. 非 page-aligned / unsupported case 有明确 fallback。 +``` + +--- + +## 10. 后续 Phase 候选 + +Phase 5 后,如果仍然慢,下一步应集中在 runtime compute path: + +```text +- selected history page materialize,而不是 full dense materialize; +- local topk + global topk merge; +- owner-to-requester selected KV exchange; +- distributed sparse attention; +- 避免 attention compute 需要 full current KV allgather。 +``` + +这些不建议混入 Phase 5,否则同时改 split、allocator、write、PD、attention kernel,调试面会过大。