From 7284a469a2a3ce427c4ee5a256adab5e15c327fc Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Thu, 11 Jun 2026 05:09:41 +0800
Subject: [PATCH] Reuse prepared HiCache load descriptors across CP prefill
 layers

CP shared-KV bs>1 cache-hit loads already merge request load ops, but the host pool still rebuilt layer-invariant mapping work from the same host/device indices. Introduce a PreparedLoadDescriptor lifecycle around begin/end load, wire MLA KV and NSA index H2D loads through tai-kernel prepared submit when available, and add timing hooks plus regression coverage for descriptor reuse and explicit fallback logging. Record the P4/P6b design and benchmark results in the advanced feature notes.

Constraint: Radix residency and allocator decisions remain synchronous; only the data-transfer descriptor is prepared for per-layer async submit.

Constraint: Production fast path must not silently fall back when tai prepared H2D support is missing.

Rejected: Cross-batch descriptor reuse | descriptor lifetime and tensor ownership are only safe within one load operation.

Rejected: Change L2->L1 scheduling to layer-ahead prefetch in this commit | that is a separate lifecycle change after descriptor reuse is stable.

Confidence: medium

Scope-risk: moderate

Directive: Keep LayerDoneCounter per-layer readiness semantics; do not replace with all-layer waits.

Tested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py

Tested: Remote g0034:cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py (88 passed)

Tested: Remote tai-kernel prepared descriptor CUDA test (6 passed) and P4 benchmark full matrix (90 rows)

Not-tested: ETE replay/GSM8K cache-hit correctness after this commit

Not-tested: Layer-ahead L2->L1 prefetch scheduling

Co-authored-by: OmX <omx@oh-my-codex.dev>
---
 ...l1_transfer_prefetch_descriptor_plan_zh.md | 773 ++++++++++++++++++
 .../sglang/srt/managers/cache_controller.py   |  65 ++
 .../sglang/srt/mem_cache/memory_pool_host.py  | 295 ++++++-
 .../managers/test_hicache_controller_cp.py    | 252 +++++-
 4 files changed, 1375 insertions(+), 10 deletions(-)
 create mode 100644 docs/advanced_features/nsa_prefill_cp_hicache_bs_gt1_l2_l1_transfer_prefetch_descriptor_plan_zh.md

