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

Reduce CP shared-KV batch overhead without reverting bs1 planning

Browse Source 
CP shared-KV bs>1 exposed two separate overhead sources: HiCache load-back could thrash near capacity, and partial-current sync compose could all-reduce row-major page-table gaps between request prefixes. Keep the intended batch-plan path for bs=1, but make host/L1 free-room handling less reactive and teach the MLA/index sync compose path to use exact per-request prefix slot spans instead of one bounding span.\n\nThe exact-span path preserves the single-span IPC fast path for bs=1/single-span cases, while avoiding over-communication for heterogeneous cache-hit batches. The HiCache metadata tests cover host/L1 free-room propagation and load-back batching behavior.\n\nConstraint: bs=1 using CPSharedKVBatchPlan is the expected steady-state path and must not be treated as a regression.\nConstraint: Remote production-like validation runs inside g0034 container /sgl-workspace/sglang-tai.\nRejected: Disable batch-plan for bs=1 | user confirmed this is intended behavior and it would hide the actual bs>1 overhead.\nRejected: Keep one bounding prefix span for bs>1 | row-major page tables can include current/gap slots and inflate per-layer all-reduce work.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not replace exact prefix spans with a single row-major bounding span unless ETE data proves collective launch count dominates gap over-communication.\nTested: g0034 docker py_compile for changed Python/test files.\nTested: g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py — 241 passed, 5 warnings, 2 subtests passed.\nNot-tested: Full ETE bs>1 throughput after the exact-prefix-span change.\nNot-tested: CUDA kernel benchmark and live traffic replay.
This commit is contained in:
laoyao0822
2026-06-06 02:03:16 +08:00
parent 1b99de7459
commit 81bfd4bfd3
10 changed files with 544 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
+39
View File
@@ -359,6 +359,45 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase):
self.assertEqual(target.value.tolist(), loaded.tolist())
self.assertIn(target.id, cache.ongoing_load_back)
def test_load_back_does_not_synchronously_evict_for_l1_free_room_when_exact_capacity_fits(self):
allocator = _make_allocator()
# Owner lane 0 has exactly enough capacity for the load-back target, but
# not enough to satisfy the configured free-room target. Load-back must
# not synchronously evict just to refill free room; doing so puts heavy
# eviction planning on the scheduler hot path.
allocator.free_pages = torch.tensor([1], dtype=torch.int64)
cache = _make_cache(allocator)
cache.hicache_l1_free_room_ratio = 0.5
cache.hicache_l1_free_room_trigger_ratio = 0.25
victim = _make_node(
22,
220,
[0],
value=torch.tensor([8, 9, 10, 11], dtype=torch.int64),
priority=0,
)
_attach_child(cache, cache.root_node, victim)
cache.evictable_leaves.add(victim)
cache.evictable_size_ = len(victim.key)
target = _make_node(23, 320, [0], value=None, priority=10)
_attach_child(cache, cache.root_node, target)
plan = cache._build_cp_load_back_plan([target], node_id=target.id)
self.assertEqual(plan.deficit_by_owner, [0, 0, 0, 0])
# Existing free-room accounting reports all lanes below the watermark.
# Load-back keeps that signal advisory and does not synchronously act on it.
self.assertEqual(plan.free_room_deficit_by_owner, [2, 2, 2, 2])
loaded = cache.load_back(target, mem_quota=100)
self.assertIsNotNone(loaded)
self.assertEqual(cache.cache_controller.load_calls, 1)
self.assertEqual(cache.cache_controller.evicted_device_indices, [])
self.assertEqual(target.value.tolist(), loaded.tolist())
self.assertIn(target.id, cache.ongoing_load_back)
def test_owner_lane_evict_params_choose_deficit_contributing_victim(self):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+5 -2
View File
@@ -2839,7 +2839,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
class TestHiRadixCacheCPLoadBack(CustomTestCase):
def test_cp_load_back_plan_uses_l1_free_room_target(self):
def test_cp_load_back_plan_reports_l1_free_room_without_blocking_exact_fit(self):
class FreeRoomAllocator:
def compute_owner_lane_stats(self, _page_owners):
return [1, 0], [0, 8], [1, 0]
@@ -2869,8 +2869,11 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertEqual(plan.required_by_owner, [1, 0])
self.assertEqual(plan.available_by_owner, [0, 8])
# Synchronous load-back admission must only block on exact capacity.
self.assertEqual(plan.deficit_by_owner, [1, 0])
# The free-room target is still reported for observability/proactive policy.
# required=1 page, available=0, target_room=ceil(8*0.5)=4 pages.
self.assertEqual(plan.deficit_by_owner, [5, 0])
self.assertEqual(plan.free_room_deficit_by_owner, [5, 0])
def test_cp_load_back_uses_host_len_not_host_value(self):
cache = HiRadixCache.__new__(HiRadixCache)
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+113 -11
View File
@@ -749,7 +749,7 @@ 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):
def test_mla_partial_current_sync_uses_batch_prefix_slot_spans(self):
from pathlib import Path
source = (
@@ -762,8 +762,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
branch_source = source[branch_start:branch_end]
local_locs_source = source[method_start:branch_end]
self.assertIn("build_batch_prefix_slot_span", source)
self.assertIn("prefix_slot_span=", branch_source)
self.assertIn("build_batch_prefix_slot_spans", source)
self.assertIn("prefix_slot_spans=", branch_source)
self.assertIn("get_cp_shared_kv_local_out_cache_loc", local_locs_source)
self.assertNotIn("current_locs = forward_batch.out_cache_loc", branch_source)
self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)
@@ -1372,7 +1372,7 @@ 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):
def test_batch_prefix_slot_span_covers_bounding_prefix_range(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
logical_pages = torch.tensor(
@@ -1408,6 +1408,42 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
(0, 0),
)
def test_batch_prefix_slot_spans_keep_request_prefix_ranges_exact(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_spans(
logical_pages=logical_pages,
prefix_lens_cpu=[8, 4],
page_size=4,
),
[(0, 2), (3, 4)],
)
self.assertEqual(
runtime.build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=[0, 4],
page_size=4,
),
[(3, 4)],
)
self.assertEqual(
runtime.build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=[0, 0],
page_size=4,
),
[],
)
def test_batch_current_slot_spans_follow_prefix_and_extend_pages(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -1468,7 +1504,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
)
prefix_slot_span = runtime.build_batch_prefix_slot_span(
prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
logical_pages=remap_logical_pages,
prefix_lens_cpu=[8, 4],
page_size=page_size,
@@ -1493,7 +1529,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_span=prefix_slot_span,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
)
)
@@ -1505,6 +1541,72 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
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_spans_do_not_reduce_row_gaps(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_spans = runtime.build_batch_prefix_slot_spans(
logical_pages=remap_logical_pages,
prefix_lens_cpu=[8, 4],
page_size=page_size,
)
current_slot_spans = runtime.build_batch_current_slot_spans(
logical_pages=remap_logical_pages,
prefix_lens_cpu=[8, 4],
extend_lens_cpu=[2, 2],
page_size=page_size,
)
reduced_ranges = []
def record_all_reduce(buffer, cp_size, start, end, **kwargs):
reduced_ranges.append((start, end))
return buffer
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", record_all_reduce
):
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_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
)
self.assertEqual(reduced_ranges[:2], [(4, 12), (16, 20)])
self.assertNotIn((4, 20), reduced_ranges)
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
@@ -1524,7 +1626,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
logical_pages,
layout,
)
prefix_slot_span = runtime.build_batch_prefix_slot_span(
prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=[8, 4],
page_size=page_size,
@@ -1555,7 +1657,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=0,
prefix_slot_span=prefix_slot_span,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
layer_id=2,
)
@@ -1583,7 +1685,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
)
def test_index_partial_current_sync_uses_batch_prefix_slot_span(self):
def test_index_partial_current_sync_uses_batch_prefix_slot_spans(self):
from pathlib import Path
source = (
@@ -1598,8 +1700,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
branch_source = source[branch_start:branch_end]
branch_compact = "".join(branch_source.split())
self.assertIn("build_batch_prefix_slot_span", source)
self.assertIn("prefix_slot_span=", branch_source)
self.assertIn("build_batch_prefix_slot_spans", source)
self.assertIn("prefix_slot_spans=", branch_source)
self.assertIn("get_cp_shared_kv_local_out_cache_loc", branch_source)
self.assertNotIn("current_locs=forward_batch.out_cache_loc", branch_compact)
self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source)