Enable CP shared-KV compute padding without inflating cache state

Tiny extend requests can leave most CP lanes without query work, which has been tied to hangs and accept-length regressions. This change introduces a dual valid/compute metadata contract: forward paths may materialize compute-padded rows, while cache, current reuse, direct write, HiCache backup, and load remain valid/page based.

The implementation keeps radix/HiCache/device allocation on real page extents, filters dummy compute rows before MLA/index cache writes and current reuse, makes top-k/index consume compute rows while compacting valid rows, and opens tiny CP shared-KV in-seq split through compute padding. The accompanying plan document records the contract and P1-P7 evidence.

Constraint: CP shared KV and HiCache must stay page-granular; dummy compute rows must not allocate, write, backup, or load KV cache.

Constraint: Avoid silent fallback and avoid adding collectives on hot paths.

Rejected: Pad cache allocations to cp_size pages | would waste KV capacity and pollute radix/HiCache state.

Rejected: Keep tiny suffixes out of CP split | preserves the zero-lane behavior that compute padding is meant to remove.

Confidence: medium

Scope-risk: broad

Directive: Do not route compute-padded dummy rows into out_cache_loc, current reuse, HiCache reservation, or backup descriptors; keep valid/cache metadata explicit.

Tested: Remote g0034 container targeted P7 tests: 3 passed, 3 warnings.

Tested: Remote g0034 container full unit slice: PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py => 214 passed, 5 warnings, 2 subtests passed.

Tested: Local py_compile for touched P7 test file.

Not-tested: Latest CUDA/ETE traffic validation for dummy top-k rows, accept len, output len, and detokenizer hang behavior.

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

@@ -67,12 +67,14 @@ from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.nsa.utils import (
     cp_all_gather_rerange_output,
+    get_cp_shared_kv_batch_plan,
     get_cp_shared_kv_local_out_cache_loc,
     get_cp_shared_kv_local_physical_out_cache_loc,
     is_nsa_enable_prefill_cp,
     is_nsa_prefill_cp_in_seq_split,
     nsa_use_prefill_cp,
     raise_cp_shared_kv_direct_write_error,
+    select_cp_local_valid_rows_for_cache_write,
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.layers.communicator import ScatterMode
@@ -358,29 +360,57 @@ class Indexer(MultiPlatformOp):
                     f"logical_page_table_shape={tuple(logical_page_table.shape)} "
                     f"page_size={page_size}"
                 )
-            current_locs = forward_batch.out_cache_loc
-            valid_current_rows = current_extend_kv_rows_for_reuse(
-                forward_batch,
-                current_index_kv[0],
-                current_index_kv[1],
+            batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+            compute_padding_current = batch_plan is not None and bool(
+                getattr(batch_plan, "compute_padding_enabled", False)
             )
-            if valid_current_rows is None:
-                raise RuntimeError(
-                    "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
-                    "CP shared KV index partial-current compose received "
-                    "current_index_kv that does not satisfy current reuse "
-                    "metadata. "
-                    f"cp_rank={layout.cp_rank} layer_id={layer_id} "
-                    f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
-                    f"current_k_shape={tuple(current_index_kv[0].shape)} "
-                    f"current_scale_shape={tuple(current_index_kv[1].shape)} "
-                    f"out_cache_loc_shape={tuple(current_locs.shape)}"
+            if compute_padding_current:
+                current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+                if current_locs is None:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose requires local "
+                        "valid out_cache_loc when compute padding is enabled. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id}"
+                    )
+                valid_current_rows = int(current_locs.numel())
+                if (
+                    int(current_index_kv[0].shape[0]) != valid_current_rows
+                    or int(current_index_kv[1].shape[0]) != valid_current_rows
+                ):
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose received "
+                        "local current rows that do not match local valid locs. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id} "
+                        f"current_k_shape={tuple(current_index_kv[0].shape)} "
+                        f"current_scale_shape={tuple(current_index_kv[1].shape)} "
+                        f"current_locs_shape={tuple(current_locs.shape)}"
+                    )
+            else:
+                current_locs = forward_batch.out_cache_loc
+                valid_current_rows = current_extend_kv_rows_for_reuse(
+                    forward_batch,
+                    current_index_kv[0],
+                    current_index_kv[1],
+                )
+                if valid_current_rows is None:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose received "
+                        "current_index_kv that does not satisfy current reuse "
+                        "metadata. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id} "
+                        f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
+                        f"current_k_shape={tuple(current_index_kv[0].shape)} "
+                        f"current_scale_shape={tuple(current_index_kv[1].shape)} "
+                        f"out_cache_loc_shape={tuple(current_locs.shape)}"
+                    )
+                current_locs = current_locs[:valid_current_rows]
+                current_index_kv = (
+                    current_index_kv[0][:valid_current_rows],
+                    current_index_kv[1][:valid_current_rows],
                 )
-            current_locs = current_locs[:valid_current_rows]
-            current_index_kv = (
-                current_index_kv[0][:valid_current_rows],
-                current_index_kv[1][:valid_current_rows],
-            )
             prefix_slot_span = None
             if len(prefix_lens_cpu) == 1:
                 prefix_pages = int(prefix_lens_cpu[0]) // page_size
@@ -1614,7 +1644,10 @@ class Indexer(MultiPlatformOp):
         current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> torch.Tensor:
         assert forward_batch.nsa_cp_metadata is not None
-        if int(getattr(forward_batch.nsa_cp_metadata, "batch_size", 1) or 1) > 1:
+        if (
+            int(getattr(forward_batch.nsa_cp_metadata, "batch_size", 1) or 1) > 1
+            or get_cp_shared_kv_batch_plan(forward_batch) is not None
+        ):
             return self._get_topk_in_seq_cp_pair_batch(
                 forward_batch,
                 layer_id,
@@ -1648,9 +1681,8 @@ class Indexer(MultiPlatformOp):
         shared_index_buffer = None
         shared_block_tables = None
         current_index_kv_for_topk = current_index_kv
-        if current_index_kv is not None and not is_current_only_extend_batch(
-            forward_batch
-        ):
+        current_only_batch = is_current_only_extend_batch(forward_batch)
+        if current_index_kv is not None and not current_only_batch:
             current_index_kv_for_topk = None
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
@@ -1714,35 +1746,84 @@ class Indexer(MultiPlatformOp):
         cp_metadata = forward_batch.nsa_cp_metadata
         assert cp_metadata is not None
         batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1)
+        batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+        compute_padding_enabled = bool(
+            getattr(cp_metadata, "compute_padding_enabled", False)
+            or bool(getattr(batch_plan, "compute_padding_enabled", False))
+        )
+
+        def metadata_list(name: str, fallback_name: Optional[str] = None) -> List[int]:
+            values = getattr(cp_metadata, name, None)
+            if values is None and batch_plan is not None:
+                values = getattr(batch_plan, name, None)
+            if values is None and fallback_name is not None:
+                values = getattr(cp_metadata, fallback_name, None)
+                if values is None and batch_plan is not None:
+                    values = getattr(batch_plan, fallback_name, None)
+            return list(values or [])
 
         request_kv_len_prev = list(getattr(cp_metadata, "request_kv_len_prev", []) or [])
         request_kv_len_next = list(getattr(cp_metadata, "request_kv_len_next", []) or [])
