From 251a48fb0ab9c0126ea5fcae79ed6b5729f16e67 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 31 May 2026 01:26:07 +0800 Subject: [PATCH] Stabilize CP HiCache page-first direct transfers on CUDA 13 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. --- ..._prefill_cp_page_aligned_cache_contract.md | 447 ++++++++++++++++++ .../sglang/srt/managers/cache_controller.py | 106 +++-- .../sglang/srt/mem_cache/memory_pool_host.py | 96 +++- .../managers/test_hicache_controller_cp.py | 104 ++++ .../unit/mem_cache/test_nsa_pool_host_unit.py | 153 +++++- 5 files changed, 843 insertions(+), 63 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index dd25f511d..cb7a4de5d 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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//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` -> `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. diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 81f1e012c..1770ef49c 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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( diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index fbacc88b4..d75c735df 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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: diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 6e7fd205a..31232bf56 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -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)) diff --git a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py index b35798485..08ce9b65d 100644 --- a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py @@ -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()