chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat(nsa): port IndexCache (indexer topk sharing across layers)

Browse Source 
Port upstream IndexCache for DeepSeek-V3.2 / GLM-5 (#21405 + #27114 gate fix) to
our diverged tree, matching upstream HEAD's merged form. The NSA indexer topk is
computed every `index_topk_freq` layers; `skip_topk` layers reuse the previous
layer's topk via prev_topk_indices threaded through the model forward. Default
index_topk_freq=1 => no sharing => zero behavior change until the model config opts in.

- forward_mla.py: gate the two indexer call sites
  (`if not skip_topk or (is_nextn and prev_topk_indices is None)`), else reuse
  prev_topk_indices; forward_absorb_core returns (output, topk_indices) when
  next_skip_topk is set (topk_indices already threaded prepare->core via inner_state).
- deepseek_v2.py: AttentionMLA.__init__ skip_topk/next_skip_topk setup
  (freq/pattern/offset; is_nextn=True/True); thread prev_topk_indices through
  forward/forward_prepare/op_core; DecoderLayer returns a 3-tuple (tuple-unpack
  placed AFTER the CP shared-KV finally); Model.forward loop threads topk_indices.
- deepseek_nextn.py: 3-tuple decoder unpack.
- server_args.py: port the #27114 guard - raise on --enable-two-batch-overlap with
  index-topk sharing (the TBO op path does not propagate topk across layers, so
  shared layers would run sparse attention with no indices).

All DecoderLayer.forward callers covered (Model loop, nextn, glm4_moe_lite /
mistral_large_3_eagle inherit, TBO op-path discards via op_core). Did NOT port
#24392 (orthogonal indexer-topk capture/output infra). Import-validated on the
B300 / torch-2.11 image. Enabling it requires setting `index_topk_freq` in the
model config (Zhipu-confirmed for GLM-5.1).

Refs: WI-2026-06-07-001

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit af2a41b8d396033e378b6bdd5f6baebf2a98c301)
This commit is contained in:
孙远航
2026-06-07 19:26:50 +00:00
committed by laoyao0822
parent 4cf8885a74
commit ce3b5bc874
5 changed files with 162 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
View File

