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

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -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
+    K/V direct transfers must be dispatched separately because the direct
-    shared ``item_size`` in paths that are not safe for asymmetric K/V, so they
+    kernels derive copy sizes from each call's first tensor.
-    raise instead of silently corrupting V copies.
     """
 
     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 = (
+            buffer_index = (
-                k_base_ptr
+                indices[index] // self.page_size
-                + indices[index]
+                if self.layout == "page_first_direct"
-                * self.layer_num
+                else indices[index]
-                * 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
             )
+            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

@@ -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
+                # asymmetric K/V allocation. Both kernel/page_first and
-                # path has a safe split K/V implementation.
+                # direct/page_first_direct have split K/V transfer paths.
-                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"
         elif (
             "Step3p5ForCausalLM" in model_arch
             or "Step3p7ForConditionalGeneration" in model_arch

test/registered/jit/test_kvcacheio_asymmetric.py

@@ -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;
+# the real sgl-kernel transfer ops.
-# direct/page_first_direct is intentionally rejected in the CPU dispatch tests.
 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(
+        torch.zeros(k_dims, dtype=dtype).pin_memory(),
-            TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, K_HEAD_DIM, dtype=dtype
+        torch.zeros(v_dims, dtype=dtype).pin_memory(),
-        ).pin_memory(),
-        torch.zeros(
-            TOTAL_ITEMS, NUM_LAYERS, HEAD_NUM, V_HEAD_DIM, 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

@@ -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

@@ -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

@@ -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):
+    def test_direct_load_splits_k_and_v_for_page_first_direct(self):
-        # Direct single-buffer D2H is not reliable for asymmetric K/V in the
+        # Direct kernels derive copy size from each call's first tensor, so K/V
-        # current sgl-kernel fast path, so the asymmetric host pool is kernel-only.
+        # must be dispatched separately when their head dims differ.
-        host = _make_host("page_first")
+        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):
+        self.assertEqual(transfer.call_count, 2)
-        # Same restriction for D2H backup: asymmetric MHA uses the kernel path
+        k_call, v_call = transfer.call_args_list
-        # until the direct kernel has an explicit safe asymmetric mode.
+        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'"):
+        with self.assertRaisesRegex(ValueError, "expected 'page_first_direct'"):
-            host.backup_from_device_all_layer(
+            host.load_to_device_per_layer(
-                device_pool, host_indices, device_indices, io_backend="direct"
+                device_pool,
+                host_indices,
+                device_indices,
+                layer_id=2,
+                io_backend="direct",
             )