# NSA Prefill CP HiCache layer_page_first direct 支持计划

## 目标

在 CP shared KV + HiCache 场景下，为 `--hicache-io-backend direct` 增加 `--hicache-mem-layout layer_page_first`。目标是让 per-layer L2(host) <-> L1(device) backup/load 使用 TAI `cudaMemcpyBatchAsync` direct path，减少 page-first-direct 在固定 layer 访问时的 descriptor 数量和提交开销。

当前计划只覆盖 host-only CP HiCache，也就是现有 `enable_nsa_prefill_cp_shared_kv` 路径。当前 CP shared KV 已禁止 `hicache_storage_backend`，因此 L3 / Mooncake Store / NIXL / HF3FS zero-copy 不作为第一阶段支持目标；遇到这些组合必须 fail-fast，不能静默 fallback。

## 背景结论

### 当前 page_first_direct 布局

当前 direct 生产配置使用：

```text
--hicache-io-backend direct
--hicache-mem-layout page_first_direct
```

host KV layout：

```text
MLA/NSA KV:      [page, layer, page_size, 1, kv_cache_dim]
MHA K/V:         [2, page, layer, page_size, head_num, head_dim]
NSA indexer:     [page, layer, 1, indexer_page_stride_size]
```

固定 layer 做 per-layer transfer 时，连续 page 的同一 layer 数据在 host 内存中被其他 layer 隔开。即使 page id 连续，也容易形成多个 memcpy descriptor。

### 目标 layer_page_first 布局

目标 host layout：

```text
MLA/NSA KV:      [layer, page, page_size, 1, kv_cache_dim]
MHA K/V:         [2, layer, page, page_size, head_num, head_dim]
NSA indexer:     [layer, page, 1, indexer_page_stride_size]
```

固定 layer 下，连续 host page 可以成为连续内存 run；TAI LPF direct path 会把每个连续 page run 合并成一个 `cudaMemcpyBatchAsync` descriptor。

### CP zigzag 与 rank-local physical page

CP request 的 logical page owner 仍由 zigzag / owner 规则决定；但 backup/load 进入 HostKVCache 前会把 logical token loc 映射到 rank-local physical loc：

```python
physical_page = (logical_page - 1) // cp_size + 1
```

相关代码：

- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py:79-101`
- `python/sglang/srt/managers/cache_controller.py:1002-1003`

因此对于单个 rank，本地 physical pages 通常是 compact 的。LPF 收益主要取决于 rank-local host pages 和 device pages 是否形成连续 run，而不是 global logical page 是否 zigzag。

### 已有 benchmark 证据

历史文档 `nsa_prefill_cp_page_aligned_cache_contract.md` 的 C115/C116 记录了远端 benchmark：

- 192k BF16 fragmented H2D：`page_first 6.290 ms` -> `LPF same 4.739 ms`，约 `1.33x`。
- 192k BF16 fragmented D2H：`page_first 6.226 ms` -> `LPF same 4.724 ms`，约 `1.32x`。
- 192k FP8 fragmented H2D：`page_first 4.612 ms` -> `LPF same 2.737 ms`，约 `1.68x`。
- 192k FP8 fragmented D2H：`page_first 4.560 ms` -> `LPF same 2.709 ms`，约 `1.68x`。

结论：LPF direct 对 per-layer HiCache transfer 有实际收益，FP8 下收益更明显。但收益必须绑定 host extent allocation；如果 host pages 完全 random，LPF same 接近 neutral，compact LPF 才稳定获益。

## 当前代码状态

### TAI kernel 已具备 LPF direct API

`tai-kernel` 已存在：

- `tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lf_lpf`
  - D2H backup，device layer-first -> host `[layer, page, page_size, ...]`。
- `tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lpf_lf`
  - H2D load，host `[layer, page, page_size, ...]` -> device layer-first。

代码位置：

- `tai-kernel/python/tai_kernel/nsa_prefill/kvcacheio.py:83-111`
- `tai-kernel/python/tai_kernel/nsa_prefill/kvcacheio.py:251-276`
- `tai-kernel/python/tai_kernel/nsa_prefill/csrc/kvcacheio_lf_pf.cu:982-1158`
- `tai-kernel/python/tai_kernel/nsa_prefill/csrc/kvcacheio_lf_pf.cu:1160-1336`

约束：

- `src_indices` / `dst_indices` 必须是 CPU int64 contiguous。
- indices 必须是完整 page spans。
- 依赖 CUDA `cudaMemcpyBatchAsync`，缺失时 fail-fast。
- 不允许隐藏 CUDA index -> CPU index 拷贝。

### SGLang 尚未接入 layer_page_first

当前缺口：

1. `server_args.py` 未允许 `layer_page_first`。
   - `--hicache-mem-layout` choices 缺失。
   - CP HiCache validation 缺失 `("direct", "layer_page_first")`。
   - direct/page_first 自动归一逻辑未覆盖 LPF。

2. `memory_pool_host.py` 未定义 LPF buffer layout。
   - `MHATokenToKVPoolHost.init_kv_buffer()` 不支持。
   - `MLATokenToKVPoolHost.init_kv_buffer()` 不支持。
   - `NSATokenToKVPoolHost._init_indexer_buffers()` 不支持。

3. direct transfer dispatch 未接 LPF。
   - MHA `load_to_device_per_layer()` / `backup_from_device_per_layer()` / `backup_from_device_all_layer()` 缺失 LPF 分支。
   - MLA 同上。
   - NSA indexer `_load_indexer_to_device_per_layer()` / `_backup_indexer_from_device_per_layer()` / `_backup_indexer_from_device_all_layer()` 缺失 LPF 分支。

4. `HiCacheController.move_indices()` 未允许 direct + LPF。
   - 当前 direct 只接受 `layer_first` 和 `page_first_direct`。
   - LPF 应与 `page_first_direct` 一样返回 `host_indices` CPU、`device_indices.cpu()`，不能排序 host indices。

5. page-level storage metadata 不适合直接复用。
   - `get_page_buffer_meta()` 当前对 page-first 假设一个 page 的所有 layer 连续。
   - LPF 下一个 page 跨 layer 不连续，不能返回一个 page-level pointer + full-page element size。
   - 第一阶段应在 storage backend / zero-copy path fail-fast，不支持 LPF。

## 设计选择

### 选择 A：只支持 CP host-only direct + LPF

这是推荐方案。

范围：

```text
CP shared KV + HiCache + direct + layer_page_first + no hicache_storage_backend
```

行为：

- per-layer backup/load 使用 TAI LPF direct API。
- all-layer backup 仍循环 per-layer TAI LPF API，避免 sgl-kernel CUDA 13 direct ABI 风险。
- storage backend / zero-copy path 遇到 LPF 直接 fail-fast。
- `kernel + layer_page_first` 不支持。
- `direct + layer_page_first` 缺少 TAI kernel 时直接报错，不 fallback 到 SM copy。

优点：

- 与当前线上 CP host-only HiCache 目标一致。
- 改动边界清晰。
- 不破坏 Mooncake/NIXL/HF3FS page-level zero-copy contract。
- 最快验证 LPF 对 per-layer backup/load 的收益。

缺点：

- 暂时不能用于通用 HiCache storage backend。
- 需要新增 unit / remote CUDA / ETE 验证。

### 选择 B：同时重写 storage page metadata 支持 LPF

不推荐第一阶段做。

原因：

- LPF 下 page 跨 layer 不是单段连续内存；storage zero-copy 的 key/ptr/size contract 要从 page-level 改成 layer-page-level。
- Mooncake Store / NIXL / HF3FS 都需要跟着改。
- 风险远大于当前 CP host-only 目标。

### 选择 C：只保留 page_first_direct，不接 LPF

不推荐。

原因：

- 已有 TAI benchmark 显示 LPF 在 fragmented/contiguous rank-local pages 下有明确收益。
- 当前 per-layer backup/load 已经是主要优化方向，host layout 应该匹配 per-layer access pattern。

## 实现计划

### P0：补 server args 与 validation

修改文件：

- `python/sglang/srt/server_args.py`
- `test/registered/unit/server_args/test_server_args.py`

步骤：

1. 在 `--hicache-mem-layout` choices 中加入 `layer_page_first`。
2. 在 `_handle_cp_hicache_layout_validation()` 中加入支持对：

```python
("direct", "layer_page_first")
```

3. 保持 `kernel + layer_page_first` 不支持，报错信息明确写出仅 direct 支持。
4. 保持 `direct + page_first` 自动归一到 `page_first_direct`，不要自动改成 LPF。
5. 新增单测：
   - direct + layer_page_first + CP shared KV + HiCache 通过 validation。
   - kernel + layer_page_first 报错。
   - non-CP 普通启动如果指定 LPF，仍由 HostKVCache 初始化或 storage validation 负责报错；不要在 server_args 中误改旧行为。

验收命令：

```bash
PYTHONPATH=python python -m pytest -q test/registered/unit/server_args/test_server_args.py
```

### P1：补 TAI LPF loader 和 fail-fast 文案

修改文件：

- `python/sglang/srt/mem_cache/memory_pool_host.py`

新增 loader：

```python
_load_tai_transfer_kv_per_layer_direct_lf_lpf()
_load_tai_transfer_kv_per_layer_direct_lpf_lf()
```

要求：

- 缺失 TAI API 时抛 `RuntimeError`。
- 错误前缀使用醒目的 fail-fast：

```text
[CP_HICACHE_FAILFAST][missing_tai_layer_page_first_direct_lf_lpf]
[CP_HICACHE_FAILFAST][missing_tai_layer_page_first_direct_lpf_lf]
```

- 明确说明不 fallback 到 SM-consuming kernel。

单测：

- patch import 失败路径或 patch loader 返回 fake direct op，验证 LPF 分支调用正确 loader。

### P2：补 host KV buffer layout

修改文件：

- `python/sglang/srt/mem_cache/memory_pool_host.py`
- `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py`

新增 layout：

MHA：

```python
# current page_first_direct: [2, page, layer, page_size, head_num, head_dim]
# new layer_page_first:     [2, layer, page, page_size, head_num, head_dim]
dims = (2, self.layer_num, self.page_num, self.page_size, self.head_num, self.head_dim)
```

MLA/NSA KV：

```python
# current page_first_direct: [page, layer, page_size, 1, kv_cache_dim]
# new layer_page_first:     [layer, page, page_size, 1, kv_cache_dim]
dims = (self.layer_num, self.page_num, self.page_size, 1, self.kv_cache_dim)
```

NSA indexer：

```python
# current page_first_direct: [page, layer, 1, indexer_page_stride_size]
# new layer_page_first:     [layer, page, 1, indexer_page_stride_size]
dims = (self.layer_num, self.indexer_page_num, 1, self.indexer_page_stride_size)
```

注意：

- `get_size_per_token()` 不变；只是物理维度顺序变化。
- `page_num` / `size` 计算不变。
- target pool 和 draft pool 必须同样支持 LPF。

单测：

- 构造 MHA / MLA / NSA host pool，断言 `kv_buffer.shape` 和 `index_k_with_scale_buffer.shape` 正确。
- dtype 覆盖 bf16 与 fp8_e4m3 的 NSA pool roundtrip 测试留到 P6 remote CUDA。

### P3：补 per-layer direct load/backup dispatch

修改文件：

- `python/sglang/srt/mem_cache/memory_pool_host.py`
- `test/registered/unit/managers/test_hicache_controller_cp.py`
- `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py`

新增分支：

#### MHA load H2D

```python
elif self.layout == "layer_page_first":
    _load_tai_transfer_kv_per_layer_direct_lpf_lf()(
        src_ptrs=[self.k_buffer, self.v_buffer],
        dst_ptrs=[device_pool.k_buffer[layer_id], device_pool.v_buffer[layer_id]],
        src_indices=host_indices,
        dst_indices=device_indices,
        layer_id=layer_id,
        page_size=self.page_size,
    )
