Enable CP HiCache direct transfers to use layer-page host layout

CP shared-KV HiCache transfers are per-layer, so the host layout should match the access pattern instead of forcing page-major strides through every layer. This adds a direct-only layer_page_first layout, routes per-layer KV and NSA index backup/load through the TAI LF<->LPF direct kernels, and keeps storage/page-buffer metadata paths fail-fast until their page-level contract is redesigned.\n\nThe direct controller keeps host indices in caller order for both page_first_direct and layer_page_first because the TAI direct path requires CPU index descriptors and owns descriptor coalescing. All-layer backup intentionally loops over per-layer direct kernels rather than using the sgl-kernel all-layer direct ABI.\n\nConstraint: layer_page_first is currently host-only CP HiCache; storage backends assume page-major contiguous page metadata.\nConstraint: TAI LPF direct kernels require CPU int64 page indices and complete page spans.\nRejected: silently fallback to SM copy when TAI LPF kernels are missing | that hides production performance regressions.\nRejected: support storage page metadata in this commit | LPF requires a layer-page-level storage contract, not a one-pointer-per-page contract.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not enable storage or kernel backend for layer_page_first without redesigning page-buffer metadata and adding remote ETE coverage.\nTested: local py_compile for touched runtime files.\nTested: remote py_compile in g0034 container for touched runtime files.\nTested: remote targeted pytest: 5 passed for parser/storage/layout/move_indices smoke coverage.\nNot-tested: full CP HiCache ETE with --hicache-mem-layout layer_page_first after this commit step.\nNot-tested: combined CUDA roundtrip tests in one pytest process; previous independent runs passed but combined run exposed a host-memory registration lifecycle issue.
2026-06-10 02:03:07 +08:00
parent 5e22279670
commit 24da983ff5
8 changed files with 1796 additions and 9 deletions
@@ -17,6 +17,7 @@ from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
    build_in_seq_page_compute_owners,
 )
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
 from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
 from sglang.srt.mem_cache.memory_pool_host import (
    ALLOC_MEMORY_FUNCS,
@@ -32,6 +33,184 @@ from sglang.test.test_utils import CustomTestCase
 register_cuda_ci(est_time=3, suite="stage-b-test-1-gpu-small")


+class _DummyHiCacheStorage(HiCacheStorage):
+    def get(self, key, target_location=None, target_sizes=None):
+        return None
+
+    def batch_get(self, keys, target_locations=None, target_sizes=None):
+        return []
+
+    def set(self, key, value=None, target_location=None, target_sizes=None):
+        return False
+
+    def batch_set(self, keys, values=None, target_locations=None, target_sizes=None):
+        return False
+
+    def exists(self, key):
+        return False
+
+
+class TestLayerPageFirstDirectHostLayout(CustomTestCase):
+    def test_mha_layer_page_first_direct_host_layout_is_layer_page_major(self):
+        device_pool = SimpleNamespace(
+            store_dtype=torch.float16,
+            size=8,
+            start_layer=0,
+            end_layer=3,
+            device="cpu",
+            head_num=2,
+            head_dim=8,
+            layer_num=3,
+        )
+
+        host_pool = MHATokenToKVPoolHost(
+            device_pool=device_pool,
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=4,
+            layout="layer_page_first",
+            pin_memory=False,
+            device="cpu",
+            host_token_capacity=16,
+        )
+
+        self.assertEqual(tuple(host_pool.kv_buffer.shape), (2, 3, 4, 4, 2, 8))
+        self.assertEqual(tuple(host_pool.k_buffer.shape), (3, 4, 4, 2, 8))
+        self.assertEqual(len(host_pool.k_data_refs), 3)
+        self.assertEqual(tuple(host_pool.k_data_refs[0].shape), (4, 4, 2, 8))
+
+    def test_mla_layer_page_first_direct_host_layout_is_layer_page_major(self):
+        device_pool = SimpleNamespace(
+            store_dtype=torch.float16,
+            size=8,
+            start_layer=0,
+            end_layer=3,
+            device="cpu",
+            kv_lora_rank=16,
+            qk_rope_head_dim=4,
+            layer_num=3,
+        )
+
+        host_pool = MLATokenToKVPoolHost(
+            device_pool=device_pool,
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=4,
+            layout="layer_page_first",
+            pin_memory=False,
+            device="cpu",
+            host_token_capacity=16,
+        )
+
+        self.assertEqual(tuple(host_pool.kv_buffer.shape), (3, 4, 4, 1, 20))
+        self.assertEqual(len(host_pool.data_refs), 3)
+        self.assertEqual(tuple(host_pool.data_refs[0].shape), (4, 4, 1, 20))
+
+    def test_nsa_layer_page_first_direct_indexer_layout_is_layer_page_major(self):
+        indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype
+        device_pool = SimpleNamespace(
+            store_dtype=torch.float16,
+            size=8,
+            start_layer=0,
+            end_layer=3,
+            device="cpu",
+            kv_lora_rank=16,
+            qk_rope_head_dim=4,
+            kv_cache_dim=24,
+            layer_num=3,
+            index_head_dim=16,
+            quant_block_size=8,
+            index_k_with_scale_buffer=[
+                torch.empty((5, 96), dtype=indexer_dtype) for _ in range(3)
+            ],
+        )
+
+        host_pool = NSATokenToKVPoolHost(
+            device_pool=device_pool,
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=4,
+            layout="layer_page_first",
+            pin_memory=False,
+            device="cpu",
+            host_token_capacity=16,
+        )
+
+        self.assertEqual(tuple(host_pool.kv_buffer.shape), (3, 4, 4, 1, 24))
+        self.assertEqual(
+            tuple(host_pool.index_k_with_scale_buffer.shape),
+            (3, host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size),
+        )
+        self.assertEqual(len(host_pool.index_k_data_refs), 3)
+        self.assertEqual(
+            tuple(host_pool.index_k_data_refs[0].shape),
+            (host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size),
+        )
+
+    def test_mha_layer_page_first_page_buffer_meta_fails_fast_for_storage(self):
+        device_pool = SimpleNamespace(
+            store_dtype=torch.float16,
+            size=8,
+            start_layer=0,
+            end_layer=2,
+            device="cpu",
+            head_num=2,
+            head_dim=8,
+            layer_num=2,
+        )
+        host_pool = MHATokenToKVPoolHost(
+            device_pool=device_pool,
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=4,
+            layout="layer_page_first",
+            pin_memory=False,
+            device="cpu",
+            host_token_capacity=16,
+        )
+
+        with self.assertRaisesRegex(
+            RuntimeError, "layer_page_first_page_buffer_meta_unsupported"
+        ):
+            host_pool.get_page_buffer_meta(torch.arange(0, 4, dtype=torch.int64))
+
+    def test_mla_layer_page_first_page_buffer_meta_fails_fast_for_storage(self):
+        device_pool = SimpleNamespace(
+            store_dtype=torch.float16,
+            size=8,
+            start_layer=0,
+            end_layer=2,
+            device="cpu",
+            kv_lora_rank=16,
+            qk_rope_head_dim=4,
+            layer_num=2,
+        )
+        host_pool = MLATokenToKVPoolHost(
+            device_pool=device_pool,
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=4,
+            layout="layer_page_first",
+            pin_memory=False,
+            device="cpu",
+            host_token_capacity=16,
+        )
+
+        with self.assertRaisesRegex(
+            RuntimeError, "layer_page_first_page_buffer_meta_unsupported"
+        ):
+            host_pool.get_page_buffer_meta(torch.arange(0, 4, dtype=torch.int64))
+
+    def test_storage_registration_fails_fast_for_layer_page_first(self):
+        storage = _DummyHiCacheStorage()
+        mem_pool_host = SimpleNamespace(layout="layer_page_first")
+
+        with self.assertRaisesRegex(
+            RuntimeError, "layer_page_first_storage_unsupported"
+        ):
+            storage.register_mem_pool_host(mem_pool_host)
+
+
 class TestNSAHiCacheTransfer(CustomTestCase):
    def setUp(self):
        if not torch.cuda.is_available():
@@ -1263,7 +1442,9 @@ class TestNSAHiCacheTransfer(CustomTestCase):
            io_backend="direct", layout="page_first_direct"
        )

