Preserve CP narrow output while planning real batches

Batch-size support needs request-first CP metadata; treating a batch as one long sequence breaks page ownership, top-k ranges, and phase1 compact output collection. This adds a batch CP plan that records per-request page-aligned splits, rank-local offsets, kv/actual-seq metadata, last-token owners, and flattened descriptors for downstream allocator/runtime workstreams.

The scalar full-rerange path now fail-fasts for batch metadata so bs>1 cannot silently discard the narrow-output optimization or restore hidden states with single-request assumptions.

Constraint: CP shared-KV cache state is page-owned and must preserve request boundaries under bs>1.

Rejected: Let bs>1 fall back to scalar full hidden rerange | it loses the phase1 communication reduction and uses wrong single-request metadata.

Rejected: Add a collective to confirm batch plans | all ranks can derive the same plan from CPU metadata and config.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove batch fail-fast guards until W2/W3 consumers use CPSharedKVBatchPlan end-to-end.

Tested: python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py

Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py -> 39 passed

Not-tested: full ETE bs>1 CP shared-KV runtime; W2/W3 allocator/direct-write consumers are not implemented yet

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

@@ -242,6 +242,73 @@ class NSAContextParallelMetadata:
     extend_padded_pages: int = 0
     extend_padded_tokens: int = 0
     extend_padding_tokens: int = 0
+    batch_size: int = 1
+    request_token_offsets: List[int] = None
+    request_padded_token_offsets: List[int] = None
+    request_page_offsets: List[int] = None
+    request_extend_lens: List[int] = None
+    request_prefix_lens: List[int] = None
+    request_padded_pages: List[int] = None
+    request_padded_tokens: List[int] = None
+    request_padding_tokens: List[int] = None
+    request_split_lists: List[List[int]] = None
+    request_zigzag_indices: List[List[int]] = None
+    request_segment_page_starts: List[List[int]] = None
+    request_segment_page_ends: List[List[int]] = None
+    request_rank_local_tokens: List[int] = None
+    request_rank_local_offsets: List[int] = None
+    request_kv_len_prev: List[int] = None
+    request_kv_len_next: List[int] = None
+    request_actual_seq_q_prev: List[int] = None
+    request_actual_seq_q_next: List[int] = None
+    request_kv_len_prev_tensor: torch.Tensor = None
+    request_kv_len_next_tensor: torch.Tensor = None
+    request_actual_seq_q_prev_tensor: torch.Tensor = None
+    request_actual_seq_q_next_tensor: torch.Tensor = None
+    request_actual_seq_q_prev_cu_tensor: torch.Tensor = None
+    request_actual_seq_q_next_cu_tensor: torch.Tensor = None
+    request_last_token_owner: List[int] = None
+    request_last_token_local_offset: List[int] = None
+    output_collect_mode: str = None
+    flat_split_list: List[int] = None
+    flat_zigzag_index: List[int] = None
+    flat_segment_request_ids: List[int] = None
+    flat_segment_offsets: List[int] = None
+    batch_plan: object = None
+
+
+@dataclass(frozen=True)
+class CPSharedKVBatchPlan:
+    batch_size: int
+    page_size: int
+    cp_size: int
+    cp_rank: int
+    request_token_offsets: List[int]
+    request_padded_token_offsets: List[int]
+    request_page_offsets: List[int]
+    request_extend_lens: List[int]
+    request_prefix_lens: List[int]
+    request_padded_pages: List[int]
+    request_padded_tokens: List[int]
+    request_padding_tokens: List[int]
+    request_split_infos: List[PageAlignedInSeqSplitInfo]
+    request_split_lists: List[List[int]]
+    request_zigzag_indices: List[List[int]]
+    request_segment_page_starts: List[List[int]]
+    request_segment_page_ends: List[List[int]]
+    request_rank_local_tokens: List[int]
+    request_rank_local_offsets: List[int]
+    request_kv_len_prev: List[int]
+    request_kv_len_next: List[int]
+    request_actual_seq_q_prev: List[int]
+    request_actual_seq_q_next: List[int]
+    request_last_token_owner: List[int]
+    request_last_token_local_offset: List[int]
+    output_collect_mode: str
+    flat_split_list: List[int]
+    flat_zigzag_index: List[int]
+    flat_segment_request_ids: List[int]
+    flat_segment_offsets: List[int]
 
 
 def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]:
@@ -256,6 +323,264 @@ def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List
     return [base + (1 if i < remainder else 0) for i in range(cp_segment_num)]
 
 
+def _prefix_offsets(lengths: List[int]) -> List[int]:
+    offsets: List[int] = []
+    cursor = 0
+    for length in lengths:
+        offsets.append(cursor)
+        cursor += int(length)
+    return offsets
+
+
+def build_batch_page_aligned_in_seq_split_plan(
+    *,
+    extend_lens: List[int],
+    prefix_lens: List[int],
+    page_size: int,
+    cp_size: int,
+    cp_rank: int,
+) -> CPSharedKVBatchPlan:
+    """Build per-request page-aligned in-seq CP metadata for a real batch.
+
+    The contract is intentionally request-first: each request is split and
+    page-rounded independently, then flattened for runtime consumers.  This
+    preserves phase1 narrow-output collection for bs>1 without treating the
+    batch as one long sequence.
+    """
+
+    if len(extend_lens) != len(prefix_lens):
+        raise ValueError(
+            "extend_lens and prefix_lens must have the same length, "
+            f"got {len(extend_lens)} and {len(prefix_lens)}"
+        )
+    if cp_size <= 0:
+        raise ValueError(f"cp_size must be positive, got {cp_size}")
+    if cp_rank < 0 or cp_rank >= cp_size:
+        raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}")
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    request_extend_lens = [int(x) for x in extend_lens]
+    request_prefix_lens = [int(x) for x in prefix_lens]
+    for req_id, (extend_len, prefix_len) in enumerate(
+        zip(request_extend_lens, request_prefix_lens)
+    ):
+        if extend_len <= 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_empty_extend] "
+                f"CP shared-KV batch planning requires positive extend lens. "
+                f"req_id={req_id} extend_len={extend_len}"
+            )
+        if prefix_len < 0:
+            raise ValueError(f"prefix_len must be non-negative, got {prefix_len}")
+        if prefix_len % page_size != 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_non_page_aligned_prefix] "
+                "CP shared-KV batch planning requires page-aligned prefixes. "
+                "The radix/HiCache match path should floor cache hits to the "
+                "previous page boundary before planning. "
+                f"req_id={req_id} prefix_len={prefix_len} "
+                f"extend_len={extend_len} page_size={page_size}"
+            )
+
+    cp_segment_num = cp_size * 2
+    request_split_infos: List[PageAlignedInSeqSplitInfo] = []
+    request_split_lists: List[List[int]] = []
+    request_zigzag_indices: List[List[int]] = []
+    request_segment_page_starts: List[List[int]] = []
+    request_segment_page_ends: List[List[int]] = []
+    request_padded_pages: List[int] = []
+    request_padded_tokens: List[int] = []
+    request_padding_tokens: List[int] = []
+    request_rank_local_tokens: List[int] = []
+    request_kv_len_prev: List[int] = []
+    request_kv_len_next: List[int] = []
+    request_actual_seq_q_prev: List[int] = []
+    request_actual_seq_q_next: List[int] = []
+    request_last_token_owner: List[int] = []
+    request_last_token_local_offset: List[int] = []
+    flat_split_list: List[int] = []
+    flat_zigzag_index: List[int] = []
+    flat_segment_request_ids: List[int] = []
+    flat_segment_offsets: List[int] = []
+
+    for req_id, (extend_len, prefix_len) in enumerate(
+        zip(request_extend_lens, request_prefix_lens)
+    ):
+        split_list, split_info = build_page_aligned_in_seq_split_list(
+            total_len=extend_len,
+            extend_len=extend_len,
+            extend_prefix_len=prefix_len,
+            page_size=page_size,
+            cp_size=cp_size,
+        )
+        if not split_info.page_aligned:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_page_split_fallback] "
+                "CP shared-KV batch planning must not use token-balanced "
+                f"fallback. req_id={req_id} prefix_len={prefix_len} "
+                f"extend_len={extend_len} page_size={page_size}"
+            )
+
+        owner, local_offset = _get_in_seq_last_token_owner_and_offset(
+            split_list=split_list,
+            cp_size=cp_size,
+            actual_token_count=extend_len,
+        )
+        zigzag_index = [cp_rank, cp_segment_num - cp_rank - 1]
+        prefix_sum_list = list(accumulate(split_list))
+        mirror_idx = cp_segment_num - cp_rank - 1
+        rank_local_tokens = (
+            split_list[cp_rank] + split_list[mirror_idx]
+        )
+        split_prefix_list = [0] + prefix_sum_list[:-1]
+
+        request_split_infos.append(split_info)
+        request_split_lists.append(split_list)
+        request_zigzag_indices.append(zigzag_index)
+        request_segment_page_starts.append(split_info.segment_page_starts)
+        request_segment_page_ends.append(split_info.segment_page_ends)
+        request_padded_pages.append(split_info.extend_padded_pages)
+        request_padded_tokens.append(split_info.extend_padded_tokens)
+        request_padding_tokens.append(split_info.extend_padding_tokens)
+        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_last_token_owner.append(owner)
+        request_last_token_local_offset.append(local_offset)
+        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)
+        flat_segment_request_ids.extend([req_id] * cp_segment_num)
+        flat_segment_offsets.extend(split_prefix_list)
+
+    return CPSharedKVBatchPlan(
+        batch_size=len(request_extend_lens),
+        page_size=page_size,
+        cp_size=cp_size,
+        cp_rank=cp_rank,
+        request_token_offsets=_prefix_offsets(request_extend_lens),
+        request_padded_token_offsets=_prefix_offsets(request_padded_tokens),
+        request_page_offsets=_prefix_offsets(request_padded_pages),
+        request_extend_lens=request_extend_lens,
+        request_prefix_lens=request_prefix_lens,
+        request_padded_pages=request_padded_pages,
+        request_padded_tokens=request_padded_tokens,
+        request_padding_tokens=request_padding_tokens,
+        request_split_infos=request_split_infos,
+        request_split_lists=request_split_lists,
+        request_zigzag_indices=request_zigzag_indices,
+        request_segment_page_starts=request_segment_page_starts,
+        request_segment_page_ends=request_segment_page_ends,
+        request_rank_local_tokens=request_rank_local_tokens,
+        request_rank_local_offsets=_prefix_offsets(request_rank_local_tokens),
+        request_kv_len_prev=request_kv_len_prev,
+        request_kv_len_next=request_kv_len_next,
+        request_actual_seq_q_prev=request_actual_seq_q_prev,
+        request_actual_seq_q_next=request_actual_seq_q_next,
+        request_last_token_owner=request_last_token_owner,
+        request_last_token_local_offset=request_last_token_local_offset,
+        output_collect_mode="narrow_last_token",
+        flat_split_list=flat_split_list,
+        flat_zigzag_index=flat_zigzag_index,
+        flat_segment_request_ids=flat_segment_request_ids,
+        flat_segment_offsets=flat_segment_offsets,
+    )
+
+
+def get_cp_shared_kv_batch_plan(forward_batch: "ForwardBatch"):
+    metadata = getattr(forward_batch, "nsa_cp_metadata", None)
+    if metadata is None:
+        return None
+    plan = getattr(metadata, "batch_plan", None)
+    if plan is not None:
+        return plan
+    if isinstance(metadata, CPSharedKVBatchPlan):
+        return metadata
+    if getattr(metadata, "batch_size", 1) > 1:
+        return metadata
+    return None
+
+
+def split_tensor_by_cp_batch_plan(
+    tensor: torch.Tensor,
+    plan,
+    *,
+    mode: str = "data",
+) -> 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.
+    """
+
+    if mode not in ("1d", "data", "position"):
+        raise ValueError(f"unsupported CP batch split mode={mode!r}")
+
+    request_extend_lens = getattr(plan, "request_extend_lens", None)
+    request_split_lists = getattr(plan, "request_split_lists", None)
+    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_split_lists is None
+        or request_zigzag_indices is None
+        or len(request_extend_lens) != batch_size
+        or len(request_split_lists) != batch_size
+        or len(request_zigzag_indices) != batch_size
+    ):
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_split_metadata] "
+            "CP shared-KV batch split requires per-request split metadata."
+        )
+
+    expected_tokens = sum(int(x) for x in request_extend_lens)
+    if int(tensor.shape[0]) != expected_tokens:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch] "
+            f"input tokens={int(tensor.shape[0])} expected={expected_tokens}"
+        )
+
+    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
+    ):
+        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)
+
+    if not local_chunks:
+        return tensor.new_empty((0, *tensor.shape[1:]))
+    return torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:])
+
+
+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,
+    )
+
+    owners: List[int] = []
+    for req_id, (extend_len, prefix_len) in enumerate(
+        zip(plan.request_extend_lens, plan.request_prefix_lens)
+    ):
+        request_owners = build_in_seq_page_compute_owners(
+            extend_len=int(extend_len),
+            extend_prefix_len=int(prefix_len),
+            page_size=int(plan.page_size),
+            cp_size=int(plan.cp_size),
+        )
+        if request_owners is None:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_page_owner_plan_unavailable] "
+                f"req_id={req_id} extend_len={extend_len} prefix_len={prefix_len} "
+                f"page_size={plan.page_size} cp_size={plan.cp_size}"
+            )
+        owners.extend(request_owners)
+    return owners
+
+
 def _fallback_page_aligned_split_info(
     *,
     page_size: int,
@@ -389,6 +714,142 @@ def _build_in_seq_split_for_forward_batch(
     )
 
 
+def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
+    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:
+            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
+    max_rank_len = [max_rank_token for _ in range(plan.cp_size)]
+
+    # 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_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_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_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
+
+    return NSAContextParallelMetadata(
+        split_list=first_split,
+        split_list_tensor=torch.tensor(
+            plan.flat_split_list, device="cuda", dtype=torch.int32
+        ),
+        split_prefix_tensor=torch.tensor(
+            [0] + list(accumulate(plan.flat_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,
+        kv_len_prev=first_kv_len_prev,
+        kv_len_next=first_kv_len_next,
+        actual_seq_q_prev=first_actual_seq_q_prev,
+        actual_seq_q_next=first_actual_seq_q_next,
+        kv_len_prev_tensor=torch.tensor(
+            first_kv_len_prev, device="cuda", dtype=torch.int32
+        ),
+        kv_len_next_tensor=torch.tensor(
+            first_kv_len_next, device="cuda", dtype=torch.int32
+        ),
+        actual_seq_q_prev_tensor=torch.tensor(
+            first_actual_seq_q_prev, device="cuda", dtype=torch.int32
+        ),
+        actual_seq_q_next_tensor=torch.tensor(
+            first_actual_seq_q_next, device="cuda", dtype=torch.int32
+        ),
+        actual_seq_q_prev_cu_tensor=torch.tensor(
+            [0, first_actual_seq_q_prev], device="cuda", dtype=torch.int32
+        ),
+        actual_seq_q_next_cu_tensor=torch.tensor(
+            [0, first_actual_seq_q_next], device="cuda", dtype=torch.int32
+        ),
+        total_seq_lens=torch.tensor(max_rank_token * plan.cp_size),
+        page_aligned=True,
+        page_size=plan.page_size,
+        extend_prefix_len=(
+            first_info.extend_prefix_len if first_info is not None else 0
+        ),
+        segment_page_starts=(
+            first_info.segment_page_starts if first_info is not None else []
+        ),
+        segment_page_ends=(
+            first_info.segment_page_ends if first_info is not None else []
+        ),
+        extend_valid_tokens=(
+            first_info.extend_valid_tokens if first_info is not None else 0
+        ),
+        extend_padded_pages=(
+            first_info.extend_padded_pages if first_info is not None else 0
+        ),
+        extend_padded_tokens=(
+            first_info.extend_padded_tokens if first_info is not None else 0
+        ),
+        extend_padding_tokens=(
+            first_info.extend_padding_tokens if first_info is not None else 0
+        ),
+        batch_size=plan.batch_size,
+        request_token_offsets=plan.request_token_offsets,
+        request_padded_token_offsets=plan.request_padded_token_offsets,
+        request_page_offsets=plan.request_page_offsets,
+        request_extend_lens=plan.request_extend_lens,
+        request_prefix_lens=plan.request_prefix_lens,
+        request_padded_pages=plan.request_padded_pages,
+        request_padded_tokens=plan.request_padded_tokens,
+        request_padding_tokens=plan.request_padding_tokens,
+        request_split_lists=plan.request_split_lists,
+        request_zigzag_indices=plan.request_zigzag_indices,
+        request_segment_page_starts=plan.request_segment_page_starts,
+        request_segment_page_ends=plan.request_segment_page_ends,
+        request_rank_local_tokens=plan.request_rank_local_tokens,
+        request_rank_local_offsets=plan.request_rank_local_offsets,
+        request_kv_len_prev=plan.request_kv_len_prev,
+        request_kv_len_next=plan.request_kv_len_next,
+        request_actual_seq_q_prev=plan.request_actual_seq_q_prev,
+        request_actual_seq_q_next=plan.request_actual_seq_q_next,
+        request_kv_len_prev_tensor=torch.tensor(
+            plan.request_kv_len_prev, device="cuda", dtype=torch.int32
+        ),
+        request_kv_len_next_tensor=torch.tensor(
+            plan.request_kv_len_next, device="cuda", dtype=torch.int32
+        ),
+        request_actual_seq_q_prev_tensor=torch.tensor(
+            plan.request_actual_seq_q_prev, device="cuda", dtype=torch.int32
+        ),
+        request_actual_seq_q_next_tensor=torch.tensor(
+            plan.request_actual_seq_q_next, device="cuda", dtype=torch.int32
+        ),
+        request_actual_seq_q_prev_cu_tensor=torch.tensor(
+            [0] + list(accumulate(plan.request_actual_seq_q_prev)),
+            device="cuda",
+            dtype=torch.int32,
+        ),
+        request_actual_seq_q_next_cu_tensor=torch.tensor(
+            [0] + list(accumulate(plan.request_actual_seq_q_next)),
+            device="cuda",
+            dtype=torch.int32,
+        ),
+        request_last_token_owner=plan.request_last_token_owner,
+        request_last_token_local_offset=plan.request_last_token_local_offset,
+        output_collect_mode=plan.output_collect_mode,
+        flat_split_list=plan.flat_split_list,
+        flat_zigzag_index=plan.flat_zigzag_index,
+        flat_segment_request_ids=plan.flat_segment_request_ids,
+        flat_segment_offsets=plan.flat_segment_offsets,
+        batch_plan=plan,
+    )
+
+
 def should_use_replicated_compute_for_short_radix_hit(
     forward_batch: "ForwardBatch",
     cp_size: int,
@@ -495,11 +956,15 @@ 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:
+        return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
+
     input_list = list(
-        torch.split(input_, forward_batch.nsa_cp_metadata.split_list, dim=0)
+        torch.split(input_, metadata.split_list, dim=0)
     )
     result = torch.cat(
-        [input_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index], dim=0
+        [input_list[i] for i in metadata.zigzag_index], dim=0
     ).view(-1, input_.shape[-1])
     return result
 
@@ -512,14 +977,26 @@ def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor):
     if round_robin_split:
         return nsa_cp_round_robin_split_data(input_)
 
+    metadata = forward_batch.nsa_cp_metadata
+    if getattr(metadata, "batch_size", 1) > 1:
+        return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
+
     input_list = list(
-        torch.split(input_, forward_batch.nsa_cp_metadata.split_list, dim=0)
+        torch.split(input_, metadata.split_list, dim=0)
     )
     return torch.cat(
-        [input_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index], dim=0
+        [input_list[i] for i in metadata.zigzag_index], dim=0
     ).view(-1)
 
 
+def _cp_split_and_rebuild_batch_in_seq(forward_batch, input_: torch.Tensor):
+    return split_tensor_by_cp_batch_plan(
+        input_,
+        get_cp_shared_kv_batch_plan(forward_batch),
+        mode="1d" if input_.dim() == 1 else "data",
+    )
+
+
 def get_cp_local_embedding_padded_token_count(forward_batch, local_num_tokens: int):
     metadata = getattr(forward_batch, "nsa_cp_metadata", None)
     max_rank_len = getattr(metadata, "max_rank_len", None)
@@ -1001,6 +1478,14 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
     | token0, token1, token2, token3, token4, token5, token6, token7, ...
     |   +-------------------------+
     """
