[UnifiedTree]: Rollback mamba hicache test to direct io backend (#28904)

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -1405,11 +1405,15 @@ class MambaPoolHost(HostKVCache):
     ):
         self.device_pool = device_pool
         self.page_size = 1
-        assert layout in [
+
-            "page_first",
+        # TODO: Mamba pool is currently incompatible with write-back staging
-            "page_first_direct",
+        # kernel; only allow 'page_first_direct' + 'direct' for now.
-            "layer_first",
+        # Relax this restriction once the staging bug is fixed.
-        ], f"Unsupported layout: {layout}"
+        if layout != "page_first_direct":
+            raise ValueError(
+                f"MambaPoolHost only supports layout='page_first_direct', "
+                f"got '{layout}'."
+            )
 
         self.layout = layout
         self.pin_memory = pin_memory
@@ -1767,6 +1771,11 @@ class MambaPoolHost(HostKVCache):
         layer_id,
         io_backend="kernel",
     ):
+        if io_backend != "direct":
+            raise ValueError(
+                f"MambaPoolHost only supports io_backend='direct', "
+                f"got '{io_backend}'."
+            )
         if self.layout in ["page_first", "page_first_direct"]:
             self._copy_tensor_pf_lf(
                 src=self.temporal_buffer,
@@ -1807,6 +1816,11 @@ class MambaPoolHost(HostKVCache):
     def backup_from_device_all_layer(
         self, device_pool, host_indices, device_indices, io_backend="kernel"
     ):
+        if io_backend != "direct":
+            raise ValueError(
+                f"MambaPoolHost only supports io_backend='direct', "
+                f"got '{io_backend}'."
+            )
         if self.layout in ["page_first", "page_first_direct"]:
             self._copy_tensor_all_layers_lf_pf(
                 src_layers=device_pool.mamba_cache.temporal,

test/registered/hicache/test_qwen35_hicache.py

@@ -65,6 +65,8 @@ class TestQwen35WithHiCache(CustomTestCase):
                 '{"enable_multithread_load": true,"num_threads": 64}',
                 "--hicache-mem-layout",
                 "page_first_direct",
+                "--hicache-io-backend",
+                "direct",
                 "--enable-hierarchical-cache",
                 "--hicache-ratio",
                 "2",

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py

@@ -108,9 +108,9 @@ class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
                 "--hicache-write-policy",
                 "write_through",
                 "--hicache-io-backend",
-                "kernel",
+                "direct",
                 "--hicache-mem-layout",
-                "page_first",
+                "page_first_direct",
                 "--max-total-tokens",
                 "12000",
                 "--max-mamba-cache-size",
@@ -169,9 +169,9 @@ class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
                 "--hicache-storage-prefetch-policy",
                 "wait_complete",
                 "--hicache-io-backend",
-                "kernel",
+                "direct",
                 "--hicache-mem-layout",
-                "page_first",
+                "page_first_direct",
                 "--hicache-storage-backend",
                 "file",
                 "--max-mamba-cache-size",

test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py

@@ -353,6 +353,10 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
                 torch.equal(host.kv_buffer[host_indices, layer_id], expected[layer_id])
             )
 
+    @unittest.skip(
+        "TODO: Mamba pool is currently incompatible with write-back staging "
+        "kernel; re-enable once the staging bug is fixed."
+    )
     def test_mamba_backup_then_load_roundtrip_uses_staged(self):
         num_layers = 2
         host_indices = _indices(0, 4)