diff --git a/docs/advanced_features/nsa_prefill_cp_hicache_bs_gt1_l2_l1_transfer_prefetch_descriptor_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_hicache_bs_gt1_l2_l1_transfer_prefetch_descriptor_plan_zh.md
new file mode 100644
index 000000000..24c0bc00d
--- /dev/null
+++ b/docs/advanced_features/nsa_prefill_cp_hicache_bs_gt1_l2_l1_transfer_prefetch_descriptor_plan_zh.md
@@ -0,0 +1,773 @@
+# NSA Prefill CP HiCache bs>1 L2->L1 Transfer Prefetch Descriptor 计划
+
+> 日期：2026-06-10  
+> 分支：`cjy-cp-refactor`  
+> 范围：CP shared-KV + HiCache + bs>1 cache-hit 场景下，L2/host -> L1/device 数据搬运的 batch-aware per-layer prefetch 与 descriptor 复用。  
+> 相关文档：
+> - `docs/advanced_features/nsa_prefill_cp_hicache_load_prefetch_overlap_notes.md`
+> - `docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md`
+> - `docs/advanced_features/nsa_prefill_cp_hicache_layer_page_first_direct_plan.md`
+
+## 0. 目标和非目标
+
+### 0.1 目标
+
+1. 在 `--enable-cp-shared-kv-prefill-bs-gt1` 下，让多个 cache-hit request 的 L2->L1 transfer 以 batch 级 descriptor 组织。
+2. 降低每层重复构造 descriptor / Python wrapper / extension submit 造成的 CPU overhead。
+3. 保留现有正确性合同：
+   - radix / residency / allocator reservation 同步完成；
+   - data transfer 异步执行；
+   - target / draft KV 跟随同一个 logical cache node；
+   - 每层消费前精确等待本层 transfer event。
+4. 保持当前 per-layer H2D overlap 能力，并为后续“提前两层 L2->L1 transfer prefetch”打基础。
+5. 支持当前生产参数：
+   - `--hicache-io-backend direct`
+   - `--hicache-mem-layout page_first_direct`
+   - `--kv-cache-dtype fp8_e4m3`
+   - CP owner-lane / zigzag / bs>1 batch plan。
+
+### 0.2 非目标
+
+1. 第一阶段不改 radix tree 语义，不让异步 transfer 决定 cache node 是否可见。
+2. 第一阶段不做跨 batch descriptor 复用。descriptor 只在当前 load op / 当前 batch 生命周期内复用。
+3. 第一阶段不改 decode KV transfer / Mooncake transfer。
+4. 第一阶段不强制切换到 `layer_page_first`。先让 `page_first_direct` 路径有 prepared descriptor，再评估 LPF 切换收益。
+5. 第一阶段不新增 collective。descriptor 一致性必须来自本地 deterministic metadata，不依赖 all-reduce/all-gather。
+
+---
+
+## 1. 当前代码事实
+
+### C1. 当前已有 batch 内 op 粗粒度合并
+
+`python/sglang/srt/managers/cache_controller.py`
+
+```text
+CacheOperation.merge_ops(load_queue)
+  host_indices = torch.cat([op.host_indices for op in ops])
+  device_indices = torch.cat([op.device_indices for op in ops])
+```
+
+结论：当前 bs>1 下多个 request 的 load-back op 已经会合并成一个 `CacheOperation`。  
+但这个合并只是 tensor concat，不是 transfer descriptor compact。
+
+### C2. 当前 per-layer transfer 每层都会重新进入 host pool API
+
+`HiCacheController.start_loading()`：
+
+```text
+begin_load_to_device_op(host_indices, device_indices, io_backend)
+for layer_id in range(layer_num):
+  mem_pool_host.load_to_device_per_layer(
+      mem_pool_device,
+      host_indices,
+      device_indices,
+      layer_id,
+      io_backend,
+  )
+end_load_to_device_op()
+```
+
+结论：indices 是同一份，但每层仍会重新调用 Python method / extension wrapper。  
+direct backend 下，底层 TAI op 仍可能每层重新构造 H2D copy descriptor。
+
+### C3. `begin_load_to_device_op()` 已经是 descriptor 预处理入口
+
+`python/sglang/srt/mem_cache/memory_pool_host.py`
+
+base host pool:
+
+```python
+def begin_load_to_device_op(self, host_indices, device_indices, io_backend):
+    """Prepare layer-invariant metadata for one host->device load op."""
+```
+
+NSA host pool 已经使用该入口预计算 indexer page indices：
+
+```text
+NSATokenToKVPoolHost.begin_load_to_device_op()
+  _active_load_indexer_page_indices = _get_indexer_page_indices(...)
+```
+
+结论：代码结构已经允许“per load op prepare once, per layer reuse”。  
+当前 MLA KV 主体还没有充分利用这个入口。
+
+### C4. 当前 H2D readiness 是 per-layer event
+
+`LayerDoneCounter` 为每个 producer 维护每层 event：
+
+```text
+LayerLoadingEvent.complete(layer_id)
+LayerDoneCounter.wait_until_on_stream(layer_id, stream)
+```
+
+KV pool 访问路径会等待：
+
+```text
+get_key_buffer_for_prefetch(layer_id, stream)
+  -> wait_layer_transfer_on_stream(layer_id, stream)
+```
+
+结论：即使 descriptor 合并，仍必须保留 per-layer complete/wait。不能退回 all-layer wait。
+
+### C5. 当前 L2->L1 load 启动点仍偏晚
+
+当前时序：
+
+```text
+PrefillAdder.add_one_req()
+  -> init_load_back()
+  -> load_cp() queue load op
+
+ScheduleBatch 创建后
+  -> ready_to_load_host_cache()
+  -> start_loading()
+  -> load_stream per-layer enqueue
+
+forward consume
+  -> wait layer event
+```
+
+这已经是 async per-layer load，但不是完整意义上的“提前两层 prefetch”。  
+第一层和早期层仍可能等待 H2D。
+
+---
+
+## 2. 问题定义
+
+当前瓶颈不是简单的“load op 没合并”，而是合并层级不够：
+
+```text
+已合并：
+  多 req host_indices/device_indices concat 成一个 CacheOperation
+
+未合并：
+  page 连续区间未 compact
+  每层 transfer descriptor 未复用
+  每层 Python/extension submit 仍线性增长
+  L2->L1 transfer 未提前到 layer-k+2 的 prefetch window
+```
+
+在 cache hit + bs>1 场景下，请求通常形态是：
+
+```text
+prefix 很大：L2/host hit
+extend 很短：200~2000 tokens
+batch size：2~10
+layer 数：约 78
+```
+
+如果每层都重新构造 descriptor，那么 CPU overhead 约随：
+
+```text
+O(layer_num * descriptor_build_cost)
+```
+
+增长。这个 overhead 很难被 GPU compute overlap，因为它发生在 transfer submit/control path。
+
+目标是把可复用部分改成：
+
+```text
+O(descriptor_build_cost + layer_num * cheap_submit_cost)
+```
+
+并为后续：
+
+```text
+layer L-2 启动 L2->L1(layer L)
+layer L-1 启动 L1 shared-KV prefix prefetch(layer L)
+layer L consume 只 wait event
+```
+
+提供基础。
+
+---
+
+## 3. 目标架构
+
+### 3.1 同步控制面
+
+保持现有同步控制面：
+
+```text
+scheduler admission
+  match_prefix()
+  init_load_back()
+  alloc_pages_with_owners(page_owners)
+  update prefix_indices / extend_input_len
+  enqueue CacheOperation
+```
+
+这些操作仍在 scheduler 线程完成，因为它们决定：
+
+- radix node 是否 device-resident；
+- fresh device page owner pattern；
+- request prefix/extend metadata；
+- batch plan 输入。
+
+### 3.2 异步数据面
+
+把当前 data plane 从：
+
+```text
+for each layer:
+  build/submit transfer using host_indices/device_indices
+```
+
+演进为：
+
+```text
+begin_load_to_device_op:
+  build PreparedLoadDescriptor once
+
+for each layer:
+  submit layer transfer using PreparedLoadDescriptor + layer_id
+
+end_load_to_device_op:
+  release descriptor lifetime refs
+```
+
+### 3.3 descriptor 生命周期
+
+descriptor 生命周期与 `start_loading()` 的 producer id 绑定：
+
+```text
+start_loading()
+  producer_id = update_producer()
+  descriptor = prepare(...)
+  enqueue layer transfers
+  record_stream / hold tensors
+  ack_load_queue.append(HiCacheAck(...))
+```
+
+descriptor 不跨 batch 复用。  
+descriptor 必须持有所有异步 transfer 需要的 tensor / pinned metadata 引用，直到 load stream 使用完。
+
+---
+
+## 4. Prepared descriptor 设计
+
+### 4.1 第一版 Python-side descriptor
+
+先增加轻量 Python-side descriptor，降低重复 page index 准备和 wrapper 参数构造。
+
+建议 dataclass：
+
+```python
+@dataclass
+class PreparedLoadDescriptor:
+    host_indices: torch.Tensor
+    device_indices: torch.Tensor
+    host_page_indices: torch.Tensor | None
+    device_page_indices: torch.Tensor | None
+    page_size: int
+    io_backend: str
+    layout: str
+    num_tokens: int
+    num_pages: int
+```
+
+对 MLA KV：
+
+- `host_indices/device_indices` 是 token/page slot 级映射；
+- `page_first_direct` / `layer_page_first` 传给 TAI direct op；
+- 第一版不改变底层 TAI API，只把 layer-invariant tensor 准备和参数绑定集中到 `begin_load_to_device_op()`。
+
+对 NSA index：
+
+- 复用现有 `_active_load_indexer_page_indices`；
+- 把它纳入同一个 descriptor 生命周期，避免 target KV 和 index KV 分别散落状态。
+
+### 4.2 第二版 compact segment descriptor
+
+在 Python-side descriptor 稳定后，增加 segment compact：
+
+```python
+@dataclass
+class TransferSegment:
+    host_start: int
+    device_start: int
+    length: int
+```
+
+compact 规则：
+
+```text
+如果 host_indices 和 device_indices 同时连续:
+  合并为一个 segment
+否则:
+  保持 page/token 级 entries
+```
+
+收益：
+
+- 减少 cudaMemcpyBatch / TAI descriptor entries；
+- 与 L1/L2 allocator 连续分配优化协同；
+- 对 `layer_page_first` 更有价值。
+
+### 4.3 第三版 TAI prepared descriptor API
+
+最终把 descriptor 下沉到 tai-kernel：
+
+```python
+desc = tai_kernel.nsa_prefill.prepare_h2d_page_descriptor(
+    host_indices,
+    device_indices,
+    page_size,
+    layout,
+)
+
+tai_kernel.nsa_prefill.submit_h2d_layer(
+    desc,
+    src_ptr,
+    dst_ptr,
+    layer_id,
+)
+```
+
+TAI descriptor 内部可以选择：
+
+- cudaMemcpyBatchAsync entries；
+- compacted ranges；
+- future 0SM / copy-engine queue；
+- future CUDA driver batch copy API fallback。
+
+SGLang 侧只负责 descriptor 生命周期和 per-layer event。
+
+---
+
+## 5. Phase 计划
+
+### P0. 记录现状和保护合同
+
+**目标：** 明确当前行为，避免后续优化误改语义。
+
+**工作：**
+
+1. 在本文件记录：
+   - current batch `CacheOperation.merge_ops` 已经做 concat；
+   - `start_loading()` 每层调用 `load_to_device_per_layer()`；
+   - `LayerDoneCounter` 是 per-layer wait 合同；
+   - descriptor 只能跨 layer，不能跨 batch。
+2. 不改代码。
+
+**验证：**
+
+- 文档自检无 “TBD/TODO/后续补充” 占位。
+
+### P1. 为 MLA KV 主体增加 Python-side prepared descriptor
+
+**目标：** 使用已有 `begin_load_to_device_op()` / `end_load_to_device_op()` 入口，为 MLA KV 主体缓存 layer-invariant 状态。
+
+**涉及文件：**
+
+- `python/sglang/srt/mem_cache/memory_pool_host.py`
+
+**设计：**
+
+1. 在 host pool base / MLA host pool 上添加 `_active_load_descriptor` 字段。
+2. `begin_load_to_device_op()` 创建 descriptor。
+3. `load_to_device_per_layer()` 优先使用 descriptor。
+4. `end_load_to_device_op()` 清理 descriptor。
+5. 若 descriptor 缺失，保留当前路径作为 fail-fast 或 warning fallback；生产 fast path 不应静默 fallback。
+
+**测试：**
+
+- unit test 覆盖：
+  - begin 后 descriptor 存在；
+  - per-layer 调用复用 descriptor；
+  - end 后 descriptor 清空；
+  - missing descriptor 的 fallback/warning 行为明确。
+
+### P2. 把 NSA index load descriptor 生命周期合并到统一结构
+
+**目标：** 当前 index 已经有 `_active_load_indexer_page_indices`，但状态分散。把 index descriptor 纳入统一 prepared descriptor。
+
+**涉及文件：**
+
+- `python/sglang/srt/mem_cache/memory_pool_host.py`
+
+**设计：**
+
+1. descriptor 内包含：
+   - MLA KV token/page mapping；
+   - NSA index page mapping；
+   - active index layer ids。
+2. `NSATokenToKVPoolHost.begin_load_to_device_op()` 调 base prepare 后补充 index fields。
+3. `_load_indexer_to_device_per_layer()` 从 descriptor 读取 prepared index page indices。
+
+**测试：**
+
+- active index layer 被正确 load；
+- inactive index layer 不访问 compact index cache；
+- `index_topk_freq > 1` 时 descriptor 不请求 inactive layers。
+
+### P3. 增加 descriptor build/submit timing
+
+**目标：** 用现有 timing env 验证 CPU overhead 是否下降，不新增长期 noisy log。
+
+**涉及文件：**
+
+- `python/sglang/srt/managers/cache_controller.py`
+- `python/sglang/srt/mem_cache/memory_pool_host.py`
+
+**设计：**
+
+复用：
+
+```text
+SGLANG_CP_SHARED_KV_BS_GT1_TIMING
+SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT
+SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS
+```
+
+记录阶段：
+
+```text
+prepare_load_descriptor
+submit_h2d_layer_loop
+submit_h2d_layer_per_call_slow
+end_load_descriptor
+```
+
+**测试：**
+
+- env disabled 时不输出；
+- env enabled 且 slow threshold 命中时输出；
+- 不新增独立 debug env。
+
+### P4. TAI microbenchmark：descriptor reuse baseline
+
+**目标：** 独立验证 descriptor reuse 是否降低 CPU submit overhead。
+
+**涉及仓库：**
+
+- `tai-kernel`
+
+**benchmark 范围：**
+
+1. tokens/pages：
+   - 4k
+   - 16k
+   - 64k
+   - 120k
+   - 192k
+2. batch shape：
+   - bs=1 large prefix
+   - bs=5 mixed prefix
+   - bs=10 small extend / large prefix
+3. layout：
+   - `page_first_direct`
+   - `layer_page_first`
+4. page 分布：
+   - contiguous
+   - allocator-like owner-lane
+   - random fragmented
+
+**指标：**
+
+```text
+descriptor build ms
+per-layer submit total ms
+effective H2D GB/s
+entries/pages/segments
+CPU time per layer
+```
+
+### P5. Segment compact
+
+**目标：** 在 descriptor 内合并连续 host/device ranges，减少 bottom-level transfer entries。
+
+**涉及文件：**
+
+- SGLang descriptor builder；
+- tai-kernel benchmark；
+- 后续 tai-kernel runtime API。
+
+**设计：**
+
+先在 CPU 上做线性 scan：
+
+```text
+prev_host + 1 == cur_host
+prev_device + 1 == cur_device
+```
+
+连续则合并。  
+不做复杂 merge/sort；不改变 order；不引入额外 O(n log n)。
+
+**测试：**
+
+- contiguous pages 合并成少量 segments；
+- fragmented pages 不错误合并；
+- device/host 只有一侧连续时不合并；
+- segment 展开后与原 indices 完全一致。
+
+### P6. TAI prepared descriptor API
+
+**目标：** 将 descriptor 进一步下沉，减少 Python/C++ wrapper 每层重复工作。
+
+**涉及仓库：**
+
+- `tai-kernel`
+- `sglang-dev`
+
+**设计：**
+
+1. tai-kernel 提供 prepare / submit / destroy 三段式 API。
+2. SGLang `begin_load_to_device_op()` 调 prepare。
+3. `load_to_device_per_layer()` 调 submit。
+4. `end_load_to_device_op()` 调 destroy。
+5. 如果 tai-kernel 不支持当前 layout/dtype，warning + 明确 fallback 到旧 direct path。
+
+**测试：**
+
+- bf16 / fp8_e4m3；
+- `page_first_direct`；
+- `layer_page_first`；
+- target-only；
+- target+draft；
+- index active-layer skip。
+
+### P7. 提前两层 L2->L1 transfer prefetch
+
+**目标：** 在 descriptor reuse 稳定后，把 transfer 启动点从 “batch 创建后立刻提交所有层” 改成可控的 layer-ahead prefetch。
+
+**设计方向：**
+
+第一版不改变 scheduler admission，只改变 submit scheduling：
+
+```text
+start_loading()
+  prepare descriptor once
+  enqueue/submit first N warmup layers
+
+layer end hook 或 forward progress hook:
+  submit layer_id + 2 transfer
+```
+
+约束：
+
+1. 不阻塞 forward stream 发起 transfer。
+2. 每层仍通过 `LayerDoneCounter.complete(layer)` 通知 readiness。
+3. 如果 layer-ahead submit 来不及，consume wait 仍保证 correctness。
+4. draft target 的 layer submit 顺序一致。
+
+**风险：**
+
+- 当前 `start_loading()` 一次性提交所有层，简单且正确；
+- 改成逐层提交需要一个可靠 hook，不应复用占 SM 的 shared-KV materialize hook；
+- 第一版可以先保留一次性提交，只做 descriptor reuse；提前两层作为独立阶段。
+
+### P8. ETE 验证
+
+**目标：** 证明性能和正确性没有回退。
+
+**验证项：**
+
+1. 单测：
+   - `test_prefill_adder.py`
+   - CP shared KV runtime/layout tests
+   - 新增 descriptor tests
+2. 远端 py_compile。
+3. TAI benchmark。
+4. GSM8K：
+   - 第一轮 cold/cache-miss；
+   - 第二轮 cache-hit；
+   - 精度不掉点。
+5. replay workload：
+   - cache hit 下吞吐；
+   - accept len；
+   - prefill failed count；
+   - timing log 中 descriptor build/submit 是否下降。
+
+---
+
+## 6. 关键不变量
+
+1. `prefix_indices` 更新必须发生在 scheduler admission 阶段，不能等 transfer 完成后异步更新。
+2. `device_indices` 分配后，即使 layer 数据尚未 load 完，也可以作为 future-ready L1 slot 暴露给 batch metadata；消费前必须 wait。
+3. `LayerDoneCounter` 的 per-layer event 不能被 all-layer event 替代。
+4. descriptor 可跨 layer 复用，不可跨 batch 复用。
+5. descriptor 不能持有会在 load stream 完成前释放的临时 tensor。
+6. zero-owned rank 仍需要 logical no-op ack，保持 CP ranks 的 load ack 顺序一致。
+7. target/draft KV 如果共享 HiCache logical node，则 descriptor 生命周期必须覆盖两者；不能 target 成功而 draft 失败后继续可见。
+
+---
+
+## 7. 建议实施顺序
+
+推荐先做：
+
+```text
+P1 -> P2 -> P3 -> P4
+```
+
+原因：
+
+- 低风险，不改变 transfer 启动时机；
+- 能直接验证 CPU overhead 是否来自 descriptor 重建；
+- 如果收益明显，再进入 P5/P6；
+- 如果收益不明显，说明主要瓶颈可能在 H2D bandwidth、early layer wait、shared-KV L1 prefetch/reduce 或 decode transfer。
+
+不建议一开始就做 P7。  
+P7 会改 forward/layer hook 生命周期，风险高于 descriptor reuse，应在 P1-P4 有数据后推进。
+
+
+---
+
+## 8. P4-P6 实施记录（2026-06-11）
+
+### R1. P4 benchmark 已补齐并完成一轮全矩阵验证
+
+新增 `tai-kernel/benchmark/nsa_prefill/benchmark_hicache_h2d_prepared_descriptor.py`。
+
+覆盖维度：
+
+- token 默认：`4096, 16384, 65536, 120000, 192000`；
+- batch shape 默认：`1, 5, 10`；
+- layout：`page_first_direct`, `layer_page_first`；
+- page pattern：`contiguous`, `owner_lane`, `random`，并支持 `fragmented/strided/reverse`；
+- 输出：descriptor build ms、per-layer submit ms、submit per layer us、total ms、effective GB/s、segments。
+
+远端 quick smoke：
+
+```text
+page_first_direct,contiguous,bs=1,tokens=4096,layers=2,segments=64,total=0.373190ms
+layer_page_first,contiguous,bs=1,tokens=4096,layers=2,segments=1,total=0.081027ms
+```
+
+P4 全矩阵远端验证：
+
+```text
+容器：g0034:cjy-glm5-new
+命令：benchmark_hicache_h2d_prepared_descriptor.py
+参数：tokens=4096,16384,65536,120000,192000
+参数：bs=1,5,10
+参数：patterns=contiguous,owner_lane,random
+参数：layouts=page_first_direct,layer_page_first
+参数：warmup=1 repeat=2
+CSV：/mnt/beegfs/cjy/log/hicache_h2d_prepared_descriptor_p4_20260610_210053.csv
+行数：90 = 2 layouts * 3 patterns * 3 bs * 5 token sizes
+```
+
+代表性结果：
+
+```text
+page_first_direct contiguous bs=1 tokens=65536:
+  segments=1024 build=1.224ms submit=74.595ms total=107.546ms gbps=54.756
+
+layer_page_first contiguous bs=1 tokens=65536:
+  segments=1 build=0.882ms submit=0.918ms total=106.841ms gbps=55.117
+
+page_first_direct owner_lane bs=10 tokens=192000:
+  segments=3000 build=4.158ms submit=283.340ms total=315.447ms gbps=54.692
+
+layer_page_first owner_lane bs=10 tokens=192000:
+  segments=3000 build=3.636ms submit=283.182ms total=315.782ms gbps=54.634
+
+page_first_direct random bs=10 tokens=192000:
+  segments=3000 build=4.129ms submit=283.274ms total=315.489ms gbps=54.685
+
+layer_page_first random bs=10 tokens=192000:
+  segments=2998 build=3.665ms submit=282.752ms total=315.268ms gbps=54.723
+```
+
+结论：
+
+1. benchmark 已能覆盖 owner-lane 大样本，不再因为 `pool_tokens` 不足失败。实现上按 pattern 动态估算 host page 池，并按 layout lazy 分配 host pinned tensor，避免同时持有 PFD/LPF 两份大 buffer。
+2. `layer_page_first` 对“跨 request/跨 page 真实连续”的 contiguous 场景能显著降低 submit CPU：64k 时 `74.6ms -> 0.9ms`，segments `1024 -> 1`。
+3. `owner_lane/random` 下 LPF 基本 neutral，因为物理 page 不连续，segments 仍约等于 pages；这说明 LPF 的收益强依赖 host/L2 allocator 能提供按 owner-lane/request 连续的物理 page。
+4. 大 transfer 的有效带宽稳定约 `54-55 GB/s`；当前 P4 主要暴露的是 submit/descriptor 开销，不是 H2D 带宽不足。
+
+### R2. P5 segment compact 已实现为保守线性扫描
+
+新增 tai-kernel API：
+
+```python
+prepare_h2d_page_descriptor(src_indices, dst_indices, page_size=..., layout=...)
+```
+
+合同：
+
+- indices 必须是 CPU int64 contiguous；CUDA indices fail-fast；
+- token indices 必须 page aligned；
+- `page_first_direct` 固定 layer H2D 不合并相邻 page，因为 host 物理布局是 `[page, layer, page_size, ...]`，相邻 page 的同一 layer 中间隔着其他 layer；
+- `layer_page_first` 只在 host 和 device token start 都连续时合并 segment；不 sort、不 reorder、不做 O(n log n)。
+
+新增单测：
+
+```text
+tests/nsa_prefill/test_kvcacheio_prepared_descriptor.py
+```
+
+验证：
+
+- `page_first_direct` contiguous pages 仍保持每页一个 segment；
+- `layer_page_first` contiguous run 合并；
+- 只有一侧连续不合并；
+- 非 page aligned / CUDA indices fail-fast；
+- destroy 后 submit fail-fast。
+
+### R3. P6 tai-kernel prepared API + SGLang 接入完成
+
+新增 tai-kernel public API：
+
+```python
+H2DPageDescriptor
+H2DPageSegment
+prepare_h2d_page_descriptor(...)
+submit_h2d_layer(desc, src_ptrs, dst_ptrs, layer_id=...)
+destroy_h2d_page_descriptor(desc)
+```
+
+SGLang 接入：
+
+- `HostKVCache.begin_load_to_device_op()` 会为 direct + `page_first_direct/layer_page_first` 预构建 `tai_h2d_descriptor`；
+- `NSATokenToKVPoolHost.begin_load_to_device_op()` 会额外预构建 index `tai_index_h2d_descriptor`（`page_size=1`）；
+- `MLATokenToKVPoolHost.load_to_device_per_layer()` 优先走 `submit_h2d_layer()`；
+- NSA index load 优先走 index prepared descriptor；
+- `end_load_to_device_op()` 统一 destroy；
+- tai API 缺失时 warning fallback 到旧 direct path，不 silent fallback；prepare API 存在但合同不满足时 fail-fast。
+
+当前 P6 已完成到 P6b：
+
+1. Python API 层保留 `H2DPageDescriptor` 合同；
+2. tai-kernel C++ 新增：
+   - `prepare_h2d_page_descriptor_data(...)`
+   - `submit_h2d_layer_prepared(...)`
+3. prepare 阶段在 C++ 中完成 page-aligned 校验、page starts 和 segment starts/pages 构建；
+4. submit 阶段直接使用 C++ prepared tensors，不再每层从 Python wrapper 重新做 validation/segment 构造；
+5. SGLang 在 direct + `page_first_direct/layer_page_first` 下会检查 `cpp_prepared`，如果 tai Python API 存在但 C++ prepared submit 数据不可用，会输出明确 fallback warning，不 silent fallback。
+
+仍未做的是“真正 opaque C++ custom descriptor / 预生成每层 pointer vector”。当前 C++ prepared submit 仍会在每层构造 `cudaMemcpyBatchAsync` 的 src/dst/size vector；但已把重复 page contract 校验和 segment build 从每层热路径移出。若 ETE timing 仍显示 submit CPU overhead 高，下一步才需要继续做 custom descriptor 或 copy-engine queue。
+
+### R4. 验证记录
+
+本地：
+
+```text
+python -m py_compile sglang memory_pool_host.py / tai kvcacheio.py / benchmark script: OK
+PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_kvcacheio_prepared_descriptor.py: 4 passed
+```
+
+远端 `g0034:cjy-glm5-new`：
+
+```text
+/mnt/beegfs/cjy/tai-kernel: test_kvcacheio_prepared_descriptor.py -> 4 passed
+/mnt/beegfs/cjy/sglang-dev: test_hicache_controller_cp.py -> 88 passed
+benchmark_hicache_h2d_prepared_descriptor.py quick smoke -> OK
+```
+
+补充验证（P6b 完成后）：
+
+```text
+/mnt/beegfs/cjy/tai-kernel: test_kvcacheio_prepared_descriptor.py -> 6 passed
+/mnt/beegfs/cjy/sglang-dev: test_hicache_controller_cp.py -> 88 passed
+P4 full matrix benchmark -> 90 rows completed
+```
+
+未验证：
+
+- ETE replay / GSM8K cache-hit 正确性；
+- 底层 C++ opaque descriptor / copy-engine queue 的进一步 CPU submit 降幅（尚未实现）。
diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py
index c86510eb0..51bf710cd 100644
--- a/python/sglang/srt/managers/cache_controller.py
+++ b/python/sglang/srt/managers/cache_controller.py
@@ -78,6 +78,23 @@ def _cp_shared_kv_bs_gt1_cache_timing(
     )
 
 
