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

Enable compute-owner KV layout by page-aligning NSA CP split

Browse Source 
Phase 5 needs each current KV page to have exactly one CP compute owner before local KV/index direct writes can be safe. This change teaches in-seq NSA prefill CP to produce page-aligned split metadata under shared-KV mode, threads page size into the metadata builders, and fixes local pair splitting so unequal page-aligned zigzag segments do not corrupt topk inputs.

Constraint: Phase 5 direct-write layout requires page ownership to be expressible at page granularity
Constraint: Short page-unit batches remain on the token-balanced fallback to avoid zero-page segment risk
Rejected: Split local q/weights by half | page-aligned zigzag segments can have unequal token counts
Confidence: medium
Scope-risk: moderate
Directive: Do not enable compute-owner direct writes unless nsa_cp_metadata.page_aligned is true and local loc ownership is verified
Tested: python3 -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py python/sglang/srt/models/deepseek_v2.py python/sglang/srt/models/deepseek_nextn.py test/registered/unit/layers/test_nsa_cp_utils.py
Not-tested: Local pytest collection is blocked in this environment by missing pybase64; container/runtime tests were not rerun during this commit step
This commit is contained in:
laoyao0822
2026-05-01 00:54:16 +08:00
parent 47bd2fdf1f
commit 91fa31bcac
6 changed files with 344 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -54,6 +54,7 @@ from sglang.srt.layers.attention.nsa.utils import (
cp_all_gather_rerange_output,
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
split_in_seq_cp_local_pair,
)
from sglang.srt.layers.communicator import ScatterMode
from sglang.srt.layers.linear import ReplicatedLinear
@@ -1374,11 +1375,17 @@ class Indexer(MultiPlatformOp):
# cp_batch_seq_index_prev = forward_batch.nsa_cp_metadata["cp_batch_seq_index_prev"]
# cp_batch_seq_index_next = forward_batch.nsa_cp_metadata["cp_batch_seq_index_next"]
# TODO prev, next, combined into a single call
q_fp8_prev, q_fp8_next = torch.split(
q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0
q_fp8_prev, q_fp8_next = split_in_seq_cp_local_pair(
q_fp8,
actual_seq_q_prev,
actual_seq_q_next,
name="q_fp8",
)
weights_prev, weights_next = torch.split(
weights, (weights.shape[0] + 1) // 2, dim=0
weights_prev, weights_next = split_in_seq_cp_local_pair(
weights,
actual_seq_q_prev,
actual_seq_q_next,
name="weights",
)
topk_result_prev = self._get_topk_ragged_with_cp(
forward_batch,

208
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -135,6 +135,15 @@ def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
return nsa_cache_seqlens
@dataclass
class PageAlignedInSeqSplitInfo:
page_aligned: bool = False
page_size: int = 1
extend_prefix_len: int = 0
segment_page_starts: List[int] = None
segment_page_ends: List[int] = None
@dataclass
class NSAContextParallelMetadata:
split_list: List[int] = None
@@ -152,6 +161,146 @@ class NSAContextParallelMetadata:
actual_seq_q_prev_tensor: torch.Tensor = None
actual_seq_q_next_tensor: torch.Tensor = None
total_seq_lens: torch.Tensor = None
page_aligned: bool = False
page_size: int = 1
extend_prefix_len: int = 0
segment_page_starts: List[int] = None
segment_page_ends: List[int] = None
def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]:
if cp_size <= 0:
raise ValueError(f"cp_size must be positive, got {cp_size}")
if total_len < 0:
raise ValueError(f"total_len must be non-negative, got {total_len}")
cp_segment_num = cp_size * 2
base = total_len // cp_segment_num
remainder = total_len % cp_segment_num
return [base + (1 if i < remainder else 0) for i in range(cp_segment_num)]
def _fallback_page_aligned_split_info(
*,
page_size: int,
extend_prefix_len: int,
) -> PageAlignedInSeqSplitInfo:
return PageAlignedInSeqSplitInfo(
page_aligned=False,
page_size=page_size,
extend_prefix_len=extend_prefix_len,
segment_page_starts=[],
segment_page_ends=[],
)
def build_page_aligned_in_seq_split_list(
*,
total_len: int,
extend_len: int,
extend_prefix_len: int,
page_size: int,
cp_size: int,
) -> Tuple[List[int], PageAlignedInSeqSplitInfo]:
"""Build an in-seq split list whose real-token boundaries do not cut pages.
Phase 4 deliberately uses a conservative gate: at least `2 * cp_size` page
units are required so every zigzag segment has at least one page unit. When
the gate does not hold, this helper falls back to the existing token-balanced
split and marks the result as not page-aligned.
"""
if extend_len < 0:
raise ValueError(f"extend_len must be non-negative, got {extend_len}")
if total_len < extend_len:
raise ValueError(
f"total_len must be >= extend_len, got total_len={total_len} "
f"extend_len={extend_len}"
)
fallback_split = build_token_balanced_in_seq_split_list(total_len, cp_size)
fallback_info = _fallback_page_aligned_split_info(
page_size=page_size,
extend_prefix_len=extend_prefix_len,
)
if page_size <= 1 or extend_len <= 0 or extend_prefix_len % page_size != 0:
return fallback_split, fallback_info
full_pages = extend_len // page_size
tail_tokens = extend_len % page_size
num_page_units = full_pages + (1 if tail_tokens > 0 else 0)
cp_segment_num = cp_size * 2
if num_page_units < cp_segment_num:
return fallback_split, fallback_info
base_units = num_page_units // cp_segment_num
remainder_units = num_page_units % cp_segment_num
unit_counts = [
base_units + (1 if i < remainder_units else 0)
for i in range(cp_segment_num)
]
split_list: List[int] = []
segment_page_starts: List[int] = []
segment_page_ends: List[int] = []
unit_cursor = 0
base_page = extend_prefix_len // page_size
for unit_count in unit_counts:
segment_page_starts.append(base_page + unit_cursor)
token_count = 0
for _ in range(unit_count):
if unit_cursor < full_pages:
token_count += page_size
else:
token_count += tail_tokens
unit_cursor += 1
segment_page_ends.append(base_page + unit_cursor)
split_list.append(token_count)
padding_tokens = total_len - extend_len
if padding_tokens > 0:
split_list[-1] += padding_tokens
return split_list, PageAlignedInSeqSplitInfo(
page_aligned=True,
page_size=page_size,
extend_prefix_len=extend_prefix_len,
segment_page_starts=segment_page_starts,
segment_page_ends=segment_page_ends,
)
def _build_in_seq_split_for_forward_batch(
*,
total_len: int,
cp_size: int,
forward_batch: "ForwardBatch" = None,
page_size: int = None,
) -> Tuple[List[int], PageAlignedInSeqSplitInfo]:
use_page_aligned_split = (
forward_batch is not None
and getattr(forward_batch, "uses_cp_shared_kv", False)
and page_size is not None
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
and len(forward_batch.extend_prefix_lens_cpu) == 1
)
if use_page_aligned_split:
return build_page_aligned_in_seq_split_list(
total_len=total_len,
extend_len=int(forward_batch.extend_seq_lens_cpu[0]),
extend_prefix_len=int(forward_batch.extend_prefix_lens_cpu[0]),
page_size=int(page_size),
cp_size=cp_size,
)
split_list = build_token_balanced_in_seq_split_list(total_len, cp_size)
return split_list, _fallback_page_aligned_split_info(
page_size=int(page_size or 1),
extend_prefix_len=0,
)
def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
@@ -214,6 +363,34 @@ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
return positions
def split_in_seq_cp_local_pair(
input_: torch.Tensor,
prev_len: int,
next_len: int,
*,
name: str = "input",
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Split a local in-seq CP tensor by its two logical segment lengths.
In-seq CP assigns each rank two logical segments: `rank` and
`2 * cp_size - rank - 1`. The legacy token-balanced split made these
segments almost equal, so splitting the local tensor in half happened to
work. Page-aligned split can intentionally make the two segment lengths
different, so consumers must split by metadata lengths instead of by half.
"""
prev_len = int(prev_len)
next_len = int(next_len)
expected = prev_len + next_len
actual = int(input_.shape[0])
if actual != expected:
raise RuntimeError(
f"{name} local in-seq CP length mismatch: actual={actual}, "
f"expected={expected}, prev_len={prev_len}, next_len={next_len}"
)
return torch.split(input_, (prev_len, next_len), dim=0)
@triton.jit
def nsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
@@ -456,6 +633,9 @@ def prepare_input_dp_with_cp_dsa(
cp_rank,
cp_size,
seqs_len,
*,
forward_batch: "ForwardBatch" = None,
page_size: int = None,
):
if is_nsa_prefill_cp_round_robin_split():
return True
@@ -497,19 +677,17 @@ def prepare_input_dp_with_cp_dsa(
- To mitigate uneven load, the input hissenstate needs to be sliced by cp_size*2 and rearranged.
"""
# just support batch = 1
kv_len = torch.tensor(kv_len)
kv_len_int = int(kv_len)
kv_len = torch.tensor(kv_len_int)
bs_per_cp_group = 1
kv_len_origin = kv_len
# get zigzag index
cp_segment_num = cp_size * 2
seq_per_batch = kv_len // cp_segment_num # seq_len for each batch and segment
split_list = seq_per_batch.repeat_interleave(cp_segment_num).int().tolist()
remainder = kv_len % (cp_segment_num)
if remainder > 0:
split_list[:remainder] = [x + 1 for x in split_list[:remainder]]
seq_max_rank_len = (kv_len + cp_size - 1) // cp_size
max_rank_len = seq_max_rank_len.repeat_interleave(cp_size).int().tolist()
split_list, page_split_info = _build_in_seq_split_for_forward_batch(
total_len=kv_len_int,
cp_size=cp_size,
forward_batch=forward_batch,
page_size=page_size,
)
zigzag_index = list(
range(cp_rank, cp_rank + bs_per_cp_group * cp_segment_num, cp_segment_num)
) + list(
@@ -523,6 +701,9 @@ def prepare_input_dp_with_cp_dsa(
per_rank_actual_token = list(
split_list[i] + split_list[cp_size * 2 - i - 1] for i in range(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(cp_size)]
comm_total_seq_lens = torch.tensor(max_rank_token * cp_size)
reverse_split_len = [
element
for i in range(cp_size)
@@ -573,7 +754,12 @@ def prepare_input_dp_with_cp_dsa(
kv_len_next_tensor=kv_len_next_tensor,
actual_seq_q_prev_tensor=actual_seq_q_prev_tensor,
actual_seq_q_next_tensor=actual_seq_q_next_tensor,
total_seq_lens=kv_len_origin,
total_seq_lens=comm_total_seq_lens,
page_aligned=page_split_info.page_aligned,
page_size=page_split_info.page_size,
extend_prefix_len=page_split_info.extend_prefix_len,
segment_page_starts=page_split_info.segment_page_starts,
segment_page_ends=page_split_info.segment_page_ends,
)
return nsa_cp_metadata

6
python/sglang/srt/models/deepseek_nextn.py
View File

@@ -252,6 +252,12 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
forward_batch=forward_batch,
page_size=getattr(
getattr(forward_batch, "token_to_kv_pool", None),
"page_size",
None,
),
)
hidden_states = self.model(input_ids, positions, forward_batch)
return self.logits_processor(

6
python/sglang/srt/models/deepseek_v2.py
View File

@@ -2204,6 +2204,12 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
forward_batch=forward_batch,
page_size=getattr(
getattr(forward_batch, "token_to_kv_pool", None),
"page_size",
None,
),
)
with get_attn_tp_context().maybe_input_scattered(forward_batch):