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

Reuse prepared HiCache load descriptors across CP prefill layers

Browse Source 
CP shared-KV bs>1 cache-hit loads already merge request load ops, but the host pool still rebuilt layer-invariant mapping work from the same host/device indices. Introduce a PreparedLoadDescriptor lifecycle around begin/end load, wire MLA KV and NSA index H2D loads through tai-kernel prepared submit when available, and add timing hooks plus regression coverage for descriptor reuse and explicit fallback logging. Record the P4/P6b design and benchmark results in the advanced feature notes.

Constraint: Radix residency and allocator decisions remain synchronous; only the data-transfer descriptor is prepared for per-layer async submit.

Constraint: Production fast path must not silently fall back when tai prepared H2D support is missing.

Rejected: Cross-batch descriptor reuse | descriptor lifetime and tensor ownership are only safe within one load operation.

Rejected: Change L2->L1 scheduling to layer-ahead prefetch in this commit | that is a separate lifecycle change after descriptor reuse is stable.

Confidence: medium

Scope-risk: moderate

Directive: Keep LayerDoneCounter per-layer readiness semantics; do not replace with all-layer waits.

Tested: python -m py_compile python/sglang/srt/mem_cache/memory_pool_host.py python/sglang/srt/managers/cache_controller.py

Tested: Remote g0034:cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py (88 passed)

Tested: Remote tai-kernel prepared descriptor CUDA test (6 passed) and P4 benchmark full matrix (90 rows)

Not-tested: ETE replay/GSM8K cache-hit correctness after this commit

