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

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)
2026-06-07 19:26:50 +00:00
parent 4cf8885a74
commit ce3b5bc874
5 changed files with 162 additions and 20 deletions
--- a/gov/work/2026-06-07-rebase-runtime-env-to-torch-2-11-dev-cu13-image-deepgemm-sgl-deep-gemm-port-compat.toml
+++ b/gov/work/2026-06-07-rebase-runtime-env-to-torch-2-11-dev-cu13-image-deepgemm-sgl-deep-gemm-port-compat.toml
@@ -20,6 +20,31 @@ date = "2026-06-07"
 scope = "deepgemm-port"
 content = "Complete-migration sweep (user: 要迁移就完整迁移). Verified EVERY deep_gemm symbol our code calls against HEAD/0.1.2. UNCHANGED+confirmed: fp8_mqa_logits (ks/ke cu_seqlens, not paged ctx - no 2D issue), fp8_gemm_nt, bf16_gemm_nt/nn, get_mk_alignment_for_contiguous_layout (no-arg), transform_sf_into_required_layout (same kwargs), get/set_num_sms, fp8_m_grouped_gemm_nt_masked (HEAD only ADDS optional recipe_a/b NVFP4 + _ensure_cuda, not required for FP8). PORTED beyond the 4 compat edits: (5) #26839 SBO masked-return unpack guard in moe_runner/deep_gemm.py (our line 388 had the unguarded unpack; our launch enables --enable-single-batch-overlap). (6) #23979 PDL-on-by-default in entrypoint.py (hasattr-guarded set_pdl(True), get_bool_env_var SGLANG_DEEPGEMM_PDL default true - perf default matching new wheel). (7) engine.py sglang-kernel assert 0.4.0->0.4.3 (migrated code requires sgl-deep-gemm-era wheel; (N_total,1) ctx would misbehave on old bundled DeepGEMM). SKIPPED with rationale (not on GLM-5.1-FP8/NSA/deepgemm/H200-SM90 path): #26025 fallback-unsupported-shapes (N/A - our tree has no _varlen_deep_gemm_silu_mul_quant/SGLANG_OPT_USE_JIT_EP_ACTIVATION), #25286 Gemma4 triton_scaled_mm scale layout (triton path not deepgemm), #22300 Minimax fp16/flashinfer-trtllm fallback, #26473/#17392 BF16 features, #26496/#24692 SM120/NVFP4, #26238/#25884 dsv4, AMD/ROCm. All 7 edited files py_compile clean. Runtime validation pending the torch-2.11 image (Phase 2/3)."

