[HiCache] Unified Mooncake Registration for Logical Anchors and Draft Pools

python/sglang/srt/managers/cache_controller.py

@@ -835,7 +835,13 @@ def set_draft_kv_pool(self, draft_device_pool, draft_host_pool) -> None:
         self._maybe_register_draft_with_storage()
 
     def _maybe_register_draft_with_storage(self) -> None:
-        """Pick the draft L3 IO implementation."""
+        """Pick the draft L3 IO implementation.
+
+        For Mooncake (v2 path), checks is_mooncake_registerable first to avoid
+        registration failures when draft pools use non-Mooncake allocators
+        (e.g., EAGLE + mmap). When non-registerable, L3 is skipped but L2
+        (host DRAM) continues.
+        """
         self.draft_page_get_func = None
         self.draft_page_set_func = None
         if not self.has_draft or not self.enable_storage:
@@ -851,6 +857,13 @@ def _maybe_register_draft_with_storage(self) -> None:
                     "supported on the mooncake v2 path."
                 )
                 return
+            if not self.mem_pool_host_draft.is_mooncake_registerable:
+                logger.warning(
+                    "Draft L3 disabled: draft pool allocator (%s) is not "
+                    "Mooncake-compatible. L2 (host DRAM) still works.",
+                    type(self.mem_pool_host_draft.allocator).__name__,
+                )
+                return
             self.storage_backend.register_mem_host_pool_v2(
                 self.mem_pool_host_draft, PoolName.DRAFT
             )

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -2283,6 +2283,11 @@ class LogicalHostPool:
     compressed side pools use these logical FULL indices as stable page anchors.
     """
 
+    @property
+    def is_mooncake_registerable(self) -> bool:
+        """Logical anchors have no physical buffer to register."""
+        return False
+
     def __init__(self, size: int, page_size: int, layout: str = "layer_first"):
         if size % page_size != 0:
             raise ValueError(
@@ -3170,6 +3175,11 @@ def start_layer(self):
     def end_layer(self):
         return self.anchor_entry.host_pool.end_layer
 
+    @property
+    def is_mooncake_registerable(self) -> bool:
+        """Delegate to anchor pool's Mooncake registration capability."""
+        return self.anchor_entry.host_pool.is_mooncake_registerable
+
     def get_ksize_per_token(self):
         return self.anchor_entry.host_pool.get_ksize_per_token()
 

python/sglang/srt/mem_cache/pool_host/base.py

@@ -77,6 +77,25 @@ def wrapper(self, *args, **kwargs):
 
 
 class HostKVCache(abc.ABC):
