# NSA Prefill CP Phase 8: MLA / index prefix prefetch

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 时。

本阶段不引入 `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。

---


## 0. 当前实现状态

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 attention 返回时不强制等待 prefetch
Layer L+1 consume prefetched buffer 时 wait event，再补齐 suffix/current pages
```

后续在 HiCache 分支补充了 batch-scoped slot remap 复用：`metadata.real_page_table`
和 index page table 在同一个 forward 内跨 layer 不变，因此只在第一次 materialize /
prefetch 时构造 `slot_logical_pages`、`page_inverse`、paged `dense_pages`，后续 layer
复用该计划；每层仍重新 materialize 本层 KV/index 数据并执行 CP all-reduce。这个优化
只复用“逻辑页到 dense slot 的映射计划”，不复用 dense KV 内容，避免跨层 KV 数据串用。
如果同一个 `ForwardBatch` 已经带有 remap cache 但无法复用（key mismatch 或 cache
状态不完整），运行时会打印 rate-limited warning；第一次冷启动构造 cache 不打日志。

保留的环境变量：

```text
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，默认关闭
SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-1
                                               # 小 prefix skip threshold；默认 -1 表示使用 cp_size；
                                               # 0 表示不跳过；正数表示绝对 page 数阈值
```

已移除的策略开关：

```text
SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION
```

移除原因：attention 返回前立刻 wait 会把 prefetch 重新变成当前层同步尾延迟，抵消主要 overlap 收益。Phase 8 现在固定使用“下一层消费前等待”的单一策略，减少运行时分支和环境变量读取。

---

## 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 与 index K/scale 已经在之前的 chunk/request 中写入 persistent pool；当前层 index/topk 和 attention 仍会同步 materialize 整个可见范围，包括历史 prefix 和当前 suffix。

Phase 8 的机会是：

```text
历史 prefix 部分已经存在，可以提前为下一层 materialize。
当前 suffix 部分必须等下一层 prepare/index write 写入后才能 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(...)
```

Index K/scale prefetch 在 MLA KV prefetch 之后补齐。当前约束：

1. index materialize 已经完成 Phase 6 的一次合并；
2. indexer 在 MLA attention 之前执行，调度窗口不同；
3. 当前实现先落地 MLA KV，再补 index K/scale；两者共用同一个 page-aligned slot layout；
4. index K/scale prefetch 只覆盖 PAGED topk 路径，RAGGED 仍不进入本阶段。

---

## 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 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 性能目标

当前每层 MLA attention 前都会同步做：

```text
materialize(prefix + current)
```

Phase 8 希望把它改为：

```text
consume(prefetched prefix)
materialize(current/suffix only)
```

并把下一层 prefix 的 materialize/all-reduce 提前到当前层 attention 前启动，允许它持续到下一层消费前完成：

```text
current/next layer compute window
  overlaps with
next layer prefix materialize + async CP all-reduce
next layer consume waits only if prefetch has not finished
```

### 3.3 非目标

Phase 8 不做：

- 不引入 `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
    -> wait event only when the prefetched buffer is actually needed
  materialize Layer L suffix/current pages
  start async prefetch for Layer L+1 prefix MLA KV
  run Layer L attention
  return without waiting for Layer L+1 prefetch

Layer L+1 attention:
  consume Layer L+1 prefetched prefix
    -> wait event here if async prefetch has not completed yet
  materialize Layer L+1 suffix/current pages
```

注意这个策略是 **one-layer deferred-consume**：

```text
prefetch 最多提前一层。
等待点固定在下一层 consume，而不是当前层 attention 返回前。
不跨 forward/batch 保存 dense KV。
不改变 CP collective 的 rank 间顺序；所有 rank 对同一 next layer 同步启动同一类 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
SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-1
```

默认关闭。调试时可额外开启：

```text
SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1
```

`SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES` 同时作用于 MLA KV prefix prefetch 和 NSA index K/scale prefix prefetch。默认 `-1` 使用 `cp_size` 作为阈值，跳过无法覆盖全部 CP lane 的 tiny prefix，避免为收益很小的前缀付出 dense buffer 分配、zero-fill、异步 all-reduce setup 和 event 管理开销。设置为 `0` 可恢复旧行为；设置为正数时按绝对 prefix page 数判断。

不新增/不保留：

```text
SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND
SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION
```

因为 Phase 8 v1 只有 MLA path，没有 index/index+mla 可选组合；wait 策略固定为 deferred-consume，避免 attention 返回前 wait 抵消 overlap。

### 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
  - 记录 deferred wait 状态；真实 wait 在下一层 consume
  - 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. Deferred-consume 调度

### 8.1 启动时机

不要在 model layer loop 开头启动 prefetch。