-    def test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages(self):
+    def _run_direct_roundtrip_preserves_kv_and_indexer_pages(
+        self, *, layout: str, dtype: torch.dtype
+    ):
        page_size = 64
        layer_num = 3
        size = page_size * 20
@@ -1272,7 +1453,7 @@ class TestNSAHiCacheTransfer(CustomTestCase):
            size=size,
            page_size=page_size,
            kv_lora_rank=512,
-            dtype=torch.float8_e4m3fn,
+            dtype=dtype,
            qk_rope_head_dim=64,
            layer_num=layer_num,
            device="cuda",
@@ -1285,7 +1466,7 @@ class TestNSAHiCacheTransfer(CustomTestCase):
            host_to_device_ratio=2.0,
            host_size=0,
            page_size=page_size,
-            layout="page_first_direct",
+            layout=layout,
            pin_memory=True,
            device="cpu",
        )
@@ -1373,7 +1554,8 @@ class TestNSAHiCacheTransfer(CustomTestCase):
                ]
                self.assertTrue(
                    torch.equal(got_kv, expected_kv[layer_id][page_idx]),
-                    f"KV roundtrip mismatch layer={layer_id} dst_page={dst_page}",
+                    f"KV roundtrip mismatch layout={layout} dtype={dtype} "
+                    f"layer={layer_id} dst_page={dst_page}",
                )

                got_index = device_pool.index_k_with_scale_buffer[layer_id][
@@ -1381,9 +1563,24 @@ class TestNSAHiCacheTransfer(CustomTestCase):
                ]
                self.assertTrue(
                    torch.equal(got_index, expected_index[layer_id][page_idx]),
-                    f"index roundtrip mismatch layer={layer_id} dst_page={dst_page}",
+                    f"index roundtrip mismatch layout={layout} dtype={dtype} "
+                    f"layer={layer_id} dst_page={dst_page}",
                )