+    """Base class for host-side KV cache pools.
+
+    Attributes:
+        is_mooncake_registerable: Whether this pool's buffer can be zero-copy
+            registered with Mooncake's distributed memory store. Defaults to True
+            for pools with physical buffers. Logical anchors and pools using
+            non-Mooncake allocators override this to False.
+    """
+
+    @property
+    def is_mooncake_registerable(self) -> bool:
+        """Whether this pool's buffer can be zero-copy registered with Mooncake.
+
+        Returns True if the allocator produces Mooncake-compatible tensors.
+        Returns False for:
+        - Logical anchors (no physical buffer, only page indices)
+        - Pools using non-Mooncake allocators (e.g., mmap for draft pools)
+        """
+        return self.allocator.is_mooncake_compatible
 
     def __init__(
         self,

python/sglang/srt/mem_cache/pool_host/common.py

@@ -16,6 +16,11 @@ def __init__(self):
         self.dtype = None
         self.dims = None
 
+    @property
+    def is_mooncake_compatible(self) -> bool:
+        """Whether tensors from this allocator can be registered with Mooncake."""
+        return False
+
     def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor:
         assert (
             device == "cpu"

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -39,6 +39,11 @@ def __init__(self):
         self.allocator = MooncakeHostMemAllocator()
         self.ptr = None
 
+    @property
+    def is_mooncake_compatible(self) -> bool:
+        """Mooncake allocator is compatible with Mooncake registration."""
+        return True
+
     def allocate(
         self, dims: tuple, dtype: torch.dtype, device: str = "cpu"
     ) -> torch.Tensor:
@@ -643,10 +648,19 @@ def warmup(self):
 
     def register_mem_pool_host(self, mem_pool_host: HostKVCache):
         super().register_mem_pool_host(mem_pool_host)
+
         if getattr(self.mem_pool_host, "kv_buffer", None) is None:
-            # Hybrid logical anchors only own allocation indices. Their physical
-            # tensors are registered through register_mem_host_pool_v2().
+            # Logical anchor: no physical buffer to register.
+            # Use marker keys for existence tracking in batch_exists/batch_get/batch_set.
+            self._logical_kv_anchor = True
+            self.gb_per_page = 0.0
+            logger.info(
+                "Detected logical KV anchor pool (%s), "
+                "using marker keys for existence tracking",
+                type(mem_pool_host).__name__,
+            )
             return
+
         try:
             for buffer in self._iter_host_pool_buffers(self.mem_pool_host):
                 super().register_buffer(buffer)
@@ -1001,6 +1015,10 @@ def batch_get_v1(
         # Apply extra_backend_tag prefix if available
         keys = self._tag_keys(keys)
 
+        # Logical anchor: check marker keys instead of physical KV data
+        if getattr(self, "_logical_kv_anchor", False):
+            return self._batch_get_logical_anchor(keys)
+
         key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
 
         start_time = time.perf_counter()
@@ -1030,6 +1048,10 @@ def batch_set_v1(
         # Apply extra_backend_tag prefix if available
         keys = self._tag_keys(keys)
 
+        # Logical anchor: set marker keys instead of physical KV data
+        if getattr(self, "_logical_kv_anchor", False):
+            return self._batch_set_logical_anchor(keys)
+
         key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
         key_multiplier = len(key_strs) // len(keys)
         group_ids = (
@@ -1204,6 +1226,15 @@ def batch_exists(
         # Apply extra_backend_tag prefix if available
         keys = self._tag_keys(keys)
 
+        # Logical anchor: check marker keys instead of physical KV keys
+        if getattr(self, "_logical_kv_anchor", False):
+            query_keys = [self._logical_anchor_key(key) for key in keys]
+            exist_result = self._batch_exist(query_keys)
+            for i in range(len(query_keys)):
+                if exist_result[i] != 1:
+                    return i
+            return len(query_keys)
+
         if self.is_mla_backend:
             query_keys = [f"{key}_{self.mla_suffix}_k" for key in keys]
             key_multiplier = 1
@@ -1276,6 +1307,34 @@ def _get_batch_zero_copy_impl(
     def _batch_exist(self, key_strs: List[str]) -> List[int]:
         return self.store.batch_is_exist(key_strs)
 
+    @staticmethod
+    def _logical_anchor_key(key: str) -> str:
+        """Generate marker key for logical anchor existence tracking."""
+        return f"{key}_logical_kv"
+
+    def _batch_set_logical_anchor(self, keys: List[str]) -> List[bool]:
+        """Set marker keys for logical anchor pages.
+
+        Uses 1-byte markers since Mooncake put() rejects zero-length payloads.
+        Only writes keys that don't already exist to avoid redundant writes.
+        """
+        key_strs = [self._logical_anchor_key(key) for key in keys]
+        exist_result = self._batch_exist(key_strs)
+        results = [state == 1 for state in exist_result]
+
+        # Mooncake put() rejects zero-length payloads, so use one byte
+        marker = b"1"
+        for i, key in enumerate(key_strs):
+            if results[i]:
+                continue
+            results[i] = self.store.put(key, marker) == 0
+        return results
+
+    def _batch_get_logical_anchor(self, keys: List[str]) -> List[bool]:
+        """Check marker key existence for logical anchor pages."""
+        key_strs = [self._logical_anchor_key(key) for key in keys]
+        return [state == 1 for state in self._batch_exist(key_strs)]
+
     def get_stats(self):
         storage_metrics = StorageMetrics()
         storage_metrics.prefetch_pgs.extend(self.prefetch_pgs)

test/registered/unit/mem_cache/test_draft_pool_graceful_degradation.py

@@ -0,0 +1,180 @@
+"""Tests for draft pool graceful degradation with storage backend."""
+
+import pytest
+from unittest.mock import Mock, MagicMock, patch
+import logging
+
+
+class TestDraftPoolGracefulDegradation:
+    """Test draft pool registration graceful degradation."""
+
+    @pytest.fixture
+    def cache_controller(self):
+        """Create a minimal HiCacheController instance for testing."""
+        from sglang.srt.managers.cache_controller import HiCacheController
+
+        ctrl = HiCacheController.__new__(HiCacheController)
+        ctrl.enable_storage = True
+        ctrl.has_draft = False
+        ctrl.storage_backend_type = "mooncake"
+        ctrl.storage_config = Mock()
+        ctrl.storage_config.should_split_heads = False
+        ctrl.draft_page_get_func = None
+        ctrl.draft_page_set_func = None
+        return ctrl
+
+    def test_maybe_register_draft_no_storage(self, cache_controller):
+        """Test that registration is skipped when storage is disabled."""
+        cache_controller.enable_storage = False
+        cache_controller.has_draft = True
+
+        cache_controller._maybe_register_draft_with_storage()
+        assert cache_controller.draft_page_get_func is None
+        assert cache_controller.draft_page_set_func is None
+
+    def test_maybe_register_draft_no_draft(self, cache_controller):
+        """Test that registration is skipped when draft is not enabled."""
+        cache_controller.enable_storage = True
+        cache_controller.has_draft = False
+
+        cache_controller._maybe_register_draft_with_storage()
+        assert cache_controller.draft_page_get_func is None
+        assert cache_controller.draft_page_set_func is None
+
+    def test_maybe_register_draft_mooncake_success(self, cache_controller):
+        """Test successful Mooncake draft pool registration wires up v2 funcs."""
+        cache_controller.enable_storage = True
+        cache_controller.has_draft = True
+        cache_controller.storage_backend_type = "mooncake"
+
+        mock_draft_pool = Mock()
+        mock_draft_pool.size = 1024
+        mock_draft_pool.page_size = 64
+        mock_draft_pool.is_mooncake_registerable = True
+        cache_controller.mem_pool_host_draft = mock_draft_pool
+
+        mock_backend = Mock()
+        mock_backend.register_mem_host_pool_v2 = Mock()
+        cache_controller.storage_backend = mock_backend
+
+        cache_controller._maybe_register_draft_with_storage()
+
+        mock_backend.register_mem_host_pool_v2.assert_called_once()
+        assert cache_controller.draft_page_get_func == cache_controller._draft_page_get_v2
+        assert cache_controller.draft_page_set_func == cache_controller._draft_page_set_v2
+
+    def test_maybe_register_draft_mooncake_non_registerable(self, cache_controller, caplog):
+        """Test graceful degradation when draft pool is not Mooncake-registerable (EAGLE + mmap)."""
+        cache_controller.enable_storage = True
+        cache_controller.has_draft = True
+        cache_controller.storage_backend_type = "mooncake"
+
+        mock_draft_pool = Mock()
+        mock_draft_pool.size = 1024
+        mock_draft_pool.page_size = 64
+        mock_draft_pool.is_mooncake_registerable = False
+        mock_draft_pool.allocator = Mock()
+        mock_draft_pool.allocator.__class__.__name__ = "HostTensorAllocator"
+        cache_controller.mem_pool_host_draft = mock_draft_pool
+
+        mock_backend = Mock()
+        cache_controller.storage_backend = mock_backend
+
+        with caplog.at_level(logging.WARNING):
+            cache_controller._maybe_register_draft_with_storage()
+
+        # Registration should not be attempted
+        mock_backend.register_mem_host_pool_v2.assert_not_called()
+        assert cache_controller.draft_page_get_func is None
+        assert cache_controller.draft_page_set_func is None
+        # Should log warning about non-Mooncake-compatible allocator
+        assert any(
+            "not Mooncake-compatible" in record.message
+            for record in caplog.records
+        )
+
+    def test_maybe_register_draft_split_heads_disabled(self, cache_controller, caplog):
+        """Test that split_heads config disables draft L3 on Mooncake."""
+        cache_controller.enable_storage = True
+        cache_controller.has_draft = True
+        cache_controller.storage_backend_type = "mooncake"
+        cache_controller.storage_config.should_split_heads = True
+
+        mock_draft_pool = Mock()
+        mock_draft_pool.is_mooncake_registerable = True
+        cache_controller.mem_pool_host_draft = mock_draft_pool
+        cache_controller.storage_backend = Mock()
+
+        with caplog.at_level(logging.WARNING):
+            cache_controller._maybe_register_draft_with_storage()
+
+        assert cache_controller.draft_page_get_func is None
+        assert cache_controller.draft_page_set_func is None
+        assert any("should_split_heads" in record.message for record in caplog.records)
+
+    def test_maybe_register_draft_generic_backend(self, cache_controller):
+        """Test generic backends wire up generic funcs."""
+        cache_controller.enable_storage = True
+        cache_controller.has_draft = True
+        cache_controller.storage_backend_type = "hf3fs"
+
+        mock_draft_pool = Mock()
+        mock_draft_pool.is_mooncake_registerable = True
+        cache_controller.mem_pool_host_draft = mock_draft_pool
+        cache_controller.storage_backend = Mock()
+
+        with pytest.MonkeyPatch.context() as m:
+            # hf3fs is in the unsupported list, should warn and leave funcs None
+            cache_controller._maybe_register_draft_with_storage()
+            assert cache_controller.draft_page_get_func is None
+            assert cache_controller.draft_page_set_func is None
+
+    def test_maybe_register_draft_generic_fallback(self, cache_controller):
+        """Test unknown backends fall through to generic funcs."""
+        cache_controller.enable_storage = True
+        cache_controller.has_draft = True
+        cache_controller.storage_backend_type = "some_new_backend"
+
+        mock_draft_pool = Mock()
+        mock_draft_pool.is_mooncake_registerable = True
+        cache_controller.mem_pool_host_draft = mock_draft_pool
+        cache_controller.storage_backend = Mock()
+
+        cache_controller._maybe_register_draft_with_storage()
+
+        assert cache_controller.draft_page_get_func == cache_controller._draft_page_get_generic
+        assert cache_controller.draft_page_set_func == cache_controller._draft_page_set_generic
+
+    def test_draft_page_set_no_func_is_noop(self, cache_controller):
+        """Test _draft_page_set is a no-op when draft_page_set_func is None."""
+        cache_controller.draft_page_set_func = None
+        # Should not raise
+        cache_controller._draft_page_set(["hash1"], [0, 1, 2])
+
+    def test_draft_page_get_no_func_is_noop(self, cache_controller):
+        """Test _draft_page_get is a no-op when draft_page_get_func is None."""
+        cache_controller.draft_page_get_func = None
+        # Should not raise
+        cache_controller._draft_page_get(["hash1"], [0, 1, 2])
+
+    def test_draft_page_set_exception_swallowed(self, cache_controller):
+        """Test _draft_page_set swallows exceptions (best-effort)."""
+        def failing_func(hashes, indices):
+            raise RuntimeError("storage backend error")
+
+        cache_controller.draft_page_set_func = failing_func
+        # Should not raise
+        cache_controller._draft_page_set(["hash1"], [0, 1, 2])
+
+    def test_draft_page_get_exception_swallowed(self, cache_controller):
+        """Test _draft_page_get swallows exceptions (best-effort)."""
+        def failing_func(hashes, indices):
+            raise RuntimeError("storage backend error")
+
+        cache_controller.draft_page_get_func = failing_func
+        # Should not raise
+        cache_controller._draft_page_get(["hash1"], [0, 1, 2])
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])