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

[stmatengss]

This PR unifies the handling of two Mooncake registration failure modes that share the same root cause: pools that cannot be registered with Mooncake's zero-copy interface.

Motivation

• EAGLE draft pools (PR fix: graceful degradation when EAGLE draft pool fails to register with Mooncake standalone storage #29035): Draft pools using mmap allocator fail during register_buffer() because the buffer is outside Mooncake's shared memory segment
• DeepSeek-V4 logical anchors (PR Fix Mooncake HiCache registration for DeepSeek-V4 logical KV anchor #26649): Logical anchor pools have no physical buffer (kv_buffer=None), causing registration to fail

Both currently require separate workarounds or cause crashes.

Modifications

Introduce a unified is_mooncake_registerable property that allows pools to declare whether they support Mooncake registration:

1. Allocator compatibility check: Added is_mooncake_compatible property to HostTensorAllocator (False for mmap, True for MooncakeHostTensorAllocator)
2. Property delegation: HostKVCache.is_mooncake_registerable delegates to self.allocator.is_mooncake_compatible
3. Logical anchor override: LogicalHostPool overrides to False (no buffer exists)
4. Proactive registration check: _maybe_register_draft_with_storage() checks the property before attempting registration, avoiding exceptions entirely

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #28534083240
Latest PR Test (Extra): ❌ Run #28534083143

[gemini-code-assist]

Code Review

This pull request introduces the is_mooncake_registerable property for host pools and is_mooncake_compatible for allocators to determine Mooncake compatibility. It also implements logical anchor handling in MooncakeStore using marker keys for existence tracking and adds corresponding unit tests. The review feedback highlights three key improvements: wrapping self.store.put in a try-except block to prevent unhandled exceptions from crashing the backup thread, fixing an AttributeError in the unit tests caused by attempting to write to the read-only __class__.__name__ attribute, and using getattr to safely access is_mooncake_compatible on the allocator.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The self.store.put call is executed in a loop and can raise exceptions if there are network or storage issues. Since the backup thread does not catch general exceptions, any unhandled exception here will crash the backup thread. Wrapping the put call in a try-except block ensures robustness.

    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
            try:
                results[i] = self.store.put(key, marker) == 0
            except Exception as e:
                logger.error("Failed to put logical anchor key %s: %s", key, e)
                results[i] = False
        return results

[gemini-code-assist]

In Python, the __name__ attribute of a class/type object is read-only. Attempting to write to mock_draft_pool.allocator.__class__.__name__ will raise an AttributeError: attribute '__name__' of 'type' objects is not writable and cause the test to fail. Using a simple dummy class instead of a mock avoids this issue.

Suggested change

	mock_draft_pool.allocator = Mock()
	mock_draft_pool.allocator.__class__.__name__ = "HostTensorAllocator"
	class HostTensorAllocator:
	pass
	mock_draft_pool.allocator = HostTensorAllocator()

[gemini-code-assist]

To prevent potential AttributeError if self.allocator is None or does not have the is_mooncake_compatible attribute, it is safer to use getattr with a default value of False.

Suggested change

	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
	@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 getattr(self.allocator, "is_mooncake_compatible", False)