diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 4017f4da3..4569ce529 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -16,6 +16,9 @@ from sglang.jit_kernel.hicache import ( from sglang.jit_kernel.hicache import ( transfer_hicache_all_layer_mla as jit_transfer_hicache_all_layer_mla, ) +from sglang.jit_kernel.hicache import ( + transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf, +) from sglang.jit_kernel.hicache import ( transfer_hicache_all_layer_staged_lf_pf as jit_transfer_hicache_all_layer_staged_lf_pf, ) @@ -941,6 +944,51 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost): kernels derive copy sizes from each call's first tensor. """ + def _init_write_back_staging_buffers(self): + self.staging_page_capacity = 0 + self.staging_token_capacity = 0 + self.staging_k_buffer = None + self.staging_v_buffer = None + self.can_use_write_back_jit = False + if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps): + return + + # K and V have different element sizes. Use the single-buffer staged + # kernel for each side, which specializes to its native stride. + can_use_staged_jit = (_is_cuda or _is_hip) and all( + can_use_write_back_jit_kernel(element_size=element_size) + for element_size in ( + self._k_token_stride_size(), + self._v_token_stride_size(), + ) + ) + if not can_use_staged_jit: + return + + self.can_use_write_back_jit = True + self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK) + self.staging_token_capacity = self.staging_page_capacity * self.page_size + self.staging_k_buffer = torch.empty( + ( + self.staging_token_capacity, + self.layer_num, + self.head_num, + self.head_dim, + ), + dtype=self.dtype, + device=self.device_pool.device, + ) + self.staging_v_buffer = torch.empty( + ( + self.staging_token_capacity, + self.layer_num, + self.head_num, + self.v_head_dim, + ), + dtype=self.dtype, + device=self.device_pool.device, + ) + def get_size_per_token(self): self.head_num = self.device_pool.head_num self.head_dim = self.device_pool.head_dim @@ -1092,24 +1140,42 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost): f"Unsupported layout for models with head_dim != v_head_dim " f"and io_backend='kernel': {self.layout}; expected 'page_first'." ) - transfer_kv_all_layer_mla_lf_pf( - src_layers=device_pool.k_data_ptrs, - dst=self.k_buffer, - src_indices=device_indices, - dst_indices=host_indices, - item_size=self._k_token_stride_size(), - dst_layout_dim=self._k_layout_dim(), - num_layers=self.layer_num, - ) - transfer_kv_all_layer_mla_lf_pf( - src_layers=device_pool.v_data_ptrs, - dst=self.v_buffer, - src_indices=device_indices, - dst_indices=host_indices, - item_size=self._v_token_stride_size(), - dst_layout_dim=self._v_layout_dim(), - num_layers=self.layer_num, - ) + if self.can_use_write_back_jit: + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=device_pool.k_data_ptrs, + src_indices=device_indices, + dst_indices=host_indices, + staging=self.staging_k_buffer, + dst=self.k_buffer, + page_size=self.page_size, + ) + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=device_pool.v_data_ptrs, + src_indices=device_indices, + dst_indices=host_indices, + staging=self.staging_v_buffer, + dst=self.v_buffer, + page_size=self.page_size, + ) + else: + transfer_kv_all_layer_mla_lf_pf( + src_layers=device_pool.k_data_ptrs, + dst=self.k_buffer, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self._k_token_stride_size(), + dst_layout_dim=self._k_layout_dim(), + num_layers=self.layer_num, + ) + transfer_kv_all_layer_mla_lf_pf( + src_layers=device_pool.v_data_ptrs, + dst=self.v_buffer, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self._v_token_stride_size(), + dst_layout_dim=self._v_layout_dim(), + num_layers=self.layer_num, + ) elif io_backend == "direct": if self.layout != "page_first_direct": raise ValueError( diff --git a/python/sglang/test/kl_test_utils.py b/python/sglang/test/kl_test_utils.py index df05f00ed..b6c5b3751 100644 --- a/python/sglang/test/kl_test_utils.py +++ b/python/sglang/test/kl_test_utils.py @@ -137,6 +137,7 @@ def _generate( return_logprob=False, logprob_start_len=-1, temperature=0.0, + routed_dp_rank=None, ): """Send generate request and return results.""" json_data = { @@ -155,11 +156,19 @@ def _generate( "logprob_start_len": logprob_start_len, } ) + if routed_dp_rank is not None: + json_data["routed_dp_rank"] = routed_dp_rank response = requests.post(base_url + "/generate", json=json_data) return response.json() -def _get_input_logprobs(base_url, new_input_ids, output_logprobs, temperature=0.0): +def _get_input_logprobs( + base_url, + new_input_ids, + output_logprobs, + temperature=0.0, + routed_dp_rank=None, +): """Run prefill to get input logprobs matching output logprobs.""" _flush_cache(base_url) results = _generate( @@ -169,6 +178,7 @@ def _get_input_logprobs(base_url, new_input_ids, output_logprobs, temperature=0. return_logprob=True, logprob_start_len=0, temperature=temperature, + routed_dp_rank=routed_dp_rank, ) assert len(results) == len(new_input_ids) diff --git a/test/registered/jit/test_kvcacheio_asymmetric.py b/test/registered/jit/test_kvcacheio_asymmetric.py index 711c9f5af..e13d99784 100644 --- a/test/registered/jit/test_kvcacheio_asymmetric.py +++ b/test/registered/jit/test_kvcacheio_asymmetric.py @@ -55,6 +55,7 @@ def make_host_pool(dtype, layout="page_first"): host.head_dim = K_HEAD_DIM host.v_head_dim = V_HEAD_DIM host.dtype = dtype + host.can_use_write_back_jit = False 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) diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py new file mode 100644 index 000000000..1eb02974e --- /dev/null +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mimo.py @@ -0,0 +1,165 @@ +"""MiMo V2.5 HiCache host load-back accuracy regression test.""" + +import random +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kl_test_utils import ( + _extract_output_logprobs, + _flush_cache, + _generate, + _get_input_logprobs, + compare_kl_divergence, +) +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +MIMO_MODEL = "XiaomiMiMo/MiMo-V2.5" +MIMO_LAUNCH_TIMEOUT = 3600 + +# MiMo V2.5 is pre-cached on the eight-H200 runner. The H200-only nightly suite +# exercises the asymmetric MHA host pool end to end without adding PR CI cost. +register_cuda_ci(est_time=1200, suite="nightly-8-gpu-h200", nightly=True) + + +class TestUnifiedMiMoHiCacheLoadBackKL(CustomTestCase): + """Verify KL accuracy after asymmetric MHA KV is evicted to and loaded from L2.""" + + page_size = 64 + prompt_len = 1024 + max_total_tokens = 4096 + # EAGLE verification and full-prefill replay use different execution paths; + # allow their expected numerical drift while still catching regressions. + kl_threshold = 0.01 + # DP ranks own independent radix trees, so cache pressure and load-back must + # target the same rank instead of following the round-robin default. + routed_dp_rank = 0 + + @classmethod + def setUpClass(cls): + cls.model = MIMO_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=MIMO_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--cuda-graph-max-bs-decode", + "64", + "--enable-multimodal", + "--tp", + "8", + "--dp", + "2", + "--enable-dp-attention", + "--mm-enable-dp-encoder", + "--attention-backend", + "fa3", + "--mm-attention-backend", + "fa3", + "--mem-fraction-static", + "0.65", + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--enable-multi-layer-eagle", + "--reasoning-parser", + "mimo", + "--page-size", + str(cls.page_size), + "--max-total-tokens", + str(cls.max_total_tokens), + "--enable-hierarchical-cache", + "--hicache-ratio", + "1.2", + "--hicache-write-policy", + "write_through", + "--hicache-io-backend", + "kernel", + "--hicache-mem-layout", + "page_first", + ], + env={ + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_USE_CUDA_IPC_TRANSPORT": "1", + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + @classmethod + def _prompt(cls, seed: int) -> list[int]: + rng = random.Random(seed) + return [rng.randint(1, 30000) for _ in range(cls.prompt_len)] + + def _generate_one(self, input_ids, max_new_tokens, return_logprob=False): + results = _generate( + self.base_url, + [input_ids], + max_new_tokens=max_new_tokens, + return_logprob=return_logprob, + temperature=0, + routed_dp_rank=self.routed_dp_rank, + ) + self.assertEqual(len(results), 1) + return results[0] + + def test_host_load_back_logprobs_match_prefill_replay(self): + """Force L2 eviction, then compare load-back output logprobs with replay.""" + base_prompt = self._prompt(1) + pressure_prompts = [self._prompt(seed) for seed in range(2, 6)] + + _flush_cache(self.base_url) + self._generate_one(base_prompt, max_new_tokens=1) + + # Four unique page-aligned prefixes fill the 4096-token L1 cache and + # evict the oldest prefix (base_prompt) to the HiCache host tier. + for prompt in pressure_prompts: + self._generate_one(prompt, max_new_tokens=1) + + load_back = self._generate_one( + base_prompt, max_new_tokens=8, return_logprob=True + ) + meta_info = load_back["meta_info"] + cached_details = meta_info.get("cached_tokens_details") or {} + host_cached_tokens = int(cached_details.get("host", 0)) + self.assertGreater( + host_cached_tokens, + 0, + "Expected the original prefix to be restored from the HiCache host tier; " + f"got cached_tokens={meta_info.get('cached_tokens')}, " + f"cached_tokens_details={cached_details}", + ) + + output_logprobs = [_extract_output_logprobs(load_back)] + replay_input_ids = [base_prompt + load_back["output_ids"]] + input_logprobs = _get_input_logprobs( + self.base_url, + replay_input_ids, + output_logprobs, + temperature=0, + routed_dp_rank=self.routed_dp_rank, + ) + compare_kl_divergence( + input_logprobs, + output_logprobs, + {self.model: {"kl_div": self.kl_threshold}}, + self.model, + "hicache_host_load_back", + ) + + +if __name__ == "__main__": + unittest.main() 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 6028c4fc2..e75b99039 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 @@ -50,6 +50,7 @@ def _make_host(layout: str) -> AsymmetricMHATokenToKVPoolHost: raise ValueError(f"Unsupported test layout: {layout}") host.kv_buffer = (torch.empty(k_dims), torch.empty(v_dims)) + host.can_use_write_back_jit = False return host @@ -79,6 +80,42 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase): get_mha_host_pool_cls(asymmetric_pool), AsymmetricMHATokenToKVPoolHost ) + def test_staged_write_back_jit_uses_separate_kv_buffers(self): + host = _make_host("page_first") + host.page_num = 4 + host.v_head_dim = 8 + host.device_pool = SimpleNamespace(device="cuda") + cpu_empty = torch.empty + + def _cpu_empty(shape, *, dtype, device): + return cpu_empty(shape, dtype=dtype) + + with ( + mock.patch("sglang.srt.mem_cache.pool_host.mha._is_cuda", True), + mock.patch("sglang.srt.mem_cache.pool_host.mha._is_hip", False), + mock.patch("sglang.srt.mem_cache.pool_host.mha._is_npu", False), + mock.patch("sglang.srt.mem_cache.pool_host.mha._is_xpu", False), + mock.patch("sglang.srt.mem_cache.pool_host.mha._is_mps", False), + mock.patch( + "sglang.srt.mem_cache.pool_host.mha.can_use_write_back_jit_kernel", + return_value=True, + ) as can_use, + mock.patch( + "sglang.srt.mem_cache.pool_host.mha.torch.empty", + side_effect=_cpu_empty, + ), + ): + host._init_write_back_staging_buffers() + + self.assertTrue(host.can_use_write_back_jit) + self.assertEqual(host.staging_page_capacity, 4) + self.assertEqual(host.staging_token_capacity, 8) + self.assertEqual(host.staging_k_buffer.shape, (8, 3, 2, 4)) + self.assertEqual(host.staging_v_buffer.shape, (8, 3, 2, 8)) + self.assertEqual( + [call.kwargs["element_size"] for call in can_use.call_args_list], [16, 32] + ) + def test_kernel_load_splits_k_and_v_with_separate_strides(self): # Dispatch-only test: the CUDA kernel is mocked; this verifies that K and # V are sent as separate single-buffer calls with their own byte strides. @@ -137,6 +174,116 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase): self.assertEqual(v_call.kwargs["item_size"], 24) self.assertEqual(v_call.kwargs["dst_layout_dim"], 72) + def test_kernel_backup_uses_staged_kernel_for_each_kv_buffer(self): + host = _make_host("page_first") + host.can_use_write_back_jit = True + host.staging_k_buffer = torch.empty(4, 3, 2, 4) + host.staging_v_buffer = torch.empty(4, 3, 2, 6) + 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.pool_host.mha.jit_transfer_hicache_all_layer_mla_staged_lf_pf" + ) as staged, + mock.patch( + "sglang.srt.mem_cache.pool_host.mha.transfer_kv_all_layer_mla_lf_pf", + create=True, + ) as fallback, + ): + host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend="kernel" + ) + + self.assertEqual(staged.call_count, 2) + self.assertEqual(fallback.call_count, 0) + k_call, v_call = staged.call_args_list + self.assertIs(k_call.kwargs["ptr_src"], device_pool.k_data_ptrs) + self.assertIs(k_call.kwargs["staging"], host.staging_k_buffer) + self.assertIs(k_call.kwargs["dst"], host.k_buffer) + self.assertIs(v_call.kwargs["ptr_src"], device_pool.v_data_ptrs) + self.assertIs(v_call.kwargs["staging"], host.staging_v_buffer) + self.assertIs(v_call.kwargs["dst"], host.v_buffer) + + def test_staged_kernel_backup_load_roundtrip_preserves_asymmetric_kv(self): + """Staged write-back must preserve both K and V values across a round trip.""" + host = _make_host("page_first") + host.can_use_write_back_jit = True + host.staging_k_buffer = torch.empty(4, 3, 2, 4) + host.staging_v_buffer = torch.empty(4, 3, 2, 6) + 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) + + for layer_id in range(host.layer_num): + device_pool.k_buffer[layer_id].copy_( + torch.arange(device_pool.k_buffer[layer_id].numel()).reshape_as( + device_pool.k_buffer[layer_id] + ) + + layer_id * 1000 + ) + device_pool.v_buffer[layer_id].copy_( + torch.arange(device_pool.v_buffer[layer_id].numel()).reshape_as( + device_pool.v_buffer[layer_id] + ) + + layer_id * 10000 + ) + + expected_k = [buffer[device_indices].clone() for buffer in device_pool.k_buffer] + expected_v = [buffer[device_indices].clone() for buffer in device_pool.v_buffer] + buffers_by_ptrs = { + tuple(device_pool.k_data_ptrs.tolist()): device_pool.k_buffer, + tuple(device_pool.v_data_ptrs.tolist()): device_pool.v_buffer, + } + + def staged_copy(*, ptr_src, src_indices, dst_indices, dst, **_): + src_buffers = buffers_by_ptrs[tuple(ptr_src.tolist())] + for layer_id, src in enumerate(src_buffers): + dst[dst_indices, layer_id] = src[src_indices] + + def per_layer_copy(*, src, dst, src_indices, dst_indices, layer_id, **_): + dst[dst_indices] = src[src_indices, layer_id] + + with ( + mock.patch( + "sglang.srt.mem_cache.pool_host.mha.jit_transfer_hicache_all_layer_mla_staged_lf_pf", + side_effect=staged_copy, + ), + mock.patch( + "sglang.srt.mem_cache.pool_host.mha.transfer_kv_per_layer_mla_pf_lf", + side_effect=per_layer_copy, + create=True, + ), + ): + host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend="kernel" + ) + for buffer in device_pool.k_buffer + device_pool.v_buffer: + buffer.zero_() + for layer_id in range(host.layer_num): + host.load_to_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id=layer_id, + io_backend="kernel", + ) + + for layer_id in range(host.layer_num): + torch.testing.assert_close( + host.k_buffer[host_indices, layer_id], expected_k[layer_id] + ) + torch.testing.assert_close( + host.v_buffer[host_indices, layer_id], expected_v[layer_id] + ) + torch.testing.assert_close( + device_pool.k_buffer[layer_id][device_indices], expected_k[layer_id] + ) + torch.testing.assert_close( + device_pool.v_buffer[layer_id][device_indices], expected_v[layer_id] + ) + 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.