```

#### MHA backup D2H

```python
elif self.layout == "layer_page_first":
    _load_tai_transfer_kv_per_layer_direct_lf_lpf()(
        src_ptrs=[device_pool.k_buffer[layer_id], device_pool.v_buffer[layer_id]],
        dst_ptrs=[self.k_buffer, self.v_buffer],
        src_indices=device_indices,
        dst_indices=host_indices,
        layer_id=layer_id,
        page_size=self.page_size,
    )
```

#### MLA / NSA KV

同 MHA，但 `src_ptrs` / `dst_ptrs` 只有一个 `kv_buffer`。

#### NSA indexer

- H2D 使用 `_load_tai_transfer_kv_per_layer_direct_lpf_lf()`。
- D2H 使用 `_load_tai_transfer_kv_per_layer_direct_lf_lpf()`。
- `page_size=1`，indices 使用 `_get_indexer_page_indices()` 预计算结果。
- `index_k_with_scale_buffer` LPF shape 为 `[layer, page, 1, stride]`。

单测：

- 对 MHA / MLA / NSA indexer 的 LPF per-layer load/backup patch fake TAI op，验证：
  - 使用 LPF loader。
  - `src_ptrs` / `dst_ptrs` 指向正确对象。
  - `layer_id` / `page_size` 正确。
  - indexer path 传入 page indices 且 `page_size=1`。

### P4：补 all-layer backup route

修改文件：

- `python/sglang/srt/mem_cache/memory_pool_host.py`
- `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py`

要求：

- `backup_from_device_all_layer(..., io_backend="direct")` 在 `layout == "layer_page_first"` 时循环调用 per-layer LPF D2H TAI op。
- 不使用 `sgl_kernel.transfer_kv_all_layer_direct_lf_pf`。
- 继续保持 CUDA 13 安全路径。

单测：

- 参考现有 `page_first_direct_all_layer_backup_uses_tai_per_layer_route`，新增 LPF 版本。
- patch 旧 all-layer direct op 为 AssertionError，确保不会误调用。

### P5：补 move_indices direct contract

修改文件：

- `python/sglang/srt/managers/cache_controller.py`
- `test/registered/unit/managers/test_hicache_controller_cp.py`

行为：

```python
elif mem_pool_host.layout in ["page_first_direct", "layer_page_first"]:
    return host_indices, device_indices.cpu()
