diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 2f7e551a4..67e1eda1d 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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 - 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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7f29f3f5e..f2f30f4c1 100644 --- a/python/sglang/srt/server_args.py +++ b/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 - # 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 diff --git a/test/registered/jit/test_kvcacheio_asymmetric.py b/test/registered/jit/test_kvcacheio_asymmetric.py index 6d06ea681..3fa8f9c33 100644 --- a/test/registered/jit/test_kvcacheio_asymmetric.py +++ b/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; -# 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"])) diff --git a/test/registered/models_e2e/test_mimo_v2.py b/test/registered/models_e2e/test_mimo_v2.py index d22019284..2c0bf3293 100644 --- a/test/registered/models_e2e/test_mimo_v2.py +++ b/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", diff --git a/test/registered/models_e2e/test_mimo_v2_flash.py b/test/registered/models_e2e/test_mimo_v2_flash.py index f2d9a0c38..931c19200 100644 --- a/test/registered/models_e2e/test_mimo_v2_flash.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py b/test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py index 6c3f0e7ac..2aa2592eb 100644 --- a/test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py +++ b/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): - # 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", )