Avoid allocating indexer state for shared NSA layers

Target-model skip-topk layers reuse the previous active layer's top-k indices and should not run local indexer modules. Centralize the layer-needs-indexer decision, skip constructing indexers on shared target layers, and skip their checkpoint tensors during load while keeping nextn/draft layers conservative for state safety.

Constraint: index skip should reduce GPU memory in both prefill and decode without changing top-k propagation semantics
Constraint: nextn/draft layers report shared top-k behavior but still need local indexer state safety
Rejected: Loader-only filtering | parameters are already allocated during model construction
Rejected: Dummy indexer modules for skipped layers | preserves most of the memory cost this change removes
Confidence: high
Scope-risk: moderate
Directive: Do not reintroduce indexer execution on skip_topk target layers without proving prev_topk propagation and weight residency semantics
Tested: remote g0034 cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/speculative/test_spec_utils.py test/registered/unit/configs/test_nsa_index_layers.py test/registered/unit/models/test_deepseek_index_skip_weight_loading.py -> 19 passed
Tested: remote g0034 cjy-glm5-new py_compile for modified runtime files
Not-tested: full GLM5 model restart memory delta measurement

python/sglang/srt/configs/nsa_index_layers.py

@@ -136,6 +136,25 @@ def nsa_index_skip_flags(
     return skip_topk, next_skip_topk
 
 
+def nsa_indexer_layer_needs_weights(
+    config, layer_id: int, *, is_nextn: bool = False
+) -> bool:
+    """Return whether this logical layer needs local Indexer parameters.
+
+    Target-model skip-topk layers reuse top-k indices produced by the previous
+    active layer. They must not run their own indexer and therefore do not need
+    checkpoint weights or parameter allocation for ``self_attn.indexer``.
+
+    Draft/nextn layers keep their indexer weights for state safety even though
+    ``nsa_index_skip_flags(..., is_nextn=True)`` reports shared top-k semantics.
+    """
+
+    if is_nextn:
+        return True
+    skip_topk, _ = nsa_index_skip_flags(config, layer_id, is_nextn=False)
+    return not skip_topk
+
+
 def build_nsa_index_layer_plan(
     config, start_layer: int, end_layer: int, *, is_nextn: bool = False
 ) -> NSAIndexLayerPlan:

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

@@ -144,14 +144,20 @@ class DeepseekMHAForwardMixin:
                     q = self.q_b_proj(q_lora)[0].view(
                         -1, self.num_local_heads, self.qk_head_dim
                     )
-                _ = self.indexer(
-                    x=hidden_states,
-                    q_lora=q_lora,
-                    positions=positions,
-                    forward_batch=forward_batch,
-                    layer_id=self.layer_id,
-                    return_indices=False,
-                )
+                if not self.skip_topk or self.is_nextn:
+                    if self.indexer is None:
+                        raise RuntimeError(
+                            f"[IndexCache] layer {self.layer_id} needs to run "
+                            "the indexer but no indexer module was constructed"
+                        )
+                    _ = self.indexer(
+                        x=hidden_states,
+                        q_lora=q_lora,
+                        positions=positions,
+                        forward_batch=forward_batch,
+                        layer_id=self.layer_id,
+                        return_indices=False,
+                    )
             elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
                 # MXFP4: fused RMSNorm + quant
                 q, _, _, _ = fused_rms_mxfp4_quant(

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

@@ -204,6 +204,11 @@ class DeepseekMLAForwardMixin:
                 # the sole intentional fallback (the nextn layer has its own
                 # weights).
                 if not self.skip_topk or (self.is_nextn and prev_topk_indices is None):
+                    if self.indexer is None:
+                        raise RuntimeError(
+                            f"[IndexCache] layer {self.layer_id} needs to run "
+                            "the indexer but no indexer module was constructed"
+                        )
                     topk_indices = self.indexer(
                         x=hidden_states,
                         q_lora=q_lora,
@@ -232,6 +237,12 @@ class DeepseekMLAForwardMixin:
                     if not self.skip_topk or (
                         self.is_nextn and prev_topk_indices is None
                     ):
+                        if self.indexer is None:
+                            raise RuntimeError(
+                                f"[IndexCache] layer {self.layer_id} needs to "
+                                "run the indexer but no indexer module was "
+                                "constructed"
+                            )
                         topk_indices = self.indexer(
                             x=hidden_states,
                             q_lora=q_lora,

python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py

@@ -22,6 +22,7 @@ import torch.nn as nn
 import tqdm
 from transformers import PretrainedConfig
 
+from sglang.srt.configs.nsa_index_layers import nsa_indexer_layer_needs_weights
 from sglang.srt.distributed.parallel_state import GroupCoordinator
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
@@ -156,6 +157,8 @@ class DeepseekV2WeightLoaderMixin:
                     )
                 ):
                     continue
+                if self._should_skip_nsa_indexer_weight(name, is_nextn=is_nextn):
+                    continue
                 if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
                     name = name.replace(
                         "mlp.shared_experts",
@@ -361,6 +364,18 @@ class DeepseekV2WeightLoaderMixin:
 
         self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
 
+    def _should_skip_nsa_indexer_weight(self, name: str, *, is_nextn: bool) -> bool:
+        if is_nextn:
+            return False
+        if ".self_attn.indexer." not in name:
+            return False
+        layer_id = get_layer_id(name)
+        if layer_id is None:
+            return False
+        return not nsa_indexer_layer_needs_weights(
+            self.config, layer_id, is_nextn=False
+        )
+
     def _initialize_nextn_conf(self, is_nextn: bool) -> NextNConfig:
         """
         Initialize the nextn configuration.

python/sglang/srt/models/deepseek_v2.py

@@ -32,7 +32,10 @@ from sglang.srt.batch_overlap.two_batch_overlap import (
     MaybeTboDeepEPDispatcher,
     model_forward_maybe_tbo,
 )
-from sglang.srt.configs.nsa_index_layers import nsa_index_skip_flags
+from sglang.srt.configs.nsa_index_layers import (
+    nsa_index_skip_flags,
+    nsa_indexer_layer_needs_weights,
+)
 from sglang.srt.configs.model_config import (
     compute_mla_mscale_scaling,
     get_nsa_index_head_dim,
@@ -1193,31 +1196,37 @@ class DeepseekV2AttentionMLA(
         self.skip_topk = None
         self.next_skip_topk = None
         if self.use_nsa:
-            is_neox_style = not getattr(config, "indexer_rope_interleave", False)
-            self.indexer = Indexer(
-                hidden_size=hidden_size,
-                index_n_heads=get_nsa_index_n_heads(config),
-                index_head_dim=get_nsa_index_head_dim(config),
-                rope_head_dim=qk_rope_head_dim,
-                index_topk=get_nsa_index_topk(config),
-                q_lora_rank=q_lora_rank,
-                max_position_embeddings=max_position_embeddings,
-                rope_theta=rope_theta,
-                scale_fmt="ue8m0",
-                block_size=128,
-                rope_scaling=rope_scaling,
-                is_neox_style=is_neox_style,
-                prefix=add_prefix("indexer", prefix),
-                quant_config=quant_config,
-                layer_id=layer_id,
-                alt_stream=alt_stream,
-            )
             # Refer: https://arxiv.org/abs/2603.12201 for more details.
             # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices.
             # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices.
             self.skip_topk, self.next_skip_topk = nsa_index_skip_flags(
                 config, layer_id, is_nextn=is_nextn
             )
+            needs_indexer_weights = nsa_indexer_layer_needs_weights(
+                config, layer_id, is_nextn=is_nextn
+            )
+            if needs_indexer_weights:
+                is_neox_style = not getattr(config, "indexer_rope_interleave", False)
+                self.indexer = Indexer(
+                    hidden_size=hidden_size,
+                    index_n_heads=get_nsa_index_n_heads(config),
+                    index_head_dim=get_nsa_index_head_dim(config),
+                    rope_head_dim=qk_rope_head_dim,
+                    index_topk=get_nsa_index_topk(config),
+                    q_lora_rank=q_lora_rank,
+                    max_position_embeddings=max_position_embeddings,
+                    rope_theta=rope_theta,
+                    scale_fmt="ue8m0",
+                    block_size=128,
+                    rope_scaling=rope_scaling,
+                    is_neox_style=is_neox_style,
+                    prefix=add_prefix("indexer", prefix),
+                    quant_config=quant_config,
+                    layer_id=layer_id,
+                    alt_stream=alt_stream,
+                )
+            else:
+                self.indexer = None
 
         self.kv_b_proj = ColumnParallelLinear(
             self.kv_lora_rank,

test/registered/unit/configs/test_nsa_index_layers.py

@@ -4,6 +4,7 @@ import pytest
 
 from sglang.srt.configs.nsa_index_layers import (
     build_nsa_index_layer_plan,
+    nsa_indexer_layer_needs_weights,
     nsa_index_skip_flags,
 )
 
@@ -134,3 +135,18 @@ def test_invalid_offset_fails_before_layer_zero_can_skip_without_prior_topk():
     cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=0)
     with pytest.raises(ValueError, match="index_skip_topk_offset"):
         nsa_index_skip_flags(cfg, 0)
+
+
+def test_indexer_weights_needed_only_for_active_target_layers():
+    cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="FSF")
+
+    assert nsa_indexer_layer_needs_weights(cfg, 0) is True
+    assert nsa_indexer_layer_needs_weights(cfg, 1) is False
+    assert nsa_indexer_layer_needs_weights(cfg, 2) is True
+
+
+def test_indexer_weights_kept_for_nextn_even_when_topk_is_shared():
+    cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1)
+
+    assert nsa_index_skip_flags(cfg, 0, is_nextn=True) == (True, True)
+    assert nsa_indexer_layer_needs_weights(cfg, 0, is_nextn=True) is True

test/registered/unit/models/test_deepseek_index_skip_weight_loading.py

@@ -0,0 +1,35 @@
+from types import SimpleNamespace
+
+from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
+    DeepseekV2WeightLoaderMixin,
+)
+
+
+class TestDeepseekIndexSkipWeightLoading:
+    def _loader(self, pattern="FSF"):
+        loader = object.__new__(DeepseekV2WeightLoaderMixin)
+        loader.config = SimpleNamespace(index_topk_freq=1, index_topk_pattern=pattern)
+        return loader
+
+    def test_skip_layer_indexer_checkpoint_weights_are_ignored(self):
+        loader = self._loader()
+
+        assert not loader._should_skip_nsa_indexer_weight(
+            "model.layers.0.self_attn.indexer.wk.weight", is_nextn=False
+        )
+        assert loader._should_skip_nsa_indexer_weight(
+            "model.layers.1.self_attn.indexer.wk.weight", is_nextn=False
+        )
+        assert not loader._should_skip_nsa_indexer_weight(
+            "model.layers.2.self_attn.indexer.wk.weight", is_nextn=False
+        )
+
+    def test_non_indexer_and_nextn_weights_are_not_skipped(self):
+        loader = self._loader()
+
+        assert not loader._should_skip_nsa_indexer_weight(
+            "model.layers.1.self_attn.q_b_proj.weight", is_nextn=False
+        )
+        assert not loader._should_skip_nsa_indexer_weight(
+            "model.layers.1.self_attn.indexer.wk.weight", is_nextn=True
+        )