```

原因：

- TAI LPF direct API 要求 CPU indices。
- LPF 与 page_first_direct 一样，不需要对 host_indices 排序；排序会破坏 host/device pair 的对应关系。
- `layer_first` direct 仍保持旧逻辑：host sort + device reorder。

单测：

- direct + LPF `move_indices()` 返回 CPU device_indices，host_indices 顺序不变。
- direct + LPF 不走 layer_first 的 sort/reorder。

### P6：处理 page data / storage metadata 边界

修改文件：

- `python/sglang/srt/mem_cache/memory_pool_host.py`
- `python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py`
- `python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py`
- `python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py`
- `python/sglang/srt/mem_cache/storage/backend_factory.py`

第一阶段策略：

1. `get_data_page()` / `set_from_flat_data_page()` 可以支持 LPF，以便测试和非-zero-copy fallback 有明确语义：
   - MHA LPF page view：`kv_buffer[:, :, real_page:real_page+1, :, :, :]`，flatten 后形状仍为 `[2, layer, page_size, head, dim]` 语义。
   - MLA LPF page view：`kv_buffer[:, real_page:real_page+1, :, :, :]`，flatten 后语义为 `[layer, page_size, 1, dim]`。
2. `get_page_buffer_meta()` 对 LPF 必须 fail-fast：

```text
[CP_HICACHE_FAILFAST][layer_page_first_page_buffer_meta_unsupported]
```

原因：LPF 下一个 page 的所有 layer 不是一个 page-major contiguous block；直接返回 page-level ptr/size 会产生错误 zero-copy。

3. storage backend register 时：
   - Mooncake Store 不把 LPF 加到 supported list。
   - NIXL `is_zero_copy` 不包含 LPF。
   - HF3FS 若需要 bytes_per_page，LPF 第一阶段直接报 unsupported，避免误算。

注意：当前 CP shared KV 已禁止 `hicache_storage_backend`，因此这些边界主要防止非 CP 或未来路径静默错用。

### P7：补远端 CUDA roundtrip 测试

修改文件：

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

新增测试：

1. `test_fp8_layer_page_first_roundtrip_preserves_kv_and_indexer_pages`
   - 参考现有 `test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages`。
   - layout 改为 `layer_page_first`。
   - D2H backup 到 host，再 H2D load 到不同 device pages。
   - 验证每层 KV 和 indexer page 内容一致。

2. `test_bf16_layer_page_first_roundtrip_preserves_kv_and_indexer_pages`
   - 同样逻辑，dtype 使用 bf16。

远端命令：

```bash
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc '
cd /sgl-workspace/sglang-tai && \
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAPoolHostUnit::test_fp8_layer_page_first_roundtrip_preserves_kv_and_indexer_pages \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAPoolHostUnit::test_bf16_layer_page_first_roundtrip_preserves_kv_and_indexer_pages
'"
```

如果类名与当前文件实际类名不同，以文件内现有 `page_first_direct` roundtrip 测试所在类为准。

### P8：补 benchmark / ETE 验证

TAI benchmark 已存在：

```bash
cd /sgl-workspace/tai-kernel
PYTHONPATH=python python benchmark/nsa_prefill/benchmark_hicache_layout_compare.py \
  --layouts page_first_direct layer_page_first \
  --directions h2d d2h \
  --tokens 4096 10240 40960 122880 196608 \
  --layers 78 \
  --item-size 576 \
  --page-size 64 \
  --src-patterns contiguous fragmented random \
  --dtypes bf16 fp8 \
  --warmup 2 \
  --repeat 5