-        request_actual_seq_q_prev = list(
-            getattr(cp_metadata, "request_actual_seq_q_prev", []) or []
+        if not request_kv_len_prev and batch_plan is not None:
+            request_kv_len_prev = list(
+                getattr(batch_plan, "request_kv_len_prev", []) or []
+            )
+        if not request_kv_len_next and batch_plan is not None:
+            request_kv_len_next = list(
+                getattr(batch_plan, "request_kv_len_next", []) or []
+            )
+        request_actual_seq_q_prev = metadata_list(
+            "request_compute_seq_q_prev"
+            if compute_padding_enabled
+            else "request_actual_seq_q_prev",
+            fallback_name="request_actual_seq_q_prev",
         )
-        request_actual_seq_q_next = list(
-            getattr(cp_metadata, "request_actual_seq_q_next", []) or []
+        request_actual_seq_q_next = metadata_list(
+            "request_compute_seq_q_next"
+            if compute_padding_enabled
+            else "request_actual_seq_q_next",
+            fallback_name="request_actual_seq_q_next",
         )
+        request_valid_seq_q_prev = metadata_list(
+            "request_valid_seq_q_prev",
+            fallback_name="request_valid_actual_seq_q_prev",
+        )
+        request_valid_seq_q_next = metadata_list(
+            "request_valid_seq_q_next",
+            fallback_name="request_valid_actual_seq_q_next",
+        )
+        if not compute_padding_enabled:
+            # Older bs>1 metadata did not have explicit valid-q aliases because
+            # actual q length was also the valid q length. Keep that path
+            # compatible while compute-padding remains fail-fast if valid
+            # lengths are missing.
+            if not request_valid_seq_q_prev:
+                request_valid_seq_q_prev = request_actual_seq_q_prev
+            if not request_valid_seq_q_next:
+                request_valid_seq_q_next = request_actual_seq_q_next
         if not (
             len(request_kv_len_prev) == batch_size
             and len(request_kv_len_next) == batch_size
             and len(request_actual_seq_q_prev) == batch_size
             and len(request_actual_seq_q_next) == batch_size
+            and len(request_valid_seq_q_prev) == batch_size
+            and len(request_valid_seq_q_next) == batch_size
         ):
             raise RuntimeError(
                 "[CP_SHARED_KV_FAIL_FAST][index_topk] "
                 "reason=batch_gt1_index_metadata_incomplete "
                 f"batch_size={batch_size} layer_id={layer_id} "
                 f"kv_prev={request_kv_len_prev} kv_next={request_kv_len_next} "
-                f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next}"
+                f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next} "
+                f"valid_q_prev={request_valid_seq_q_prev} "
+                f"valid_q_next={request_valid_seq_q_next}"
             )
 
         shared_index_buffer = None
         shared_block_tables = None
         current_index_kv_for_topk = current_index_kv
-        if current_index_kv is not None and not is_current_only_extend_batch(
-            forward_batch
-        ):
+        current_only_batch = is_current_only_extend_batch(forward_batch)
+        if current_index_kv is not None and not current_only_batch:
             current_index_kv_for_topk = None
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
@@ -1771,7 +1852,7 @@ class Indexer(MultiPlatformOp):
         compact_output_spans: List[Tuple[int, int]] = []
         current_only = (
             current_index_kv_for_topk is not None
-            and is_current_only_extend_batch(forward_batch)
+            and current_only_batch
         )
         page_table_1 = None if current_only else metadata.get_page_table_1()
 
@@ -1779,16 +1860,27 @@ class Indexer(MultiPlatformOp):
             *,
             req_id: int,
             segment_len: int,
+            valid_segment_len: int,
             kv_len: int,
         ) -> None:
             nonlocal cursor, output_cursor
             segment_len = int(segment_len)
+            valid_segment_len = int(valid_segment_len)
             kv_len = int(kv_len)
             q_segment = q_fp8[cursor : cursor + segment_len]
             weights_segment = weights[cursor : cursor + segment_len]
             cursor += segment_len
             if segment_len == 0:
                 return
+            if valid_segment_len < 0 or valid_segment_len > segment_len:
+                raise RuntimeError(
+                    "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                    "reason=batch_gt1_index_valid_q_len_mismatch "
+                    f"req_id={req_id} segment_len={segment_len} "
+                    f"valid_segment_len={valid_segment_len}"
+                )
+            if valid_segment_len == 0:
+                return
 
             seq_len = int(forward_batch.seq_lens_cpu[req_id].item())
             extend_seq_len = int(forward_batch.extend_seq_lens_cpu[req_id])
@@ -1807,13 +1899,13 @@ class Indexer(MultiPlatformOp):
                     logical_kv_limit = min(seq_len, int(page_table_1.shape[1]))
             valid_q_count = _compute_contiguous_valid_cp_query_count(
                 cp_kv_end=cp_kv_end,
-                actual_seq_q=segment_len,
+                actual_seq_q=valid_segment_len,
                 logical_kv_limit=logical_kv_limit,
             )
             if valid_q_count <= 0:
                 return
 
-            start_abs = cp_kv_end - segment_len
+            start_abs = cp_kv_end - valid_segment_len
             end_abs = start_abs + valid_q_count
             pre_chunk_offset = seq_len - extend_seq_len
             cp_index.append(
@@ -1831,12 +1923,14 @@ class Indexer(MultiPlatformOp):
             collect_segment(
                 req_id=req_id,
                 segment_len=request_actual_seq_q_prev[req_id],
+                valid_segment_len=request_valid_seq_q_prev[req_id],
                 kv_len=request_kv_len_prev[req_id],
             )
             output_cursor += int(request_actual_seq_q_prev[req_id])
             collect_segment(
                 req_id=req_id,
                 segment_len=request_actual_seq_q_next[req_id],
+                valid_segment_len=request_valid_seq_q_next[req_id],
                 kv_len=request_kv_len_next[req_id],
             )
             output_cursor += int(request_actual_seq_q_next[req_id])
@@ -2051,6 +2145,10 @@ class Indexer(MultiPlatformOp):
         local_out_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
         if local_out_loc is None:
             return False
+        local_key = select_cp_local_valid_rows_for_cache_write(
+            forward_batch,
+            local_key,
+        )
         if local_key.shape[0] != local_out_loc.numel():
             raise_cp_shared_kv_direct_write_error(
                 "index_local_shape_mismatch",
@@ -2239,10 +2337,41 @@ class Indexer(MultiPlatformOp):
 
         current_index_kv = None
         if self._can_reuse_current_index_kv(forward_batch):
-            valid_current_rows = current_extend_kv_rows_for_reuse(forward_batch, key)
-            if valid_current_rows is not None and key.shape[0] >= valid_current_rows:
+            batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+            compute_padding_current = batch_plan is not None and bool(
+                getattr(batch_plan, "compute_padding_enabled", False)
+            )
+            if compute_padding_current:
+                current_key = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch,
+                    local_key,
+                )
+                current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+                valid_current_rows = (
+                    int(current_locs.numel()) if current_locs is not None else None
+                )
+            else:
+                valid_current_rows = current_extend_kv_rows_for_reuse(
+                    forward_batch, key
+                )
+                current_key = (
+                    key[:valid_current_rows]
+                    if valid_current_rows is not None
+                    else None
+                )
+                current_locs = (
+                    forward_batch.out_cache_loc[:valid_current_rows]
+                    if valid_current_rows is not None
+                    else None
+                )
+            if (
+                valid_current_rows is not None
+                and current_key is not None
+                and current_locs is not None
+                and current_key.shape[0] == valid_current_rows
+            ):
                 current_k_fp8, current_k_scale = act_quant(
-                    key[:valid_current_rows].contiguous(),
+                    current_key.contiguous(),
                     self.block_size,
                     self.scale_fmt,
                 )
@@ -2258,9 +2387,7 @@ class Indexer(MultiPlatformOp):
                         if forward_batch.cp_shared_kv_layout is not None
                         else None,
                         layer_id,
-                        tensor_debug_summary(
-                            forward_batch.out_cache_loc[:valid_current_rows]
-                        ),
+                        tensor_debug_summary(current_locs),
                         tensor_debug_checksum(current_index_kv[0]),
                         tensor_debug_checksum(current_index_kv[1]),
                     )

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

