diff --git a/python/sglang/srt/configs/nsa_index_layers.py b/python/sglang/srt/configs/nsa_index_layers.py index 60036291e..921b8a731 100644 --- a/python/sglang/srt/configs/nsa_index_layers.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 34b839d6c..1a679175a 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/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( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index a605303ee..73614c3b2 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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, diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 12ce382ed..d8d326f5f 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/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. diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 62d5321a0..ca84d2a5d 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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, diff --git a/test/registered/unit/configs/test_nsa_index_layers.py b/test/registered/unit/configs/test_nsa_index_layers.py index 0602da8f6..02b7215af 100644 --- a/test/registered/unit/configs/test_nsa_index_layers.py +++ b/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 diff --git a/test/registered/unit/models/test_deepseek_index_skip_weight_loading.py b/test/registered/unit/models/test_deepseek_index_skip_weight_loading.py new file mode 100644 index 000000000..145848b96 --- /dev/null +++ b/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 + )