Harden IndexCache for first e2e: validation, fail-fast, prefetch gating

Pre-e2e audit fixes for NSA index-topk sharing (index_topk_pattern):

- Validate the pattern at init: F/S charset only, must start with F
  (layer 0 has nothing to share from), and length must equal
  num_hidden_layers — a short pattern previously crashed deep in pool
  init with an obscure per-layer ValueError, a long one silently
  ignored tail characters. Warn when a pattern shadows a configured
  index_topk_freq (pattern takes precedence silently otherwise).
- Fail fast when a shared (S) layer receives prev_topk_indices=None:
  both gated indexer call sites previously fell through to omitting
  topk_indices entirely, running sparse attention without indices and
  silently corrupting output if threading were ever dropped.
- Reject pipeline parallelism with index-topk sharing (same rationale
  as the existing TBO guard): topk indices are not threaded across
  PP-stage boundaries, and each stage's loop resets them to None.
- Gate the CP shared-KV index prefetcher on the next layer having an
  index-cache slot: prefetching an S layer's nonexistent buffer
  tripped the inactive-layer fail-fast and disabled the prefetcher
  for the entire run on the first F->S transition.
- Log the resolved active index layer plan at pool init so an e2e run
  can confirm IndexCache is actually on.

Tests: pattern validation cases incl. the GLM-5 reference 78-char
pattern, next_skip == skip-of-next invariant; existing 'C' pattern
test updated to upstream 'F' charset.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/configs/nsa_index_layers.py

@@ -1,8 +1,62 @@
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass
 from typing import Dict
 
+logger = logging.getLogger(__name__)
+
+# Patterns already validated this process, to avoid re-scanning and to emit
+# the freq-shadowing warning only once per distinct pattern.
+_validated_patterns: set = set()
+
+
+def validate_index_topk_config(config) -> None:
+    """Validate IndexCache pattern config against the model shape.
+
+    Raises at init time with an actionable message instead of failing deep in
+    pool setup (short pattern) or silently ignoring tail characters / the
+    whole freq setting (long pattern / pattern+freq both set).
+    """
+
+    pattern = getattr(config, "index_topk_pattern", None)
+    if pattern is None:
+        return
+    memo_key = (pattern, getattr(config, "num_hidden_layers", None))
+    if memo_key in _validated_patterns:
+        return
+
+    invalid_chars = set(pattern) - {"F", "S"}
+    if invalid_chars:
+        raise ValueError(
+            "index_topk_pattern must contain only 'F' (full: layer runs its "
+            "indexer) and 'S' (shared: layer reuses the previous F layer's "
+            f"topk); got invalid characters {sorted(invalid_chars)!r}"
+        )
+    if pattern[0] == "S":
+        raise ValueError(
+            "index_topk_pattern must start with 'F': layer 0 has no previous "
+            "layer to share topk indices from"
+        )
+
+    num_hidden_layers = getattr(config, "num_hidden_layers", None)
+    if num_hidden_layers is not None and len(pattern) != num_hidden_layers:
+        raise ValueError(
+            f"index_topk_pattern length {len(pattern)} does not match "
+            f"num_hidden_layers={num_hidden_layers}; the pattern must have "
+            "exactly one F/S character per hidden layer of this model"
+        )
+
+    index_topk_freq = getattr(config, "index_topk_freq", 1)
+    if index_topk_freq not in (None, 1):
+        logger.warning(
+            "index_topk_pattern overrides index_topk_freq=%s: the pattern "
+            "takes precedence and the freq setting is ignored",
+            index_topk_freq,
+        )
+
+    _validated_patterns.add(memo_key)
+
 
 @dataclass(frozen=True)
 class NSAIndexLayerPlan:
