fix(hicache): support staged write-back for asymmetric MHA (#30981)

Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>

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(
+            if self.can_use_write_back_jit:
-                src_layers=device_pool.k_data_ptrs,
+                jit_transfer_hicache_all_layer_mla_staged_lf_pf(
-                dst=self.k_buffer,
+                    ptr_src=device_pool.k_data_ptrs,
-                src_indices=device_indices,
+                    src_indices=device_indices,
-                dst_indices=host_indices,
+                    dst_indices=host_indices,
-                item_size=self._k_token_stride_size(),
+                    staging=self.staging_k_buffer,
-                dst_layout_dim=self._k_layout_dim(),
+                    dst=self.k_buffer,
-                num_layers=self.layer_num,
+                    page_size=self.page_size,
-            )
+                )
-            transfer_kv_all_layer_mla_lf_pf(
+                jit_transfer_hicache_all_layer_mla_staged_lf_pf(
-                src_layers=device_pool.v_data_ptrs,
+                    ptr_src=device_pool.v_data_ptrs,
-                dst=self.v_buffer,
+                    src_indices=device_indices,
-                src_indices=device_indices,
+                    dst_indices=host_indices,
-                dst_indices=host_indices,
+                    staging=self.staging_v_buffer,
-                item_size=self._v_token_stride_size(),
+                    dst=self.v_buffer,
-                dst_layout_dim=self._v_layout_dim(),
+                    page_size=self.page_size,
-                num_layers=self.layer_num,
+                )
-            )
+            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(

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)
 

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)

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()

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.