Reduce inactive NSA index-cache transfer safely

Centralize the IndexCache skip formula and thread the resulting active logical index layers into NSA KV pools. HiCache now skips only the indexer H2D/D2H payload for inactive target layers while preserving per-layer MLA KV transfer, keeping allocation shape unchanged for this phase.

Constraint: P0-P2 must not compact device or host allocation yet; prefill/decode state transfer still has no logical layer-id metadata.

Rejected: Recompute the skip formula separately in mem_cache | formula drift would corrupt cache or waste transfers when offset/pattern settings change.

Rejected: Skip whole-layer HiCache load/backup | MLA KV remains required for every attention layer.

Confidence: medium

Scope-risk: moderate

Directive: Before enabling compact state buffers or compact allocation, add layer-id metadata validation to PD transfer.

Tested: Local py_compile for touched files; remote pytest in g0034 container: test_nsa_index_layers.py and TestNSAIndexerPageIndices, 20 passed.

Not-tested: ETE replay/GSM8K with --nsa-index-topk-freq 4; PD state-transfer compaction remains unimplemented.

python/sglang/srt/configs/nsa_index_layers.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+
+@dataclass(frozen=True)
+class NSAIndexLayerPlan:
+    """Logical NSA index-cache layers used by IndexCache/top-k sharing."""
+
+    start_layer: int
+    end_layer: int
+    active_layer_ids: tuple[int, ...]
+    layer_to_slot: Dict[int, int]
+
+    def is_active(self, layer_id: int) -> bool:
+        return layer_id in self.layer_to_slot
+
+    def slot_for_layer(self, layer_id: int) -> int:
+        try:
+            return self.layer_to_slot[layer_id]
+        except KeyError as exc:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][index_cache_layer] "
+                f"inactive index layer requested: layer_id={layer_id} "
+                f"active_layer_ids={list(self.active_layer_ids)}"
+            ) from exc
+
+
+def nsa_index_skip_flags(
+    config, layer_id: int, *, is_nextn: bool = False
+) -> tuple[bool, bool]:
+    """Return `(skip_topk, next_skip_topk)` for one logical layer.
+
+    This intentionally mirrors the historical DeepseekV2AttentionMLA formula.
+    Keep this helper as the single source of truth for model-forward and cache
+    layer planning.
+    """
+
+    if is_nextn:
+        return True, True
+
+    index_topk_freq = getattr(config, "index_topk_freq", 1)
+    if index_topk_freq is None:
+        index_topk_freq = 1
+    if index_topk_freq < 1:
+        raise ValueError(f"index_topk_freq must be >= 1, got {index_topk_freq}")
+
+    index_topk_pattern = getattr(config, "index_topk_pattern", None)
+    index_skip_topk_offset = getattr(config, "index_skip_topk_offset", None)
+
+    if index_topk_pattern is None and index_skip_topk_offset is not None:
+        if index_skip_topk_offset <= 0:
+            raise ValueError(
+                "index_skip_topk_offset must be positive when configured; "
+                f"got {index_skip_topk_offset}"
+            )
+        skip_topk = (
+            max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0
+        )
+        next_skip_topk = (
+            max(layer_id - index_skip_topk_offset + 2, 0) % index_topk_freq != 0
+        )
+        return skip_topk, next_skip_topk
+
+    if index_topk_pattern is None:
+        skip_topk = max(layer_id - 1, 0) % index_topk_freq != 0
+        next_skip_topk = layer_id % index_topk_freq != 0
+        return skip_topk, next_skip_topk
+
+    if layer_id < 0 or layer_id >= len(index_topk_pattern):
+        raise ValueError(
+            f"layer_id={layer_id} outside index_topk_pattern "
+            f"length={len(index_topk_pattern)}"
+        )
+    skip_topk = index_topk_pattern[layer_id] == "S"
+    next_skip_topk = (
+        layer_id < len(index_topk_pattern) - 1
+        and index_topk_pattern[layer_id + 1] == "S"
+    )
+    return skip_topk, next_skip_topk
+
+
+def build_nsa_index_layer_plan(
+    config, start_layer: int, end_layer: int, *, is_nextn: bool = False
+) -> NSAIndexLayerPlan:
+    """Build logical-layer to active-index-slot metadata.
+
+    `end_layer` is exclusive, matching model-runner layer ranges.
+    Draft/nextn pools intentionally keep all local layers active for state
+    safety; top-k skip inside the draft forward is a separate model behavior.
+    """
+
+    if end_layer < start_layer:
+        raise ValueError(f"end_layer={end_layer} must be >= start_layer={start_layer}")
+
+    if is_nextn:
+        active_layer_ids = tuple(range(start_layer, end_layer))
+    else:
+        active_layer_ids = tuple(
+            layer_id
+            for layer_id in range(start_layer, end_layer)
+            if not nsa_index_skip_flags(config, layer_id, is_nextn=False)[0]
+        )
+
+    return NSAIndexLayerPlan(
+        start_layer=start_layer,
+        end_layer=end_layer,
+        active_layer_ids=active_layer_ids,
+        layer_to_slot={
+            layer_id: slot for slot, layer_id in enumerate(active_layer_ids)
+        },
+    )