chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Reuse prepared HiCache load descriptors across CP prefill layers

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

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

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

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

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

Confidence: medium

Scope-risk: moderate

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

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

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

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

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

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

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-11 05:09:41 +08:00
parent adf357b02c
commit 7284a469a2
4 changed files with 1375 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

773
docs/advanced_features/nsa_prefill_cp_hicache_bs_gt1_l2_l1_transfer_prefetch_descriptor_plan_zh.md Normal file
View File

@@ -0,0 +1,773 @@
# NSA Prefill CP HiCache bs>1 L2->L1 Transfer Prefetch Descriptor 计划
> 日期：2026-06-10
> 分支：`cjy-cp-refactor`
> 范围：CP shared-KV + HiCache + bs>1 cache-hit 场景下，L2/host -> L1/device 数据搬运的 batch-aware per-layer prefetch 与 descriptor 复用。
> 相关文档：
> - `docs/advanced_features/nsa_prefill_cp_hicache_load_prefetch_overlap_notes.md`
> - `docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md`
> - `docs/advanced_features/nsa_prefill_cp_hicache_layer_page_first_direct_plan.md`
## 0. 目标和非目标
### 0.1 目标
1. 在 `--enable-cp-shared-kv-prefill-bs-gt1` 下，让多个 cache-hit request 的 L2->L1 transfer 以 batch 级 descriptor 组织。
2. 降低每层重复构造 descriptor / Python wrapper / extension submit 造成的 CPU overhead。
3. 保留现有正确性合同：
- radix / residency / allocator reservation 同步完成；
- data transfer 异步执行；
- target / draft KV 跟随同一个 logical cache node；
- 每层消费前精确等待本层 transfer event。
4. 保持当前 per-layer H2D overlap 能力，并为后续“提前两层 L2->L1 transfer prefetch”打基础。
5. 支持当前生产参数：
- `--hicache-io-backend direct`
- `--hicache-mem-layout page_first_direct`
- `--kv-cache-dtype fp8_e4m3`
- CP owner-lane / zigzag / bs>1 batch plan。
### 0.2 非目标
1. 第一阶段不改 radix tree 语义，不让异步 transfer 决定 cache node 是否可见。
2. 第一阶段不做跨 batch descriptor 复用。descriptor 只在当前 load op / 当前 batch 生命周期内复用。
3. 第一阶段不改 decode KV transfer / Mooncake transfer。
4. 第一阶段不强制切换到 `layer_page_first`。先让 `page_first_direct` 路径有 prepared descriptor，再评估 LPF 切换收益。
5. 第一阶段不新增 collective。descriptor 一致性必须来自本地 deterministic metadata，不依赖 all-reduce/all-gather。
---
## 1. 当前代码事实
### C1. 当前已有 batch 内 op 粗粒度合并
`python/sglang/srt/managers/cache_controller.py`
```text
CacheOperation.merge_ops(load_queue)
host_indices = torch.cat([op.host_indices for op in ops])
device_indices = torch.cat([op.device_indices for op in ops])
```
结论：当前 bs>1 下多个 request 的 load-back op 已经会合并成一个 `CacheOperation`。
但这个合并只是 tensor concat，不是 transfer descriptor compact。
### C2. 当前 per-layer transfer 每层都会重新进入 host pool API
`HiCacheController.start_loading()`：
```text
begin_load_to_device_op(host_indices, device_indices, io_backend)
for layer_id in range(layer_num):
mem_pool_host.load_to_device_per_layer(
mem_pool_device,
host_indices,
device_indices,
layer_id,
io_backend,
)
end_load_to_device_op()
```
结论：indices 是同一份，但每层仍会重新调用 Python method / extension wrapper。
direct backend 下，底层 TAI op 仍可能每层重新构造 H2D copy descriptor。
### C3. `begin_load_to_device_op()` 已经是 descriptor 预处理入口
`python/sglang/srt/mem_cache/memory_pool_host.py`
base host pool:
```python
def begin_load_to_device_op(self, host_indices, device_indices, io_backend):
"""Prepare layer-invariant metadata for one host->device load op."""
```
NSA host pool 已经使用该入口预计算 indexer page indices：
```text
NSATokenToKVPoolHost.begin_load_to_device_op()
_active_load_indexer_page_indices = _get_indexer_page_indices(...)
```
结论：代码结构已经允许“per load op prepare once, per layer reuse”。
当前 MLA KV 主体还没有充分利用这个入口。
### C4. 当前 H2D readiness 是 per-layer event
`LayerDoneCounter` 为每个 producer 维护每层 event：
```text
LayerLoadingEvent.complete(layer_id)
LayerDoneCounter.wait_until_on_stream(layer_id, stream)
```
KV pool 访问路径会等待：
```text
get_key_buffer_for_prefetch(layer_id, stream)
-> wait_layer_transfer_on_stream(layer_id, stream)
```
结论：即使 descriptor 合并，仍必须保留 per-layer complete/wait。不能退回 all-layer wait。
### C5. 当前 L2->L1 load 启动点仍偏晚
当前时序：
```text
PrefillAdder.add_one_req()
-> init_load_back()
-> load_cp() queue load op
ScheduleBatch 创建后
-> ready_to_load_host_cache()
-> start_loading()
-> load_stream per-layer enqueue
forward consume
-> wait layer event
```
这已经是 async per-layer load，但不是完整意义上的“提前两层 prefetch”。
第一层和早期层仍可能等待 H2D。
---
## 2. 问题定义
当前瓶颈不是简单的“load op 没合并”，而是合并层级不够：
```text
已合并：
多 req host_indices/device_indices concat 成一个 CacheOperation
未合并：
page 连续区间未 compact
每层 transfer descriptor 未复用
每层 Python/extension submit 仍线性增长
L2->L1 transfer 未提前到 layer-k+2 的 prefetch window
```
在 cache hit + bs>1 场景下，请求通常形态是：
```text
prefix 很大：L2/host hit
extend 很短：200~2000 tokens
batch size：2~10
layer 数：约 78
```
如果每层都重新构造 descriptor，那么 CPU overhead 约随：
```text
O(layer_num * descriptor_build_cost)
```
增长。这个 overhead 很难被 GPU compute overlap，因为它发生在 transfer submit/control path。
目标是把可复用部分改成：
```text
O(descriptor_build_cost + layer_num * cheap_submit_cost)
```
并为后续：
```text
layer L-2 启动 L2->L1(layer L)
layer L-1 启动 L1 shared-KV prefix prefetch(layer L)
layer L consume 只 wait event
```
提供基础。
---
## 3. 目标架构
### 3.1 同步控制面
保持现有同步控制面：
```text
scheduler admission
match_prefix()
init_load_back()
alloc_pages_with_owners(page_owners)
update prefix_indices / extend_input_len
enqueue CacheOperation
```
这些操作仍在 scheduler 线程完成，因为它们决定：
- radix node 是否 device-resident；
- fresh device page owner pattern；
- request prefix/extend metadata；
- batch plan 输入。
### 3.2 异步数据面
把当前 data plane 从：
```text
for each layer:
build/submit transfer using host_indices/device_indices
```
演进为：
```text
begin_load_to_device_op:
build PreparedLoadDescriptor once
for each layer:
submit layer transfer using PreparedLoadDescriptor + layer_id
end_load_to_device_op:
release descriptor lifetime refs
```
### 3.3 descriptor 生命周期
descriptor 生命周期与 `start_loading()` 的 producer id 绑定：
```text
start_loading()
producer_id = update_producer()
descriptor = prepare(...)
enqueue layer transfers
record_stream / hold tensors
ack_load_queue.append(HiCacheAck(...))
```
descriptor 不跨 batch 复用。
descriptor 必须持有所有异步 transfer 需要的 tensor / pinned metadata 引用，直到 load stream 使用完。
---
## 4. Prepared descriptor 设计
### 4.1 第一版 Python-side descriptor
先增加轻量 Python-side descriptor，降低重复 page index 准备和 wrapper 参数构造。
建议 dataclass：
```python
@dataclass
class PreparedLoadDescriptor:
host_indices: torch.Tensor
device_indices: torch.Tensor
host_page_indices: torch.Tensor | None
device_page_indices: torch.Tensor | None
page_size: int
io_backend: str
layout: str
num_tokens: int
num_pages: int
```
对 MLA KV：
- `host_indices/device_indices` 是 token/page slot 级映射；
- `page_first_direct` / `layer_page_first` 传给 TAI direct op；
- 第一版不改变底层 TAI API，只把 layer-invariant tensor 准备和参数绑定集中到 `begin_load_to_device_op()`。
对 NSA index：
- 复用现有 `_active_load_indexer_page_indices`；
- 把它纳入同一个 descriptor 生命周期，避免 target KV 和 index KV 分别散落状态。
### 4.2 第二版 compact segment descriptor
在 Python-side descriptor 稳定后，增加 segment compact：
```python
@dataclass
class TransferSegment:
host_start: int
device_start: int
length: int
```
compact 规则：
```text
如果 host_indices 和 device_indices 同时连续:
合并为一个 segment
否则:
保持 page/token 级 entries
```
收益：
- 减少 cudaMemcpyBatch / TAI descriptor entries；
- 与 L1/L2 allocator 连续分配优化协同；
- 对 `layer_page_first` 更有价值。
### 4.3 第三版 TAI prepared descriptor API
最终把 descriptor 下沉到 tai-kernel：
```python
desc = tai_kernel.nsa_prefill.prepare_h2d_page_descriptor(
host_indices,
device_indices,
page_size,
layout,
)
tai_kernel.nsa_prefill.submit_h2d_layer(
desc,
src_ptr,
dst_ptr,
layer_id,
)
```
TAI descriptor 内部可以选择：
- cudaMemcpyBatchAsync entries；
- compacted ranges；
- future 0SM / copy-engine queue；
- future CUDA driver batch copy API fallback。
SGLang 侧只负责 descriptor 生命周期和 per-layer event。
---
## 5. Phase 计划
### P0. 记录现状和保护合同
**目标：** 明确当前行为，避免后续优化误改语义。
**工作：**
1. 在本文件记录：
- current batch `CacheOperation.merge_ops` 已经做 concat；
- `start_loading()` 每层调用 `load_to_device_per_layer()`；
- `LayerDoneCounter` 是 per-layer wait 合同；
- descriptor 只能跨 layer，不能跨 batch。
2. 不改代码。
**验证：**
- 文档自检无 “TBD/TODO/后续补充” 占位。
### P1. 为 MLA KV 主体增加 Python-side prepared descriptor
**目标：** 使用已有 `begin_load_to_device_op()` / `end_load_to_device_op()` 入口，为 MLA KV 主体缓存 layer-invariant 状态。
**涉及文件：**
- `python/sglang/srt/mem_cache/memory_pool_host.py`
**设计：**
1. 在 host pool base / MLA host pool 上添加 `_active_load_descriptor` 字段。
2. `begin_load_to_device_op()` 创建 descriptor。
3. `load_to_device_per_layer()` 优先使用 descriptor。
4. `end_load_to_device_op()` 清理 descriptor。
5. 若 descriptor 缺失，保留当前路径作为 fail-fast 或 warning fallback；生产 fast path 不应静默 fallback。
**测试：**
- unit test 覆盖：
- begin 后 descriptor 存在；
- per-layer 调用复用 descriptor；
- end 后 descriptor 清空；
- missing descriptor 的 fallback/warning 行为明确。
### P2. 把 NSA index load descriptor 生命周期合并到统一结构
**目标：** 当前 index 已经有 `_active_load_indexer_page_indices`，但状态分散。把 index descriptor 纳入统一 prepared descriptor。
**涉及文件：**
- `python/sglang/srt/mem_cache/memory_pool_host.py`
**设计：**
1. descriptor 内包含：
- MLA KV token/page mapping；
- NSA index page mapping；
- active index layer ids。
2. `NSATokenToKVPoolHost.begin_load_to_device_op()` 调 base prepare 后补充 index fields。
3. `_load_indexer_to_device_per_layer()` 从 descriptor 读取 prepared index page indices。
**测试：**
- active index layer 被正确 load；
- inactive index layer 不访问 compact index cache；
- `index_topk_freq > 1` 时 descriptor 不请求 inactive layers。
### P3. 增加 descriptor build/submit timing
**目标：** 用现有 timing env 验证 CPU overhead 是否下降，不新增长期 noisy log。
**涉及文件：**
- `python/sglang/srt/managers/cache_controller.py`
- `python/sglang/srt/mem_cache/memory_pool_host.py`
**设计：**
复用：
```text
SGLANG_CP_SHARED_KV_BS_GT1_TIMING
SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT
SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS
```
记录阶段：
```text
prepare_load_descriptor
submit_h2d_layer_loop
submit_h2d_layer_per_call_slow
end_load_descriptor
```
**测试：**
- env disabled 时不输出；
- env enabled 且 slow threshold 命中时输出；
- 不新增独立 debug env。
### P4. TAI microbenchmark：descriptor reuse baseline
**目标：** 独立验证 descriptor reuse 是否降低 CPU submit overhead。
**涉及仓库：**
- `tai-kernel`
**benchmark 范围：**
1. tokens/pages：
- 4k
- 16k
- 64k
- 120k
- 192k
2. batch shape：
- bs=1 large prefix
- bs=5 mixed prefix
- bs=10 small extend / large prefix
3. layout：
- `page_first_direct`
- `layer_page_first`
4. page 分布：
- contiguous
- allocator-like owner-lane
- random fragmented
**指标：**
```text
descriptor build ms
per-layer submit total ms
effective H2D GB/s
entries/pages/segments
CPU time per layer
```
### P5. Segment compact
**目标：** 在 descriptor 内合并连续 host/device ranges，减少 bottom-level transfer entries。
**涉及文件：**
- SGLang descriptor builder；
- tai-kernel benchmark；
- 后续 tai-kernel runtime API。
**设计：**
先在 CPU 上做线性 scan：
```text
prev_host + 1 == cur_host
prev_device + 1 == cur_device
```
连续则合并。
不做复杂 merge/sort；不改变 order；不引入额外 O(n log n)。
**测试：**
- contiguous pages 合并成少量 segments；
- fragmented pages 不错误合并；
- device/host 只有一侧连续时不合并；
- segment 展开后与原 indices 完全一致。
### P6. TAI prepared descriptor API
**目标：** 将 descriptor 进一步下沉，减少 Python/C++ wrapper 每层重复工作。
**涉及仓库：**
- `tai-kernel`
- `sglang-dev`
**设计：**
1. tai-kernel 提供 prepare / submit / destroy 三段式 API。
2. SGLang `begin_load_to_device_op()` 调 prepare。
3. `load_to_device_per_layer()` 调 submit。
4. `end_load_to_device_op()` 调 destroy。
5. 如果 tai-kernel 不支持当前 layout/dtype，warning + 明确 fallback 到旧 direct path。
**测试：**
- bf16 / fp8_e4m3；
- `page_first_direct`；
- `layer_page_first`；
- target-only；
- target+draft；
- index active-layer skip。
### P7. 提前两层 L2->L1 transfer prefetch
**目标：** 在 descriptor reuse 稳定后，把 transfer 启动点从 “batch 创建后立刻提交所有层” 改成可控的 layer-ahead prefetch。
**设计方向：**
第一版不改变 scheduler admission，只改变 submit scheduling：
```text
start_loading()
prepare descriptor once
enqueue/submit first N warmup layers
layer end hook 或 forward progress hook:
submit layer_id + 2 transfer
```
约束：
1. 不阻塞 forward stream 发起 transfer。
2. 每层仍通过 `LayerDoneCounter.complete(layer)` 通知 readiness。
3. 如果 layer-ahead submit 来不及，consume wait 仍保证 correctness。
4. draft target 的 layer submit 顺序一致。
**风险：**
- 当前 `start_loading()` 一次性提交所有层，简单且正确；
- 改成逐层提交需要一个可靠 hook，不应复用占 SM 的 shared-KV materialize hook；
- 第一版可以先保留一次性提交，只做 descriptor reuse；提前两层作为独立阶段。
### P8. ETE 验证
**目标：** 证明性能和正确性没有回退。
**验证项：**
1. 单测：
- `test_prefill_adder.py`
- CP shared KV runtime/layout tests
- 新增 descriptor tests
2. 远端 py_compile。
3. TAI benchmark。
4. GSM8K：
- 第一轮 cold/cache-miss；
- 第二轮 cache-hit；
- 精度不掉点。
5. replay workload：
- cache hit 下吞吐；
- accept len；
- prefill failed count；
- timing log 中 descriptor build/submit 是否下降。
---
## 6. 关键不变量
1. `prefix_indices` 更新必须发生在 scheduler admission 阶段，不能等 transfer 完成后异步更新。
2. `device_indices` 分配后，即使 layer 数据尚未 load 完，也可以作为 future-ready L1 slot 暴露给 batch metadata；消费前必须 wait。
3. `LayerDoneCounter` 的 per-layer event 不能被 all-layer event 替代。
4. descriptor 可跨 layer 复用，不可跨 batch 复用。
5. descriptor 不能持有会在 load stream 完成前释放的临时 tensor。
6. zero-owned rank 仍需要 logical no-op ack，保持 CP ranks 的 load ack 顺序一致。
7. target/draft KV 如果共享 HiCache logical node，则 descriptor 生命周期必须覆盖两者；不能 target 成功而 draft 失败后继续可见。
---
## 7. 建议实施顺序
推荐先做：
```text
P1 -> P2 -> P3 -> P4
```
原因：
- 低风险，不改变 transfer 启动时机；
- 能直接验证 CPU overhead 是否来自 descriptor 重建；
- 如果收益明显，再进入 P5/P6；
- 如果收益不明显，说明主要瓶颈可能在 H2D bandwidth、early layer wait、shared-KV L1 prefetch/reduce 或 decode transfer。
不建议一开始就做 P7。
P7 会改 forward/layer hook 生命周期，风险高于 descriptor reuse，应在 P1-P4 有数据后推进。
---
## 8. P4-P6 实施记录（2026-06-11）
### R1. P4 benchmark 已补齐并完成一轮全矩阵验证
新增 `tai-kernel/benchmark/nsa_prefill/benchmark_hicache_h2d_prepared_descriptor.py`。
覆盖维度：
- token 默认：`4096, 16384, 65536, 120000, 192000`；
- batch shape 默认：`1, 5, 10`；
- layout：`page_first_direct`, `layer_page_first`；
- page pattern：`contiguous`, `owner_lane`, `random`，并支持 `fragmented/strided/reverse`；
- 输出：descriptor build ms、per-layer submit ms、submit per layer us、total ms、effective GB/s、segments。
远端 quick smoke：
```text
page_first_direct,contiguous,bs=1,tokens=4096,layers=2,segments=64,total=0.373190ms
layer_page_first,contiguous,bs=1,tokens=4096,layers=2,segments=1,total=0.081027ms
```
P4 全矩阵远端验证：
```text
容器：g0034:cjy-glm5-new
命令：benchmark_hicache_h2d_prepared_descriptor.py
参数：tokens=4096,16384,65536,120000,192000
参数：bs=1,5,10
参数：patterns=contiguous,owner_lane,random
参数：layouts=page_first_direct,layer_page_first
参数：warmup=1 repeat=2
CSV：/mnt/beegfs/cjy/log/hicache_h2d_prepared_descriptor_p4_20260610_210053.csv
行数：90 = 2 layouts * 3 patterns * 3 bs * 5 token sizes
```
代表性结果：
```text
page_first_direct contiguous bs=1 tokens=65536:
segments=1024 build=1.224ms submit=74.595ms total=107.546ms gbps=54.756
layer_page_first contiguous bs=1 tokens=65536:
segments=1 build=0.882ms submit=0.918ms total=106.841ms gbps=55.117
page_first_direct owner_lane bs=10 tokens=192000:
segments=3000 build=4.158ms submit=283.340ms total=315.447ms gbps=54.692
layer_page_first owner_lane bs=10 tokens=192000:
segments=3000 build=3.636ms submit=283.182ms total=315.782ms gbps=54.634
page_first_direct random bs=10 tokens=192000:
segments=3000 build=4.129ms submit=283.274ms total=315.489ms gbps=54.685
layer_page_first random bs=10 tokens=192000:
segments=2998 build=3.665ms submit=282.752ms total=315.268ms gbps=54.723
```
结论：
1. benchmark 已能覆盖 owner-lane 大样本，不再因为 `pool_tokens` 不足失败。实现上按 pattern 动态估算 host page 池，并按 layout lazy 分配 host pinned tensor，避免同时持有 PFD/LPF 两份大 buffer。
2. `layer_page_first` 对“跨 request/跨 page 真实连续”的 contiguous 场景能显著降低 submit CPU：64k 时 `74.6ms -> 0.9ms`，segments `1024 -> 1`。
3. `owner_lane/random` 下 LPF 基本 neutral，因为物理 page 不连续，segments 仍约等于 pages；这说明 LPF 的收益强依赖 host/L2 allocator 能提供按 owner-lane/request 连续的物理 page。
4. 大 transfer 的有效带宽稳定约 `54-55 GB/s`；当前 P4 主要暴露的是 submit/descriptor 开销，不是 H2D 带宽不足。
### R2. P5 segment compact 已实现为保守线性扫描
新增 tai-kernel API：
```python
prepare_h2d_page_descriptor(src_indices, dst_indices, page_size=..., layout=...)
```
合同：
- indices 必须是 CPU int64 contiguous；CUDA indices fail-fast；
- token indices 必须 page aligned；
- `page_first_direct` 固定 layer H2D 不合并相邻 page，因为 host 物理布局是 `[page, layer, page_size, ...]`，相邻 page 的同一 layer 中间隔着其他 layer；
- `layer_page_first` 只在 host 和 device token start 都连续时合并 segment；不 sort、不 reorder、不做 O(n log n)。
新增单测：
```text
tests/nsa_prefill/test_kvcacheio_prepared_descriptor.py
```
验证：
- `page_first_direct` contiguous pages 仍保持每页一个 segment；
- `layer_page_first` contiguous run 合并；
- 只有一侧连续不合并；
- 非 page aligned / CUDA indices fail-fast；
- destroy 后 submit fail-fast。
### R3. P6 tai-kernel prepared API + SGLang 接入完成
新增 tai-kernel public API：
```python
H2DPageDescriptor
H2DPageSegment
prepare_h2d_page_descriptor(...)
submit_h2d_layer(desc, src_ptrs, dst_ptrs, layer_id=...)
destroy_h2d_page_descriptor(desc)
```
SGLang 接入：
- `HostKVCache.begin_load_to_device_op()` 会为 direct + `page_first_direct/layer_page_first` 预构建 `tai_h2d_descriptor`；
- `NSATokenToKVPoolHost.begin_load_to_device_op()` 会额外预构建 index `tai_index_h2d_descriptor`（`page_size=1`）；
- `MLATokenToKVPoolHost.load_to_device_per_layer()` 优先走 `submit_h2d_layer()`；
- NSA index load 优先走 index prepared descriptor；
- `end_load_to_device_op()` 统一 destroy；
- tai API 缺失时 warning fallback 到旧 direct path，不 silent fallback；prepare API 存在但合同不满足时 fail-fast。
当前 P6 已完成到 P6b：
1. Python API 层保留 `H2DPageDescriptor` 合同；
2. tai-kernel C++ 新增：
- `prepare_h2d_page_descriptor_data(...)`
- `submit_h2d_layer_prepared(...)`
3. prepare 阶段在 C++ 中完成 page-aligned 校验、page starts 和 segment starts/pages 构建；
4. submit 阶段直接使用 C++ prepared tensors，不再每层从 Python wrapper 重新做 validation/segment 构造；
5. SGLang 在 direct + `page_first_direct/layer_page_first` 下会检查 `cpp_prepared`，如果 tai Python API 存在但 C++ prepared submit 数据不可用，会输出明确 fallback warning，不 silent fallback。
仍未做的是“真正 opaque C++ custom descriptor / 预生成每层 pointer vector”。当前 C++ prepared submit 仍会在每层构造 `cudaMemcpyBatchAsync` 的 src/dst/size vector；但已把重复 page contract 校验和 segment build 从每层热路径移出。若 ETE timing 仍显示 submit CPU overhead 高，下一步才需要继续做 custom descriptor 或 copy-engine queue。
### R4. 验证记录
本地：
```text
python -m py_compile sglang memory_pool_host.py / tai kvcacheio.py / benchmark script: OK
PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_kvcacheio_prepared_descriptor.py: 4 passed
```
远端 `g0034:cjy-glm5-new`：
```text
/mnt/beegfs/cjy/tai-kernel: test_kvcacheio_prepared_descriptor.py -> 4 passed
/mnt/beegfs/cjy/sglang-dev: test_hicache_controller_cp.py -> 88 passed
benchmark_hicache_h2d_prepared_descriptor.py quick smoke -> OK
```
补充验证（P6b 完成后）：
```text
/mnt/beegfs/cjy/tai-kernel: test_kvcacheio_prepared_descriptor.py -> 6 passed
/mnt/beegfs/cjy/sglang-dev: test_hicache_controller_cp.py -> 88 passed
P4 full matrix benchmark -> 90 rows completed
```
未验证：
- ETE replay / GSM8K cache-hit 正确性；
- 底层 C++ opaque descriptor / copy-engine queue 的进一步 CPU submit 降幅（尚未实现）。

