[mem_cache] Move the unified-memory allocators into allocator/ and split the composites out (#38072)

test/registered/unit/mem_cache/test_full_loc_fast_path.py

@@ -170,7 +170,7 @@ class TestUnifiedSWATombstoneClamp(unittest.TestCase):
     """
 
     def _make_bare_pool(self, page_size, v2p, multiplier=1):
-        from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
+        from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
         from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
 
         # A real sub-allocator (not a stand-in): the translation reads its v2p

test/registered/unit/mem_cache/test_kv_index_translator.py

@@ -48,10 +48,10 @@ from types import SimpleNamespace
 import torch
 from test_multi_ended_allocator import _FakeUnifiedSWAKVPool
 
-from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables
-from sglang.srt.mem_cache.multi_ended_allocator import (
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
     UnifiedSWATokenToKVPoolAllocator,
 )
+from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool
 

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -30,11 +30,15 @@ import unittest
 
 import torch
 
-from sglang.srt.mem_cache.multi_ended_allocator import (
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
+    UnifiedSWATokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.allocator.unified_mamba import (
+    UnifiedMambaTokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     FloatMultiEndedAllocator,
     MultiEndedAllocator,
-    UnifiedMambaTokenToKVPoolAllocator,
-    UnifiedSWATokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.unified_memory_pool import (
     MambaSubPoolSpec,
@@ -1263,7 +1267,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
 
     # 7. SWA composite joint byte-budget in page units.
     def test_paged_swa_joint_byte_budget(self):
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
             UnifiedSWATokenToKVPoolAllocator,
         )
 
@@ -1334,7 +1338,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
     # stays -1 and `translate_kv_loc(virt_token)` returns negative token
     # ids → CUDA OOB in the Triton attention kernel.
     def test_paged_alloc_extend_binds_v2p_p2v(self):
-        from sglang.srt.mem_cache import multi_ended_allocator as mea_mod
+        from sglang.srt.mem_cache.allocator import unified_sub_pool as mea_mod
 
         _, full_alloc, _, _, _ = self._build()
         PS = self.PAGE_SIZE
@@ -1417,7 +1421,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
     # prefix's tail page (num_new_pages == 0), but the page-wrapping case
     # must update tables.
     def test_paged_alloc_decode_binds_v2p_p2v_on_page_wrap(self):
-        from sglang.srt.mem_cache import multi_ended_allocator as mea_mod
+        from sglang.srt.mem_cache.allocator import unified_sub_pool as mea_mod
 
         _, full_alloc, _, _, _ = self._build()
         PS = self.PAGE_SIZE
@@ -1496,7 +1500,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
     # (the common case — the decode token reuses the prefix's tail page)
     # must NOT advance the watermark and NOT touch v2p / p2v.
     def test_paged_alloc_decode_no_op_when_no_new_page(self):
-        from sglang.srt.mem_cache import multi_ended_allocator as mea_mod
+        from sglang.srt.mem_cache.allocator import unified_sub_pool as mea_mod
 
         _, full_alloc, _, _, _ = self._build()
         PS = self.PAGE_SIZE
@@ -1699,7 +1703,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
     # `full_available_size() + allocated_tokens == static_cap` must hold for
     # the SWA composite.
     def test_paged_swa_full_available_size_in_tokens(self):
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
             UnifiedSWATokenToKVPoolAllocator,
         )
 
@@ -1785,7 +1789,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
     # (`#full token`, `full token usage`) and would have crashed Mamba+radix
     # if radix weren't auto-downgraded to page=1.
     def test_paged_mamba_size_in_tokens(self):
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_mamba import (
             UnifiedMambaTokenToKVPoolAllocator,
         )
 
@@ -1887,7 +1891,7 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
     # The instance methods in production wrap this helper, so the same
     # math is covered.
     def test_paged_pool_translate_helper_returns_physical_tokens(self):
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
             UnifiedSWATokenToKVPoolAllocator,
         )
         from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -3440,7 +3444,7 @@ class TestDcpWidening(unittest.TestCase):
                 self.assertTrue(bool((written[owned] > 0).all()))
 
     def _build_composite(self, *, page_size):
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_mamba import (
             UnifiedMambaTokenToKVPoolAllocator,
         )
 

test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py

@@ -43,7 +43,9 @@ import unittest
 import torch
 from test_multi_ended_allocator import _FakeUnifiedSWAKVPool  # sibling fixture
 
-from sglang.srt.mem_cache.multi_ended_allocator import UnifiedSWATokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
+    UnifiedSWATokenToKVPoolAllocator,
+)
 from sglang.srt.mem_cache.unified_cache.cache_action import RecoverSWAWithLockedFull
 from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
 from sglang.srt.mem_cache.unified_cache.components.swa_component import SWAComponent

test/registered/unit/mem_cache/test_unified_byte_accounting.py

@@ -42,7 +42,7 @@ from test_multi_ended_allocator import (
     TestUnifiedSWATokenToKVPoolAllocator as _SwaFixture,
 )
 
-from sglang.srt.mem_cache import multi_ended_allocator as mea
+from sglang.srt.mem_cache.allocator import unified_sub_pool as mea
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=15, suite="base-a-test-cpu")

test/registered/unit/mem_cache/test_unified_free_no_host_sync.py

@@ -44,7 +44,8 @@ from unittest import mock
 import torch
 from test_multi_ended_allocator import TestPagedMultiEndedAllocator as _PagedFixture
 
-from sglang.srt.mem_cache import multi_ended_allocator as mea
+from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba
+from sglang.srt.mem_cache.allocator import unified_sub_pool as mea
 from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -87,14 +88,18 @@ _TOMBSTONE_METHODS = [
 _NO_SYNC_TOMBSTONE_FORMS = ("index_fill_", "free_unbind_inplace")
 
 
+_UNIFIED_MODULES = (mea, unified_mamba, unified_hybrid_swa)
+
+
 def _allocators_in_module():
-    """Every allocator class DEFINED in multi_ended_allocator (not imported)."""
+    """Every allocator class DEFINED in the unified allocator modules (not imported)."""
     return sorted(
         (
             c
-            for c in vars(mea).values()
+            for mod in _UNIFIED_MODULES
+            for c in vars(mod).values()
             if isinstance(c, type)
-            and c.__module__ == mea.__name__
+            and c.__module__ == mod.__name__
             and "Allocator" in c.__name__
         ),
         key=lambda c: c.__name__,
@@ -380,8 +385,8 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
     def test_all_overridden(self):
         for cls in (
             mea.MultiEndedAllocator,
-            mea.UnifiedMambaTokenToKVPoolAllocator,
-            mea.UnifiedSWATokenToKVPoolAllocator,
+            unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
+            unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator,
         ):
             with self.subTest(cls=cls.__name__):
                 self.assertIsNot(
@@ -399,8 +404,8 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
         segment free, or `free_segment` raises inside a group."""
         for cls in (
             mea.MultiEndedAllocator,
-            mea.UnifiedMambaTokenToKVPoolAllocator,
-            mea.UnifiedSWATokenToKVPoolAllocator,
+            unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
+            unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator,
         ):
             with self.subTest(cls=cls.__name__):
                 self.assertIn("free_page_reps_group", inspect.getsource(cls))
@@ -462,7 +467,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
             def attach_allocators(self, **kwargs):
                 pass
 
-        return mea.UnifiedSWATokenToKVPoolAllocator(
+        return unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator(
             unified_buffer=pool,
             kvcache=_KV(pool),
             device="cpu",

test/registered/unit/mem_cache/test_unified_handout_zeroing.py

@@ -6,7 +6,7 @@ import unittest
 
 import torch
 
-from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
+from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
 from sglang.srt.mem_cache.unified_memory_pool import (
     MambaSubPoolSpec,
     MLASubPoolSpec,

test/registered/unit/mem_cache/test_unified_mla_views.py

@@ -41,11 +41,11 @@ import unittest
 
 import torch
 
+from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
 from sglang.srt.mem_cache.layout.page_major import (
     build_mla_views,
     mla_entry_bytes,
 )
-from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
 from sglang.srt.mem_cache.unified_memory_pool import (
     MambaSubPoolSpec,
     MLASubPoolSpec,

test/registered/unit/mem_cache/test_unified_swa_shared_virtual_ids.py

@@ -30,7 +30,7 @@ import unittest
 import torch
 from test_swa_locked_full_recover_unified import _DEV, _FakeUnifiedSWAKVPool
 
-from sglang.srt.mem_cache.multi_ended_allocator import (
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
     UnifiedSWATokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool

test/registered/unit/mem_cache/test_unified_tri_pool.py

@@ -39,11 +39,11 @@ import unittest
 
 import torch
 
-import sglang.srt.mem_cache.multi_ended_allocator as mea
-from sglang.srt.mem_cache.multi_ended_allocator import (
-    FloatMultiEndedAllocator,
+import sglang.srt.mem_cache.allocator.unified_sub_pool as mea
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
     UnifiedMambaSWATokenToKVPoolAllocator,
 )
+from sglang.srt.mem_cache.allocator.unified_sub_pool import FloatMultiEndedAllocator
 from sglang.srt.mem_cache.unified_memory_pool import (
     MambaSubPoolSpec,
     MHASubPoolSpec,
@@ -349,7 +349,7 @@ class TestUnifiedTriPool(unittest.TestCase):
         sa = allocator.swa_attn_allocator
         holes = sa._hole_pages()
         self.assertGreater(holes, 0)
-        from sglang.srt.mem_cache.multi_ended_allocator import _relieve_for_alloc
+        from sglang.srt.mem_cache.allocator.unified_sub_pool import _relieve_for_alloc
 
         _relieve_for_alloc(allocator, 1)
         self.assertEqual(sa._hole_pages(), holes)  # holes are assets, not backlog
@@ -864,7 +864,7 @@ class TestTriDeferredAbsorption(unittest.TestCase):
         alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS)
         self.assertGreater(sa._hole_pages(), 0)
         moves_before = len(sa._inverse_history)
-        from sglang.srt.mem_cache.multi_ended_allocator import _relieve_for_alloc
+        from sglang.srt.mem_cache.allocator.unified_sub_pool import _relieve_for_alloc
 
         _relieve_for_alloc(alloc, 1)  # the ladder
         self.assertEqual(sa._hole_pages(), 0)  # rung 0 ran