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

Enable batched CP shared-KV current reuse correctness

Browse Source 
Batch-size > 1 cache-hit traffic must not lose the current/partial-current reuse fast path. This change extends the target MLA/index sync-correct path to validate batched current suffix rows, compose page-aligned prefix spans, and route batched index top-k through current-only or partial-current reuse without falling back to scalar guards.\n\nThe implementation keeps page as the minimum cache unit: prefix cache coverage is page-aligned, while current suffix rows are sliced by valid extend lengths so padded tail rows are not exposed to attention. The index top-k batch path now mirrors the single-request contract: current-only reuses current index KV directly, partial-current materializes once and shares the dense buffer across request segments.\n\nConstraint: CP shared-KV supports in-seq zigzag only; round-robin remains fail-fast for shared-KV.\nConstraint: No new collectives are introduced; this is a sync correctness path, not a new communication scheme.\nRejected: Keep bs>1 current_index_kv fail-fast | disables the cache-hit path that W4 is meant to restore.\nRejected: Pad requests to a common batch length | violates the page-granular contract and exposes padded tails to consumers.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce batch_size_not_one or batch_gt1 current-reuse guards without proving an equivalent batched fast path exists.\nTested: Local py_compile for touched runtime/test files.\nTested: Remote g0034 pytest test_nsa_cp_utils.py test_cp_shared_kv_runtime.py => 157 passed, 5 warnings, 2 subtests passed.\nTested: Remote g0034 pytest test_cp_shared_kv_layout.py => 27 passed, 3 warnings.\nNot-tested: Full ETE with live traffic; CUDA kernel-backed batched descriptors; L1 prefetch, HiCache load/backup, and draft/EAGLE bs>1 current reuse.
This commit is contained in:
laoyao0822
2026-06-03 07:06:21 +08:00
parent a7472c415f
commit 50d0008705
9 changed files with 1195 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+255 -2
View File
@@ -565,6 +565,62 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64)
self.assertFalse(can_reuse_current_extend_kv(forward_batch))
def test_can_reuse_current_extend_kv_allows_batch_partial_cache_hits(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
can_reuse_current_extend_kv,
)
forward_batch = SimpleNamespace(
forward_mode=_FakeExtendForwardMode(),
batch_size=2,
extend_prefix_lens_cpu=[40384, 8192],
extend_seq_lens_cpu=[65, 128],
seq_lens_cpu=torch.tensor([40384 + 65, 8192 + 128], dtype=torch.int32),
out_cache_loc=torch.arange(256, dtype=torch.int64),
)
self.assertTrue(can_reuse_current_extend_kv(forward_batch))
forward_batch.out_cache_loc = torch.arange(192, dtype=torch.int64)
self.assertFalse(can_reuse_current_extend_kv(forward_batch))
forward_batch.out_cache_loc = torch.arange(256, dtype=torch.int64)
forward_batch.seq_lens_cpu = torch.tensor([40384 + 64, 8192 + 128], dtype=torch.int32)
self.assertFalse(can_reuse_current_extend_kv(forward_batch))
def test_current_extend_kv_rows_for_reuse_accepts_batch_bf16_and_fp8_payloads(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
forward_batch = SimpleNamespace(
forward_mode=_FakeExtendForwardMode(),
batch_size=2,
extend_prefix_lens_cpu=[0, 1024],
extend_seq_lens_cpu=[65, 128],
seq_lens_cpu=torch.tensor([65, 1024 + 128], dtype=torch.int32),
out_cache_loc=torch.arange(256, dtype=torch.int64),
)
bf16_k = torch.empty((193, 2, 4), dtype=torch.bfloat16)
fp8_like_k = torch.empty((193, 5), dtype=torch.uint8)
with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
self.assertEqual(
runtime.current_extend_kv_rows_for_reuse(
forward_batch,
bf16_k,
fp8_like_k,
),
193,
)
self.assertIsNone(
runtime.current_extend_kv_rows_for_reuse(
forward_batch,
bf16_k[:192],
fp8_like_k,
)
)
def test_should_reuse_current_extend_kv_enables_draft_partial_cache_hit_suffix(
self,
):
@@ -692,6 +748,21 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertIn("prefix_pages=0", branch_source)
self.assertNotIn("kv_cache = current_kv_cache", branch_source)
def test_mla_partial_current_sync_uses_batch_prefix_slot_span(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa_backend.py"
).read_text()
branch_start = source.index("eagle_draft_mla_branch = \"partial_current_sync\"")
branch_end = source.index("forward_partial_current_sync_compose", branch_start)
branch_source = source[branch_start:branch_end]
self.assertIn("build_batch_prefix_slot_span", source)
self.assertIn("prefix_slot_span=", branch_source)
self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
def test_nsa_backend_topk_transform_uses_effective_forward_impl(self):
from sglang.srt.layers.attention.nsa_backend import (
NativeSparseAttnBackend,
@@ -942,7 +1013,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
batch_size=2,
extend_prefix_lens_cpu=[64, 64],
extend_seq_lens_cpu=[64, 64],
seq_lens_cpu=torch.tensor([128, 128], dtype=torch.int32),
seq_lens_cpu=torch.tensor([129, 128], dtype=torch.int32),
out_cache_loc=torch.arange(128, dtype=torch.int64),
)
@@ -952,7 +1023,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
joined = "\n".join(logs.output)
self.assertIn("[CP_SHARED_KV_FALLBACK][current_reuse]", joined)
self.assertIn("batch_size_not_one", joined)
self.assertIn("prefix_extend_seq_len_mismatch_req_0", joined)
def test_tai_index_mqa_prepare_fast_path_miss_logs_warning(self):
from sglang.srt.environ import envs
@@ -1210,6 +1281,188 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]])
def test_batch_prefix_slot_span_covers_request_prefix_pages_only(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
logical_pages = torch.tensor(
[
[1, 2, 5],
[9, 11, 0],
],
dtype=torch.int64,
)
self.assertEqual(
runtime.build_batch_prefix_slot_span(
logical_pages=logical_pages,
prefix_lens_cpu=[8, 4],
page_size=4,
),
(0, 4),
)
self.assertEqual(
runtime.build_batch_prefix_slot_span(
logical_pages=logical_pages,
prefix_lens_cpu=[0, 4],
page_size=4,
),
(3, 4),
)
self.assertEqual(
runtime.build_batch_prefix_slot_span(
logical_pages=logical_pages,
prefix_lens_cpu=[0, 0],
page_size=4,
),
(0, 0),
)
def test_materialize_batch_prefix_span_and_reuse_current_kv_page_slots(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
page_size = 4
layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0)
kv_cache = torch.arange(0, 64, dtype=torch.float32).view(64, 1, 1)
logical_locs = torch.tensor(
[
[4, 8, 20, 21, -1, -1],
[36, 44, 45, -1, -1, -1],
],
dtype=torch.int64,
)
current_locs = torch.tensor([20, 21, 44, 45], dtype=torch.int64)
current_kv = torch.arange(100, 104, dtype=torch.float32).view(4, 1, 1)
remap_logical_pages = torch.tensor(
[
[1, 2, 5],
[9, 11, 0],
],
dtype=torch.int64,
)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
prefix_slot_span = runtime.build_batch_prefix_slot_span(
logical_pages=remap_logical_pages,
prefix_lens_cpu=[8, 4],
page_size=page_size,
)
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
):
mixed_kv, mixed_locs = (
runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_kv_cache=current_kv,
current_locs=current_locs,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_span=prefix_slot_span,
)
)
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1], [16, 20, 21, -1, -1, -1]])
self.assertTrue(torch.equal(mixed_kv[4:8], kv_cache[4:8]))
self.assertTrue(torch.equal(mixed_kv[8:12], kv_cache[8:12]))
self.assertTrue(torch.equal(mixed_kv[16:20], kv_cache[36:40]))
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv[:2]))
self.assertTrue(torch.equal(mixed_kv[20:22], current_kv[2:]))
def test_materialize_batch_prefix_span_and_reuse_current_index_page_slots(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
page_size = 4
index_head_dim = 4
scale_bytes = 4
page_bytes = page_size * index_head_dim + page_size * scale_bytes
layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0)
page_buffer = torch.zeros((16, page_bytes), dtype=torch.uint8)
page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8)
page_buffer[2] = torch.arange(40, 40 + page_bytes, dtype=torch.uint8)
page_buffer[9] = torch.arange(70, 70 + page_bytes, dtype=torch.uint8)
logical_pages = torch.tensor([[1, 2, 5], [9, 11, 0]], dtype=torch.int64)
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
page_buffer,
logical_pages,
layout,
)
prefix_slot_span = runtime.build_batch_prefix_slot_span(
logical_pages=logical_pages,
prefix_lens_cpu=[8, 4],
page_size=page_size,
)
current_k = torch.tensor(
[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]],
dtype=torch.uint8,
)
current_scale = torch.tensor([[1.25], [2.5], [3.5], [4.5]], dtype=torch.float32)
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
):
dense_page_buffer, dense_pages = (
runtime.materialize_prefix_and_reuse_current_index_page_slots(
page_buffer=page_buffer,
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=torch.tensor([20, 21, 44, 45], dtype=torch.int64),
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=0,
prefix_slot_span=prefix_slot_span,
layer_id=2,
)
)
scale_offset = page_size * index_head_dim
self.assertEqual(dense_pages.tolist(), [[1, 2, 3], [4, 5, 0]])
self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1]))
self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[2]))
self.assertTrue(torch.equal(dense_page_buffer[4], page_buffer[9]))
self.assertTrue(torch.equal(dense_page_buffer[3, 0:4], current_k[0]))
self.assertTrue(torch.equal(dense_page_buffer[3, 4:8], current_k[1]))
self.assertTrue(torch.equal(dense_page_buffer[5, 0:4], current_k[2]))
self.assertTrue(torch.equal(dense_page_buffer[5, 4:8], current_k[3]))
expected_scale_bytes = (
current_scale[:2].contiguous().view(torch.uint8).reshape(2, scale_bytes)
)
self.assertTrue(
torch.equal(
dense_page_buffer[
3,
scale_offset : scale_offset + page_size * scale_bytes,
].view(page_size, scale_bytes)[:2],
expected_scale_bytes,
)
)
def test_index_partial_current_sync_uses_batch_prefix_slot_span(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
).read_text()
branch_start = source.index(" if current_index_kv is not None:")
branch_end = source.index(" return materialized, dense_pages", branch_start)
branch_source = source[branch_start:branch_end]
self.assertIn("build_batch_prefix_slot_span", source)
self.assertIn("prefix_slot_span=", branch_source)
self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_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