+def _cp_shared_kv_bs_gt1_cache_timing_start() -> Optional[float]:
+    if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
+        return None
+    return time.perf_counter()
+
+
+def _cp_shared_kv_bs_gt1_cache_timing_log(
+    key: str,
+    start_time: Optional[float],
+    message: str,
+    *args,
+) -> None:
+    if start_time is None:
+        return
+    _cp_shared_kv_bs_gt1_cache_timing(key, start_time, message, *args)
+
+
 class LayerLoadingEvent:
     def __init__(self, num_layers: int):
         self._num_layers = num_layers
@@ -1752,6 +1769,7 @@ class HiCacheController:
         try:
             with device_module.stream(self.load_stream):
                 producer_event.start_event.wait(self.load_stream)
+                prepare_start = _cp_shared_kv_bs_gt1_cache_timing_start()
                 if op is not None:
                     target_load_op_started = begin_load_op(
                         self.mem_pool_host, host_indices, device_indices
@@ -1762,10 +1780,25 @@ class HiCacheController:
                         draft_host_indices,
                         draft_device_indices,
                     )
+                _cp_shared_kv_bs_gt1_cache_timing_log(
+                    "prepare_load_descriptor",
+                    prepare_start,
+                    "target_tokens=%s draft_tokens=%s target_started=%s draft_started=%s",
+                    int(host_indices.numel()) if host_indices is not None else 0,
+                    int(draft_host_indices.numel())
+                    if draft_host_indices is not None
+                    else 0,
+                    target_load_op_started,
+                    draft_load_op_started,
+                )
                 draft_layer_num = (
                     self.draft_mem_pool_device.layer_num if draft_op is not None else 0
                 )
