Keep CP compute padding out of sparse MoE

CP shared-KV compute padding creates per-request lane slots, so valid rows are not a simple prefix/suffix mask. DeepEP MoE was still seeing dummy rows and using scalar non-padded semantics, which let padding participate in gate/topk and corrupted cache-hit tiny-extend inference.\n\nThe fix compacts CP-local valid rows before MoE dispatch and restores the compact output back to the compute-padded row layout before downstream layer communication. The local GSM8K investigation ledger is now removed from the tracked tree and ignored so future debug notes stay local.\n\nConstraint: CP shared-KV compute-padding layout must keep downstream communicator shapes stable.\nRejected: Disable bs>1/current reuse/cache-hit fast paths | hides the semantic bug and loses the intended performance path.\nRejected: Use num_token_non_padded for MoE under compute padding | valid rows are interleaved with dummy lane slots, not suffix-padded.\nConfidence: high\nScope-risk: moderate\nDirective: Do not feed compute-padded dummy rows into sparse MoE gate/topk; compact valid rows at the MoE boundary and restore shape afterward.\nTested: python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/models/deepseek_v2.py\nTested: remote focused CP utils tests passed, 4 tests.\nTested: remote GSM8K 50-question smoke accuracy 0.960; 200-question runs accuracy 0.955 and 0.965; full 1319-question run accuracy 0.952.\nNot-tested: Long-running production traffic beyond GSM8K after this commit.

.gitignore

@@ -280,3 +280,6 @@ sgl-kernel/csrc/**/*_musa/
 # (kept on disk for local work, never committed)
 docs_internal/
 tai-kernel/
+
+# Local CP shared-KV/GSM8K debug ledger
+docs/advanced_features/nsa_prefill_cp_gsm8k_cachehit_temp_findings.md

python/sglang/srt/layers/attention/nsa/utils.py

