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 用例已通过。
|
||||
@@ -1697,7 +1697,7 @@ class HiCacheController:
|
||||
device_indices = device_indices.cpu()
|
||||
host_indices, idx = host_indices.sort()
|
||||
return host_indices, device_indices.index_select(0, idx)
|
||||
elif mem_pool_host.layout == "page_first_direct":
|
||||
elif mem_pool_host.layout in ("page_first_direct", "layer_page_first"):
|
||||
return host_indices, device_indices.cpu()
|
||||
else:
|
||||
raise ValueError(
|
||||
|
||||
@@ -78,6 +78,14 @@ class HiCacheStorage(ABC):
|
||||
# todo, the page size of storage backend does not have to be the same as the same as host memory pool
|
||||
|
||||
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
|
||||
if getattr(mem_pool_host, "layout", None) == "layer_page_first":
|
||||
raise RuntimeError(
|
||||
"[CP_HICACHE_FAILFAST][layer_page_first_storage_unsupported] "
|
||||
"layer_page_first is currently a host-only direct CP HiCache "
|
||||
"layout. Storage backends are not supported for this layout "
|
||||
"because their page-level zero-copy/metadata contracts assume "
|
||||
"page-major contiguous pages."
|
||||
)
|
||||
self.mem_pool_host = mem_pool_host
|
||||
|
||||
def batch_get_v1(
|
||||
|
||||
@@ -83,6 +83,22 @@ def _load_tai_transfer_kv_per_layer_direct_lf_pf():
|
||||
) from exc
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_transfer_kv_per_layer_direct_lf_lpf():
|
||||
try:
|
||||
from tai_kernel.nsa_prefill import transfer_kv_per_layer_direct_lf_lpf
|
||||
|
||||
return transfer_kv_per_layer_direct_lf_lpf
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"[CP_HICACHE_FAILFAST][missing_tai_layer_page_first_direct_lf_lpf] "
|
||||
"direct+layer_page_first per-layer D2H backup requires "
|
||||
"tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lf_lpf. "
|
||||
"Build/sync tai-kernel; this path intentionally does not fall back "
|
||||
"to an SM-consuming copy kernel."
|
||||
) from exc
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_transfer_kv_per_layer_direct_pf_lf():
|
||||
try:
|
||||
@@ -99,6 +115,32 @@ def _load_tai_transfer_kv_per_layer_direct_pf_lf():
|
||||
) from exc
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _load_tai_transfer_kv_per_layer_direct_lpf_lf():
|
||||
try:
|
||||
from tai_kernel.nsa_prefill import transfer_kv_per_layer_direct_lpf_lf
|
||||
|
||||
return transfer_kv_per_layer_direct_lpf_lf
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"[CP_HICACHE_FAILFAST][missing_tai_layer_page_first_direct_lpf_lf] "
|
||||
"direct+layer_page_first per-layer H2D load requires "
|
||||
"tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lpf_lf. "
|
||||
"Build/sync tai-kernel; this path intentionally does not fall back "
|
||||
"to an SM-consuming copy kernel."
|
||||
) from exc
|
||||
|
||||
|
||||
def _raise_layer_page_first_page_buffer_meta_unsupported() -> None:
|
||||
raise RuntimeError(
|
||||
"[CP_HICACHE_FAILFAST][layer_page_first_page_buffer_meta_unsupported] "
|
||||
"layer_page_first is currently a host-only direct CP HiCache layout. "
|
||||
"Storage/zero-copy page buffer metadata is not supported because a "
|
||||
"single page across all layers is not stored as one page-major "
|
||||
"contiguous block. Disable storage or use page_first_direct."
|
||||
)
|
||||
|
||||
|
||||
def synchronized(func):
|
||||
@wraps(func)
|
||||
def wrapper(self, *args, **kwargs):
|
||||
@@ -768,6 +810,15 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
self.head_num,
|
||||
self.head_dim,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
dims = (
|
||||
2,
|
||||
self.layer_num,
|
||||
self.page_num,
|
||||
self.page_size,
|
||||
self.head_num,
|
||||
self.head_dim,
|
||||
)
|
||||
elif self.layout == "page_head":
|
||||
dims = (
|
||||
2,
|
||||
@@ -882,6 +933,18 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
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,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
elif io_backend == "kernel_ascend":
|
||||
@@ -986,6 +1049,20 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf()
|
||||
for layer_id in range(self.layer_num):
|
||||
tai_transfer(
|
||||
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,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
elif io_backend == "kernel_ascend":
|
||||
@@ -1056,6 +1133,18 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
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,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
elif io_backend == "kernel_ascend":
|
||||
@@ -1084,6 +1173,9 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
elif self.layout in ["page_first_direct", "page_head"]:
|
||||
real_index = index // self.page_size
|
||||
data_page = self.kv_buffer[:, real_index : real_index + 1, :, :, :, :]
|
||||
elif self.layout == "layer_page_first":
|
||||
real_index = index // self.page_size
|
||||
data_page = self.kv_buffer[:, :, real_index : real_index + 1, :, :, :]
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
if flat:
|
||||
@@ -1122,6 +1214,13 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
2, 1, self.layer_num, self.page_size, self.head_num, self.head_dim
|
||||
)
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
real_index = index // self.page_size
|
||||
self.kv_buffer[:, :, real_index : real_index + 1, :, :, :] = (
|
||||
data_page.reshape(
|
||||
2, self.layer_num, 1, self.page_size, self.head_num, self.head_dim
|
||||
)
|
||||
)
|
||||
elif self.layout == "page_head":
|
||||
real_index = index // self.page_size
|
||||
self.kv_buffer[:, real_index : real_index + 1, :, :, :, :] = (
|
||||
@@ -1238,6 +1337,8 @@ class MHATokenToKVPoolHost(HostKVCache):
|
||||
* self.head_dim
|
||||
)
|
||||
element_size_list = [element_size] * len(ptr_list)
|
||||
elif self.layout == "layer_page_first":
|
||||
_raise_layer_page_first_page_buffer_meta_unsupported()
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
return ptr_list, element_size_list
|
||||
@@ -1313,6 +1414,14 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
1,
|
||||
self.kv_cache_dim,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
dims = (
|
||||
self.layer_num,
|
||||
self.page_num,
|
||||
self.page_size,
|
||||
1,
|
||||
self.kv_cache_dim,
|
||||
)
|
||||
# Ascend-specific: Aligns with NPUMLATokenToKVPool layout
|
||||
# Separately allocate k_buffer and v_buffer for easier data transfer.
|
||||
elif self.layout == "page_first_kv_split":
|
||||
@@ -1406,6 +1515,15 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
|
||||
src_ptrs=[self.kv_buffer],
|
||||
dst_ptrs=[device_pool.kv_buffer[layer_id]],
|
||||
src_indices=host_indices,
|
||||
dst_indices=device_indices,
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
elif io_backend == "kernel_ascend":
|
||||
@@ -1478,6 +1596,17 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf()
|
||||
for layer_id in range(self.layer_num):
|
||||
tai_transfer(
|
||||
src_ptrs=[device_pool.kv_buffer[layer_id]],
|
||||
dst_ptrs=[self.kv_buffer],
|
||||
src_indices=device_indices,
|
||||
dst_indices=host_indices,
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
elif io_backend == "kernel_ascend":
|
||||
@@ -1546,6 +1675,15 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
_load_tai_transfer_kv_per_layer_direct_lf_lpf()(
|
||||
src_ptrs=[device_pool.kv_buffer[layer_id]],
|
||||
dst_ptrs=[self.kv_buffer],
|
||||
src_indices=device_indices,
|
||||
dst_indices=host_indices,
|
||||
layer_id=layer_id,
|
||||
page_size=self.page_size,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
elif io_backend == "kernel_ascend":
|
||||
@@ -1576,6 +1714,9 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
elif self.layout == "page_first_direct":
|
||||
real_index = index // self.page_size
|
||||
data_page = self.kv_buffer[real_index : real_index + 1, :, :, :, :]
|
||||
elif self.layout == "layer_page_first":
|
||||
real_index = index // self.page_size
|
||||
data_page = self.kv_buffer[:, real_index : real_index + 1, :, :, :]
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
if flat:
|
||||
@@ -1619,6 +1760,15 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
1,
|
||||
self.kv_cache_dim,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
real_index = index // self.page_size
|
||||
self.kv_buffer[:, real_index : real_index + 1, :, :, :] = data_page.reshape(
|
||||
self.layer_num,
|
||||
1,
|
||||
self.page_size,
|
||||
1,
|
||||
self.kv_cache_dim,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
|
||||
@@ -1658,6 +1808,8 @@ class MLATokenToKVPoolHost(HostKVCache):
|
||||
* self.kv_cache_dim
|
||||
)
|
||||
element_size_list = [element_size] * len(ptr_list)
|
||||
elif self.layout == "layer_page_first":
|
||||
_raise_layer_page_first_page_buffer_meta_unsupported()
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
return ptr_list, element_size_list
|
||||
@@ -1754,6 +1906,27 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
||||
pin_memory=self.pin_memory,
|
||||
allocator=self.allocator,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
self.index_k_with_scale_buffer = alloc_func(
|
||||
(
|
||||
self.layer_num,
|
||||
self.indexer_page_num,
|
||||
1,
|
||||
self.indexer_page_stride_size,
|
||||
),
|
||||
dtype=self.indexer_dtype,
|
||||
device=self.device,
|
||||
pin_memory=self.pin_memory,
|
||||
allocator=self.allocator,
|
||||
)
|
||||
self.index_k_data_refs = [
|
||||
self.index_k_with_scale_buffer[i] for i in range(self.layer_num)
|
||||
]
|
||||
self.index_k_data_ptrs = torch.tensor(
|
||||
[x.data_ptr() for x in self.index_k_data_refs],
|
||||
dtype=torch.uint64,
|
||||
device=self.device_pool.device,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
|
||||
@@ -1833,6 +2006,15 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
||||
layer_id=layer_id,
|
||||
page_size=1,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
|
||||
src_ptrs=[self.index_k_with_scale_buffer],
|
||||
dst_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
|
||||
src_indices=host_page_indices,
|
||||
dst_indices=device_page_indices,
|
||||
layer_id=layer_id,
|
||||
page_size=1,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
else:
|
||||
@@ -1892,6 +2074,17 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
||||
layer_id=layer_id,
|
||||
page_size=1,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf()
|
||||
for layer_id in range(self.layer_num):
|
||||
tai_transfer(
|
||||
src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
|
||||
dst_ptrs=[self.index_k_with_scale_buffer],
|
||||
src_indices=device_page_indices,
|
||||
dst_indices=host_page_indices,
|
||||
layer_id=layer_id,
|
||||
page_size=1,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
else:
|
||||
@@ -1943,6 +2136,15 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
|
||||
layer_id=layer_id,
|
||||
page_size=1,
|
||||
)
|
||||
elif self.layout == "layer_page_first":
|
||||
_load_tai_transfer_kv_per_layer_direct_lf_lpf()(
|
||||
src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]],
|
||||
dst_ptrs=[self.index_k_with_scale_buffer],
|
||||
src_indices=device_page_indices,
|
||||
dst_indices=host_page_indices,
|
||||
layer_id=layer_id,
|
||||
page_size=1,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported layout: {self.layout}")
|
||||
else:
|
||||
|
||||
@@ -962,6 +962,7 @@ class ServerArgs:
|
||||
("kernel", "page_first"),
|
||||
("direct", "layer_first"),
|
||||
("direct", "page_first_direct"),
|
||||
("direct", "layer_page_first"),
|
||||
}
|
||||
current_pair = (self.hicache_io_backend, self.hicache_mem_layout)
|
||||
if current_pair in supported_pairs:
|
||||
@@ -970,12 +971,14 @@ class ServerArgs:
|
||||
raise ValueError(
|
||||
"CP shared KV HiCache supports only "
|
||||
"kernel/layer_first, kernel/page_first, direct/layer_first, "
|
||||
"and direct/page_first_direct after HiCache normalization. "
|
||||
"direct/page_first_direct, and direct/layer_page_first after "
|
||||
"HiCache normalization. "
|
||||
f"Got hicache_io_backend={self.hicache_io_backend!r} "
|
||||
f"hicache_mem_layout={self.hicache_mem_layout!r}. "
|
||||
"page_head is MHA/Mooncake-specific, page_first_kv_split is "
|
||||
"kernel_ascend-specific, and kernel_ascend CP shared KV HiCache "
|
||||
"is not implemented in the CUDA NSA host path."
|
||||
"is not implemented in the CUDA NSA host path. "
|
||||
"layer_page_first is supported only with the direct IO backend."
|
||||
)
|
||||
|
||||
def _handle_deprecated_args(self):
|
||||
@@ -3027,6 +3030,17 @@ class ServerArgs:
|
||||
)
|
||||
|
||||
def _resolve_storage_layout_compatibility(self):
|
||||
if (
|
||||
self.hicache_storage_backend is not None
|
||||
and self.hicache_mem_layout == "layer_page_first"
|
||||
):
|
||||
raise ValueError(
|
||||
"[CP_HICACHE_FAILFAST][layer_page_first_storage_unsupported] "
|
||||
"layer_page_first is currently a host-only direct CP HiCache "
|
||||
"layout and does not support storage backends. Disable "
|
||||
"hicache_storage_backend or use page_first_direct."
|
||||
)
|
||||
|
||||
if (
|
||||
self.hicache_storage_backend != "mooncake"
|
||||
or self.hicache_mem_layout != "layer_first"
|
||||
@@ -5225,6 +5239,7 @@ class ServerArgs:
|
||||
"layer_first",
|
||||
"page_first",
|
||||
"page_first_direct",
|
||||
"layer_page_first",
|
||||
"page_first_kv_split",
|
||||
"page_head",
|
||||
],
|
||||
|
||||
@@ -105,7 +105,7 @@ for _schema in (
|
||||
if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
|
||||
raise
|
||||
|
||||
from sglang.srt.managers.cache_controller import HiCacheController
|
||||
from sglang.srt.managers.cache_controller import CacheOperation, HiCacheController
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
|
||||
from sglang.srt.mem_cache.memory_pool_host import (
|
||||
@@ -449,6 +449,139 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
|
||||
self.assertEqual(kwargs["layer_id"], 1)
|
||||
self.assertEqual(kwargs["page_size"], 1)
|
||||
|
||||
def test_mla_layer_page_first_per_layer_backup_uses_direct_lf_lpf(self):
|
||||
calls = []
|
||||
|
||||
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
|
||||
calls.append(
|
||||
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
|
||||
)
|
||||
|
||||
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((3, 8, 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_transfer_kv_per_layer_direct_lf_lpf",
|
||||
return_value=fake_direct,
|
||||
):
|
||||
host_pool.backup_from_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=2,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
|
||||
self.assertEqual(len(src_ptrs), 1)
|
||||
self.assertEqual(src_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
|
||||
self.assertEqual(len(dst_ptrs), 1)
|
||||
self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
|
||||
self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
|
||||
self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
|
||||
self.assertEqual(kwargs["layer_id"], 2)
|
||||
self.assertEqual(kwargs["page_size"], 4)
|
||||
|
||||
def test_mha_layer_page_first_per_layer_backup_uses_direct_lf_lpf(self):
|
||||
calls = []
|
||||
|
||||
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
|
||||
calls.append(
|
||||
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
|
||||
)
|
||||
|
||||
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((2, 3, 8, 4, 2, 8), dtype=torch.uint8)
|
||||
device_pool = type("DevicePool", (), {})()
|
||||
device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
|
||||
device_pool.v_buffer = torch.empty((3, 32, 2, 8), 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_transfer_kv_per_layer_direct_lf_lpf",
|
||||
return_value=fake_direct,
|
||||
):
|
||||
host_pool.backup_from_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=2,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
|
||||
self.assertEqual(len(src_ptrs), 2)
|
||||
self.assertEqual(src_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
|
||||
self.assertEqual(src_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
|
||||
self.assertEqual(len(dst_ptrs), 2)
|
||||
self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
|
||||
self.assertEqual(dst_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
|
||||
self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
|
||||
self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
|
||||
self.assertEqual(kwargs["layer_id"], 2)
|
||||
self.assertEqual(kwargs["page_size"], 4)
|
||||
|
||||
def test_nsa_indexer_layer_page_first_per_layer_backup_uses_direct_lf_lpf(self):
|
||||
calls = []
|
||||
|
||||
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
|
||||
calls.append(
|
||||
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
|
||||
)
|
||||
|
||||
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.index_k_with_scale_buffer = torch.empty(
|
||||
(3, 8, 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_transfer_kv_per_layer_direct_lf_lpf",
|
||||
return_value=fake_direct,
|
||||
):
|
||||
host_pool._backup_indexer_from_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=1,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
|
||||
self.assertEqual(len(src_ptrs), 1)
|
||||
self.assertEqual(
|
||||
src_ptrs[0].data_ptr(),
|
||||
device_pool.index_k_with_scale_buffer[1].data_ptr(),
|
||||
)
|
||||
self.assertEqual(len(dst_ptrs), 1)
|
||||
self.assertEqual(
|
||||
dst_ptrs[0].data_ptr(),
|
||||
host_pool.index_k_with_scale_buffer.data_ptr(),
|
||||
)
|
||||
self.assertEqual(src_indices.tolist(), [3])
|
||||
self.assertEqual(dst_indices.tolist(), [1])
|
||||
self.assertEqual(kwargs["layer_id"], 1)
|
||||
self.assertEqual(kwargs["page_size"], 1)
|
||||
|
||||
def test_mla_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
|
||||
calls = []
|
||||
|
||||
@@ -582,6 +715,139 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
|
||||
self.assertEqual(kwargs["layer_id"], 1)
|
||||
self.assertEqual(kwargs["page_size"], 1)
|
||||
|
||||
def test_mla_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf(self):
|
||||
calls = []
|
||||
|
||||
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
|
||||
calls.append(
|
||||
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
|
||||
)
|
||||
|
||||
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((3, 8, 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_transfer_kv_per_layer_direct_lpf_lf",
|
||||
return_value=fake_direct,
|
||||
):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=2,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
|
||||
self.assertEqual(len(src_ptrs), 1)
|
||||
self.assertEqual(src_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
|
||||
self.assertEqual(len(dst_ptrs), 1)
|
||||
self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
|
||||
self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
|
||||
self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
|
||||
self.assertEqual(kwargs["layer_id"], 2)
|
||||
self.assertEqual(kwargs["page_size"], 4)
|
||||
|
||||
def test_mha_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf(self):
|
||||
calls = []
|
||||
|
||||
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
|
||||
calls.append(
|
||||
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
|
||||
)
|
||||
|
||||
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.kv_buffer = torch.empty((2, 3, 8, 4, 2, 8), dtype=torch.uint8)
|
||||
device_pool = type("DevicePool", (), {})()
|
||||
device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
|
||||
device_pool.v_buffer = torch.empty((3, 32, 2, 8), 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_transfer_kv_per_layer_direct_lpf_lf",
|
||||
return_value=fake_direct,
|
||||
):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=2,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
|
||||
self.assertEqual(len(src_ptrs), 2)
|
||||
self.assertEqual(src_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
|
||||
self.assertEqual(src_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
|
||||
self.assertEqual(len(dst_ptrs), 2)
|
||||
self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
|
||||
self.assertEqual(dst_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
|
||||
self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
|
||||
self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
|
||||
self.assertEqual(kwargs["layer_id"], 2)
|
||||
self.assertEqual(kwargs["page_size"], 4)
|
||||
|
||||
def test_nsa_indexer_layer_page_first_per_layer_load_uses_tai_direct_lpf_lf(self):
|
||||
calls = []
|
||||
|
||||
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
|
||||
calls.append(
|
||||
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
|
||||
)
|
||||
|
||||
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.index_k_with_scale_buffer = torch.empty(
|
||||
(3, 8, 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_transfer_kv_per_layer_direct_lpf_lf",
|
||||
return_value=fake_direct,
|
||||
):
|
||||
host_pool._load_indexer_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
device_indices,
|
||||
layer_id=1,
|
||||
io_backend="direct",
|
||||
)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
|
||||
self.assertEqual(len(src_ptrs), 1)
|
||||
self.assertEqual(
|
||||
src_ptrs[0].data_ptr(),
|
||||
host_pool.index_k_with_scale_buffer.data_ptr(),
|
||||
)
|
||||
self.assertEqual(len(dst_ptrs), 1)
|
||||
self.assertEqual(
|
||||
dst_ptrs[0].data_ptr(),
|
||||
device_pool.index_k_with_scale_buffer[1].data_ptr(),
|
||||
)
|
||||
self.assertEqual(src_indices.tolist(), [1])
|
||||
self.assertEqual(dst_indices.tolist(), [3])
|
||||
self.assertEqual(kwargs["layer_id"], 1)
|
||||
self.assertEqual(kwargs["page_size"], 1)
|
||||
|
||||
def test_nsa_indexer_load_reuses_precomputed_page_indices_across_layers(self):
|
||||
calls = []
|
||||
|
||||
@@ -1845,6 +2111,23 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
|
||||
self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
|
||||
self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])
|
||||
|
||||
def test_direct_layer_page_first_move_indices_keeps_host_order_and_cpu_device_indices(self):
|
||||
host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
|
||||
host_pool.layout = "layer_page_first"
|
||||
controller = self.make_controller(host_pool, cp_rank=1)
|
||||
op = CacheOperation(
|
||||
host_indices=torch.tensor([12, 8, 9, 10], dtype=torch.int64),
|
||||
device_indices=torch.tensor([32, 28, 29, 30], dtype=torch.int64),
|
||||
node_id=7,
|
||||
)
|
||||
|
||||
host_indices, device_indices = controller.move_indices(op, host_pool)
|
||||
|
||||
self.assertEqual(host_indices.tolist(), [12, 8, 9, 10])
|
||||
self.assertEqual(host_indices.device.type, "cpu")
|
||||
self.assertEqual(device_indices.tolist(), [32, 28, 29, 30])
|
||||
self.assertEqual(device_indices.device.type, "cpu")
|
||||
|
||||
def test_cp_load_rejects_non_contiguous_physical_device_page(self):
|
||||
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
|
||||
allocator = FakeAllocator(
|
||||
|
||||
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
||||
build_in_seq_page_compute_owners,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
|
||||
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
|
||||
from sglang.srt.mem_cache.memory_pool_host import (
|
||||
ALLOC_MEMORY_FUNCS,
|
||||
@@ -32,6 +33,184 @@ from sglang.test.test_utils import CustomTestCase
|
||||
register_cuda_ci(est_time=3, suite="stage-b-test-1-gpu-small")
|
||||
|
||||
|
||||
class _DummyHiCacheStorage(HiCacheStorage):
|
||||
def get(self, key, target_location=None, target_sizes=None):
|
||||
return None
|
||||
|
||||
def batch_get(self, keys, target_locations=None, target_sizes=None):
|
||||
return []
|
||||
|
||||
def set(self, key, value=None, target_location=None, target_sizes=None):
|
||||
return False
|
||||
|
||||
def batch_set(self, keys, values=None, target_locations=None, target_sizes=None):
|
||||
return False
|
||||
|
||||
def exists(self, key):
|
||||
return False
|
||||
|
||||
|
||||
class TestLayerPageFirstDirectHostLayout(CustomTestCase):
|
||||
def test_mha_layer_page_first_direct_host_layout_is_layer_page_major(self):
|
||||
device_pool = SimpleNamespace(
|
||||
store_dtype=torch.float16,
|
||||
size=8,
|
||||
start_layer=0,
|
||||
end_layer=3,
|
||||
device="cpu",
|
||||
head_num=2,
|
||||
head_dim=8,
|
||||
layer_num=3,
|
||||
)
|
||||
|
||||
host_pool = MHATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=4,
|
||||
layout="layer_page_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
host_token_capacity=16,
|
||||
)
|
||||
|
||||
self.assertEqual(tuple(host_pool.kv_buffer.shape), (2, 3, 4, 4, 2, 8))
|
||||
self.assertEqual(tuple(host_pool.k_buffer.shape), (3, 4, 4, 2, 8))
|
||||
self.assertEqual(len(host_pool.k_data_refs), 3)
|
||||
self.assertEqual(tuple(host_pool.k_data_refs[0].shape), (4, 4, 2, 8))
|
||||
|
||||
def test_mla_layer_page_first_direct_host_layout_is_layer_page_major(self):
|
||||
device_pool = SimpleNamespace(
|
||||
store_dtype=torch.float16,
|
||||
size=8,
|
||||
start_layer=0,
|
||||
end_layer=3,
|
||||
device="cpu",
|
||||
kv_lora_rank=16,
|
||||
qk_rope_head_dim=4,
|
||||
layer_num=3,
|
||||
)
|
||||
|
||||
host_pool = MLATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=4,
|
||||
layout="layer_page_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
host_token_capacity=16,
|
||||
)
|
||||
|
||||
self.assertEqual(tuple(host_pool.kv_buffer.shape), (3, 4, 4, 1, 20))
|
||||
self.assertEqual(len(host_pool.data_refs), 3)
|
||||
self.assertEqual(tuple(host_pool.data_refs[0].shape), (4, 4, 1, 20))
|
||||
|
||||
def test_nsa_layer_page_first_direct_indexer_layout_is_layer_page_major(self):
|
||||
indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype
|
||||
device_pool = SimpleNamespace(
|
||||
store_dtype=torch.float16,
|
||||
size=8,
|
||||
start_layer=0,
|
||||
end_layer=3,
|
||||
device="cpu",
|
||||
kv_lora_rank=16,
|
||||
qk_rope_head_dim=4,
|
||||
kv_cache_dim=24,
|
||||
layer_num=3,
|
||||
index_head_dim=16,
|
||||
quant_block_size=8,
|
||||
index_k_with_scale_buffer=[
|
||||
torch.empty((5, 96), dtype=indexer_dtype) for _ in range(3)
|
||||
],
|
||||
)
|
||||
|
||||
host_pool = NSATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=4,
|
||||
layout="layer_page_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
host_token_capacity=16,
|
||||
)
|
||||
|
||||
self.assertEqual(tuple(host_pool.kv_buffer.shape), (3, 4, 4, 1, 24))
|
||||
self.assertEqual(
|
||||
tuple(host_pool.index_k_with_scale_buffer.shape),
|
||||
(3, host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size),
|
||||
)
|
||||
self.assertEqual(len(host_pool.index_k_data_refs), 3)
|
||||
self.assertEqual(
|
||||
tuple(host_pool.index_k_data_refs[0].shape),
|
||||
(host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size),
|
||||
)
|
||||
|
||||
def test_mha_layer_page_first_page_buffer_meta_fails_fast_for_storage(self):
|
||||
device_pool = SimpleNamespace(
|
||||
store_dtype=torch.float16,
|
||||
size=8,
|
||||
start_layer=0,
|
||||
end_layer=2,
|
||||
device="cpu",
|
||||
head_num=2,
|
||||
head_dim=8,
|
||||
layer_num=2,
|
||||
)
|
||||
host_pool = MHATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=4,
|
||||
layout="layer_page_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
host_token_capacity=16,
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError, "layer_page_first_page_buffer_meta_unsupported"
|
||||
):
|
||||
host_pool.get_page_buffer_meta(torch.arange(0, 4, dtype=torch.int64))
|
||||
|
||||
def test_mla_layer_page_first_page_buffer_meta_fails_fast_for_storage(self):
|
||||
device_pool = SimpleNamespace(
|
||||
store_dtype=torch.float16,
|
||||
size=8,
|
||||
start_layer=0,
|
||||
end_layer=2,
|
||||
device="cpu",
|
||||
kv_lora_rank=16,
|
||||
qk_rope_head_dim=4,
|
||||
layer_num=2,
|
||||
)
|
||||
host_pool = MLATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=4,
|
||||
layout="layer_page_first",
|
||||
pin_memory=False,
|
||||
device="cpu",
|
||||
host_token_capacity=16,
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError, "layer_page_first_page_buffer_meta_unsupported"
|
||||
):
|
||||
host_pool.get_page_buffer_meta(torch.arange(0, 4, dtype=torch.int64))
|
||||
|
||||
def test_storage_registration_fails_fast_for_layer_page_first(self):
|
||||
storage = _DummyHiCacheStorage()
|
||||
mem_pool_host = SimpleNamespace(layout="layer_page_first")
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError, "layer_page_first_storage_unsupported"
|
||||
):
|
||||
storage.register_mem_pool_host(mem_pool_host)
|
||||
|
||||
|
||||
class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
def setUp(self):
|
||||
if not torch.cuda.is_available():
|
||||
@@ -1263,7 +1442,9 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
io_backend="direct", layout="page_first_direct"
|
||||
)
|
||||
|
||||
def test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages(self):
|
||||
def _run_direct_roundtrip_preserves_kv_and_indexer_pages(
|
||||
self, *, layout: str, dtype: torch.dtype
|
||||
):
|
||||
page_size = 64
|
||||
layer_num = 3
|
||||
size = page_size * 20
|
||||
@@ -1272,7 +1453,7 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
size=size,
|
||||
page_size=page_size,
|
||||
kv_lora_rank=512,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
dtype=dtype,
|
||||
qk_rope_head_dim=64,
|
||||
layer_num=layer_num,
|
||||
device="cuda",
|
||||
@@ -1285,7 +1466,7 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=page_size,
|
||||
layout="page_first_direct",
|
||||
layout=layout,
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
)
|
||||
@@ -1373,7 +1554,8 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
]
|
||||
self.assertTrue(
|
||||
torch.equal(got_kv, expected_kv[layer_id][page_idx]),
|
||||
f"KV roundtrip mismatch layer={layer_id} dst_page={dst_page}",
|
||||
f"KV roundtrip mismatch layout={layout} dtype={dtype} "
|
||||
f"layer={layer_id} dst_page={dst_page}",
|
||||
)
|
||||
|
||||
got_index = device_pool.index_k_with_scale_buffer[layer_id][
|
||||
@@ -1381,9 +1563,24 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
]
|
||||
self.assertTrue(
|
||||
torch.equal(got_index, expected_index[layer_id][page_idx]),
|
||||
f"index roundtrip mismatch layer={layer_id} dst_page={dst_page}",
|
||||
f"index roundtrip mismatch layout={layout} dtype={dtype} "
|
||||
f"layer={layer_id} dst_page={dst_page}",
|
||||
)
|
||||
|
||||
def test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages(self):
|
||||
self._run_direct_roundtrip_preserves_kv_and_indexer_pages(
|
||||
layout="page_first_direct", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
def test_fp8_layer_page_first_roundtrip_preserves_kv_and_indexer_pages(self):
|
||||
self._run_direct_roundtrip_preserves_kv_and_indexer_pages(
|
||||
layout="layer_page_first", dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
def test_bf16_layer_page_first_roundtrip_preserves_kv_and_indexer_pages(self):
|
||||
self._run_direct_roundtrip_preserves_kv_and_indexer_pages(
|
||||
layout="layer_page_first", dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
class TestPageFirstDirectAllLayerBackupRoute(CustomTestCase):
|
||||
def test_mla_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
|
||||
@@ -1576,6 +1773,126 @@ class TestNSAIndexerPageIndices(CustomTestCase):
|
||||
self.assertEqual(call["dst_indices"].tolist(), [0, 1])
|
||||
self.assertEqual(call["page_size"], 1)
|
||||
|
||||
def test_mla_layer_page_first_all_layer_backup_uses_tai_per_layer_route(self):
|
||||
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.layer_num = 2
|
||||
host_pool.kv_buffer = "host-mla-layer-page-first"
|
||||
device_pool = type(
|
||||
"FakeDevicePool",
|
||||
(),
|
||||
{"kv_buffer": ["device-mla-layer-0", "device-mla-layer-1"]},
|
||||
)()
|
||||
calls = []
|
||||
|
||||
def fake_tai_transfer(**kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
|
||||
return_value=fake_tai_transfer,
|
||||
):
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool,
|
||||
torch.tensor([0, 1, 2, 3], dtype=torch.int64),
|
||||
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
|
||||
"direct",
|
||||
)
|
||||
|
||||
self.assertEqual([call["layer_id"] for call in calls], [0, 1])
|
||||
for layer_id, call in enumerate(calls):
|
||||
self.assertEqual(call["src_ptrs"], [f"device-mla-layer-{layer_id}"])
|
||||
self.assertEqual(call["dst_ptrs"], ["host-mla-layer-page-first"])
|
||||
self.assertEqual(call["src_indices"].tolist(), [8, 9, 10, 11])
|
||||
self.assertEqual(call["dst_indices"].tolist(), [0, 1, 2, 3])
|
||||
self.assertEqual(call["page_size"], 4)
|
||||
|
||||
def test_mha_layer_page_first_all_layer_backup_uses_tai_per_layer_route(self):
|
||||
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.page_size = 4
|
||||
host_pool.layer_num = 2
|
||||
host_pool.kv_buffer = ["host-k-layer-page-first", "host-v-layer-page-first"]
|
||||
device_pool = type(
|
||||
"FakeDevicePool",
|
||||
(),
|
||||
{
|
||||
"k_buffer": ["device-k-layer-0", "device-k-layer-1"],
|
||||
"v_buffer": ["device-v-layer-0", "device-v-layer-1"],
|
||||
},
|
||||
)()
|
||||
calls = []
|
||||
|
||||
def fake_tai_transfer(**kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
|
||||
return_value=fake_tai_transfer,
|
||||
):
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool,
|
||||
torch.tensor([0, 1, 2, 3], dtype=torch.int64),
|
||||
torch.tensor([8, 9, 10, 11], dtype=torch.int64),
|
||||
"direct",
|
||||
)
|
||||
|
||||
self.assertEqual([call["layer_id"] for call in calls], [0, 1])
|
||||
for layer_id, call in enumerate(calls):
|
||||
self.assertEqual(
|
||||
call["src_ptrs"],
|
||||
[f"device-k-layer-{layer_id}", f"device-v-layer-{layer_id}"],
|
||||
)
|
||||
self.assertEqual(
|
||||
call["dst_ptrs"],
|
||||
["host-k-layer-page-first", "host-v-layer-page-first"],
|
||||
)
|
||||
self.assertEqual(call["src_indices"].tolist(), [8, 9, 10, 11])
|
||||
self.assertEqual(call["dst_indices"].tolist(), [0, 1, 2, 3])
|
||||
self.assertEqual(call["page_size"], 4)
|
||||
|
||||
def test_layer_page_first_all_layer_indexer_backup_uses_tai_per_layer_route(self):
|
||||
host_pool = self.make_host_pool_stub(page_size=4)
|
||||
host_pool.layout = "layer_page_first"
|
||||
host_pool.indexer_page_stride_size = 8
|
||||
host_pool.layer_num = 3
|
||||
host_pool.index_k_with_scale_buffer = "host-layer-page-first-indexer"
|
||||
device_pool = type(
|
||||
"FakeDevicePool",
|
||||
(),
|
||||
{
|
||||
"index_k_with_scale_buffer": [
|
||||
"device-layer-0",
|
||||
"device-layer-1",
|
||||
"device-layer-2",
|
||||
]
|
||||
},
|
||||
)()
|
||||
calls = []
|
||||
|
||||
def fake_tai_transfer(**kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
|
||||
return_value=fake_tai_transfer,
|
||||
):
|
||||
host_pool._backup_indexer_from_device_all_layer(
|
||||
device_pool,
|
||||
torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.int64),
|
||||
torch.tensor([8, 9, 10, 11, 12, 13, 14, 15], dtype=torch.int64),
|
||||
"direct",
|
||||
)
|
||||
|
||||
self.assertEqual([call["layer_id"] for call in calls], [0, 1, 2])
|
||||
for layer_id, call in enumerate(calls):
|
||||
self.assertEqual(call["src_ptrs"], [f"device-layer-{layer_id}"])
|
||||
self.assertEqual(call["dst_ptrs"], ["host-layer-page-first-indexer"])
|
||||
self.assertEqual(call["src_indices"].tolist(), [2, 3])
|
||||
self.assertEqual(call["dst_indices"].tolist(), [0, 1])
|
||||
self.assertEqual(call["page_size"], 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -75,6 +75,23 @@ def test_cp_shared_kv_prefill_bs_gt1_parser_limits():
|
||||
assert args.cp_shared_kv_prefill_max_total_extend_tokens == 8192
|
||||
|
||||
|
||||
def test_hicache_mem_layout_parser_accepts_layer_page_first():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
raw_args = parser.parse_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--hicache-mem-layout",
|
||||
"layer_page_first",
|
||||
]
|
||||
)
|
||||
args = ServerArgs.from_cli_args(raw_args)
|
||||
assert args.hicache_mem_layout == "layer_page_first"
|
||||
|
||||
|
||||
class TestLoadBalanceMethod(unittest.TestCase):
|
||||
def test_non_pd_defaults_to_round_robin(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")
|
||||
@@ -556,6 +573,7 @@ class TestHiCacheArgs(CustomTestCase):
|
||||
("kernel", "page_first"),
|
||||
("direct", "layer_first"),
|
||||
("direct", "page_first_direct"),
|
||||
("direct", "layer_page_first"),
|
||||
]
|
||||
|
||||
for io_backend, mem_layout in cases:
|
||||
@@ -598,6 +616,26 @@ class TestHiCacheArgs(CustomTestCase):
|
||||
hicache_mem_layout="page_first_kv_split",
|
||||
)
|
||||
|
||||
def test_cp_hicache_rejects_kernel_layer_page_first_layout(self):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "CP shared KV HiCache.*kernel.*layer_page_first"
|
||||
):
|
||||
self._normalize_and_validate_cp_hicache_args(
|
||||
hicache_io_backend="kernel",
|
||||
hicache_mem_layout="layer_page_first",
|
||||
)
|
||||
|
||||
def test_hicache_storage_rejects_layer_page_first_layout(self):
|
||||
args = self._make_args(
|
||||
enable_hierarchical_cache=True,
|
||||
hicache_storage_backend="mooncake",
|
||||
hicache_io_backend="direct",
|
||||
hicache_mem_layout="layer_page_first",
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "layer_page_first.*storage"):
|
||||
args._handle_hicache()
|
||||
|
||||
def test_cp_hicache_rejects_kernel_ascend_backend(self):
|
||||
with self.assertRaisesRegex(ValueError, "CP shared KV HiCache.*kernel_ascend"):
|
||||
self._normalize_and_validate_cp_hicache_args(
|
||||
|
||||
Reference in New Issue
Block a user