From 014948c5d79119b59791a5337d1c4a1000798494 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 28 Jun 2026 14:09:58 +0000 Subject: [PATCH] CP shared-KV: global-position-rotated zigzag owner-lane balancing (default on) In NSA in-seq-split prefill-CP shared-KV the owner builder restarted every chunk/request at owner 0, starving physical lane 0 -> load-back livelock (g0057: ~64k cp_load_back_owner_lane_capacity_failed + ~4k owner_lane_exhausted). Rotate the owner/compute assignment by phase=((extend_prefix_len//page_size)+req_index) %cp_size so chunks/requests/draft pages spread across lanes, keeping owner==compute (the local_loc_owner_mismatch invariant) by rotating the zigzag compute split by the same phase. Cache-safe: reload replays recorded page_owners, never recomputes phase. - owner builder + helpers (cp_owner_lane_phase two-term prefix+req_index, cp_in_seq_reverse_map, cp_rotated_per_rank_actual_token) in cp_shared_kv_compute_owner.py - zigzag compute split, gather-back reverse map, bs>1 torch rerange, last-token owner, request_phases on NSAContextParallelMetadata, scalar-path + last-token fail-loud guards in nsa/utils.py - mqa-logits prefill buffer estimator made rotation-aware (per-request phase) so it no longer under-sizes the admission budget - flag SGLANG_CP_SHARED_KV_OWNER_ROTATION (default True; =0 is the legacy escape hatch -- phase forced to 0 = byte-for-byte legacy, needs no new kernel) - startup require_tai_kernel_version("0.0.2") gate (utils/common.py, model_runner) - unit tests (11): tiny-chunk pathology [40,0x7]->[5x8], synchronized-batch spread [0x8]->[0..7], gather-back reassembly, per-rank-token Requires tai-kernel >= 0.0.2 (phase-aware in_seq_all_gather_rerange). Reviewed by 3 adversarial agents (owner==compute consistency, gather-back equivalence, integration safety): SHIP-WITH-NITS, all findings fixed. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/environ.py | 11 + .../sglang/srt/layers/attention/nsa/utils.py | 209 +++++++++++----- .../cp_shared_kv_prefill_buffer_estimator.py | 20 +- .../mem_cache/cp_shared_kv_compute_owner.py | 104 +++++++- .../sglang/srt/model_executor/model_runner.py | 23 ++ python/sglang/srt/utils/common.py | 52 ++++ .../mem_cache/test_cp_owner_lane_rotation.py | 234 ++++++++++++++++++ 7 files changed, 587 insertions(+), 66 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_cp_owner_lane_rotation.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 20da16d75..6d81dcf34 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -509,6 +509,17 @@ class Envs: # DeepSeek MHA Optimization SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192) + # CP shared-KV owner-lane balancing (global-position-rotated zigzag). DEFAULT ON: + # the in-seq-split owner/compute assignment is rotated by + # phase=((extend_prefix_len//page_size)+req_index)%cp_size so chunks/requests/draft + # pages spread across owner lanes instead of always restarting at owner 0 (which + # starves physical lane 0 and livelocks load-back). Requires the phase-aware + # tai-kernel in_seq_all_gather_rerange (>=0.0.2); a shared-KV server boot fails + # loud via require_tai_kernel_version if the kernel is older. Set to 0 only as a + # legacy escape hatch -- it reproduces the unrotated behavior byte-for-byte + # (phase forced to 0) and does not require the new kernel. + SGLANG_CP_SHARED_KV_OWNER_ROTATION = EnvBool(True) + # DeepEP SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 5c7d23a94..7b113cc35 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -23,6 +23,13 @@ from sglang.srt.layers.dp_attention import ( get_attention_dp_rank, is_allocation_symmetric, ) +from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( + build_in_seq_page_compute_owners, + cp_in_seq_reverse_map, + cp_owner_lane_phase, + cp_owner_lane_rotation_enabled, + cp_rotated_per_rank_actual_token, +) from sglang.srt.server_args import get_global_server_args from sglang.srt.utils.common import ceil_align, ceil_div @@ -376,6 +383,10 @@ class NSAContextParallelMetadata: request_page_offsets: List[int] = None request_extend_lens: List[int] = None request_prefix_lens: List[int] = None + # Per-request owner-lane rotation phase = (prefix_len//page_size)%cp_size; 0 == + # legacy (unrotated). Carried so the gather-back rerange can map rank<->segment + # consistently with the rotated compute split (owner == compute). + request_phases: List[int] = None request_padded_pages: List[int] = None request_padded_tokens: List[int] = None request_padding_tokens: List[int] = None @@ -698,19 +709,27 @@ def build_batch_page_aligned_in_seq_split_plan( compute_page_starts = split_info.segment_page_starts compute_page_ends = split_info.segment_page_ends compute_tokens = sum(int(token_count) for token_count in compute_split_list) + # Global-position-rotated owner-lane assignment: this rank computes group + # g (segments g and 2N-1-g) for this request; phase==0 is the legacy + # identity (g==cp_rank). The owner builder rotates by the SAME phase, so + # owner == compute (see _get_in_seq_last_token_owner_and_offset + the owner + # builder in cp_shared_kv_compute_owner.py). + phase = cp_owner_lane_phase(prefix_len, page_size, cp_size, req_id) + g = (cp_rank - phase) % cp_size + mirror_idx = cp_segment_num - g - 1 owner, local_offset = _get_in_seq_last_token_owner_and_offset( split_list=compute_split_list, cp_size=cp_size, actual_token_count=extend_len, + phase=phase, ) - zigzag_index = [cp_rank, cp_segment_num - cp_rank - 1] + zigzag_index = [g, mirror_idx] 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_list[g] + split_list[mirror_idx] ) compute_rank_local_tokens = ( - compute_split_list[cp_rank] + compute_split_list[mirror_idx] + compute_split_list[g] + compute_split_list[mirror_idx] ) split_prefix_list = [0] + prefix_sum_list[:-1] @@ -723,9 +742,9 @@ def build_batch_page_aligned_in_seq_split_plan( 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_prev.append(prefix_sum_list[g]) request_kv_len_next.append(prefix_sum_list[mirror_idx]) - request_actual_seq_q_prev.append(compute_split_list[cp_rank]) + request_actual_seq_q_prev.append(compute_split_list[g]) 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) @@ -736,14 +755,14 @@ def build_batch_page_aligned_in_seq_split_plan( 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_prev.append(split_list[g]) 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_prev.append(split_list[g]) 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]), + (0, split_list[g]), + (compute_split_list[g], split_list[mirror_idx]), ] ) request_compute_split_lists.append(compute_split_list) @@ -753,9 +772,9 @@ def build_batch_page_aligned_in_seq_split_plan( request_compute_padded_tokens.append(compute_tokens) request_compute_padding_tokens.append(compute_tokens - 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_prev.append(compute_split_list[g]) 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_prev.append(compute_split_list[g]) request_compute_seq_q_next.append(compute_split_list[mirror_idx]) flat_split_list.extend(split_list) segment_base = req_id * cp_segment_num @@ -1369,10 +1388,6 @@ def _pad_cp_request_tensor_for_split( 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) @@ -1382,6 +1397,7 @@ def build_flat_page_owner_plan(plan) -> List[int]: extend_prefix_len=int(prefix_len), page_size=int(plan.page_size), cp_size=int(plan.cp_size), + req_index=req_id, ) if request_owners is None: raise RuntimeError( @@ -1532,13 +1548,18 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan): 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 communication_split_lists: - rank_tokens += split_list[rank] + split_list[mirror] - per_rank_actual_token.append(rank_tokens) + cp_segment_num = plan.cp_size * 2 + # Per-request rotation phase (identical to the owner builder's): rank r computes + # group g=(r-phase)%cp_size for each request. phase==0 is the legacy identity, so + # every expression below reduces exactly to the unrotated behavior when rotation + # is off. + request_phases = [ + cp_owner_lane_phase(prefix_len, plan.page_size, plan.cp_size, req_index) + for req_index, prefix_len in enumerate(plan.request_prefix_lens) + ] + per_rank_actual_token = cp_rotated_per_rank_actual_token( + communication_split_lists, request_phases, plan.cp_size + ) 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)] @@ -1554,32 +1575,31 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan): ) 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_phase = request_phases[0] if request_phases else 0 + first_g = (plan.cp_rank - first_phase) % plan.cp_size + first_mirror = cp_segment_num - first_g - 1 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_valid_prefix_sum[first_g] if first_valid_prefix_sum else 0 ) - first_mirror = plan.cp_size * 2 - plan.cp_rank - 1 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_prev = first_split[first_g] 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 [] - ) + # Gather-back reverse map for the rank-major all-gather of the first request + # (rotated by first_phase; reduces to the legacy map at first_phase==0). + if first_split: + first_reverse_split_len, first_cp_reverse_index = cp_in_seq_reverse_map( + first_split, plan.cp_size, first_phase + ) + else: + first_reverse_split_len, first_cp_reverse_index = [], [] return NSAContextParallelMetadata( split_list=first_split, @@ -1648,6 +1668,7 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan): request_page_offsets=plan.request_page_offsets, request_extend_lens=plan.request_extend_lens, request_prefix_lens=plan.request_prefix_lens, + request_phases=request_phases, request_padded_pages=plan.request_padded_pages, request_padded_tokens=plan.request_padded_tokens, request_padding_tokens=plan.request_padding_tokens, @@ -2504,16 +2525,31 @@ def _try_tai_in_seq_all_gather_rerange( return None try: - return in_seq_all_gather_rerange( - input_tensor_all, - split_lens, - split_prefix, + rerange_kwargs = dict( total_tokens=total_tokens, hidden_size=hidden_size, max_segment_len=max_segment_len, max_rank_token=max_rank_token, cp_size=cp_size, ) + # Pass the bs=1 request's rotation phase only when rotation is enabled, so an + # older (phase-unaware) tai-kernel keeps working unchanged with the flag off. + # The phase-aware kernel is enforced at startup by require_tai_kernel_version. + if cp_owner_lane_rotation_enabled(): + request_phases = getattr(metadata, "request_phases", None) + if not request_phases: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][owner_rotation_missing_phase] " + "owner-lane rotation is on but nsa_cp_metadata.request_phases is " + "missing/empty for the bs=1 gather-back rerange" + ) + rerange_kwargs["phase"] = int(request_phases[0]) + return in_seq_all_gather_rerange( + input_tensor_all, + split_lens, + split_prefix, + **rerange_kwargs, + ) except Exception as exc: _log_tai_in_seq_rerange_fallback( "kernel_failed", @@ -2574,6 +2610,7 @@ def _torch_batch_in_seq_all_gather_rerange( metadata, "request_compute_split_lists", None ) max_rank_len = getattr(metadata, "max_rank_len", None) + request_phases = getattr(metadata, "request_phases", None) if metadata is None: _raise_batch_rerange_error("missing_metadata", "nsa_cp_metadata is missing") if batch_size <= 1: @@ -2672,14 +2709,30 @@ def _torch_batch_in_seq_all_gather_rerange( cp_size, ) + # Per-request rotation phase (identical to the compute split's); 0 == legacy. + # Rank r holds group g=(r-phase)%cp_size for each request, so its rows are + # split[g] + split[2N-1-g]. This MUST match the rotated owner/compute split. + if request_phases is not None and len(request_phases) == batch_size: + phases = [int(p) % cp_size for p in request_phases] + elif cp_owner_lane_rotation_enabled(): + _raise_batch_rerange_error( + "missing_request_phases", + "owner-lane rotation is on but request_phases is missing/mismatched. " + "have=%s batch_size=%s", + None if request_phases is None else len(request_phases), + batch_size, + ) + else: + phases = [0] * batch_size rank_request_offsets: List[List[int]] = [] for source_rank in range(cp_size): - mirror = cp_size * 2 - source_rank - 1 offsets: List[int] = [] cursor = 0 - for split_list in source_split_lists: + for req_id, split_list in enumerate(source_split_lists): + g = (source_rank - phases[req_id]) % cp_size + mirror = cp_size * 2 - g - 1 offsets.append(cursor) - cursor += split_list[source_rank] + split_list[mirror] + cursor += split_list[g] + split_list[mirror] if cursor > max_rank_token: _raise_batch_rerange_error( "rank_payload_exceeds_max", @@ -2705,12 +2758,16 @@ def _torch_batch_in_seq_all_gather_rerange( for segment_id, segment_len in enumerate(output_split): if segment_len <= 0: continue - if segment_id < cp_size: - source_rank = segment_id - source_segment_offset = 0 - else: - source_rank = cp_size * 2 - segment_id - 1 - source_segment_offset = source_split[source_rank] + # The group's low segment (== zigzag_base) sets the within-rank offset; + # the source RANK is that group rotated by this request's phase + # (owner == compute). phase==0 reduces to the legacy identity. + low_segment = ( + segment_id if segment_id < cp_size else cp_size * 2 - segment_id - 1 + ) + source_rank = (low_segment + phases[req_id]) % cp_size + source_segment_offset = ( + 0 if segment_id < cp_size else source_split[low_segment] + ) source_start = ( source_rank * max_rank_token + rank_request_offsets[source_rank][req_id] @@ -2955,7 +3012,20 @@ def prepare_input_dp_with_cp_dsa( ) return _build_batch_metadata_from_plan(batch_plan) - # scalar compatibility path + # scalar compatibility path (legacy fallback, NOT the shared-KV batch path). + # The owner-lane rotation is wired only through the batch-plan path above; if a + # shared-KV forward reaches this fallback with rotation on, the rotated + # allocation and this unrotated compute split would disagree and trip + # local_loc_owner_mismatch -- fail loud rather than corrupt. + if cp_owner_lane_rotation_enabled() and getattr( + forward_batch, "uses_cp_shared_kv", False + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][owner_rotation_scalar_path_unsupported] " + "owner-lane rotation requires the batch-plan in-seq metadata path, but " + "the scalar compatibility fallback was reached for a shared-KV forward " + "(missing extend_seq_lens_cpu/extend_prefix_lens_cpu)." + ) kv_len_int = int(kv_len) kv_len = torch.tensor(kv_len_int) bs_per_cp_group = 1 @@ -3123,14 +3193,16 @@ def _get_in_seq_last_token_owner_and_offset( split_list: List[int], cp_size: int, actual_token_count: int, + phase: int = 0, ) -> Tuple[int, int]: """Return the CP rank and local offset for the real last token. - In in-seq split, each rank owns two logical segments: - rank r gets segment r followed by segment (2 * cp_size - r - 1). - `split_list` is built from the padded model input length, while - `actual_token_count` is the non-padded extend length. The real last token - can therefore sit in a middle segment instead of rank 0's trailing segment. + In in-seq split, each rank owns two logical segments: under the rotated + assignment rank r gets the group g=(r-phase)%cp_size, i.e. segment g followed + by segment (2*cp_size-g-1). `split_list` is built from the padded model input + length, while `actual_token_count` is the non-padded extend length. The real + last token can therefore sit in a middle segment instead of the trailing one. + ``phase == 0`` is the legacy (unrotated) identity. """ if cp_size <= 0: raise ValueError(f"cp_size must be positive, got {cp_size}") @@ -3166,11 +3238,16 @@ def _get_in_seq_last_token_owner_and_offset( f"actual={actual_token_count} split_list={split_list}" ) + # The group's "low" segment (== zigzag_base of segment_idx) sets the + # within-rank offset; the owning RANK is that group rotated by phase so + # owner == compute. phase == 0 reduces to the legacy identity. + low_segment = ( + segment_idx if segment_idx < cp_size else cp_segment_num - segment_idx - 1 + ) + owner = (low_segment + phase) % cp_size if segment_idx < cp_size: - return segment_idx, offset_in_segment - - owner = cp_segment_num - segment_idx - 1 - local_offset = split_list[owner] + offset_in_segment + return owner, offset_in_segment + local_offset = split_list[low_segment] + offset_in_segment return owner, local_offset @@ -3190,6 +3267,18 @@ def _in_seq_collect_last_token( if bs == 1 and metadata is not None: actual_token_count = sum(int(x) for x in forward_batch.extend_seq_lens_cpu) + # This bs=1 scalar fallback is the non-shared-KV path (shared-KV always has a + # batch plan and routes above), so phase=0 is correct here. Guard against a + # future routing change sending a rotated shared-KV forward through with the + # unrotated phase=0 -- fail loud rather than mis-place the gathered last token. + if cp_owner_lane_rotation_enabled() and getattr( + forward_batch, "uses_cp_shared_kv", False + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][owner_rotation_last_token_scalar_path] " + "owner-lane rotation requires the batch last-token path, but the bs=1 " + "scalar path was reached for a shared-KV forward" + ) owner, local_offset = _get_in_seq_last_token_owner_and_offset( split_list=metadata.split_list, cp_size=cp_size, diff --git a/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py b/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py index e2ea9bfb5..d12819293 100644 --- a/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py +++ b/python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py @@ -197,13 +197,27 @@ def _estimate_cp_split_mqa_logits_peak_bytes( ) return row_cap * k_rows * 4 + from sglang.srt.mem_cache.cp_shared_kv_compute_owner import cp_owner_lane_phase + cp_segment_num = cp_size * 2 + # Owner-lane rotation: rank cp_rank computes group g=(cp_rank-phase[req])%cp_size + # for each request (NOT the fixed group cp_rank), so a single rank can accumulate + # the heaviest group across requests. Mirror the runtime phase here so the peak + # estimate is exact -- otherwise it under-sizes this admission budget. phase==0 + # (rotation off) reduces to the legacy fixed (cp_rank, mirror) diagonal. + request_phases = [ + cp_owner_lane_phase(int(prefix_len), page_size, cp_size, req_index) + for req_index, prefix_len in enumerate(prefix_lens) + ] max_peak_bytes = 0 for cp_rank in range(cp_size): q_rows = 0 k_rows = 0 - mirror_idx = cp_segment_num - cp_rank - 1 - for prefix_len, extend_len in zip(prefix_lens, extend_lens): + for req_index, (prefix_len, extend_len) in enumerate( + zip(prefix_lens, extend_lens) + ): + g = (cp_rank - request_phases[req_index]) % cp_size + mirror_idx = cp_segment_num - g - 1 split_list = _cp_segment_valid_lengths( extend_len=int(extend_len), page_size=page_size, @@ -212,7 +226,7 @@ def _estimate_cp_split_mqa_logits_peak_bytes( segment_end = 0 for segment_idx, segment_len in enumerate(split_list): segment_end += int(segment_len) - if segment_idx not in (cp_rank, mirror_idx) or segment_len <= 0: + if segment_idx not in (g, mirror_idx) or segment_len <= 0: continue q_rows += int(segment_len) # Match the fused CP MQA path: each owned segment materializes a diff --git a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py index 5eea7a77e..80e6f61bd 100644 --- a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py +++ b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py @@ -3,6 +3,43 @@ from __future__ import annotations from typing import List, Optional +def cp_owner_lane_rotation_enabled() -> bool: + """Global-position-rotated zigzag owner assignment is enabled. + + OFF (default) reproduces the legacy unrotated behavior byte-for-byte. See + docs_internal/cp_shared_kv_owner_lane_rotation_design.md. + """ + from sglang.srt.environ import envs + + return bool(envs.SGLANG_CP_SHARED_KV_OWNER_ROTATION.get()) + + +def cp_owner_lane_phase( + extend_prefix_len: int, page_size: int, cp_size: int, req_index: int = 0 +) -> int: + """Per-request rotation phase = ((extend_prefix_len//page_size) + req_index) % cp_size. + + Returns 0 (identity, legacy behavior) when rotation is disabled or inputs are + degenerate. Two terms, both rank-uniform: + * ``extend_prefix_len // page_size`` -- the page's global page position in its + sequence (for EAGLE draft-extend this is the accepted length). Spreads the + chunks of one sequence across lanes as it grows. + * ``req_index`` -- the request's position in the current batch. Spreads a + SYNCHRONIZED batch of short extends (e.g. many brand-new prefix=0 requests), + which the prefix term alone cannot (they would all map to owner 0). + Phase need NOT be cache-stable: reload replays the recorded ``page_owners`` (it + never recomputes the phase), so any rank-uniform deterministic value is correct. + The SAME phase MUST be applied to the owner builder and to every zigzag + compute-split binding (the owner == compute invariant), with the SAME req_index + -- alloc and compute iterate the batch in the same order. + """ + if page_size <= 0 or cp_size <= 0: + return 0 + if not cp_owner_lane_rotation_enabled(): + return 0 + return ((int(extend_prefix_len) // page_size) + int(req_index)) % cp_size + + def get_in_seq_page_compute_owner_unavailable_reason( *, extend_len: int, @@ -27,12 +64,63 @@ def get_in_seq_page_compute_owner_unavailable_reason( return None +def cp_rotated_per_rank_actual_token( + split_lists: List[List[int]], phases: List[int], cp_size: int +) -> List[int]: + """Per-rank total owned tokens across a batch, under per-request rotation. + + Rank r computes group g=(r-phase[req])%cp_size for each request, owning segments + g and 2*cp_size-1-g; this sums those across requests (drives max_rank_token). + phase==0 reduces to split[r]+split[2N-1-r] summed -- the legacy behavior. + """ + cp_segment_num = cp_size * 2 + out: List[int] = [] + for rank in range(cp_size): + total = 0 + for req_index, split_list in enumerate(split_lists): + g = (rank - phases[req_index]) % cp_size + total += split_list[g] + split_list[cp_segment_num - 1 - g] + out.append(total) + return out + + +def cp_in_seq_reverse_map( + split_list: List[int], cp_size: int, phase: int +) -> tuple[List[int], List[int]]: + """(reverse_split_len, cp_reverse_index) for the rank-major all-gather reassembly + of one request's 2*cp_size segments under rotation ``phase``. + + The gathered buffer is rank-major: rank j holds group g_j=(j-phase)%cp_size, i.e. + pieces (segment g_j, segment 2N-1-g_j) of lengths (split[g_j], split[2N-1-g_j]). + ``cp_reverse_index`` reorders those 2N pieces into sequence segment order: segment + s sits in rank (zigzag_base(s)+phase)%cp_size's buffer (its low piece if s Optional[List[int]]: """Return compute-owner CP rank for each newly allocated current page. @@ -65,13 +153,20 @@ def build_in_seq_page_compute_owners( base_units = num_page_units // cp_segment_num remainder_units = num_page_units % cp_segment_num + # Rotate the owner of each zigzag group by the page's global sequence position + # so successive chunks/requests do not all dump their first segment onto owner 0. + # phase == 0 (rotation disabled) is the legacy identity mapping. The compute + # split must rotate by the SAME phase (rank r computes group (r-phase)%cp_size) + # so owner == compute holds; see the zigzag bindings in nsa/utils.py. + phase = cp_owner_lane_phase(extend_prefix_len, page_size, cp_size, req_index) owners: List[int] = [] for segment_idx in range(cp_segment_num): unit_count = base_units + (1 if segment_idx < remainder_units else 0) if segment_idx < cp_size: - owner = segment_idx + base_owner = segment_idx else: - owner = cp_segment_num - segment_idx - 1 + base_owner = cp_segment_num - segment_idx - 1 + owner = (base_owner + phase) % cp_size owners.extend([owner] * unit_count) return owners @@ -91,12 +186,15 @@ def build_batch_in_seq_page_compute_owners( ) batch_owners: List[int] = [] - for extend_len, extend_prefix_len in zip(extend_lens, extend_prefix_lens): + for req_index, (extend_len, extend_prefix_len) in enumerate( + zip(extend_lens, extend_prefix_lens) + ): owners = build_in_seq_page_compute_owners( extend_len=int(extend_len), extend_prefix_len=int(extend_prefix_len), page_size=page_size, cp_size=cp_size, + req_index=req_index, ) if owners is None: return None diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 63b1ef54b..64e12ae18 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -341,6 +341,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.physical_max_total_num_tokens = None self.uses_cp_shared_kv = server_args.enable_nsa_prefill_cp_shared_kv + self._validate_cp_owner_lane_rotation_prereqs() self.cp_shared_kv_layout = None self.is_hybrid_swa = model_config.is_hybrid_swa self.is_hybrid_swa_compress = model_config.is_hybrid_swa_compress @@ -447,6 +448,28 @@ class ModelRunner(ModelRunnerKVCacheMixin): self._model_update_group = {} self._weights_send_group = {} + def _validate_cp_owner_lane_rotation_prereqs(self) -> None: + # The global-position-rotated owner-lane assignment (off by default, + # SGLANG_CP_SHARED_KV_OWNER_ROTATION) needs the phase-aware tai-kernel + # in_seq_all_gather_rerange for the bs=1 gather-back; an older kernel + # hardcodes the unrotated rank<->segment map and would scramble the + # gathered KV. Gate at startup so an outdated kernel fails loud here, not + # mid-inference. + from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( + cp_owner_lane_rotation_enabled, + ) + from sglang.srt.utils import require_tai_kernel_version + + if not (self.uses_cp_shared_kv and cp_owner_lane_rotation_enabled()): + return + # Min tai-kernel version that ships the phase-aware cp_in_seq_rerange kernel. + require_tai_kernel_version( + "0.0.2", + reason="CP shared-KV owner-lane rotation " + "(SGLANG_CP_SHARED_KV_OWNER_ROTATION) needs the phase-aware " + "in_seq_all_gather_rerange kernel", + ) + def init_mindspore_runner(self): # Init the mindspore runner # for now, there is only some communication initialization work diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 643efdaf8..0d4ae99c6 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3867,3 +3867,55 @@ def bind_to_closest_numa_node_cuda(): if is_numa_available() and nvgpu_available(): node_id = get_current_device_numa_node_cuda() numa_bind_to_node(node_id) + + +def get_tai_kernel_version() -> Optional[str]: + """Best-effort installed tai-kernel version string, or None if unavailable.""" + try: + import tai_kernel + except Exception: + return None + version = getattr(tai_kernel, "__version__", None) + if version: + return str(version) + try: + from importlib.metadata import version as _pkg_version + + return _pkg_version("tai_kernel") + except Exception: + return None + + +def require_tai_kernel_version(min_version: str, *, reason: str) -> None: + """Fail loud unless the installed tai-kernel is at least ``min_version``. + + General tai-kernel version gate (not feature-specific): a feature that depends + on a newer tai-kernel calls this at startup so an outdated kernel raises a clear + error instead of silently mis-executing. Version strings are simple dotted + numerics (e.g. "0.0.2"); non-numeric suffixes are ignored in the comparison. + """ + + def _parse(v: str) -> tuple: + parts = [] + for token in str(v).split("."): + num = "" + for ch in token: + if ch.isdigit(): + num += ch + else: + break + parts.append(int(num) if num else 0) + return tuple(parts) + + have = get_tai_kernel_version() + if have is None: + raise RuntimeError( + f"tai-kernel >= {min_version} is required ({reason}), but tai-kernel is " + "not installed / its version could not be determined. Build or upgrade " + "tai-kernel." + ) + if _parse(have) < _parse(min_version): + raise RuntimeError( + f"tai-kernel >= {min_version} is required ({reason}), but the installed " + f"version is {have}. Rebuild/upgrade tai-kernel." + ) diff --git a/test/registered/unit/mem_cache/test_cp_owner_lane_rotation.py b/test/registered/unit/mem_cache/test_cp_owner_lane_rotation.py new file mode 100644 index 000000000..8266beb8e --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_owner_lane_rotation.py @@ -0,0 +1,234 @@ +"""Unit tests for the global-position-rotated zigzag owner assignment. + +Pure-Python (no torch): drives `build_in_seq_page_compute_owners` directly. + +Root cause it guards (see docs_internal/cp_shared_kv_owner_lane_rotation_design.md): +the legacy owner builder restarts every chunk/request at owner 0, so under chunked +prefill / bs>1 / EAGLE every first-segment page piles onto physical lane 0 and +starves it. The rotation makes owner a function of the page's global sequence +position (phase = extend_prefix_len // page_size % cp_size), balancing all lanes. +""" + +import importlib.util +import os +import unittest + +_MODULE_PATH = os.path.join( + os.path.dirname(__file__), + "../../../../python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py", +) + + +def _load_owner_module(): + spec = importlib.util.spec_from_file_location( + "cp_shared_kv_compute_owner_under_test", os.path.abspath(_MODULE_PATH) + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +cow = _load_owner_module() +PAGE = 64 +CP = 8 + + +def _set_rotation(enabled: bool): + # Patch the flag accessor directly so the test stays torch/env-free. + cow.cp_owner_lane_rotation_enabled = lambda: enabled + + +def _seq_owner_histogram(total_tokens, chunk_tokens, *, rotate, page=PAGE, cp=CP): + _set_rotation(rotate) + owners = [] + prefix = 0 + while prefix < total_tokens: + ext = min(chunk_tokens, total_tokens - prefix) + # A new owned page is minted only at a page boundary (page-floored prefix). + if prefix % page == 0: + o = cow.build_in_seq_page_compute_owners( + extend_len=ext, extend_prefix_len=prefix, page_size=page, cp_size=cp + ) + if o: + owners += o + prefix += ext + hist = [0] * cp + for o in owners: + hist[o] += 1 + return hist + + +class TestCpOwnerLaneRotation(unittest.TestCase): + def test_phase_equals_prefix_pages_mod_cp(self): + _set_rotation(True) + for p_pages in range(0, 3 * CP): + prefix = p_pages * PAGE + owners = cow.build_in_seq_page_compute_owners( + extend_len=PAGE, extend_prefix_len=prefix, page_size=PAGE, cp_size=CP + ) + self.assertEqual(owners, [p_pages % CP], msg=f"prefix_pages={p_pages}") + + def test_disabled_is_legacy_identity(self): + # phase forced to 0 -> every chunk starts at owner 0 (the legacy behavior). + _set_rotation(False) + owners = cow.build_in_seq_page_compute_owners( + extend_len=PAGE, extend_prefix_len=10 * PAGE, page_size=PAGE, cp_size=CP + ) + self.assertEqual(owners, [0]) + + def test_tiny_chunks_balance_under_rotation(self): + # The production pathology: many 1-page continuation chunks. + off = _seq_owner_histogram(PAGE * 40, PAGE, rotate=False) + on = _seq_owner_histogram(PAGE * 40, PAGE, rotate=True) + self.assertEqual(off, [40, 0, 0, 0, 0, 0, 0, 0]) # legacy starves lane 0.. + self.assertEqual(max(on) - min(on), 0) # ..rotation spreads perfectly + self.assertEqual(on, [5] * CP) + + def test_balanced_for_any_chunking(self): + # For any chunk size the rotated histogram is (a) never worse than legacy + # and (b) bounded by the unavoidable intra-chunk zigzag imbalance + # ceil(chunk_pages / cp_size) -- the legacy code instead grows the lane-0 + # deficit without bound as chunks accumulate. + total = PAGE * 30 + for chunk_pages in (1, 2, 3, 5, 7, 13): + on = _seq_owner_histogram(total, PAGE * chunk_pages, rotate=True) + off = _seq_owner_histogram(total, PAGE * chunk_pages, rotate=False) + bound = -(-chunk_pages // CP) # ceil(chunk_pages / cp_size) + self.assertLessEqual( + max(on) - min(on), bound, msg=f"chunk_pages={chunk_pages} hist={on}" + ) + self.assertLessEqual( + max(on) - min(on), + max(off) - min(off), + msg=f"rotation worse than legacy: chunk_pages={chunk_pages} on={on} off={off}", + ) + + def test_single_big_forward_unchanged(self): + # A single large forward already spreads owners; rotation must not regress it. + off = _seq_owner_histogram(PAGE * 40, PAGE * 40, rotate=False) + on = _seq_owner_histogram(PAGE * 40, PAGE * 40, rotate=True) + self.assertEqual(max(off) - min(off), 0) + self.assertEqual(max(on) - min(on), 0) + + def test_batch_builder_matches_per_request(self): + _set_rotation(True) + extend_lens = [PAGE, 2 * PAGE, PAGE] + prefix_lens = [3 * PAGE, 5 * PAGE, 0] + batch = cow.build_batch_in_seq_page_compute_owners( + extend_lens=extend_lens, + extend_prefix_lens=prefix_lens, + page_size=PAGE, + cp_size=CP, + ) + expected = [] + for req_index, (e, p) in enumerate(zip(extend_lens, prefix_lens)): + expected += cow.build_in_seq_page_compute_owners( + extend_len=e, + extend_prefix_len=p, + page_size=PAGE, + cp_size=CP, + req_index=req_index, + ) + self.assertEqual(batch, expected) + # req0 prefix3+idx0=phase3 -> [3]; req1 prefix5+idx1=phase6, 2 pages -> [6,7]; + # req2 prefix0+idx2=phase2 -> [2]. + self.assertEqual(batch, [3, 6, 7, 2]) + + def test_synchronized_short_requests_spread(self): + # The case the prefix term alone cannot fix: a batch of brand-new (prefix=0) + # 1-page requests. Legacy/prefix-only maps them all to owner 0; the req_index + # term spreads them round-robin across lanes. + _set_rotation(True) + on = cow.build_batch_in_seq_page_compute_owners( + extend_lens=[PAGE] * CP, + extend_prefix_lens=[0] * CP, + page_size=PAGE, + cp_size=CP, + ) + self.assertEqual(on, list(range(CP))) # owners 0,1,...,cp-1 + _set_rotation(False) + off = cow.build_batch_in_seq_page_compute_owners( + extend_lens=[PAGE] * CP, + extend_prefix_lens=[0] * CP, + page_size=PAGE, + cp_size=CP, + ) + self.assertEqual(off, [0] * CP) # the pathology + + +class TestCpInSeqReverseMap(unittest.TestCase): + """The gather-back reverse map used by _build_batch_metadata_from_plan.""" + + def _reassemble(self, split_list, cp, phase): + # Build the rank-major all-gather buffer (rank j holds group g_j=(j-phase)%cp), + # apply the real helper's (reverse_split_len, cp_reverse_index), and return the + # reconstructed flat segment-marker sequence. + from itertools import accumulate # noqa: F401 + + N2 = cp * 2 + rsl, cri = cow.cp_in_seq_reverse_map(split_list, cp, phase) + trimmed = [] + for j in range(cp): + g = (j - phase) % cp + for seg in (g, N2 - 1 - g): + trimmed += [seg] * split_list[seg] + pieces, off = [], 0 + for length in rsl: + pieces.append(trimmed[off : off + length]) + off += length + out = [] + for i in cri: + out += pieces[i] + return out + + def test_reconstructs_sequence_order_all_phases(self): + import random + + for cp in (2, 4, CP): + N2 = cp * 2 + for trial in range(60): + rng = random.Random(trial * 13 + cp) + split_list = [rng.randint(0, 5) for _ in range(N2)] + expected = [] + for s in range(N2): + expected += [s] * split_list[s] + for phase in range(cp): + self.assertEqual( + self._reassemble(split_list, cp, phase), + expected, + msg=f"cp={cp} phase={phase} split={split_list}", + ) + + def test_identity_at_phase_zero(self): + split_list = [1, 2, 1, 3, 2, 1, 2, 1] # cp=4 + rsl, cri = cow.cp_in_seq_reverse_map(split_list, 4, 0) + self.assertEqual(cri, [0, 2, 4, 6, 7, 5, 3, 1]) # legacy reverse index + self.assertEqual( + rsl, + [split_list[0], split_list[7], split_list[1], split_list[6], + split_list[2], split_list[5], split_list[3], split_list[4]], + ) + + +class TestCpRotatedPerRankActualToken(unittest.TestCase): + def test_phase_shifts_which_rank_owns_tokens(self): + # req0 (zigzag-balanced -> 9 per rank) + req1 whose only tokens are in segment 0 + # (10 tokens, owned by group 0 = the rank at (0+phase)). The phase moves those 10. + sls = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 0, 0, 0, 0, 0, 0, 0]] + self.assertEqual( + cow.cp_rotated_per_rank_actual_token(sls, [0, 1], 4), [9, 19, 9, 9] + ) + self.assertEqual( + cow.cp_rotated_per_rank_actual_token(sls, [0, 0], 4), [19, 9, 9, 9] + ) + + def test_identity_at_phase_zero(self): + N2 = CP * 2 + sls = [[i + 1 for i in range(N2)], [1] * N2] + got = cow.cp_rotated_per_rank_actual_token(sls, [0] * len(sls), CP) + expected = [sum(sl[r] + sl[N2 - 1 - r] for sl in sls) for r in range(CP)] + self.assertEqual(got, expected) + + +if __name__ == "__main__": + unittest.main()