@@ -320,6 +320,33 @@ class NSAContextParallelMetadata:
     flat_zigzag_index: List[int] = None
     flat_segment_request_ids: List[int] = None
     flat_segment_offsets: List[int] = None
+    compute_padding_enabled: bool = False
+    request_valid_split_lists: List[List[int]] = None
+    request_valid_segment_page_starts: List[List[int]] = None
+    request_valid_segment_page_ends: List[List[int]] = None
+    request_valid_padded_pages: List[int] = None
+    request_valid_padded_tokens: List[int] = None
+    request_valid_padding_tokens: List[int] = None
+    request_valid_rank_local_tokens: List[int] = None
+    request_valid_rank_local_offsets: List[int] = None
+    request_valid_actual_seq_q_prev: List[int] = None
+    request_valid_actual_seq_q_next: List[int] = None
+    request_valid_seq_q_prev: List[int] = None
+    request_valid_seq_q_next: List[int] = None
+    request_valid_query_row_spans: List[List[Tuple[int, int]]] = None
+    request_compute_split_lists: List[List[int]] = None
+    request_compute_segment_page_starts: List[List[int]] = None
+    request_compute_segment_page_ends: List[List[int]] = None
+    request_compute_padded_pages: List[int] = None
+    request_compute_padded_tokens: List[int] = None
+    request_compute_padding_tokens: List[int] = None
+    request_compute_padded_token_offsets: List[int] = None
+    request_compute_rank_local_tokens: List[int] = None
+    request_compute_rank_local_offsets: List[int] = None
+    request_compute_actual_seq_q_prev: List[int] = None
+    request_compute_actual_seq_q_next: List[int] = None
+    request_compute_seq_q_prev: List[int] = None
+    request_compute_seq_q_next: List[int] = None
     batch_plan: object = None
 
 
@@ -355,6 +382,33 @@ class CPSharedKVBatchPlan:
     flat_zigzag_index: List[int]
     flat_segment_request_ids: List[int]
     flat_segment_offsets: List[int]
+    compute_padding_enabled: bool
+    request_valid_split_lists: List[List[int]]
+    request_valid_segment_page_starts: List[List[int]]
+    request_valid_segment_page_ends: List[List[int]]
+    request_valid_padded_pages: List[int]
+    request_valid_padded_tokens: List[int]
+    request_valid_padding_tokens: List[int]
+    request_valid_rank_local_tokens: List[int]
+    request_valid_rank_local_offsets: List[int]
+    request_valid_actual_seq_q_prev: List[int]
+    request_valid_actual_seq_q_next: List[int]
+    request_valid_seq_q_prev: List[int]
+    request_valid_seq_q_next: List[int]
+    request_valid_query_row_spans: List[List[Tuple[int, int]]]
+    request_compute_split_lists: List[List[int]]
+    request_compute_segment_page_starts: List[List[int]]
+    request_compute_segment_page_ends: List[List[int]]
+    request_compute_padded_pages: List[int]
+    request_compute_padded_tokens: List[int]
+    request_compute_padding_tokens: List[int]
+    request_compute_padded_token_offsets: List[int]
+    request_compute_rank_local_tokens: List[int]
+    request_compute_rank_local_offsets: List[int]
+    request_compute_actual_seq_q_prev: List[int]
+    request_compute_actual_seq_q_next: List[int]
+    request_compute_seq_q_prev: List[int]
+    request_compute_seq_q_next: List[int]
 
 
 def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]:
@@ -378,6 +432,54 @@ def _prefix_offsets(lengths: List[int]) -> List[int]:
     return offsets
 
 
