[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",
-            "page_first_direct",
-            "layer_first",
-        ], f"Unsupported layout: {layout}"
+
+        # TODO: Mamba pool is currently incompatible with write-back staging
+        # kernel; only allow 'page_first_direct' + 'direct' for now.
+        # Relax this restriction once the staging bug is fixed.
+        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,