```

SGLang ETE 验证分两轮：

1. 正确性：
   - 启动 CP shared KV + HiCache + `direct + layer_page_first`。
   - 跑 GSM8K 全量两轮，第二轮 cache hit 后准确率不得明显下降。
   - 参考通过标准：`Accuracy` 接近 `0.955`，第二轮不能掉到 `~0.68` 这类 cache-hit corruption 特征。

2. 性能：
   - 使用线上 replay workload。
   - 对比 `page_first_direct` 与 `layer_page_first`。
   - 关注端到端 latency、prefill throughput、HiCache load/backup timing、是否出现 fallback/fail-fast。

启动参数差异：

```text
--hicache-io-backend direct
--hicache-mem-layout layer_page_first
```

其余参数保持当前稳定配置。

## 风险与防线

### 风险 1：storage zero-copy 误用 LPF

防线：`get_page_buffer_meta()` 对 LPF fail-fast；storage backend 不把 LPF 加入 zero-copy supported layout。

### 风险 2：indices 不在 CPU 导致隐藏 D2H copy

防线：复用 TAI fail-fast；SGLang `move_indices()` 对 LPF 明确返回 CPU device_indices。

### 风险 3：all-layer backup 误走 sgl-kernel direct path

防线：LPF all-layer backup 只循环 TAI per-layer direct op；单测 patch 旧 op 抛错。

### 风险 4：draft KV pool 漏接 LPF

防线：draft host pool 与 target host pool 共用 HostKVCache 子类；P2/P3/P4 必须覆盖 draft 可走的同一 dispatch。ETE 启动保留 `SGLANG_CP_DRAFT_SHARED_KV=1`。

### 风险 5：NSA indexer layout 漏接

防线：NSA indexer 单独支持 `[layer, page, 1, stride]`，并补 load/backup 单测和 CUDA roundtrip。

### 风险 6：host pages 过度碎片化导致 LPF 收益不明显

防线：当前已有 `alloc_contiguous_preferred()` 和 host free-room/trigger 参数。ETE 性能对比必须同时检查 host allocation/evict 日志和 transfer timing。如果 random residency 下 LPF same 只是 neutral，不能宣称收益；需要后续 host compact allocator 或 transfer-time gather/staging。

## 建议任务切分

1. **W1 server/config gate**
   - 改 server args choices / validation / 单测。

2. **W2 HostKVCache layout + dispatch**
   - 改 MHA/MLA/NSA indexer layout。
   - 接 TAI LPF direct loader。
   - 接 per-layer/all-layer load/backup。

3. **W3 storage boundary + fail-fast**
   - LPF page metadata fail-fast。
   - storage backend 不误用 zero-copy。

4. **W4 tests + remote CUDA**
   - 单元测试。
   - g0034 CUDA roundtrip。
   - TAI benchmark quick smoke。

## 进展记录

### 2026-06-10：W1 server/config gate 完成

已完成：

- `--hicache-mem-layout` parser 接受 `layer_page_first`。
- CP shared KV + HiCache validation 接受 `("direct", "layer_page_first")`。
- `kernel + layer_page_first` 保持拒绝，避免误走未实现 kernel backend。

远端验证：

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py::test_hicache_mem_layout_parser_accepts_layer_page_first \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_accepts_supported_backend_layout_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_normalizes_supported_alias_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_rejects_kernel_layer_page_first_layout
```

结果：通过。完整 `test_server_args.py` 曾有 1 个 HuggingFace/DNS 配置拉取失败，与本改动无关。

### 2026-06-10：W2 HostKVCache layout + direct dispatch 完成

已完成：

- `memory_pool_host.py` 接入 TAI LPF direct loader：
  - `_load_tai_transfer_kv_per_layer_direct_lf_lpf()`：device layer-first -> host layer-page-first。
  - `_load_tai_transfer_kv_per_layer_direct_lpf_lf()`：host layer-page-first -> device layer-first。
- MHA host layout 支持 `[2, layer, page, page_size, head, dim]`。
- MLA/NSA KV host layout 支持 `[layer, page, page_size, 1, kv_cache_dim]`。
- NSA indexer host layout 支持 `[layer, page, 1, indexer_page_stride_size]`。
- MHA/MLA/NSA indexer per-layer H2D load 和 D2H backup 接入 LPF direct TAI op。
- MHA/MLA/NSA indexer all-layer backup 在 `direct + layer_page_first` 下循环调用 TAI per-layer D2H op，不走 sgl-kernel all-layer direct path。
- `HiCacheController.move_indices()` 对 `direct + layer_page_first` 与 `page_first_direct` 保持相同 contract：
  - `host_indices` 顺序不变。
  - `device_indices` 转 CPU。
  - 不排序，避免破坏 host/device pair。

设计确认：

- LPF direct 缺少 TAI op 时 fail-fast，不 fallback 到 SM-consuming copy kernel。
- `layer_first` direct 仍保留旧的 host sort + device reorder 语义。
- storage/page metadata 仍未支持 LPF；这是 W3 边界，不能在 storage backend 中静默使用。

远端验证：

```bash
python -m py_compile \
  python/sglang/srt/mem_cache/memory_pool_host.py \
  python/sglang/srt/managers/cache_controller.py \
  python/sglang/srt/server_args.py \
  test/registered/unit/managers/test_hicache_controller_cp.py \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py \
  test/registered/unit/server_args/test_server_args.py
```

