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

Stabilize CP HiCache page-first direct transfers on CUDA 13

Browse Source 
CP HiCache direct/page_first_direct all-layer backup was still able to enter sgl-kernel's stale cudaMemcpyBatchAsync path, which segfaults under CUDA 13 before Python can surface an error. The SGLang route now avoids that all-layer sgl-kernel op for page_first_direct backup and uses the TAI per-layer direct LF->PF op for MHA, MLA, and NSA indexer data.\n\nThe load-back path also prepares NSA indexer page indices once per load op and reuses them across per-layer H2D loads, preserving per-layer overlap while removing redundant page-index derivation.\n\nConstraint: Remote runtime is CUDA 13.0 where sgl-kernel's all-layer direct LF->PF op uses the wrong cudaMemcpyBatchAsync ABI.\nRejected: Patch sgl-kernel in this branch | we are converging production HiCache direct/page_first_direct paths onto tai-kernel and do not want to maintain another CUDA-ABI-sensitive copy path here.\nRejected: Collapse H2D load-back into one all-layer op | that would reduce submit count but lose per-layer completion visibility and forward overlap.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce sgl_kernel.transfer_kv_all_layer_direct_lf_pf for direct/page_first_direct HiCache backup without CUDA 13 ABI verification.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q -s test/registered/unit/mem_cache/test_nsa_pool_host_unit.py -> 10 passed, 3 warnings.\nTested: g0034 container: PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py -> 61 passed, 3 warnings.\nTested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py test/registered/unit/mem_cache/test_nsa_pool_host_unit.py test/registered/unit/managers/test_hicache_controller_cp.py\nNot-tested: Full ETE prefill/decode traffic after this commit.\nNot-tested: sgl-kernel implementation itself remains unchanged.
This commit is contained in:
laoyao0822
2026-05-31 01:26:07 +08:00
parent b328baec7c
commit 251a48fb0a
5 changed files with 843 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

447
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
View File

