diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 0799a6b73..cbbd07b3b 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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, diff --git a/test/registered/hicache/test_qwen35_hicache.py b/test/registered/hicache/test_qwen35_hicache.py index d2b478fe8..8398faaba 100644 --- a/test/registered/hicache/test_qwen35_hicache.py +++ b/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", diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py index 591707712..a46a7f2ca 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_mamba.py +++ b/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", diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 3c613338b..7e4d4fc0d 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/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)