Not-tested: Layer-ahead L2->L1 prefetch scheduling

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-11 05:09:41 +08:00
co-authored by OmX
parent adf357b02c
commit 7284a469a2
4 changed files with 1375 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_hicache_controller_cp.py
+249 -3
View File
@@ -113,6 +113,7 @@ from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
PreparedLoadDescriptor,
)
from sglang.srt.mem_cache.radix_cache import TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
@@ -206,6 +207,207 @@ class DummyHostKVCacheForAlloc(HostKVCache):
pass
class TestPreparedLoadDescriptor(CustomTestCase):
def test_host_begin_load_builds_page_aligned_descriptor(self):
host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
host_pool.page_size = 4
host_pool.layout = "page_first_direct"
host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
device_indices = torch.tensor(
[40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
)
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
desc = host_pool._active_load_descriptor
self.assertIsInstance(desc, PreparedLoadDescriptor)
self.assertTrue(torch.equal(desc.host_indices, host_indices))
self.assertTrue(torch.equal(desc.device_indices, device_indices))
self.assertEqual(desc.num_tokens, 8)
self.assertEqual(desc.num_pages, 2)
self.assertEqual(desc.layout, "page_first_direct")
self.assertEqual(desc.io_backend, "direct")
self.assertEqual(desc.host_page_indices.tolist(), [2, 5])
self.assertEqual(desc.device_page_indices.tolist(), [10, 16])
host_pool.end_load_to_device_op()
self.assertIsNone(host_pool._active_load_descriptor)
def test_nsa_begin_load_attaches_indexer_pages_to_descriptor(self):
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.page_size = 4
host_pool.layout = "page_first_direct"
host_pool.index_active_layer_ids = (0, 2)
host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
device_indices = torch.tensor(
[40, 41, 42, 43, 64, 65, 66, 67], dtype=torch.int64
)
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
desc = host_pool._active_load_descriptor
self.assertEqual(desc.index_active_layer_ids, (0, 2))
self.assertEqual(desc.index_host_page_indices.tolist(), [2, 5])
self.assertEqual(desc.index_device_page_indices.tolist(), [10, 16])
self.assertIs(host_pool._active_load_indexer_page_indices[0], desc.index_host_page_indices)
self.assertIs(
host_pool._active_load_indexer_page_indices[1],
desc.index_device_page_indices,
)
host_pool.end_load_to_device_op()
self.assertIsNone(host_pool._active_load_descriptor)
self.assertIsNone(host_pool._active_load_indexer_page_indices)
def test_missing_direct_load_descriptor_warns_once(self):
host_pool = DummyHostKVCacheForAlloc.__new__(DummyHostKVCacheForAlloc)
host_pool.page_size = 4
host_pool.layout = "page_first_direct"
with self.assertLogs(
"sglang.srt.mem_cache.memory_pool_host", level="WARNING"
) as logs:
self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
self.assertIsNone(host_pool._get_active_load_descriptor("direct"))
warnings = [
line for line in logs.output if "missing_prepared_descriptor" in line
]
self.assertEqual(len(warnings), 1)
def test_mla_direct_load_uses_prepared_tai_descriptor_when_available(self):
calls = []
fake_desc = object()
def fake_prepare(src_indices, dst_indices, **kwargs):
calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
return fake_desc
def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))
def fake_destroy(desc):
calls.append(("destroy", desc))
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((3, 32, 1, 16), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
return_value=fake_prepare,
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
return_value=fake_submit,
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_destroy_h2d_page_descriptor",
return_value=fake_destroy,
),
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
host_pool.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
host_pool.end_load_to_device_op()
self.assertEqual(calls[0][0], "prepare")
self.assertEqual(calls[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(calls[0][2].tolist(), [12, 13, 14, 15])
self.assertEqual(calls[0][3], {"page_size": 4, "layout": "page_first_direct"})
self.assertEqual(calls[1][0], "submit")
self.assertIs(calls[1][1], fake_desc)
self.assertEqual(calls[1][2][0].data_ptr(), host_pool.kv_buffer.data_ptr())
self.assertEqual(
calls[1][3][0].data_ptr(), device_pool.kv_buffer[2].data_ptr()
)
self.assertEqual(calls[1][4], {"layer_id": 2})
self.assertEqual(calls[2], ("destroy", fake_desc))
def test_nsa_index_direct_load_uses_prepared_tai_index_descriptor(self):
calls = []
fake_base_desc = object()
fake_index_desc = object()
def fake_prepare(src_indices, dst_indices, **kwargs):
calls.append(("prepare", src_indices.clone(), dst_indices.clone(), kwargs))
if kwargs["page_size"] == 1:
return fake_index_desc
return fake_base_desc
def fake_submit(desc, src_ptrs, dst_ptrs, **kwargs):
calls.append(("submit", desc, src_ptrs, dst_ptrs, kwargs))
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_active_layer_ids = (0, 1, 2)
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), dtype=torch.uint8
)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with (
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
return_value=fake_prepare,
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_submit_h2d_layer",
return_value=fake_submit,
),
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
)
host_pool._load_indexer_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="direct",
)
self.assertEqual([c[0] for c in calls[:2]], ["prepare", "prepare"])
self.assertEqual(calls[1][1].tolist(), [1])
self.assertEqual(calls[1][2].tolist(), [3])
self.assertEqual(calls[1][3], {"page_size": 1, "layout": "page_first_direct"})
self.assertEqual(calls[2][0], "submit")
self.assertIs(calls[2][1], fake_index_desc)
self.assertEqual(
calls[2][2][0].data_ptr(), host_pool.index_k_with_scale_buffer.data_ptr()
)
self.assertEqual(
calls[2][3][0].data_ptr(),
device_pool.index_k_with_scale_buffer[1].data_ptr(),
)
self.assertEqual(calls[2][4], {"layer_id": 1})
class FakeDevicePool:
device = "cpu"
layer_num = 1
@@ -878,9 +1080,15 @@ class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
host_pool._get_indexer_page_indices = counting_getter
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
with (
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_prepare_h2d_page_descriptor",
side_effect=RuntimeError("missing prepared descriptor api"),
),
patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf",
return_value=fake_direct,
),
):
host_pool.begin_load_to_device_op(
host_indices, device_indices, io_backend="direct"
@@ -1898,6 +2106,44 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
self.assertEqual(host_pool.end_calls, 1)
self.assertEqual([load[2] for load in host_pool.loads], [0, 1, 2])
def test_start_loading_emits_descriptor_timing_when_enabled(self):
host_pool = FakeHostPool(
torch.tensor([100, 101, 102, 103], dtype=torch.int64)
)
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
allocator.device_pool = FakeDevicePool(layer_num=2)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
controller.load_cp([node], node_id=114)
timing_keys = []
def record_timing(key, start_time, message, *args):
timing_keys.append(key)
with patch(
"sglang.srt.managers.cache_controller.envs."
"SGLANG_CP_SHARED_KV_BS_GT1_TIMING.get",
return_value=True,
), patch(
"sglang.srt.managers.cache_controller."
"_cp_shared_kv_bs_gt1_cache_timing",
side_effect=record_timing,
):
controller.start_loading()
self.assertIn("prepare_load_descriptor", timing_keys)
self.assertIn("submit_h2d_layer_loop", timing_keys)
self.assertIn("submit_h2d_layer_per_call_slow", timing_keys)
self.assertIn("end_load_descriptor", timing_keys)
def test_cp_load_frees_unexpected_owner_allocator_length(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))