minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
python/sglang/srt/mem_cache/memory_pool_host.py
+82 -23
View File
@@ -936,9 +936,8 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
``self.v_buffer``) instead of a single ``(2, ...)`` tensor, so each side
keeps its native stride. The kernel transfer path dispatches K and V as
independent single-buffer copies so each side uses its own ``item_size``.
Direct transfer and the flat-page L3 storage interface assume a single
shared ``item_size`` in paths that are not safe for asymmetric K/V, so they
raise instead of silently corrupting V copies.
K/V direct transfers must be dispatched separately because the direct
kernels derive copy sizes from each call's first tensor.
"""
def get_size_per_token(self):
@@ -960,10 +959,25 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
if self.layout == "page_first":
k_dims = (self.size, self.layer_num, self.head_num, self.head_dim)
v_dims = (self.size, self.layer_num, self.head_num, self.v_head_dim)
elif self.layout == "page_first_direct":
k_dims = (
self.page_num,
self.layer_num,
self.page_size,
self.head_num,
self.head_dim,
)
v_dims = (
self.page_num,
self.layer_num,
self.page_size,
self.head_num,
self.v_head_dim,
)
else:
raise ValueError(
f"Unsupported layout for models with head_dim != v_head_dim: "
f"{self.layout}; expected 'page_first'."
f"{self.layout}; expected 'page_first' or 'page_first_direct'."
)
# token_stride_size / layout_dim are intentionally NOT set: K and V
@@ -1039,10 +1053,33 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
item_size=self._v_token_stride_size(),
src_layout_dim=self._v_layout_dim(),
)
elif io_backend == "direct":
if self.layout != "page_first_direct":
raise ValueError(
f"Unsupported layout for models with head_dim != v_head_dim "
f"and io_backend='direct': {self.layout}; expected "
"'page_first_direct'."
)
transfer_kv_per_layer_direct_pf_lf(
src_ptrs=[self.k_buffer],
dst_ptrs=[device_pool.k_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=layer_id,
page_size=self.page_size,
)
transfer_kv_per_layer_direct_pf_lf(
src_ptrs=[self.v_buffer],
dst_ptrs=[device_pool.v_buffer[layer_id]],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=layer_id,
page_size=self.page_size,
)
else:
raise ValueError(
f"Unsupported IO backend for models with head_dim != v_head_dim: "
f"{io_backend}; expected 'kernel'."
f"{io_backend}; expected 'kernel' or 'direct'."
)
def backup_from_device_all_layer(
@@ -1072,10 +1109,31 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
dst_layout_dim=self._v_layout_dim(),
num_layers=self.layer_num,
)
elif io_backend == "direct":
if self.layout != "page_first_direct":
raise ValueError(
f"Unsupported layout for models with head_dim != v_head_dim "
f"and io_backend='direct': {self.layout}; expected "
"'page_first_direct'."
)
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.k_buffer,
dst_ptrs=[self.k_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=device_pool.v_buffer,
dst_ptrs=[self.v_buffer],
src_indices=device_indices,
dst_indices=host_indices,
page_size=self.page_size,
)
else:
raise ValueError(
f"Unsupported IO backend for models with head_dim != v_head_dim: "
f"{io_backend}; expected 'kernel'."
f"{io_backend}; expected 'kernel' or 'direct'."
)
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
@@ -1097,7 +1155,7 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
def get_page_buffer_meta(self, indices):
assert len(indices) % self.page_size == 0
if self.layout != "page_first":
if self.layout not in ("page_first", "page_first_direct"):
raise ValueError(
f"Unsupported layout for models with head_dim != v_head_dim: "
f"{self.layout}"
@@ -1121,29 +1179,30 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
)
ptr_list = []
element_size_list = []
if self.layout == "page_first_direct":
k_index_stride = (
self.layer_num * self.page_size * self.head_num * self.head_dim
)
v_index_stride = (
self.layer_num * self.page_size * self.head_num * self.v_head_dim
)
else:
k_index_stride = self.layer_num * self.head_num * self.head_dim
v_index_stride = self.layer_num * self.head_num * self.v_head_dim
for index in range(0, len(indices), self.page_size):
k_ptr = (
k_base_ptr
+ indices[index]
* self.layer_num
* self.head_num
* self.head_dim
* self.dtype.itemsize
)
v_ptr = (
v_base_ptr
+ indices[index]
* self.layer_num
* self.head_num
* self.v_head_dim
* self.dtype.itemsize
buffer_index = (
indices[index] // self.page_size
if self.layout == "page_first_direct"
else indices[index]
)
k_ptr = k_base_ptr + buffer_index * k_index_stride * self.dtype.itemsize
v_ptr = v_base_ptr + buffer_index * v_index_stride * self.dtype.itemsize
ptr_list.extend([k_ptr, v_ptr])
element_size_list.extend([k_element_size, v_element_size])
return ptr_list, element_size_list
def is_stride_page_aligned(self, page_size_bytes: int = 4096) -> bool:
if self.layout != "page_first":
if self.layout not in ("page_first", "page_first_direct"):
return False
k_stride = (
self.page_size
python/sglang/srt/server_args.py
+2 -15
View File
@@ -2440,21 +2440,8 @@ class ServerArgs:
)
# MiMoV2 has head_dim != v_head_dim, so the host KV pool uses
# asymmetric K/V allocation. Only the kernel/page_first transfer
# path has a safe split K/V implementation.
if self.hicache_io_backend != "kernel":
logger.warning(
f"Force hicache_io_backend to 'kernel' for MiMoV2 model "
f"(was {self.hicache_io_backend!r})."
)
self.hicache_io_backend = "kernel"
if self.hicache_mem_layout != "page_first":
logger.warning(
f"Force hicache_mem_layout to 'page_first' for "
f"MiMoV2 model (was {self.hicache_mem_layout!r}); "
f"asymmetric K/V HiCache requires kernel/page_first."
)
self.hicache_mem_layout = "page_first"
# asymmetric K/V allocation. Both kernel/page_first and
# direct/page_first_direct have split K/V transfer paths.
elif (
"Step3p5ForCausalLM" in model_arch
or "Step3p7ForConditionalGeneration" in model_arch
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",
)