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