diff --git a/python/sglang/srt/configs/nsa_index_layers.py b/python/sglang/srt/configs/nsa_index_layers.py index c485684e8..60036291e 100644 --- a/python/sglang/srt/configs/nsa_index_layers.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index ae34b2f4a..a336a0323 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/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: diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 7f09ddff4..9c5069528 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/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, 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 afbb251eb..a605303ee 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 @@ -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( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index db1a3a433..6e100e0a2 100644 --- a/python/sglang/srt/server_args.py +++ b/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: diff --git a/test/registered/unit/configs/test_nsa_index_layers.py b/test/registered/unit/configs/test_nsa_index_layers.py index fca72d1c9..0602da8f6 100644 --- a/test/registered/unit/configs/test_nsa_index_layers.py +++ b/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)