+def _build_full_page_unit_split(
+    *,
+    page_units: int,
+    extend_prefix_len: int,
+    page_size: int,
+    cp_size: int,
+) -> Tuple[List[int], List[int], List[int]]:
+    """Split full physical pages across in-seq CP segments.
+
+    This is the compute-side counterpart of
+    `build_page_aligned_in_seq_split_list`: every assigned page contributes a
+    full `page_size` rows because dummy compute rows fill any valid-token tail.
+    """
+
+    if page_units < 0:
+        raise ValueError(f"page_units must be non-negative, got {page_units}")
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+    if cp_size <= 0:
+        raise ValueError(f"cp_size must be positive, got {cp_size}")
+    if extend_prefix_len < 0 or extend_prefix_len % page_size != 0:
+        raise ValueError(
+            "extend_prefix_len must be non-negative and page-aligned, "
+            f"got extend_prefix_len={extend_prefix_len} page_size={page_size}"
+        )
+
+    cp_segment_num = cp_size * 2
+    base_units = page_units // cp_segment_num
+    remainder_units = page_units % cp_segment_num
+    unit_counts = [
+        base_units + (1 if i < remainder_units else 0)
+        for i in range(cp_segment_num)
+    ]
+
+    split_list: List[int] = []
+    segment_page_starts: List[int] = []
+    segment_page_ends: List[int] = []
+    unit_cursor = 0
+    base_page = extend_prefix_len // page_size
+    for unit_count in unit_counts:
+        segment_page_starts.append(base_page + unit_cursor)
+        unit_cursor += unit_count
+        segment_page_ends.append(base_page + unit_cursor)
+        split_list.append(unit_count * page_size)
+
+    return split_list, segment_page_starts, segment_page_ends
+
+
 def build_batch_page_aligned_in_seq_split_plan(
     *,
     extend_lens: List[int],
@@ -445,6 +547,29 @@ def build_batch_page_aligned_in_seq_split_plan(
     request_actual_seq_q_next: List[int] = []
     request_last_token_owner: List[int] = []
     request_last_token_local_offset: List[int] = []
+    request_valid_split_lists: List[List[int]] = []
+    request_valid_segment_page_starts: List[List[int]] = []
+    request_valid_segment_page_ends: List[List[int]] = []
+    request_valid_padded_pages: List[int] = []
+    request_valid_padded_tokens: List[int] = []
+    request_valid_padding_tokens: List[int] = []
+    request_valid_rank_local_tokens: List[int] = []
+    request_valid_actual_seq_q_prev: List[int] = []
+    request_valid_actual_seq_q_next: List[int] = []
+    request_valid_seq_q_prev: List[int] = []
+    request_valid_seq_q_next: List[int] = []
+    request_valid_query_row_spans: List[List[Tuple[int, int]]] = []
+    request_compute_split_lists: List[List[int]] = []
+    request_compute_segment_page_starts: List[List[int]] = []
+    request_compute_segment_page_ends: List[List[int]] = []
+    request_compute_padded_pages: List[int] = []
+    request_compute_padded_tokens: List[int] = []
+    request_compute_padding_tokens: List[int] = []
+    request_compute_rank_local_tokens: List[int] = []
+    request_compute_actual_seq_q_prev: List[int] = []
+    request_compute_actual_seq_q_next: List[int] = []
+    request_compute_seq_q_prev: List[int] = []
+    request_compute_seq_q_next: List[int] = []
     flat_split_list: List[int] = []
     flat_zigzag_index: List[int] = []
     flat_segment_request_ids: List[int] = []
@@ -468,8 +593,17 @@ def build_batch_page_aligned_in_seq_split_plan(
                 f"extend_len={extend_len} page_size={page_size}"
             )
 
+        compute_pages = max(split_info.extend_padded_pages, cp_size)
+        compute_split_list, compute_page_starts, compute_page_ends = (
+            _build_full_page_unit_split(
+                page_units=compute_pages,
+                extend_prefix_len=prefix_len,
+                page_size=page_size,
+                cp_size=cp_size,
+            )
+        )
         owner, local_offset = _get_in_seq_last_token_owner_and_offset(
-            split_list=split_list,
+            split_list=compute_split_list,
             cp_size=cp_size,
             actual_token_count=extend_len,
         )
@@ -479,6 +613,9 @@ def build_batch_page_aligned_in_seq_split_plan(
         rank_local_tokens = (
             split_list[cp_rank] + split_list[mirror_idx]
         )
+        compute_rank_local_tokens = (
+            compute_split_list[cp_rank] + compute_split_list[mirror_idx]
+        )
         split_prefix_list = [0] + prefix_sum_list[:-1]
 
         request_split_infos.append(split_info)
@@ -492,10 +629,38 @@ def build_batch_page_aligned_in_seq_split_plan(
         request_rank_local_tokens.append(rank_local_tokens)
         request_kv_len_prev.append(prefix_sum_list[cp_rank])
         request_kv_len_next.append(prefix_sum_list[mirror_idx])
-        request_actual_seq_q_prev.append(split_list[cp_rank])
-        request_actual_seq_q_next.append(split_list[mirror_idx])
+        request_actual_seq_q_prev.append(compute_split_list[cp_rank])
+        request_actual_seq_q_next.append(compute_split_list[mirror_idx])
         request_last_token_owner.append(owner)
         request_last_token_local_offset.append(local_offset)
+        request_valid_split_lists.append(split_list)
+        request_valid_segment_page_starts.append(split_info.segment_page_starts)
+        request_valid_segment_page_ends.append(split_info.segment_page_ends)
+        request_valid_padded_pages.append(split_info.extend_padded_pages)
+        request_valid_padded_tokens.append(split_info.extend_padded_tokens)
+        request_valid_padding_tokens.append(split_info.extend_padding_tokens)
+        request_valid_rank_local_tokens.append(rank_local_tokens)
+        request_valid_actual_seq_q_prev.append(split_list[cp_rank])
+        request_valid_actual_seq_q_next.append(split_list[mirror_idx])
+        request_valid_seq_q_prev.append(split_list[cp_rank])
+        request_valid_seq_q_next.append(split_list[mirror_idx])
+        request_valid_query_row_spans.append(
+            [
+                (0, split_list[cp_rank]),
+                (compute_split_list[cp_rank], split_list[mirror_idx]),
+            ]
+        )
+        request_compute_split_lists.append(compute_split_list)
+        request_compute_segment_page_starts.append(compute_page_starts)
+        request_compute_segment_page_ends.append(compute_page_ends)
+        request_compute_padded_pages.append(compute_pages)
+        request_compute_padded_tokens.append(compute_pages * page_size)
+        request_compute_padding_tokens.append(compute_pages * page_size - extend_len)
+        request_compute_rank_local_tokens.append(compute_rank_local_tokens)
+        request_compute_actual_seq_q_prev.append(compute_split_list[cp_rank])
+        request_compute_actual_seq_q_next.append(compute_split_list[mirror_idx])
+        request_compute_seq_q_prev.append(compute_split_list[cp_rank])
+        request_compute_seq_q_next.append(compute_split_list[mirror_idx])
         flat_split_list.extend(split_list)
         segment_base = req_id * cp_segment_num
         flat_zigzag_index.extend(segment_base + idx for idx in zigzag_index)
@@ -533,6 +698,42 @@ def build_batch_page_aligned_in_seq_split_plan(
         flat_zigzag_index=flat_zigzag_index,
         flat_segment_request_ids=flat_segment_request_ids,
         flat_segment_offsets=flat_segment_offsets,
+        compute_padding_enabled=any(
+            compute_tokens != valid_tokens
+            for compute_tokens, valid_tokens in zip(
+                request_compute_padded_tokens, request_valid_padded_tokens
+            )
+        ),
+        request_valid_split_lists=request_valid_split_lists,
+        request_valid_segment_page_starts=request_valid_segment_page_starts,
+        request_valid_segment_page_ends=request_valid_segment_page_ends,
+        request_valid_padded_pages=request_valid_padded_pages,
+        request_valid_padded_tokens=request_valid_padded_tokens,
+        request_valid_padding_tokens=request_valid_padding_tokens,
+        request_valid_rank_local_tokens=request_valid_rank_local_tokens,
+        request_valid_rank_local_offsets=_prefix_offsets(request_valid_rank_local_tokens),
+        request_valid_actual_seq_q_prev=request_valid_actual_seq_q_prev,
+        request_valid_actual_seq_q_next=request_valid_actual_seq_q_next,
+        request_valid_seq_q_prev=request_valid_seq_q_prev,
+        request_valid_seq_q_next=request_valid_seq_q_next,
+        request_valid_query_row_spans=request_valid_query_row_spans,
+        request_compute_split_lists=request_compute_split_lists,
+        request_compute_segment_page_starts=request_compute_segment_page_starts,
+        request_compute_segment_page_ends=request_compute_segment_page_ends,
+        request_compute_padded_pages=request_compute_padded_pages,
+        request_compute_padded_tokens=request_compute_padded_tokens,
+        request_compute_padding_tokens=request_compute_padding_tokens,
+        request_compute_padded_token_offsets=_prefix_offsets(
+            request_compute_padded_tokens
+        ),
+        request_compute_rank_local_tokens=request_compute_rank_local_tokens,
+        request_compute_rank_local_offsets=_prefix_offsets(
+            request_compute_rank_local_tokens
+        ),
+        request_compute_actual_seq_q_prev=request_compute_actual_seq_q_prev,
+        request_compute_actual_seq_q_next=request_compute_actual_seq_q_next,
+        request_compute_seq_q_prev=request_compute_seq_q_prev,
+        request_compute_seq_q_next=request_compute_seq_q_next,
     )
 
 
@@ -555,25 +756,43 @@ def split_tensor_by_cp_batch_plan(
     plan,
     *,
     mode: str = "data",
+    split_kind: str = "compute",
 ) -> torch.Tensor:
     """Split a flattened batch tensor by per-request in-seq CP plan.
 
     `mode` is kept explicit for future shape-specific kernels.  The current
     CPU/Python planner path splits along dim0 for 1d, position, and data views.
+    `split_kind="compute"` materializes padded compute rows.  Cache writes must
+    use `split_kind="valid"` so dummy compute rows never receive cache locs.
     """
 
     if mode not in ("1d", "data", "position"):
         raise ValueError(f"unsupported CP batch split mode={mode!r}")
+    if split_kind not in ("compute", "valid"):
+        raise ValueError(f"unsupported CP batch split_kind={split_kind!r}")
 
     request_extend_lens = getattr(plan, "request_extend_lens", None)
-    request_split_lists = getattr(plan, "request_split_lists", None)
+    compute_padding_enabled = bool(getattr(plan, "compute_padding_enabled", False))
+    if split_kind == "valid":
+        request_split_lists = getattr(
+            plan, "request_valid_split_lists", None
+        ) or getattr(plan, "request_split_lists", None)
+        request_target_lens = request_extend_lens
+    elif compute_padding_enabled:
+        request_split_lists = getattr(plan, "request_compute_split_lists", None)
+        request_target_lens = getattr(plan, "request_compute_padded_tokens", None)
+    else:
+        request_split_lists = getattr(plan, "request_split_lists", None)
+        request_target_lens = request_extend_lens
     request_zigzag_indices = getattr(plan, "request_zigzag_indices", None)
     batch_size = int(getattr(plan, "batch_size", 1) or 1)
     if (
         request_extend_lens is None
+        or request_target_lens is None
         or request_split_lists is None
         or request_zigzag_indices is None
         or len(request_extend_lens) != batch_size
+        or len(request_target_lens) != batch_size
         or len(request_split_lists) != batch_size
         or len(request_zigzag_indices) != batch_size
     ):
@@ -591,9 +810,28 @@ def split_tensor_by_cp_batch_plan(
 
     local_chunks = []
     request_tensors = torch.split(tensor, [int(x) for x in request_extend_lens], dim=0)
-    for req_tensor, split_list, zigzag_index in zip(
-        request_tensors, request_split_lists, request_zigzag_indices
+    for req_id, (req_tensor, target_len, split_list, zigzag_index) in enumerate(
+        zip(
+            request_tensors,
+            request_target_lens,
+            request_split_lists,
+            request_zigzag_indices,
+        )
     ):
+        req_tensor = _pad_cp_request_tensor_for_split(
+            req_tensor,
+            target_len=int(target_len),
+            mode=mode,
+            req_id=req_id,
+        )
+        split_total = sum(int(x) for x in split_list)
+        if split_total != int(req_tensor.shape[0]):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_target_len_mismatch] "
+                "request split rows must equal padded request rows. "
+                f"req_id={req_id} split_total={split_total} "
+                f"request_rows={int(req_tensor.shape[0])} mode={mode}"
+            )
         req_segments = list(torch.split(req_tensor, [int(x) for x in split_list], dim=0))
         local_chunks.extend(req_segments[int(index)] for index in zigzag_index)
 
@@ -602,6 +840,150 @@ def split_tensor_by_cp_batch_plan(
     return torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:])
 
 
+def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.device):
+    cached = getattr(forward_batch, "cp_local_valid_row_indices_for_cache_write", None)
+    cached_expected_rows = getattr(
+        forward_batch, "cp_local_valid_compute_rows_for_cache_write", None
+    )
+    if cached is not None and cached_expected_rows is not None:
+        if cached.device == device:
+            return cached, int(cached_expected_rows)
+
+    batch_size = int(getattr(plan, "batch_size", 1) or 1)
+    request_compute_split_lists = getattr(plan, "request_compute_split_lists", None)
+    request_valid_split_lists = getattr(
+        plan, "request_valid_split_lists", None
+    ) or getattr(plan, "request_split_lists", None)
+    request_zigzag_indices = getattr(plan, "request_zigzag_indices", None)
+    if (
+        request_compute_split_lists is None
+        or request_valid_split_lists is None
+        or request_zigzag_indices is None
+        or len(request_compute_split_lists) != batch_size
+        or len(request_valid_split_lists) != batch_size
+        or len(request_zigzag_indices) != batch_size
+    ):
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
+            "CP shared-KV cache writes require valid and compute split metadata "
+            "when compute padding is enabled."
+        )
+
+    chunks = []
+    local_cursor = 0
+    for req_id in range(batch_size):
+        compute_split = request_compute_split_lists[req_id]
+        valid_split = request_valid_split_lists[req_id]
+        if len(compute_split) != len(valid_split):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
+                "valid and compute split metadata disagree. "
+                f"req_id={req_id} compute_segments={len(compute_split)} "
+                f"valid_segments={len(valid_split)}"
+            )
+        for segment_id in request_zigzag_indices[req_id]:
+            segment_id = int(segment_id)
+            compute_len = int(compute_split[segment_id])
+            valid_len = int(valid_split[segment_id])
+            if valid_len > compute_len:
+                raise RuntimeError(
+                    "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
+                    "valid segment is longer than compute segment. "
+                    f"req_id={req_id} segment_id={segment_id} "
+                    f"valid_len={valid_len} compute_len={compute_len}"
+                )
+            if valid_len > 0:
+                chunks.append(
+                    torch.arange(
+                        local_cursor,
+                        local_cursor + valid_len,
+                        device=device,
+                        dtype=torch.long,
+                    )
+                )
+            local_cursor += compute_len
+
+    if chunks:
+        indices = torch.cat(chunks, dim=0)
+    else:
+        indices = torch.empty((0,), device=device, dtype=torch.long)
+    forward_batch.cp_local_valid_row_indices_for_cache_write = indices
+    forward_batch.cp_local_valid_compute_rows_for_cache_write = local_cursor
+    return indices, local_cursor
+
+
+def select_cp_local_valid_rows_for_cache_write(
+    forward_batch,
+    local_tensor: torch.Tensor,
+) -> torch.Tensor:
+    """Drop compute-padding rows before writing CP shared KV into persistent cache."""
+
+    plan = get_cp_shared_kv_batch_plan(forward_batch)
+    if plan is None or not bool(getattr(plan, "compute_padding_enabled", False)):
+        return local_tensor
+
+    indices, expected_compute_rows = _get_cp_local_valid_row_indices_cache(
+        forward_batch,
+        plan,
+        local_tensor.device,
+    )
+    local_rows = int(local_tensor.shape[0])
+    if local_rows != expected_compute_rows:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_shape] "
+            "CP shared-KV cache write tensor must contain local compute rows "
+            "before valid-row selection. "
+            f"local_rows={local_rows} expected_compute_rows={expected_compute_rows} "
+            f"valid_rows={int(indices.numel())}"
+        )
+    if indices.numel() == local_rows:
+        return local_tensor
+    if indices.numel() == 0:
+        return local_tensor.new_empty((0, *local_tensor.shape[1:]))
+    return local_tensor.index_select(0, indices)
+
+
+def _pad_cp_request_tensor_for_split(
+    tensor: torch.Tensor,
+    *,
+    target_len: int,
+    mode: str,
+    req_id: int,
+) -> torch.Tensor:
+    current_len = int(tensor.shape[0])
+    if target_len < current_len:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_len_mismatch] "
+            "target split length is shorter than request valid rows. "
+            f"req_id={req_id} target_len={target_len} current_len={current_len}"
+        )
+    pad_len = target_len - current_len
+    if pad_len == 0:
+        return tensor
+    if current_len <= 0:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_empty_request] "
+            f"cannot compute-pad an empty request. req_id={req_id}"
+        )
+
+    if mode == "position":
+        if tensor.dim() != 1:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_position_padding_rank] "
+                "position compute padding expects a 1D position tensor. "
+                f"req_id={req_id} shape={tuple(tensor.shape)}"
+            )
+        start = tensor[-1] + 1
+        padding = torch.arange(
+            pad_len,
+            device=tensor.device,
+            dtype=tensor.dtype,
+        ) + start
+    else:
+        padding = tensor.new_zeros((pad_len, *tensor.shape[1:]))
+    return torch.cat([tensor, padding], dim=0)
+
+
 def build_flat_page_owner_plan(plan) -> List[int]:
     from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
         build_in_seq_page_compute_owners,