@@ -68,6 +122,7 @@ def nsa_index_skip_flags(
         next_skip_topk = layer_id % index_topk_freq != 0
         return skip_topk, next_skip_topk
 
+    validate_index_topk_config(config)
     if layer_id < 0 or layer_id >= len(index_topk_pattern):
         raise ValueError(
             f"layer_id={layer_id} outside index_topk_pattern "
@@ -94,6 +149,9 @@ def build_nsa_index_layer_plan(
     if end_layer < start_layer:
         raise ValueError(f"end_layer={end_layer} must be >= start_layer={start_layer}")
 
+    if not is_nextn:
+        validate_index_topk_config(config)
+
     if is_nextn:
         active_layer_ids = tuple(range(start_layer, end_layer))
     else:

python/sglang/srt/layers/attention/nsa_backend.py

@@ -179,10 +179,16 @@ def _maybe_start_cp_shared_kv_attention_prefetch(
         forward_batch, "cp_shared_kv_index_prefetcher", None
     )
     if index_prefetcher is not None:
-        index_prefetcher.start_next_layer_prefix(
-            next_layer_id=next_layer_id,
-            token_to_kv_pool=token_to_kv_pool,
-        )
+        # IndexCache shared (S) layers have no index-cache slot; requesting
+        # their buffer would trip the inactive-layer fail-fast and disable the
+        # prefetcher for the whole run. Skip them; F->F transitions keep the
+        # prefetch overlap.
+        index_layer_to_slot = getattr(token_to_kv_pool, "index_layer_to_slot", None)
+        if index_layer_to_slot is None or next_layer_id in index_layer_to_slot:
+            index_prefetcher.start_next_layer_prefix(
+                next_layer_id=next_layer_id,
+                token_to_kv_pool=token_to_kv_pool,
+            )
 
     mla_prefetcher = getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None)
     if mla_prefetcher is not None:

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -511,6 +511,15 @@ class ModelRunnerKVCacheMixin:
                 self.end_layer,
                 is_nextn=self.is_draft_worker,
             )
+            num_local_layers = self.end_layer - self.start_layer
+            if len(index_layer_plan.active_layer_ids) != num_local_layers:
+                logger.info(
+                    "[IndexCache] NSA index-topk sharing enabled: %d/%d layers "
+                    "run the indexer; active_layer_ids=%s",
+                    len(index_layer_plan.active_layer_ids),
+                    num_local_layers,
+                    list(index_layer_plan.active_layer_ids),
+                )
             nsa_pool_kwargs = dict(
                 size=physical_kv_pool_size,
                 page_size=self.page_size,

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -212,6 +212,14 @@ class DeepseekMLAForwardMixin:
                         layer_id=self.layer_id,
                     )
                 else:
+                    if prev_topk_indices is None:
+                        raise RuntimeError(
+                            f"[IndexCache] shared (skip_topk) layer "
+                            f"{self.layer_id} received no prev_topk_indices: "
+                            "topk threading was dropped upstream of this layer "
+                            "(it would otherwise run sparse attention without "
+                            "indices and silently corrupt output)"
+                        )
                     topk_indices = prev_topk_indices
                 current_stream.wait_stream(self.alt_stream)
             else:
@@ -232,6 +240,14 @@ class DeepseekMLAForwardMixin:
                             layer_id=self.layer_id,
                         )
                     else:
