Reduce CP shared KV overhead without changing ownership semantics
The shared-KV path now keeps more CP metadata on-device and reuses physical out-cache locations across MLA and NSA index writes, so each layer avoids repeating logical-to-physical remaps. The in-seq CP all-gather rerange path now delegates to tai-kernel when available and falls back to the existing torch split/cat path with an explicit log. This also extends the Phase8 prefetch machinery to cover shared KV materialization metadata and keeps debug/fallback behavior gated so the fast path is not polluted by diagnostic checks. Constraint: Custom CP kernels must live in tai-kernel and be imported lazily from SGLang Constraint: Decode does not use CP; these changes target NSA prefill CP in-seq-split shared KV Rejected: Recompute physical local cache locations separately for MLA and index writes | repeats the same remap work every layer Rejected: Keep the in-seq rerange Triton code inline in SGLang | duplicates kernel ownership and blocks tai-kernel reuse Confidence: medium Scope-risk: moderate Directive: Keep CP collective ordering identical across ranks; do not add rank-local fallback decisions inside shared KV materialize paths Tested: Remote g0034 container py_compile for modified SGLang/tai-kernel files; remote pytest test/registered/unit/layers/test_nsa_cp_utils.py passed with 24 tests Not-tested: Full multi-node GLM5 prefill/decode throughput after the final commit boundary
This commit is contained in:
laoyao0822
@@ -1,8 +1,8 @@
|
||||
# NSA Prefill CP Phase 8: MLA prefix prefetch
|
||||
# NSA Prefill CP Phase 8: MLA / index prefix prefetch
|
||||
|
||||
Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 chunked prefill / radix cache hit 场景引入 **MLA KV one-layer-ahead prefix prefetch**，把历史 prefix 的 shared KV materialize 从当前层 attention 前的同步阻塞路径里移出，并把等待点延迟到下一层真正消费 prefetched KV 时。
|
||||
Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 chunked prefill / radix cache hit 场景引入 **MLA KV 与 NSA index K/scale 的 one-layer-ahead prefix prefetch**，把历史 prefix 的 shared KV/index materialize 从当前层同步阻塞路径里移出，并把等待点延迟到下一层真正消费 prefetched buffer 时。
|
||||
|
||||
本阶段只做 MLA KV prefix prefetch。暂时不做 index K/scale prefetch，不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量，也不保留 `SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION` 这类 wait 策略开关。
|
||||
本阶段不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量，也不保留 `SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION` 这类 wait 策略开关。`SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1` 现在同时打开 MLA KV prefix prefetch 和 NSA index K/scale prefix prefetch。
|
||||
|
||||
---
|
||||
|
||||
@@ -12,15 +12,16 @@ Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 ch
|
||||
Phase 8 已按 deferred-consume 策略实现：
|
||||
|
||||
```text
|
||||
Layer L index/topk 后启动 Layer L+1 prefix index K/scale prefetch
|
||||
Layer L attention 前启动 Layer L+1 prefix MLA KV prefetch
|
||||
Layer L forward_extend 返回时不等待该 prefetch
|
||||
Layer L+1 consume prefetched KV 时 wait event，再补齐 suffix/current pages
|
||||
Layer L attention 返回时不强制等待 prefetch
|
||||
Layer L+1 consume prefetched buffer 时 wait event，再补齐 suffix/current pages
|
||||
```
|
||||
|
||||
保留的环境变量：
|
||||
|
||||
```text
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 # 生产开关，默认关闭
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 # 生产开关，默认关闭；同时控制 MLA KV 与 index K/scale prefetch
|
||||
SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1 # 调试日志，只打印 probe layer，默认关闭
|
||||
```
|
||||
|
||||
@@ -48,13 +49,13 @@ owner-sharded physical MLA KV on each CP rank
|
||||
-> existing NSA attention kernel
|
||||
```
|
||||
|
||||
在 chunked prefill 的第二个 chunk 及之后，或者 radix cache 命中时，`extend_prefix_len > 0`。这部分 prefix KV 已经在之前的 chunk/request 中写入 persistent KV pool；当前层 attention 仍会同步 materialize 整个可见 KV，包括历史 prefix 和当前 suffix。
|
||||
在 chunked prefill 的第二个 chunk 及之后，或者 radix cache 命中时，`extend_prefix_len > 0`。这部分 prefix KV 与 index K/scale 已经在之前的 chunk/request 中写入 persistent pool；当前层 index/topk 和 attention 仍会同步 materialize 整个可见范围，包括历史 prefix 和当前 suffix。
|
||||
|
||||
Phase 8 的机会是：
|
||||
|
||||
```text
|
||||
历史 prefix 部分已经存在，可以提前为下一层 materialize。
|
||||
当前 suffix 部分必须等下一层 prepare 写入后才能 materialize。
|
||||
当前 suffix 部分必须等下一层 prepare/index write 写入后才能 materialize。
|
||||
```
|
||||
|
||||
因此 Phase 8 不做跨 chunk 的 dense KV cache 复用，而是做 **每次 forward 内、相邻 layer 之间的一层提前预取**。
|
||||
@@ -173,12 +174,12 @@ _get_topk_in_seq_cp_pair(...)
|
||||
-> materialize_shared_paged_buffer(...)
|
||||
```
|
||||
|
||||
Phase 8 暂时不碰 index path。原因：
|
||||
Index K/scale prefetch 在 MLA KV prefetch 之后补齐。当前约束：
|
||||
|
||||
1. index materialize 已经完成 Phase 6 的一次合并；
|
||||
2. indexer 在 MLA attention 之前执行，调度窗口不同；
|
||||
3. 本阶段目标是先验证 MLA prefix prefetch 是否能显著隐藏最大块的 KV materialize；
|
||||
4. 避免同时改变 topk/index 和 attention KV 两条路径，降低正确性风险。
|
||||
3. 当前实现先落地 MLA KV，再补 index K/scale；两者共用同一个 page-aligned slot layout；
|
||||
4. index K/scale prefetch 只覆盖 PAGED topk 路径，RAGGED 仍不进入本阶段。
|
||||
|
||||
---
|
||||
|
||||
@@ -192,9 +193,16 @@ Phase 8 暂时不碰 index path。原因：
|
||||
Layer L attention 计算期间，
|
||||
提前为 Layer L+1 materialize prefix MLA KV。
|
||||
|
||||
Layer L topk/indexer 完成后，
|
||||
提前为 Layer L+1 materialize prefix index K/scale。
|
||||
|
||||
Layer L+1 attention 前，
|
||||
复用已经 materialize 完成的 prefix dense KV，
|
||||
只同步补齐 current/suffix pages。
|
||||
|
||||
Layer L+1 topk/indexer 前，
|
||||
复用已经 materialize 完成的 prefix dense index buffer，
|
||||
只同步补齐 current/suffix pages。
|
||||
```
|
||||
|
||||
### 3.2 性能目标
|
||||
@@ -225,7 +233,6 @@ next layer consume waits only if prefetch has not finished
|
||||
|
||||
Phase 8 不做：
|
||||
|
||||
- 不做 NSA index K/scale prefetch；
|
||||
- 不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND`；
|
||||
- 不做 bandwidth throttle/page budget；
|
||||
- 不做多层 dense KV 常驻缓存；
|
||||
@@ -576,7 +583,7 @@ disable MLA prefetch and fallback sync materialize
|
||||
2. 太早会和当前层 index materialize/topk 抢带宽；
|
||||
3. 太早会增加 NCCL collective 顺序风险。
|
||||
|
||||
推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动：
|
||||
MLA KV 推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动：
|
||||
|
||||
```text
|
||||
1. 当前层需要的 kv_cache/page_table_1 已经准备好；
|
||||
@@ -614,6 +621,19 @@ if prefetcher is not None:
|
||||
return attn_output
|
||||
```
|
||||
|
||||
Index K/scale 的启动点在 `Indexer.forward_cuda(...)` 末尾：
|
||||
|
||||
```python
|
||||
topk_result = run_current_layer_topk(...)
|
||||
index_prefetcher.start_next_layer_prefix(
|
||||
next_layer_id=layer_id + 1,
|
||||
token_to_kv_pool=forward_batch.token_to_kv_pool,
|
||||
)
|
||||
return topk_result
|
||||
```
|
||||
|
||||
当前层 `_maybe_materialize_shared_index_buffer(...)` 优先消费 `cp_shared_kv_index_prefetcher.consume(...)`；miss 时回退原来的 `materialize_shared_paged_buffer(...)`，因此关闭环境变量或 prefetch 条件不满足时语义不变。
|
||||
|
||||
`wait_attention_window()` 仍在 `forward_extend(...)` 的 finally 中调用，但当前实现不做同步等待，只保留 deferred 日志/状态检查入口。真正的同步点在下一层 `consume(...)`：
|
||||
|
||||
```text
|
||||
@@ -790,9 +810,10 @@ Phase 8 的 local copy/remap helper 应尽量复用 Phase 7 的 tai materialize
|
||||
```text
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
|
||||
forward_batch.cp_shared_kv_mla_prefetcher
|
||||
forward_batch.cp_shared_kv_index_prefetcher
|
||||
```
|
||||
|
||||
### Step 2: 拆分 token materialize helper
|
||||
### Step 2: 拆分 token/page materialize helper
|
||||
|
||||
修改：
|
||||
|
||||
@@ -803,6 +824,8 @@ forward_batch.cp_shared_kv_mla_prefetcher
|
||||
```text
|
||||
slot remap builder
|
||||
range local token KV materialize into existing dense buffer
|
||||
range local paged buffer materialize into existing dense buffer
|
||||
logical page -> slot dense page remap
|
||||
sync range all-reduce
|
||||
async all-reduce wrapper
|
||||
```
|
||||
@@ -813,7 +836,7 @@ async all-reduce wrapper
|
||||
full materialize 结果 == prefix materialize + suffix materialize 结果
|
||||
```
|
||||
|
||||
### Step 3: 新增 MLA prefetcher
|
||||
### Step 3: 新增 MLA / index prefetcher
|
||||
|
||||
新增：
|
||||
|
||||
@@ -822,7 +845,7 @@ full materialize 结果 == prefix materialize + suffix materialize 结果
|
||||
职责：
|
||||
|
||||
```text
|
||||
gate
|
||||
common gate
|
||||
allocate dense buffer
|
||||
prefetch prefix range
|
||||
record event
|
||||
@@ -831,7 +854,15 @@ defer wait until consume
|
||||
fallback
|
||||
```
|
||||
|
||||
### Step 4: 接入 `nsa_backend.forward_extend`
|
||||
Index K/scale prefetch 与 MLA prefetch 使用相同 slot-layout page table：
|
||||
|
||||
```text
|
||||
real_page_table.reshape(-1)[slot] -> dense page slot + 1
|
||||
```
|
||||
|
||||
区别是 MLA KV materialize 的单位是 token rows，需要 `page_size` 展开；index K/scale materialize 的单位是 page rows，prefix/suffix range 直接对应 dense page 行。
|
||||
|
||||
### Step 4: 接入 `nsa_backend.forward_extend` / `nsa_indexer`
|
||||
|
||||
修改：
|
||||
|
||||
@@ -840,7 +871,7 @@ fallback
|
||||
接入位置：
|
||||
|
||||
```text
|
||||
shared KV PAGED path 的 MLA materialize 分支。
|
||||
shared KV PAGED path 的 MLA materialize 分支，以及 NSA indexer 的 `_maybe_materialize_shared_index_buffer(...)`。
|
||||
```
|
||||
|
||||
逻辑：
|
||||
@@ -853,6 +884,15 @@ run attention
|
||||
do not wait before return; consume waits before use
|
||||
```
|
||||
|
||||
Index K/scale 的启动点在当前层 topk/indexer 完成后：
|
||||
|
||||
```text
|
||||
Layer L writes current layer index K/scale
|
||||
Layer L materializes/uses current layer index K/scale for topk
|
||||
Layer L starts Layer L+1 prefix index K/scale prefetch
|
||||
Layer L+1 indexer consumes prefetched prefix and fills suffix
|
||||
```
|
||||
|
||||
### Step 5: 单元测试
|
||||
|
||||
修改：
|
||||
@@ -861,12 +901,13 @@ do not wait before return; consume waits before use
|
||||
|
||||
新增覆盖：
|
||||
|
||||
1. prefix/suffix range materialize 拼接结果等价于 full materialize；
|
||||
1. token/paged prefix/suffix range materialize 拼接结果等价于 full materialize；
|
||||
2. prefix page-aligned gate；
|
||||
3. prefix_len=0 不 prefetch；
|
||||
4. non-PAGED / batch_size>1 / debug enabled fallback；
|
||||
5. consume handle layer mismatch fallback；
|
||||
6. started async handle 必须 wait。
|
||||
7. index materialize 优先消费 prefetched buffer，miss 后回退 full materialize；
|
||||
|
||||
### Step 6: 远端集成验证
|
||||
|
||||
@@ -885,7 +926,7 @@ router: g0034
|
||||
3. 重复请求触发 radix hit；
|
||||
4. 检查输出质量；
|
||||
5. 检查没有 collective hang；
|
||||
6. profile 确认 `cp_shared_kv.materialize.token` 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。
|
||||
6. profile 确认 `cp_shared_kv.materialize.token` / index paged materialize 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。
|
||||
|
||||
---
|
||||
|
||||
@@ -972,4 +1013,3 @@ Fallback：
|
||||
debug enabled 时自动回到现有同步 materialize
|
||||
pynccl unavailable 时自动回到现有同步 materialize
|
||||
```
|
||||
|
||||
|
||||
Reference in New Issue
Block a user