Enable CP HiCache direct transfers to use layer-page host layout
CP shared-KV HiCache transfers are per-layer, so the host layout should match the access pattern instead of forcing page-major strides through every layer. This adds a direct-only layer_page_first layout, routes per-layer KV and NSA index backup/load through the TAI LF<->LPF direct kernels, and keeps storage/page-buffer metadata paths fail-fast until their page-level contract is redesigned.\n\nThe direct controller keeps host indices in caller order for both page_first_direct and layer_page_first because the TAI direct path requires CPU index descriptors and owns descriptor coalescing. All-layer backup intentionally loops over per-layer direct kernels rather than using the sgl-kernel all-layer direct ABI.\n\nConstraint: layer_page_first is currently host-only CP HiCache; storage backends assume page-major contiguous page metadata.\nConstraint: TAI LPF direct kernels require CPU int64 page indices and complete page spans.\nRejected: silently fallback to SM copy when TAI LPF kernels are missing | that hides production performance regressions.\nRejected: support storage page metadata in this commit | LPF requires a layer-page-level storage contract, not a one-pointer-per-page contract.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not enable storage or kernel backend for layer_page_first without redesigning page-buffer metadata and adding remote ETE coverage.\nTested: local py_compile for touched runtime files.\nTested: remote py_compile in g0034 container for touched runtime files.\nTested: remote targeted pytest: 5 passed for parser/storage/layout/move_indices smoke coverage.\nNot-tested: full CP HiCache ETE with --hicache-mem-layout layer_page_first after this commit step.\nNot-tested: combined CUDA roundtrip tests in one pytest process; previous independent runs passed but combined run exposed a host-memory registration lifecycle issue.
This commit is contained in:
laoyao0822
@@ -0,0 +1,924 @@
|
||||
# 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 用例已通过。
|
||||
Reference in New Issue
Block a user