Document MLA prefix prefetch before Phase 8 implementation
Phase 8 needs a narrow boundary because previous shared-KV phases already changed persistent layout and materialize behavior. The plan records an MLA-only, one-layer-ahead prefix prefetch path that targets chunked-prefill/radix-hit sync materialize overhead without adding index prefetch or bandwidth scheduling yet. Constraint: Phase 8 must not introduce SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND Constraint: Current request is documentation only; do not modify runtime code Rejected: Add index K/scale prefetch to the same phase | widens topk correctness and collective-order risk before MLA overlap is validated Rejected: Persist dense KV across chunks | reintroduces large full-context memory pressure Confidence: high Scope-risk: narrow Directive: Keep Phase 8 v1 MLA-only unless profiling proves index materialize is the next blocker Tested: git diff --check -- docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md Not-tested: Runtime/server execution; documentation-only commit
This commit is contained in:
@@ -0,0 +1,939 @@
|
||||
# NSA Prefill CP Phase 8: MLA 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 前的同步阻塞路径里移出,并尽量与当前层 attention compute 重叠。
|
||||
|
||||
本阶段只做 MLA KV prefix prefetch。暂时不做 index K/scale prefetch,不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量。
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景
|
||||
|
||||
Phase 2-5 已经把 persistent KV/index cache 从“每个 CP rank 保存完整逻辑 KV”改成“每个 CP rank 只保存自己 owner 的 shard”。Phase 6 把 in-seq-split 下 prev/next NSA index materialize 从两次合成了一次。Phase 7 用 tai-kernel/Triton 优化了 materialize local remap/copy 路径,减少 PyTorch tensor op、D2H/MtD sync 和 kernel launch 开销。
|
||||
|
||||
当前 shared KV read compatibility 仍然需要在每层 attention 前恢复 dense full-view:
|
||||
|
||||
```text
|
||||
owner-sharded physical MLA KV on each CP rank
|
||||
-> local materialize: owned pages copied, non-owned pages zero-filled
|
||||
-> CP all-reduce(sum)
|
||||
-> dense full-view MLA KV
|
||||
-> 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。
|
||||
|
||||
Phase 8 的机会是:
|
||||
|
||||
```text
|
||||
历史 prefix 部分已经存在,可以提前为下一层 materialize。
|
||||
当前 suffix 部分必须等下一层 prepare 写入后才能 materialize。
|
||||
```
|
||||
|
||||
因此 Phase 8 不做跨 chunk 的 dense KV cache 复用,而是做 **每次 forward 内、相邻 layer 之间的一层提前预取**。
|
||||
|
||||
---
|
||||
|
||||
## 2. 当前关键代码路径
|
||||
|
||||
### 2.1 Metadata 构造
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa_backend.py`
|
||||
|
||||
`NativeSparseAttnBackend.init_forward_metadata(...)` 构造本 forward 的 page table 和 sequence metadata:
|
||||
|
||||
```text
|
||||
forward_batch.req_to_token_pool.req_to_token
|
||||
-> metadata.page_table_1
|
||||
-> metadata.real_page_table
|
||||
-> metadata.nsa_extend_seq_lens_list
|
||||
```
|
||||
|
||||
相关字段:
|
||||
|
||||
```text
|
||||
forward_batch.extend_prefix_lens_cpu
|
||||
forward_batch.extend_seq_lens_cpu
|
||||
metadata.page_table_1 # token-level logical loc table, page_size=1
|
||||
metadata.real_page_table # page-level logical page table, page_size=token_to_kv_pool.page_size
|
||||
```
|
||||
|
||||
Phase 8 的 prefix prefetch gate 主要依赖 `extend_prefix_lens_cpu[0] > 0` 和 prefix page alignment。
|
||||
|
||||
### 2.2 MLA prepare/write
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
|
||||
- `python/sglang/srt/models/deepseek_v2.py`
|
||||
|
||||
当前 MLA prepare 流程:
|
||||
|
||||
```text
|
||||
forward_absorb_prepare(...)
|
||||
-> compute q/k/v latent
|
||||
-> indexer(...)
|
||||
-> compute k_nope/k_pe
|
||||
-> _maybe_write_cp_shared_local_mla_kv(...)
|
||||
-> rebuild_cp_kv_cache(...)
|
||||
```
|
||||
|
||||
其中:
|
||||
|
||||
```text
|
||||
_maybe_write_cp_shared_local_mla_kv(...)
|
||||
```
|
||||
|
||||
在 Phase 5 direct-write 可用时,把本 rank 计算出的 local MLA KV 直接写入本 rank physical KV pool。
|
||||
|
||||
```text
|
||||
rebuild_cp_kv_cache(...)
|
||||
```
|
||||
|
||||
仍会为了当前 attention 计算路径构造 current full KV view。这一步不是 Phase 8 的目标。
|
||||
|
||||
### 2.3 MLA shared KV materialize
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa_backend.py`
|
||||
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
|
||||
|
||||
当前 `NativeSparseAttnBackend.forward_extend(...)` 中,PAGED topk path 会在 attention kernel 前执行:
|
||||
|
||||
```python
|
||||
kv_cache, page_table_1 = materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=page_table_1,
|
||||
remap_logical_locs=metadata.page_table_1,
|
||||
remap_logical_pages=metadata.real_page_table,
|
||||
layout=forward_batch.cp_shared_kv_layout,
|
||||
page_size=forward_batch.token_to_kv_pool.page_size,
|
||||
)
|
||||
```
|
||||
|
||||
`materialize_shared_token_kv_buffer(...)` 的逻辑:
|
||||
|
||||
```text
|
||||
1. validate/filter logical_locs
|
||||
2. build slot page remap from metadata.real_page_table
|
||||
3. remap topk logical locs -> dense locs
|
||||
4. materialize_local_token_kv_page_slots(...)
|
||||
5. _all_reduce_materialized_buffer(...)
|
||||
6. return dense_kv_cache, dense_locs
|
||||
```
|
||||
|
||||
Phase 8 要拆分的是第 4/5 步:
|
||||
|
||||
```text
|
||||
prefix pages: 可以提前 materialize + async all-reduce
|
||||
suffix/current pages: 必须在当前层 prepare 写完后同步补齐
|
||||
```
|
||||
|
||||
### 2.4 Index materialize
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
|
||||
|
||||
当前 Phase 6 已经把 prev/next index materialize 合成一次:
|
||||
|
||||
```text
|
||||
_get_topk_in_seq_cp_pair(...)
|
||||
-> _maybe_materialize_shared_index_buffer(...)
|
||||
-> materialize_shared_paged_buffer(...)
|
||||
```
|
||||
|
||||
Phase 8 暂时不碰 index path。原因:
|
||||
|
||||
1. index materialize 已经完成 Phase 6 的一次合并;
|
||||
2. indexer 在 MLA attention 之前执行,调度窗口不同;
|
||||
3. 本阶段目标是先验证 MLA prefix prefetch 是否能显著隐藏最大块的 KV materialize;
|
||||
4. 避免同时改变 topk/index 和 attention KV 两条路径,降低正确性风险。
|
||||
|
||||
---
|
||||
|
||||
## 3. Phase 8 目标
|
||||
|
||||
### 3.1 核心目标
|
||||
|
||||
对 `extend_prefix_len > 0` 的 shared KV CP prefill batch:
|
||||
|
||||
```text
|
||||
Layer L attention 计算期间,
|
||||
提前为 Layer L+1 materialize prefix MLA KV。
|
||||
|
||||
Layer L+1 attention 前,
|
||||
复用已经 materialize 完成的 prefix dense KV,
|
||||
只同步补齐 current/suffix pages。
|
||||
```
|
||||
|
||||
### 3.2 性能目标
|
||||
|
||||
当前每层 MLA attention 前都会同步做:
|
||||
|
||||
```text
|
||||
materialize(prefix + current)
|
||||
```
|
||||
|
||||
Phase 8 希望把它改为:
|
||||
|
||||
```text
|
||||
consume(prefetched prefix)
|
||||
materialize(current/suffix only)
|
||||
```
|
||||
|
||||
并把下一层 prefix 的 materialize/all-reduce 放到当前 attention compute 窗口内:
|
||||
|
||||
```text
|
||||
current layer attention compute
|
||||
overlaps with
|
||||
next layer prefix materialize + async CP all-reduce
|
||||
```
|
||||
|
||||
### 3.3 非目标
|
||||
|
||||
Phase 8 不做:
|
||||
|
||||
- 不做 NSA index K/scale prefetch;
|
||||
- 不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND`;
|
||||
- 不做 bandwidth throttle/page budget;
|
||||
- 不做多层 dense KV 常驻缓存;
|
||||
- 不修改 persistent KV layout;
|
||||
- 不修改 PD transfer 协议;
|
||||
- 不修改 radix cache eviction;
|
||||
- 不修改 NSA topk 语义;
|
||||
- 不让 attention kernel 直接读取 owner-sharded KV;
|
||||
- 不支持 `round-robin-split`;
|
||||
- 不支持 RAGGED topk transform path;
|
||||
- 不支持 decode CP。
|
||||
|
||||
---
|
||||
|
||||
## 4. Phase 8 适用条件
|
||||
|
||||
Phase 8 v1 只在保守条件下启用:
|
||||
|
||||
```text
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1
|
||||
forward_batch.uses_cp_shared_kv == True
|
||||
forward_batch.forward_mode.is_context_parallel_extend() == True
|
||||
nsa_prefill_cp_mode == in-seq-split
|
||||
topk_transform_method == PAGED
|
||||
batch_size == 1
|
||||
extend_prefix_len > 0
|
||||
extend_prefix_len % page_size == 0
|
||||
metadata.real_page_table is not None
|
||||
metadata.page_table_1 is not None
|
||||
not CUDA graph capture
|
||||
```
|
||||
|
||||
其中:
|
||||
|
||||
```text
|
||||
extend_prefix_len > 0
|
||||
```
|
||||
|
||||
统一覆盖:
|
||||
|
||||
- chunked prefill 的第二个 chunk 及之后;
|
||||
- radix cache hit;
|
||||
- prefix sharing hit。
|
||||
|
||||
如果 gate 不满足,直接走现有同步 `materialize_shared_token_kv_buffer(...)`,不改变行为。
|
||||
|
||||
---
|
||||
|
||||
## 5. 数据流设计
|
||||
|
||||
### 5.1 当前数据流
|
||||
|
||||
```text
|
||||
Layer L prepare:
|
||||
hidden_states
|
||||
-> MLA q/k/v
|
||||
-> write local MLA KV shard
|
||||
-> rebuild current CP KV
|
||||
|
||||
Layer L attention:
|
||||
materialize full MLA KV(prefix + current)
|
||||
-> local copy/zero
|
||||
-> CP all-reduce
|
||||
attention kernel
|
||||
|
||||
Layer L MLP / EP
|
||||
```
|
||||
|
||||
问题:
|
||||
|
||||
```text
|
||||
materialize full MLA KV 是同步阻塞。
|
||||
chunked/radix-hit 下 prefix 往往占大头。
|
||||
每层都在 attention 前重复等待 prefix materialize。
|
||||
```
|
||||
|
||||
### 5.2 Phase 8 数据流
|
||||
|
||||
```text
|
||||
Layer L attention:
|
||||
consume Layer L prefetched prefix if available
|
||||
materialize Layer L suffix/current pages
|
||||
run Layer L attention
|
||||
|
||||
before Layer L attention kernel:
|
||||
start async prefetch for Layer L+1 prefix MLA KV
|
||||
|
||||
after Layer L attention kernel and before returning:
|
||||
wait Layer L+1 prefetch complete
|
||||
|
||||
Layer L MLP / EP:
|
||||
no outstanding CP shared KV prefetch collective
|
||||
```
|
||||
|
||||
注意这个策略是 **attention-bounded**:
|
||||
|
||||
```text
|
||||
prefetch 只允许与当前 attention compute 重叠。
|
||||
不能跨出 attention backend 返回边界。
|
||||
不能与后面的 prepare_mlp / EP / MoE A2A 竞争不可控带宽。
|
||||
不能改变 CP collective 的跨层顺序。
|
||||
```
|
||||
|
||||
### 5.3 为什么不是“复用上一 chunk dense KV”
|
||||
|
||||
Phase 8 不保存整个上下文 dense KV,也不跨 chunk 复用 dense materialized buffer。
|
||||
|
||||
原因:
|
||||
|
||||
1. dense full-view KV 接近原始上下文大小,跨 chunk 保存会重新引入接近 full KV 的显存压力;
|
||||
2. 每层 dense KV 都不同,保存多层会快速膨胀;
|
||||
3. radix/cache eviction 会让 dense view 的生命周期难以和 persistent logical page 生命周期对齐;
|
||||
4. 当前目标是隐藏同步开销,不是改变 attention kernel 读 layout。
|
||||
|
||||
Phase 8 只允许额外占用 **一层的 transient dense workspace**。
|
||||
|
||||
---
|
||||
|
||||
## 6. Prefix/suffix 切分
|
||||
|
||||
Phase 8 v1 只支持 batch size 1,因此 prefix/suffix 可以按 page table 的前缀切:
|
||||
|
||||
```text
|
||||
page_size = forward_batch.token_to_kv_pool.page_size
|
||||
extend_prefix_len = forward_batch.extend_prefix_lens_cpu[0]
|
||||
prefix_pages = extend_prefix_len // page_size
|
||||
total_pages = metadata.real_page_table.numel()
|
||||
```
|
||||
|
||||
在 `metadata.real_page_table.reshape(-1)` 的 slot layout 下:
|
||||
|
||||
```text
|
||||
slot 0 .. prefix_pages-1:
|
||||
historical prefix pages
|
||||
|
||||
slot prefix_pages .. total_pages-1:
|
||||
current/suffix pages
|
||||
```
|
||||
|
||||
dense token KV layout 保持与现有 slot remap 一致:
|
||||
|
||||
```text
|
||||
dense page 0:
|
||||
dummy/padding page
|
||||
|
||||
dense page i + 1:
|
||||
metadata.real_page_table.reshape(-1)[i]
|
||||
```
|
||||
|
||||
对应 token row:
|
||||
|
||||
```text
|
||||
prefix rows:
|
||||
[page_size, (prefix_pages + 1) * page_size)
|
||||
|
||||
suffix rows:
|
||||
[(prefix_pages + 1) * page_size, (total_pages + 1) * page_size)
|
||||
```
|
||||
|
||||
Phase 8 不能改变 dense page id,否则现有 `page_table_1` / `topk_indices` remap 语义会错。
|
||||
|
||||
---
|
||||
|
||||
## 7. Runtime 设计
|
||||
|
||||
### 7.1 新环境变量
|
||||
|
||||
只新增一个环境变量:
|
||||
|
||||
```text
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1
|
||||
```
|
||||
|
||||
默认关闭。
|
||||
|
||||
不新增:
|
||||
|
||||
```text
|
||||
SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND
|
||||
```
|
||||
|
||||
因为 Phase 8 v1 只有 MLA path,没有 index/index+mla 可选组合。
|
||||
|
||||
### 7.2 ForwardBatch 挂载 prefetcher
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/model_executor/forward_batch_info.py`
|
||||
|
||||
计划新增:
|
||||
|
||||
```python
|
||||
cp_shared_kv_mla_prefetcher: Optional[Any] = None
|
||||
```
|
||||
|
||||
生命周期:
|
||||
|
||||
```text
|
||||
每个 forward batch 初始化一次。
|
||||
只在当前 forward 的 layer loop 内有效。
|
||||
forward 结束后自然释放。
|
||||
```
|
||||
|
||||
### 7.3 新增 prefetch runtime 模块
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
|
||||
|
||||
核心对象:
|
||||
|
||||
```text
|
||||
CpSharedKVMlaPrefetcher
|
||||
CpSharedKVMlaPrefetchHandle
|
||||
```
|
||||
|
||||
建议职责:
|
||||
|
||||
```text
|
||||
CpSharedKVMlaPrefetcher:
|
||||
- 根据 forward_batch/metadata 判断 gate
|
||||
- 管理 prefetch stream
|
||||
- 管理 layer_id -> handle
|
||||
- consume 当前 layer prefix handle
|
||||
- start 下一 layer prefix prefetch
|
||||
- wait attention window 内已启动的 prefetch
|
||||
- fallback 到同步 materialize
|
||||
|
||||
CpSharedKVMlaPrefetchHandle:
|
||||
- layer_id
|
||||
- dense_kv_cache
|
||||
- dense_locs 或 remap 所需 page_inverse
|
||||
- slot_logical_pages
|
||||
- prefix_pages
|
||||
- total_pages
|
||||
- CUDA event
|
||||
- valid/consumed 状态
|
||||
```
|
||||
|
||||
### 7.4 cp_shared_kv_runtime 新 helper
|
||||
|
||||
文件:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
|
||||
|
||||
需要把当前 full materialize 拆成 range-capable helper。
|
||||
|
||||
建议新增:
|
||||
|
||||
```python
|
||||
def build_slot_token_kv_remap(
|
||||
logical_locs: torch.Tensor,
|
||||
remap_logical_pages: torch.Tensor,
|
||||
*,
|
||||
kv_cache: torch.Tensor,
|
||||
layout: CpSharedKVLayout,
|
||||
page_size: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, int]:
|
||||
"""Return slot_logical_pages, dense_locs, logical_page_capacity."""
|
||||
```
|
||||
|
||||
用于复用现有:
|
||||
|
||||
```text
|
||||
build_slot_page_inverse(...)
|
||||
remap_logical_locs_to_slot_dense_locs(...)
|
||||
tai remap helper if available
|
||||
```
|
||||
|
||||
建议新增 range materialize:
|
||||
|
||||
```python
|
||||
def materialize_local_token_kv_page_slots_into(
|
||||
*,
|
||||
kv_cache: torch.Tensor,
|
||||
dense_kv_cache: torch.Tensor,
|
||||
slot_logical_pages: torch.Tensor,
|
||||
layout: CpSharedKVLayout,
|
||||
page_size: int,
|
||||
start_slot: int,
|
||||
end_slot: int,
|
||||
) -> None:
|
||||
"""Materialize slot range [start_slot, end_slot) into existing dense_kv_cache."""
|
||||
```
|
||||
|
||||
异步 all-reduce helper:
|
||||
|
||||
```python
|
||||
def all_reduce_materialized_buffer_async(
|
||||
buffer: torch.Tensor,
|
||||
*,
|
||||
cp_size: int,
|
||||
stream: torch.cuda.Stream,
|
||||
) -> torch.cuda.Event:
|
||||
"""Enqueue CP all-reduce on stream and return a completion event."""
|
||||
```
|
||||
|
||||
同步 suffix helper 可以先复用现有 `_all_reduce_materialized_buffer(...)`,或新增 range sync helper:
|
||||
|
||||
```python
|
||||
def all_reduce_materialized_buffer_range(
|
||||
buffer_slice: torch.Tensor,
|
||||
*,
|
||||
cp_size: int,
|
||||
) -> torch.Tensor:
|
||||
"""Synchronously reduce a contiguous slice of dense_kv_cache."""
|
||||
```
|
||||
|
||||
### 7.5 Async collective 约束
|
||||
|
||||
现有 `_all_reduce_materialized_buffer(...)` 对 fp8/uint8 会走:
|
||||
|
||||
```text
|
||||
torch.distributed.all_reduce(...)
|
||||
```
|
||||
|
||||
这通常不是理想的 async prefetch primitive。
|
||||
|
||||
Phase 8 async prefetch 应优先使用:
|
||||
|
||||
```text
|
||||
get_attention_cp_group().pynccl_comm.all_reduce(comm_buffer, stream=prefetch_stream)
|
||||
```
|
||||
|
||||
如果 pynccl 不可用:
|
||||
|
||||
```text
|
||||
disable MLA prefetch and fallback sync materialize
|
||||
```
|
||||
|
||||
不建议在 v1 中用 blocking `torch.distributed.all_reduce` 假装 prefetch,否则会增加复杂度但不隐藏开销。
|
||||
|
||||
---
|
||||
|
||||
## 8. Attention-bounded 调度
|
||||
|
||||
### 8.1 启动时机
|
||||
|
||||
不要在 model layer loop 开头启动 prefetch。
|
||||
|
||||
原因:
|
||||
|
||||
1. 太早会和当前层 `prepare_attn` / CP hidden gather 抢带宽;
|
||||
2. 太早会和当前层 index materialize/topk 抢带宽;
|
||||
3. 太早会增加 NCCL collective 顺序风险。
|
||||
|
||||
推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动:
|
||||
|
||||
```text
|
||||
1. 当前层需要的 kv_cache/page_table_1 已经准备好;
|
||||
2. 当前层 materialize/reuse 决策已经完成;
|
||||
3. attention kernel 即将开始;
|
||||
4. 此时启动下一层 prefix prefetch;
|
||||
5. 当前 attention kernel 提供 overlap window;
|
||||
6. attention 返回前 wait prefetch event。
|
||||
```
|
||||
|
||||
### 8.2 伪代码
|
||||
|
||||
```python
|
||||
prefetcher = getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None)
|
||||
|
||||
if shared_kv_paged_path:
|
||||
consumed = prefetcher.consume(layer.layer_id, logical_locs=page_table_1)
|
||||
if consumed is not None:
|
||||
kv_cache, page_table_1 = consumed
|
||||
kv_cache, page_table_1 = materialize_suffix_into_prefetched_buffer(...)
|
||||
else:
|
||||
kv_cache, page_table_1 = materialize_shared_token_kv_buffer(...)
|
||||
|
||||
prefetcher.start_next_layer_prefix(
|
||||
next_layer_id=layer.layer_id + 1,
|
||||
metadata=metadata,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
|
||||
attn_output = run_attention_kernel(...)
|
||||
|
||||
if prefetcher is not None:
|
||||
prefetcher.wait_attention_window()
|
||||
|
||||
return attn_output
|
||||
```
|
||||
|
||||
`wait_attention_window()` 必须在 `forward_extend(...)` 返回前调用,保证:
|
||||
|
||||
```text
|
||||
没有 outstanding prefetch collective 跨入 prepare_mlp / EP / 下一层 prepare。
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 9. Correctness 约束
|
||||
|
||||
### 9.1 所有 CP ranks 必须一致参与 collective
|
||||
|
||||
Phase 8 gate 不能依赖 rank-local 条件。所有 CP rank 必须对同一个 layer 做相同顺序的 collective:
|
||||
|
||||
```text
|
||||
rank 0: start prefetch layer L+1
|
||||
rank 1: start prefetch layer L+1
|
||||
...
|
||||
rank N: start prefetch layer L+1
|
||||
```
|
||||
|
||||
如果某个 rank 独立 fallback,而其他 rank 继续 async all-reduce,会造成 collective mismatch/hang。
|
||||
|
||||
因此 fallback 策略:
|
||||
|
||||
```text
|
||||
gate 不满足:所有 rank 不 prefetch
|
||||
prefetch runtime 缺失:所有 rank 不 prefetch
|
||||
unsupported shape:所有 rank 不 prefetch
|
||||
```
|
||||
|
||||
不允许某个 rank 单独跳过已经启动的 collective。
|
||||
|
||||
### 9.2 Prefix 必须 page aligned
|
||||
|
||||
Phase 8 只 materialize prefix 的完整 page:
|
||||
|
||||
```text
|
||||
extend_prefix_len % page_size == 0
|
||||
```
|
||||
|
||||
如果 prefix 不 page-aligned,边界 page 同时包含历史 prefix token 和 current suffix token。这个 page 不能在 current suffix 写入前安全提前 materialize。
|
||||
|
||||
### 9.3 Prefetch 只覆盖 prefix
|
||||
|
||||
对于 Layer L+1:
|
||||
|
||||
```text
|
||||
prefix pages:
|
||||
已经由之前 chunk/request 写入 persistent KV,可以提前读。
|
||||
|
||||
current/suffix pages:
|
||||
必须等 Layer L+1 forward_absorb_prepare 写入本层 persistent KV 后再 materialize。
|
||||
```
|
||||
|
||||
如果错误地提前读取 current/suffix page,会读到旧值或未初始化值,导致模型输出异常。
|
||||
|
||||
### 9.4 Dense slot layout 必须与现有 remap 一致
|
||||
|
||||
prefetched dense KV 必须保持:
|
||||
|
||||
```text
|
||||
dense page 0 = dummy
|
||||
dense page i+1 = slot i in metadata.real_page_table.reshape(-1)
|
||||
```
|
||||
|
||||
这样 `page_table_1` 经 remap 后才能正确索引 dense KV。
|
||||
|
||||
### 9.5 Debug path
|
||||
|
||||
当:
|
||||
|
||||
```text
|
||||
SGLANG_DEBUG_CP_SHARED_KV=1
|
||||
```
|
||||
|
||||
建议禁用 async prefetch,走现有同步 materialize path。
|
||||
|
||||
原因:
|
||||
|
||||
1. debug path 需要同步 tensor summary/checksum;
|
||||
2. 异步 prefetch 会让错误定位更难;
|
||||
3. 现有 debug assert 主要围绕同步 materialize 设计。
|
||||
|
||||
---
|
||||
|
||||
## 10. Failure/fallback 策略
|
||||
|
||||
Phase 8 fallback 必须保守:
|
||||
|
||||
```text
|
||||
prefetch disabled
|
||||
-> existing sync materialize
|
||||
|
||||
prefetch gate failed
|
||||
-> existing sync materialize
|
||||
|
||||
prefetched handle missing
|
||||
-> existing sync materialize
|
||||
|
||||
prefetched handle layer_id mismatch
|
||||
-> existing sync materialize
|
||||
|
||||
prefetched handle shape mismatch
|
||||
-> existing sync materialize
|
||||
|
||||
async pynccl unavailable before any collective starts
|
||||
-> existing sync materialize
|
||||
```
|
||||
|
||||
如果 prefetch collective 已经启动:
|
||||
|
||||
```text
|
||||
必须 wait 完成;
|
||||
不能丢弃未完成 collective;
|
||||
完成后如果校验失败,再 fallback 后续 layer。
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 11. 与 Phase 3/5/7 的关系
|
||||
|
||||
### 11.1 Phase 3 current reuse
|
||||
|
||||
Phase 3 的 current reuse 只适用于:
|
||||
|
||||
```text
|
||||
extend_prefix_len == 0
|
||||
current-only batch
|
||||
```
|
||||
|
||||
Phase 8 只适用于:
|
||||
|
||||
```text
|
||||
extend_prefix_len > 0
|
||||
prefix exists
|
||||
```
|
||||
|
||||
两者互补,不冲突。
|
||||
|
||||
### 11.2 Phase 5 compute-owner direct write
|
||||
|
||||
Phase 8 假设 Phase 5 direct write 已经把 persistent MLA KV 写到 owner physical pool。
|
||||
|
||||
如果 direct write fallback 到 legacy path,Phase 8 仍可以 correctness fallback 到同步 materialize;但性能收益会下降。
|
||||
|
||||
### 11.3 Phase 7 tai materialize
|
||||
|
||||
Phase 8 的 local copy/remap helper 应尽量复用 Phase 7 的 tai materialize kernel。
|
||||
|
||||
但 Phase 8 v1 的核心不是替换 local copy kernel,而是:
|
||||
|
||||
```text
|
||||
把 prefix materialize 的时机提前,并用 async collective overlap。
|
||||
```
|
||||
|
||||
如果 tai path 不可用,Phase 8 可以先复用 torch range materialize;性能收益会变小,但语义仍可验证。
|
||||
|
||||
---
|
||||
|
||||
## 12. 实现计划
|
||||
|
||||
### Step 1: 环境变量和 ForwardBatch 字段
|
||||
|
||||
修改:
|
||||
|
||||
- `python/sglang/srt/environ.py`
|
||||
- `python/sglang/srt/model_executor/forward_batch_info.py`
|
||||
|
||||
新增:
|
||||
|
||||
```text
|
||||
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
|
||||
forward_batch.cp_shared_kv_mla_prefetcher
|
||||
```
|
||||
|
||||
### Step 2: 拆分 token materialize helper
|
||||
|
||||
修改:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
|
||||
|
||||
新增:
|
||||
|
||||
```text
|
||||
slot remap builder
|
||||
range local token KV materialize into existing dense buffer
|
||||
sync range all-reduce
|
||||
async all-reduce wrapper
|
||||
```
|
||||
|
||||
要求:
|
||||
|
||||
```text
|
||||
full materialize 结果 == prefix materialize + suffix materialize 结果
|
||||
```
|
||||
|
||||
### Step 3: 新增 MLA prefetcher
|
||||
|
||||
新增:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
|
||||
|
||||
职责:
|
||||
|
||||
```text
|
||||
gate
|
||||
allocate dense buffer
|
||||
prefetch prefix range
|
||||
record event
|
||||
consume handle
|
||||
wait attention window
|
||||
fallback
|
||||
```
|
||||
|
||||
### Step 4: 接入 `nsa_backend.forward_extend`
|
||||
|
||||
修改:
|
||||
|
||||
- `python/sglang/srt/layers/attention/nsa_backend.py`
|
||||
|
||||
接入位置:
|
||||
|
||||
```text
|
||||
shared KV PAGED path 的 MLA materialize 分支。
|
||||
```
|
||||
|
||||
逻辑:
|
||||
|
||||
```text
|
||||
try consume prefetched prefix for current layer
|
||||
materialize current/suffix into prefetched buffer
|
||||
start next layer prefix prefetch
|
||||
run attention
|
||||
wait prefetch before return
|
||||
```
|
||||
|
||||
### Step 5: 单元测试
|
||||
|
||||
修改:
|
||||
|
||||
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
|
||||
|
||||
新增覆盖:
|
||||
|
||||
1. 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。
|
||||
|
||||
### Step 6: 远端集成验证
|
||||
|
||||
不在本地跑重测试。远端容器验证:
|
||||
|
||||
```text
|
||||
prefill: g0034 container /sgl-workspace/sglang-tai
|
||||
decode: g0035/g0036 container /sgl-workspace/sglang-tai
|
||||
router: g0034
|
||||
```
|
||||
|
||||
验证场景:
|
||||
|
||||
1. `SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0` baseline;
|
||||
2. `SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1` long prompt chunked prefill;
|
||||
3. 重复请求触发 radix hit;
|
||||
4. 检查输出质量;
|
||||
5. 检查没有 collective hang;
|
||||
6. profile 确认 `cp_shared_kv.materialize.token` 同步时间下降或被 attention overlap。
|
||||
|
||||
---
|
||||
|
||||
## 13. 风险
|
||||
|
||||
### 13.1 Collective 顺序风险
|
||||
|
||||
async CP all-reduce 插入层间路径,最大风险是不同 rank collective 顺序不一致。
|
||||
|
||||
缓解:
|
||||
|
||||
```text
|
||||
gate 必须全 rank 一致
|
||||
prefetch 只在 attention backend 内启动
|
||||
forward_extend 返回前 wait
|
||||
不跨入 prepare_mlp / EP
|
||||
```
|
||||
|
||||
### 13.2 带宽竞争
|
||||
|
||||
prefetch 会消耗 NVLink/NCCL 带宽。
|
||||
|
||||
Phase 8 暂时不做 bandwidth 控制,但通过 attention-bounded wait 限制影响范围:
|
||||
|
||||
```text
|
||||
只与 current attention compute overlap
|
||||
不与 EP/MoE A2A overlap
|
||||
不与下一层 prepare_attn overlap
|
||||
```
|
||||
|
||||
如果 profiling 显示 attention 本身也被明显拖慢,再进入后续 Phase 做 bandwidth throttle/page budget。
|
||||
|
||||
### 13.3 显存增加
|
||||
|
||||
Phase 8 会额外持有一层 prefetched dense KV workspace。
|
||||
|
||||
预期额外占用:
|
||||
|
||||
```text
|
||||
one layer dense MLA KV materialize buffer
|
||||
```
|
||||
|
||||
不会保存所有层,也不会保存整个上下文跨 forward 的 dense KV。
|
||||
|
||||
### 13.4 Layer 0 无法隐藏
|
||||
|
||||
第一层没有上一层 attention compute 可以用来隐藏它的 prefix materialize。
|
||||
|
||||
因此 Phase 8 v1 的收益主要来自:
|
||||
|
||||
```text
|
||||
layer 1 .. last_layer
|
||||
```
|
||||
|
||||
Layer 0 仍可能同步 materialize。
|
||||
|
||||
---
|
||||
|
||||
## 14. 验收标准
|
||||
|
||||
Correctness:
|
||||
|
||||
```text
|
||||
prefetch off 与 prefetch on 输出一致或在采样随机性范围内一致
|
||||
long prompt chunked prefill 正常
|
||||
radix cache hit 正常
|
||||
无 collective hang
|
||||
无 device assert
|
||||
```
|
||||
|
||||
Performance:
|
||||
|
||||
```text
|
||||
prefetch on 时,MLA KV prefix materialize 的同步阻塞减少
|
||||
Nsight 中可看到 next-layer prefix materialize/all-reduce 与 current attention 有 overlap
|
||||
总 TTFT / chunked prefill latency 有下降
|
||||
```
|
||||
|
||||
Fallback:
|
||||
|
||||
```text
|
||||
不满足 gate 时自动回到现有同步 materialize
|
||||
debug enabled 时自动回到现有同步 materialize
|
||||
pynccl unavailable 时自动回到现有同步 materialize
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user