原因：

1. 太早会和当前层 `prepare_attn` / CP hidden gather 抢带宽；
2. 太早会和当前层 index materialize/topk 抢带宽；
3. 太早会增加 NCCL collective 顺序风险。

MLA KV 推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动：

```text
1. 当前层需要的 kv_cache/page_table_1 已经准备好；
2. 当前层 materialize/reuse 决策已经完成；
3. attention kernel 即将开始；
4. 此时启动下一层 prefix prefetch；
5. 当前 attention 以及后续到下一层 consume 前的计算提供 overlap window；
6. attention 返回前不 wait；下一层 consume prefetched KV 时 wait 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()  # historical name; now only records deferred wait state/log

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
torch.cuda.current_stream().wait_event(prefetch_event)
```

这样避免 attention 返回前 wait 把 prefetch 收益吃掉。

---

## 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)
SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
forward_batch.cp_shared_kv_mla_prefetcher
forward_batch.cp_shared_kv_index_prefetcher
```

### Step 2: 拆分 token/page 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
range local paged buffer materialize into existing dense buffer
logical page -> slot dense page remap
sync range all-reduce
async all-reduce wrapper
```

要求：

```text
full materialize 结果 == prefix materialize + suffix materialize 结果
```

### Step 3: 新增 MLA / index prefetcher

新增：

- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`

职责：

```text
common gate
allocate dense buffer
prefetch prefix range
record event
consume handle
defer wait until consume
fallback
```

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`

修改：

- `python/sglang/srt/layers/attention/nsa_backend.py`

接入位置：

```text
shared KV PAGED path 的 MLA materialize 分支，以及 NSA indexer 的 `_maybe_materialize_shared_index_buffer(...)`。
```

逻辑：

```text
try consume prefetched prefix for current layer
materialize current/suffix into prefetched buffer
start next layer prefix prefetch
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: 单元测试

修改：

- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`

新增覆盖：

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: 远端集成验证

不在本地跑重测试。远端容器验证：

```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` / index paged materialize 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。

---

## 13. 风险

### 13.1 Collective 顺序风险

async CP all-reduce 插入层间路径，最大风险是不同 rank collective 顺序不一致。

缓解：

```text
gate 必须全 rank 一致
prefetch 只在 attention backend 内启动
每个 batch 最多 outstanding 一个 next-layer prefetch
下一层 consume 前强制 wait event
不跨 forward/batch 保存 dense KV
```

### 13.2 带宽竞争

prefetch 会消耗 NVLink/NCCL 带宽。

Phase 8 暂时不做 bandwidth 控制。deferred-consume 策略会扩大 overlap 窗口，但也可能让 prefetch 与 MLP/EP/下一层 prepare 存在带宽竞争：

```text
最多 one-layer-ahead outstanding prefetch
下一层 consume 前必须 wait
不做多层预取
```

如果 profiling 显示带宽竞争明显，再进入后续 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 被提前，并在下一层 consume 前只剩必要 wait
总 TTFT / chunked prefill latency 有下降
```

Fallback：

```text
不满足 gate 时自动回到现有同步 materialize
debug enabled 时自动回到现有同步 materialize
pynccl unavailable 时自动回到现有同步 materialize
```

## 2026-06-01：attention-front prefetch OOM observation

Remote log:

- `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260531_161740.log`
- Failure time: `2026-05-31 16:21:21`
- Stack: target prefill `eagle_worker.forward_target_extend()` -> `DeepseekV2AttentionMLA.forward_prepare()` -> `Indexer._get_topk_in_seq_cp_pair()` -> `_get_topk_ragged_with_cp()` -> `deep_gemm.fp8_mqa_logits()`.
- Error: CUDA OOM while allocating `7.58 GiB`; GPU0 had `6.67 GiB` free, process used `132.84 GiB`.
- Immediately preceding successful batch included very large requests, e.g. `#new-token: 140672`, `#cached-token: 13824`, `#inflight-req: 2`; next failed request prepared backup for `logical_len=178945`.

Confirmed root cause direction:

- The OOM is not a host-cache capacity error; it is a CUDA activation/temp-buffer peak inside current-layer NSA index MQA/topk.
- CP in-seq `_get_topk_ragged_with_cp()` calls `deep_gemm.fp8_mqa_logits()` without the q-dimension chunking guard that the non-CP ragged path already has.
- For `logical_len=178945`, CP0's zigzag tail segment is about `11184` query rows against about `178945` KV rows, so the fp32 logits allocation is about `7.45 GiB`, matching the `7.58 GiB` OOM.
- Do not solve this by limiting max prefetch size. The follow-up fix should add chunked MQA logits handling to the CP in-seq path.

Do not re-debug from scratch before checking this section.