@@ -761,11 +1143,16 @@ def _build_in_seq_split_for_forward_batch(
 
 
 def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
+    communication_split_lists = (
+        plan.request_compute_split_lists
+        if plan.compute_padding_enabled
+        else plan.request_split_lists
+    )
     per_rank_actual_token = []
     for rank in range(plan.cp_size):
         rank_tokens = 0
         mirror = plan.cp_size * 2 - rank - 1
-        for split_list in plan.request_split_lists:
+        for split_list in communication_split_lists:
             rank_tokens += split_list[rank] + split_list[mirror]
         per_rank_actual_token.append(rank_tokens)
     max_rank_token = max(per_rank_actual_token) if per_rank_actual_token else 0
@@ -773,31 +1160,58 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
 
     # Scalar fields remain populated for compatibility, but scalar-only
     # consumers must not use them when batch_size > 1.
-    first_split = plan.request_split_lists[0] if plan.request_split_lists else []
+    first_split = communication_split_lists[0] if communication_split_lists else []
+    first_valid_split = (
+        plan.request_valid_split_lists[0]
+        if plan.request_valid_split_lists
+        else plan.request_split_lists[0]
+        if plan.request_split_lists
+        else []
+    )
     first_info = plan.request_split_infos[0] if plan.request_split_infos else None
     first_zigzag = plan.request_zigzag_indices[0] if plan.request_zigzag_indices else []
-    first_prefix_sum = list(accumulate(first_split))
-    first_kv_len_prev = first_prefix_sum[plan.cp_rank] if first_prefix_sum else 0
+    first_valid_prefix_sum = list(accumulate(first_valid_split))
+    first_kv_len_prev = (
+        first_valid_prefix_sum[plan.cp_rank] if first_valid_prefix_sum else 0
+    )
     first_mirror = plan.cp_size * 2 - plan.cp_rank - 1
-    first_kv_len_next = first_prefix_sum[first_mirror] if first_prefix_sum else 0
+    first_kv_len_next = (
+        first_valid_prefix_sum[first_mirror] if first_valid_prefix_sum else 0
+    )
     first_actual_seq_q_prev = first_split[plan.cp_rank] if first_split else 0
     first_actual_seq_q_next = first_split[first_mirror] if first_split else 0
+    flat_communication_split_list = [
+        token_count
+        for split_list in communication_split_lists
+        for token_count in split_list
+    ]
+    first_reverse_split_len = [
+        element
+        for i in range(plan.cp_size)
+        for element in (first_split[i], first_split[plan.cp_size * 2 - i - 1])
+    ]
+    first_cp_reverse_index = (
+        list(range(0, plan.cp_size * 2, 2))
+        + list(range(plan.cp_size * 2 - 1, 0, -2))
+        if first_split
+        else []
+    )
 
     return NSAContextParallelMetadata(
         split_list=first_split,
         split_list_tensor=torch.tensor(
-            plan.flat_split_list, device="cuda", dtype=torch.int32
+            flat_communication_split_list, device="cuda", dtype=torch.int32
         ),
         split_prefix_tensor=torch.tensor(
-            [0] + list(accumulate(plan.flat_split_list))[:-1],
+            [0] + list(accumulate(flat_communication_split_list))[:-1],
             device="cuda",
             dtype=torch.int32,
         ),
         max_rank_len=max_rank_len,
         zigzag_index=first_zigzag,
         per_rank_actual_token=per_rank_actual_token,
-        reverse_split_len=None,
-        cp_reverse_index=None,
+        reverse_split_len=first_reverse_split_len,
+        cp_reverse_index=first_cp_reverse_index,
         kv_len_prev=first_kv_len_prev,
         kv_len_next=first_kv_len_next,
         actual_seq_q_prev=first_actual_seq_q_prev,
@@ -892,6 +1306,33 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
         flat_zigzag_index=plan.flat_zigzag_index,
         flat_segment_request_ids=plan.flat_segment_request_ids,
         flat_segment_offsets=plan.flat_segment_offsets,
+        compute_padding_enabled=plan.compute_padding_enabled,
+        request_valid_split_lists=plan.request_valid_split_lists,
+        request_valid_segment_page_starts=plan.request_valid_segment_page_starts,
+        request_valid_segment_page_ends=plan.request_valid_segment_page_ends,
+        request_valid_padded_pages=plan.request_valid_padded_pages,
+        request_valid_padded_tokens=plan.request_valid_padded_tokens,
+        request_valid_padding_tokens=plan.request_valid_padding_tokens,
+        request_valid_rank_local_tokens=plan.request_valid_rank_local_tokens,
+        request_valid_rank_local_offsets=plan.request_valid_rank_local_offsets,
+        request_valid_actual_seq_q_prev=plan.request_valid_actual_seq_q_prev,
+        request_valid_actual_seq_q_next=plan.request_valid_actual_seq_q_next,
+        request_valid_seq_q_prev=plan.request_valid_seq_q_prev,
+        request_valid_seq_q_next=plan.request_valid_seq_q_next,
+        request_valid_query_row_spans=plan.request_valid_query_row_spans,
+        request_compute_split_lists=plan.request_compute_split_lists,
+        request_compute_segment_page_starts=plan.request_compute_segment_page_starts,
+        request_compute_segment_page_ends=plan.request_compute_segment_page_ends,
+        request_compute_padded_pages=plan.request_compute_padded_pages,
+        request_compute_padded_tokens=plan.request_compute_padded_tokens,
+        request_compute_padding_tokens=plan.request_compute_padding_tokens,
+        request_compute_padded_token_offsets=plan.request_compute_padded_token_offsets,
+        request_compute_rank_local_tokens=plan.request_compute_rank_local_tokens,
+        request_compute_rank_local_offsets=plan.request_compute_rank_local_offsets,
+        request_compute_actual_seq_q_prev=plan.request_compute_actual_seq_q_prev,
+        request_compute_actual_seq_q_next=plan.request_compute_actual_seq_q_next,
+        request_compute_seq_q_prev=plan.request_compute_seq_q_prev,
+        request_compute_seq_q_next=plan.request_compute_seq_q_next,
         batch_plan=plan,
     )
 
@@ -914,20 +1355,26 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent(
     forward_batch: "ForwardBatch",
     cp_size: int,
 ) -> bool:
-    """Avoid in-seq CP split when a shared-KV suffix has too few pages.
+    """Compatibility hook for the old tiny-suffix skip gate.
 
-    CP shared KV is page-owned.  A cache-hit suffix with fewer physical pages
-    than CP lanes creates mostly-zero in-seq segments; the distributed NSA path
-    has repeatedly shown hangs on that shape.  Keep the page cache contract by
-    running those tiny suffixes without NSA in-seq CP split instead of falling
-    back to token-balanced page-splitting.
+    Compute padding now handles suffixes with fewer physical pages than CP
+    lanes, so this function validates the page-aligned shared-KV contract and
+    no longer requests a skip.
     """
 
-    if (
-        forward_batch is None
-        or not getattr(forward_batch, "uses_cp_shared_kv", False)
-        or cp_size <= 1
-    ):
+    if not _is_cp_shared_kv_forward_batch(forward_batch) or cp_size <= 1:
+        return False
+    _validate_cp_shared_kv_cp_split_plan_inputs(forward_batch, cp_size)
+    return False
+
+
+def _validate_cp_shared_kv_cp_split_plan_inputs(
+    forward_batch: "ForwardBatch",
+    cp_size: int,
+) -> bool:
+    if not _is_cp_shared_kv_forward_batch(forward_batch):
+        return False
+    if cp_size <= 1:
         return False
 
     extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
@@ -936,32 +1383,55 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent(
     page_size = int(getattr(token_to_kv_pool, "page_size", 0) or 0)
     if (
         extend_seq_lens_cpu is None
-        or len(extend_seq_lens_cpu) != 1
         or extend_prefix_lens_cpu is None
-        or len(extend_prefix_lens_cpu) != 1
         or page_size <= 0
     ):
-        return False
-
-    extend_len = int(extend_seq_lens_cpu[0])
-    if extend_len <= 0:
-        return False
-
-    prefix_len = int(extend_prefix_lens_cpu[0])
-    if prefix_len < 0:
-        return False
-    if prefix_len % page_size != 0:
         raise RuntimeError(
-            "[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] "
-            "CP shared KV NSA in-seq split requires a page-aligned prefix. "
-            "The radix/HiCache match path should floor cache hits to the "
-            "previous page boundary before CP split planning. "
-            f"prefix_len={prefix_len} extend_len={extend_len} "
+            "[CP_SHARED_KV_FAIL_FAST][cp_split_missing_page_plan_inputs] "
+            "CP shared KV NSA in-seq split requires extend lengths, prefix "
+            "lengths, and token_to_kv_pool.page_size before planning. "
+            f"extend_seq_lens_cpu={extend_seq_lens_cpu} "
+            f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} "
             f"page_size={page_size} cp_size={cp_size}"
         )
+    if len(extend_seq_lens_cpu) != len(extend_prefix_lens_cpu):
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][cp_split_length_mismatch] "
+            "CP shared KV NSA in-seq split requires one prefix length per "
+            "extend length. "
+            f"extend_seq_lens_cpu={extend_seq_lens_cpu} "
+            f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} "
+            f"cp_size={cp_size}"
+        )
+    if len(extend_seq_lens_cpu) == 0:
+        return False
 
-    padded_pages = ceil_div(extend_len, page_size)
-    return padded_pages < cp_size
+    has_extend = False
+    for req_id, (extend_len_raw, prefix_len_raw) in enumerate(
+        zip(extend_seq_lens_cpu, extend_prefix_lens_cpu)
+    ):
+        extend_len = int(extend_len_raw)
+        prefix_len = int(prefix_len_raw)
+        if extend_len < 0 or prefix_len < 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][cp_split_negative_len] "
+                "CP shared KV NSA in-seq split received a negative length. "
+                f"req_id={req_id} prefix_len={prefix_len} "
+                f"extend_len={extend_len} page_size={page_size} "
+                f"cp_size={cp_size}"
+            )
+        if prefix_len % page_size != 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] "
+                "CP shared KV NSA in-seq split requires a page-aligned prefix. "
+                "The radix/HiCache match path should floor cache hits to the "
+                "previous page boundary before CP split planning. "
+                f"req_id={req_id} prefix_len={prefix_len} "
+                f"extend_len={extend_len} page_size={page_size} "
+                f"cp_size={cp_size}"
+            )
+        has_extend = has_extend or extend_len > 0
+    return has_extend
 
 
 def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
@@ -977,18 +1447,28 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
         # the seq data needs to be divided and recombined at twice the size of cp_size.
         if should_use_replicated_compute_for_short_radix_hit(forward_batch, cp_size):
             return False
-        if should_skip_cp_shared_kv_cp_split_for_short_page_extent(
-            forward_batch, cp_size
-        ):
-            return False
-        cur_cp_seq_len = seq_len // (cp_size * 2)
+        if _is_cp_shared_kv_forward_batch(forward_batch):
+            cur_cp_seq_len = (
+                1
+                if _validate_cp_shared_kv_cp_split_plan_inputs(
+                    forward_batch, cp_size
+                )
+                else 0
+            )
+        else:
+            cur_cp_seq_len = seq_len // (cp_size * 2)
+    extend_token_count = sum(forward_batch.extend_seq_lens_cpu)
+    if _is_cp_shared_kv_forward_batch(forward_batch):
+        min_extend_token_count = 1
+    else:
+        min_extend_token_count = cp_size
     if (
         cur_cp_seq_len != 0
         and cp_size > 1
         and use_nsa
         and forward_batch.forward_mode.is_context_parallel_extend()
         and is_nsa_enable_prefill_cp()
-        and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
+        and extend_token_count >= min_extend_token_count
     ):
         return True
     else:
@@ -1005,7 +1485,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
         return nsa_cp_round_robin_split_data(input_)
 
     metadata = forward_batch.nsa_cp_metadata
-    if getattr(metadata, "batch_size", 1) > 1:
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None:
         return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
 
     input_list = list(
@@ -1027,7 +1507,7 @@ def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor):
         return nsa_cp_round_robin_split_data(input_)
 
     metadata = forward_batch.nsa_cp_metadata
-    if getattr(metadata, "batch_size", 1) > 1:
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None:
         return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
 
     input_list = list(
@@ -1135,7 +1615,7 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
             "nsa_prefill_cp_mode is not in-seq-split",
         )
     batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
-    if batch_plan is not None and int(getattr(batch_plan, "batch_size", 1) or 1) > 1:
+    if batch_plan is not None:
         split_tokens = sum(int(x) for x in getattr(batch_plan, "request_extend_lens", []))
         mismatch_reason = "batch_split_out_cache_len_mismatch"
     else:
@@ -1150,10 +1630,18 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
             out_cache_tokens,
         )
 
-    local_out_cache_loc = cp_split_and_rebuild_1d(
-        forward_batch,
-        out_cache_loc.contiguous(),
-    )
+    if batch_plan is not None:
+        local_out_cache_loc = split_tensor_by_cp_batch_plan(
+            out_cache_loc.contiguous(),
+            batch_plan,
+            mode="1d",
+            split_kind="valid",
+        )
+    else:
+        local_out_cache_loc = cp_split_and_rebuild_1d(
+            forward_batch,
+            out_cache_loc.contiguous(),
+        )
     if local_out_cache_loc.numel() == 0:
         forward_batch.cp_local_out_cache_loc = local_out_cache_loc
         return local_out_cache_loc
@@ -1231,6 +1719,13 @@ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
         )
         return nsa_cp_round_robin_split_data(positions)
 
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None:
+        return split_tensor_by_cp_batch_plan(
+            positions,
+            get_cp_shared_kv_batch_plan(forward_batch),
+            mode="position",
+        )
+
     position_id_list = list(
         torch.split(positions, forward_batch.nsa_cp_metadata.split_list, dim=-1)
     )
@@ -1855,11 +2350,10 @@ def prepare_input_dp_with_cp_dsa(
         and getattr(forward_batch, "uses_cp_shared_kv", False)
         and getattr(forward_batch, "extend_seq_lens_cpu", None) is not None
         and getattr(forward_batch, "extend_prefix_lens_cpu", None) is not None
-        and len(forward_batch.extend_seq_lens_cpu) > 1
     ):
         if page_size is None:
             raise RuntimeError(
-                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_page_size] "
+                "[CP_SHARED_KV_FAIL_FAST][cp_shared_missing_page_size] "
                 "CP shared-KV batch planning requires token_to_kv_pool.page_size"
             )
         batch_plan = build_batch_page_aligned_in_seq_split_plan(
@@ -2098,7 +2592,7 @@ def _in_seq_collect_last_token(
     cp_rank = get_attention_cp_rank()
     bs = len(forward_batch.extend_seq_lens_cpu)
     metadata = getattr(forward_batch, "nsa_cp_metadata", None)
-    if bs > 1:
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None or bs > 1:
         return _in_seq_collect_last_token_batch(hidden_states, metadata, cp_size, cp_rank, bs)
 
     owner = 0
@@ -2134,9 +2628,25 @@ def _in_seq_collect_last_token_batch(
             "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_metadata] "
             "CP in-seq bs>1 last-token collection requires batch metadata."
         )
-    owners = getattr(metadata, "request_last_token_owner", None)
-    local_offsets = getattr(metadata, "request_last_token_local_offset", None)
-    rank_offsets = getattr(metadata, "request_rank_local_offsets", None)
+    plan = getattr(metadata, "batch_plan", None)
+    owners = _get_cp_last_token_metadata_list(
+        metadata, plan, "request_last_token_owner"
+    )
+    local_offsets = _get_cp_last_token_metadata_list(
+        metadata, plan, "request_last_token_local_offset"
+    )
+    compute_padding_enabled = bool(
+        getattr(metadata, "compute_padding_enabled", False)
+        or bool(getattr(plan, "compute_padding_enabled", False))
+    )
+    if compute_padding_enabled:
+        rank_offsets = _get_cp_last_token_metadata_list(
+            metadata, plan, "request_compute_rank_local_offsets"
+        )
+    else:
+        rank_offsets = _get_cp_last_token_metadata_list(
+            metadata, plan, "request_rank_local_offsets"
+        )
     if (
         owners is None
         or local_offsets is None
@@ -2181,3 +2691,16 @@ def _in_seq_collect_last_token_batch(
         dtype=torch.long,
     )
     return gathered.index_select(0, gather_indices)
+
+
+def _get_cp_last_token_metadata_list(
+    metadata: NSAContextParallelMetadata,
+    plan,
+    field_name: str,
+):
+    value = getattr(metadata, field_name, None)
+    if value is not None:
+        return value
+    if plan is not None:
+        return getattr(plan, field_name, None)
+    return None

python/sglang/srt/layers/attention/nsa_backend.py

@@ -51,11 +51,14 @@ from sglang.srt.layers.attention.nsa.transform_index import (
 from sglang.srt.layers.attention.nsa.utils import (
     can_nsa_prefill_cp_round_robin_split,
     compute_nsa_seqlens,
+    get_cp_shared_kv_batch_plan,
+    get_cp_shared_kv_local_out_cache_loc,
     is_nsa_enable_prefill_cp,
     nsa_cp_round_robin_split_data,
     nsa_cp_round_robin_split_q_seqs,
     nsa_use_prefill_cp,
     pad_nsa_cache_seqlens,
+    select_cp_local_valid_rows_for_cache_write,
 )
 from sglang.srt.layers.attention.utils import (
     concat_mla_absorb_q_general,
@@ -1807,12 +1810,41 @@ class NativeSparseAttnBackend(
             mla_prefetcher = getattr(
                 forward_batch, "cp_shared_kv_mla_prefetcher", None
             )
-            current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse(
-                forward_batch,
-                k,
-                k_rope,
+            batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+            compute_padding_current = batch_plan is not None and bool(
+                getattr(batch_plan, "compute_padding_enabled", False)
             )
-            can_reuse_current_kv = current_kv_rows_for_reuse is not None
+            current_locs_for_reuse = None
+            current_k_nope = None
+            current_k_rope = None
+            if compute_padding_current:
+                assert k is not None and k_rope is not None
+                current_k_nope = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch, k
+                )
+                current_k_rope = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch, k_rope
+                )
+                current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc(
+                    forward_batch
+                )
+                current_kv_rows_for_reuse = (
+                    int(current_locs_for_reuse.numel())
+                    if current_locs_for_reuse is not None
+                    else None
+                )
+                can_reuse_current_kv = (
+                    current_kv_rows_for_reuse is not None
+                    and int(current_k_nope.shape[0]) == current_kv_rows_for_reuse
+                    and int(current_k_rope.shape[0]) == current_kv_rows_for_reuse
+                )
+            else:
+                current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse(
+                    forward_batch,
+                    k,
+                    k_rope,
+                )
+                can_reuse_current_kv = current_kv_rows_for_reuse is not None
             if cp_shared_kv_mla_prefetch_log_enabled():
                 if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id):
                     prefix_lens_cpu = getattr(
@@ -1850,8 +1882,15 @@ class NativeSparseAttnBackend(
                 assert k is not None and k_rope is not None
                 assert current_kv_rows_for_reuse is not None
                 valid_current_rows = int(current_kv_rows_for_reuse)
-                current_k_nope = k[:valid_current_rows]
-                current_k_rope = k_rope[:valid_current_rows]
+                if not compute_padding_current:
+                    current_k_nope = k[:valid_current_rows]
+                    current_k_rope = k_rope[:valid_current_rows]
+                    current_locs_for_reuse = forward_batch.out_cache_loc[
+                        :valid_current_rows
+                    ]
+                assert current_k_nope is not None
+                assert current_k_rope is not None
+                assert current_locs_for_reuse is not None
                 if is_packed_fp8_mla_kv_cache(kv_cache):
                     current_kv_cache = pack_current_mla_kv_for_reuse(
                         current_k_nope,
@@ -1860,9 +1899,6 @@ class NativeSparseAttnBackend(
                     )
                 else:
                     current_kv_cache = _cat([current_k_nope, current_k_rope], dim=-1)
-                current_locs_for_reuse = forward_batch.out_cache_loc[
-                    :valid_current_rows
-                ]
                 logical_page_table_1 = page_table_1
                 current_remap_page_size, current_remap_logical_page_capacity = (
                     current_loc_remap_fast_path_args(forward_batch)

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     log_cp_draft_shared_kv_debug,
     nsa_use_prefill_cp,
     raise_cp_shared_kv_direct_write_error,
+    select_cp_local_valid_rows_for_cache_write,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     should_reuse_current_extend_kv,
@@ -557,6 +558,8 @@ class DeepseekMLAForwardMixin:
         local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
         if local_out_cache_loc is None:
             return False
+        k_nope = select_cp_local_valid_rows_for_cache_write(forward_batch, k_nope)
+        k_pe = select_cp_local_valid_rows_for_cache_write(forward_batch, k_pe)
         if (
             k_nope.shape[0] != local_out_cache_loc.numel()
             or k_pe.shape[0] != local_out_cache_loc.numel()