+                layer_loop_start = _cp_shared_kv_bs_gt1_cache_timing_start()
                 for i in range(max(self.layer_num, draft_layer_num)):
+                    layer_submit_start = _cp_shared_kv_bs_gt1_cache_timing_start()
+                    submitted_target = False
+                    submitted_draft = False
                     if draft_op is not None and i < draft_layer_num:
                         if len(draft_host_indices) > 0:
                             self.draft_mem_pool_host.load_to_device_per_layer(
@@ -1775,6 +1808,7 @@ class HiCacheController:
                                 i,
                                 self.io_backend,
                             )
+                            submitted_draft = True
                     if op is not None and i < self.layer_num:
                         if len(host_indices) > 0:
                             self.mem_pool_host.load_to_device_per_layer(
@@ -1784,9 +1818,32 @@ class HiCacheController:
                                 i,
                                 self.io_backend,
                             )
+                            submitted_target = True
                         producer_event.complete(i)
                     elif op is None and i < self.layer_num:
                         producer_event.complete(i)
+                    _cp_shared_kv_bs_gt1_cache_timing_log(
+                        "submit_h2d_layer_per_call_slow",
+                        layer_submit_start,
+                        "layer_id=%s target=%s draft=%s target_tokens=%s draft_tokens=%s",
+                        i,
+                        submitted_target,
+                        submitted_draft,
+                        int(host_indices.numel()) if host_indices is not None else 0,
+                        int(draft_host_indices.numel())
+                        if draft_host_indices is not None
+                        else 0,
+                    )
+                _cp_shared_kv_bs_gt1_cache_timing_log(
+                    "submit_h2d_layer_loop",
+                    layer_loop_start,
+                    "layers=%s target_tokens=%s draft_tokens=%s",
+                    max(self.layer_num, draft_layer_num),
+                    int(host_indices.numel()) if host_indices is not None else 0,
+                    int(draft_host_indices.numel())
+                    if draft_host_indices is not None
+                    else 0,
+                )
                 # NOTE: We must save the host indices and device indices here,
                 # this is because we need to guarantee that these tensors are
                 # still alive when the load stream is executing.