结果：`PY_COMPILE_OK`。

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel::test_mla_layer_page_first_per_layer_backup_uses_direct_lf_lpf \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel::test_mha_layer_page_first_per_layer_backup_uses_direct_lf_lpf \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel::test_nsa_indexer_layer_page_first_per_layer_backup_uses_direct_lf_lpf \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel::test_mla_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel::test_mha_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel::test_nsa_indexer_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestHiCacheControllerCPLoad::test_direct_layer_page_first_move_indices_keeps_host_order_and_cpu_device_indices \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestLayerPageFirstDirectHostLayout \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAIndexerPageIndices::test_mla_layer_page_first_all_layer_backup_uses_tai_per_layer_route \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAIndexerPageIndices::test_mha_layer_page_first_all_layer_backup_uses_tai_per_layer_route \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAIndexerPageIndices::test_layer_page_first_all_layer_indexer_backup_uses_tai_per_layer_route
```

结果：`13 passed`。

扩展相关验证：

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py::test_hicache_mem_layout_parser_accepts_layer_page_first \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_accepts_supported_backend_layout_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_normalizes_supported_alias_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_rejects_kernel_layer_page_first_layout \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestLayerPageFirstDirectHostLayout \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestPageFirstDirectAllLayerBackupRoute \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAIndexerPageIndices
```

结果：`32 passed, 5 subtests passed`。

未完成：

- W4：远端 CUDA roundtrip、TAI benchmark quick smoke、ETE GSM8K/replay 验证。

### 2026-06-10：W4 CUDA roundtrip 与 TAI benchmark quick smoke 完成

#### CUDA roundtrip

新增并远端验证了 `layer_page_first + direct` 的真实 CUDA D2H/H2D roundtrip：

- `test_fp8_layer_page_first_roundtrip_preserves_kv_and_indexer_pages`
- `test_bf16_layer_page_first_roundtrip_preserves_kv_and_indexer_pages`

验证内容：

1. 在 device L1 KV/indexer buffer 写入已知数据。
2. 通过 `backup_from_device_all_layer(..., io_backend="direct")` 写入 LPF host pool。
3. 清空目标 device pages。
4. 通过 `load_to_device_per_layer(..., io_backend="direct")` 从 LPF host pool 读回。
5. 逐 layer 比较 KV 与 indexer page 内容。