+    def test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages(self):
+        self._run_direct_roundtrip_preserves_kv_and_indexer_pages(
+            layout="page_first_direct", dtype=torch.float8_e4m3fn
+        )
+
+    def test_fp8_layer_page_first_roundtrip_preserves_kv_and_indexer_pages(self):
+        self._run_direct_roundtrip_preserves_kv_and_indexer_pages(
+            layout="layer_page_first", dtype=torch.float8_e4m3fn
+        )
+
+    def test_bf16_layer_page_first_roundtrip_preserves_kv_and_indexer_pages(self):
+        self._run_direct_roundtrip_preserves_kv_and_indexer_pages(
+            layout="layer_page_first", dtype=torch.bfloat16
+        )

 class TestPageFirstDirectAllLayerBackupRoute(CustomTestCase):
    def test_mla_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
@@ -1576,6 +1773,126 @@ class TestNSAIndexerPageIndices(CustomTestCase):
            self.assertEqual(call["dst_indices"].tolist(), [0, 1])
            self.assertEqual(call["page_size"], 1)

+    def test_mla_layer_page_first_all_layer_backup_uses_tai_per_layer_route(self):
+        host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
+        host_pool.layout = "layer_page_first"
+        host_pool.page_size = 4
+        host_pool.layer_num = 2
+        host_pool.kv_buffer = "host-mla-layer-page-first"
+        device_pool = type(
+            "FakeDevicePool",
+            (),
+            {"kv_buffer": ["device-mla-layer-0", "device-mla-layer-1"]},
+        )()
+        calls = []
+
+        def fake_tai_transfer(**kwargs):
+            calls.append(kwargs)
+
+        with patch(
+            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
+            return_value=fake_tai_transfer,
+        ):
+            host_pool.backup_from_device_all_layer(
+                device_pool,
+                torch.tensor([0, 1, 2, 3], dtype=torch.int64),
+                torch.tensor([8, 9, 10, 11], dtype=torch.int64),
+                "direct",
+            )
+
+        self.assertEqual([call["layer_id"] for call in calls], [0, 1])
+        for layer_id, call in enumerate(calls):
+            self.assertEqual(call["src_ptrs"], [f"device-mla-layer-{layer_id}"])
+            self.assertEqual(call["dst_ptrs"], ["host-mla-layer-page-first"])
+            self.assertEqual(call["src_indices"].tolist(), [8, 9, 10, 11])
+            self.assertEqual(call["dst_indices"].tolist(), [0, 1, 2, 3])
+            self.assertEqual(call["page_size"], 4)
+
+    def test_mha_layer_page_first_all_layer_backup_uses_tai_per_layer_route(self):
+        host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
+        host_pool.layout = "layer_page_first"
+        host_pool.page_size = 4
+        host_pool.layer_num = 2
+        host_pool.kv_buffer = ["host-k-layer-page-first", "host-v-layer-page-first"]
+        device_pool = type(
+            "FakeDevicePool",
+            (),
+            {
+                "k_buffer": ["device-k-layer-0", "device-k-layer-1"],
+                "v_buffer": ["device-v-layer-0", "device-v-layer-1"],
+            },
+        )()
+        calls = []
+
+        def fake_tai_transfer(**kwargs):
+            calls.append(kwargs)
+
+        with patch(
+            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
+            return_value=fake_tai_transfer,
+        ):
+            host_pool.backup_from_device_all_layer(
+                device_pool,
+                torch.tensor([0, 1, 2, 3], dtype=torch.int64),
+                torch.tensor([8, 9, 10, 11], dtype=torch.int64),
+                "direct",
+            )
+
+        self.assertEqual([call["layer_id"] for call in calls], [0, 1])
+        for layer_id, call in enumerate(calls):
+            self.assertEqual(
+                call["src_ptrs"],
+                [f"device-k-layer-{layer_id}", f"device-v-layer-{layer_id}"],
+            )
+            self.assertEqual(
+                call["dst_ptrs"],
+                ["host-k-layer-page-first", "host-v-layer-page-first"],
+            )
+            self.assertEqual(call["src_indices"].tolist(), [8, 9, 10, 11])
+            self.assertEqual(call["dst_indices"].tolist(), [0, 1, 2, 3])
+            self.assertEqual(call["page_size"], 4)
+
+    def test_layer_page_first_all_layer_indexer_backup_uses_tai_per_layer_route(self):
+        host_pool = self.make_host_pool_stub(page_size=4)
+        host_pool.layout = "layer_page_first"
+        host_pool.indexer_page_stride_size = 8
+        host_pool.layer_num = 3
+        host_pool.index_k_with_scale_buffer = "host-layer-page-first-indexer"
+        device_pool = type(
+            "FakeDevicePool",
+            (),
+            {
+                "index_k_with_scale_buffer": [
+                    "device-layer-0",
+                    "device-layer-1",
+                    "device-layer-2",
+                ]
+            },
+        )()
+        calls = []
+
+        def fake_tai_transfer(**kwargs):
+            calls.append(kwargs)
+
+        with patch(
+            "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_lpf",
+            return_value=fake_tai_transfer,
+        ):
+            host_pool._backup_indexer_from_device_all_layer(
+                device_pool,
+                torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype=torch.int64),
+                torch.tensor([8, 9, 10, 11, 12, 13, 14, 15], dtype=torch.int64),
+                "direct",
+            )
+
+        self.assertEqual([call["layer_id"] for call in calls], [0, 1, 2])
+        for layer_id, call in enumerate(calls):
+            self.assertEqual(call["src_ptrs"], [f"device-layer-{layer_id}"])
+            self.assertEqual(call["dst_ptrs"], ["host-layer-page-first-indexer"])
+            self.assertEqual(call["src_indices"].tolist(), [2, 3])
+            self.assertEqual(call["dst_indices"].tolist(), [0, 1])
+            self.assertEqual(call["page_size"], 1)
+

 if __name__ == "__main__":
    unittest.main()