+    metadata = getattr(forward_batch, "nsa_cp_metadata", None)
+    if getattr(metadata, "batch_size", 1) > 1:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_full_rerange_unsupported] "
+            "CP shared-KV bs>1 must not use scalar full hidden/KV rerange. "
+            "Use batch-aware narrow output collection or add batch-aware full "
+            "rerange metadata/kernels for this consumer."
+        )
     if is_nsa_prefill_cp_round_robin_split():
         with use_symmetric_memory(
             get_attention_cp_group(), disabled=not is_allocation_symmetric()
@@ -1139,7 +1624,28 @@ def prepare_input_dp_with_cp_dsa(
         * Last rank may focus on more tokens (more computation)
     - To mitigate uneven load, the input hissenstate needs to be sliced by cp_size*2 and rearranged.
     """
-    # just support batch = 1
+    if (
+        forward_batch is not None
+        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 batch planning requires token_to_kv_pool.page_size"
+            )
+        batch_plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[int(x) for x in forward_batch.extend_seq_lens_cpu],
+            prefix_lens=[int(x) for x in forward_batch.extend_prefix_lens_cpu],
+            page_size=int(page_size),
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+        )
+        return _build_batch_metadata_from_plan(batch_plan)
+
+    # scalar compatibility path
     kv_len_int = int(kv_len)
     kv_len = torch.tensor(kv_len_int)
     bs_per_cp_group = 1
@@ -1259,20 +1765,44 @@ def _round_robin_collect_last_token(
     forward_batch: "ForwardBatch",
     cp_size: int,
 ) -> torch.Tensor:
-    total_tokens = sum(forward_batch.extend_seq_lens_cpu)
-    owner = (total_tokens - 1) % cp_size
     cp_rank = get_attention_cp_rank()
     bs = len(forward_batch.extend_seq_lens_cpu)
-    cp_group = get_attention_cp_group()
 
-    if cp_rank == owner and hidden_states.shape[0] > 0:
-        local_last = hidden_states[-bs:].contiguous()
-    else:
-        local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
+    request_offsets = _prefix_offsets([int(x) for x in forward_batch.extend_seq_lens_cpu])
+    owners = []
+    local_offsets = []
+    for req_offset, req_len in zip(request_offsets, forward_batch.extend_seq_lens_cpu):
+        req_len = int(req_len)
+        if req_len <= 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_empty_extend] "
+                f"round-robin last-token collect got non-positive req_len={req_len}"
+            )
+        global_last = int(req_offset) + req_len - 1
+        owners.append(global_last % cp_size)
+        local_offsets.append(global_last // cp_size)
+
+    local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
+    for req_id, (owner, local_offset) in enumerate(zip(owners, local_offsets)):
+        if cp_rank != owner:
+            continue
+        if local_offset >= int(hidden_states.shape[0]):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_last_token_oob] "
+                "round-robin last-token owner metadata points past local hidden. "
+                f"req_id={req_id} local_offset={local_offset} "
+                f"local_tokens={int(hidden_states.shape[0])}"
+            )
+        local_last[req_id] = hidden_states[local_offset]
 
     gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
     attn_cp_all_gather_into_tensor(gathered, local_last)
-    return gathered[owner * bs : owner * bs + bs]
+    gather_indices = torch.tensor(
+        [owner * bs + req_id for req_id, owner in enumerate(owners)],
+        device=hidden_states.device,
+        dtype=torch.long,
+    )
+    return gathered.index_select(0, gather_indices)
 
 
 def _get_in_seq_last_token_owner_and_offset(
@@ -1337,13 +1867,17 @@ def _in_seq_collect_last_token(
 ) -> torch.Tensor:
     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:
+        return _in_seq_collect_last_token_batch(hidden_states, metadata, cp_size, cp_rank, bs)
+
     owner = 0
     local_offset = hidden_states.shape[0] - bs
 
-    if bs == 1 and forward_batch.nsa_cp_metadata is not None:
+    if bs == 1 and metadata is not None:
         actual_token_count = sum(int(x) for x in forward_batch.extend_seq_lens_cpu)
         owner, local_offset = _get_in_seq_last_token_owner_and_offset(
-            split_list=forward_batch.nsa_cp_metadata.split_list,
+            split_list=metadata.split_list,
             cp_size=cp_size,
             actual_token_count=actual_token_count,
         )
@@ -1356,3 +1890,64 @@ def _in_seq_collect_last_token(
     gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
     attn_cp_all_gather_into_tensor(gathered, local_last)
     return gathered[owner * bs : owner * bs + bs]
+
+
+def _in_seq_collect_last_token_batch(
+    hidden_states: torch.Tensor,
+    metadata: NSAContextParallelMetadata,
+    cp_size: int,
+    cp_rank: int,
+    bs: int,
+) -> torch.Tensor:
+    if metadata is None:
+        raise RuntimeError(
+            "[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)
+    if (
+        owners is None
+        or local_offsets is None
+        or rank_offsets is None
+        or len(owners) != bs
+        or len(local_offsets) != bs
+        or len(rank_offsets) != bs
+    ):
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_last_token_metadata] "
+            "CP in-seq bs>1 narrow output requires per-request owner, "
+            "local offset, and rank-local offset metadata."
+        )
+
+    local_last = hidden_states.new_zeros((bs, hidden_states.shape[1]))
+    for req_id, (owner, local_offset, rank_offset) in enumerate(
+        zip(owners, local_offsets, rank_offsets)
+    ):
+        owner = int(owner)
+        if owner < 0 or owner >= cp_size:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_invalid_last_token_owner] "
+                f"req_id={req_id} owner={owner} cp_size={cp_size}"
+            )
+        if cp_rank != owner:
+            continue
+        global_local_offset = int(rank_offset) + int(local_offset)
+        if global_local_offset >= int(hidden_states.shape[0]):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_last_token_oob] "
+                "in-seq last-token metadata points past local hidden. "
+                f"req_id={req_id} owner={owner} rank_offset={rank_offset} "
+                f"local_offset={local_offset} local_tokens={int(hidden_states.shape[0])}"
+            )
+        local_last[req_id] = hidden_states[global_local_offset]
+
+    gathered = hidden_states.new_empty((cp_size * bs, hidden_states.shape[1]))
+    attn_cp_all_gather_into_tensor(gathered, local_last)
+    gather_indices = torch.tensor(
+        [int(owner) * bs + req_id for req_id, owner in enumerate(owners)],
+        device=hidden_states.device,
+        dtype=torch.long,
+    )
+    return gathered.index_select(0, gather_indices)