远端独立进程验证：

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_layer_page_first_roundtrip_preserves_kv_and_indexer_pages

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_bf16_layer_page_first_roundtrip_preserves_kv_and_indexer_pages

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages
```

结果：三个独立进程均 `1 passed`。

风险记录：

- 同一 pytest 进程连续跑 `page_first_direct fp8 -> layer_page_first fp8 -> layer_page_first bf16`
  时，第三个测试触发：

```text
cudaErrorHostMemoryAlreadyRegistered
part or all of the requested memory range is already mapped
```

- `CUDA_LAUNCH_BLOCKING=1` 后仍显示第三个测试初始化 CUDA tensor 时暴露该错误。
- LPF fp8/bf16 独立进程都通过，因此这不是单个 LPF layout 的数据正确性问题。
- 更像 direct path/pinned host registration 在同一进程多次 roundtrip 后的生命周期或测试隔离问题。
  生产 server 进程长时间使用 direct path 时也可能相关，后续如果出现同类 CUDA host
  registration 报错，需要优先检查 TAI direct kernel 是否重复 register 已 pinned/mapped
  host range、是否需要 unregister 或注册缓存。

#### TAI benchmark quick smoke

运行位置：

```text
g0034:/mnt/beegfs/cjy/tai-kernel/benchmark/nsa_prefill
```

输出 CSV：

```text
/mnt/beegfs/cjy/log/hicache_lpf_quick_bf16_cf.csv
/mnt/beegfs/cjy/log/hicache_lpf_quick_bf16_owner4k.csv
/mnt/beegfs/cjy/log/hicache_lpf_quick_fp8_cf.csv
/mnt/beegfs/cjy/log/hicache_lpf_quick_fp8_owner4k.csv
```

参数：

- `directions`: H2D, D2H
- `layouts`: `page_first_direct`, `layer_page_first`
- `layout_page_modes`: `same`, `compact`
- `tokens`: 4k, 40k for contiguous/fragmented
- `tokens`: 4k for owner_lane
- `layers`: 78
- `page_size`: 64
- `item_size`: 656
- `dtype`: bf16, fp8_e4m3fn
- `warmup=1`, `repeat=2`

注意：

- `fragmented` 40k 需要约 2x page pool；`pool_tokens=65536` 不够，已改为
  `pool_tokens=131072`。
- `owner_lane` 40k 按 `cp_size=8` stride 需要更大 pool；quick smoke 只跑 4k。

关键结果摘录：

BF16，40k contiguous：

```text
H2D page_first_direct: total_median_ms=1.752396, effective_gbps=30.666
H2D layer_page_first same: total_median_ms=1.214203, effective_gbps=44.259
H2D layer_page_first compact: total_median_ms=1.207671, effective_gbps=44.498
D2H page_first_direct: total_median_ms=1.708416, effective_gbps=31.456
D2H layer_page_first same: total_median_ms=1.200778, effective_gbps=44.754
D2H layer_page_first compact: total_median_ms=1.198237, effective_gbps=44.849
```

BF16，40k fragmented：

```text
H2D page_first_direct: total_median_ms=1.734709, effective_gbps=30.979
H2D layer_page_first same: total_median_ms=1.267374, effective_gbps=42.402
H2D layer_page_first compact: total_median_ms=1.206922, effective_gbps=44.526
D2H page_first_direct: total_median_ms=1.720521, effective_gbps=31.234
D2H layer_page_first same: total_median_ms=1.267231, effective_gbps=42.407
D2H layer_page_first compact: total_median_ms=1.199161, effective_gbps=44.814
```

FP8，40k contiguous：

```text
H2D page_first_direct: total_median_ms=1.071615, effective_gbps=25.074
H2D layer_page_first same: total_median_ms=0.657684, effective_gbps=40.855
H2D layer_page_first compact: total_median_ms=0.659018, effective_gbps=40.772
D2H page_first_direct: total_median_ms=1.023117, effective_gbps=26.263
D2H layer_page_first same: total_median_ms=0.656052, effective_gbps=40.957
D2H layer_page_first compact: total_median_ms=0.655805, effective_gbps=40.972
```

FP8，40k fragmented：

```text
H2D page_first_direct: total_median_ms=1.063277, effective_gbps=25.271
H2D layer_page_first same: total_median_ms=0.710445, effective_gbps=37.821
H2D layer_page_first compact: total_median_ms=0.659661, effective_gbps=40.733
D2H page_first_direct: total_median_ms=1.020919, effective_gbps=26.319
D2H layer_page_first same: total_median_ms=0.708986, effective_gbps=37.899
D2H layer_page_first compact: total_median_ms=0.660234, effective_gbps=40.697
```

4k owner_lane：

- BF16:
  - H2D page_first_direct `0.281816 ms` vs LPF compact `0.171370 ms`
  - D2H page_first_direct `0.225308 ms` vs LPF compact `0.171644 ms`
- FP8:
  - H2D page_first_direct `0.227589 ms` vs LPF compact `0.124067 ms`
  - D2H page_first_direct `0.168431 ms` vs LPF compact `0.121730 ms`

结论：

- TAI LPF H2D/D2H kernel 可运行。
- 40k contiguous/fragmented 下 LPF 明显优于 `page_first_direct`。
- owner_lane same 由于物理 pages 仍是 stride，收益有限；compact LPF 明显更好。
  因此生产要吃满 LPF 收益，仍需要 host allocator 尽量给 CP owner/rank 分配
  compact physical pages。

仍未完成：

- ETE：启动 SGLang `--hicache-mem-layout layer_page_first` 后跑 GSM8K/replay。

### 2026-06-10：W3 storage/page metadata fail-fast 边界完成

用户确认：storage 暂时不开启，`layer_page_first` 第一阶段不支持 storage。

已完成：

- `ServerArgs._resolve_storage_layout_compatibility()` 在 `hicache_storage_backend != None`
  且 `hicache_mem_layout == "layer_page_first"` 时直接 fail-fast。
- `HiCacheStorage.register_mem_pool_host()` 在任何 storage backend 尝试注册
  `layer_page_first` host pool 时 fail-fast，覆盖运行时 attach storage 的绕过路径。
- `MHATokenToKVPoolHost.get_page_buffer_meta()` 和
  `MLATokenToKVPoolHost.get_page_buffer_meta()` 对 `layer_page_first` 显式
  fail-fast：
  - 错误前缀：`[CP_HICACHE_FAILFAST][layer_page_first_page_buffer_meta_unsupported]`
  - 原因：LPF 下“一个 page 跨所有 layer”的内容不是 page-major contiguous block，
    不能复用 storage/zero-copy 的 page-level pointer + size contract。

设计确认：

- 保留 `get_data_page()` / `set_from_flat_data_page()` 的 LPF 支持，用于非
  storage 测试/调试语义；但 storage 注册和 page metadata 入口不允许进入。
- 不把 `layer_page_first` 加入 Mooncake/NIXL/HF3FS zero-copy supported list。
- 不做慢 fallback 到 flatten/copy storage，因为这会隐藏生产路径配置错误。

远端 RED 验证：

新增 4 个测试在实现前失败：

- storage + `layer_page_first` server args 未报错。
- MHA/MLA `get_page_buffer_meta()` 只是 generic unsupported layout。
- `HiCacheStorage.register_mem_pool_host()` 允许 LPF storage 注册。

远端 GREEN 验证：

```bash
python -m py_compile \
  python/sglang/srt/server_args.py \
  python/sglang/srt/mem_cache/memory_pool_host.py \
  python/sglang/srt/mem_cache/hicache_storage.py \
  test/registered/unit/server_args/test_server_args.py \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py
