From af20334f4c60ac768e54f7a118d221923aa28c7e Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 10 Jun 2026 05:53:05 +0800 Subject: [PATCH] Size CP HiCache host budget from compact index layers The CP shared target+draft HiCache budget was still estimating NSA target bytes per token as if every target layer had an index cache. After index skip compacts L1 and L2 index buffers to active layers, that stale estimate kept host/L2 token capacity artificially low even though the physical host index allocation was already smaller. Use the pool's index_active_layer_ids when estimating NSA HiCache bytes per token, with the existing full-layer fallback for pools that do not expose compact metadata. This keeps the shared target+draft capacity computation consistent with the actual L2 host layout. Constraint: Index skip compacts NSA index cache by active logical layers, but MLA KV remains full-layer. Rejected: Keep full-layer HiCache estimate | preserves correctness but wastes the L2 capacity recovered by compact index allocation. Confidence: high Scope-risk: narrow Directive: Keep HiCache byte estimates aligned with NSATokenToKVPool allocation metadata; otherwise target+draft token budgets will drift from real host memory use. Tested: Remote container targeted pytest for active-index HiCache estimate and shared-budget tests -> 3 passed. Tested: Remote container P3-P6 regression suite plus new capacity tests -> 203 passed, 2 subtests passed. Not-tested: Full ETE replay/GSM8K after this budget correction. Co-authored-by: OmX --- ...ll_cp_index_skip_cache_communication_plan.md | 1 + python/sglang/srt/mem_cache/hiradix_cache.py | 9 ++++++++- .../unit/mem_cache/test_cp_hicache_metadata.py | 17 +++++++++++++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md b/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md index 48ff09622..8cf48772f 100644 --- a/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_index_skip_cache_communication_plan.md @@ -1135,6 +1135,7 @@ Completion criteria: - Test-only/legacy host stubs without compact metadata keep dense-slot behavior; real compact host pools still fail fast if an inactive logical index layer is requested. - Direct HiCache tests now follow the production CPU-index contract before calling TAI `cudaMemcpyBatchAsync` paths. - Host-register allocation now checks `cudaHostRegister` errors and unregisters host tensors on object finalization; this prevents repeated CUDA unit tests from poisoning later CUDA ops with stale `cudaErrorHostMemoryAlreadyRegistered`. + - Follow-up correction: target+draft shared `--hicache-size` budgeting must also use active index-layer count in `_estimate_hicache_size_per_token`; otherwise the host/L2 buffer is physically compact but the L2 token capacity remains conservatively sized as if all target layers still had index cache. - Verification: remote container P3-P6 suite passed: `test_nsa_pool_host_unit.py`, `test_model_runner_kv_cache_mixin.py`, `test_cp_shared_kv_transfer_mapping.py`, `test_pd_state_layer_ids.py`, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 7c6e30450..0ca4a4cb4 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -75,10 +75,17 @@ def _estimate_hicache_size_per_token(kv_cache) -> int: kv_cache.index_head_dim + kv_cache.index_head_dim // kv_cache.quant_block_size * 4 ) + index_layer_count = len( + getattr( + kv_cache, + "index_active_layer_ids", + range(kv_cache.layer_num), + ) + ) return ( kv_cache.kv_cache_dim * dtype_size * kv_cache.layer_num + indexer_size_per_token - * kv_cache.layer_num + * index_layer_count * NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize ) if isinstance(kv_cache, MLATokenToKVPool): diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 3301f744f..b0ed171ae 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -209,6 +209,23 @@ class TestHiRadixCacheCPDraftHostPool(CustomTestCase): # the paged index buffer: 128 uint8 K bytes + 4 uint8 scale bytes. self.assertEqual(size_per_token, (656 + 128 + 4) * 78) + def test_fp8_nsa_hicache_size_estimate_uses_active_index_layers(self): + pool = object.__new__(NSATokenToKVPool) + pool.store_dtype = torch.uint8 + pool.kv_cache_dim = 656 + pool.layer_num = 78 + pool.index_head_dim = 128 + pool.quant_block_size = 128 + pool.index_active_layer_ids = tuple(range(0, 78, 4)) + + size_per_token = hiradix_cache._estimate_hicache_size_per_token(pool) + + indexer_size_per_token = 128 + 4 + self.assertEqual( + size_per_token, + 656 * 78 + indexer_size_per_token * len(pool.index_active_layer_ids), + ) + def test_fp8_shared_budget_matches_target_and_one_layer_draft_capacity(self): bytes_per_layer = 656 + 128 + 4 target_size_per_token = bytes_per_layer * 78