Expand prefill CP KV capacity by sharding persistent NSA KV
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer. Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary. Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization Constraint: Decode side remains non-CP and receives full KV through Mooncake Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode Confidence: medium Scope-risk: broad Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics Tested: Remote py_compile for shared-KV touched Python files in g0034 container Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected Not-tested: Full GLM5 multi-node throughput/regression run after final doc update Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
This commit is contained in:
laoyao0822
@@ -0,0 +1,61 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
filter_kv_pages_for_cp_shared_kv,
|
||||
select_pages_by_request_positions,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestCPSharedKVTransferMapping(unittest.TestCase):
|
||||
def test_filter_kv_pages_for_cp_shared_kv(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
logical_pages = [5, 6, 7, 8, 9, 10, 11]
|
||||
src_pages, positions = filter_kv_pages_for_cp_shared_kv(
|
||||
layout=layout,
|
||||
logical_pages=logical_pages,
|
||||
chunk_page_start=100,
|
||||
)
|
||||
self.assertEqual(src_pages.tolist(), [2, 3])
|
||||
self.assertEqual(positions.tolist(), [102, 106])
|
||||
|
||||
def test_filter_kv_pages_for_cp_shared_kv_rejects_negative_pages(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"CP shared KV transfer got negative logical_pages",
|
||||
):
|
||||
filter_kv_pages_for_cp_shared_kv(
|
||||
layout=layout,
|
||||
logical_pages=np.array([5, -1, 7], dtype=np.int32),
|
||||
chunk_page_start=100,
|
||||
)
|
||||
|
||||
def test_filter_nsa_state_pages_uses_physical_source_and_logical_destination(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
logical_state_pages = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.int32)
|
||||
decode_state_pages = np.array([41, 42, 43, 44, 45, 46, 47, 48, 49], dtype=np.int32)
|
||||
|
||||
prefill_physical_pages, logical_positions = filter_kv_pages_for_cp_shared_kv(
|
||||
layout=layout,
|
||||
logical_pages=logical_state_pages,
|
||||
chunk_page_start=0,
|
||||
)
|
||||
selected_decode_pages = select_pages_by_request_positions(
|
||||
decode_state_pages,
|
||||
logical_positions,
|
||||
)
|
||||
|
||||
self.assertEqual(prefill_physical_pages.tolist(), [1, 2])
|
||||
self.assertEqual(logical_positions.tolist(), [2, 6])
|
||||
self.assertEqual(selected_decode_pages.tolist(), [43, 47])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,41 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
_get_in_seq_last_token_owner_and_offset,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestNSAInSeqCPUtils(unittest.TestCase):
|
||||
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
|
||||
# segment 2, not in rank 0's trailing padded segment.
|
||||
split_list = [4] * 16
|
||||
|
||||
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
|
||||
split_list=split_list,
|
||||
cp_size=8,
|
||||
actual_token_count=11,
|
||||
)
|
||||
|
||||
self.assertEqual(owner, 2)
|
||||
self.assertEqual(local_offset, 2)
|
||||
|
||||
def test_last_token_owner_keeps_existing_unpadded_fast_path_location(self):
|
||||
split_list = [4] * 16
|
||||
|
||||
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
|
||||
split_list=split_list,
|
||||
cp_size=8,
|
||||
actual_token_count=64,
|
||||
)
|
||||
|
||||
self.assertEqual(owner, 0)
|
||||
self.assertEqual(local_offset, 7)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,210 @@
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa_backend import (
|
||||
NSAIndexerMetadata,
|
||||
NSAMetadata,
|
||||
TopkTransformMethod,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestNSATopkTransform(unittest.TestCase):
|
||||
def test_paged_topk_transform_raises_when_fused_output_is_not_from_page_table(self):
|
||||
page_table = torch.tensor(
|
||||
[
|
||||
[10, 11, 12, 13, 14, 15, 16],
|
||||
[20, 21, 22, 23, 24, 25, 26],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
lengths_seen = {}
|
||||
|
||||
def fake_fast_topk_transform_fused(**kwargs):
|
||||
lengths_seen["value"] = kwargs["lengths"].clone()
|
||||
return torch.tensor(
|
||||
[
|
||||
[10, 1_039_799_618, 11, 12],
|
||||
[20, 21, 22, 23],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
metadata = NSAMetadata(
|
||||
page_size=1,
|
||||
cache_seqlens_int32=torch.tensor([7], dtype=torch.int32),
|
||||
max_seq_len_q=1,
|
||||
max_seq_len_k=7,
|
||||
cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, 7], dtype=torch.int32),
|
||||
page_table_1=page_table,
|
||||
real_page_table=page_table,
|
||||
nsa_cache_seqlens_int32=torch.tensor([4, 7], dtype=torch.int32),
|
||||
nsa_cu_seqlens_q=torch.tensor([0, 1, 2], dtype=torch.int32),
|
||||
nsa_cu_seqlens_k=torch.tensor([0, 4, 11], dtype=torch.int32),
|
||||
nsa_extend_seq_lens_list=[2],
|
||||
nsa_seqlens_expanded=torch.tensor([4, 7], dtype=torch.int32),
|
||||
)
|
||||
indexer_metadata = NSAIndexerMetadata(
|
||||
attn_metadata=metadata,
|
||||
topk_transform_method=TopkTransformMethod.PAGED,
|
||||
validate_paged_topk=True,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"sgl_kernel.fast_topk_transform_fused",
|
||||
side_effect=fake_fast_topk_transform_fused,
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_NSA_FUSE_TOPK.get", return_value=True
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get", return_value=True
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"NSA PAGED fused topk_transform produced values outside page_table_1",
|
||||
):
|
||||
indexer_metadata.topk_transform(
|
||||
logits=torch.zeros((2, 7), dtype=torch.float32),
|
||||
topk=4,
|
||||
cu_seqlens_q=torch.tensor([1, 1], dtype=torch.int32),
|
||||
)
|
||||
|
||||
self.assertEqual(lengths_seen["value"].tolist(), [4, 7])
|
||||
|
||||
def test_paged_topk_transform_rejects_lengths_exceeding_page_table_width(self):
|
||||
page_table = torch.tensor([[10, 11, 12]], dtype=torch.int32)
|
||||
metadata = NSAMetadata(
|
||||
page_size=1,
|
||||
cache_seqlens_int32=torch.tensor([3], dtype=torch.int32),
|
||||
max_seq_len_q=1,
|
||||
max_seq_len_k=3,
|
||||
cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, 3], dtype=torch.int32),
|
||||
page_table_1=page_table,
|
||||
real_page_table=page_table,
|
||||
nsa_cache_seqlens_int32=torch.tensor([4], dtype=torch.int32),
|
||||
nsa_cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
nsa_cu_seqlens_k=torch.tensor([0, 4], dtype=torch.int32),
|
||||
nsa_extend_seq_lens_list=[1],
|
||||
nsa_seqlens_expanded=torch.tensor([4], dtype=torch.int32),
|
||||
)
|
||||
indexer_metadata = NSAIndexerMetadata(
|
||||
attn_metadata=metadata,
|
||||
topk_transform_method=TopkTransformMethod.PAGED,
|
||||
validate_paged_topk=True,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"sgl_kernel.fast_topk_transform_fused",
|
||||
side_effect=AssertionError("fused kernel should not be called"),
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_NSA_FUSE_TOPK.get", return_value=True
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get", return_value=True
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"NSA PAGED fused topk lengths exceed page_table width",
|
||||
):
|
||||
indexer_metadata.topk_transform(
|
||||
logits=torch.zeros((1, 4), dtype=torch.float32),
|
||||
topk=4,
|
||||
)
|
||||
|
||||
def test_paged_topk_transform_skips_validation_during_cuda_graph_capture(self):
|
||||
page_table = torch.tensor([[10, 11, 12]], dtype=torch.int32)
|
||||
|
||||
def fake_fast_topk_transform_fused(**kwargs):
|
||||
return torch.tensor([[1_039_799_618]], dtype=torch.int32)
|
||||
|
||||
metadata = NSAMetadata(
|
||||
page_size=1,
|
||||
cache_seqlens_int32=torch.tensor([3], dtype=torch.int32),
|
||||
max_seq_len_q=1,
|
||||
max_seq_len_k=3,
|
||||
cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, 3], dtype=torch.int32),
|
||||
page_table_1=page_table,
|
||||
real_page_table=page_table,
|
||||
nsa_cache_seqlens_int32=torch.tensor([3], dtype=torch.int32),
|
||||
nsa_cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
nsa_cu_seqlens_k=torch.tensor([0, 3], dtype=torch.int32),
|
||||
nsa_extend_seq_lens_list=[1],
|
||||
nsa_seqlens_expanded=torch.tensor([3], dtype=torch.int32),
|
||||
)
|
||||
indexer_metadata = NSAIndexerMetadata(
|
||||
attn_metadata=metadata,
|
||||
topk_transform_method=TopkTransformMethod.PAGED,
|
||||
validate_paged_topk=True,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"sgl_kernel.fast_topk_transform_fused",
|
||||
side_effect=fake_fast_topk_transform_fused,
|
||||
), patch(
|
||||
"sglang.srt.layers.attention.nsa_backend._is_cuda_stream_capturing",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_NSA_FUSE_TOPK.get", return_value=True
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get", return_value=True
|
||||
):
|
||||
out = indexer_metadata.topk_transform(
|
||||
logits=torch.zeros((1, 3), dtype=torch.float32),
|
||||
topk=1,
|
||||
)
|
||||
|
||||
self.assertEqual(out.tolist(), [[1_039_799_618]])
|
||||
|
||||
def test_paged_topk_transform_skips_validation_when_cp_shared_debug_disabled(self):
|
||||
page_table = torch.tensor([[10, 11, 12]], dtype=torch.int32)
|
||||
|
||||
def fake_fast_topk_transform_fused(**kwargs):
|
||||
return torch.tensor([[1_039_799_618]], dtype=torch.int32)
|
||||
|
||||
metadata = NSAMetadata(
|
||||
page_size=1,
|
||||
cache_seqlens_int32=torch.tensor([3], dtype=torch.int32),
|
||||
max_seq_len_q=1,
|
||||
max_seq_len_k=3,
|
||||
cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, 3], dtype=torch.int32),
|
||||
page_table_1=page_table,
|
||||
real_page_table=page_table,
|
||||
nsa_cache_seqlens_int32=torch.tensor([3], dtype=torch.int32),
|
||||
nsa_cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32),
|
||||
nsa_cu_seqlens_k=torch.tensor([0, 3], dtype=torch.int32),
|
||||
nsa_extend_seq_lens_list=[1],
|
||||
nsa_seqlens_expanded=torch.tensor([3], dtype=torch.int32),
|
||||
)
|
||||
indexer_metadata = NSAIndexerMetadata(
|
||||
attn_metadata=metadata,
|
||||
topk_transform_method=TopkTransformMethod.PAGED,
|
||||
validate_paged_topk=True,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"sgl_kernel.fast_topk_transform_fused",
|
||||
side_effect=fake_fast_topk_transform_fused,
|
||||
), patch(
|
||||
"sglang.srt.layers.attention.nsa_backend._is_cuda_stream_capturing",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_NSA_FUSE_TOPK.get", return_value=True
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get", return_value=False
|
||||
):
|
||||
out = indexer_metadata.topk_transform(
|
||||
logits=torch.zeros((1, 3), dtype=torch.float32),
|
||||
topk=1,
|
||||
)
|
||||
|
||||
self.assertEqual(out.tolist(), [[1_039_799_618]])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,39 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.transform_index import (
|
||||
transform_index_page_table_decode_ref,
|
||||
transform_index_page_table_prefill_ref,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestNSATransformIndex(unittest.TestCase):
|
||||
def test_decode_ref_masks_high_out_of_bounds_topk_indices(self):
|
||||
page_table = torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32)
|
||||
topk_indices = torch.tensor([[0, 2, 3, -1], [1, 4, 0, -1]], dtype=torch.int64)
|
||||
|
||||
result = transform_index_page_table_decode_ref(page_table, topk_indices)
|
||||
|
||||
self.assertEqual(result.tolist(), [[10, 12, -1, -1], [21, -1, 20, -1]])
|
||||
|
||||
def test_prefill_ref_masks_high_out_of_bounds_topk_indices(self):
|
||||
page_table = torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32)
|
||||
topk_indices = torch.tensor(
|
||||
[[0, 3, -1], [2, 1, 9], [1, 0, -1]], dtype=torch.int64
|
||||
)
|
||||
|
||||
result = transform_index_page_table_prefill_ref(
|
||||
page_table=page_table,
|
||||
topk_indices=topk_indices,
|
||||
extend_lens_cpu=[2, 1],
|
||||
)
|
||||
|
||||
self.assertEqual(result.tolist(), [[10, -1, -1], [12, 11, -1], [21, 20, -1]])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,78 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestCpSharedKVLayout(unittest.TestCase):
|
||||
def test_page_owner_skips_dummy_page(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=0)
|
||||
pages = torch.tensor([-2, -1, 0, 1, 2, 3, 4, 5, 8, 9], dtype=torch.int64)
|
||||
owners = layout.owner_for_logical_pages(pages)
|
||||
self.assertEqual(owners.tolist(), [-1, -1, -1, 0, 1, 2, 3, 0, 3, 0])
|
||||
|
||||
def test_logical_to_physical_pages_keeps_dummy_zero(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=0)
|
||||
pages = torch.tensor([0, 1, 2, 3, 4, 5, 8, 9], dtype=torch.int64)
|
||||
physical = layout.logical_pages_to_physical(pages)
|
||||
self.assertEqual(physical.tolist(), [0, 1, 1, 1, 1, 2, 2, 3])
|
||||
|
||||
def test_owned_mask_and_loc_translation(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
locs = torch.tensor([0, 64, 128, 192, 256, 320, 384], dtype=torch.int64)
|
||||
mask = layout.owned_by_this_rank(locs)
|
||||
self.assertEqual(
|
||||
mask.tolist(), [False, False, False, True, False, False, False]
|
||||
)
|
||||
physical = layout.logical_locs_to_physical(locs[mask])
|
||||
self.assertEqual(physical.tolist(), [64])
|
||||
|
||||
def test_negative_and_dummy_locs_are_never_owned(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
locs = torch.tensor([-1, 0, 1, 63, 64, 128, 192], dtype=torch.int64)
|
||||
mask = layout.owned_by_this_rank(locs)
|
||||
self.assertEqual(
|
||||
mask.tolist(), [False, False, False, False, False, False, True]
|
||||
)
|
||||
|
||||
def test_numpy_filter_returns_request_absolute_positions(self):
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=1)
|
||||
logical_pages = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.int32)
|
||||
request_positions = np.arange(10, 19, dtype=np.int32)
|
||||
src_physical, positions = layout.filter_owned_pages_np(
|
||||
logical_pages, request_positions
|
||||
)
|
||||
self.assertEqual(src_physical.tolist(), [1, 2])
|
||||
self.assertEqual(positions.tolist(), [11, 15])
|
||||
|
||||
|
||||
class TestCPSharedPagedAllocator(unittest.TestCase):
|
||||
def test_shared_allocator_exposes_logical_capacity(self):
|
||||
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
|
||||
|
||||
allocator = CPSharedPagedTokenToKVPoolAllocator(
|
||||
logical_size=64 * 8,
|
||||
physical_size=64 * 2,
|
||||
page_size=64,
|
||||
dtype=torch.bfloat16,
|
||||
device="cpu",
|
||||
kvcache=None,
|
||||
need_sort=False,
|
||||
cp_size=4,
|
||||
cp_rank=0,
|
||||
)
|
||||
self.assertEqual(allocator.available_size(), 64 * 8)
|
||||
locs = allocator.alloc(64 * 2)
|
||||
self.assertEqual(locs.numel(), 64 * 2)
|
||||
self.assertEqual(allocator.available_size(), 64 * 6)
|
||||
allocator.free(locs)
|
||||
self.assertEqual(allocator.available_size(), 64 * 8)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,435 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
def test_all_reduce_uses_group_fast_path_for_float_buffers(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
class DummyGroup:
|
||||
device_group = object()
|
||||
|
||||
def __init__(self):
|
||||
self.called = False
|
||||
|
||||
def all_reduce(self, tensor):
|
||||
self.called = True
|
||||
|
||||
dummy_group = DummyGroup()
|
||||
buffer = torch.ones((2, 3), dtype=torch.float32)
|
||||
|
||||
with patch.object(
|
||||
runtime, "get_attention_cp_group", return_value=dummy_group
|
||||
), patch("torch.distributed.all_reduce") as dist_all_reduce:
|
||||
out = runtime._all_reduce_materialized_buffer(buffer, cp_size=2)
|
||||
|
||||
self.assertIs(out, buffer)
|
||||
self.assertTrue(dummy_group.called)
|
||||
dist_all_reduce.assert_not_called()
|
||||
|
||||
def test_all_reduce_copies_out_of_place_group_result_for_float_buffers(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
class DummyGroup:
|
||||
device_group = object()
|
||||
|
||||
def all_reduce(self, tensor):
|
||||
return tensor + 7
|
||||
|
||||
dummy_group = DummyGroup()
|
||||
buffer = torch.ones((2, 3), dtype=torch.float32)
|
||||
|
||||
with patch.object(
|
||||
runtime, "get_attention_cp_group", return_value=dummy_group
|
||||
), patch("torch.distributed.all_reduce") as dist_all_reduce:
|
||||
out = runtime._all_reduce_materialized_buffer(buffer, cp_size=2)
|
||||
|
||||
self.assertIs(out, buffer)
|
||||
self.assertTrue(torch.equal(buffer, torch.full((2, 3), 8.0)))
|
||||
dist_all_reduce.assert_not_called()
|
||||
|
||||
def test_all_reduce_bypasses_custom_group_for_byte_buffers(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
class DummyGroup:
|
||||
device_group = object()
|
||||
|
||||
def all_reduce(self, tensor):
|
||||
raise AssertionError("byte buffers must bypass custom all-reduce")
|
||||
|
||||
dummy_group = DummyGroup()
|
||||
buffer = torch.ones((2, 3), dtype=torch.uint8)
|
||||
|
||||
with patch.object(
|
||||
runtime, "get_attention_cp_group", return_value=dummy_group
|
||||
), patch("torch.distributed.all_reduce") as dist_all_reduce:
|
||||
out = runtime._all_reduce_materialized_buffer(buffer, cp_size=2)
|
||||
|
||||
self.assertIs(out, buffer)
|
||||
dist_all_reduce.assert_called_once()
|
||||
self.assertIs(dist_all_reduce.call_args.args[0], buffer)
|
||||
self.assertIs(dist_all_reduce.call_args.kwargs["group"], dummy_group.device_group)
|
||||
|
||||
def test_build_dense_page_remap_preserves_sentinels(self):
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
build_dense_page_remap,
|
||||
)
|
||||
|
||||
logical_pages = torch.tensor([0, 5, 2, -1, 5, 9, 2], dtype=torch.int32)
|
||||
unique_pages, dense_pages = build_dense_page_remap(logical_pages)
|
||||
|
||||
self.assertEqual(unique_pages.tolist(), [2, 5, 9])
|
||||
self.assertEqual(dense_pages.tolist(), [0, 2, 1, -1, 2, 3, 1])
|
||||
|
||||
def test_remap_logical_locs_to_dense_locs(self):
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
remap_logical_locs_to_dense_locs,
|
||||
)
|
||||
|
||||
logical_locs = torch.tensor([0, 1, 8, 9, -1, 20], dtype=torch.int32)
|
||||
unique_pages = torch.tensor([2, 5], dtype=torch.int32)
|
||||
|
||||
dense_locs = remap_logical_locs_to_dense_locs(
|
||||
logical_locs,
|
||||
unique_logical_pages=unique_pages,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_locs.tolist(), [0, 1, 4, 5, -1, 8])
|
||||
|
||||
def test_materialize_local_token_kv_pages(self):
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
build_dense_page_remap,
|
||||
materialize_local_token_kv_pages,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=4, cp_rank=1)
|
||||
kv_cache = torch.arange(0, 16 * 2, dtype=torch.float32).view(16, 1, 2)
|
||||
logical_pages = torch.tensor([1, 2, 5, 7, 10], dtype=torch.int32)
|
||||
unique_pages, _ = build_dense_page_remap(logical_pages)
|
||||
|
||||
dense_kv = materialize_local_token_kv_pages(
|
||||
kv_cache=kv_cache,
|
||||
unique_logical_pages=unique_pages,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(list(dense_kv.shape), [24, 1, 2])
|
||||
self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[4:8]))
|
||||
self.assertTrue(torch.equal(dense_kv[20:24], kv_cache[12:16]))
|
||||
self.assertEqual(float(dense_kv[4:8].abs().sum().item()), 0.0)
|
||||
|
||||
def test_materialize_local_paged_index_buffer(self):
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
build_dense_page_remap,
|
||||
materialize_local_paged_buffer,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=4, cp_rank=1)
|
||||
page_buffer = torch.arange(0, 5 * 3, dtype=torch.uint8).view(5, 3)
|
||||
logical_pages = torch.tensor([1, 2, 5, 7, 10], dtype=torch.int32)
|
||||
unique_pages, dense_pages = build_dense_page_remap(logical_pages)
|
||||
|
||||
dense_page_buffer = materialize_local_paged_buffer(
|
||||
page_buffer=page_buffer,
|
||||
unique_logical_pages=unique_pages,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_pages.tolist(), [1, 2, 3, 4, 5])
|
||||
self.assertEqual(list(dense_page_buffer.shape), [6, 3])
|
||||
self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[1]))
|
||||
self.assertTrue(torch.equal(dense_page_buffer[5], page_buffer[3]))
|
||||
self.assertEqual(int(dense_page_buffer[1].sum().item()), 0)
|
||||
|
||||
def test_filter_owned_logical_locs_for_persistent_writes(self):
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
filter_owned_logical_locs,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
logical_locs = torch.tensor(
|
||||
[0, 64, 128, 192, 256, 320, 384, 448, 449],
|
||||
dtype=torch.int64,
|
||||
)
|
||||
|
||||
owned_mask, physical_locs = filter_owned_logical_locs(logical_locs, layout)
|
||||
|
||||
self.assertEqual(
|
||||
owned_mask.tolist(),
|
||||
[False, False, False, True, False, False, False, True, True],
|
||||
)
|
||||
self.assertEqual(physical_locs.tolist(), [64, 128, 129])
|
||||
|
||||
def test_materialize_token_kv_raises_on_locs_outside_physical_pool(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=4, cp_rank=1)
|
||||
kv_cache = torch.arange(0, 12 * 2, dtype=torch.float32).view(12, 1, 2)
|
||||
logical_locs = torch.tensor([4, 8, 20, 40, -1], dtype=torch.int64)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=True
|
||||
), patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV materialize got logical token locs outside the physical pool",
|
||||
):
|
||||
runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
def test_materialize_token_kv_skips_physical_pool_validation_when_debug_disabled(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=4, cp_rank=1)
|
||||
kv_cache = torch.arange(0, 12 * 2, dtype=torch.float32).view(12, 1, 2)
|
||||
# Logical page 9 maps outside the local physical pool, but is owned by
|
||||
# rank 0. With debug disabled, this source-finding validation should
|
||||
# not synchronize or raise on rank 1.
|
||||
logical_locs = torch.tensor([4, 36, -1], dtype=torch.int64)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=False
|
||||
), patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x):
|
||||
_, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_locs.tolist(), [4, 8, -1])
|
||||
|
||||
def test_debug_materialize_token_kv_allows_minus_one_sentinel(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0)
|
||||
kv_cache = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1)
|
||||
logical_locs = torch.tensor([4, -1, 8], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=True):
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x):
|
||||
_, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_locs.tolist(), [4, -1, 8])
|
||||
|
||||
def test_debug_materialize_token_kv_rejects_below_minus_one_locs(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0)
|
||||
kv_cache = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1)
|
||||
logical_locs = torch.tensor([4, -2, 8], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=True):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV token materialize got logical_locs below -1",
|
||||
):
|
||||
runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
def test_debug_materialize_paged_buffer_rejects_negative_pages(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0)
|
||||
page_buffer = torch.arange(0, 4 * 3, dtype=torch.uint8).view(4, 3)
|
||||
logical_pages = torch.tensor([1, -1, 2], dtype=torch.int32)
|
||||
|
||||
with patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=True):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV paged materialize got negative logical_pages",
|
||||
):
|
||||
runtime.materialize_shared_paged_buffer(
|
||||
page_buffer=page_buffer,
|
||||
logical_pages=logical_pages,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
def test_debug_filter_owned_logical_locs_rejects_negative_write_locs(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=64, cp_size=4, cp_rank=2)
|
||||
logical_locs = torch.tensor([192, -1], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=True):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV persistent write got negative logical_locs",
|
||||
):
|
||||
runtime.filter_owned_logical_locs(logical_locs, layout)
|
||||
|
||||
def test_materialize_token_kv_uses_remap_source_for_target_subset(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0)
|
||||
kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1)
|
||||
logical_locs = torch.tensor([8, 20, -1], dtype=torch.int64)
|
||||
remap_source_locs = torch.tensor([4, 8, 12, 16, 20], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x):
|
||||
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
remap_logical_locs=remap_source_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_locs.tolist(), [8, 20, -1])
|
||||
self.assertEqual(list(dense_kv.shape), [24, 1, 1])
|
||||
self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[8:12]))
|
||||
self.assertTrue(torch.equal(dense_kv[20:24], kv_cache[20:24]))
|
||||
|
||||
def test_materialize_token_kv_keeps_dense_shape_for_shared_remap_source(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0)
|
||||
kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1)
|
||||
remap_source_locs = torch.tensor([4, 8, 12, 16, 20], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x):
|
||||
dense_kv_a, dense_locs_a = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=torch.tensor([4, 20], dtype=torch.int64),
|
||||
remap_logical_locs=remap_source_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
dense_kv_b, dense_locs_b = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=torch.tensor([8, 16], dtype=torch.int64),
|
||||
remap_logical_locs=remap_source_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(list(dense_kv_a.shape), [24, 1, 1])
|
||||
self.assertEqual(list(dense_kv_b.shape), [24, 1, 1])
|
||||
self.assertEqual(dense_locs_a.tolist(), [4, 20])
|
||||
self.assertEqual(dense_locs_b.tolist(), [8, 16])
|
||||
|
||||
def test_materialize_paged_buffer_raises_on_pages_outside_physical_pool(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=4, cp_rank=1)
|
||||
page_buffer = torch.arange(0, 3 * 4, dtype=torch.uint8).view(3, 4)
|
||||
logical_pages = torch.tensor([1, 2, 6, 10], dtype=torch.int32)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=True
|
||||
), patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV materialize got logical pages outside the physical page buffer",
|
||||
):
|
||||
runtime.materialize_shared_paged_buffer(
|
||||
page_buffer=page_buffer,
|
||||
logical_pages=logical_pages,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
|
||||
class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
|
||||
def test_mla_write_filter_does_not_build_debug_summaries_when_debug_disabled(self):
|
||||
from sglang.srt.layers.attention import nsa_backend
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
forward_mode="test",
|
||||
)
|
||||
cache_loc = torch.tensor([4, 8, 12, 16], dtype=torch.int64)
|
||||
k = torch.arange(0, 8, dtype=torch.float32).view(4, 1, 2)
|
||||
k_rope = torch.arange(0, 4, dtype=torch.float32).view(4, 1, 1)
|
||||
|
||||
with patch.object(
|
||||
nsa_backend, "tensor_debug_summary", side_effect=AssertionError
|
||||
), patch.object(
|
||||
nsa_backend, "tensor_debug_checksum", side_effect=AssertionError
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get",
|
||||
return_value=False,
|
||||
):
|
||||
physical_locs, k_to_write, k_rope_to_write = (
|
||||
nsa_backend.NativeSparseAttnBackend._maybe_filter_shared_mla_kv_write(
|
||||
None,
|
||||
forward_batch,
|
||||
cache_loc,
|
||||
k,
|
||||
k_rope,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(physical_locs.tolist(), [4, 8])
|
||||
self.assertEqual(k_to_write.shape[0], 2)
|
||||
self.assertEqual(k_rope_to_write.shape[0], 2)
|
||||
|
||||
def test_index_write_filter_does_not_build_debug_summaries_when_debug_disabled(self):
|
||||
from sglang.srt.layers.attention.nsa import nsa_indexer
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
forward_mode="test",
|
||||
out_cache_loc=torch.tensor([4, 8, 12, 16], dtype=torch.int64),
|
||||
)
|
||||
key = torch.arange(0, 8, dtype=torch.float32).view(4, 2)
|
||||
|
||||
with patch.object(
|
||||
nsa_indexer, "tensor_debug_summary", side_effect=AssertionError
|
||||
), patch.object(
|
||||
nsa_indexer, "tensor_debug_checksum", side_effect=AssertionError
|
||||
), patch(
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get",
|
||||
return_value=False,
|
||||
):
|
||||
physical_locs, key_to_write = nsa_indexer.Indexer._filter_shared_index_write(
|
||||
None,
|
||||
forward_batch,
|
||||
key,
|
||||
)
|
||||
|
||||
self.assertEqual(physical_locs.tolist(), [4, 8])
|
||||
self.assertEqual(key_to_write.shape[0], 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -34,6 +34,25 @@ class TestPrepareServerArgs(CustomTestCase):
|
||||
)
|
||||
|
||||
|
||||
def test_enable_nsa_prefill_cp_shared_kv_parser_flag():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
raw_args = parser.parse_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--enable-nsa-prefill-context-parallel",
|
||||
"--enable-nsa-prefill-cp-shared-kv",
|
||||
"--nsa-prefill-cp-mode",
|
||||
"in-seq-split",
|
||||
]
|
||||
)
|
||||
args = ServerArgs.from_cli_args(raw_args)
|
||||
assert args.enable_nsa_prefill_cp_shared_kv is True
|
||||
|
||||
|
||||
class TestLoadBalanceMethod(unittest.TestCase):
|
||||
def test_non_pd_defaults_to_round_robin(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")
|
||||
|
||||
Reference in New Issue
Block a user