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

python/sglang/srt/managers/schedule_policy.py

@@ -53,6 +53,12 @@ from sglang.srt.mem_cache.allocator.swa import (
     PureSWATokenToKVPoolAllocator,
     SWATokenToKVPoolAllocator,
 )
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
+    UnifiedMambaSWATokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.allocator.unified_mamba import (
+    UnifiedMambaTokenToKVPoolAllocator,
+)
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
     InitLoadBackParams,
@@ -60,10 +66,6 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchPrefixParams,
     zero_match_result,
 )
-from sglang.srt.mem_cache.multi_ended_allocator import (
-    UnifiedMambaSWATokenToKVPoolAllocator,
-    UnifiedMambaTokenToKVPoolAllocator,
-)
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
 
 if TYPE_CHECKING:

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -21,11 +21,11 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import (
     SchedulerPoolStatsObserver,
 )
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
-from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
-from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
-from sglang.srt.mem_cache.multi_ended_allocator import (
     UnifiedMambaSWATokenToKVPoolAllocator,
 )
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.observability.scheduler_stage_metrics import (
     SCHEDULER_STAGE_SANITY_CHECK_CACHE,
     SchedulerStageMetricsRecorder,

python/sglang/srt/managers/scheduler_components/pool_stats_observer.py

@@ -11,7 +11,7 @@ from typing import (
     Tuple,
 )
 
-from sglang.srt.mem_cache.multi_ended_allocator import (
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
     UnifiedMambaSWATokenToKVPoolAllocator,
 )
 

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

@@ -52,36 +52,26 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
         return self.size
 
     # -- scheduler-facing capacity hooks --
-    # The scheduler calls these UNCONDITIONALLY (zero feature branches on its
+    # The scheduler calls these unconditionally, with no allocator-type branches
-    # side); the defaults reproduce the historical token behavior exactly, and
+    # on its side; byte-accounted composites override the token-count defaults.
-    # unified composites override them with byte-denominated logic.
 
     def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None:
-        """Ask the prefix cache to evict unlocked entries until this allocator
+        """Evict unlocked prefix-cache entries until this allocator can serve
-        can serve ``num_tokens`` (or nothing evictable remains). Default = the
+        ``num_tokens`` or nothing evictable remains."""
-        shared token-count eviction; joint-byte composites override (evicting
-        one multi-lifetime tree node frees bytes on several sides at once).
-        """
         from sglang.srt.mem_cache.common import evict_from_tree_cache
 
         evict_from_tree_cache(tree_cache, num_tokens)
 
     def check_decode_capacity(self, *, num_tokens: int, tree_cache) -> bool:
-        """Whether the NEXT decode step's ``num_tokens`` allocation fits,
+        """Whether the next decode step's ``num_tokens`` allocation fits after
-        evicting reclaimable cache first. The retract loop converges on this
+        evicting reclaimable cache. The retract loop converges on this same
-        same check, so allocator-side shortfalls retract gracefully instead of
+        check, so a shortfall here retracts instead of failing in alloc."""
-        tripping fail-loud alloc errors. Default reproduces the historical
-        ``ScheduleBatch.check_decode_mem`` body; unified composites override
-        with byte gates + per-step reservations of their own.
-        """
         self.evict_to_free_tokens(tree_cache, num_tokens)
         return self.available_size() >= num_tokens
 
     def verify_byte_accounting(self) -> list:
-        """Idle-time conservation diagnostic: recompute this allocator's
+        """Idle-time diagnostic: recompute byte/slot accounting and return
-        byte/slot accounting and return human-readable violation strings
+        violation strings, empty when healthy. Static pools have no byte model."""
-        (empty == healthy). Default: static pools have no byte model.
-        """
         return []
 
     def debug_print(self) -> str:
@@ -126,11 +116,9 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
     def translate_kv_indices_for_transfer(
         self, kv_indices: torch.Tensor
     ) -> torch.Tensor:
-        """Token ids as the PD-disaggregation transfer engine addresses them.
+        """Token ids as the PD transfer engine addresses them. Identity here
-
+        because a static pool's ids index its registered buffers directly;
-        Identity here: a static pool's token ids index its registered buffers
+        virtual-id pools must override."""
-        directly. Virtual-id pools must override.
-        """
         return kv_indices
 
     def get_cpu_copy(self, indices, mamba_indices=None):
@@ -166,11 +154,8 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
         raise NotImplementedError()
 
     def free_full(self, free_index: torch.Tensor):
-        """Free slots whose SWA peers the caller already released.
+        """Free full-attention slots whose paired SWA slots the caller already
-
+        released. A single pool has no SWA peer, so this is a plain free()."""
-        A hybrid SWA allocator pairs each full-attention slot with an SWA slot
-        that can die first; this releases the full side alone. A single pool has
-        no peer, so it is a plain free()."""
         self.free(free_index)
 
     def free_segment(self, free_index: torch.Tensor, *, start_pos: int):
@@ -178,7 +163,7 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
 
         In page units the segment is ``[start_pos // ps, ceil(end / ps))``:
         ``start_pos`` sits on a page boundary, the end may fall mid-page, and
-        the whole last page is released. Default: plain free()."""
+        the whole last page is released."""
         assert start_pos % self.page_size == 0, (
             f"segment start {start_pos} is not page-aligned"
         )
@@ -186,18 +171,15 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
 
     def free_segments(self, segments):
         """Free several ``(free_index, start_pos)`` segments of one request's
-        kv row.
+        kv row. Each covers pages ``[start_pos // ps, ceil(end / ps))``; starts
-
+        are page-aligned and consecutive page ranges do not overlap, so every
-        Each segment covers the pages ``[start_pos // ps, ceil(end / ps))``.
+        page is released exactly once."""
-        Starts sit on page boundaries, ends may fall mid-page, and the page
-        ranges of consecutive segments do not overlap -- so in page units the
-        segments are aligned and disjoint, and every page is released once."""
         for free_index, start_pos in self._page_disjoint(segments):
             self.free_segment(free_index, start_pos=start_pos)
 
     def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int):
         """free_full() for a kv-row segment; same start-alignment contract as
-        free_segment(). Default: plain free_full()."""
+        free_segment()."""
         assert start_pos % self.page_size == 0, (
             f"segment start {start_pos} is not page-aligned"
         )

python/sglang/srt/mem_cache/allocator/hisparse.py

@@ -112,11 +112,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         last_loc: torch.Tensor,
         extend_num_tokens: int,
     ):
-        """Allocate only logical indices without hisparse device indices.
+        """Allocate only logical indices without hisparse device indices."""
-
-        Used in the direct-to-host transfer path where KV data is written
-        directly to host memory by the prefill node, skipping GPU staging.
-        """
         return self.logical_attn_allocator.alloc_extend(
             prefix_lens,
             prefix_lens_cpu,
@@ -131,9 +127,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         # clear original reference and isolate the buffer from outside addressing, allocate new buffer if needed
         hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices]
         self.full_to_hisparse_device_index_mapping[allocated_indices] = 0
-        # Filter valid (non-zero) hisparse indices.
+        # Zero means unmapped; after alloc_logical_only the mapping is all zeros.
-        # In the direct-to-host path, mapping is all zeros since no hisparse
-        # device indices were pre-allocated.
         hisparse_indices = hisparse_indices[hisparse_indices > 0]
         if len(hisparse_indices) >= need_size:
             buffer_indices = hisparse_indices[:need_size]
@@ -241,7 +235,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def clear(self):
         self.logical_attn_allocator.clear()
         self.hisparse_attn_allocator.clear()
-        # Note: the last item is -1, we don't clear it, see the comment in __init__
+        # Keep the trailing -1: it is what a last_loc of -1 translates to.
         self.full_to_hisparse_device_index_mapping[:-1].fill_(0)
         self.free_group = None
 

python/sglang/srt/mem_cache/allocator/mamba.py

@@ -28,20 +28,13 @@ import torch
 
 
 class MambaSlotAllocator:
-    """Manages the free-list of Mamba pool slot indices.
+    """Free-list of Mamba pool slot indices. Deliberately not a subclass of
-
+    ``BaseTokenToKVPoolAllocator``: slots are per request, not per token."""
-    Unlike ``BaseTokenToKVPoolAllocator`` which is designed for per-token KV
-    pages, Mamba slots are request-level (typically 1 slot per request).
-    We keep the interface minimal and do NOT inherit the KV base class.
-    """
 
     def __init__(self, size: int, device: str):
         self.size = size
         self.device = device
-        # Active preallocated batch for `alloc_group_begin` / `alloc_group_end`.
+        # Set by alloc_group_begin(); alloc(1) drains it until alloc_group_end().
-        # When non-None, `alloc(1)` consumes the next slot from this iterator
-        # instead of calling `_do_alloc(1)` per request. Reset to None outside
-        # a group window so `alloc` falls through to the per-call path.
         self._alloc_iter: Optional[Iterator] = None
         self.clear()
 
@@ -49,10 +42,8 @@ class MambaSlotAllocator:
         return len(self.free_slots)
 
     def schedulable_available_size(self) -> int:
-        """Planner-facing free count. Identity to ``available_size`` for the
+        """Planner-facing free count. Same as ``available_size`` for a static pool;
-        static pool (slot-count and byte-coordinated views coincide); the shared
+        byte-coordinated allocators return their byte-limited view instead."""
-        ``UnifiedMambaSlotAllocator`` overrides it with the byte-coordinated view.
-        Lets ``alloc_req_slots`` call it uniformly without a getattr fallback."""
         return self.available_size()
 
     def alloc_group_begin(self, num_reqs: int):

python/sglang/srt/mem_cache/allocator/paged.py

@@ -119,11 +119,8 @@ def alloc_extend_naive(
 
 
 class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
-    """
+    """Same interface as `TokenToKVPoolAllocator`, but the indices handed to one
-    An allocator managing the indices to kv cache data.
+    request are always page-aligned.
-
-    This class has the same interface as `TokenToKVPoolAllocator` but the output
-    of one request is always page-aligned.
 
     TODO: fuse last_loc into the kernel.
     """
@@ -141,18 +138,8 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.num_pages = size // page_size
         self.debug_mode = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
 
-        # Pre-warm the torch.unique HIP kernel used in free(). When a request
+        # Pre-warm the torch.unique used by free(): on ROCm the first call
-        # finishes with a prompt that already exists in the radix tree (e.g.
+        # JIT-compiles rocPRIM sort/unique kernels and costs ~200ms.
-        # bench_serving sending the same warmup+measured prompt), the radix
-        # cache's _insert_helper frees the duplicate KV indices via
-        # token_to_kv_pool_allocator.free(value[start:prefix_len]). That call
-        # path runs `torch.unique(free_index // self.page_size)` on a
-        # ~prompt_len-sized int64 tensor. The first such call on AMD ROCm
-        # JIT-compiles rocPRIM sort/unique kernels and costs ~200ms, which
-        # shows up as a mysterious "second-request slow" (Run 1) for
-        # repeated-prompt benchmarks. Running it once at init time moves
-        # that JIT cost to startup. This is a ROCm-only JIT cost, so the
-        # warm-up is gated on _is_hip and skipped on other platforms.
         if _is_hip and torch.cuda.is_available():
             try:
                 _warmup = torch.arange(1024, dtype=torch.int64, device=device)
@@ -298,13 +285,8 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self._debug_check_no_duplicate_pages()
 
     def free_segment(self, free_index: torch.Tensor, *, start_pos: int):
-        """Fixed-shape counterpart of free().
+        """Fixed-shape free(): page-aligned start plus contiguous per-page tokens
-
+        make ``free_index[::page_size]`` hit each page once; no torch.unique sync."""
-        The segment starts on a page boundary and a page's tokens sit
-        consecutively in the kv row, so ``free_index[::page_size]`` is one
-        token from each page the segment covers -- including a partial last
-        page. No torch.unique, whose data-dependent output shape forces a
-        device sync. Contract: see base."""
         if free_index.numel() == 0:
             return
 

python/sglang/srt/mem_cache/allocator/swa.py

@@ -84,9 +84,8 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
                 swa_kv_pool,
                 need_sort,
             )
-        # Note: append one more item of value -1 in the end so -1 maps to -1.
+        # Trailing -1: a last_loc of -1 (no prefix) indexes it, so alloc_extend and
-        # It is needed for the last_loc in alloc_extend, where the first full_last_loc
+        # alloc_decode see -1 on the SWA side as well.
-        # is -1, and we need to map it to swa_last_loc -1 as well.
         self.full_to_swa_index_mapping = torch.cat(
             [
                 torch.zeros(
@@ -233,12 +232,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         extend_num_tokens: int,
         swa_tail_len: int,
     ):
-        """Allocate full KV for the whole extend and SWA KV only for the tail.
+        """Allocate full KV for the whole extend and SWA KV only for the tail."""
-
-        This is used by disaggregated decode preallocation: decode receives full
-        prompt KV for full-attention layers, but only the sliding-window state is
-        transferred for SWA layers.
-        """
         assert self.page_size > 1
         assert len(seq_lens_cpu) == 1, "SWA tail allocation currently supports bs=1"
         assert len(prefix_lens_cpu) == 1
@@ -334,10 +328,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def set_full_to_swa_mapping(
         self, full_indices: torch.Tensor, swa_indices: torch.Tensor
     ) -> None:
-        """Write full_to_swa_index_mapping[full_indices[i]] = swa_indices[i].
-
-        Used by HiCache load-back path to rebuild the mapping after FULL and SWA device alloc.
-        """
         if full_indices.numel() == 0:
             return
         assert full_indices.numel() == swa_indices.numel()

python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py

@@ -0,0 +1,918 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Unified-memory composites for hybrid SWA models: the full-attention and SWA
+sub-pools of one `UnifiedKVPool`, and the tri-pool variant that adds mamba state."""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional, Sequence
+
+import torch
+from torch.profiler import record_function
+
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.unified_sub_pool import (
+    FloatMultiEndedAllocator,
+    MultiEndedAllocator,
+    _chain_byte_accounting_violations,
+    _end_pair_chain,
+    _float_open_short_side,
+    _relieve_for_alloc,
+)
+from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
+from sglang.srt.utils.common import get_num_new_pages
+
+logger = logging.getLogger(__name__)
+
+
+class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
+    """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools).
+
+    One alloc(N) binds N pages on BOTH sides under the same virtual id, so
+    `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check.
+    """
+
+    # Parent's `size` property has no setter but base init does `self.size = size`;
+    # the no-op setter below absorbs that write.
+    @property
+    def size(self) -> int:
+        return min(self._size_full, self._size_swa)
+
+    @size.setter
+    def size(self, value) -> None:
+        pass
+
+    def __init__(
+        self,
+        *,
+        unified_buffer: UnifiedKVPool,
+        kvcache,  # UnifiedSWAKVPool
+        device: str,
+        full_max_total_num_tokens: int,
+        swa_max_total_num_tokens: int,
+        page_size: int = 1,
+        need_sort: bool = False,
+        forward_stream: Optional[torch.cuda.Stream] = None,
+        lazy_compaction: bool = False,
+    ):
+        # Set _size_full / _size_swa BEFORE base init (read during it). STATIC
+        # partition caps -- the slot-conservation value the leak invariant expects.
+        self._size_full = full_max_total_num_tokens
+        self._size_swa = swa_max_total_num_tokens
+        self._full_max_total_num_tokens = full_max_total_num_tokens
+        self._swa_max_total_num_tokens = swa_max_total_num_tokens
+        self.page_size = page_size
+
+        # The parent is inherited only for the isinstance contract: skip its
+        # static-partition sub-pool allocation, which the unified pool replaces.
+        BaseTokenToKVPoolAllocator.__init__(
+            self,
+            size=full_max_total_num_tokens,
+            page_size=page_size,
+            dtype=unified_buffer.mha_spec("full").store_dtype,
+            device=device,
+            kvcache=kvcache,
+            need_sort=need_sort,
+        )
+        self.unified_buffer = unified_buffer
+        self._kvcache = kvcache
+        self.lazy_compaction = lazy_compaction
+
+        self.full_attn_allocator = MultiEndedAllocator(
+            kvcache=kvcache.full_kv_pool,
+            unified_buffer=unified_buffer,
+            sub_pool_name="full",
+            device=device,
+            is_id_owner=True,
+            page_size=page_size,
+            need_sort=need_sort,
+            forward_stream=forward_stream,
+            lazy_compaction=lazy_compaction,
+        )
+        self.swa_attn_allocator = self._build_swa_attn_allocator(
+            kvcache=kvcache.swa_kv_pool,
+            unified_buffer=unified_buffer,
+            device=device,
+            page_size=page_size,
+            need_sort=need_sort,
+            forward_stream=forward_stream,
+            lazy_compaction=lazy_compaction,
+            # swa binds the virtual pages full mints, so it must address
+            # full's whole id space.
+            virtual_num_pages=self.full_attn_allocator.num_virtual_ids,
+        )
+        self._wire_peers()
+
+        # Epoch-keyed memo for the joint capacity view (any chain member's
+        # mutation invalidates -- see `MultiEndedAllocator._chain_capacity_epoch`).
+        self._joint_avail_memo_epoch: Optional[int] = None
+        self._joint_avail_memo_tokens: int = 0
+
+        # The full/SWA KV pools need no allocator wiring (write locations resolved
+        # in attention metadata); the composite keeps allocators for read-path translates.
+        kvcache.attach_allocators(
+            full_allocator=self.full_attn_allocator,
+            swa_allocator=self.swa_attn_allocator,
+        )
+
+        self.free_group = None
+        self.free_page_reps_group: Optional[List[torch.Tensor]] = None
+        self.full_free_group: List[torch.Tensor] = []
+        # Empty (not None) for the leak checker.
+        self.free_pages = torch.empty(0, dtype=torch.int64, device=device)
+        self.release_pages = torch.empty(0, dtype=torch.int64, device=device)
+
+        logger.info(
+            "[unified-memory-pool] UnifiedSWATokenToKVPoolAllocator ready: "
+            "full max_slots=%d (min_slot_index=%d, entry_bytes=%d), "
+            "swa max_slots=%d (min_slot_index=%d, entry_bytes=%d), "
+            "static caps full=%d swa=%d, joint available=%d",
+            self.full_attn_allocator.max_slots,
+            self.full_attn_allocator.min_slot_index,
+            self.full_attn_allocator.entry_bytes,
+            self.swa_attn_allocator.max_slots,
+            self.swa_attn_allocator.min_slot_index,
+            self.swa_attn_allocator.entry_bytes,
+            self._full_max_total_num_tokens,
+            self._swa_max_total_num_tokens,
+            self.available_size(),
+        )
+
+    # -- construction hooks (the tri-pool subclass overrides both) --
+
+    def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator:
+        """The swa sub-allocator: an END pool in the 2-pool pair."""
+        return MultiEndedAllocator(
+            sub_pool_name="swa",
+            is_id_owner=False,  # non-owner; consumes virtuals minted by full
+            **kwargs,
+        )
+
+    def _wire_peers(self) -> None:
+        self.full_attn_allocator.bind_peer(self.swa_attn_allocator)
+        self.swa_attn_allocator.bind_peer(self.full_attn_allocator)
+
+    # -- capacity reporting (three-way split) --
+
+    def available_size(self) -> int:
+        """Tokens available for `alloc(N)` / `alloc_extend(N)` (TOKENS)."""
+        epoch = self.full_attn_allocator._chain_capacity_epoch()
+        if self._joint_avail_memo_epoch != epoch:
+            self._joint_avail_memo_tokens = self._compute_available_size()
+            self._joint_avail_memo_epoch = epoch
+        return self._joint_avail_memo_tokens
+
+    def _compute_available_size(self) -> int:
+        """Joint byte budget in TOKENS: each composite alloc(1) consumes one
+        full-side AND one swa-side page under the same virtual id."""
+        fa, sa = self.full_attn_allocator, self.swa_attn_allocator
+        e_f = fa.entry_bytes_per_page
+        e_s = sa.entry_bytes_per_page
+        # Direction-agnostic shared gap: the free byte band between the two pools.
+        if fa.grow_direction == "up":
+            gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier())
+        else:
+            gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier())
+        R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
+        R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
+
+        if not self.lazy_compaction:
+            pages_by_bytes = gap_bytes // (e_f + e_s)
+            return min(pages_by_bytes, R_f, R_s) * self.page_size
+
+        H_f = len(fa._free_phys_pages)
+        H_s = len(sa._free_phys_pages)
+
+        K1 = min(H_f, H_s)  # Phase 1: both drain
+
+        # Phase 2: fewer-holes side extends; more-holes side keeps draining.
+        if H_f <= H_s:
+            e_phase2 = e_f
+            K_phase2_max = H_s
+        else:
+            e_phase2 = e_s
+            K_phase2_max = H_f
+        K2_room = K_phase2_max - K1
+        K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room
+        gap_bytes -= K2 * e_phase2
+
+        K3 = gap_bytes // (e_f + e_s)  # Phase 3: both extend
+
+        K_total = K1 + K2 + K3
+        K_total = min(K_total, H_f + R_f, H_s + R_s)  # index-space caps
+        return K_total * self.page_size
+
+    # Slot-conservation views for the leak invariant only; the byte-coordinated
+    # value would flag spurious leaks. `allocated_count()` is in TOKENS.
+    def _conserve_full_available_size(self) -> int:
+        return (
+            self._full_max_total_num_tokens - self.full_attn_allocator.allocated_count()
+        )
+
+    def _conserve_swa_available_size(self) -> int:
+        return (
+            self._swa_max_total_num_tokens - self.swa_attn_allocator.allocated_count()
+        )
+
+    # Per-side views read by scheduling / eviction: the static-conserve cap bounds
+    # the lending side, `schedulable_*` the side grown into the shared gap.
+    def full_available_size(self) -> int:
+        return min(
+            self._conserve_full_available_size(),
+            self.schedulable_full_available_size(),
+        )
+
+    def swa_available_size(self) -> int:
+        return min(
+            self._conserve_swa_available_size(),
+            self.schedulable_swa_available_size(),
+        )
+
+    # Leak-invariant aliases; schedulers take the `min(...)` views above, whose
+    # byte term dips below the conserve cap when bytes are lent to a peer.
+    def conserve_full_available_size(self) -> int:
+        return self._conserve_full_available_size()
+
+    def conserve_swa_available_size(self) -> int:
+        return self._conserve_swa_available_size()
+
+    # Byte-coordinated, realizable-with-compaction views (peer drainable holes
+    # credited -- see `MultiEndedAllocator.schedulable_available_size`).
+    def schedulable_full_available_size(self) -> int:
+        return self.full_attn_allocator.schedulable_available_size()
+
+    def schedulable_swa_available_size(self) -> int:
+        return self.swa_attn_allocator.schedulable_available_size()
+
+    def _flush_targets(self):
+        """Flush ALL members, including ones that are not short themselves: a
+        one-sided hole is unusable, and compacting it yields SHARED gap."""
+        return (self.full_attn_allocator, self.swa_attn_allocator)
+
+    def _ask_float_for_room(self, need_tokens: int) -> None:
+        """No float in a two-END chain -- nothing can slide."""
+        return None
+
+    # `size_full` / `size_swa` are inherited and read the static caps; reporting
+    # `max_slots - 1` here would be ~= full_max + swa_max and over-promise.
+
+    @property
+    def draft_virtual_id_space(self) -> int:
+        return self.full_attn_allocator.max_slots - 1
+
+    def debug_print(self) -> str:
+        return (
+            f"#full-available={self.full_attn_allocator.available_size()}, "
+            f"#swa-available={self.swa_attn_allocator.available_size()}, "
+            f"#joint-available={self.available_size()}"
+        )
+
+    def get_kvcache(self):
+        return self._kvcache
+
+    def translate_kv_loc(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids.
+        ``out=`` writes in place, for cuda-graph buffer stability."""
+        result = self.full_attn_allocator.translate_kv_loc(loc, out=out)
+        return result
+
+    def translate_loc_from_full_to_swa(
+        self,
+        kv_indices: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids."""
+        return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out)
+
+    @property
+    def kernel_page_multiplier(self) -> int:
+        return self.full_attn_allocator.kernel_page_multiplier
+
+    @property
+    def full_v2p_page_table(self) -> torch.Tensor:
+        """Page-level virtual->physical table of the full sub-pool."""
+        return self.full_attn_allocator.virtual_to_physical
+
+    @property
+    def full_p2v_page_table(self) -> torch.Tensor:
+        """Page-level physical->virtual table of the full sub-pool."""
+        return self.full_attn_allocator.physical_to_virtual
+
+    def translate_kv_loc_for_kernel(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full-pool virtual TOKEN ids -> kernel-facing ids."""
+        return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
+
+    def translate_write_loc_for_kernel(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this
+        composite at argument validation, so it coincides with the read translate."""
+        return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
+
+    @property
+    def swa_kernel_page_multiplier(self) -> int:
+        return self.swa_attn_allocator.kernel_page_multiplier
+
+    @property
+    def swa_v2p_page_table(self) -> torch.Tensor:
+        """Page-level virtual->physical table of the SWA sub-pool."""
+        return self.swa_attn_allocator.virtual_to_physical
+
+    # -- alloc --
+
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
+        with record_function("UnifiedSWAAlloc.alloc"):
+            # Joint pre-check. Both sides are mutual peers (each side's compaction
+            # opens gap for the other), so flush BOTH on shortfall.
+            if need_size > self.available_size():
+                if not _relieve_for_alloc(self, need_size):
+                    return None
+            # Snapshot the virtual PAGES full will consume, to bind them on swa too.
+            num_pages = need_size // self.page_size
+            fa = self.full_attn_allocator
+            new_virtual_pages = fa.free_virtual_ids[:num_pages].clone()
+
+            v_tokens = fa.alloc(need_size)
+            # Post-pre-check failure can only be internal-state inconsistency.
+            assert v_tokens is not None, (
+                "UnifiedSWA.alloc: full.alloc returned None after joint "
+                "pre-check passed — internal-state inconsistency"
+            )
+            self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages)
+            return v_tokens
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+    ) -> Optional[torch.Tensor]:
+        """Paged extend; returns virtual TOKEN ids. The same virtual page maps to
+        full- and swa-physical, so swa binds exactly what the full kernel consumed."""
+        with record_function("UnifiedSWAAlloc.alloc_extend"):
+            num_new_pages = get_num_new_pages(
+                seq_lens=seq_lens_cpu,
+                page_size=self.page_size,
+                prefix_lens=prefix_lens_cpu,
+            )
+            need_tokens = num_new_pages * self.page_size
+            if need_tokens > self.available_size():
+                if not _relieve_for_alloc(self, need_tokens):
+                    return None
+
+            # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps
+            # its view after the slice is consumed.
+            fa = self.full_attn_allocator
+            new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone()
+
+            out_indices = fa.alloc_extend(
+                prefix_lens,
+                prefix_lens_cpu,
+                seq_lens,
+                seq_lens_cpu,
+                last_loc,
+                extend_num_tokens,
+                num_new_pages=num_new_pages,
+            )
+            assert out_indices is not None, (
+                "UnifiedSWA.alloc_extend: full.alloc_extend returned None "
+                "after joint pre-check passed — internal-state inconsistency"
+            )
+            self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages)
+            return out_indices  # virtual TOKEN ids
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        """Paged decode: one new token per request, consuming a page only when the
+        decode wraps."""
+        with record_function("UnifiedSWAAlloc.alloc_decode"):
+            num_new_pages = get_num_new_pages(
+                seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True
+            )
+            need_tokens = num_new_pages * self.page_size
+            if need_tokens > self.available_size():
+                if not _relieve_for_alloc(self, need_tokens):
+                    return None
+
+            fa = self.full_attn_allocator
+            new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone()
+
+            out_indices = fa.alloc_decode(seq_lens, seq_lens_cpu, last_loc)
+            assert out_indices is not None, (
+                "UnifiedSWA.alloc_decode: full.alloc_decode returned None "
+                "after joint pre-check passed — internal-state inconsistency"
+            )
+
+            if new_virtual_pages.numel() > 0:
+                self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages)
+
+            return out_indices  # virtual TOKEN ids
+
+    def is_slot_allocated(self, slot: int) -> bool:
+        """Token-slot surface = the full side (which owns the virtual ids)."""
+        return self.full_attn_allocator.is_slot_allocated(slot)
+
+    def allocator_state_str(self) -> str:
+        return self.full_attn_allocator.allocator_state_str()
+
+    # -- free --
+
+    def free(self, free_index: torch.Tensor) -> None:
+        with record_function("UnifiedSWAAlloc.free"):
+            if free_index is None or free_index.numel() == 0:
+                return
+            if self.free_group is not None:
+                self.free_group.append(self._copy_for_free_group(free_index))
+                return
+            # Order is not load-bearing: the per-sub-pool v2p IS the mapping. Only
+            # swa needs the tombstone filter; full owns the ids, so all are bound.
+            v = free_index.detach().to(torch.int64)
+            v_pages = v // self.page_size
+            swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages]
+            # `> 0` strict: -1 = tombstoned, 0 = padding-sink page; both skipped.
+            live_token_mask = swa_v2p_pages > 0
+            live_tokens = v[live_token_mask]
+            if live_tokens.numel() > 0:
+                self.swa_attn_allocator.free(live_tokens)
+            self.full_attn_allocator.free(v)
+            self.full_attn_allocator.clear_inverse_history()
+            self.swa_attn_allocator.clear_inverse_history()
+
+    def free_swa(
+        self, free_index: torch.Tensor, *, start_pos: Optional[int] = None
+    ) -> None:
+        """SWA tombstone path: release swa-physical, keep the virtual id and
+        full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone."""
+        if free_index is None or free_index.numel() == 0:
+            return
+        v = free_index.detach().to(torch.int64)
+        ps = self.page_size
+        # `start_pos` promises a contiguous ascending range starting at that prefix
+        # position, so page reps come from stride arithmetic, not `torch.unique`.
+        if start_pos is not None and ps > 1:
+            reps = self.swa_attn_allocator._page_reps(v, start_pos)
+            # Keep only pages still bound on swa; freeing a tombstoned one would
+            # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink.
+            rep_pages = reps // ps
+            swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages]
+            live_reps = reps[swa_v2p_pages > 0]
+            if live_reps.numel() == 0:
+                return
+            self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps)
+            self.swa_attn_allocator.clear_inverse_history()
+            return
+        v_pages = v // ps
+        # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable).
+        swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages]
+        live = v[swa_v2p_pages > 0]
+        if live.numel() == 0:
+            return
+        if ps == 1:
+            # token == page and the live filter just deduped against the v2p
+            # table, so these ARE unique page ids -- same skip as `_free_lazy`.
+            self.swa_attn_allocator.free(live, _pages=live)
+        else:
+            self.swa_attn_allocator.free(live)
+        self.swa_attn_allocator.clear_inverse_history()
+
+    def free_full(self, free_index: torch.Tensor) -> None:
+        """Release the full-physical page and the virtual id, leaving the swa
+        side alone -- the caller already tombstoned it (`swa.v2p_page == -1`)."""
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.free_group is not None:
+            self.full_free_group.append(self._copy_for_free_group(free_index))
+            return
+        self.full_attn_allocator.free(free_index.detach().to(torch.int64))
+        self.full_attn_allocator.clear_inverse_history()
+
+    def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            # token == page: free_full already frees by exact ids, no dedup.
+            self.free_full(free_index)
+            return
+        # The swa v2p is the mapping, so a tombstoned swa page drops out of the
+        # two-sided segment path by itself; full-only is the same call.
+        self.free_segment(free_index, start_pos=start_pos)
+
+    def set_full_to_swa_mapping(
+        self, full_indices: torch.Tensor, swa_indices: torch.Tensor
+    ) -> None:
+        """No-op stub for HiCache load-back: in shared mode the swa v2p IS the
+        mapping, and HiCache for shared SWA is out of scope."""
+        return
+
+    def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None:
+        # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor.
+        return
+
+    # -- free-group --
+
+    # Not the SWA parent's hooks: those open the parent's paged full allocator
+    # as a free group, and this composite's sub-pools defer on their own.
+    def free_group_begin(self) -> None:
+        BaseTokenToKVPoolAllocator.free_group_begin(self)
+        self.free_page_reps_group = []
+        self.full_free_group = []
+
+    def free_group_end(self) -> None:
+        pending, self.free_page_reps_group = self.free_page_reps_group, None
+        full_free_group, self.full_free_group = self.full_free_group, []
+        BaseTokenToKVPoolAllocator.free_group_end(self)
+        if full_free_group:
+            self.full_attn_allocator.free(torch.cat(full_free_group))
+            self.full_attn_allocator.clear_inverse_history()
+        if pending:
+            self._release_page_reps(pending)
+
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`.
+        Both sides share one page-rep derivation instead of dedup'ing twice."""
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            self.free(free_index)
+            return
+        reps = self.full_attn_allocator._page_reps(
+            free_index.detach().to(torch.int64), start_pos
+        )
+        if self.free_page_reps_group is None:
+            self._release_page_reps((reps,))
+        else:
+            self.free_page_reps_group.append(reps)
+
+    def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None:
+        reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces))
+        v_pages = reps // self.page_size
+        # Same tombstone filter as `free`, but at PAGE granularity (page_size
+        # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink.
+        swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages]
+        live_pages = v_pages[swa_v2p_pages > 0]
+        if live_pages.numel() > 0:
+            self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages)
+        self.full_attn_allocator.free(reps, _pages=v_pages)
+        self.full_attn_allocator.clear_inverse_history()
+        self.swa_attn_allocator.clear_inverse_history()
+
+    def verify_byte_accounting(self) -> List[str]:
+        return (
+            _chain_byte_accounting_violations(
+                _end_pair_chain(self.full_attn_allocator, self.swa_attn_allocator)
+            )
+            + self._joint_capacity_memo_violations()
+        )
+
+    def _joint_capacity_memo_violations(self) -> List[str]:
+        """Idle-time twin of `MultiEndedAllocator._capacity_memo_violations`
+        for the composite joint view. Empty == healthy."""
+        if (
+            self._joint_avail_memo_epoch
+            != self.full_attn_allocator._chain_capacity_epoch()
+        ):
+            return []
+        actual = self._compute_available_size()
+        if self._joint_avail_memo_tokens == actual:
+            return []
+        return [
+            f"[joint] stale available_size memo: "
+            f"cached={self._joint_avail_memo_tokens}, actual={actual}"
+        ]
+
+    def clear(self) -> None:
+        self.full_attn_allocator.clear()
+        self.swa_attn_allocator.clear()
+        self.free_group = None
+        self.free_page_reps_group = None
+        self.full_free_group = []
+
+    # -- Lazy compaction hooks --
+
+    def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None:
+        """Forward the per-batch `forward_done` event to BOTH sub-allocators."""
+        with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"):
+            self.full_attn_allocator.set_latest_forward_done_event(event)
+            self.swa_attn_allocator.set_latest_forward_done_event(event)
+
+    def set_inflight_forward(
+        self,
+        forward_done: torch.cuda.Event,
+        out_cache_loc_virtual: Optional[torch.Tensor],
+    ) -> None:
+        """Hand the forward's metadata to BOTH sub-pools; each materializes its own
+        write-set via its OWN v2p, and the forward writes both sides per token."""
+        with record_function("UnifiedSWAAlloc.set_inflight_forward"):
+            self.full_attn_allocator.set_inflight_forward(
+                forward_done, out_cache_loc_virtual
+            )
+            self.swa_attn_allocator.set_inflight_forward(
+                forward_done, out_cache_loc_virtual
+            )
+
+    def flush_opportunistic(self) -> int:
+        """Non-urgent flush of BOTH sub-allocators; sync-free."""
+        with record_function("UnifiedSWAAlloc.flush_opportunistic"):
+            fa = self.full_attn_allocator
+            sa = self.swa_attn_allocator
+            if (
+                fa._free_phys_pages.numel() == 0
+                and not fa._pending_reuse
+                and sa._free_phys_pages.numel() == 0
+                and not sa._pending_reuse
+            ):
+                return 0
+            return fa.flush_opportunistic() + sa.flush_opportunistic()
+
+
+class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
+    """Tri-pool composite for models with full KV + SWA KV + mamba/conv state
+    (both `mambaish_config` and `is_hybrid_swa`).
+
+    Chain (low byte -> high byte):
+
+        [ mamba/conv (grow-up END) | swa (FLOAT middle) | full (grow-down END) ]
+
+    The ends never relocate, so they take the per-request state pool and the
+    unbounded per-step grower; SWA is window-capped with the cheapest slots to
+    move, and its out-of-window tombstones become float holes recycled in place.
+    Per-request state is served through `mamba_allocator`, wrapped by
+    `UnifiedMambaSlotAllocator`.
+    """
+
+    def __init__(
+        self,
+        *,
+        unified_buffer: UnifiedKVPool,
+        kvcache,  # UnifiedSWAKVPool
+        mamba_kvcache,  # UnifiedMambaPool (req_to_token_pool.mamba_pool)
+        device: str,
+        full_max_total_num_tokens: int,
+        swa_max_total_num_tokens: int,
+        page_size: int = 1,
+        need_sort: bool = False,
+        forward_stream: Optional[torch.cuda.Stream] = None,
+        lazy_compaction: bool = False,
+    ):
+        super().__init__(
+            unified_buffer=unified_buffer,
+            kvcache=kvcache,
+            device=device,
+            full_max_total_num_tokens=full_max_total_num_tokens,
+            swa_max_total_num_tokens=swa_max_total_num_tokens,
+            page_size=page_size,
+            need_sort=need_sort,
+            forward_stream=forward_stream,
+            lazy_compaction=lazy_compaction,
+        )
+        # Per-request state END pool (grow-up; page_size=1 -- state is
+        # per-request, orthogonal to KV paging).
+        self.mamba_allocator = MultiEndedAllocator(
+            kvcache=mamba_kvcache,
+            unified_buffer=unified_buffer,
+            sub_pool_name="mamba",
+            device=device,
+            is_id_owner=True,
+            page_size=1,
+            need_sort=need_sort,
+            forward_stream=forward_stream,
+            lazy_compaction=lazy_compaction,
+        )
+        # Chain wiring: mamba <-> swa(float) <-> full.
+        self.mamba_allocator.bind_high_peer(self.swa_attn_allocator)
+        self.swa_attn_allocator.bind_low_peer(self.mamba_allocator)
+        self.swa_attn_allocator.bind_high_peer(self.full_attn_allocator)
+        self.full_attn_allocator.bind_low_peer(self.swa_attn_allocator)
+
+        # None, not empty: `free_pages is None` is the leak checker's documented
+        # skip contract; its mamba census would mix physical and virtual ids.
+        self.free_pages = None
+        self.release_pages = None
+
+        logger.info(
+            "[unified-memory-pool] UnifiedMambaSWATokenToKVPoolAllocator ready: "
+            "chain=[mamba(up) | swa(float) | full(down)], "
+            "mamba max_slots=%d (entry_bytes=%d), joint available=%d",
+            self.mamba_allocator.max_slots,
+            self.mamba_allocator.entry_bytes,
+            self.available_size(),
+        )
+
+    # -- construction hooks --
+
+    def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator:
+        # The swa side is the FLOAT middle: it never runs the lazy event pipeline
+        # regardless of the composite's flag (frees mark holes, allocs reuse them).
+        kwargs["lazy_compaction"] = False
+        return FloatMultiEndedAllocator(
+            sub_pool_name="swa",
+            is_id_owner=False,  # non-owner; consumes virtuals minted by full
+            **kwargs,
+        )
+
+    def _wire_peers(self) -> None:
+        # Chain wired in __init__ once the mamba end exists.
+        return
+
+    # -- capacity --
+
+    def _compute_available_size(self) -> int:
+        """Joint TOKENS for `alloc(N)`: N costs N full pages AND N swa pages, drawn
+        from DIFFERENT bands -- full extends only into the high band, the float into
+        either side but only ONE per batch alloc. Feasibility is monotone in N, so
+        binary search; the order matches the alloc path (full takes the high band).
+        """
+        fa, sa = self.full_attn_allocator, self.swa_attn_allocator
+        e_f, e_s = fa.entry_bytes_per_page, sa.entry_bytes_per_page
+        # full is grow-down: its chain gap IS the high band.
+        b_high = fa._current_gap_bytes()
+        if sa._is_frontier_transparent():
+            b_low = 0
+        else:
+            b_low = max(
+                0,
+                sa._byte_low_frontier() - sa._chain_high_frontier_below_bytes(),
+            )
+        h_f = len(fa._free_phys_pages) if fa.lazy_compaction else 0
+        h_s = sa._hole_pages()
+        r_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
+        r_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
+
+        def feasible(n: int) -> bool:
+            if n > h_f + r_f or n > h_s + r_s:
+                return False
+            ext_f = max(0, n - h_f)
+            if ext_f * e_f > b_high:
+                return False
+            ext_s = max(0, n - h_s)
+            # On the float's page grid, never in raw bytes: a byte budget
+            # credits a page `take_physical_pages` cannot yield.
+            full_low_after = fa._byte_low_frontier() - ext_f * e_f
+            if sa._is_frontier_transparent():
+                room = sa.pages_in_band(
+                    low_byte=sa._chain_high_frontier_below_bytes(),
+                    high_byte=full_low_after,
+                )
+                return ext_s <= room
+            p_low = sa.pages_in_band(
+                low_byte=sa._chain_high_frontier_below_bytes(),
+                high_byte=sa._byte_low_frontier(),
+            )
+            p_high = sa.pages_in_band(
+                low_byte=sa._byte_high_frontier(),
+                high_byte=full_low_after,
+            )
+            return ext_s <= max(p_low, p_high)
+
+        lo_n, hi_n = 0, min(h_f + r_f, h_s + r_s)
+        while lo_n < hi_n:
+            mid = (lo_n + hi_n + 1) // 2
+            if feasible(mid):
+                lo_n = mid
+            else:
+                hi_n = mid - 1
+        return lo_n * self.page_size
+
+    def _flush_targets(self):
+        """All three members, float FIRST: its zero-copy boundary absorption must
+        land before the deficit math prices a relocation it already covered."""
+        return (
+            self.swa_attn_allocator,
+            self.full_attn_allocator,
+            self.mamba_allocator,
+        )
+
+    def _alloc_demand(self, need_tokens: int):
+        """Demand VECTOR for one composite allocation, in PAGES per band. A token
+        never draws a state slot, so mamba is an explicit 0, not an omission."""
+        need_n = -(-need_tokens // self.page_size)
+        return {
+            self.full_attn_allocator: need_n,
+            self.swa_attn_allocator: need_n,
+            self.mamba_allocator: 0,
+        }
+
+    def _ask_float_for_room(self, need_tokens: int) -> None:
+        """Composite shortfall: hand the demand vector to the shared policy;
+        the float is whichever demanded band floats."""
+        demand = self._alloc_demand(need_tokens)
+        flt = None
+        for b in demand:
+            if isinstance(b, FloatMultiEndedAllocator):
+                flt = b
+        _float_open_short_side(flt, demand)
+
+    def mamba_slot_full_token_cost(self) -> int:
+        """Full-token-equivalents one mamba/conv slot removes from the shared buffer:
+        a tri-pool token costs e_f + e_s bytes, and the quotient is rounded UP."""
+        e_tok = (
+            self.full_attn_allocator.entry_bytes + self.swa_attn_allocator.entry_bytes
+        )
+        return -(-self.mamba_allocator.entry_bytes_per_page // e_tok)
+
+    def debug_print(self) -> str:
+        sa = self.swa_attn_allocator
+        return (
+            super().debug_print()
+            + f", #mamba-available={self.mamba_allocator.available_size()}"
+            + f", swa-float span=[{sa.low_wm_page},{sa.high_wm_page}) "
+            + f"holes={sa._hole_pages()}"
+        )
+
+    # -- lifecycle fanout (adds the mamba end) --
+
+    def clear(self) -> None:
+        super().clear()
+        self.mamba_allocator.clear()
+
+    def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None:
+        super().set_latest_forward_done_event(event)
+        self.mamba_allocator.set_latest_forward_done_event(event)
+
+    def set_inflight_forward(
+        self,
+        forward_done: torch.cuda.Event,
+        out_cache_loc_virtual: Optional[torch.Tensor],
+    ) -> None:
+        # The mamba state is written by the conv kernels, not through
+        # `out_cache_loc`, so its in-flight write-set is None.
+        super().set_inflight_forward(forward_done, out_cache_loc_virtual)
+        self.mamba_allocator.set_inflight_forward(forward_done, None)
+
+    def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None:
+        """Joint-aware eviction: one tri-lifetime node frees bytes on several sides
+        at once, so re-check the JOINT gate instead of the per-side shortfall."""
+        from sglang.srt.mem_cache.common import evict_from_tree_cache
+
+        for _ in range(4):
+            before = self.available_size()
+            if before >= num_tokens:
+                return
+            evict_from_tree_cache(tree_cache, num_tokens)
+            if self.available_size() <= before:
+                return  # no progress
+
+    def verify_byte_accounting(self) -> List[str]:
+        return (
+            _chain_byte_accounting_violations(
+                [
+                    self.mamba_allocator,
+                    self.swa_attn_allocator,
+                    self.full_attn_allocator,
+                ]
+            )
+            + self._joint_capacity_memo_violations()
+        )
+
+    def flush_opportunistic(self) -> int:
+        """Per-step reclaim across the whole chain. The float participates for its
+        deferred boundary absorption, which is where its single D2H is paid."""
+        fa, ma = self.full_attn_allocator, self.mamba_allocator
+        sa = self.swa_attn_allocator
+        if (
+            fa._free_phys_pages.numel() == 0
+            and not fa._pending_reuse
+            and ma._free_phys_pages.numel() == 0
+            and not ma._pending_reuse
+            and sa._free_phys_pages.numel() == 0
+        ):
+            return 0
+        return (
+            fa.flush_opportunistic()
+            + ma.flush_opportunistic()
+            + sa.flush_opportunistic()
+        )

python/sglang/srt/mem_cache/allocator/unified_mamba.py

@@ -0,0 +1,390 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Unified-memory composite for hybrid Mamba models: the full-attention and
+mamba-state end pools of one `UnifiedKVPool`."""
+
+from __future__ import annotations
+
+import logging
+from typing import Callable, List, Optional, Sequence
+
+import torch
+from torch.profiler import record_function
+
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.unified_sub_pool import (
+    MultiEndedAllocator,
+    _chain_byte_accounting_violations,
+    _end_pair_chain,
+)
+from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
+from sglang.srt.runtime_context import get_parallel
+
+logger = logging.getLogger(__name__)
+
+
+class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
+    """Composite allocator for the MHA (full-attn) + Mamba hybrid pair.
+
+    The token-slot surface is the full-attn side; the mamba sub-pool's per-request
+    `alloc(1)` is driven separately by `UnifiedHybridReqToTokenPool`. The two
+    sub-allocators own independent virtual-id spaces.
+    """
+
+    def __init__(
+        self,
+        *,
+        unified_buffer: UnifiedKVPool,
+        kvcache,  # HybridLinearKVPool
+        device: str,
+        page_size: int = 1,
+        need_sort: bool = False,
+        forward_stream: Optional[torch.cuda.Stream] = None,
+        lazy_compaction: bool = False,
+    ):
+        full_max = unified_buffer.max_slots("full")
+        dcp_size = get_parallel().attn_dcp_size
+        super().__init__(
+            size=(full_max - 1) * dcp_size,
+            page_size=page_size * dcp_size,
+            dtype=unified_buffer.spec("full").get_dtype(),
+            device=device,
+            kvcache=kvcache,
+            need_sort=need_sort,
+        )
+        self.unified_buffer = unified_buffer
+        self._kvcache = kvcache
+        # Widened under DCP, matching the full sub-allocator; see its __init__.
+        self.page_size = page_size * dcp_size
+        self.lazy_compaction = lazy_compaction
+
+        # Only FULL shards under DCP; the mamba state is replicated on every rank
+        # and stays page_size=1, orthogonal to the full side's per-token paging.
+        self.full_attn_allocator = MultiEndedAllocator(
+            kvcache=kvcache.full_kv_pool,
+            unified_buffer=unified_buffer,
+            sub_pool_name="full",
+            device=device,
+            is_id_owner=True,
+            page_size=page_size,
+            shards_under_dcp=True,
+            need_sort=need_sort,
+            forward_stream=forward_stream,
+            lazy_compaction=lazy_compaction,
+        )
+        self.mamba_allocator = MultiEndedAllocator(
+            kvcache=kvcache.mamba_pool,
+            unified_buffer=unified_buffer,
+            sub_pool_name="mamba",
+            device=device,
+            is_id_owner=True,
+            page_size=1,  # Mamba state stays slot-granular (1-per-req)
+            need_sort=need_sort,
+            forward_stream=forward_stream,
+            lazy_compaction=lazy_compaction,
+        )
+        self.full_attn_allocator.bind_peer(self.mamba_allocator)
+        self.mamba_allocator.bind_peer(self.full_attn_allocator)
+
+        # `init_unified_mamba_pools` later wraps `self.mamba_allocator` in a
+        # `UnifiedMambaSlotAllocator` owning the v2p translate; the KV pools get no
+        # allocator (write locations resolve in the attention metadata).
+
+        self.free_group = None
+        self.free_page_reps_group: Optional[List[torch.Tensor]] = None
+        # Base init left these None; we use watermark math, not free-lists.
+        self.free_pages = torch.empty(0, dtype=torch.int64, device=device)
+        self.release_pages = torch.empty(0, dtype=torch.int64, device=device)
+
+        logger.info(
+            "[unified-memory-pool] UnifiedMambaTokenToKVPoolAllocator ready: "
+            "full max_slots=%d (min_slot_index=%d, page_size=%d, "
+            "num_pages=%d), mamba max_slots=%d (min_slot_index=%d), "
+            "full_available=%d, mamba_available=%d",
+            self.full_attn_allocator.max_slots,
+            self.full_attn_allocator.min_slot_index,
+            self.full_attn_allocator.page_size,
+            self.full_attn_allocator.num_pages,
+            self.mamba_allocator.max_slots,
+            self.mamba_allocator.min_slot_index,
+            self.full_attn_allocator.available_size(),
+            self.mamba_allocator.available_size(),
+        )
+
+    # -- size: dynamic --
+    @property
+    def size(self) -> int:
+        # TOKENS. MUST use the SAME available view as `available_size()`, so the
+        # available term cancels out of the leak invariant.
+        return (
+            self.full_attn_allocator.schedulable_available_size()
+            + self.full_attn_allocator.allocated_count()
+        )
+
+    @size.setter
+    def size(self, value) -> None:
+        pass  # base init writes here; computed dynamically
+
+    # -- token-slot surface: MHA side --
+
+    # Realizable-with-compaction view, so the retract gate / evict / schedule_policy
+    # do not over-retract while the mamba peer holds drainable holes.
+    def available_size(self) -> int:
+        return self.full_attn_allocator.schedulable_available_size()
+
+    def full_available_size(self) -> int:
+        return self.full_attn_allocator.schedulable_available_size()
+
+    def mamba_slot_full_token_cost(self) -> int:
+        """Full-token-equivalents of shared-gap bytes ONE mamba state consumes; the
+        prefill planner reserves this so admission stays inside the JOINT budget,
+        rounded UP. The `dcp_size` factor is there because that budget is in widened
+        tokens, one of which is `entry_bytes / dcp_size` local bytes.
+        """
+        return -(
+            -self.mamba_allocator.entry_bytes_per_page
+            * get_parallel().attn_dcp_size
+            // self.full_attn_allocator.entry_bytes
+        )
+
+    @property
+    def size_full(self) -> int:
+        # Widened like `size`: a logical token capacity, not a row count.
+        return (self.full_attn_allocator.max_slots - 1) * get_parallel().attn_dcp_size
+
+    @property
+    def draft_virtual_id_space(self) -> int:
+        return self.size_full
+
+    @property
+    def size_mamba(self) -> int:
+        return self.mamba_allocator.max_slots - 1
+
+    def debug_print(self) -> str:
+        return (
+            f"#full-available={self.full_attn_allocator.available_size()}, "
+            f"#mamba-available={self.mamba_allocator.available_size()}"
+        )
+
+    def get_kvcache(self):
+        return self._kvcache
+
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
+        with record_function("UnifiedMambaAlloc.alloc"):
+            return self.full_attn_allocator.alloc(need_size)
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        num_new_pages: Optional[int] = None,
+    ) -> Optional[torch.Tensor]:
+        """Paged extend. Mamba state is per-request (doesn't advance per-token),
+        so forward only to the full sub-allocator."""
+        with record_function("UnifiedMambaAlloc.alloc_extend"):
+            return self.full_attn_allocator.alloc_extend(
+                prefix_lens,
+                prefix_lens_cpu,
+                seq_lens,
+                seq_lens_cpu,
+                last_loc,
+                extend_num_tokens,
+                num_new_pages=num_new_pages,
+            )
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        """Paged decode. Mamba side stays untouched per-decode."""
+        with record_function("UnifiedMambaAlloc.alloc_decode"):
+            return self.full_attn_allocator.alloc_decode(
+                seq_lens, seq_lens_cpu, last_loc
+            )
+
+    def translate_kv_loc(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full-pool virtual TOKEN ids -> physical TOKEN ids; `-1` passes through as
+        `-1` (padding downstream). ``out=`` supports cuda-graph buffer stability."""
+        result = self.full_attn_allocator.translate_kv_loc(loc, out=out)
+        return result
+
+    @property
+    def kernel_page_multiplier(self) -> int:
+        return self.full_attn_allocator.kernel_page_multiplier
+
+    @property
+    def full_v2p_page_table(self) -> torch.Tensor:
+        """Page-level virtual->physical table of the full sub-pool. Kernels that
+        build the MLA block table straight from req_to_token gather through this,
+        then scale by `kernel_page_multiplier` to reach the per-page block."""
+        return self.full_attn_allocator.virtual_to_physical
+
+    @property
+    def full_p2v_page_table(self) -> torch.Tensor:
+        """Page-level physical->virtual table of the full sub-pool."""
+        return self.full_attn_allocator.physical_to_virtual
+
+    def translate_kv_loc_for_kernel(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full-pool virtual TOKEN ids -> kernel-facing ids."""
+        return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out)
+
+    def translate_write_loc_for_kernel(
+        self,
+        loc: torch.Tensor,
+        *,
+        out: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Widened virtual WRITE loc -> DENSE id; see the sub-allocator's copy."""
+        return self.full_attn_allocator.translate_write_loc_for_kernel(loc, out=out)
+
+    def translate_kv_indices_for_transfer(
+        self, kv_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Virtual TOKEN ids -> PHYSICAL token ids for the PD transfer engine.
+        PHYSICAL, not kernel-facing: the transfer registers page ENVELOPES (see
+        `UnifiedMLATokenToKVPool.get_contiguous_buf_infos`)."""
+        # Defensive: `_validate_unified_memory_dcp` rejects this pairing at
+        # argument validation, so reaching it means a config path got past that.
+        assert get_parallel().attn_dcp_size == 1, (
+            "PD-disaggregation transfer with the unified memory pool does not "
+            "support decode context parallelism: the transfer ships whole page "
+            "envelopes, which hold only this rank's shard of each widened page."
+        )
+        return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
+
+    def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None:
+        """Install the PD-disaggregation move gate on both sub-allocators."""
+        assert self.lazy_compaction, (
+            "PD disaggregation with the unified memory pool requires lazy "
+            "compaction (eager free-path compaction moves pages under "
+            "in-flight transfers)."
+        )
+        self.full_attn_allocator.disagg_move_gate = gate
+        self.mamba_allocator.disagg_move_gate = gate
+
+    def is_slot_allocated(self, slot: int) -> bool:
+        return self.full_attn_allocator.is_slot_allocated(slot)
+
+    def allocator_state_str(self) -> str:
+        return self.full_attn_allocator.allocator_state_str()
+
+    def free(self, free_index: torch.Tensor) -> None:
+        with record_function("UnifiedMambaAlloc.free"):
+            if free_index is None or free_index.numel() == 0:
+                return
+            if self.free_group is not None:
+                self.free_group.append(self._copy_for_free_group(free_index))
+                return
+            self.full_attn_allocator.free(free_index)
+            self.full_attn_allocator.clear_inverse_history()
+            self.mamba_allocator.clear_inverse_history()
+
+    def clear(self) -> None:
+        self.full_attn_allocator.clear()
+        self.mamba_allocator.clear()
+        self.free_group = None
+
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`.
+        The mamba sub-pool is slot-granular and untouched by a token free."""
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            self.free(free_index)
+            return
+        reps = self.full_attn_allocator._page_reps(
+            free_index.detach().to(torch.int64), start_pos
+        )
+        if self.free_page_reps_group is None:
+            self._release_page_reps((reps,))
+        else:
+            self.free_page_reps_group.append(reps)
+
+    def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None:
+        reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces))
+        self.full_attn_allocator.free(reps, _pages=reps // self.page_size)
+        self.full_attn_allocator.clear_inverse_history()
+        self.mamba_allocator.clear_inverse_history()
+
+    def verify_byte_accounting(self) -> List[str]:
+        return _chain_byte_accounting_violations(
+            _end_pair_chain(self.mamba_allocator, self.full_attn_allocator)
+        )
+
+    def free_group_begin(self) -> None:
+        super().free_group_begin()
+        self.free_page_reps_group = []
+
+    def free_group_end(self) -> None:
+        pending, self.free_page_reps_group = self.free_page_reps_group, None
+        super().free_group_end()
+        if pending:
+            self._release_page_reps(pending)
+
+    def clear(self) -> None:
+        self.full_attn_allocator.clear()
+        self.mamba_allocator.clear()
+        self.free_group = None
+        self.free_page_reps_group = None
+
+    # -- Lazy compaction hooks --
+
+    def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None:
+        """Forward the per-batch `forward_done` event to BOTH sub-allocators."""
+        with record_function("UnifiedMambaAlloc.set_latest_forward_done_event"):
+            self.full_attn_allocator.set_latest_forward_done_event(event)
+            self.mamba_allocator.set_latest_forward_done_event(event)
+
+    def set_inflight_forward(
+        self,
+        forward_done: torch.cuda.Event,
+        out_cache_loc_virtual: Optional[torch.Tensor],
+    ) -> None:
+        """Hand the forward's metadata to BOTH sub-pools; the mamba state is written
+        by mamba kernels, not `set_kv_buffer`, so its write-set is `None`."""
+        with record_function("UnifiedMambaAlloc.set_inflight_forward"):
+            self.full_attn_allocator.set_inflight_forward(
+                forward_done, out_cache_loc_virtual
+            )
+            self.mamba_allocator.set_inflight_forward(forward_done, None)
+
+    def flush_opportunistic(self) -> int:
+        """Non-urgent flush of BOTH sub-allocators; sync-free."""
+        with record_function("UnifiedMambaAlloc.flush_opportunistic"):
+            fa = self.full_attn_allocator
+            ma = self.mamba_allocator
+            if (
+                fa._free_phys_pages.numel() == 0
+                and not fa._pending_reuse
+                and ma._free_phys_pages.numel() == 0
+                and not ma._pending_reuse
+            ):
+                return 0
+            return fa.flush_opportunistic() + ma.flush_opportunistic()

python/sglang/srt/mem_cache/common.py

@@ -101,9 +101,9 @@ def free_swa_out_of_window_slots(
         free_slots = req_to_token_pool.req_to_token[
             req.kv.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen
         ]
-        # Local import: multi_ended_allocator imports this module lazily for
+        # Local import: the unified allocators import this module lazily for
         # eviction; a module-level import here would be a cycle hazard.
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
             UnifiedSWATokenToKVPoolAllocator,
         )
 

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -47,6 +47,12 @@ from sglang.srt.mem_cache.allocator.swa import (
     PureSWATokenToKVPoolAllocator,
     SWATokenToKVPoolAllocator,
 )
+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.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
 from sglang.srt.mem_cache.memory_pool import (
@@ -64,10 +70,6 @@ from sglang.srt.mem_cache.memory_pool import (
     PageMajorMHATokenToKVPool,
     ReqToTokenPool,
 )
-from sglang.srt.mem_cache.multi_ended_allocator import (
-    UnifiedMambaTokenToKVPoolAllocator,
-    UnifiedSWATokenToKVPoolAllocator,
-)
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (

python/sglang/srt/mem_cache/kv_index_translator.py

@@ -64,11 +64,13 @@ from sglang.kernels.ops.kvcache.kv_read_table import (
     build_kv_read_table,
     build_kv_read_table_packed,
 )
-from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
-from sglang.srt.mem_cache.multi_ended_allocator import (
-    UnifiedMambaTokenToKVPoolAllocator,
     UnifiedSWATokenToKVPoolAllocator,
 )
+from sglang.srt.mem_cache.allocator.unified_mamba import (
+    UnifiedMambaTokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
 from sglang.srt.runtime_context import get_parallel
 
 

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -31,6 +31,9 @@ from sglang.srt.mem_cache.allocator import (
     PagedTokenToKVPoolAllocator,
     TokenToKVPoolAllocator,
 )
+from sglang.srt.mem_cache.allocator.unified_mamba import (
+    UnifiedMambaTokenToKVPoolAllocator,
+)
 from sglang.srt.mem_cache.base_prefix_cache import (
     BasePrefixCache,
     DecLockRefParams,
@@ -45,9 +48,6 @@ from sglang.srt.mem_cache.base_prefix_cache import (
 )
 from sglang.srt.mem_cache.events import KVCacheEventRecorder
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
-from sglang.srt.mem_cache.multi_ended_allocator import (
-    UnifiedMambaTokenToKVPoolAllocator,
-)
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.utils import split_node_hash_value
 from sglang.srt.runtime_context import (

python/sglang/srt/mem_cache/unified_cache/components/swa_component.py

@@ -182,7 +182,7 @@ class SWAComponent(TreeComponent):
 
     def _unified_allocator(self):
         """The unified SWA composite, or None when running on the static pool."""
-        from sglang.srt.mem_cache.multi_ended_allocator import (
+        from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
             UnifiedSWATokenToKVPoolAllocator,
         )
 

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -1193,7 +1193,7 @@ def init_unified_mamba_pools(
     unified_total_bytes: Optional[int] = None,
 ) -> UnifiedPoolBundle:
     """Build the Mamba-hybrid unified-memory-pool stack."""
-    from sglang.srt.mem_cache.multi_ended_allocator import (
+    from sglang.srt.mem_cache.allocator.unified_mamba import (
         UnifiedMambaTokenToKVPoolAllocator,
     )
 
@@ -1676,7 +1676,7 @@ def init_unified_swa_pools(
     sliding_window_size: Optional[int] = None,
 ) -> UnifiedSWAPoolBundle:
     """Build the SWA-hybrid unified-memory-pool stack."""
-    from sglang.srt.mem_cache.multi_ended_allocator import (
+    from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
         UnifiedSWATokenToKVPoolAllocator,
     )
 
@@ -1860,7 +1860,7 @@ def init_unified_mamba_swa_pools(
     fed until the byte configurator lands); the buffer budget is their byte
     sum and the runtime split floats.
     """
