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 <omx@oh-my-codex.dev>

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`,

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):

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