@@ -1799,10 +1856,18 @@ class HiCacheController:
                 if draft_device_indices is not None and draft_device_indices.is_cuda:
                     draft_device_indices.record_stream(self.load_stream)
         finally:
+            end_start = _cp_shared_kv_bs_gt1_cache_timing_start()
             if draft_load_op_started:
                 end_load_op(self.draft_mem_pool_host)
             if target_load_op_started:
                 end_load_op(self.mem_pool_host)
+            _cp_shared_kv_bs_gt1_cache_timing_log(
+                "end_load_descriptor",
+                end_start,
+                "target_started=%s draft_started=%s",
+                target_load_op_started,
+                draft_load_op_started,
+            )
 
         self.ack_load_queue.append(
             HiCacheAck(
diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py
index 3957c05c4..ba1b4ee30 100644
--- a/python/sglang/srt/mem_cache/memory_pool_host.py
+++ b/python/sglang/srt/mem_cache/memory_pool_host.py
@@ -5,8 +5,9 @@ import logging
 import threading
 import weakref
 from collections import defaultdict
+from dataclasses import dataclass
 from functools import lru_cache, wraps
-from typing import Optional
+from typing import Optional, Tuple
 
 import psutil
 import torch
@@ -54,6 +55,26 @@ if _is_npu:
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class PreparedLoadDescriptor:
+    """Layer-invariant metadata for one host->device HiCache load op."""
+
+    host_indices: torch.Tensor
+    device_indices: torch.Tensor
+    host_page_indices: Optional[torch.Tensor]
+    device_page_indices: Optional[torch.Tensor]
+    page_size: int
+    io_backend: str
+    layout: str
+    num_tokens: int
+    num_pages: int
+    index_host_page_indices: Optional[torch.Tensor] = None
+    index_device_page_indices: Optional[torch.Tensor] = None
+    index_active_layer_ids: Tuple[int, ...] = ()
+    tai_h2d_descriptor: Optional[object] = None
+    tai_index_h2d_descriptor: Optional[object] = None
+
+
 @lru_cache(maxsize=1)
 def _load_tai_transfer_kv_per_layer_mla_lf_pf():
     try:
@@ -132,6 +153,47 @@ def _load_tai_transfer_kv_per_layer_direct_lpf_lf():
         ) from exc
 
 
+@lru_cache(maxsize=1)
+def _load_tai_prepare_h2d_page_descriptor():
+    try:
+        from tai_kernel.nsa_prefill import prepare_h2d_page_descriptor
+
+        return prepare_h2d_page_descriptor
+    except Exception as exc:
+        raise RuntimeError(
+            "[CP_HICACHE_FALLBACK][missing_tai_prepared_h2d_descriptor] "
+            "tai_kernel.nsa_prefill.prepare_h2d_page_descriptor is unavailable. "
+            "Falling back to legacy per-layer direct H2D submit."
+        ) from exc
+
+
+@lru_cache(maxsize=1)
+def _load_tai_submit_h2d_layer():
+    try:
+        from tai_kernel.nsa_prefill import submit_h2d_layer
+
+        return submit_h2d_layer
+    except Exception as exc:
+        raise RuntimeError(
+            "[CP_HICACHE_FALLBACK][missing_tai_submit_h2d_layer] "
+            "tai_kernel.nsa_prefill.submit_h2d_layer is unavailable. "
+            "Falling back to legacy per-layer direct H2D submit."
+        ) from exc
+
+
+@lru_cache(maxsize=1)
+def _load_tai_destroy_h2d_page_descriptor():
+    try:
+        from tai_kernel.nsa_prefill import destroy_h2d_page_descriptor
+
+        return destroy_h2d_page_descriptor
+    except Exception as exc:
+        raise RuntimeError(
+            "[CP_HICACHE_FALLBACK][missing_tai_destroy_h2d_page_descriptor] "
+            "tai_kernel.nsa_prefill.destroy_h2d_page_descriptor is unavailable."
+        ) from exc
+
+
 def _raise_layer_page_first_page_buffer_meta_unsupported() -> None:
     raise RuntimeError(
         "[CP_HICACHE_FAILFAST][layer_page_first_page_buffer_meta_unsupported] "
@@ -315,6 +377,9 @@ class HostKVCache(abc.ABC):
             )
 
         self.kv_buffer = self.init_kv_buffer()
+        self._active_load_descriptor: Optional[PreparedLoadDescriptor] = None
+        self._missing_load_descriptor_warning_emitted = False
+        self._missing_tai_prepared_h2d_warning_emitted = False
 
         # A lock for synchronized operations on memory allocation and state transitions.
         self.lock = threading.RLock()
@@ -341,9 +406,151 @@ class HostKVCache(abc.ABC):
         self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
     ) -> None:
         """Prepare layer-invariant metadata for one host->device load op."""