-    from sglang.srt.mem_cache.multi_ended_allocator import (
+    from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
         UnifiedMambaSWATokenToKVPoolAllocator,
     )
 

test/registered/unit/disaggregation/test_unified_memory_move_gate.py

@@ -17,7 +17,7 @@ from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
     unified_memory_disagg_move_gate,
 )
-from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
+from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 

test/registered/unit/layers/attention/test_kv_translate_ownership.py

@@ -7,7 +7,7 @@ so a backend that forgets a translate -- or does one twice -- reads the wrong
 rows and nothing crashes. This scan makes both unrepresentable.
 
 Out of scope, deliberately: the allocator-internal implementations
-(`multi_ended_allocator` / `unified_memory_pool`), which ARE the mechanism the
+(`allocator/unified_*` / `unified_memory_pool`), which ARE the mechanism the
 translator calls; the PD transfer plane's `translate_kv_indices_for_transfer`,
 which stages for RDMA outside the forward path; and the STATIC SWA pool's
 legacy full->swa slot map, a different mapping kind with no virtual/physical

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.allocator.unified_hybrid_swa import (
-from sglang.srt.mem_cache.multi_ended_allocator 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,
+            unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
-            mea.UnifiedSWATokenToKVPoolAllocator,
+            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,
+            unified_mamba.UnifiedMambaTokenToKVPoolAllocator,
-            mea.UnifiedSWATokenToKVPoolAllocator,
+            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
+import sglang.srt.mem_cache.allocator.unified_sub_pool as mea
-from sglang.srt.mem_cache.multi_ended_allocator import (
+from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
-    FloatMultiEndedAllocator,
     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