From e4cf8d18b4cc1b4ebec6e387e8024c267d229dd6 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 3 Jun 2026 01:21:43 +0800 Subject: [PATCH] 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 --- ...shared_kv_bs_gt1_implementation_plan_zh.md | 15 + ...hared_kv_bs_gt1_parallel_workstreams_zh.md | 10 + .../sglang/srt/layers/attention/nsa/utils.py | 625 +++++++++++++++++- .../unit/layers/test_nsa_cp_utils.py | 177 +++++ 4 files changed, 812 insertions(+), 15 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md index 20158e761..39017794d 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md @@ -272,8 +272,15 @@ NSAContextParallelMetadata request_kv_len_next: List[int] request_actual_seq_q_prev: List[int] request_actual_seq_q_next: List[int] + request_actual_seq_q_prev_cu_tensor: Tensor[batch_size + 1] + request_actual_seq_q_next_cu_tensor: Tensor[batch_size + 1] + request_rank_local_offsets: 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] ``` 兼容规则： @@ -347,9 +354,14 @@ build_batch_page_aligned_in_seq_split_plan( - per-request split list； - per-request page start/end； - per-rank local valid token count； + - per-request `kv_len_prev/next` 与 `actual_seq_q_prev/next`； + - per-request last-token owner/local offset； + - per-request rank-local offset，用于 batch compact hidden collect； - flattened segment offsets。 4. 更新 `prepare_input_dp_with_cp_dsa()`，让 CP shared-KV bs>1 使用新 metadata。 5. 保留 bs=1 的旧字段兼容。 +6. `cp_collect_last_token_hidden()` 必须继续走 phase1 compact/narrow output；不能因为 bs>1 回到 full hidden gather。 +7. bs>1 full rerange 在 batch-aware full output metadata/kernels 完成前 fail-fast。 测试重点： @@ -371,6 +383,9 @@ cp_size=8 退出标准： - metadata 能表达 `batch_size=2`，且不会把两个 request 当成一条长序列。 +- 普通 prefill bs>1 的 output collection 能按 request order 返回 compact last hidden。 +- 需要 full output 语义的 bs>1 请求不会误用 scalar full rerange。 +- 对外 helper 能提供 batch plan、按 plan split tensor、以及 flat page-owner plan。 ### Phase 2：batch-aware owner-lane allocation diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index 216d0d4c0..a5111a268 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md @@ -146,6 +146,14 @@ CPSharedKVBatchPlan request_actual_seq_q_prev request_actual_seq_q_next request_rank_local_tokens + request_rank_local_offsets + request_actual_seq_q_prev_cu_tensor + request_actual_seq_q_next_cu_tensor + + # phase1 narrow output collection + request_last_token_owner + request_last_token_local_offset + output_collect_mode # narrow_last_token | full_rerange | unsupported_fail_fast flat_split_list flat_zigzag_index @@ -161,6 +169,8 @@ CPSharedKVBatchPlan 4. zero-token segment 合法。 5. batch flattened plan 不能跨 request 合并 segment。 6. bs=1 继续兼容现有 scalar fields。 +7. phase1 narrow-output 优化不能因为 bs>1 回退：普通 prefill 必须通过 `request_last_token_owner/local_offset` 只收集每个 request 的 last hidden。 +8. bs>1 的 full rerange/logprob/hidden capture 在 batch-aware full rerange 完成前必须 fail-fast，不能使用 scalar `cp_all_gather_rerange_output()`。 ### 对外接口 diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 81f005db3..241072f33 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/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) diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 24bb40e70..e39e7353b 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -8,16 +8,23 @@ from unittest.mock import patch from sglang.srt.layers.attention.nsa.utils import ( NSAContextParallelMetadata, PageAlignedCacheExtent, + build_flat_page_owner_plan, + build_batch_page_aligned_in_seq_split_plan, build_page_aligned_cache_extent, _get_in_seq_last_token_owner_and_offset, build_page_aligned_in_seq_split_list, build_token_balanced_in_seq_split_list, can_cp_split, + cp_all_gather_rerange_output, + cp_collect_last_token_hidden, cp_split_and_rebuild_1d, + cp_split_and_rebuild_data, + get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, get_cp_shared_kv_local_physical_out_cache_loc, get_cp_local_embedding_padded_token_count, pad_cp_local_input_ids_for_embedding, + split_tensor_by_cp_batch_plan, split_in_seq_cp_local_pair, ) from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -407,6 +414,135 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(owner, 0) self.assertEqual(local_offset, 7) + def test_batch_page_aligned_plan_keeps_request_boundaries_and_last_token_owners( + self, + ): + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[4, 9], + prefix_lens=[0, 8], + page_size=4, + cp_size=2, + cp_rank=0, + ) + + self.assertEqual(plan.batch_size, 2) + self.assertEqual(plan.request_split_lists, [[4, 0, 0, 0], [4, 4, 1, 0]]) + self.assertEqual(plan.request_padded_pages, [1, 3]) + self.assertEqual(plan.request_padded_tokens, [4, 12]) + self.assertEqual(plan.request_token_offsets, [0, 4]) + self.assertEqual(plan.request_padded_token_offsets, [0, 4]) + self.assertEqual(plan.request_page_offsets, [0, 1]) + self.assertEqual(plan.request_last_token_owner, [0, 1]) + self.assertEqual(plan.request_last_token_local_offset, [3, 4]) + self.assertEqual(plan.request_rank_local_tokens, [4, 4]) + self.assertEqual(plan.request_rank_local_offsets, [0, 4]) + self.assertEqual(plan.request_kv_len_prev, [4, 4]) + self.assertEqual(plan.request_kv_len_next, [4, 9]) + self.assertEqual(plan.request_actual_seq_q_prev, [4, 4]) + self.assertEqual(plan.request_actual_seq_q_next, [0, 0]) + self.assertEqual(plan.flat_segment_request_ids, [0, 0, 0, 0, 1, 1, 1, 1]) + self.assertEqual(plan.flat_segment_offsets, [0, 4, 4, 4, 0, 4, 8, 9]) + + def test_batch_plan_stable_helpers_split_and_build_page_owner_plan(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[4, 9], + prefix_lens=[0, 8], + page_size=4, + cp_size=2, + cp_rank=1, + ) + forward_batch = SimpleNamespace(nsa_cp_metadata=SimpleNamespace(batch_plan=plan)) + + self.assertIs(get_cp_shared_kv_batch_plan(forward_batch), plan) + self.assertEqual(build_flat_page_owner_plan(plan), [0, 0, 1, 1]) + + local_1d = split_tensor_by_cp_batch_plan(torch.arange(13), plan, mode="1d") + self.assertEqual(local_1d.tolist(), [8, 9, 10, 11, 12]) + + local_data = split_tensor_by_cp_batch_plan( + torch.arange(13 * 2).view(13, 2), plan, mode="data" + ) + self.assertEqual(local_data[:, 0].tolist(), list(range(16, 26, 2))) + + def test_collect_last_token_hidden_uses_batch_owner_metadata(self): + import torch + + hidden_states = torch.tensor( + [[10.0], [11.0], [12.0], [13.0], [20.0], [21.0], [22.0], [23.0]] + ) + forward_batch = SimpleNamespace( + extend_seq_lens_cpu=[4, 9], + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_last_token_owner=[0, 1], + request_last_token_local_offset=[3, 4], + request_rank_local_offsets=[0, 4], + ), + ) + + def fake_all_gather(output, local_last): + self.assertEqual(local_last.tolist(), [[13.0], [0.0]]) + output.copy_(torch.tensor([[13.0], [0.0], [0.0], [99.0]])) + + with ( + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", + return_value=0, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", + side_effect=fake_all_gather, + ), + ): + collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2) + + self.assertEqual(collected.tolist(), [[13.0], [99.0]]) + + def test_collect_last_token_hidden_fails_fast_without_batch_owner_metadata(self): + import torch + + forward_batch = SimpleNamespace( + extend_seq_lens_cpu=[4, 9], + nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2), + ) + + with ( + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank", + return_value=0, + ), + self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_missing_last_token_metadata\]", + ), + ): + cp_collect_last_token_hidden(torch.zeros((8, 1)), forward_batch, 2) + + def test_full_rerange_fails_fast_for_batch_metadata(self): + import torch + + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2) + ) + + with self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_full_rerange_unsupported\]", + ): + cp_all_gather_rerange_output( + torch.zeros((8, 1)), 2, forward_batch, stream=None + ) + def test_local_pair_split_uses_metadata_lengths_not_half_split(self): import torch @@ -438,6 +574,47 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(local_locs.tolist(), [2, 3, 12, 13]) + def test_cp_split_and_rebuild_data_keeps_batch_request_boundaries(self): + import torch + + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_extend_lens=[4, 9], + request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], + request_zigzag_indices=[[0, 3], [0, 3]], + ) + ) + tensor = torch.arange(13 * 2).view(13, 2) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local = cp_split_and_rebuild_data(forward_batch, tensor) + + self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2))) + + def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self): + import torch + + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_extend_lens=[4, 9], + request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], + request_zigzag_indices=[[1, 2], [1, 2]], + ) + ) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local = cp_split_and_rebuild_1d(forward_batch, torch.arange(13)) + + self.assertEqual(local.tolist(), [8, 9, 10, 11, 12]) + def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self): from types import SimpleNamespace