From 96bf7a25949474e083c8bbed76863f91a5db9750 Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Thu, 7 May 2026 03:43:26 +0800
Subject: [PATCH] Reduce CP shared-KV prepare overhead without diagnostic log
 noise

The CP shared-KV path now has a gated tai-kernel replacement for NSA index
K/scale plus MQA range preparation, and Phase8 prefetch can skip tiny prefixes
that do not cover all CP lanes. The Phase9 plan documents the next scheduler
work for overlapping CP communication with peer-request attention windows.

Temporary diagnostic logs added while validating prefetch ownership and fused
index prepare routing were removed before committing so the runtime path does
not add log-only synchronization, log counters, or shape-reporting overhead.

Constraint: Production profiling showed small per-request CPU/GPU overhead from diagnostic logging and sync-prone debug counters.
Rejected: Keep fused-index prepare fallback/used logs behind a new env var | it leaves another runtime branch and logging surface for a path that should be benchmarked with profiler evidence instead.
Rejected: Keep owned page-count prefetch logs | they require sync-prone tensor reductions and were only useful for one-off diagnosis.
Confidence: medium
Scope-risk: moderate
Directive: Reintroduce CP shared-KV diagnostics only behind explicit debug paths, and avoid .item()/shape-heavy logging in hot prefill paths.
Tested: git diff --check for staged sglang-dev changes.
Tested: AST parse for environ.py, cp_shared_kv_prefetch.py, cp_shared_kv_runtime.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py.
Not-tested: Full unit test suite.
Not-tested: Multi-node GLM5 prefill/decode/router runtime after this exact commit.
---
 ...efill_cp_phase7_materialize_triton_plan.md |  38 +
 ...fill_cp_phase8_mla_prefix_prefetch_plan.md |   7 +
 ...fill_cp_phase9_two_request_overlap_plan.md | 654 ++++++++++++++++++
 python/sglang/srt/environ.py                  |   2 +
 .../attention/nsa/cp_shared_kv_prefetch.py    |   7 +
 .../attention/nsa/cp_shared_kv_runtime.py     |  73 ++
 .../srt/layers/attention/nsa/nsa_indexer.py   |  59 +-
 .../mem_cache/test_cp_shared_kv_runtime.py    |  18 +
 8 files changed, 839 insertions(+), 19 deletions(-)
 create mode 100644 docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md

diff --git a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md b/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md
index 0032c0f14..567b6049f 100644
--- a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md
+++ b/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md
@@ -518,6 +518,44 @@ forward_mla.py
 5. 该路径不替代后续 attention 前的 shared KV materialize，只减少 persistent
    write 阶段的 remap/quant/store kernel 碎片。
 
+### P0 fused NSA index MQA prepare 接入
+
+在 shared NSA index K/scale materialize 之后、`deep_gemm.fp8_mqa_logits(...)`
+之前，CP in-seq-split pair path 原来还会执行：
+
+```text
+GetK(index_buffer, block_tables[0], kv_len)
+GetS(index_buffer, block_tables[0], kv_len)
+torch.arange(ke_start, ke_start + valid_q_count)
+torch.zeros_like(ke)
+```
+
+对应的 tai-kernel fast path 已通过实验开关接入：
+
+```text
+SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE=0/1
+```
+
+默认关闭。开启后只影响 `_get_topk_ragged_with_cp(...)` 的单序列
+`cp_index is None` 且 `current_index_kv is None` 路径：
+
+```text
+nsa_indexer.py
+  _get_topk_ragged_with_cp(...)
+    -> try_tai_prepare_cp_mqa_index(...)
+       -> tai_kernel.nsa_prefill.prepare_cp_mqa_kv_and_range(
+              index_buffer, block_tables[0], kv_len,
+              valid_q_count, ke_start, page_size=64, index_head_dim=128)
+    -> fallback GetK + GetS + torch arange/zeros
+```
+
+接入约束：
+
+1. 仅支持 CUDA、`page_size=64`、`index_head_dim=128`、`index_buffer.dtype=uint8`。
+2. 不覆盖 multi-batch `cp_index is not None` path。
+3. 不覆盖 `current_index_kv` reuse path。
+4. 不融合 `fp8_mqa_logits` 本身，也不改变 distributed top-k 语义。
+
 ## Benchmark plan
 
 新增 benchmark 只测本地 materialize，不包含 CP all-reduce：
diff --git a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md
index d9c393187..70218741f 100644
--- a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md
+++ b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md
@@ -31,6 +31,9 @@ prefetch 时构造 `slot_logical_pages`、`page_inverse`、paged `dense_pages`
 ```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 数阈值
 ```
 
 已移除的策略开关：
@@ -411,6 +414,7 @@ Phase 8 不能改变 dense page id，否则现有 `page_table_1` / `topk_indices
 
 ```text
 SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1
+SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-1
 ```
 
 默认关闭。调试时可额外开启：
@@ -419,6 +423,8 @@ SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1
 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
@@ -817,6 +823,7 @@ Phase 8 的 local copy/remap helper 应尽量复用 Phase 7 的 tai materialize
 
 ```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
 ```
diff --git a/docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md b/docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md
new file mode 100644
index 000000000..7e1d681ca
--- /dev/null
+++ b/docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md
@@ -0,0 +1,654 @@
+# NSA Prefill CP Phase9：多请求重叠执行计划
+
+## 目标
+
+Phase9 为 NSA prefill CP + shared KV 增加一条 CP 专用的多请求 overlap 路径。
+目标是在同一个 CP group 内同时保留多个独立 prefill 请求，用一个请求的本地计算、attention 或 MLP/MoE 工作，隐藏另一个请求在 attention 前发生的 shared-KV 通信与 materialize 等待。
+
+Phase9 主要隐藏两类 attention 前窗口：
+
+1. **NSA index K/scale materialize**：top-k 选择前需要用到的 index cache。
+2. **MLA KV materialize**：最终 MQA / NSA attention 前需要用到的 MLA KV cache。
+
+这个阶段不减少总通信字节数，也不改变 shared KV layout。Phase9 是 latency hiding：让一个请求的 CP 通信在另一个请求的有效工作期间发生。
+
+初版可以从 2 个 slot 起步，但设计不能写死为 pair-only。`SGLANG_CP_TBO_MAX_SLOTS > 2` 表示多个独立 request slot，而不是一个 true `ForwardBatch(batch_size=N)`。
+
+## 非目标
+
+- 初版不把多个请求合并成一个 true `ForwardBatch(batch_size=N)`。
+  当前 shared-KV 路径在 `nsa_cp_metadata`、page-aligned split、local out-cache loc ownership、direct-write、materialize remap、top-k 等位置仍有 batch-size-1 假设。
+- 不改变 Phase4/Phase5 已定义的 shared-KV ownership/layout。
+- 不支持 decode CP。decode 当前不开 CP。
+- 初版不支持 `round-robin` CP mode。
+- 不直接复用现有 EP two-batch-overlap 作为 CP 调度器。现有 EP TBO 是拆一个 parent batch；Phase9 调度多个独立 request slot。
+- 初版不支持 CUDA graph capture/replay。
+- 初版不要求 slot 动态 refill。第一版可以 admitted slots 全部跑完后再 admission 下一组；后续再做 finished-slot refill。
+- 初版不重新启用 Phase8 prefetch。Phase9 初版必须先稳定 collective ordering 与 async handle 边界；Phase8 prefetch 的重新接入放到后置子阶段。
+
+## 术语：slot 与 batch size
+
+Phase9 的并发单位是 **slot**。
+
+一个 slot 拥有一个独立的：
+
+- `ScheduleBatch`
+- `ModelWorkerBatch`
+- `ForwardBatch`
+- hidden/residual/positions 状态
+- attention backend metadata
+- CP materialize handles
+
+初版每个 slot 内部仍保持 `batch_size == 1`。
+
+区别：
+
+- `SGLANG_CP_TBO_MAX_SLOTS=4`：同时交错 4 个独立请求流。这是 Phase9 初版支持 `>2` 并发的方式。
+- true `ForwardBatch.batch_size > 1`：一个 slot 内有多个 sequence。这需要 batch-aware NSA CP metadata，是后续阶段。
+
+## 当前单请求数据流
+
+单个 prefill 请求在每层内大致如下：
+
+```text
+hidden local slice
+  -> indexer local q/k/proj/rope/rotate
+  -> CP gather / rerange / materialize index K/scale
+  -> top-k transform
+  -> shared index direct write
+  -> MLA q/k/v prepare + shared MLA direct write
+  -> CP gather / materialize MLA KV
+  -> MQA / NSA attention
+  -> residual / MLP / MoE
+```
+
+当前 profile 中较明显的 bubble 在：
+
+1. index K/scale materialize / all-reduce / all-gather 之后、top-k/MQA 之前；
+2. MLA KV materialize / all-reduce 之后、attention 之前。
+
+Phase8 prefetch 能隐藏单请求 pipeline 内的部分 prefix materialize，但它仍是单请求内部的一层提前预取。Phase9 要利用另一个请求的计算窗口来隐藏这些等待。
+
+## Phase9 目标数据流
+
+Phase9 维护一个固定大小 slot table。每个 slot 是一个独立请求状态机。不同 slot 可以处在不同 layer，但所有 CP rank 的 collective enqueue 顺序必须一致。
+
+Phase9 的调度应当是 **attention-window driven**，不是“slot0 发起 CP comm，然后随便拿 slot1 的 local prepare 来 overlap”。真正需要隐藏的是 attention 前的 CP materialize/gather 等待，所以更合理的窗口是：
+
+```text
+一个 slot 正在跑 MQA / NSA attention
+另一个 slot 在 CP comm stream 上发起 index 或 MLA materialize/gather
+```
+
+也就是说，CP comm 应优先和另一个请求的 attention overlap。local prepare 可以作为铺垫阶段，但不应该作为主要 overlap 目标。原因是 local prepare 通常更碎、更短，并且可能和后续 top-k/vector kernels 抢 compute stream；attention 窗口更稳定，且正好位于下一个请求需要消费 KV 前。
+
+每个 slot 的 layer-local stage：
+
+```text
+LOCAL_INDEX_PREP
+INDEX_COMM_READY_TO_LAUNCH
+INDEX_COMM_LAUNCHED
+INDEX_READY
+TOPK_DONE
+MLA_COMM_READY_TO_LAUNCH
+MLA_COMM_LAUNCHED
+MLA_READY
+ATTENTION_READY
+ATTENTION_RUNNING
+ATTENTION_DONE
+MLP_DONE
+NEXT_LAYER 或 DONE
+```
+
+两个 slot 的目标执行形态应类似下面这样。重点是：slot1 的 CP comm 挂在 slot0 的 attention 窗口下，而不是只和 slot0/slot1 的 local prepare overlap。
+
+```text
+# warmup：先让 slot0 准备到可以跑 attention
+slot0 layer L   local index prepare -> index comm -> top-k -> MLA comm -> MLA ready
+
+# 为 slot1 做不含重 CP wait 的本地铺垫
+slot1 layer M   local index prepare
+slot1 layer M   index/top-k 或 MLA launch 条件准备好
+
+# 主 overlap 窗口
+slot1 layer M   launch index/MLA CP comm on CP comm stream
+slot0 layer L   run MQA / NSA attention on compute stream        overlaps slot1 CP comm
+
+slot1 layer M   wait CP handle only when data is consumed
+slot1 layer M   run MQA / NSA attention                          overlaps another slot's CP comm
+
+slot0 layer L   MLP/MoE -> advance L+1
+slot1 layer M   MLP/MoE -> advance M+1
+```
+
+如果 slot1 的 MLA materialize 依赖 top-k，那么调度器应该尽量把 slot1 的 index/top-k 提前到 peer attention 前完成，使 peer attention 窗口主要隐藏 slot1 的 MLA materialize。若 index 仍未 ready，则可以在 peer attention 中先 launch index comm，但这只是次优路径，因为 index 完成后还要 top-k 才能 launch MLA comm。
+
+因此 Phase9 的优先级应是：
+
+1. **最优**：peer attention 窗口隐藏另一个 slot 的 MLA materialize/gather。
+2. **次优**：peer attention 窗口隐藏另一个 slot 的 index materialize/gather。
+3. **保底**：CP comm 与另一个 slot 的 MLP/dense compute overlap。
+4. **不作为目标**：CP comm 仅与另一个 slot 的 local prepare overlap。
+
+关键不变量：
+
+```text
+所有 CP rank 必须以完全一致的顺序 enqueue collective。
+```
+
+例如所有 rank 都必须看到同样的序列：
+
+```text
+slot0.reqA.index_comm(layer=10)
+slot1.reqB.index_comm(layer=8)
+slot0.reqA.mla_comm(layer=10)
+slot2.reqC.index_comm(layer=3)
+slot1.reqB.mla_comm(layer=8)
+```
+
+不能出现 rank0 enqueue `slot0.index, slot1.index`，rank1 enqueue `slot1.index, slot0.index`。调度决策不能依赖 rank-local CUDA event readiness、rank-local fallback、local token ownership 或本地时间。
+
+## 为什么不能直接复用现有 TBO
+
+现有 SGLang TBO 相关文件：
+
+- `python/sglang/srt/batch_overlap/two_batch_overlap.py`
+- `python/sglang/srt/batch_overlap/operations.py`
+- `python/sglang/srt/layers/attention/tbo_backend.py`
+
+现有 TBO 的语义是：把一个 parent `ForwardBatch` 拆成两个 child batch，用 `execute_overlapped_operations(...)` overlap EP/MoE stages。
+
+Phase9 需要的是：
+
+- 多个独立 `ScheduleBatch` / `ForwardBatch`；
+- 每个请求独立 radix/cache/PD-transfer 状态；
+- 每个请求独立 mutable attention metadata；
+- 跨 rank 一致的 CP collective order；
+- slot 之间允许不同 layer；
+- 不依赖 parent batch token/sequence split。
+
+因此现有 TBO 的 child backend 思路可以借鉴，但 batch splitting 机制不能直接复用。
+
+## 当前代码地图与涉及模块
+
+### 1. Admission 与结果回传
+
+- `python/sglang/srt/disaggregation/prefill.py`
+  - `event_loop_overlap_disagg_prefill()` 是当前 PD prefill overlap loop。
+  - 当前流程：pop bootstrapped requests -> 构造一个 `ScheduleBatch` -> `run_batch(batch)` -> `result_queue.append((batch.copy(), result))` -> 处理上一轮结果。
+  - Phase9 在这里加 gated slot branch。
+  - 不能绕过 PD bootstrap。admitted request 必须已经从 `disagg_prefill_bootstrap_queue` 出来。
+
+- `python/sglang/srt/managers/scheduler.py`
+  - `get_new_batch_prefill()` / `_get_new_batch_prefill_raw()` 管 waiting queue mutation、`PrefillAdder`、chunked-prefill、radix cache、KV allocation。
+  - `run_batch()` 管 model-worker conversion、overlap stream、result copy、scheduler-side cache update。
+  - Phase9 不应手写 request construction。应该加窄的 slot-admission wrapper，复用当前 allocation 路径，并在初版强制每个 slot `max_prefill_bs=1`。
+
+- `python/sglang/srt/managers/schedule_batch.py`
+  - `ScheduleBatch.copy()` 只复制 `process_batch_result()` 需要的字段。
+  - Phase9 每个 slot 完成后都要独立 append `(batch.copy(), result)`。
+  - 不要构造 synthetic combined batch，否则 PD transfer、radix cache 与 request bookkeeping 容易错。
+
+### 2. Worker 与 ForwardBatch 构建
+
+- `python/sglang/srt/managers/tp_worker.py`
+  - `forward_batch_generation()` 当前只处理一个 `ModelWorkerBatch`。
+  - Phase9 需要 sibling entrypoint：`forward_batch_generation_slots(...)`。
+
+- `python/sglang/srt/model_executor/forward_batch_info.py`
+  - `ForwardBatch` 已有 CP shared-KV 字段：
+    - `nsa_cp_metadata`
+    - `uses_cp_shared_kv`
+    - `cp_shared_kv_layout`
+    - `cp_local_out_cache_loc`
+    - `cp_local_physical_out_cache_loc`
+    - `cp_shared_kv_mla_prefetcher`
+    - `cp_shared_kv_index_prefetcher`
+  - Phase9 初版不要让一个 `ForwardBatch` 表示多个请求。最多加 slot identity/state 字段。
+
+- `python/sglang/srt/model_executor/model_runner.py`
+  - `init_attention_backend()` 当前在 `--enable-two-batch-overlap` 时选择现有 EP `TboAttnBackend`。
+  - `forward_extend()` 当前初始化一次 attention metadata，然后调用 monolithic model forward。
+  - Phase9 需要 CP 专用 slot forward path，不要复用 EP TBO flag。
+
+### 3. DeepSeek / GLM MLA 层执行
+
+- `python/sglang/srt/models/deepseek_v2.py`
+  - `DeepseekV2ForCausalLM.forward()` 在 `can_cp_split(...)` 成功时创建 `nsa_cp_metadata`。
+  - `DeepseekV2Model.forward()` 做 embedding、CP token split、逐层 forward，并在 MoE 部分可进入现有 EP TBO。
+  - `DeepseekV2DecoderLayer.forward()` 是 layer 级边界，包含 pre-attn comm、attention、pre-MLP comm、MLP/MoE、postprocess。
+  - Phase9 需要 model-specific resumable slot runner，保存：`hidden_states`、`residual`、`positions`、`zero_allocator`、`layer_idx`。
+
+- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
+  - `forward_absorb_prepare()` 当前包含 local MLA projection、RoPE、indexer/top-k、shared MLA direct write、current-rank CP KV rebuild。
+  - `forward_absorb_core()` 调用 `attn_mqa(...)`，进入 NSA backend，在那里 materialize shared MLA KV。
+  - Phase9 需要在这里拆 stage hook：
+    1. local index/key prepare；
+    2. index materialize launch/wait；
+    3. top-k；
+    4. local MLA KV write；
+    5. MLA materialize launch/wait；
+    6. MQA/attention。
+
+### 4. NSA / shared KV / CP primitives
+
+- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
+  - `_maybe_materialize_shared_index_buffer()` 同步调用 `materialize_shared_paged_buffer(...)`。
+  - `_get_q_k_bf16()` 同步调用 `cp_all_gather_rerange_output(...)`。
+  - `_get_topk_in_seq_cp_pair()` 仍有 single-batch 假设，所以 Phase9 初版必须保持每 slot batch-size-1。
+
+- `python/sglang/srt/layers/attention/nsa_backend.py`
+  - `forward_extend()` 将 top-k 转成 `page_table_1`，然后在 shared KV path 调用 `materialize_shared_token_kv_buffer(...)`。
+  - 这是 Phase9 第二个主要 overlap 目标。
+
+- `python/sglang/srt/layers/attention/nsa/utils.py`
+  - `_cp_attn_tp_all_gather_padded_tensor()` 分配 padded gather output 并 enqueue `cp_all_gather_into_tensor_async(...)`。
+  - `cp_all_gather_rerange_output()` 当前立即 rerange 并返回最终 tensor。
+  - Phase9 需要拆成 launch/finish。
+
+- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
+  - `materialize_shared_token_kv_buffer()` 构造 dense MLA KV 后 all-reduce。
+  - `materialize_shared_paged_buffer()` 构造 dense index K/scale pages 后 all-reduce。
+  - Phase9 需要给两者增加 handle API。
+
+- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py`
+  - Phase8 prefetch 已经有 stream/event handle 思路。
+  - Phase9 初版不能直接让它自主 enqueue CP comm。后置阶段要把 Phase8 prefetch 接入 Phase9 scheduler，由 Phase9 统一管理 collective order。
+
+## 模块边界设计
+
+### 新增模块
+
+#### `python/sglang/srt/batch_overlap/cp_multi_request_overlap.py`
+
+职责：
+
+- CP-TBO env/config snapshot；
+- slot eligibility 与 fallback reason；
+- `CpTboStage`；
+- `CpTboStageClass`；
+- `CpTboRequestState`；
+- `CpTboSlotResult`；
+- rank-consistent deterministic scheduler；
+- EP overlap policy。
+
+建议核心数据结构：
+
+```python
+@dataclass
+class CpTboRequestState:
+    slot_id: int
+    forward_batch: ForwardBatch
+    hidden_states: torch.Tensor
+    residual: torch.Tensor | None
+    positions: torch.Tensor
+    layer_idx: int
+    layer_end: int
+    stage: CpTboStage
+    pending_index_handle: object | None = None
+    pending_mla_handle: object | None = None
+    done: bool = False
+```
+
+#### `python/sglang/srt/layers/attention/cp_tbo_backend.py`
+
+职责：
+
+- `CpTboAttnBackend`；
+- 每个 slot 一个 child backend；
+- slot stage 执行前绑定 `forward_batch.attn_backend`；
+- 隔离 `NsaBackend.forward_metadata`。
+
+### 修改模块
+
+- `python/sglang/srt/disaggregation/prefill.py`
+  - 加 gated slot admission。
+- `python/sglang/srt/managers/scheduler.py`
+  - 加 slot-aware admission 与 `run_batch_slots(...)`。
+- `python/sglang/srt/managers/tp_worker.py`
+  - 加 `forward_batch_generation_slots(...)`。
+- `python/sglang/srt/model_executor/model_runner.py`
+  - 加 CP slot forward entrypoint。
+- `python/sglang/srt/models/deepseek_v2.py`
+  - 加 DeepSeek/GLM MLA slot runner。
+- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
+  - 拆 MLA prepare/core 内部阶段。
+- `python/sglang/srt/layers/attention/nsa/utils.py`
+  - 加 CP all-gather rerange launch/finish。
+- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
+  - 加 shared index / shared MLA materialize launch/finish。
+- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
+  - 支持外部传入 materialized index 或 handle result。
+- `python/sglang/srt/layers/attention/nsa_backend.py`
+  - 支持外部传入 materialized MLA KV 或 handle result。
+- `python/sglang/srt/environ.py`
+  - 加私有 env flag，初期不加 CLI args。
+
+## 启用条件与环境变量
+
+初版只在全部条件满足时启用：
+
+- `SGLANG_CP_TWO_BATCH_OVERLAP=1`；
+- disaggregation mode 是 prefill；
+- `--enable-nsa-prefill-context-parallel`；
+- `--nsa-prefill-cp-mode in-seq-split`；
+- shared KV active；
+- 每个 slot `batch_size == 1`；
+- extend/prefill forward mode，不是 decode/idle/spec；
+- 不走 return-logprob；
+- 不走 hidden-state capture；
+- 不走 CUDA graph；
+- admitted slots 在所有 rank 上 eligibility 一致；
+- token length ratio 不超过阈值；
+- workspace 足够容纳多个 live hidden/residual/materialize buffer。
+
+初始 env：
+
+```text
+SGLANG_CP_TWO_BATCH_OVERLAP=1
+SGLANG_CP_TBO_MAX_SLOTS=2
+SGLANG_CP_TBO_MAX_INFLIGHT_COMM=1
+SGLANG_CP_TBO_DEBUG=1
+SGLANG_CP_TBO_TOKEN_RATIO=2.0
+SGLANG_CP_TBO_EP_OVERLAP_POLICY=avoid_ep_comm
+```
+
+后续可以增加 alias：
+
+```text
+SGLANG_CP_MULTI_REQUEST_OVERLAP=1
+```
+
+## EP overlap 策略
+
+Phase9 可以和 EP work overlap，但不能初版就盲目叠加 CP all-gather 与 DeepEP all-to-all。
+
+分三层：
+
+1. **自然 slot overlap**：slot0 等 CP，slot1 跑 attention/MLP。
+2. **带宽感知 CP/EP scheduling**：stage 分类为 `CP_COMM`、`ATTN_COMPUTE`、`MLP_GEMM`、`EP_COMM`。默认 `avoid_ep_comm`。
+3. **与现有 EP TBO 深度融合**：等 CP-TBO 稳定后再做。
+
+推荐 Phase9 默认实现第 1、2 层，不做第 3 层。
+
+## 子阶段计划
+
+### P9.1：slot framework，顺序执行
+
+目的：验证调度、admission、result plumbing，不改变计算顺序。
+
+任务：
+
+1. 增加 CP-TBO env gates，包括 `SGLANG_CP_TBO_MAX_SLOTS`。
+2. 增加 slot eligibility helper。
+3. 增加 `get_next_disagg_prefill_slots_to_run(max_slots)`。
+4. 增加 `run_batch_slots()`，初版内部顺序调用 `run_batch(slot_i)`。
+5. 每个 slot 独立 `batch.copy()` 并进入 result queue。
+6. 增加 slot admission/fallback debug log。
+7. 增加 eligibility、FIFO/fairness、`max_slots > 2` admission 单测。
+
+预期：无加速；功能上与单请求路径一致。
+
+### P9.2：slot state machine 与 metadata 隔离
+
+目的：让多个 request state 可以安全交错，包括不同 layer 的 slot。
+
+任务：
+
+1. 增加 `CpTboRequestState`、`CpTboStage`、`CpTboStageClass`。
+2. 增加 deterministic slot scheduler。
+3. 增加 CP slot backend wrapper。
+4. 每个 `ForwardBatch` 使用自己的 NSA metadata/backend。
+5. fake backend 单测验证 slot1 不覆盖 slot0 metadata。
+6. backend 隔离不可用时 fallback 到 P9.1 顺序执行。
+
+预期：可以表达 layer-skew，不发生 metadata corruption。
+
+### P9.3：index gather 与 MLA materialize 的 launch/wait handle
+
+目的：把同步通信/materialize 拆成可调度接口。
+
+任务：
+
+1. `cp_all_gather_rerange_output(...)` 拆出 launch/finish。
+2. `materialize_shared_paged_buffer(...)` 拆出 launch/finish。
+3. `materialize_shared_token_kv_buffer(...)` 拆出 launch/finish。
+4. 保留原同步 wrapper：`launch` 后立即 `finish`。
+5. 增加 launch/wait NVTX。
+6. 增加 fake handle 单测验证 collective order。
+
+预期：调度器可以控制通信何时发起、何时等待。
+
+### P9.4：async multi-slot layer executor
+
+目的：实现真正的跨请求 overlap。
+
+任务：
+
+1. 增加 `List[CpTboRequestState]` layer loop。
+2. 支持至少两个 slot，且 slot 可处于不同 `layer_idx`。
+3. 用同一 slot table 支持 `SGLANG_CP_TBO_MAX_SLOTS > 2`。
+4. 每 rank 使用一个 ordered CP comm stream。
+5. 只在真实 consumer 前 wait CUDA event。
+6. Phase9 active 时禁用 Phase8 prefetch。
+7. unsupported condition fallback 到 P9.1。
+
+预期：Nsight 可看到 slotA CP materialize/gather 主要被 slotB MQA/NSA attention 覆盖；attention 不可用时才退到 dense compute/MLP overlap。
+
+### P9.5：EP-aware scheduling
+
+目的：避免 CP comm 与 DeepEP comm 争抢带宽导致反向变慢。
+
+任务：
+
+1. 标记 stage class：`CP_COMM`、`ATTN_COMPUTE`、`MLP_GEMM`、`EP_COMM`。
+2. 默认策略 `SGLANG_CP_TBO_EP_OVERLAP_POLICY=avoid_ep_comm`。
+3. `avoid_ep_comm` 下，优先在另一个 slot 处于 attention/dense compute 时发 CP comm。
+4. 避免在 DeepEP dispatch/combine 期间启动新的 heavy CP comm。
+5. 增加 debug override policy 供 profile。
+6. 增加 CP-TBO stage decision / DeepEP overlap NVTX。
+
+预期：减少 pre-attn bubble，同时不引入更大的 fabric contention bubble。
+
+### P9.6：profile、调参、运行稳定化
+
+目的：让功能具备实际部署价值。
+
+任务：
+
+1. 增加 `SGLANG_CP_TBO_DEBUG=1` 下简洁日志。
+2. 增加 NVTX：
+   - `cp_tbo.slotN.index.launch/wait`
+   - `cp_tbo.slotN.mla.launch/wait`
+   - `cp_tbo.slotN.attn`
+   - `cp_tbo.slotN.mlp`
+   - `cp_tbo.scheduler.decision`
+3. 对比 CP-TBO off/on。
+4. 分别测试 2、3、4 个 concurrent long prompts。
+5. 调整 token-ratio、slot count、in-flight comm、EP policy。
+
+预期：端到端 concurrent long-prompt throughput 比 Phase8 单请求 pipeline 更好。
+
+### P9.7：重新接入 Phase8 prefetch
+
+目的：在 Phase9 初版稳定后，让单请求 layer-prefetch 和多请求 overlap 可以共存。
+
+这个阶段不在初版实现中展开具体代码位置。原因是 Phase9 初版完成前，最终的 stage boundary、handle ownership、CP comm stream、collective ordering API 仍可能调整。现在只定义接入原则和验收条件。
+
+接入原则：
+
+1. Phase8 prefetch 不能自主 enqueue 可能改变 collective order 的 CP 通信。
+2. Phase8 prefetch 必须作为 Phase9 scheduler 管理下的一个 stage/handle。
+3. prefetch 的 launch order 必须进入全 rank 一致的 CP collective schedule。
+4. prefetch 只允许在 Phase9 明确判定安全的 bandwidth window 发起。
+5. prefetch hit/miss 不允许造成不同 rank 走不同 collective 序列。
+6. Phase8 prefetch 与 CP-TBO 共同启用时，fallback 必须清晰打日志。
+
+后续设计方向：
+
+```text
+Phase8 current:
+  layer L attention 后，自主 prefetch layer L+1 prefix
+
+Phase9-compatible:
+  CP-TBO scheduler 决定是否为 slotN layer L+1 发起 prefix prefetch
+  prefetch handle 进入 slot state
+  slotN 到达 layer L+1 时 consume scheduler-owned handle
+```
+
+新增 env 建议后续再定名，初步倾向：
+
+```text
+SGLANG_CP_TBO_ENABLE_LAYER_PREFETCH=1
+```
+
+验收条件：
+
+- CP-TBO off + Phase8 prefetch 行为不变；
+- CP-TBO on + prefetch off 行为不变；
+- CP-TBO on + prefetch on 时 collective order 全 rank 一致；
+- Nsight 显示 prefetch 没有和 CP-TBO 主路径重复 materialize 同一窗口；
+- 没有 fallback spam、没有 hang、输出不劣化。
+
+### P9.8：true `ForwardBatch.batch_size > 1` 后续扩展
+
+目的：在 slot scheduler 稳定后，再支持一个 slot 内多个 sequence。
+
+前置条件：
+
+- batch-aware `NSAContextParallelMetadata`；
+- batch-aware page-aligned split；
+- batch-aware local out-cache loc ownership；
+- top-k CP pair 逻辑支持 multi-batch；
+- shared KV materialize remap 支持 per-request mapping。
+
+这不是 Phase9 初版目标。
+
+## 正确性风险
+
+### CP collective order mismatch
+
+如果不同 rank enqueue 顺序不同，可能 hang 或产生错误结果。所有 admission、stage scheduling、fallback 都必须 rank-consistent。
+
+### mutable backend metadata
+
+`NsaBackend.forward_metadata` 是 mutable state。多个 slot 共享一个 backend 会互相覆盖 metadata。必须先做 backend slot isolation。
+
+### hidden synchronous fallback
+
+如果 `cp_all_gather_into_tensor_async` 在某些 backend 下退化为同步路径，CP-TBO 的 overlap 会消失。必须检测并日志提示，必要时 fallback。
+
+### 额外显存压力
+
+多个 slot 会同时持有 hidden/residual、index dense buffer、MLA dense buffer、pending handle。admission 不只检查 KV pool，也要检查 workspace budget。
+
+### CP/EP 带宽争用
+
+CP materialize/gather、DeepEP A2A、attention 可能争抢 NVLink/IB。默认策略不能盲目叠加通信。
+
+### Phase8 prefetch 重新接入风险
+
+Phase8 prefetch 原本是单请求局部优化。重新接入 Phase9 后，如果不纳入全局 CP collective order，可能导致 rank 间 collective 序列不一致。因此 P9.7 必须在 P9.1-P9.6 初版稳定后再做。
+
+## 测试计划
+
+### 单元测试
+
+新增/扩展测试文件：
+
+- 新增 `test/registered/unit/batch_overlap/test_cp_multi_request_overlap.py`
+  - env/config parsing；
+  - slot eligibility；
+  - deterministic schedule generation；
+  - EP-overlap policy。
+
+- 新增 `test/registered/unit/layers/test_cp_tbo_backend.py`
+  - child backend creation；
+  - per-slot metadata isolation；
+  - unsupported backend fallback。
+
+- 扩展 `test/registered/unit/layers/test_nsa_cp_utils.py`
+  - CP all-gather rerange launch/finish wrapper 与旧同步 wrapper 等价；
+  - fake handle order。
+
+- 扩展 `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
+  - shared paged materialize launch/finish wrapper 等价；
+  - shared token MLA materialize launch/finish wrapper 等价；
+  - async all-reduce unavailable fallback reason。
+
+远端容器中运行，不在本地跑：
+
+```bash
+cd /sgl-workspace/sglang-tai
+PYTHONPATH=/sgl-workspace/sglang-tai/python:/mnt/beegfs/cjy/tai-kernel/python \
+pytest -q \
+  test/registered/unit/batch_overlap/test_cp_multi_request_overlap.py \
+  test/registered/unit/layers/test_cp_tbo_backend.py \
+  test/registered/unit/layers/test_nsa_cp_utils.py \
+  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+```
+
+必须覆盖：
+
+- 只接受 supported CP shared-KV in-seq-split prefill batch；
+- 拒绝 return-logprob、hidden capture、CUDA graph、speculative、rank-local fallback 不一致；
+- fewer than 2 eligible requests 时 fallback single-batch；
+- `max_slots > 2` admission 不构造 true multi-sequence `ForwardBatch`；
+- fake handle 验证不同 layer slot 的 collective order；
+- fake backend 验证 slot0/slot1/slot2 metadata 隔离；
+- `avoid_ep_comm` 不在 marked DeepEP comm window 启动新 CP comm；
+- P9.7 时验证 Phase8 prefetch 的 handle 由 CP-TBO scheduler 管理。
+
+### 远端功能测试
+
+GLM5 CP prefill 环境：
+
+1. CP-TBO disabled，发送两个 concurrent long-prompt 请求。
+2. CP-TBO enabled：
+
+```text
+SGLANG_CP_TWO_BATCH_OVERLAP=1
+SGLANG_CP_TBO_MAX_SLOTS=2
+```
+
+发送相同请求。
+
+3. 测试：
+
+```text
+SGLANG_CP_TBO_MAX_SLOTS=3
+SGLANG_CP_TBO_MAX_SLOTS=4
+```
+
+并发送对应数量 concurrent long prompts。
+
+4. deterministic generation 设置下对比输出。
+5. 检查：
+   - 无 CP shared KV fallback spam；
+   - 无 health-check hang；
+   - 无 request loss；
+   - 输出正常。
+
+### Profiling 检查
+
+Nsight 应看到：
+
+- MQA 前 large idle bubble 减少；
+- CP index/MLA communication ranges 优先被其他 slot 的 MQA/NSA attention 覆盖；attention 窗口不足时才被 dense compute/允许的 MLP work 覆盖；
+- slot 可处在不同 layer；
+- CP collective order 没有乱序；
+- 初版没有 Phase8 prefetch 的额外自主通信；
+- P9.7 后 Phase8 prefetch range 纳入 CP-TBO scheduler；
+- 默认 EP policy 下没有 CP/DeepEP 带宽 pile-up。
+
+## 成功标准
+
+Phase9 成功条件：
+
+1. 两个及以上 concurrent long-prompt prefill 请求能在 CP shared KV 下正确完成。
+2. 至少两个 slot 可处于不同 layer id，且不破坏 metadata 与 collective order。
+3. `SGLANG_CP_TBO_MAX_SLOTS > 2` 表示多个独立 request slot，并可正常运行。
+4. fallback path 安全，保持单请求行为不变。
+5. Nsight 显示两个 attention 前 communication windows 主要被其他 slot 的 MQA/NSA attention 窗口隐藏。
+6. concurrent long-prompt 端到端吞吐高于 Phase8 单请求 pipeline。
+7. P9.7 后 Phase8 prefetch 可在 CP-TBO scheduler 管理下重新启用，不引入 collective order 风险。
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 0edd9de4b..929ddf3d5 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -208,8 +208,10 @@ class Envs:
     SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
+    SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)
     SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
     SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
+    SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)
diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
index 8abd6e190..c945d7f0c 100644
--- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     cp_shared_kv_debug_enabled,
     cp_shared_kv_mla_prefetch_enabled,
     cp_shared_kv_mla_prefetch_log,
+    cp_shared_kv_mla_prefetch_min_prefix_pages,
     cp_shared_kv_mla_prefetch_should_log_layer,
     filter_locs_mappable_to_physical_pool,
     filter_pages_mappable_to_physical_pool,
@@ -181,6 +182,9 @@ class CpSharedKVMlaPrefetcher:
                 int(real_page_table.numel()),
             )
             return None
+        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size)
+        if prefix_pages < min_prefix_pages:
+            return None
 
         cp_group = get_attention_cp_group()
         if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
@@ -595,6 +599,9 @@ class CpSharedKVIndexPrefetcher:
                 int(real_page_table.numel()),
             )
             return None
+        min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size)
+        if prefix_pages < min_prefix_pages:
+            return None
 
         cp_group = get_attention_cp_group()
         if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
index 6bd3644a4..150f4fbb9 100644
--- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
@@ -41,6 +41,10 @@ def cp_shared_kv_tai_fused_mla_store_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_FUSED_MLA_STORE.get()
 
 
+def cp_shared_kv_tai_index_mqa_prepare_enabled() -> bool:
+    return envs.SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE.get()
+
+
 def cp_shared_kv_mla_prefetch_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get()
 
@@ -49,6 +53,20 @@ def cp_shared_kv_mla_prefetch_log_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get()
 
 
+def cp_shared_kv_mla_prefetch_min_prefix_pages(cp_size: int) -> int:
+    """Minimum prefix pages required to enable Phase8 prefetch.
+
+    Negative env values mean "use cp_size" so the default skips tiny prefixes
+    that cannot cover all CP lanes.  Set the env to 0 to disable the gate, or to
+    a positive absolute page count for workload-specific tuning.
+    """
+
+    configured = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.get()
+    if configured < 0:
+        return max(int(cp_size), 0)
+    return max(int(configured), 0)
+
+
 def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None:
     if cp_shared_kv_mla_prefetch_log_enabled():
         logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args)
@@ -219,6 +237,16 @@ def _load_tai_fused_mla_store_kernel():
         return None
 
 
+@lru_cache(maxsize=1)
+def _load_tai_index_mqa_prepare_kernel():
+    try:
+        from tai_kernel.nsa_prefill import prepare_cp_mqa_kv_and_range
+
+        return prepare_cp_mqa_kv_and_range
+    except Exception:
+        return None
+
+
 def _log_tai_materialize_fallback(
     key: str,
     message: str,
@@ -249,6 +277,51 @@ def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor:
     return tensor if tensor.is_contiguous() else tensor.contiguous()
 
 
+def try_tai_prepare_cp_mqa_index(
+    *,
+    index_buffer: torch.Tensor,
+    page_indices: torch.Tensor,
+    kv_len: int,
+    valid_q_count: int,
+    ke_start: int,
+    page_size: int,
+    index_head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
+    """Try the TAI fused GetK/GetS + MQA range prepare path.
+
+    The fallback path in SGLang launches separate GetK/GetS kernels and then
+    creates `ks`/`ke` tensors with torch elementwise ops.  The TAI kernel fuses
+    those preparation steps into one launch for the single-sequence CP pair path.
+    """
+
+    if not cp_shared_kv_tai_index_mqa_prepare_enabled():
+        return None
+
+    kernel = _load_tai_index_mqa_prepare_kernel()
+    if kernel is None:
+        return None
+
+    if index_buffer.dtype != torch.uint8:
+        return None
+    if not index_buffer.is_contiguous():
+        return None
+    if index_head_dim != 128 or page_size != 64:
+        return None
+
+    try:
+        return kernel(
+            index_buffer,
+            _contiguous_for_tai(page_indices),
+            kv_len=int(kv_len),
+            valid_q_count=int(valid_q_count),
+            ke_start=int(ke_start),
+            page_size=int(page_size),
+            index_head_dim=int(index_head_dim),
+        )
+    except Exception:
+        return None
+
+
 def try_tai_fused_mla_store(
     *,
     token_to_kv_pool,
diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
index 18f75eb6e..659e8c02c 100644
--- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
@@ -24,6 +24,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     materialize_shared_paged_buffer,
     tensor_debug_checksum,
     tensor_debug_summary,
+    try_tai_prepare_cp_mqa_index,
 )
 from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
 from sglang.srt.layers.layernorm import LayerNorm
@@ -1097,39 +1098,59 @@ class Indexer(MultiPlatformOp):
                 )
 
             ke_start = cp_kv_end - actual_seq_q + 1
-            ke_offset = torch.arange(
-                ke_start,
-                ke_start + valid_q_count,
-                dtype=torch.int32,
-                device=q_fp8.device,
-            )
 
             q_fp8 = q_fp8[:valid_q_count]
             weights = weights[:valid_q_count]
             kv_len = min(cp_kv_end, logical_kv_limit)
             if current_index_kv is None:
                 assert index_buffer is not None
-                k_fp8 = index_buf_accessor.GetK.execute(
-                    forward_batch.token_to_kv_pool,
-                    index_buffer,
-                    seq_len=kv_len,
-                    page_indices=block_tables[0],
-                )
-                k_scale = index_buf_accessor.GetS.execute(
-                    forward_batch.token_to_kv_pool,
-                    index_buffer,
-                    seq_len=kv_len,
+                tai_prepared = try_tai_prepare_cp_mqa_index(
+                    index_buffer=index_buffer,
                     page_indices=block_tables[0],
+                    kv_len=kv_len,
+                    valid_q_count=valid_q_count,
+                    ke_start=ke_start,
+                    page_size=page_size,
+                    index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
                 )
+                if tai_prepared is not None:
+                    k_fp8_u8, k_scale, ks, ke_offset = tai_prepared
+                    k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
+                else:
+                    ke_offset = torch.arange(
+                        ke_start,
+                        ke_start + valid_q_count,
+                        dtype=torch.int32,
+                        device=q_fp8.device,
+                    )
+                    k_fp8 = index_buf_accessor.GetK.execute(
+                        forward_batch.token_to_kv_pool,
+                        index_buffer,
+                        seq_len=kv_len,
+                        page_indices=block_tables[0],
+                    )
+                    k_scale = index_buf_accessor.GetS.execute(
+                        forward_batch.token_to_kv_pool,
+                        index_buffer,
+                        seq_len=kv_len,
+                        page_indices=block_tables[0],
+                    )
 
-                k_fp8 = k_fp8.view(torch.float8_e4m3fn)
-                k_scale = k_scale.view(torch.float32).squeeze(-1)
+                    k_fp8 = k_fp8.view(torch.float8_e4m3fn)
+                    k_scale = k_scale.view(torch.float32).squeeze(-1)
+                    ks = torch.zeros_like(ke_offset)
             else:
+                ke_offset = torch.arange(
+                    ke_start,
+                    ke_start + valid_q_count,
+                    dtype=torch.int32,
+                    device=q_fp8.device,
+                )
                 k_fp8, k_scale = current_index_kv
                 k_fp8 = k_fp8[:kv_len].contiguous()
                 k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous()
+                ks = torch.zeros_like(ke_offset)
             kv_fp8 = (k_fp8, k_scale)
-            ks = torch.zeros_like(ke_offset)
             ke = ke_offset
 
             with self._with_real_sm_count():
diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
index 89d3c7588..a6ec657d3 100644
--- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
@@ -458,6 +458,24 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True):
             self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled())
 
+    def test_mla_prefetch_min_prefix_pages_defaults_to_cp_size_and_can_override(self):
+        from sglang.srt.environ import envs
+        from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+            cp_shared_kv_mla_prefetch_min_prefix_pages,
+        )
+
+        envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.clear()
+        self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 8)
+
+        with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(0):
+            self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 0)
+
+        with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(16):
+            self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 16)
+
+        with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(-2):
+            self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(4), 4)
+
     def test_fused_mla_store_uses_tai_kernel_when_enabled(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
         from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout