From 91fa31bcac937d57b06c5fc28f4df1232a29fbea Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 1 May 2026 00:54:16 +0800 Subject: [PATCH] Enable compute-owner KV layout by page-aligning NSA CP split 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 --- ...efill_cp_phase4_page_aligned_split_plan.md | 9 +- .../srt/layers/attention/nsa/nsa_indexer.py | 15 +- .../sglang/srt/layers/attention/nsa/utils.py | 208 +++++++++++++++++- python/sglang/srt/models/deepseek_nextn.py | 6 + python/sglang/srt/models/deepseek_v2.py | 6 + .../unit/layers/test_nsa_cp_utils.py | 116 ++++++++++ 6 files changed, 344 insertions(+), 16 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md index b4492e3d6..d65d45f13 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md @@ -234,7 +234,14 @@ Phase 4 MVP 建议： fallback 到旧 token-average split ``` -推荐先使用保守 gate，跑通后再放宽。 +当前实现采用保守 gate： + +```text +如果 num_units < 2 * cp_size: + fallback 到旧 token-average split +``` + +原因是 CP 本身不适合短序列；真实长上下文场景下 page unit 数通常远大于 `2 * cp_size`，先保证每个 zigzag segment 至少拿到一个完整 page unit，可以避免 zero-token segment 给通信、attention kernel 和后续 compute-owner layout 带来额外风险。后续如果要让短序列直接不走 CP，应单独收紧 `can_cp_split(...)` 的启用阈值，而不是混入 Phase 4 的 page-aligned split 逻辑。 --- diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index b444ccc66..803ebd8a4 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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, diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index a37dac0ef..5b18b1fb3 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index d57eb8822..207c58624 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -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( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index a4d00ae00..1d6a5178e 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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): diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index e7e486fdc..be84e5b41 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -2,6 +2,9 @@ import unittest from sglang.srt.layers.attention.nsa.utils import ( _get_in_seq_last_token_owner_and_offset, + build_page_aligned_in_seq_split_list, + build_token_balanced_in_seq_split_list, + split_in_seq_cp_local_pair, ) from sglang.test.ci.ci_register import register_cpu_ci @@ -9,6 +12,103 @@ register_cpu_ci(est_time=1, suite="stage-a-test-cpu") class TestNSAInSeqCPUtils(unittest.TestCase): + def assert_page_aligned_boundaries( + self, split_list, *, extend_prefix_len, extend_len, page_size + ): + cursor = 0 + for segment_len in split_list[:-1]: + cursor += segment_len + if cursor < extend_len: + self.assertEqual((extend_prefix_len + cursor) % page_size, 0) + + def test_page_aligned_split_keeps_boundaries_on_pages(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=32768, + extend_len=32768, + extend_prefix_len=0, + page_size=64, + cp_size=8, + ) + + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 32768) + self.assertEqual(len(split_list), 16) + self.assertTrue(all(segment_len > 0 for segment_len in split_list)) + self.assert_page_aligned_boundaries( + split_list, extend_prefix_len=0, extend_len=32768, page_size=64 + ) + + def test_page_aligned_split_uses_prefix_for_boundary_alignment(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=1024, + extend_len=1024, + extend_prefix_len=128, + page_size=64, + cp_size=8, + ) + + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 1024) + self.assert_page_aligned_boundaries( + split_list, extend_prefix_len=128, extend_len=1024, page_size=64 + ) + + def test_page_aligned_split_keeps_tail_partial_page_unsplit(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=1100, + extend_len=1100, + extend_prefix_len=0, + page_size=64, + cp_size=8, + ) + + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 1100) + self.assertEqual(split_list[-1], 12) + self.assert_page_aligned_boundaries( + split_list, extend_prefix_len=0, extend_len=1100, page_size=64 + ) + + def test_page_aligned_split_falls_back_when_prefix_is_not_page_aligned(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=1024, + extend_len=1024, + extend_prefix_len=1, + page_size=64, + cp_size=8, + ) + + self.assertFalse(info.page_aligned) + self.assertEqual(split_list, build_token_balanced_in_seq_split_list(1024, 8)) + + def test_page_aligned_split_falls_back_when_page_units_are_too_short(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=512, + extend_len=512, + extend_prefix_len=0, + page_size=64, + cp_size=8, + ) + + self.assertFalse(info.page_aligned) + self.assertEqual(split_list, build_token_balanced_in_seq_split_list(512, 8)) + + def test_page_aligned_split_adds_padding_tokens_to_last_segment(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=1040, + extend_len=1024, + extend_prefix_len=0, + page_size=64, + cp_size=8, + ) + + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 1040) + self.assertEqual(split_list[-1], 80) + self.assert_page_aligned_boundaries( + split_list, extend_prefix_len=0, extend_len=1024, page_size=64 + ) + def test_last_token_owner_uses_actual_token_count_when_batch_is_padded(self): # Padded prefill can have 64 model tokens while the real prompt has only # 11 tokens. In in-seq split with cp_size=8, the real last token is in @@ -36,6 +136,22 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(owner, 0) self.assertEqual(local_offset, 7) + def test_local_pair_split_uses_metadata_lengths_not_half_split(self): + import torch + + tensor = torch.arange(9) + + prev, next_ = split_in_seq_cp_local_pair(tensor, 6, 3) + + self.assertEqual(prev.tolist(), [0, 1, 2, 3, 4, 5]) + self.assertEqual(next_.tolist(), [6, 7, 8]) + + def test_local_pair_split_rejects_stale_metadata(self): + import torch + + with self.assertRaisesRegex(RuntimeError, "local in-seq CP length mismatch"): + split_in_seq_cp_local_pair(torch.arange(9), 5, 5, name="q_fp8") + if __name__ == "__main__": unittest.main()