+                        if prev_topk_indices is None:
+                            raise RuntimeError(
+                                f"[IndexCache] shared (skip_topk) layer "
+                                f"{self.layer_id} received no prev_topk_indices: "
+                                "topk threading was dropped upstream of this "
+                                "layer (it would otherwise run sparse attention "
+                                "without indices and silently corrupt output)"
+                            )
                         topk_indices = prev_topk_indices
         else:
             q = self.q_proj(hidden_states)[0].view(

python/sglang/srt/server_args.py

@@ -1661,10 +1661,10 @@ class ServerArgs:
                 # with no indices. Reject the combination. (upstream PR #27114)
                 index_topk_freq = getattr(hf_config, "index_topk_freq", 1)
                 index_topk_pattern = getattr(hf_config, "index_topk_pattern", None)
-                if self.enable_two_batch_overlap and (
-                    index_topk_freq > 1
-                    or (index_topk_pattern is not None and "S" in index_topk_pattern)
-                ):
+                index_topk_sharing = index_topk_freq > 1 or (
+                    index_topk_pattern is not None and "S" in index_topk_pattern
+                )
+                if self.enable_two_batch_overlap and index_topk_sharing:
                     raise ValueError(
                         "--enable-two-batch-overlap is not supported with NSA/DSA "
                         "index-topk sharing (index_topk_freq > 1 or an "
@@ -1672,6 +1672,15 @@ class ServerArgs:
                         "path does not propagate topk indices across layers, so "
                         "shared layers would run sparse attention without indices."
                     )
+                if self.pp_size > 1 and index_topk_sharing:
+                    raise ValueError(
+                        "--pipeline-parallel-size > 1 is not supported with "
+                        "NSA/DSA index-topk sharing: topk indices are not "
+                        "threaded across pipeline-stage boundaries "
+                        "(PPProxyTensors), so a stage whose first layer is a "
+                        "shared (S) layer would run sparse attention without "
+                        "indices."
+                    )
 
                 if not is_npu():  # CUDA or ROCm GPU
                     if self.enable_nsa_prefill_context_parallel:

test/registered/unit/configs/test_nsa_index_layers.py

@@ -51,11 +51,64 @@ def test_freq_four_with_offset_one_uses_layers_zero_four_eight():
 
 
 def test_pattern_marks_non_shared_layers_active():
-    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="CSSSCSS")
+    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FSSSFSS")
     plan = build_nsa_index_layer_plan(cfg, 0, len(cfg.index_topk_pattern))
     assert plan.active_layer_ids == (0, 4)
 
 
+def test_pattern_next_skip_matches_skip_of_next_layer():
+    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FFSFSSSF")
+    flags = [nsa_index_skip_flags(cfg, i) for i in range(8)]
+    for i in range(7):
+        assert flags[i][1] == flags[i + 1][0], f"layer {i}"
+    assert flags[0][0] is False
+    assert flags[7][1] is False  # last layer: nothing follows
+
+
+def test_pattern_rejects_invalid_characters():
+    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FCSSF")
+    with pytest.raises(ValueError, match="invalid characters"):
+        build_nsa_index_layer_plan(cfg, 0, 5)
+
+
+def test_pattern_rejects_leading_shared_layer():
+    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="SFFFF")
+    with pytest.raises(ValueError, match="must start with 'F'"):
+        build_nsa_index_layer_plan(cfg, 0, 5)
+
+
+def test_pattern_length_must_match_num_hidden_layers():
+    cfg = SimpleNamespace(
+        index_topk_freq=1, index_topk_pattern="FFSF", num_hidden_layers=6
+    )
+    with pytest.raises(ValueError, match="does not match"):
+        build_nsa_index_layer_plan(cfg, 0, 6)
+
+
+def test_pattern_length_matching_num_hidden_layers_accepted():
+    cfg = SimpleNamespace(
+        index_topk_freq=1, index_topk_pattern="FFSFSS", num_hidden_layers=6
+    )
+    plan = build_nsa_index_layer_plan(cfg, 0, 6)
+    assert plan.active_layer_ids == (0, 1, 3)
+
+
+def test_glm5_reference_pattern_accepted():
+    pattern = (
+        "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"
+    )
+    cfg = SimpleNamespace(
+        index_topk_freq=1,
+        index_topk_pattern=pattern,
+        num_hidden_layers=len(pattern),
+    )
+    plan = build_nsa_index_layer_plan(cfg, 0, len(pattern))
+    assert len(plan.active_layer_ids) == pattern.count("F")
+    for layer_id in range(len(pattern)):
+        skip, _ = nsa_index_skip_flags(cfg, layer_id)
+        assert skip == (pattern[layer_id] == "S")
+
+
 def test_nextn_keeps_all_draft_layers_active_for_state_safety():
     cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
     plan = build_nsa_index_layer_plan(cfg, 0, 1, is_nextn=True)