[Test] Fix optimistic prefill disaggregation test after mamba radix cache removal (#40469)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

test/registered/disaggregation/test_disaggregation_optimistic_prefill.py

@@ -27,9 +27,10 @@ from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
-from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool
 from sglang.srt.mem_cache.radix_cache import RadixKey
+from sglang.srt.mem_cache.unified_cache.components.base import ComponentType
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import get_device
@@ -354,16 +355,8 @@ class TestOptimisticPrefillMambaAdmission(CustomTestCase):
 
 
 class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
-    """Optimistic retry cleanup must not treat the donated Mamba checkpoint
-    as a second donation.
-
-    Bug mechanism: the retry path first inserts the unfinished prefix, which
-    donates the tracked checkpoint and clears ``mamba_last_track_seqlen``.
-    Releasing with ``is_insert=True`` afterwards re-enters the donation path
-    with the cleared marker, inserting a zero-length radix entry that pins a
-    clone of the request's live state. Releasing with ``is_insert=False``
-    retains the donated prefix/checkpoint and frees only the uncached tail
-    and the request-owned Mamba buffers.
+    """An optimistic-prefill retry leaves the donated prefix and exactly one
+    Mamba checkpoint in the tree -- not zero, and not a second pinned clone.
     """
 
     SIZE = 128
@@ -373,7 +366,7 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
 
     def _setup_mamba_tree(self):
         server_args = ServerArgs(model_path="dummy", page_size=1)
-        # MambaRadixCache reads mamba_cache_chunk_size, whose property
+        # The mamba component reads mamba_cache_chunk_size, whose property
         # otherwise loads the HF config for the dummy model.
         server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
         set_global_server_args_for_scheduler(server_args)
@@ -427,13 +420,14 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
             kvcache=pool,
             need_sort=False,
         )
-        tree = MambaRadixCache(
+        tree = UnifiedRadixCache(
             params=CacheInitParams(
                 disable=False,
                 req_to_token_pool=req_to_token_pool,
                 token_to_kv_pool_allocator=allocator,
                 page_size=1,
                 enable_mamba_extra_buffer=True,
+                tree_components=(ComponentType.FULL, ComponentType.MAMBA),
             )
         )
         return tree, allocator, req_to_token_pool
@@ -459,7 +453,7 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
         req.kv.kv_committed_len = len(self.PROMPT)
         req.kv.kv_allocated_len = len(self.PROMPT)
         req.kv.mamba_last_track_seqlen = self.TRACK_SEQLEN
-        req.last_node = tree.root_node
+        req.last_node = tree.root_node_handle()
 
         scheduler = SimpleNamespace(
             tree_cache=tree,
@@ -478,13 +472,15 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
                 scheduler, req
             )
 
-        # The donated prefix and exactly one checkpoint stay in the tree; the
-        # old double-donation path either asserts or pins a second state.
         self.assertEqual(tree.total_size(), (self.TRACK_SEQLEN, 1))
         match = tree.match_prefix(
             MatchPrefixParams(key=RadixKey(array("q", self.PROMPT)))
         )
-        self.assertIsNotNone(match.last_device_node.mamba_value)
+        self.assertIsNotNone(
+            tree.tree_core.get_component_device_value(
+                match.last_device_node, ComponentType.MAMBA
+            )
+        )
 
         # Only the uncached tail and the request-owned Mamba buffers are
         # freed; the tree keeps the donated checkpoint slot.