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

Preserve CP narrow output while planning real batches

Browse Source 
Batch-size support needs request-first CP metadata; treating a batch as one long sequence breaks page ownership, top-k ranges, and phase1 compact output collection. This adds a batch CP plan that records per-request page-aligned splits, rank-local offsets, kv/actual-seq metadata, last-token owners, and flattened descriptors for downstream allocator/runtime workstreams.

The scalar full-rerange path now fail-fasts for batch metadata so bs>1 cannot silently discard the narrow-output optimization or restore hidden states with single-request assumptions.

Constraint: CP shared-KV cache state is page-owned and must preserve request boundaries under bs>1.

Rejected: Let bs>1 fall back to scalar full hidden rerange | it loses the phase1 communication reduction and uses wrong single-request metadata.

Rejected: Add a collective to confirm batch plans | all ranks can derive the same plan from CPU metadata and config.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove batch fail-fast guards until W2/W3 consumers use CPSharedKVBatchPlan end-to-end.

Tested: python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py

Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py -> 39 passed

Not-tested: full ETE bs>1 CP shared-KV runtime; W2/W3 allocator/direct-write consumers are not implemented yet
This commit is contained in:
laoyao0822
2026-06-03 01:21:43 +08:00
parent 0158e28689
commit e4cf8d18b4
4 changed files with 812 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

177
test/registered/unit/layers/test_nsa_cp_utils.py
View File

@@ -8,16 +8,23 @@ from unittest.mock import patch
from sglang.srt.layers.attention.nsa.utils import (
NSAContextParallelMetadata,
PageAlignedCacheExtent,
build_flat_page_owner_plan,
build_batch_page_aligned_in_seq_split_plan,
build_page_aligned_cache_extent,
_get_in_seq_last_token_owner_and_offset,
build_page_aligned_in_seq_split_list,
build_token_balanced_in_seq_split_list,
can_cp_split,
cp_all_gather_rerange_output,
cp_collect_last_token_hidden,
cp_split_and_rebuild_1d,
cp_split_and_rebuild_data,
get_cp_shared_kv_batch_plan,
get_cp_shared_kv_local_out_cache_loc,
get_cp_shared_kv_local_physical_out_cache_loc,
get_cp_local_embedding_padded_token_count,
pad_cp_local_input_ids_for_embedding,
split_tensor_by_cp_batch_plan,
split_in_seq_cp_local_pair,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -407,6 +414,135 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(owner, 0)
self.assertEqual(local_offset, 7)
def test_batch_page_aligned_plan_keeps_request_boundaries_and_last_token_owners(
self,
):
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4, 9],
prefix_lens=[0, 8],
page_size=4,
cp_size=2,
cp_rank=0,
)
self.assertEqual(plan.batch_size, 2)
self.assertEqual(plan.request_split_lists, [[4, 0, 0, 0], [4, 4, 1, 0]])
self.assertEqual(plan.request_padded_pages, [1, 3])
self.assertEqual(plan.request_padded_tokens, [4, 12])
self.assertEqual(plan.request_token_offsets, [0, 4])
self.assertEqual(plan.request_padded_token_offsets, [0, 4])
self.assertEqual(plan.request_page_offsets, [0, 1])
self.assertEqual(plan.request_last_token_owner, [0, 1])
self.assertEqual(plan.request_last_token_local_offset, [3, 4])
self.assertEqual(plan.request_rank_local_tokens, [4, 4])
self.assertEqual(plan.request_rank_local_offsets, [0, 4])
self.assertEqual(plan.request_kv_len_prev, [4, 4])
self.assertEqual(plan.request_kv_len_next, [4, 9])
self.assertEqual(plan.request_actual_seq_q_prev, [4, 4])
self.assertEqual(plan.request_actual_seq_q_next, [0, 0])
self.assertEqual(plan.flat_segment_request_ids, [0, 0, 0, 0, 1, 1, 1, 1])
self.assertEqual(plan.flat_segment_offsets, [0, 4, 4, 4, 0, 4, 8, 9])
def test_batch_plan_stable_helpers_split_and_build_page_owner_plan(self):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4, 9],
prefix_lens=[0, 8],
page_size=4,
cp_size=2,
cp_rank=1,
)
forward_batch = SimpleNamespace(nsa_cp_metadata=SimpleNamespace(batch_plan=plan))
self.assertIs(get_cp_shared_kv_batch_plan(forward_batch), plan)
self.assertEqual(build_flat_page_owner_plan(plan), [0, 0, 1, 1])
local_1d = split_tensor_by_cp_batch_plan(torch.arange(13), plan, mode="1d")
self.assertEqual(local_1d.tolist(), [8, 9, 10, 11, 12])
local_data = split_tensor_by_cp_batch_plan(
torch.arange(13 * 2).view(13, 2), plan, mode="data"
)
self.assertEqual(local_data[:, 0].tolist(), list(range(16, 26, 2)))
def test_collect_last_token_hidden_uses_batch_owner_metadata(self):
import torch
hidden_states = torch.tensor(
[[10.0], [11.0], [12.0], [13.0], [20.0], [21.0], [22.0], [23.0]]
)
forward_batch = SimpleNamespace(
extend_seq_lens_cpu=[4, 9],
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=2,
request_last_token_owner=[0, 1],
request_last_token_local_offset=[3, 4],
request_rank_local_offsets=[0, 4],
),
)
def fake_all_gather(output, local_last):
self.assertEqual(local_last.tolist(), [[13.0], [0.0]])
output.copy_(torch.tensor([[13.0], [0.0], [0.0], [99.0]]))
with (
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
return_value=0,
),
patch(
"sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor",
side_effect=fake_all_gather,
),
):
collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2)
self.assertEqual(collected.tolist(), [[13.0], [99.0]])
def test_collect_last_token_hidden_fails_fast_without_batch_owner_metadata(self):
import torch
forward_batch = SimpleNamespace(
extend_seq_lens_cpu=[4, 9],
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2),
)
with (
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
return_value=0,
),
self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_missing_last_token_metadata\]",
),
):
cp_collect_last_token_hidden(torch.zeros((8, 1)), forward_batch, 2)
def test_full_rerange_fails_fast_for_batch_metadata(self):
import torch
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2)
)
with self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_full_rerange_unsupported\]",
):
cp_all_gather_rerange_output(
torch.zeros((8, 1)), 2, forward_batch, stream=None
)
def test_local_pair_split_uses_metadata_lengths_not_half_split(self):
import torch
@@ -438,6 +574,47 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local_locs.tolist(), [2, 3, 12, 13])
def test_cp_split_and_rebuild_data_keeps_batch_request_boundaries(self):
import torch
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=2,
request_extend_lens=[4, 9],
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
request_zigzag_indices=[[0, 3], [0, 3]],
)
)
tensor = torch.arange(13 * 2).view(13, 2)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local = cp_split_and_rebuild_data(forward_batch, tensor)
self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2)))
def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self):
import torch
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=2,
request_extend_lens=[4, 9],
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
request_zigzag_indices=[[1, 2], [1, 2]],
)
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local = cp_split_and_rebuild_1d(forward_batch, torch.arange(13))
self.assertEqual(local.tolist(), [8, 9, 10, 11, 12])
def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
from types import SimpleNamespace