@@ -1192,6 +1192,58 @@ def select_cp_local_valid_rows_for_cache_write(
     return selected
 
 
+def restore_cp_local_valid_rows_for_moe(
+    forward_batch,
+    compact_tensor: torch.Tensor,
+    local_compute_reference: torch.Tensor,
+) -> torch.Tensor:
+    """Restore compact CP valid rows to the local compute-padded row layout.
+
+    CP shared-KV compute padding pads each request to a full CP lane set so
+    attention/index kernels keep a page/lane-stable layout.  Sparse MoE must not
+    route those dummy rows: they are not suffix padding and cannot be represented
+    by ``num_token_non_padded``.  The MoE path therefore runs on compact valid
+    rows and scatters the result back to the original local compute layout before
+    the layer communicator sees it.
+    """
+
+    plan = get_cp_shared_kv_batch_plan(forward_batch)
+    if plan is None or not bool(getattr(plan, "compute_padding_enabled", False)):
+        return compact_tensor
+
+    indices, expected_compute_rows = _get_cp_local_valid_row_indices_cache(
+        forward_batch,
+        plan,
+        compact_tensor.device,
+    )
+    local_rows = int(local_compute_reference.shape[0])
+    valid_rows = int(indices.numel())
+    compact_rows = int(compact_tensor.shape[0])
+    if local_rows != expected_compute_rows:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][moe_restore_compute_rows_shape] "
+            "CP shared-KV MoE restore reference must contain local compute rows. "
+            f"local_rows={local_rows} expected_compute_rows={expected_compute_rows} "
+            f"valid_rows={valid_rows}"
+        )
+    if compact_rows != valid_rows:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][moe_restore_valid_rows_shape] "
+            "CP shared-KV MoE compact output must contain local valid rows. "
+            f"compact_rows={compact_rows} valid_rows={valid_rows} "
+            f"local_rows={local_rows}"
+        )
+    if valid_rows == local_rows:
+        return compact_tensor
+    restored = local_compute_reference.new_zeros(
+        (local_rows, *compact_tensor.shape[1:]),
+        dtype=compact_tensor.dtype,
+    )
+    if valid_rows > 0:
+        restored.index_copy_(0, indices, compact_tensor)
+    return restored
+
+
 def select_cp_current_valid_rows_for_reuse(
     forward_batch,
     current_tensor: torch.Tensor,

python/sglang/srt/models/deepseek_v2.py

@@ -60,9 +60,11 @@ from sglang.srt.layers.attention.nsa.utils import (
     cp_collect_last_token_hidden,
     cp_split_and_rebuild_data,
     cp_split_and_rebuild_position,
+    restore_cp_local_valid_rows_for_moe,
     is_nsa_enable_prefill_cp,
     nsa_use_prefill_cp,
     prepare_input_dp_with_cp_dsa,
+    select_cp_local_valid_rows_for_cache_write,
 )
 from sglang.srt.layers.communicator import (
     LayerCommunicator,
@@ -755,6 +757,22 @@ class DeepseekV2MoE(nn.Module):
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
+        local_compute_hidden_states = None
+        if nsa_use_prefill_cp(forward_batch):
+            plan = getattr(
+                getattr(forward_batch, "nsa_cp_metadata", None),
+                "batch_plan",
+                None,
+            )
+            if plan is not None and bool(
+                getattr(plan, "compute_padding_enabled", False)
+            ):
+                local_compute_hidden_states = hidden_states
+                hidden_states = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch,
+                    hidden_states,
+                )
+
         shared_output = None
         sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn
         sbo_overlap_dispatch_flag = (
@@ -959,6 +977,13 @@ class DeepseekV2MoE(nn.Module):
             ):
                 final_hidden_states *= self.routed_scaling_factor
 
+        if local_compute_hidden_states is not None:
+            final_hidden_states = restore_cp_local_valid_rows_for_moe(
+                forward_batch,
+                final_hidden_states,
+                local_compute_hidden_states,
+            )
+
         return final_hidden_states
 
     def _forward_shared_experts(

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -1905,6 +1905,56 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(selected.tolist(), [[50.0, 51.0]])
 
+    def test_restore_cp_local_valid_rows_for_moe_keeps_dummy_rows_zero(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa.utils import (
+            restore_cp_local_valid_rows_for_moe,
+        )
+
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[29, 9, 60, 9, 17],
+            prefix_lens=[704, 704, 640, 704, 704],
+            page_size=64,
+            cp_size=8,
+            cp_rank=0,
+        )
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=5,
+                batch_plan=plan,
+            )
+        )
+        local_compute_rows = torch.zeros(
+            (sum(plan.request_compute_rank_local_tokens), 2), dtype=torch.float32
+        )
+        compact_valid = torch.arange(
+            sum(plan.request_valid_rank_local_tokens) * 2, dtype=torch.float32
+        ).view(-1, 2)
+
+        restored = restore_cp_local_valid_rows_for_moe(
+            forward_batch,
+            compact_valid,
+            local_compute_rows,
+        )
+
+        self.assertEqual(tuple(restored.shape), tuple(local_compute_rows.shape))
+        valid_rows = select_cp_local_valid_rows_for_cache_write(
+            forward_batch,
+            restored,
+        )
+        self.assertTrue(torch.equal(valid_rows, compact_valid))
+
+        valid_mask = torch.zeros(restored.shape[0], dtype=torch.bool)
+        cursor = 0
+        for compute_len, valid_len in zip(
+            plan.request_compute_rank_local_tokens,
+            plan.request_valid_rank_local_tokens,
+        ):
+            valid_mask[cursor : cursor + valid_len] = True
+            cursor += compute_len
+        self.assertTrue(torch.equal(restored[~valid_mask], torch.zeros_like(restored[~valid_mask])))
+
     def test_select_cp_current_valid_rows_accepts_global_rows_under_compute_padding(
         self,
     ):