chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

CP shared-KV: global-position-rotated zigzag owner-lane balancing (default on)

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-06-28 14:09:58 +00:00
co-authored by Claude Opus 4.8
parent ce0e486327
commit 014948c5d7
7 changed files with 587 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+11
View File
@@ -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)
python/sglang/srt/layers/attention/nsa/utils.py
+149 -60
View File
@@ -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,
python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py
+17 -3
View File
@@ -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
python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py
+101 -3
View File
@@ -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<cp_size
else its high piece). phase==0 is the legacy [0,2,..,2N-2, 2N-1,2N-3,..,1] map.
"""
cp_segment_num = cp_size * 2
reverse_split_len = [
element
for j in range(cp_size)
for element in (
split_list[(j - phase) % cp_size],
split_list[cp_segment_num - 1 - ((j - phase) % cp_size)],
)
]
cp_reverse_index = [
2 * ((s + phase) % cp_size) for s in range(cp_size)
] + [
2 * ((cp_segment_num - 1 - s + phase) % cp_size) + 1
for s in range(cp_size, cp_segment_num)
]
return reverse_split_len, cp_reverse_index
def build_in_seq_page_compute_owners(
*,
extend_len: int,
extend_prefix_len: int,
page_size: int,
cp_size: int,
req_index: int = 0,
) -> 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
python/sglang/srt/model_executor/model_runner.py
+23
View File
@@ -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
python/sglang/srt/utils/common.py
+52
View File
@@ -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."
)
test/registered/unit/mem_cache/test_cp_owner_lane_rotation.py
+234
View File
@@ -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()