minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[HiCache]Asymmetric pool support direct backend (#28446)

Browse Source
This commit is contained in:
huangtingwei
2026-06-16 13:17:57 -07:00
committed by GitHub
parent c0a6c3ce66
commit 9b4432fe18
6 changed files with 228 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/test_kvcacheio_asymmetric.py
+72 -14
View File
@@ -10,8 +10,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
# These tests use AsymmetricMHATokenToKVPoolHost methods and let that class call
# the real sgl-kernel transfer ops. The asymmetric host pool is kernel-only;
# direct/page_first_direct is intentionally rejected in the CPU dispatch tests.
# the real sgl-kernel transfer ops.
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available(), reason="asymmetric host-pool tests require CUDA."
)
@@ -45,22 +44,27 @@ def fill_with_offset(tensor, offset):
tensor.copy_((data + offset).view_as(tensor))
def make_host_pool(dtype):
def make_host_pool(dtype, layout="page_first"):
host = AsymmetricMHATokenToKVPoolHost.__new__(AsymmetricMHATokenToKVPoolHost)
host.layout = "page_first"
host.layout = layout
host.page_size = PAGE_SIZE
host.page_num = TOTAL_ITEMS // PAGE_SIZE
host.layer_num = NUM_LAYERS
host.head_num = HEAD_NUM
host.head_dim = K_HEAD_DIM
host.v_head_dim = V_HEAD_DIM
host.dtype = dtype
if layout == "page_first":
k_dims = (TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, K_HEAD_DIM)
v_dims = (TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, V_HEAD_DIM)
elif layout == "page_first_direct":
k_dims = (host.page_num, NUM_LAYERS, PAGE_SIZE, HEAD_NUM, K_HEAD_DIM)
v_dims = (host.page_num, NUM_LAYERS, PAGE_SIZE, HEAD_NUM, V_HEAD_DIM)
else:
raise ValueError(f"Unsupported layout: {layout}")
host.kv_buffer = (
torch.zeros(
TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, K_HEAD_DIM, dtype=dtype
).pin_memory(),
torch.zeros(
TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, V_HEAD_DIM, dtype=dtype
).pin_memory(),
torch.zeros(k_dims, dtype=dtype).pin_memory(),
torch.zeros(v_dims, dtype=dtype).pin_memory(),
)
return host
@@ -90,14 +94,28 @@ def make_device_pool(dtype):
)
def _host_k_tokens(host, indices, layer_id):
if host.layout == "page_first":
return host.k_buffer[indices, layer_id]
pages = indices[::PAGE_SIZE] // PAGE_SIZE
return host.k_buffer[pages, layer_id].reshape(-1, HEAD_NUM, K_HEAD_DIM)
def _host_v_tokens(host, indices, layer_id):
if host.layout == "page_first":
return host.v_buffer[indices, layer_id]
pages = indices[::PAGE_SIZE] // PAGE_SIZE
return host.v_buffer[pages, layer_id].reshape(-1, HEAD_NUM, V_HEAD_DIM)
def assert_backup_matches_device(host, device_pool, host_indices_host, device_indices):
for layer_id in range(NUM_LAYERS):
torch.testing.assert_close(
host.k_buffer[host_indices_host, layer_id],
_host_k_tokens(host, host_indices_host, layer_id),
device_pool.k_buffer[layer_id][device_indices].cpu(),
)
torch.testing.assert_close(
host.v_buffer[host_indices_host, layer_id],
_host_v_tokens(host, host_indices_host, layer_id),
device_pool.v_buffer[layer_id][device_indices].cpu(),
)
@@ -106,11 +124,11 @@ def assert_load_matches_host(host, device_pool, host_indices_host, load_indices)
for layer_id in range(NUM_LAYERS):
torch.testing.assert_close(
device_pool.k_buffer[layer_id][load_indices],
host.k_buffer[host_indices_host, layer_id].to(DEVICE),
_host_k_tokens(host, host_indices_host, layer_id).to(DEVICE),
)
torch.testing.assert_close(
device_pool.v_buffer[layer_id][load_indices],
host.v_buffer[host_indices_host, layer_id].to(DEVICE),
_host_v_tokens(host, host_indices_host, layer_id).to(DEVICE),
)
@@ -149,5 +167,45 @@ def test_asymmetric_mha_kernel_page_first_roundtrip(dtype):
assert_load_matches_host(host, device_pool, host_indices_host, load_indices_host)
@pytest.mark.parametrize("dtype", DTYPES)
def test_asymmetric_mha_direct_page_first_direct_roundtrip(dtype):
# Covers D2H backup + H2D load through AsymmetricMHATokenToKVPoolHost using
# page_first_direct/direct. K and V are copied through separate direct calls
# because their per-token strides differ.
host = make_host_pool(dtype, layout="page_first_direct")
device_pool = make_device_pool(dtype)
direct_stream = torch.cuda.Stream()
device_pages = torch.tensor([1, 2, 3], dtype=torch.int64)
host_pages = torch.tensor([0, 1, 2], dtype=torch.int64)
load_pages = torch.tensor([4, 5, 6], dtype=torch.int64)
device_indices_host = token_indices_for_pages(device_pages)
host_indices_host = token_indices_for_pages(host_pages)
load_indices_host = token_indices_for_pages(load_pages)
with torch.cuda.stream(direct_stream):
host.backup_from_device_all_layer(
device_pool, host_indices_host, device_indices_host, io_backend="direct"
)
direct_stream.synchronize()
assert_backup_matches_device(
host, device_pool, host_indices_host, device_indices_host
)
with torch.cuda.stream(direct_stream):
for layer_id in range(NUM_LAYERS):
device_pool.k_buffer[layer_id].zero_()
device_pool.v_buffer[layer_id].zero_()
host.load_to_device_per_layer(
device_pool,
host_indices_host,
load_indices_host,
layer_id,
io_backend="direct",
)
direct_stream.synchronize()
assert_load_matches_host(host, device_pool, host_indices_host, load_indices_host)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/models_e2e/test_mimo_v2.py
+2 -2
View File
@@ -25,9 +25,9 @@ MIMO_V2_OTHER_ARGS = [
"--hicache-ratio",
"1.5",
"--hicache-mem-layout",
"page_first",
"page_first_direct",
"--hicache-io-backend",
"kernel",
"direct",
]
MIMO_V2_MTP_OTHER_ARGS = MIMO_V2_OTHER_ARGS + [
"--speculative-algorithm",
test/registered/models_e2e/test_mimo_v2_flash.py
+2 -2
View File
@@ -47,9 +47,9 @@ class TestMiMoV2Flash(GSM8KMixin, SpecDecodingMixin, DefaultServerBase):
"--hicache-ratio",
"1.5",
"--hicache-mem-layout",
"page_first",
"page_first_direct",
"--hicache-io-backend",
"kernel",
"direct",
]
bs_1_speed_thres = 170
test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py
+68 -11
View File
@@ -30,6 +30,22 @@ def _make_host(layout: str) -> AsymmetricMHATokenToKVPoolHost:
if layout == "page_first":
k_dims = (8, host.layer_num, host.head_num, host.head_dim)
v_dims = (8, host.layer_num, host.head_num, host.v_head_dim)
elif layout == "page_first_direct":
host.page_num = 4
k_dims = (
host.page_num,
host.layer_num,
host.page_size,
host.head_num,
host.head_dim,
)
v_dims = (
host.page_num,
host.layer_num,
host.page_size,
host.head_num,
host.v_head_dim,
)
else:
raise ValueError(f"Unsupported test layout: {layout}")
@@ -121,15 +137,18 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase):
self.assertEqual(v_call.kwargs["item_size"], 24)
self.assertEqual(v_call.kwargs["dst_layout_dim"], 72)
def test_direct_load_is_rejected(self):
# Direct single-buffer D2H is not reliable for asymmetric K/V in the
# current sgl-kernel fast path, so the asymmetric host pool is kernel-only.
host = _make_host("page_first")
def test_direct_load_splits_k_and_v_for_page_first_direct(self):
# Direct kernels derive copy size from each call's first tensor, so K/V
# must be dispatched separately when their head dims differ.
host = _make_host("page_first_direct")
device_pool = _make_device_pool(host)
host_indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
device_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
with self.assertRaisesRegex(ValueError, "expected 'kernel'"):
with mock.patch(
"sglang.srt.mem_cache.memory_pool_host.transfer_kv_per_layer_direct_pf_lf",
create=True,
) as transfer:
host.load_to_device_per_layer(
device_pool,
host_indices,
@@ -138,17 +157,55 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase):
io_backend="direct",
)
def test_direct_backup_is_rejected(self):
# Same restriction for D2H backup: asymmetric MHA uses the kernel path
# until the direct kernel has an explicit safe asymmetric mode.
self.assertEqual(transfer.call_count, 2)
k_call, v_call = transfer.call_args_list
self.assertEqual(len(k_call.kwargs["src_ptrs"]), 1)
self.assertEqual(len(k_call.kwargs["dst_ptrs"]), 1)
self.assertIs(k_call.kwargs["src_ptrs"][0], host.k_buffer)
self.assertIs(k_call.kwargs["dst_ptrs"][0], device_pool.k_buffer[2])
self.assertEqual(len(v_call.kwargs["src_ptrs"]), 1)
self.assertEqual(len(v_call.kwargs["dst_ptrs"]), 1)
self.assertIs(v_call.kwargs["src_ptrs"][0], host.v_buffer)
self.assertIs(v_call.kwargs["dst_ptrs"][0], device_pool.v_buffer[2])
def test_direct_backup_splits_k_and_v_for_page_first_direct(self):
host = _make_host("page_first_direct")
device_pool = _make_device_pool(host)
host_indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
device_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
with mock.patch(
"sglang.srt.mem_cache.memory_pool_host.transfer_kv_all_layer_direct_lf_pf",
create=True,
) as transfer:
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="direct"
)
self.assertEqual(transfer.call_count, 2)
k_call, v_call = transfer.call_args_list
self.assertEqual(len(k_call.kwargs["src_ptrs"]), host.layer_num)
self.assertEqual(len(k_call.kwargs["dst_ptrs"]), 1)
self.assertIs(k_call.kwargs["src_ptrs"][0], device_pool.k_buffer[0])
self.assertIs(k_call.kwargs["dst_ptrs"][0], host.k_buffer)
self.assertEqual(len(v_call.kwargs["src_ptrs"]), host.layer_num)
self.assertEqual(len(v_call.kwargs["dst_ptrs"]), 1)
self.assertIs(v_call.kwargs["src_ptrs"][0], device_pool.v_buffer[0])
self.assertIs(v_call.kwargs["dst_ptrs"][0], host.v_buffer)
def test_direct_requires_page_first_direct_layout(self):
host = _make_host("page_first")
device_pool = _make_device_pool(host)
host_indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
device_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
with self.assertRaisesRegex(ValueError, "expected 'kernel'"):
host.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="direct"
with self.assertRaisesRegex(ValueError, "expected 'page_first_direct'"):
host.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)