```

结果：`PY_COMPILE_OK`。

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_hicache_storage_rejects_layer_page_first_layout \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestLayerPageFirstDirectHostLayout::test_mha_layer_page_first_page_buffer_meta_fails_fast_for_storage \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestLayerPageFirstDirectHostLayout::test_mla_layer_page_first_page_buffer_meta_fails_fast_for_storage \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestLayerPageFirstDirectHostLayout::test_storage_registration_fails_fast_for_layer_page_first
```

结果：`4 passed`。

扩展相关验证：

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py::test_hicache_mem_layout_parser_accepts_layer_page_first \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_accepts_supported_backend_layout_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_normalizes_supported_alias_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_rejects_kernel_layer_page_first_layout \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_hicache_storage_rejects_layer_page_first_layout \
  test/registered/unit/managers/test_hicache_controller_cp.py::TestPageFirstPerLayerBackupTaiKernel \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestLayerPageFirstDirectHostLayout \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestPageFirstDirectAllLayerBackupRoute \
  test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAIndexerPageIndices
```

结果：`36 passed, 5 subtests passed`。

未完成：

- W4：远端 CUDA roundtrip、TAI benchmark quick smoke、ETE GSM8K/replay 验证。

5. **W5 ETE correctness/performance**
   - GSM8K 两轮。
   - replay workload 对比 page_first_direct 与 layer_page_first。

## 第一阶段完成定义

第一阶段完成必须同时满足：

- `direct + layer_page_first` 能通过 server args validation。
- MHA / MLA / NSA indexer host pool 能初始化 LPF buffer。
- CP HiCache target + draft per-layer backup/load 都调用 TAI LPF direct op。
- all-layer backup 不调用 sgl-kernel direct all-layer path。
- LPF 不支持的 storage zero-copy 路径 fail-fast。
- 远端 CUDA roundtrip 覆盖 bf16/fp8 KV + NSA indexer。
- ETE GSM8K cache-hit 第二轮不掉点。
- replay workload 没有 silent fallback，日志中没有 `[CP_HICACHE_FALLBACK]` 热路径。

## 实施记录

### W1 server/config gate

状态：已完成。

改动：

- `--hicache-mem-layout` parser choices 增加 `layer_page_first`。
- CP shared KV + HiCache validation 增加支持对：

```python
("direct", "layer_page_first")
```

- `kernel + layer_page_first` 保持不支持，错误信息明确说明 `layer_page_first` 只支持 direct IO backend。
- 保持旧 normalization 不变：
  - `kernel + page_first_direct` 仍归一为 `direct + page_first_direct`。
  - `direct + page_first` 仍归一为 `direct + page_first_direct`。
  - 不自动把 `page_first_direct` 改成 `layer_page_first`。

新增测试：

- `test_hicache_mem_layout_parser_accepts_layer_page_first`
- `TestHiCacheArgs.test_cp_hicache_accepts_supported_backend_layout_pairs`
  - 覆盖 `direct + layer_page_first`
- `TestHiCacheArgs.test_cp_hicache_rejects_kernel_layer_page_first_layout`

验证：

```bash
python -m py_compile \
  python/sglang/srt/server_args.py \
  test/registered/unit/server_args/test_server_args.py
```

通过。

远端定向测试：

```bash
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc '
cd /sgl-workspace/sglang-tai && \
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py::test_hicache_mem_layout_parser_accepts_layer_page_first \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_accepts_supported_backend_layout_pairs \
  test/registered/unit/server_args/test_server_args.py::TestHiCacheArgs::test_cp_hicache_rejects_kernel_layer_page_first_layout
'"
```

结果：`3 passed, 5 subtests passed`。

远端完整 `test_server_args.py`：

结果：`45 passed, 1 failed, 9 subtests passed`。

失败项：

```text
TestPrepareServerArgs.test_prepare_server_args
```

原因：容器内解析 `Qwen/Qwen2.5-1.5B-Instruct` 配置时访问 HuggingFace 失败，报 DNS / client closed；和本次 `layer_page_first` gate 无关。本次新增和相关 HiCache validation 用例已通过。