+[[content.journal]]
+date = "2026-06-07"
+scope = "env-rebase"
+content = "Phase-2 prep (tai-kernel rebuild characterization). tai-kernel has NO prebuilt .so / no setup.py — it JIT-compiles its CUDA (.cu/.cc, at::Tensor + TORCH_LIBRARY_FRAGMENT(tai_kernel_ops)) via torch.utils.cpp_extension.load at first import, caching to ~/.cache/tai-kernel/ops/<digest>. GOOD: auto-recompiles against the runtime torch (2.11) headers on first import — no manual wheel rebuild. GOTCHA: the cache digest = md5(source bytes + torch.version.cuda + ldflags) does NOT include torch.__version__; torch 2.9.1+cu130 -> 2.11.0+cu130 keeps cuda='13.0' so the digest is unchanged -> a persisted ~/.cache/tai-kernel would be WRONGLY REUSED (ABI-incompatible) against 2.11. Mitigation: fresh container has empty ~/.cache (auto-fresh compile); if cache is on a mounted/persistent volume, rm -rf ~/.cache/tai-kernel before first import. Phase-2 risk = whether the .cu/.cc compile cleanly against torch-2.11 ATen/c10 headers (empirical, verify in image). Phase 1 deep_gemm migration committed 32818784d3."
+
+[[content.journal]]
+date = "2026-06-07"
+scope = "env-rebase"
+content = "Phase 2/3 VALIDATED on B300 (g0036/g0037, 8x B300 SXM6 each, driver 590.48, image sglang:dev-cu13 = torch 2.11.0+cu130, sgl-kernel 0.4.3, sgl-deep-gemm 0.1.2, flashinfer 0.6.12, transformers 5.8.1, xgrammar 0.2.1, mooncake-transfer-engine-cuda13 0.3.11.post1, deep_ep 1.2.1, sglang-router 0.3.2 - EXACTLY the predicted target). Synced our branch (Phase-1 port) to shared beegfs /mnt/beegfs/syh/sglang-stable (036<->037 shared). In a privileged --gpus all container: pip install -e python + tai-kernel (--no-deps) OK; sglang 0.5.11 editable shadows image default. SMOKE: (a) deep_gemm 0.1.2 wheel has set_pdl=FALSE and get_compile_mode=FALSE -> BOTH my hasattr guards are load-bearing (would crash unguarded); get_paged_mqa/fp8_paged_mqa/contig/masked present. (b) all migrated modules import clean on torch 2.11 (entrypoint w/ PDL edit, nsa_backend w/ _to_2d, fp8_utils, moe_runner/deep_gemm). (c) DEEPGEMM_SCALE_UE8M0=True on Blackwell -> my transform_scale_ue8m0 DLPack stride fix is ACTIVE/load-bearing here (was dormant on H200 SM90). (d) tai-kernel nsa_prefill CUDA ext JIT-compiled vs torch-2.11 ABI in 46s (cleared ~/.cache/tai-kernel first per the digest-gotcha) - the big ABI risk PASSED. Remaining empirical unknowns: model load (transformers 5.8 + GLM-5.1 trust_remote_code), runtime kernels on SM100, mooncake P->D over RoCE roce_rail0..7 (host /sys GIDs read empty; multi-port roce_rail ports1-6 w/ port1 DOWN - mooncake device/port/GID TBD), deepep normal(P)/low_latency(D) on 1-node B300. Topology change: 1-node P (g0036, tp8/ep8) + 1-node D (g0037, dp8); decode drops nnodes2/dist-init. NICs: roce_rail0..7 (GPU rails), bond1 10.25.203.36/37 control plane."
+
+[[content.journal]]
+date = "2026-06-07"
+scope = "b300-bringup"
+content = "B300 networking diagnosed + RDMA verified. Newly-received cluster; inter-node RDMA appeared dead but was 2 self-inflicted red herrings: (1) per-rail VRFs (CX8_PLB-config.yaml tables 100-107) - main routing table sends 10.33.x out bond1, so naive ping=100% loss; VRF-bound ping (ip vrf exec vrf_vf_rail0 / ping -I eth_vf_rail0) = 0.15ms. (2) container has no iproute2 (ip absent) - irrelevant. Fabric = B300+CX8 dual-plane rail-optimized Spectrum-X RoCE; per-GPU rails roce_vf_rail0..7 RoCE v2 GID idx 3, 10.33.16-23.x/24, 400G, PXB 1:1 GPU (nvidia-smi topo NIC10-17). IB fabric mlx5_12-15 (Quantum-3) also present but rails are the KV path. VERIFIED RDMA 036<->037 in container: ib_write_bw -d roce_vf_rail0 -x 3, handshake to BOND 10.25.203.37, RDMA over rail -> peak 344.93 / AVG 4.53 Gb/s (sustained collapses -> PFC/ECN tuning open issue). Fixed draft scripts (roce_rail->roce_vf_rail0..7). Deliverables on /mnt/beegfs/syh/harness/: run_prefill_b300.sh, run_decode_b300.sh, run_router_b300.sh, b300_rdma_test.sh (host-runnable server/client), b300-bringup-networking.md. Doc also docs_internal/. Memory: reference_b300_cluster. Dispatched 2 opus bg agents: (A) verify prefill per-layer kernel sources on SM100, (B) survey upstream Blackwell/B300/GLM5 perf fixes to consider porting."
+
+[[content.journal]]
+date = "2026-06-07"
+scope = "b300-perf"
+content = "Two opus agents done. (A) prefill per-layer kernel-source audit on SM100: NO sm_90 lock-in anywhere; all kernels runtime-selected or JIT-built for live arch. Map: NSA attn=sgl_kernel flash_mla_sparse_fwd (heads 64->128 on Blackwell); indexer=deep_gemm fp8_mqa/paged_mqa/metadata (0.1.2); MoE=deep_gemm m_grouped_fp8_gemm_nt_contiguous + TAI three_step; CP shared-KV=tai_kernel CUDA; FP8 KV=Triton; dense block-fp8=deep_gemm fp8_gemm_nt + sgl_kernel ue8m0 quant. Validated the SECOND tai extension (quantization/three_step, the no-torch-fallback MoE hot path) JIT-compiles on cu130/SM100 in 35s (nsa_prefill was 46s). Notable non-correctness: --enable-single-batch-overlap is INERT on this config (Blackwell+deep_gemm+deepep-normal disables masked-down overlap; prefill uses contiguous); BS_GT1_DEBUG != DEBUG_CP_SHARED_KV (TAI fast paths stay active). (B) perf-fix survey -> user decisions: TOOK #25299 (cache MQA-logits budget, drop per-call mem_get_info host sync) + #22232 (data_ptr-guarded copy_ instead of rope clone; q_rope/k_rope are split views of query/key so in-place RoPE -> guard no-op). Verified provenance + correctness first-principles, import-validated on B300, committed 6bc23945e4. SKIPPED per user: #21332 (un-force MHA one-shot on Blackwell - risky), #23856 (torch.mm indexer GEMM). EVALUATE later: #27320 (kv-index parallelize), #26424 (EAGLE topk=1 skip cat/topk/sort), #25859 (MQA mem-ratio knob, pairs w/ #25299)."
+
+[[content.journal]]
+date = "2026-06-07"
+scope = "indexcache"
+content = "IndexCache (NSA indexer topk reuse across layers, user's 'every-4-layers' feature) = upstream #21405 (5a35316417 'Enable IndexCache for DeepSeek V3.2', 2026-04-05) + #27114 (393d0e169e gate bugfix + index_skip_topk_offset). NOT in our tree (no skip_topk/index_topk_freq plumbing). GLM-5.1-FP8 config.json has index_topk/index_n_heads/index_head_dim but NO index_topk_freq -> defaults to 1 (no sharing); user will set index_topk_freq=4 later (Zhipu confirmed accuracy OK, so no accuracy-validation needed on our side). Mechanism: thread prev_topk_indices through Model.forward loop -> DecoderLayer -> AttentionMLA forward/forward_prepare/forward_absorb_prepare (gate: 'if not skip_topk or (is_nextn and prev is None): indexer() else reuse') -> forward_absorb_core returns (output, topk_indices). Our forward_mla ALREADY threads topk_indices prepare->core via inner_state (350/363), so carrier exists. #24392 (capture/output infra, 18 files) is ORTHOGONAL - HEAD reuse path doesn't depend on it - SKIPPED. Dispatched opus agent in isolated WORKTREE to port (match HEAD form; find+fix ALL DecoderLayer.forward callers for the new 3-tuple return; py_compile). Will review diff + B300-validate before merging to main tree. Agent B perf-survey MISSED #21405 (in range but not surfaced) - survey not exhaustive."
+
 [[content.acceptance_criteria]]
 text = "Port deep_gemm sgl-deep-gemm 0.1.2 API compat (2D->N_total,1 ctx_lens, compile-mode hasattr guard, positional m_indices, DLPack stride fix, SBO masked unpack guard, PDL default, kernel version assert)"
 status = "pending"
--- 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
@@ -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,
--- a/python/sglang/srt/models/deepseek_nextn.py
+++ b/python/sglang/srt/models/deepseek_nextn.py
@@ -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,
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -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:
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -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(