+        host_indices = host_indices.contiguous()
+        device_indices = device_indices.contiguous()
+        host_page_indices, device_page_indices = self._prepare_load_page_indices(
+            host_indices, device_indices
+        )
+        num_tokens = int(host_indices.numel())
+        self._active_load_descriptor = PreparedLoadDescriptor(
+            host_indices=host_indices,
+            device_indices=device_indices,
+            host_page_indices=host_page_indices,
+            device_page_indices=device_page_indices,
+            page_size=int(self.page_size),
+            io_backend=io_backend,
+            layout=self.layout,
+            num_tokens=num_tokens,
+            num_pages=(num_tokens // int(self.page_size))
+            if num_tokens % int(self.page_size) == 0
+            else 0,
+        )
+        self._active_load_descriptor.tai_h2d_descriptor = (
+            self._try_prepare_tai_h2d_descriptor(
+                self._active_load_descriptor,
+                host_indices=host_indices,
+                device_indices=device_indices,
+                page_size=int(self.page_size),
+            )
+        )
 
     def end_load_to_device_op(self) -> None:
         """Release metadata prepared by ``begin_load_to_device_op``."""
+        descriptor = getattr(self, "_active_load_descriptor", None)
+        if descriptor is not None:
+            self._destroy_tai_h2d_descriptor(descriptor.tai_h2d_descriptor)
+            if descriptor.tai_index_h2d_descriptor is not descriptor.tai_h2d_descriptor:
+                self._destroy_tai_h2d_descriptor(descriptor.tai_index_h2d_descriptor)
+        self._active_load_descriptor = None
+
+    def _prepare_load_page_indices(
+        self, host_indices: torch.Tensor, device_indices: torch.Tensor
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        if self.layout not in ("page_first_direct", "layer_page_first"):
+            return None, None
+        validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
+        validate_page_aligned_token_indices(
+            device_indices, self.page_size, "device_indices"
+        )
+        host_page_indices = (
+            host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
+        )
+        device_page_indices = (
+            device_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
+        )
+        return host_page_indices.contiguous(), device_page_indices.contiguous()
+
+    def _get_active_load_descriptor(
+        self, io_backend: str
+    ) -> Optional[PreparedLoadDescriptor]:
+        descriptor = getattr(self, "_active_load_descriptor", None)
+        if descriptor is None:
+            self._log_missing_load_descriptor_once(io_backend)
+            return None
+        if descriptor.io_backend != io_backend or descriptor.layout != self.layout:
+            raise RuntimeError(
+                "[CP_HICACHE_FAILFAST][load_descriptor_contract_mismatch] "
+                f"active descriptor io/layout=({descriptor.io_backend}, {descriptor.layout}) "
+                f"but current load uses ({io_backend}, {self.layout})"
+            )
+        return descriptor
+
+    def _log_missing_load_descriptor_once(self, io_backend: str) -> None:
+        if io_backend != "direct":
+            return
+        if self.layout not in ("page_first_direct", "layer_page_first"):
+            return
+        if getattr(self, "_missing_load_descriptor_warning_emitted", False):
+            return
+        self._missing_load_descriptor_warning_emitted = True
+        logger.warning(
+            "[CP_HICACHE_FALLBACK][load_descriptor] "
+            "reason=missing_prepared_descriptor pool=%s layout=%s io_backend=%s",
+            type(self).__name__,
+            self.layout,
+            io_backend,
+        )
+
+    def _try_prepare_tai_h2d_descriptor(
+        self,
+        descriptor: PreparedLoadDescriptor,
+        *,
+        host_indices: torch.Tensor,
+        device_indices: torch.Tensor,
+        page_size: int,
+    ) -> Optional[object]:
+        if descriptor.io_backend != "direct":
+            return None
+        if descriptor.layout not in ("page_first_direct", "layer_page_first"):
+            return None
+        try:
+            prepare = _load_tai_prepare_h2d_page_descriptor()
+        except RuntimeError as exc:
+            self._log_missing_tai_prepared_h2d_once(str(exc))
+            return None
+        try:
+            tai_descriptor = prepare(
+                host_indices,
+                device_indices,
+                page_size=page_size,
+                layout=descriptor.layout,
+            )
+            if not bool(getattr(tai_descriptor, "cpp_prepared", False)):
+                self._log_missing_tai_prepared_h2d_once(
+                    "tai prepared descriptor Python API is present, but the "
+                    "C++ prepared submit data is unavailable; legacy per-layer "
+                    "direct submit remains active."
+                )
+            return tai_descriptor
+        except Exception as exc:
+            raise RuntimeError(
+                "[CP_HICACHE_FAILFAST][tai_prepared_h2d_descriptor_failed] "
+                f"pool={type(self).__name__} layout={descriptor.layout} "
+                f"page_size={page_size} tokens={int(host_indices.numel())}"
+            ) from exc
+
+    def _log_missing_tai_prepared_h2d_once(self, message: str) -> None:
+        if getattr(self, "_missing_tai_prepared_h2d_warning_emitted", False):
+            return
+        self._missing_tai_prepared_h2d_warning_emitted = True
+        logger.warning(
+            "[CP_HICACHE_FALLBACK][tai_prepared_h2d_descriptor] "
+            "reason=missing_tai_api %s",
+            message,
+        )
+
+    def _destroy_tai_h2d_descriptor(self, tai_descriptor: Optional[object]) -> None:
+        if tai_descriptor is None:
+            return
+        try:
+            destroy = _load_tai_destroy_h2d_page_descriptor()
+            destroy(tai_descriptor)
+        except Exception as exc:
+            logger.warning(
+                "[CP_HICACHE_FALLBACK][tai_prepared_h2d_descriptor] "
+                "reason=destroy_failed error=%s",
+                exc,
+            )
 
     @abc.abstractmethod
     def backup_from_device_per_layer(
@@ -1503,6 +1710,10 @@ class MLATokenToKVPoolHost(HostKVCache):
     def load_to_device_per_layer(
         self, device_pool, host_indices, device_indices, layer_id, io_backend
     ):
+        descriptor = self._get_active_load_descriptor(io_backend)
+        if descriptor is not None:
+            host_indices = descriptor.host_indices
+            device_indices = descriptor.device_indices
         if io_backend == "kernel":
             if self.layout == "layer_first":
                 transfer_kv_per_layer_mla(
@@ -1534,6 +1745,14 @@ class MLATokenToKVPoolHost(HostKVCache):
                     page_size=self.page_size,
                 )
             elif self.layout == "page_first_direct":
+                if descriptor is not None and descriptor.tai_h2d_descriptor is not None:
+                    _load_tai_submit_h2d_layer()(
+                        descriptor.tai_h2d_descriptor,
+                        [self.kv_buffer],
+                        [device_pool.kv_buffer[layer_id]],
+                        layer_id=layer_id,
+                    )
+                    return
                 _load_tai_transfer_kv_per_layer_direct_pf_lf()(
                     src_ptrs=[self.kv_buffer],
                     dst_ptrs=[device_pool.kv_buffer[layer_id]],
@@ -1543,6 +1762,14 @@ class MLATokenToKVPoolHost(HostKVCache):
                     page_size=self.page_size,
                 )
             elif self.layout == "layer_page_first":
+                if descriptor is not None and descriptor.tai_h2d_descriptor is not None:
+                    _load_tai_submit_h2d_layer()(
+                        descriptor.tai_h2d_descriptor,
+                        [self.kv_buffer],
+                        [device_pool.kv_buffer[layer_id]],
+                        layer_id=layer_id,
+                    )
+                    return
                 _load_tai_transfer_kv_per_layer_direct_lpf_lf()(
                     src_ptrs=[self.kv_buffer],
                     dst_ptrs=[device_pool.kv_buffer[layer_id]],
@@ -2033,13 +2260,32 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
     def begin_load_to_device_op(
         self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
     ) -> None:
+        super().begin_load_to_device_op(host_indices, device_indices, io_backend)
         self._active_load_indexer_page_indices = None
-        self._active_load_indexer_page_indices = self._get_indexer_page_indices(
-            host_indices, device_indices
+        descriptor = self._get_active_load_descriptor(io_backend)
+        host_page_indices, device_page_indices = self._get_indexer_page_indices(
+            descriptor.host_indices,
+            descriptor.device_indices,
+        )
+        descriptor.index_host_page_indices = host_page_indices
+        descriptor.index_device_page_indices = device_page_indices
+        descriptor.index_active_layer_ids = tuple(
+            getattr(self, "index_active_layer_ids", ())
+        )
+        descriptor.tai_index_h2d_descriptor = self._try_prepare_tai_h2d_descriptor(
+            descriptor,
+            host_indices=host_page_indices,
+            device_indices=device_page_indices,
+            page_size=1,
+        )
+        self._active_load_indexer_page_indices = (
+            descriptor.index_host_page_indices,
+            descriptor.index_device_page_indices,
         )
 
     def end_load_to_device_op(self) -> None:
         self._active_load_indexer_page_indices = None
+        super().end_load_to_device_op()
 
     def _load_indexer_to_device_per_layer(
         self, device_pool, host_indices, device_indices, layer_id, io_backend
@@ -2048,13 +2294,26 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
             return
         device_layer_slot = self._device_index_layer_slot(device_pool, layer_id)
         host_layer_slot = self._host_index_layer_slot(layer_id)
-        prepared_indices = getattr(self, "_active_load_indexer_page_indices", None)
-        if prepared_indices is None:
+        descriptor = self._get_active_load_descriptor(io_backend)
+        if (
+            descriptor is not None
+            and descriptor.index_host_page_indices is not None
+            and descriptor.index_device_page_indices is not None
+        ):
+            host_page_indices = descriptor.index_host_page_indices
+            device_page_indices = descriptor.index_device_page_indices
+        else:
+            prepared_indices = getattr(self, "_active_load_indexer_page_indices", None)
+            if prepared_indices is None:
+                host_page_indices, device_page_indices = self._get_indexer_page_indices(
+                    host_indices, device_indices
+                )
+            else:
+                host_page_indices, device_page_indices = prepared_indices
+        if host_page_indices is None or device_page_indices is None:
             host_page_indices, device_page_indices = self._get_indexer_page_indices(
                 host_indices, device_indices
             )
-        else:
-            host_page_indices, device_page_indices = prepared_indices
         use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
         if use_kernel:
             if self.layout == "layer_first":
@@ -2089,6 +2348,17 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
                     page_size=1,
                 )
             elif self.layout == "page_first_direct":
+                if (
+                    descriptor is not None
+                    and descriptor.tai_index_h2d_descriptor is not None
+                ):
+                    _load_tai_submit_h2d_layer()(
+                        descriptor.tai_index_h2d_descriptor,
+                        [self.index_k_with_scale_buffer],
+                        [device_pool.index_k_with_scale_buffer[device_layer_slot]],
+                        layer_id=host_layer_slot,
+                    )
+                    return
                 _load_tai_transfer_kv_per_layer_direct_pf_lf()(
                     src_ptrs=[self.index_k_with_scale_buffer],
                     dst_ptrs=[
@@ -2100,6 +2370,17 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
                     page_size=1,
                 )
             elif self.layout == "layer_page_first":
+                if (
+                    descriptor is not None
+                    and descriptor.tai_index_h2d_descriptor is not None
+                ):
+                    _load_tai_submit_h2d_layer()(
+                        descriptor.tai_index_h2d_descriptor,
+                        [self.index_k_with_scale_buffer],
+                        [device_pool.index_k_with_scale_buffer[device_layer_slot]],
+                        layer_id=host_layer_slot,
+                    )
+                    return
                 _load_tai_transfer_kv_per_layer_direct_lpf_lf()(
                     src_ptrs=[self.index_k_with_scale_buffer],
                     dst_ptrs=[
diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py
index 2ef383fbb..7a6975aa0 100644
--- a/test/registered/unit/managers/test_hicache_controller_cp.py
+++ b/test/registered/unit/managers/test_hicache_controller_cp.py
@@ -113,6 +113,7 @@ from sglang.srt.mem_cache.memory_pool_host import (
     MHATokenToKVPoolHost,
     MLATokenToKVPoolHost,
     NSATokenToKVPoolHost,
+    PreparedLoadDescriptor,
 )
 from sglang.srt.mem_cache.radix_cache import TreeNode
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -206,6 +207,207 @@ class DummyHostKVCacheForAlloc(HostKVCache):
         pass
 
 
+class TestPreparedLoadDescriptor(CustomTestCase):
+    def test_host_begin_load_builds_page_aligned_descriptor(self):
+        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
+        host_pool.page_size = 4
+        host_pool.layout = "page_first_direct"
+
+        host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
+        device_indices = torch.tensor(
+            [40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
+        )
+
+        host_pool.begin_load_to_device_op(
+            host_indices, device_indices, io_backend="direct"
+        )
+
+        desc = host_pool._active_load_descriptor
+        self.assertIsInstance(desc, PreparedLoadDescriptor)
+        self.assertTrue(torch.equal(desc.host_indices, host_indices))
+        self.assertTrue(torch.equal(desc.device_indices, device_indices))
+        self.assertEqual(desc.num_tokens, 8)
+        self.assertEqual(desc.num_pages, 2)
+        self.assertEqual(desc.layout, "page_first_direct")
+        self.assertEqual(desc.io_backend, "direct")
+        self.assertEqual(desc.host_page_indices.tolist(), [2, 5])
+        self.assertEqual(desc.device_page_indices.tolist(), [10, 16])
+
+        host_pool.end_load_to_device_op()
+        self.assertIsNone(host_pool._active_load_descriptor)
+
+    def test_nsa_begin_load_attaches_indexer_pages_to_descriptor(self):
+        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
+        host_pool.page_size = 4
+        host_pool.layout = "page_first_direct"
+        host_pool.index_active_layer_ids = (0, 2)
+
+        host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
+        device_indices = torch.tensor(
+            [40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
+        )
+
+        host_pool.begin_load_to_device_op(
+            host_indices, device_indices, io_backend="direct"
+        )
+
+        desc = host_pool._active_load_descriptor
+        self.assertEqual(desc.index_active_layer_ids, (0, 2))
+        self.assertEqual(desc.index_host_page_indices.tolist(), [2, 5])
+        self.assertEqual(desc.index_device_page_indices.tolist(), [10, 16])
+        self.assertIs(host_pool._active_load_indexer_page_indices[0], desc.index_host_page_indices)
+        self.assertIs(
+            host_pool._active_load_indexer_page_indices[1],
+            desc.index_device_page_indices,
+        )
+
+        host_pool.end_load_to_device_op()
+        self.assertIsNone(host_pool._active_load_descriptor)
+        self.assertIsNone(host_pool._active_load_indexer_page_indices)
+
+    def test_missing_direct_load_descriptor_warns_once(self):
+        host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
+        host_pool.page_size = 4
+        host_pool.layout = "page_first_direct"
+
+        with self.assertLogs(
+            "sglang.srt.mem_cache.memory_pool_host", level="WARNING"
+        ) as logs:
+            self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
+            self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
+
+        warnings = [
+            line for line in logs.output if "missing_prepared_descriptor" in line
+        ]
+        self.assertEqual(len(warnings), 1)
+
+    def test_mla_direct_load_uses_prepared_tai_descriptor_when_available(self):
+        calls = []
+        fake_desc = object()
+
+        def fake_prepare(src_indices, dst_indices, **kwargs):
+            calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
+            return fake_desc
+
+        def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
+            calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))
+
+        def fake_destroy(desc):
+            calls.append(("destroy", desc))
+
+        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
+        host_pool.layout = "page_first_direct"
+        host_pool.page_size = 4
+        host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
+        device_pool = type("DevicePool", (), {})()
+        device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
+        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
+        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
+
+        with (
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
+                return_value=fake_prepare,
+            ),
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
+                return_value=fake_submit,
+            ),
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_destroy_h2d_page_descriptor",
+                return_value=fake_destroy,
+            ),
+        ):
+            host_pool.begin_load_to_device_op(
+                host_indices, device_indices, io_backend="direct"
+            )
+            host_pool.load_to_device_per_layer(
+                device_pool,
+                host_indices,
+                device_indices,
+                layer_id=2,
+                io_backend="direct",
+            )
+            host_pool.end_load_to_device_op()
+
+        self.assertEqual(calls[0][0], "prepare")
+        self.assertEqual(calls[0][1].tolist(), [4, 5, 6, 7])
+        self.assertEqual(calls[0][2].tolist(), [12, 13, 14, 15])
+        self.assertEqual(calls[0][3], {"page_size": 4, "layout": "page_first_direct"})
+        self.assertEqual(calls[1][0], "submit")
+        self.assertIs(calls[1][1], fake_desc)
+        self.assertEqual(calls[1][2][0].data_ptr(), host_pool.kv_buffer.data_ptr())
+        self.assertEqual(
+            calls[1][3][0].data_ptr(), device_pool.kv_buffer[2].data_ptr()
+        )
+        self.assertEqual(calls[1][4], {"layer_id": 2})
+        self.assertEqual(calls[2], ("destroy", fake_desc))
+
+    def test_nsa_index_direct_load_uses_prepared_tai_index_descriptor(self):
+        calls = []
+        fake_base_desc = object()
+        fake_index_desc = object()
+
+        def fake_prepare(src_indices, dst_indices, **kwargs):
+            calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
+            if kwargs["page_size"] == 1:
+                return fake_index_desc
+            return fake_base_desc
+
+        def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
+            calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))
+
+        host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
+        host_pool.layout = "page_first_direct"
+        host_pool.page_size = 4
+        host_pool.index_active_layer_ids = (0, 1, 2)
+        host_pool.index_k_with_scale_buffer = torch.empty(
+            (8, 3, 1, 32), dtype=torch.uint8
+        )
+        device_pool = type("DevicePool", (), {})()
+        device_pool.index_k_with_scale_buffer = torch.empty(
+            (3, 8, 32), dtype=torch.uint8
+        )
+        host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
+        device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
+
+        with (
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
+                return_value=fake_prepare,
+            ),
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
+                return_value=fake_submit,
+            ),
+        ):
+            host_pool.begin_load_to_device_op(
+                host_indices, device_indices, io_backend="direct"
+            )
+            host_pool._load_indexer_to_device_per_layer(
+                device_pool,
+                host_indices,
+                device_indices,
+                layer_id=1,
+                io_backend="direct",
+            )
+
+        self.assertEqual([c[0] for c in calls[:2]], ["prepare", "prepare"])
+        self.assertEqual(calls[1][1].tolist(), [1])
+        self.assertEqual(calls[1][2].tolist(), [3])
+        self.assertEqual(calls[1][3], {"page_size": 1, "layout": "page_first_direct"})
+        self.assertEqual(calls[2][0], "submit")
+        self.assertIs(calls[2][1], fake_index_desc)
+        self.assertEqual(
+            calls[2][2][0].data_ptr(), host_pool.index_k_with_scale_buffer.data_ptr()
+        )
+        self.assertEqual(
+            calls[2][3][0].data_ptr(),
+            device_pool.index_k_with_scale_buffer[1].data_ptr(),
+        )
+        self.assertEqual(calls[2][4], {"layer_id": 1})
+
+
 class FakeDevicePool:
     device = "cpu"
     layer_num = 1
@@ -878,9 +1080,15 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
 
         host_pool._get_indexer_page_indices = counting_getter
 
-        with patch(
-            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
-            return_value=fake_direct,
+        with (
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
+                side_effect=RuntimeError("missing prepared descriptor api"),
+            ),
+            patch(
+                "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
+                return_value=fake_direct,
+            ),
         ):
             host_pool.begin_load_to_device_op(
                 host_indices, device_indices, io_backend="direct"
@@ -1898,6 +2106,44 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         self.assertEqual(host_pool.end_calls, 1)
         self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])
 
+    def test_start_loading_emits_descriptor_timing_when_enabled(self):
+        host_pool = FakeHostPool(
+            torch.tensor([100, 101, 102, 103], dtype=torch.int64)
+        )
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
+        allocator.device_pool = FakeDevicePool(layer_num=2)
+        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
+        node = TreeNode()
+        node.host_len = 16
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=16,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
+        )
+        controller.load_cp([node], node_id=114)
+        timing_keys = []
+
+        def record_timing(key, start_time, message, *args):
+            timing_keys.append(key)
+
+        with patch(
+            "sglang.srt.managers.cache_controller.envs."
+            "SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get",
+            return_value=True,
+        ), patch(
+            "sglang.srt.managers.cache_controller."
+            "_cp_shared_kv_bs_gt1_cache_timing",
+            side_effect=record_timing,
+        ):
+            controller.start_loading()
+
+        self.assertIn("prepare_load_descriptor", timing_keys)
+        self.assertIn("submit_h2d_layer_loop", timing_keys)
+        self.assertIn("submit_h2d_layer_per_call_slow", timing_keys)
+        self.assertIn("end_load_descriptor", timing_keys)
+
     def test_cp_load_frees_unexpected_owner_allocator_length(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))