@@ -2918,3 +2918,450 @@ Runtime validation still needed:
- Start a fresh prefill/decode process with the synced code and `SGLANG_EAGLE_ACCEPT_DEBUG=1`.
- Expected prefill evidence: draft cache-hit suffix rows should log `draft_partial_current_reuse_disabled`; draft `branch=partial_current_sync` should disappear except for old logs.
- Expected decode evidence: accept distribution should move away from the `avg≈0.068 / median=0` collapse if C68 is the active root cause.
### C69. 2026-05-30 fresh remote run: parameters correct, draft partial splice disabled, accept still collapsed
Checked runtime:
- Prefill log: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260530_143043.log`.
- Decode logs: `/mnt/beegfs/cjy/log/decode0_20260530_143046.log`, `/mnt/beegfs/cjy/log/decode1_20260530_143049.log`.
- Prefill process args include `--enable-nsa-prefill-cp-shared-kv`, `--attn-cp-size 8`, `--enable-hierarchical-cache`, `--hicache-size 150`, `--hicache-io-backend direct`, and `--hicache-mem-layout page_first_direct`.
- Prefill env includes `SGLANG_CP_DRAFT_SHARED_KV=1`, `SGLANG_CP_SHARED_KV_CURRENT_REUSE=1`, `SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1`, `SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1`, `SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE=1`, `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1`, and `SGLANG_EAGLE_ACCEPT_DEBUG=1`.
Findings:
- C68 gate is active: prefill marker counts are `draft_partial_current_reuse_disabled=64`, `branch=partial_current_sync=0`, `branch=current_only=64`, `metadata_set=128`, `cp_draft_hidden=256`.
- Draft cache-hit suffixes now use `branch=full_materialize` with `can_current_reuse=False`, as intended by C68.
- EAGLE handoff metadata is non-zero after warmup and decode `metadata_get` matches `prebuilt_state` for sampled rows; the prebuilt-state copy path is not obviously corrupting metadata.
- Accept length did **not** recover:
- `decode0_20260530_143046.log`: `n=34784`, `avg=0.075`, `median=0.0`, `last100=0.235`, `zero=32516`.
- `decode1_20260530_143049.log`: `n=30638`, `avg=0.060`, `median=0.0`, `last100=0.082`, `zero=29006`.
Conclusion:
- C68 restored the intended draft full-materialize path but is not the accept-collapse root cause by itself.
- The next likely surfaces are draft KV content/row-order correctness during prefill-to-decode transfer, or a decode-side EAGLE interpretation issue after `prebuilt_state` receives valid non-zero metadata.
### C70. 2026-05-30 accept collapse root-cause candidate: TAI current-slot fill assumes contiguous logical page ids
Finding:
- After the page-granular/no-short-fallback change, the hot target MLA partial-current path uses `fill_current_kv_page_slots_and_remap_locs()` to place current suffix KV into page-slot dense buffers.
- With `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1`, this dispatches to `tai_kernel.nsa_prefill.fill_current_token_kv_page_slots_and_remap_locs`.
- The TAI masking kernel `_mask_current_page_slack_locs_direct_kernel` treats all pages between `first_current_page` and `last_current_page` as current pages and derives `candidate_current_row` by page ordinal:
`candidate = (logical_page - first_current_page) * page_size + token_offset - first_current_offset`.
- CP/paged KV allocator locs are page-based cache-pool addresses; current suffix pages are not guaranteed to have contiguous logical page ids. A suffix can use pages `[5, 25]` while unrelated prefix/query locs include page `10`.
Remote proof on g0034 container:
```text
page_size=4
page_inverse: logical page 5 -> dense page 1, page 10 -> dense page 2, page 25 -> dense page 3
logical_locs=[20,21,40,41,100,101,102,103]
current_locs=[20,21,100,101]
TAI mixed_locs=[[4, 5, -1, -1, -1, -1, -1, -1]]
TAI current_mask=[[True, True, False, False, False, False, False, False]]
expected mixed_locs should preserve unrelated prefix page 10 and mark current page 25 rows:
[[4, 5, 8, 9, 12, 13, -1, -1]]
expected current_mask:
[[True, True, False, False, True, True, False, False]]
```
Impact:
- This directly explains why accept length collapses while metadata transfer still looks valid: target prefill attention can silently hide valid prefix/current KV rows in the partial-current path, producing wrong next-token/hidden handoff for EAGLE.
- The bug is not in the Python torch reference: Python uses exact page membership (`torch.unique(current_pages)`) and exact current-loc membership. The bug is the TAI kernel's contiguous-page assumption.
- This is consistent with the regression appearing after the page-slot current-compose path became hot and TAI materialize was enabled.
Next fix:
- Fix TAI `fill_current_token_kv_page_slots_and_remap_locs` so current/slack detection uses exact `page_inverse` page membership plus exact `current_locs` membership, or a search over `current_locs`, not a contiguous page interval.
- Add a CUDA unit test with non-contiguous current pages and an unrelated prefix page between them.
- Until that test passes, do not use this TAI kernel as evidence that page-aligned CP shared-KV semantics are correct.
C70 implementation update:
- `tai-kernel` now builds exact dense-row and dense-page masks from `current_locs` after remapping them through `page_inverse`.
- The slack-mask kernel no longer infers current membership from a contiguous logical page interval. It now uses:
- `current_row_mask[dense_row]` for exact current-token membership;
- `current_page_mask[dense_page]` for page-slack masking.
- This preserves unrelated prefix pages whose logical page ids fall numerically between two sparse current suffix pages.
C70 verification:
```text
remote g0034 container, /mnt/beegfs/cjy/tai-kernel:
PYTHONPATH=python python -m pytest -q \
tests/nsa_prefill/test_cp_shared_kv_materialize.py::test_fill_current_token_kv_page_slots_handles_sparse_current_pages -q
-> passed
PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_cp_shared_kv_materialize.py
-> 10 passed
runtime import check:
tai_kernel.nsa_prefill.cp_shared_kv_materialize.__file__
-> /mnt/beegfs/cjy/tai-kernel/python/tai_kernel/nsa_prefill/cp_shared_kv_materialize.py
hasattr(..., "_zero_current_dense_masks_kernel") -> True
```
Runtime validation still needed:
- Restart prefill/decode with the updated `tai-kernel` and the same `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1` setting.
- Expected prefill evidence: target `forward_partial_current_reuse` can remain active, but accept length should no longer collapse to `avg≈0.06 / median=0` if C70 was the dominant root cause.
### C71. 2026-05-31 post-C70 runtime observation: aggregate decode accept recovered, sampled EAGLE rows still zero-heavy
Runtime observed:
- Prefill process on `g0034` started at `2026-05-30 23:52:29 CST` with:
- `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1`
- `SGLANG_CP_SHARED_KV_CURRENT_REUSE=1`
- `SGLANG_EAGLE_ACCEPT_DEBUG=1`
- `--enable-nsa-prefill-cp-shared-kv`
- `--hicache-io-backend direct --hicache-mem-layout page_first_direct`
- The synced TAI file timestamp is `2026-05-30 22:57:41 CST`, before this restart.
- Container import resolves to `/mnt/beegfs/cjy/tai-kernel/python/tai_kernel/nsa_prefill/cp_shared_kv_materialize.py` and exposes `_zero_current_dense_masks_kernel`, so the new process should be using the C70 fixed module path.
Fresh log evidence at `2026-05-31 00:00-00:01 CST`:
```text
current decode0_20260530_155232.log:
Decode batch accept: n=123 avg=2.2476 median=2.25 p10=1.93 p90=2.52 zero=0 last100=2.2534
EAGLE_ACCEPT_DEBUG avg_accept: n=1875 avg=0.1229 median=0 zero=1657 last100=0.2900
sampled accepted_draft: n=1875 avg=0.0245 median=0 zero=1848 last100=0.0000 output_zero=0
current decode1_20260530_155236.log:
Decode batch accept: n=118 avg=2.1740 median=2.15 p10=1.88 p90=2.45 zero=0 last100=2.1797
EAGLE_ACCEPT_DEBUG avg_accept: n=1903 avg=0.0851 median=0 zero=1741 last100=0.0300
sampled accepted_draft: n=1903 avg=0.0231 median=0 zero=1880 last100=0.0100 output_zero=0
previous bad decode0_20260530_143046.log:
Decode batch accept: n=1773 avg=1.3924 median=1.2 zero=0 last100=1.6584
EAGLE_ACCEPT_DEBUG avg_accept: n=52797 avg=0.0548 median=0 zero=48785 last100=0.0300
previous bad decode1_20260530_143049.log:
Decode batch accept: n=1729 avg=1.4446 median=1.3 zero=0 last100=1.7419
EAGLE_ACCEPT_DEBUG avg_accept: n=49015 avg=0.0666 median=0 zero=44342 last100=0.0000
```
Interpretation:
- The primary aggregate `Decode batch, accept len` signal has recovered above both the immediate bad run and the older no-debug baseline-ish run (`~1.3`). It is not showing the severe accept-collapse pattern.
- No `Traceback`, `Health check failed`, `detokenizer hang`, `KVCapacityWaitError`, or output-zero sample was found in the current prefill/decode logs during this observation window.
- The per-step `EAGLE_ACCEPT_DEBUG` sample remains zero-heavy because it reports sampled rows/steps, not the same aggregate as `Decode batch accept len`. Treat it as secondary evidence; do not use it alone to claim accept collapse while aggregate decode accept remains healthy.
- Remaining visible fallbacks are still present and should be cleaned separately:
- `prefix_not_page_aligned` during zero-prefix/early requests;
- `draft_partial_current_reuse_disabled` by the current draft safety gate.
Status:
- C70 is supported by runtime evidence as the dominant cause of the previous aggregate accept-length regression.
- Continue longer ETE observation before calling the whole run stable, especially under host-cache-full pressure.
C71 continued observation at `2026-05-31 00:03 CST`:
```text
current decode0_20260530_155232.log:
Decode batch accept: n=307 avg=2.2528 median=2.25 p10=1.93 p90=2.55 zero=0 last100=2.2494
EAGLE_ACCEPT_DEBUG avg_accept: n=4828 avg=0.1984 median=0 zero=3902 last100=0.7492
sampled accepted_draft output_zero=0
current decode1_20260530_155236.log:
Decode batch accept: n=309 avg=2.2042 median=2.20 p10=1.90 p90=2.52 zero=0 last100=2.2055
EAGLE_ACCEPT_DEBUG avg_accept: n=5021 avg=0.1782 median=0 zero=4152 last100=0.3650
sampled accepted_draft output_zero=0
```
No new runtime error was found in the current logs in this observation window except warmup-time max-m reduction warnings.
### C72. 2026-05-31 second runtime log recheck: accept remains stable after longer run
Observed at `2026-05-31 00:17 CST`:
- Prefill `g0034` and decode `g0035/g0036` launch_server processes are still alive.
- Latest logs remain:
- `sglang_cp_hicache_20260530_155230.log`
- `decode0_20260530_155232.log`
- `decode1_20260530_155236.log`
Fresh decode aggregate accept evidence:
```text
decode0_20260530_155232.log:
Decode batch accept: n=1228 avg=2.2488 median=2.23 p10=1.93 p90=2.58 zero=0 last100=2.1639
EAGLE_ACCEPT_DEBUG avg_accept: n=18811 avg=0.2070 median=0 p90=1 zero=14988 last100=0.0200
sampled accepted_draft output_zero=0
decode1_20260530_155236.log:
Decode batch accept: n=1211 avg=2.2405 median=2.23 p10=1.93 p90=2.57 zero=0 last100=2.2233
EAGLE_ACCEPT_DEBUG avg_accept: n=18480 avg=0.1854 median=0 p90=1 zero=15160 last100=0.0700
sampled accepted_draft output_zero=0
```
Error scan:
- No `Traceback`, `Exception`, `ERROR`, `CRITICAL`, `Health check failed`, `detokenizer hang`, `KVCapacityWaitError`, request invalidation, or completion/output-zero signature was found in the current observation window.
- Only warmup-time max-m reduction warnings matched the broad error grep.
Prefill marker counts:
```text
forward_partial_current_reuse 7984
forward_partial_current_sync_compose 1968
branch=full_materialize 7984
branch=partial_current_sync 0
branch=current_only 56
draft_partial_current_reuse_disabled 64
CP_SHARED_KV_FALLBACK 208
prefix_not_page_aligned 144
owner-lane evict END 3144
_evict_for_compute_owner_lanes 6288
KVCapacityWaitError 0
Traceback 0
Health check failed 0
detokenizer hang 0
```
Interpretation:
- Aggregate decode accept remains stable around `2.24` after a longer run, so C70 remains supported as the accept-collapse fix.
- Device owner-lane evictions are occurring frequently but not causing OOM/hang in this window.
- Remaining visible work is performance/fallback cleanup, not an immediate correctness crash in this run.
### C73. 2026-05-31 current fallback log classification
Current runtime `sglang_cp_hicache_20260530_155230.log` fallback reasons:
```text
120 x [CP_SHARED_KV_FALLBACK][mla_prefetch] reason=prefix_not_page_aligned prefix_len=0 page_size=64
120 x [CP_SHARED_KV_FALLBACK][index_prefetch] reason=prefix_not_page_aligned prefix_len=0 page_size=64
64 x [CP_SHARED_KV_FALLBACK][current_reuse] reason=draft_partial_current_reuse_disabled prefix_lens=[320] ...
```
Interpretation:
- `prefix_not_page_aligned` is currently all `prefix_len=0`. It is a prefetch optimization skip for current-only/no-cache-prefix batches, not a non-page-tail cache-hit fallback. Since there is no host/L1 prefix to prefetch, creating MLA/index prefetcher returns `None`. The warning label is misleading for this case; it should likely be classified as `create_skip/no_prefix` or gated before constructing prefetchers.
- `draft_partial_current_reuse_disabled` is the deliberate C68 EAGLE/NextN safety gate. Draft cache-hit partial-current splice is disabled, so draft uses full materialize while target partial-current reuse remains active.
Impact:
- These current fallback logs do not indicate the old page-tail fallback path reappearing.
- They are still noisy and should be cleaned/classified before finalizing logs, but they are not the current accept-length correctness regression.
### C74. 2026-05-31 L2->L1 HiCache load-back async check
Question: whether current L2(host)->L1(device) HiCache load-back/prefetch is truly async.
Code evidence:
- `schedule_policy.py` calls `tree_cache.init_load_back()` when `req.host_hit_length > 0`; this prepares load-back metadata and queues load operations, but does not itself wait for all KV bytes to land.
- `scheduler.py` calls `tree_cache.ready_to_load_host_cache()` before `prepare_for_extend()`; this maps to `HiCacheController.start_loading()` and returns `hicache_consumer_index` to the batch.
- `HiCacheController.start_loading()`:
- merges pending load ops;
- records a producer start event;
- enters `self.load_stream`;
- calls `mem_pool_host.load_to_device_per_layer(...)` for each layer;
- records one per-layer event via `producer_event.complete(i)`;
- appends an ack with the final per-layer event.
- Device KV pool access waits only on the needed layer event:
- normal `get_key_buffer()/get_value_buffer()` waits via `LayerDoneCounter.wait_until(layer_id)`;
- prefetch getter `get_key_buffer_for_prefetch(..., stream)` waits the same event on the prefetch stream using `wait_until_on_stream`.
- With current launch `--hicache-io-backend direct --hicache-mem-layout page_first_direct`, `NSATokenToKVPoolHost.load_to_device_per_layer()` calls `tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_pf_lf` for MLA KV and indexer buffers.
- The TAI H2D op uses `cudaMemcpyBatchAsync` on `at::cuda::getCurrentCUDAStream(device_id)`, so when called under `self.load_stream` it enqueues copy-engine work on the load stream and returns without CPU-side copy completion synchronization.
Runtime evidence from current log:
```text
[HiCache-load] load_back CP: 5424
[HiCache-load] load_back CP SUCCESS: 5424
CP_HICACHE_FALLBACK][cp_load_back: 0
```
Interpretation:
- The bulk L2->L1 KV byte transfer is async with respect to model compute: it is queued on a separate `load_stream`, and model streams wait per layer only when consuming that layer.
- It is not a fully zero-cost async path. There is still synchronous CPU preparation before the async H2D copies are enqueued:
- `load_cp()` does owner-lane planning/allocation/metadata concatenation;
- direct mode moves `device_indices` to CPU before enqueue;
- the TAI op builds CPU-side `cudaMemcpyBatchAsync` pointer/size vectors per layer.
- Current design is therefore best described as: async DMA/copy-engine transfer with per-layer consume waits, but synchronous CPU submit/preparation overhead. If performance still regresses on cache-hit loads, the remaining suspect is enqueue/planning overhead or insufficient lookahead, not blocking H2D memcpy completion.
### C75 — runtime L2→L1 prefetch/load-back async contract re-check
Finding date: 2026-05-31.
Current remote prefill process on g0034 was inspected via `/proc/<pid>/cmdline` and is currently launched with:
- `--hicache-io-backend kernel`
- `--hicache-mem-layout page_first`
- `SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1`
- `SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE=1`
- `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1`
This means current L2(host HiCache) → L1(device KV pool) load-back is **not** using the TAI `page_first_direct`/`cudaMemcpyBatchAsync` 0-SM copy path. It is using the SGLang `kernel + page_first` path:
- `HiRadixCache.load_back()` queues `CacheOperation` via `cache_controller.load_cp()`.
- `Scheduler._get_new_batch_prefill_raw()` calls `tree_cache.ready_to_load_host_cache()` before `prepare_for_extend()`.
- `HiCacheController.start_loading()` switches to `self.load_stream` and calls `mem_pool_host.load_to_device_per_layer()` per layer.
- `MLATokenToKVPoolHost.load_to_device_per_layer()` with `io_backend == "kernel"` and `layout == "page_first"` calls `sgl_kernel.kvcacheio.transfer_kv_per_layer_mla_pf_lf()`.
- `sgl-kernel/csrc/kvcacheio/transfer.cu` launches the transfer kernel on `at::cuda::getCurrentCUDAStream()`, which is `load_stream` in this context, and records per-layer completion events.
- Device KV consumers wait per layer through `LayerDoneCounter.wait_until()` / `wait_until_on_stream()` before reading normal KV or prefetch KV buffers.
Conclusion:
- It is **async in the CUDA stream sense**: load-back kernels are enqueued on a separate load stream, and forward/prefetch consumers wait only at per-layer consumption.
- It is **not async in the 0-SM/DMA-copy-engine sense** for the current launch: the `kernel + page_first` path runs an SM-consuming transfer CUDA kernel.
- The TAI direct path (`--hicache-io-backend direct --hicache-mem-layout page_first_direct`) would enqueue `cudaMemcpyBatchAsync` and fail fast if unavailable, but that is not the current runtime configuration.
Risk / action:
- If we expect L2→L1 load-back to avoid SM interference with attention/index/MLA kernels, current launch parameters must switch back to `direct + page_first_direct`, and remote logs should confirm no `CP_HICACHE_FALLBACK][cp_load_back` and the TAI direct path is actually installed/imported.
- Even with direct mode, CPU-side planning/submission remains synchronous; only bulk bytes are async.
### C76 — correction: C75 inspected a non-target remote process
Finding date: 2026-05-31.
Correction: C75's `g0034` process inspection matched a non-target/stale prefill process launched with `--hicache-io-backend kernel --hicache-mem-layout page_first`. The user's intended launch command for the active experiment is:
- `--hicache-io-backend direct`
- `--hicache-mem-layout page_first_direct`
Therefore C75 must not be used as evidence about the intended experiment's L2→L1 path.
For the intended launch, the code path is:
- `HiCacheController.start_loading()` enters `self.load_stream`.
- `NSATokenToKVPoolHost.load_to_device_per_layer()` / `MLATokenToKVPoolHost.load_to_device_per_layer()` with `io_backend == "direct"` and `layout == "page_first_direct"` call `_load_tai_transfer_kv_per_layer_direct_pf_lf()` for both MLA KV and NSA indexer buffers.
- `tai_kernel.nsa_prefill.kvcacheio.transfer_kv_per_layer_direct_pf_lf()` calls `torch.ops.tai_kernel_ops.transfer_kv_per_layer_direct_pf_lf`.
- The TAI C++ implementation uses `cudaMemcpyBatchAsync` on the current CUDA stream, which is `load_stream` under `start_loading()`, and fails fast if CUDA/runtime support is missing.
Conclusion for the intended launch:
- L2→L1 bulk transfer should be async and copy-engine based, not SM-transfer-kernel based.
- It still has synchronous CPU-side prep/submission overhead before the async copies are queued.
- Runtime verification must match the intended process/log, not a stale/non-target prefill process.
### C77 — L2→L1 load-back CPU overhead granularity and per-request conversion feasibility
Finding date: 2026-05-31.
Question: whether the remaining CPU-side L2→L1 load-back overheads can be moved from per-layer to per-request.
Current granularity:
- Per request / per cache-hit load-back:
- radix/HiCache load planning: `HiRadixCache.init_load_back()` → `load_back()` → `_build_cp_load_back_plan()`;
- owner-lane capacity check: `CPSharedPagedTokenToKVPoolAllocator.compute_owner_lane_stats(page_owners)`;
- optional owner-lane eviction planning: `_plan_cp_load_back_owner_lane_evictions()` / `_evict_cp_load_back_owner_lanes()`;
- owner-lane allocation: `alloc_pages_with_owners(page_owners)`;
- host/device token-index concatenation and queueing: `CacheController.load_cp()`;
- direct-mode `device_indices.cpu()` in `CacheController.move_indices()`.
- Per layer:
- `HiCacheController.start_loading()` calls `load_to_device_per_layer()` once per target layer and draft layer;
- NSA target direct/page_first_direct calls TAI H2D once for MLA KV and once for indexer for each layer;
- TAI `transfer_kv_per_layer_direct_pf_lf()` rebuilds `std::vector` memcpy descriptors and calls `cudaMemcpyBatchAsync()` on every layer;
- `NSATokenToKVPoolHost._get_indexer_page_indices()` recomputes host/device page indices inside every indexer layer load, although the result is layer-invariant.
Feasibility:
1. Safe/low-risk per-request conversion:
- Precompute indexer page indices once per load op and pass/reuse them for all indexer layers.
- Precompute reusable page-pair metadata (`host_page`, `dst_start`) once per load op so the per-layer TAI call does not redo token-index CPU conversion/validation/page extraction.
- Preserve per-layer submit/events so forward overlap semantics remain unchanged.
2. Medium-risk conversion:
- Add a TAI H2D API that accepts precomputed page-pair metadata and layer range, then internally enqueues one `cudaMemcpyBatchAsync` per layer and optionally records/returns per-layer completion ordering. This reduces Python/C++ descriptor preparation churn but must preserve per-layer readiness.
3. High-risk / not recommended as first step:
- Collapse all layers into one giant all-layer `cudaMemcpyBatchAsync` and only signal a final completion event. This makes the submit path per-request, but destroys per-layer load/forward overlap because layer 0 would effectively wait behind all later-layer copies unless a more complex event scheme is added.
Conclusion:
- Yes, some per-layer CPU work can be moved to per-request without changing correctness: indexer page-index derivation and page-pair metadata preparation are the first targets.
- The actual per-layer async submit should initially remain per-layer to preserve overlap and existing `LayerDoneCounter` semantics.
- A full per-request H2D op is possible only if it still preserves per-layer completion visibility; otherwise it trades CPU overhead for lost overlap.
### C78 — implemented per-request indexer page-index preparation for L2→L1 load-back
Finding date: 2026-05-31.
Change implemented:
- `HostKVCache` now exposes optional load-op lifecycle hooks:
- `begin_load_to_device_op(host_indices, device_indices, io_backend)`
- `end_load_to_device_op()`
- `HiCacheController.start_loading()` calls these hooks once per merged target load op and once per merged draft load op, then preserves the existing per-layer `load_to_device_per_layer()` loop and per-layer completion events.
- `NSATokenToKVPoolHost.begin_load_to_device_op()` computes indexer host/device page indices once using `_get_indexer_page_indices()` and stores them in `_active_load_indexer_page_indices`.
- `NSATokenToKVPoolHost._load_indexer_to_device_per_layer()` reuses the prepared indices when present; direct calls without a begin hook still compute indices locally, preserving the old API behavior.
Important stream correction:
- The begin hook is called inside `with device_module.stream(self.load_stream)`, before the per-layer loop.
- This is intentional: in `kernel + page_first` mode the prepared page-index tensors can be CUDA tensors, so deriving them on the load stream avoids introducing a cross-stream producer/consumer race. In `direct + page_first_direct` mode the tensors are CPU, so this has no downside.
What this optimizes:
- It removes the per-layer repeated `reshape(... page_size)[:, 0] // page_size` indexer page-index derivation for load-back.
- It does not yet remove per-layer TAI `cudaMemcpyBatchAsync` descriptor construction for the actual KV/indexer copies.
- It preserves per-layer overlap and `LayerDoneCounter` semantics.
Verification:
- RED on remote before implementation:
- new `test_nsa_indexer_load_reuses_precomputed_page_indices_across_layers` failed with missing `begin_load_to_device_op`;
- new `test_start_loading_prepares_load_op_once_for_all_layers` failed because `start_loading()` did not call the begin hook before per-layer loads.
- GREEN after implementation on g0034 container:
- `python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py test/registered/unit/managers/test_hicache_controller_cp.py`
- `PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py` → `61 passed, 3 warnings`.
- `PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAIndexerPageIndices` → `4 passed, 3 warnings`.
Known verification gap:
- Full `test_nsa_pool_host_unit.py` exited with code 139 after the first CUDA test in the remote container. This is not proven related to the current Python load-op hook change; the CPU-only page-index subset passes. Do not claim the CUDA host-transfer suite is clean until this is isolated separately.
### C79 — 2026-05-30 full `test_nsa_pool_host_unit.py` CUDA exit 139 root cause
Finding: the CUDA 139 is a native segfault in the old `sgl_kernel.kvcacheio.transfer_kv_all_layer_direct_lf_pf` path, not a Python exception and not the new per-request H2D indexer-page precompute hook.
Evidence:
- Remote isolated run: `test_device_to_host_indexer_kernel_layer_first` passed, `test_device_to_host_indexer_direct_layer_first` passed, then `test_device_to_host_indexer_direct_page_first_direct` crashed with `Segmentation fault (core dumped)`.
- Fault stack enters `cuMemcpyBatchAsync_v2` -> `cudaMemcpyBatchAsync` -> `transfer_kv_page_first_direct_impl<true>` -> `transfer_kv_all_layer_direct_lf_pf`.
- Remote container reports `torch.version.cuda == 13.0`, so the CUDA 13 `cudaMemcpyBatchAsync` v2 ABI is active.
- `sgl-kernel/csrc/kvcacheio/transfer.cu` still casts `cudaMemcpyBatchAsync` to the pre-CUDA-13 signature with a `fail_idx` argument in `transfer_kv_page_first_direct_impl`, while `tai-kernel/python/tai_kernel/nsa_prefill/csrc/kvcacheio_lf_pf.cu` already uses a `CUDA_VERSION >= 13000` signature gate without `fail_idx`.
- Re-running the same NSA page-first-direct D2H path with `pin_memory=True` still exits 139, so pinned/pageable host allocation is not the root cause.
Impact:
- The failing unit covers the all-layer D2H backup indexer path for `direct + page_first_direct`; production per-layer async backup/load should prefer the TAI kernel path and avoid this stale sgl-kernel all-layer direct path on CUDA 13.
- Do not use this failing full suite as evidence that the Python H2D precompute change is wrong; the isolated Python/CPU indexer tests and `test_hicache_controller_cp.py` still pass.
Next correction target:
- Either route NSA page-first-direct indexer all-layer D2H backup away from `sgl_kernel.transfer_kv_all_layer_direct_lf_pf` on CUDA 13, or port the CUDA-13-safe `cudaMemcpyBatchAsync` ABI handling from tai-kernel into sgl-kernel. Since this branch intentionally avoids maintaining sgl-kernel for our path, the cleaner SGLang-side fix is to avoid the stale all-layer sgl-kernel op for NSA page-first-direct indexer backup.
### C80 — 2026-05-30 SGLang-side route away from sgl-kernel all-layer direct LF->PF
Decision: follow option 1 from C79. Do not patch `sgl-kernel` for this branch; route `direct + page_first_direct` all-layer D2H backup through the CUDA-13-safe TAI per-layer direct LF->PF op.
Implementation:
- `MHATokenToKVPoolHost.backup_from_device_all_layer(..., io_backend="direct")` now loops layers and calls `tai_kernel.nsa_prefill.transfer_kv_per_layer_direct_lf_pf` for K/V when `layout == "page_first_direct"`.
- `MLATokenToKVPoolHost.backup_from_device_all_layer(..., io_backend="direct")` now loops layers and calls the same TAI op for MLA KV when `layout == "page_first_direct"`.
- `NSATokenToKVPoolHost._backup_indexer_from_device_all_layer(..., io_backend="direct")` now loops layers and calls the same TAI op for NSA indexer pages with `page_size=1`.
Why MHA/MLA had to change too:
- Avoiding only the NSA indexer path was insufficient. `NSATokenToKVPoolHost.backup_from_device_all_layer()` first calls the MLA base backup, and the base MLA `page_first_direct` path still entered `sgl_kernel.transfer_kv_all_layer_direct_lf_pf`, causing the same CUDA 13 segfault before the indexer route ran.
Test corrections:
- `test_nsa_pool_host_unit.py` now has route tests that patch `transfer_kv_all_layer_direct_lf_pf` to fail if the stale sgl-kernel path is used for MHA, MLA, or NSA indexer page-first-direct all-layer backup.
- The CUDA transfer test synchronizes after direct backup before reading CPU host buffers because the TAI memcpy-batch path is asynchronous.
- The page-first-direct NSA indexer expected shape uses `[host_page, layer_id, 0]`; the host indexer buffer layout is `(page, layer, 1, stride)`, not `(page, layer, stride)`.
Verification on g0034 container:
- RED before implementation:
- `TestPageFirstDirectAllLayerBackupRoute` failed because MHA/MLA still called `sgl_kernel.transfer_kv_all_layer_direct_lf_pf`.
- Full `test_nsa_pool_host_unit.py` exited 139 in `cuMemcpyBatchAsync_v2` before this route change.
- GREEN after implementation:
- `PYTHONPATH=python python -m pytest -q -s test/registered/unit/mem_cache/test_nsa_pool_host_unit.py` → `10 passed, 3 warnings`.
- `PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py` → `61 passed, 3 warnings`.
Remaining risk:
- This makes the all-layer direct/page-first backup API submit one TAI per-layer memcpy-batch op instead of one stale sgl-kernel all-layer op. It is correct and CUDA-13-safe, but it can increase CPU submit count on all-layer backup paths. Production forward-overlap paths already operate per-layer, so this is acceptable for the current branch goal.

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

@@ -1517,6 +1517,18 @@ class HiCacheController:
if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0:
return -1
def begin_load_op(mem_pool_host, host_indices, device_indices) -> bool:
begin = getattr(mem_pool_host, "begin_load_to_device_op", None)
if begin is None:
return False
begin(host_indices, device_indices, self.io_backend)
return True
def end_load_op(mem_pool_host) -> None:
end = getattr(mem_pool_host, "end_load_to_device_op", None)
if end is not None:
end()
producer_id = self.layer_done_counter.update_producer()
op = CacheOperation.merge_ops(self.load_queue) if self.load_queue else None
draft_op = (
@@ -1540,44 +1552,62 @@ class HiCacheController:
producer_event = self.layer_done_counter.events[producer_id]
producer_event.start_event.record()
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
draft_layer_num = (
self.draft_mem_pool_device.layer_num if draft_op is not None else 0
)
for i in range(max(self.layer_num, draft_layer_num)):
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(
self.draft_mem_pool_device,
draft_host_indices,
draft_device_indices,
i,
self.io_backend,
)
if op is not None and i < self.layer_num:
if len(host_indices) > 0:
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
device_indices,
i,
self.io_backend,
)
producer_event.complete(i)
elif op is None and i < self.layer_num:
producer_event.complete(i)
# 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.
if host_indices is not None and host_indices.is_cuda:
host_indices.record_stream(self.load_stream)
if device_indices is not None and device_indices.is_cuda:
device_indices.record_stream(self.load_stream)
if draft_host_indices is not None and draft_host_indices.is_cuda:
draft_host_indices.record_stream(self.load_stream)
if draft_device_indices is not None and draft_device_indices.is_cuda:
draft_device_indices.record_stream(self.load_stream)
target_load_op_started = False
draft_load_op_started = False
try:
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
if op is not None:
target_load_op_started = begin_load_op(
self.mem_pool_host, host_indices, device_indices
)
if draft_op is not None:
draft_load_op_started = begin_load_op(
self.draft_mem_pool_host,
draft_host_indices,
draft_device_indices,
)
draft_layer_num = (
self.draft_mem_pool_device.layer_num if draft_op is not None else 0
)
for i in range(max(self.layer_num, draft_layer_num)):
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(
self.draft_mem_pool_device,
draft_host_indices,
draft_device_indices,
i,
self.io_backend,
)
if op is not None and i < self.layer_num:
if len(host_indices) > 0:
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
device_indices,
i,
self.io_backend,
)
producer_event.complete(i)
elif op is None and i < self.layer_num:
producer_event.complete(i)
# 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.
if host_indices is not None and host_indices.is_cuda:
host_indices.record_stream(self.load_stream)
if device_indices is not None and device_indices.is_cuda:
device_indices.record_stream(self.load_stream)
if draft_host_indices is not None and draft_host_indices.is_cuda:
draft_host_indices.record_stream(self.load_stream)
if draft_device_indices is not None and draft_device_indices.is_cuda:
draft_device_indices.record_stream(self.load_stream)
finally:
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)
self.ack_load_queue.append(
HiCacheAck(

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

@@ -266,6 +266,14 @@ class HostKVCache(abc.ABC):
"""
raise NotImplementedError()
def begin_load_to_device_op(
self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
) -> None:
"""Prepare layer-invariant metadata for one host->device load op."""
def end_load_to_device_op(self) -> None:
"""Release metadata prepared by ``begin_load_to_device_op``."""
@abc.abstractmethod
def backup_from_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
@@ -605,13 +613,23 @@ class MHATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.k_buffer + device_pool.v_buffer,
dst_ptrs=[self.k_buffer, self.v_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
# Avoid sgl-kernel's all-layer direct LF->PF path. CUDA 13
# changes the cudaMemcpyBatchAsync ABI, and the stale all-layer
# op can segfault before Python sees an error. Reuse the TAI
# per-layer direct op, which owns the version-specific ABI.
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf()
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":
@@ -1090,13 +1108,20 @@ class MLATokenToKVPoolHost(HostKVCache):
page_size=self.page_size,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.kv_buffer,
dst_ptrs=[self.kv_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
# Avoid sgl-kernel's all-layer direct LF->PF path. CUDA 13
# changes the cudaMemcpyBatchAsync ABI, and the stale all-layer
# op can segfault before Python sees an error. Reuse the TAI
# per-layer direct op, which owns the version-specific ABI.
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf()
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":
@@ -1391,12 +1416,27 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
)
return host_page_indices, device_page_indices
def begin_load_to_device_op(
self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str
) -> None:
self._active_load_indexer_page_indices = None
self._active_load_indexer_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
def end_load_to_device_op(self) -> None:
self._active_load_indexer_page_indices = None
def _load_indexer_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
host_page_indices, device_page_indices = self._get_indexer_page_indices(
host_indices, device_indices
)
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
use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
if use_kernel:
if self.layout == "layer_first":
@@ -1481,13 +1521,21 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
page_size=1,
)
elif self.layout == "page_first_direct":
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.index_k_with_scale_buffer,
dst_ptrs=[self.index_k_with_scale_buffer],
src_indices=device_page_indices,
dst_indices=host_page_indices,
page_size=1,
)
# Avoid sgl-kernel's all-layer direct LF->PF path here. On CUDA 13
# that op still uses the pre-v2 cudaMemcpyBatchAsync ABI and can
# segfault before surfacing a Python exception. Route NSA indexer
# page-first-direct all-layer backup through the TAI per-layer op,
# which owns the CUDA-version-specific memcpy-batch ABI handling.
tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf()
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:

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

@@ -537,6 +537,61 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
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 = []
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 = "page_first_direct"
host_pool.page_size = 4
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)
original_getter = host_pool._get_indexer_page_indices
getter_calls = []
def counting_getter(h, d):
getter_calls.append((h.clone(), d.clone()))
return original_getter(h, d)
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,
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
try:
for layer_id in range(3):
host_pool._load_indexer_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=layer_id,
io_backend="direct",
)
finally:
host_pool.end_load_to_device_op()
self.assertEqual(len(getter_calls), 1)
self.assertEqual(len(calls), 3)
self.assertEqual([call[4]["layer_id"] for call in calls], [0, 1, 2])
self.assertEqual([call[2].tolist() for call in calls], [[1], [1], [1]])
self.assertEqual([call[3].tolist() for call in calls], [[3], [3], [3]])
class FakeAllocator:
def __init__(self, alloc_result=None):
@@ -1109,6 +1164,55 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
def test_start_loading_prepares_load_op_once_for_all_layers(self):
class PreparingFakeHostPool(FakeHostPool):
def __init__(self, alloc_result):
super().__init__(alloc_result)
self.begin_calls = []
self.end_calls = 0
self.active = False
def begin_load_to_device_op(self, host_indices, device_indices, io_backend):
self.begin_calls.append(
(host_indices.clone(), device_indices.clone(), io_backend)
)
self.active = True
def end_load_to_device_op(self):
self.end_calls += 1
self.active = False
def load_to_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
assert self.active
super().load_to_device_per_layer(
device_pool, host_indices, device_indices, layer_id, io_backend
)
host_pool = PreparingFakeHostPool(
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=3)
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=113)
controller.start_loading()
self.assertEqual(len(host_pool.begin_calls), 1)
self.assertEqual(host_pool.end_calls, 1)
self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])
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))

153
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py
View File

@@ -1,10 +1,13 @@
import unittest
from unittest.mock import patch
import torch
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import (
ALLOC_MEMORY_FUNCS,
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
alloc_with_pin_memory,
)
@@ -101,13 +104,16 @@ class TestNSAHiCacheTransfer(CustomTestCase):
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend
)
torch.cuda.synchronize()
for layer_id in range(layer_num):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
if layout == "page_first_direct":
got = host_pool.index_k_with_scale_buffer[host_page, layer_id].cpu()
got = host_pool.index_k_with_scale_buffer[
host_page, layer_id, 0
].cpu()
else:
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
expected = device_pool.index_k_with_scale_buffer[layer_id][
@@ -143,6 +149,102 @@ class TestNSAHiCacheTransfer(CustomTestCase):
)
class TestPageFirstDirectAllLayerBackupRoute(CustomTestCase):
def test_mla_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.layer_num = 2
host_pool.kv_buffer = "host-mla-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.transfer_kv_all_layer_direct_lf_pf",
side_effect=AssertionError(
"MLA page_first_direct all-layer backup must not use sgl-kernel direct LF->PF"
),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
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-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_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.layer_num = 2
host_pool.kv_buffer = ["host-k-page-first", "host-v-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.transfer_kv_all_layer_direct_lf_pf",
side_effect=AssertionError(
"MHA page_first_direct all-layer backup must not use sgl-kernel direct LF->PF"
),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
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-page-first", "host-v-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)
class TestNSAIndexerPageIndices(CustomTestCase):
def make_host_pool_stub(self, page_size: int):
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
@@ -189,6 +291,55 @@ class TestNSAIndexerPageIndices(CustomTestCase):
torch.tensor([16, 17, 18, 19], dtype=torch.int64),
)
def test_page_first_direct_all_layer_indexer_backup_uses_tai_per_layer_route(self):
host_pool = self.make_host_pool_stub(page_size=4)
host_pool.layout = "page_first_direct"
host_pool.indexer_page_stride_size = 8
host_pool.layer_num = 3
host_pool.index_k_with_scale_buffer = "host-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.transfer_kv_all_layer_direct_lf_pf",
side_effect=AssertionError(
"NSA page_first_direct indexer all-layer backup must not use sgl-kernel direct LF->PF"
),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf",
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-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()