@@ -102,6 +102,7 @@ class DeepseekMLAForwardMixin:
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
llama_4_scaling: Optional[torch.Tensor] = None,
prev_topk_indices: Optional[torch.Tensor] = None,
):
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
@@ -193,18 +194,16 @@ class DeepseekMLAForwardMixin:
q = self.q_b_proj(q)[0].view(
-1, self.num_local_heads, self.qk_head_dim
)
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
current_stream.wait_stream(self.alt_stream)
else:
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
if q_lora is not None:
# skip_topk (shared) layers carry no indexer weights in the
# checkpoint, so they must reuse the carried topk and never run
# the indexer. Do NOT widen this to `or prev_topk_indices is
# None` (the original #21405 gate): that recomputes with an
# uninitialized indexer whenever cross-layer propagation is
# unavailable (e.g. the TBO op path drops topk_indices),
# reintroducing the >index_topk garbling. The is_nextn clause is
# 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):
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
@@ -212,6 +211,28 @@ class DeepseekMLAForwardMixin:
forward_batch=forward_batch,
layer_id=self.layer_id,
)
else:
topk_indices = prev_topk_indices
current_stream.wait_stream(self.alt_stream)
else:
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
if q_lora is not None:
# See the skip_topk note above: shared layers have no
# indexer weights, so this gate must not fall back to
# computing when prev_topk_indices is None.
if not self.skip_topk or (
self.is_nextn and prev_topk_indices is None
):
topk_indices = self.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=self.layer_id,
)
else:
topk_indices = prev_topk_indices
else:
q = self.q_proj(hidden_states)[0].view(
-1, self.num_local_heads, self.qk_head_dim
@@ -547,7 +568,14 @@ class DeepseekMLAForwardMixin:
)
output, _ = self.o_proj(attn_bmm_output)
return output
if self.next_skip_topk is None:
return output
# Return topk_indices for the next layer when enabling index cache
if not self.next_skip_topk:
return output, None
else:
return output, topk_indices
def _maybe_write_cp_shared_local_mla_kv(
self: DeepseekV2AttentionMLA,

2
python/sglang/srt/models/deepseek_nextn.py
View File

@@ -465,7 +465,7 @@ class DeepseekModelNextN(nn.Module):
)
residual = None
with get_global_expert_distribution_recorder().disable_this_region():
hidden_states, residual = self.decoder(
hidden_states, residual, topk_indices = self.decoder(
positions,
hidden_states,
forward_batch,

82
python/sglang/srt/models/deepseek_v2.py
View File

@@ -1093,9 +1093,11 @@ class DeepseekV2AttentionMLA(
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
skip_rope: bool = False,
is_nextn: bool = False,
) -> None:
super().__init__()
self.layer_id = layer_id
self.is_nextn = is_nextn
self.hidden_size = hidden_size
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
@@ -1162,6 +1164,8 @@ class DeepseekV2AttentionMLA(
prefix=add_prefix("kv_a_proj_with_mqa", prefix),
)
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(
@@ -1182,6 +1186,47 @@ class DeepseekV2AttentionMLA(
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.
if is_nextn:
self.skip_topk = True
self.next_skip_topk = True
else:
self.index_topk_freq = getattr(config, "index_topk_freq", 1)
self.index_topk_pattern = getattr(config, "index_topk_pattern", None)
self.index_skip_topk_offset = getattr(
config, "index_skip_topk_offset", None
)
if (
self.index_topk_pattern is None
and self.index_skip_topk_offset is not None
):
assert self.index_skip_topk_offset > 0, (
"index_skip_topk_offset must be positive; offset <= 0 "
"marks layer 0 as skip_topk with no prior topk to reuse"
)
self.skip_topk = (
max(layer_id - self.index_skip_topk_offset + 1, 0)
% self.index_topk_freq
!= 0
)
self.next_skip_topk = (
max(layer_id - self.index_skip_topk_offset + 2, 0)
% self.index_topk_freq
!= 0
)
elif self.index_topk_pattern is None:
self.skip_topk = max(layer_id - 1, 0) % self.index_topk_freq != 0
self.next_skip_topk = layer_id % self.index_topk_freq != 0
else:
self.skip_topk = self.index_topk_pattern[layer_id] == "S"
if layer_id < len(self.index_topk_pattern) - 1:
self.next_skip_topk = (
self.index_topk_pattern[layer_id + 1] == "S"
)
else:
self.next_skip_topk = False
self.kv_b_proj = ColumnParallelLinear(
self.kv_lora_rank,
@@ -1313,9 +1358,14 @@ class DeepseekV2AttentionMLA(
)
def op_core(self, state):
state.hidden_states_after_attn = self.forward_core(
state.pop("attn_intermediate_state")
)
result = self.forward_core(state.pop("attn_intermediate_state"))
# forward_core may return (hidden_states, topk_indices) for NSA models
# with index cache enabled. In the TBO path, topk_indices is not
# propagated between layers, so we discard it here.
if isinstance(result, tuple):
state.hidden_states_after_attn = result[0]
else:
state.hidden_states_after_attn = result
def forward(
self,
@@ -1325,6 +1375,7 @@ class DeepseekV2AttentionMLA(
zero_allocator: BumpAllocator,
layer_scatter_modes: LayerScatterModes = None,
llama_4_scaling: Optional[torch.Tensor] = None,
prev_topk_indices: Optional[torch.Tensor] = None,
):
s = self.forward_prepare(
positions=positions,
@@ -1333,6 +1384,7 @@ class DeepseekV2AttentionMLA(
zero_allocator=zero_allocator,
layer_scatter_modes=layer_scatter_modes,
llama_4_scaling=llama_4_scaling,
prev_topk_indices=prev_topk_indices,
)
return self.forward_core(s)
@@ -1344,6 +1396,7 @@ class DeepseekV2AttentionMLA(
zero_allocator: BumpAllocator,
layer_scatter_modes: LayerScatterModes = None,
llama_4_scaling: Optional[torch.Tensor] = None,
prev_topk_indices: Optional[torch.Tensor] = None,
):
if self.attn_mha.kv_b_proj is None:
self.attn_mha.kv_b_proj = self.kv_b_proj
@@ -1383,7 +1436,12 @@ class DeepseekV2AttentionMLA(
)
elif attn_forward_method == AttnForwardMethod.MLA:
inner_state = self.forward_absorb_prepare(
positions, hidden_states, forward_batch, zero_allocator, llama_4_scaling
positions,
hidden_states,
forward_batch,
zero_allocator,
llama_4_scaling,
prev_topk_indices,
)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM:
inner_state = self.forward_absorb_fused_mla_rope_prepare(
@@ -1543,6 +1601,7 @@ class DeepseekV2DecoderLayer(nn.Module):
reduce_results=False,
prefix=add_prefix("self_attn", prefix),
alt_stream=alt_stream,
is_nextn=is_nextn,
)
if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get():
raise ValueError(
@@ -1645,6 +1704,7 @@ class DeepseekV2DecoderLayer(nn.Module):
zero_allocator: BumpAllocator,
gemm_output_zero_allocator: BumpAllocator = None,
llama_4_scaling: Optional[torch.Tensor] = None,
prev_topk_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
quant_format = (
"mxfp4"
@@ -1694,6 +1754,7 @@ class DeepseekV2DecoderLayer(nn.Module):
zero_allocator=zero_allocator,
llama_4_scaling=llama_4_scaling,
layer_scatter_modes=self.layer_scatter_modes,
prev_topk_indices=prev_topk_indices,
)
finally:
if previous_cp_shared_kv_num_model_layers is None:
@@ -1703,6 +1764,13 @@ class DeepseekV2DecoderLayer(nn.Module):
previous_cp_shared_kv_num_model_layers
)
# forward returns (hidden_states, topk_indices) for NSA models with the
# index cache enabled; otherwise a plain hidden_states tensor.
if isinstance(hidden_states, tuple):
hidden_states, topk_indices = hidden_states
else:
topk_indices = None
hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch
)
@@ -1738,7 +1806,7 @@ class DeepseekV2DecoderLayer(nn.Module):
)
self._notify_cp_hicache_layer_end(forward_batch)
return hidden_states, residual
return hidden_states, residual, topk_indices
def op_comm_prepare_attn(
self,
@@ -2018,6 +2086,7 @@ class DeepseekV2Model(nn.Module):
elif self.first_k_dense_replace < normal_start_layer:
normal_end_layer = normal_start_layer = 0
aux_hidden_states = []
topk_indices = None
for i in range(normal_start_layer, normal_end_layer):
# NOTE: torch dynamo does not support graph break in context manager
ctx = (
@@ -2035,7 +2104,7 @@ class DeepseekV2Model(nn.Module):
else:
aux_hidden_states.append(hidden_states + residual)
layer = self.layers[i]
hidden_states, residual = layer(
hidden_states, residual, topk_indices = layer(
positions,
hidden_states,
forward_batch,
@@ -2043,6 +2112,7 @@ class DeepseekV2Model(nn.Module):
zero_allocator,
gemm_output_zero_allocator,
llama_4_scaling,
prev_topk_indices=topk_indices,
)
if normal_end_layer != self.end_layer:

19
python/sglang/srt/server_args.py
View File

@@ -1584,6 +1584,25 @@ class ServerArgs:
self.attention_backend = "nsa"
logger.info("Use nsa attention backend for DeepSeek with DSA.")
# IndexCache (NSA/DSA indexer topk sharing) reuses one layer's topk
# across the following layers, threaded through the monolithic decoder
# forward. The two-batch-overlap op path does NOT propagate topk_indices
# between layers, so shared (skip_topk) layers would run sparse attention
# 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)
):
raise ValueError(
"--enable-two-batch-overlap is not supported with NSA/DSA "
"index-topk sharing (index_topk_freq > 1 or an "
"index_topk_pattern containing shared layers): the TBO op "
"path does not propagate topk indices across layers, so "
"shared layers would run sparse attention without indices."
)
if not is_npu(): # CUDA or ROCm GPU
if self.enable_nsa_prefill_context_parallel:
logger.warning(