65
python/sglang/srt/managers/cache_controller.py
View File

@@ -78,6 +78,23 @@ def _cp_shared_kv_bs_gt1_cache_timing(
)
def _cp_shared_kv_bs_gt1_cache_timing_start() -> Optional[float]:
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get():
return None
return time.perf_counter()
def _cp_shared_kv_bs_gt1_cache_timing_log(
key: str,
start_time: Optional[float],
message: str,
*args,
) -> None:
if start_time is None:
return
_cp_shared_kv_bs_gt1_cache_timing(key, start_time, message, *args)
class LayerLoadingEvent:
def __init__(self, num_layers: int):
self._num_layers = num_layers
@@ -1752,6 +1769,7 @@ class HiCacheController:
try:
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
prepare_start = _cp_shared_kv_bs_gt1_cache_timing_start()
if op is not None:
target_load_op_started = begin_load_op(
self.mem_pool_host, host_indices, device_indices
@@ -1762,10 +1780,25 @@ class HiCacheController:
draft_host_indices,
draft_device_indices,
)
_cp_shared_kv_bs_gt1_cache_timing_log(
"prepare_load_descriptor",
prepare_start,
"target_tokens=%s draft_tokens=%s target_started=%s draft_started=%s",
int(host_indices.numel()) if host_indices is not None else 0,
int(draft_host_indices.numel())
if draft_host_indices is not None
else 0,
target_load_op_started,
draft_load_op_started,
)
draft_layer_num = (
self.draft_mem_pool_device.layer_num if draft_op is not None else 0
)
layer_loop_start = _cp_shared_kv_bs_gt1_cache_timing_start()
for i in range(max(self.layer_num, draft_layer_num)):
layer_submit_start = _cp_shared_kv_bs_gt1_cache_timing_start()
submitted_target = False
submitted_draft = False
if draft_op is not None and i < draft_layer_num:
if len(draft_host_indices) > 0:
self.draft_mem_pool_host.load_to_device_per_layer(
@@ -1775,6 +1808,7 @@ class HiCacheController:
i,
self.io_backend,
)
submitted_draft = True
if op is not None and i < self.layer_num:
if len(host_indices) > 0:
self.mem_pool_host.load_to_device_per_layer(
@@ -1784,9 +1818,32 @@ class HiCacheController:
i,
self.io_backend,
)
submitted_target = True
producer_event.complete(i)
elif op is None and i < self.layer_num:
producer_event.complete(i)
_cp_shared_kv_bs_gt1_cache_timing_log(
"submit_h2d_layer_per_call_slow",
layer_submit_start,
"layer_id=%s target=%s draft=%s target_tokens=%s draft_tokens=%s",
i,
submitted_target,
submitted_draft,
int(host_indices.numel()) if host_indices is not None else 0,
int(draft_host_indices.numel())
if draft_host_indices is not None
else 0,
)
_cp_shared_kv_bs_gt1_cache_timing_log(
"submit_h2d_layer_loop",
layer_loop_start,
"layers=%s target_tokens=%s draft_tokens=%s",
max(self.layer_num, draft_layer_num),
int(host_indices.numel()) if host_indices is not None else 0,
int(draft_host_indices.numel())
if draft_host_indices is not None
else 0,
)
# NOTE: We must save the host indices and device indices here,
# this is because we need to guarantee that these tensors are
# still alive when the load stream is executing.
@@ -1799,10 +1856,18 @@ class HiCacheController:
if draft_device_indices is not None and draft_device_indices.is_cuda:
draft_device_indices.record_stream(self.load_stream)
finally:
end_start = _cp_shared_kv_bs_gt1_cache_timing_start()
if draft_load_op_started:
end_load_op(self.draft_mem_pool_host)
if target_load_op_started:
end_load_op(self.mem_pool_host)
_cp_shared_kv_bs_gt1_cache_timing_log(
"end_load_descriptor",
end_start,
"target_started=%s draft_started=%s",
target_load_op_started,
draft_load_op_started,
)
self.ack_load_queue.append(
HiCacheAck(

295
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -5,8 +5,9 @@ import logging
import threading
import weakref
from collections import defaultdict
from dataclasses import dataclass
from functools import lru_cache, wraps
from typing import Optional
from typing import Optional, Tuple
import psutil
import torch
@@ -54,6 +55,26 @@ if _is_npu:
logger = logging.getLogger(__name__)
@dataclass
class PreparedLoadDescriptor:
"""Layer-invariant metadata for one host->device HiCache load op."""
host_indices: torch.Tensor
device_indices: torch.Tensor
host_page_indices: Optional[torch.Tensor]
device_page_indices: Optional[torch.Tensor]
page_size: int
io_backend: str
layout: str
num_tokens: int
num_pages: int
index_host_page_indices: Optional[torch.Tensor] = None
index_device_page_indices: Optional[torch.Tensor] = None
index_active_layer_ids: Tuple[int, ...] = ()
tai_h2d_descriptor: Optional[object] = None
tai_index_h2d_descriptor: Optional[object] = None
@lru_cache(maxsize=1)
def _load_tai_transfer_kv_per_layer_mla_lf_pf():
try:
@@ -132,6 +153,47 @@ def _load_tai_transfer_kv_per_layer_direct_lpf_lf():
) from exc
@lru_cache(maxsize=1)
def _load_tai_prepare_h2d_page_descriptor():
try:
from tai_kernel.nsa_prefill import prepare_h2d_page_descriptor
return prepare_h2d_page_descriptor
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FALLBACK][missing_tai_prepared_h2d_descriptor] "
"tai_kernel.nsa_prefill.prepare_h2d_page_descriptor is unavailable. "
"Falling back to legacy per-layer direct H2D submit."
) from exc
@lru_cache(maxsize=1)
def _load_tai_submit_h2d_layer():
try:
from tai_kernel.nsa_prefill import submit_h2d_layer
return submit_h2d_layer
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FALLBACK][missing_tai_submit_h2d_layer] "
"tai_kernel.nsa_prefill.submit_h2d_layer is unavailable. "
"Falling back to legacy per-layer direct H2D submit."
) from exc
@lru_cache(maxsize=1)
def _load_tai_destroy_h2d_page_descriptor():
try:
from tai_kernel.nsa_prefill import destroy_h2d_page_descriptor
return destroy_h2d_page_descriptor
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FALLBACK][missing_tai_destroy_h2d_page_descriptor] "
"tai_kernel.nsa_prefill.destroy_h2d_page_descriptor is unavailable."
) from exc
def _raise_layer_page_first_page_buffer_meta_unsupported() -> None:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][layer_page_first_page_buffer_meta_unsupported] "
@@ -315,6 +377,9 @@ class HostKVCache(abc.ABC):
)
self.kv_buffer = self.init_kv_buffer()
self._active_load_descriptor: Optional[PreparedLoadDescriptor] = None
self._missing_load_descriptor_warning_emitted = False
self._missing_tai_prepared_h2d_warning_emitted = False
# A lock for synchronized operations on memory allocation and state transitions.
self.lock = threading.RLock()
@@ -341,9 +406,151 @@ class HostKVCache(abc.ABC):
self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
) -> None:
"""Prepare layer-invariant metadata for one host->device load op."""
host_indices = host_indices.contiguous()
device_indices = device_indices.contiguous()
host_page_indices, device_page_indices = self._prepare_load_page_indices(
host_indices, device_indices
)
num_tokens = int(host_indices.numel())
self._active_load_descriptor = PreparedLoadDescriptor(
host_indices=host_indices,
device_indices=device_indices,
host_page_indices=host_page_indices,
device_page_indices=device_page_indices,
page_size=int(self.page_size),
io_backend=io_backend,
layout=self.layout,
num_tokens=num_tokens,
num_pages=(num_tokens // int(self.page_size))
if num_tokens % int(self.page_size) == 0
else 0,
)
self._active_load_descriptor.tai_h2d_descriptor = (
self._try_prepare_tai_h2d_descriptor(
self._active_load_descriptor,
host_indices=host_indices,
device_indices=device_indices,
page_size=int(self.page_size),
)
)
def end_load_to_device_op(self) -> None:
"""Release metadata prepared by ``begin_load_to_device_op``."""
descriptor = getattr(self, "_active_load_descriptor", None)
if descriptor is not None:
self._destroy_tai_h2d_descriptor(descriptor.tai_h2d_descriptor)
if descriptor.tai_index_h2d_descriptor is not descriptor.tai_h2d_descriptor:
self._destroy_tai_h2d_descriptor(descriptor.tai_index_h2d_descriptor)
self._active_load_descriptor = None
def _prepare_load_page_indices(
self, host_indices: torch.Tensor, device_indices: torch.Tensor
) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
if self.layout not in ("page_first_direct", "layer_page_first"):
return None, None
validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
validate_page_aligned_token_indices(
device_indices, self.page_size, "device_indices"
)
host_page_indices = (
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
)
device_page_indices = (
device_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
)
return host_page_indices.contiguous(), device_page_indices.contiguous()
def _get_active_load_descriptor(
self, io_backend: str
) -> Optional[PreparedLoadDescriptor]:
descriptor = getattr(self, "_active_load_descriptor", None)
if descriptor is None:
self._log_missing_load_descriptor_once(io_backend)
return None
if descriptor.io_backend != io_backend or descriptor.layout != self.layout:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][load_descriptor_contract_mismatch] "
f"active descriptor io/layout=({descriptor.io_backend}, {descriptor.layout}) "
f"but current load uses ({io_backend}, {self.layout})"
)
return descriptor
def _log_missing_load_descriptor_once(self, io_backend: str) -> None:
if io_backend != "direct":
return
if self.layout not in ("page_first_direct", "layer_page_first"):
return
if getattr(self, "_missing_load_descriptor_warning_emitted", False):
return
self._missing_load_descriptor_warning_emitted = True
logger.warning(
"[CP_HICACHE_FALLBACK][load_descriptor] "
"reason=missing_prepared_descriptor pool=%s layout=%s io_backend=%s",
type(self).__name__,
self.layout,
io_backend,
)
def _try_prepare_tai_h2d_descriptor(
self,
descriptor: PreparedLoadDescriptor,
*,
host_indices: torch.Tensor,
device_indices: torch.Tensor,
page_size: int,
) -> Optional[object]:
if descriptor.io_backend != "direct":
return None
if descriptor.layout not in ("page_first_direct", "layer_page_first"):
return None
try:
prepare = _load_tai_prepare_h2d_page_descriptor()
except RuntimeError as exc:
self._log_missing_tai_prepared_h2d_once(str(exc))
return None
try:
tai_descriptor = prepare(
host_indices,
device_indices,
page_size=page_size,
layout=descriptor.layout,
)
if not bool(getattr(tai_descriptor, "cpp_prepared", False)):
self._log_missing_tai_prepared_h2d_once(
"tai prepared descriptor Python API is present, but the "
"C++ prepared submit data is unavailable; legacy per-layer "
"direct submit remains active."
)
return tai_descriptor
except Exception as exc:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][tai_prepared_h2d_descriptor_failed] "
f"pool={type(self).__name__} layout={descriptor.layout} "
f"page_size={page_size} tokens={int(host_indices.numel())}"
) from exc
def _log_missing_tai_prepared_h2d_once(self, message: str) -> None:
if getattr(self, "_missing_tai_prepared_h2d_warning_emitted", False):
return
self._missing_tai_prepared_h2d_warning_emitted = True
logger.warning(
"[CP_HICACHE_FALLBACK][tai_prepared_h2d_descriptor] "
"reason=missing_tai_api %s",
message,
)
def _destroy_tai_h2d_descriptor(self, tai_descriptor: Optional[object]) -> None:
if tai_descriptor is None:
return
try:
destroy = _load_tai_destroy_h2d_page_descriptor()
destroy(tai_descriptor)
except Exception as exc:
logger.warning(
"[CP_HICACHE_FALLBACK][tai_prepared_h2d_descriptor] "
"reason=destroy_failed error=%s",
exc,
)
@abc.abstractmethod
def backup_from_device_per_layer(
@@ -1503,6 +1710,10 @@ class MLATokenToKVPoolHost(HostKVCache):
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
descriptor = self._get_active_load_descriptor(io_backend)
if descriptor is not None:
host_indices = descriptor.host_indices
device_indices = descriptor.device_indices
if io_backend == "kernel":
if self.layout == "layer_first":
transfer_kv_per_layer_mla(
@@ -1534,6 +1745,14 @@ class MLATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
if descriptor is not None and descriptor.tai_h2d_descriptor is not None:
_load_tai_submit_h2d_layer()(
descriptor.tai_h2d_descriptor,
[self.kv_buffer],
[device_pool.kv_buffer[layer_id]],
layer_id=layer_id,
)
return
_load_tai_transfer_kv_per_layer_direct_pf_lf()(
src_ptrs=[self.kv_buffer],
dst_ptrs=[device_pool.kv_buffer[layer_id]],
@@ -1543,6 +1762,14 @@ class MLATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "layer_page_first":
if descriptor is not None and descriptor.tai_h2d_descriptor is not None:
_load_tai_submit_h2d_layer()(
descriptor.tai_h2d_descriptor,
[self.kv_buffer],
[device_pool.kv_buffer[layer_id]],
layer_id=layer_id,
)
return
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
src_ptrs=[self.kv_buffer],
dst_ptrs=[device_pool.kv_buffer[layer_id]],
@@ -2033,13 +2260,32 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
def begin_load_to_device_op(
self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
) -> None:
super().begin_load_to_device_op(host_indices, device_indices, io_backend)
self._active_load_indexer_page_indices = None
self._active_load_indexer_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
descriptor = self._get_active_load_descriptor(io_backend)
host_page_indices, device_page_indices = self._get_indexer_page_indices(
descriptor.host_indices,
descriptor.device_indices,
)
descriptor.index_host_page_indices = host_page_indices
descriptor.index_device_page_indices = device_page_indices
descriptor.index_active_layer_ids = tuple(
getattr(self, "index_active_layer_ids", ())
)
descriptor.tai_index_h2d_descriptor = self._try_prepare_tai_h2d_descriptor(
descriptor,
host_indices=host_page_indices,
device_indices=device_page_indices,
page_size=1,
)
self._active_load_indexer_page_indices = (
descriptor.index_host_page_indices,
descriptor.index_device_page_indices,
)
def end_load_to_device_op(self) -> None:
self._active_load_indexer_page_indices = None
super().end_load_to_device_op()
def _load_indexer_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
@@ -2048,13 +2294,26 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
return
device_layer_slot = self._device_index_layer_slot(device_pool, layer_id)
host_layer_slot = self._host_index_layer_slot(layer_id)
prepared_indices = getattr(self, "_active_load_indexer_page_indices", None)
if prepared_indices is None:
descriptor = self._get_active_load_descriptor(io_backend)
if (
descriptor is not None
and descriptor.index_host_page_indices is not None
and descriptor.index_device_page_indices is not None
):
host_page_indices = descriptor.index_host_page_indices
device_page_indices = descriptor.index_device_page_indices
else:
prepared_indices = getattr(self, "_active_load_indexer_page_indices", None)
if prepared_indices is None:
host_page_indices, device_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
else:
host_page_indices, device_page_indices = prepared_indices
if host_page_indices is None or device_page_indices is None:
host_page_indices, device_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
else:
host_page_indices, device_page_indices = prepared_indices
use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
if use_kernel:
if self.layout == "layer_first":
@@ -2089,6 +2348,17 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
page_size=1,
)
elif self.layout == "page_first_direct":
if (
descriptor is not None
and descriptor.tai_index_h2d_descriptor is not None
):
_load_tai_submit_h2d_layer()(
descriptor.tai_index_h2d_descriptor,
[self.index_k_with_scale_buffer],
[device_pool.index_k_with_scale_buffer[device_layer_slot]],
layer_id=host_layer_slot,
)
return
_load_tai_transfer_kv_per_layer_direct_pf_lf()(
src_ptrs=[self.index_k_with_scale_buffer],
dst_ptrs=[
@@ -2100,6 +2370,17 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
page_size=1,
)
elif self.layout == "layer_page_first":
if (
descriptor is not None
and descriptor.tai_index_h2d_descriptor is not None
):
_load_tai_submit_h2d_layer()(
descriptor.tai_index_h2d_descriptor,
[self.index_k_with_scale_buffer],
[device_pool.index_k_with_scale_buffer[device_layer_slot]],
layer_id=host_layer_slot,
)
return
_load_tai_transfer_kv_per_layer_direct_lpf_lf()(
src_ptrs=[self.index_k_with_scale_buffer],
dst_ptrs=[

252
test/registered/unit/managers/test_hicache_controller_cp.py
View File

@@ -113,6 +113,7 @@ from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
PreparedLoadDescriptor,
)
from sglang.srt.mem_cache.radix_cache import TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
@@ -206,6 +207,207 @@ class DummyHostKVCacheForAlloc(HostKVCache):
pass
class TestPreparedLoadDescriptor(CustomTestCase):
def test_host_begin_load_builds_page_aligned_descriptor(self):
host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
host_pool.page_size = 4
host_pool.layout = "page_first_direct"
host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
device_indices = torch.tensor(
[40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
)
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
desc = host_pool._active_load_descriptor
self.assertIsInstance(desc, PreparedLoadDescriptor)
self.assertTrue(torch.equal(desc.host_indices, host_indices))
self.assertTrue(torch.equal(desc.device_indices, device_indices))
self.assertEqual(desc.num_tokens, 8)
self.assertEqual(desc.num_pages, 2)
self.assertEqual(desc.layout, "page_first_direct")
self.assertEqual(desc.io_backend, "direct")
self.assertEqual(desc.host_page_indices.tolist(), [2, 5])
self.assertEqual(desc.device_page_indices.tolist(), [10, 16])
host_pool.end_load_to_device_op()
self.assertIsNone(host_pool._active_load_descriptor)
def test_nsa_begin_load_attaches_indexer_pages_to_descriptor(self):
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.page_size = 4
host_pool.layout = "page_first_direct"
host_pool.index_active_layer_ids = (0, 2)
host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
device_indices = torch.tensor(
[40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
)
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
desc = host_pool._active_load_descriptor
self.assertEqual(desc.index_active_layer_ids, (0, 2))
self.assertEqual(desc.index_host_page_indices.tolist(), [2, 5])
self.assertEqual(desc.index_device_page_indices.tolist(), [10, 16])
self.assertIs(host_pool._active_load_indexer_page_indices[0], desc.index_host_page_indices)
self.assertIs(
host_pool._active_load_indexer_page_indices[1],
desc.index_device_page_indices,
)
host_pool.end_load_to_device_op()
self.assertIsNone(host_pool._active_load_descriptor)
self.assertIsNone(host_pool._active_load_indexer_page_indices)
def test_missing_direct_load_descriptor_warns_once(self):
host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
host_pool.page_size = 4
host_pool.layout = "page_first_direct"
with self.assertLogs(
"sglang.srt.mem_cache.memory_pool_host", level="WARNING"
) as logs:
self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
warnings = [
line for line in logs.output if "missing_prepared_descriptor" in line
]
self.assertEqual(len(warnings), 1)
def test_mla_direct_load_uses_prepared_tai_descriptor_when_available(self):
calls = []
fake_desc = object()
def fake_prepare(src_indices, dst_indices, **kwargs):
calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
return fake_desc
def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))
def fake_destroy(desc):
calls.append(("destroy", desc))
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
return_value=fake_prepare,
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
return_value=fake_submit,
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_destroy_h2d_page_descriptor",
return_value=fake_destroy,
),
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
host_pool.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
host_pool.end_load_to_device_op()
self.assertEqual(calls[0][0], "prepare")
self.assertEqual(calls[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(calls[0][2].tolist(), [12, 13, 14, 15])
self.assertEqual(calls[0][3], {"page_size": 4, "layout": "page_first_direct"})
self.assertEqual(calls[1][0], "submit")
self.assertIs(calls[1][1], fake_desc)
self.assertEqual(calls[1][2][0].data_ptr(), host_pool.kv_buffer.data_ptr())
self.assertEqual(
calls[1][3][0].data_ptr(), device_pool.kv_buffer[2].data_ptr()
)
self.assertEqual(calls[1][4], {"layer_id": 2})
self.assertEqual(calls[2], ("destroy", fake_desc))
def test_nsa_index_direct_load_uses_prepared_tai_index_descriptor(self):
calls = []
fake_base_desc = object()
fake_index_desc = object()
def fake_prepare(src_indices, dst_indices, **kwargs):
calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
if kwargs["page_size"] == 1:
return fake_index_desc
return fake_base_desc
def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_active_layer_ids = (0, 1, 2)
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), dtype=torch.uint8
)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
return_value=fake_prepare,
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
return_value=fake_submit,
),
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
host_pool._load_indexer_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="direct",
)
self.assertEqual([c[0] for c in calls[:2]], ["prepare", "prepare"])
self.assertEqual(calls[1][1].tolist(), [1])
self.assertEqual(calls[1][2].tolist(), [3])
self.assertEqual(calls[1][3], {"page_size": 1, "layout": "page_first_direct"})
self.assertEqual(calls[2][0], "submit")
self.assertIs(calls[2][1], fake_index_desc)
self.assertEqual(
calls[2][2][0].data_ptr(), host_pool.index_k_with_scale_buffer.data_ptr()
)
self.assertEqual(
calls[2][3][0].data_ptr(),
device_pool.index_k_with_scale_buffer[1].data_ptr(),
)
self.assertEqual(calls[2][4], {"layer_id": 1})
class FakeDevicePool:
device = "cpu"
layer_num = 1
@@ -878,9 +1080,15 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
host_pool._get_indexer_page_indices = counting_getter
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
with (
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
side_effect=RuntimeError("missing prepared descriptor api"),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
),
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
@@ -1898,6 +2106,44 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
self.assertEqual(host_pool.end_calls, 1)
self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])
def test_start_loading_emits_descriptor_timing_when_enabled(self):
host_pool = FakeHostPool(
torch.tensor([100, 101, 102, 103], dtype=torch.int64)
)
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
allocator.device_pool = FakeDevicePool(layer_num=2)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
controller.load_cp([node], node_id=114)
timing_keys = []
def record_timing(key, start_time, message, *args):
timing_keys.append(key)
with patch(
"sglang.srt.managers.cache_controller.envs."
"SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get",
return_value=True,
), patch(
"sglang.srt.managers.cache_controller."
"_cp_shared_kv_bs_gt1_cache_timing",
side_effect=record_timing,
):
controller.start_loading()
self.assertIn("prepare_load_descriptor", timing_keys)
self.assertIn("submit_h2d_layer_loop", timing_keys)
self.assertIn("submit_h2d_layer_per_call_slow", timing_keys)
self.assertIn("end_load_descriptor", timing_keys)
def test_cp_load_frees_unexpected_owner_allocator_length(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))