feat(unified-memory): three sub-pools for mamba + hybrid-SWA models (#35177)

Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

python/sglang/srt/managers/schedule_policy.py

@@ -61,6 +61,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     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
@@ -547,15 +548,17 @@ class PrefillAdder:
 
         self.rem_swa_token_offset = 0
 
-        # Unified-pool joint budget: a new mamba state consumes shared-gap bytes
-        # that `rem_total_tokens` (full KV) otherwise counts as free, so reserve
-        # the gap per new mamba slot or admission over-commits. Gate on the
-        # ALLOCATOR being the unified Mamba composite, NOT on `is_hybrid_ssm_cache`
-        # (False for `ChunkCache`, which would skip the reservation on the
-        # chunk-cache path): the gap coupling is a property of the byte buffer.
+        # A new state slot eats shared-gap bytes that `rem_total_tokens` counts
+        # as free, so reserve per slot or admission over-commits. Gate on the
+        # ALLOCATOR, not `is_hybrid_ssm_cache`: that is False for `ChunkCache`,
+        # which would skip the reservation on the chunk-cache path.
         self._mamba_slot_cost = 0
         if isinstance(
-            self.token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator
+            self.token_to_kv_pool_allocator,
+            (
+                UnifiedMambaTokenToKVPoolAllocator,
+                UnifiedMambaSWATokenToKVPoolAllocator,
+            ),
         ):
             self._mamba_slot_cost = (
                 self.token_to_kv_pool_allocator.mamba_slot_full_token_cost()

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

@@ -23,6 +23,9 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import (
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.mem_cache.multi_ended_allocator import (
+    UnifiedMambaSWATokenToKVPoolAllocator,
+)
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import (
     ceil_align,
@@ -116,9 +119,14 @@ class SchedulerInvariantChecker:
                 // allocator.page_size
                 * allocator.page_size
             )
+        full_available = ps.full_available_size
+        if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator):
+            # Pair the static per-layer total with the conserve view, never the
+            # byte-coordinated one -- see `conserve_full_available_size`.
+            full_available = allocator.conserve_full_available_size()
         leak, msg = self._check_pool_invariant(
             "full",
-            ps.full_available_size,
+            full_available,
             full_evictable_size,
             protected,
             session_held,
@@ -134,9 +142,15 @@ class SchedulerInvariantChecker:
         return leak, msg
 
     def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
+        allocator = self.token_to_kv_pool_allocator
+        swa_available = ps.swa_available_size
+        if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator):
+            # Tri-pool: same floating-boundary phantom as the full pool -- use the
+            # slot-conservation view, not the byte-coordinated min (see _check_full_pool).
+            swa_available = allocator.conserve_swa_available_size()
         return self._check_pool_invariant(
             "swa",
-            ps.swa_available_size,
+            swa_available,
             ps.swa_evictable_size,
             self.tree_cache.swa_protected_size(),
             self.pool_stats_observer.session_held_swa_tokens(),

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

@@ -11,6 +11,10 @@ from typing import (
     Tuple,
 )
 
+from sglang.srt.mem_cache.multi_ended_allocator import (
+    UnifiedMambaSWATokenToKVPoolAllocator,
+)
+
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
     from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
@@ -284,9 +288,18 @@ class SchedulerPoolStatsObserver:
         )
 
     def _get_swa_token_info(self) -> PoolStats:
-        full_available_size = self.token_to_kv_pool_allocator.full_available_size()
+        # `*_num_used` is `static_cap - (available + evictable)`, so the
+        # available term must match the static cap's denomination: the conserve
+        # view, never the byte-coordinated one (see
+        # `conserve_full_available_size`). Measured ~25-90x inflated otherwise.
+        allocator = self.token_to_kv_pool_allocator
+        if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator):
+            full_available_size = allocator.conserve_full_available_size()
+            swa_available_size = allocator.conserve_swa_available_size()
+        else:
+            full_available_size = allocator.full_available_size()
+            swa_available_size = allocator.swa_available_size()
         full_evictable_size = self.tree_cache.full_evictable_size()
-        swa_available_size = self.token_to_kv_pool_allocator.swa_available_size()
         swa_evictable_size = self.tree_cache.swa_evictable_size()
         full_num_used = self.full_tokens_per_layer - (
             full_available_size + full_evictable_size

python/sglang/srt/mem_cache/common.py

@@ -101,7 +101,21 @@ 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
         ]
-        token_to_kv_pool_allocator.free_swa(free_slots)
+        # Local import: multi_ended_allocator imports this module lazily for
+        # eviction; a module-level import here would be a cycle hazard.
+        from sglang.srt.mem_cache.multi_ended_allocator import (
+            UnifiedSWATokenToKVPoolAllocator,
+        )
+
+        if isinstance(token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator):
+            # Contiguous range with host-int bounds: hand the composite its
+            # start position so the free stays host-sync-free (`free_segment`
+            # derives page reps by stride math instead of `torch.unique`).
+            token_to_kv_pool_allocator.free_swa(
+                free_slots, start_pos=req.kv.swa_evicted_seqlen
+            )
+        else:
+            token_to_kv_pool_allocator.free_swa(free_slots)
         req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
 
 

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -409,7 +409,30 @@ class KVCacheConfigurator:
         # (req_to_token_pool is None); supports hybrid Mamba and hybrid SWA (not DSV4).
         if get_memory().enable_unified_memory and req_to_token_pool is None:
             pd_enabled = get_disagg().disaggregation_mode != "null"
-            if self.mambaish_config is not None:
+            is_dsv4 = is_deepseek_v4(self.model_config.hf_config)
+            # Order matters: an Inkling-class model is BOTH mambaish and
+            # hybrid-SWA, and the mamba pair would store every SWA layer's KV at
+            # FULL lifetime -- its branch reads the HF config's
+            # full_attention_layer_ids, which for Inkling is ALL layers.
+            if self.mambaish_config is not None and self.is_hybrid_swa and not is_dsv4:
+                if pd_enabled:
+                    # Same limitation as the 2-pool SWA branch below: the
+                    # tri-pool carries an SWA sub-pool, and there is no
+                    # whole-envelope transfer scheme for it.
+                    raise ValueError(
+                        "--enable-unified-memory with PD disaggregation does "
+                        "not support hybrid-SWA models yet (no whole-envelope "
+                        "transfer scheme for the SWA sub-pool); this model "
+                        "routes to the mamba+SWA tri-pool, which has one. Drop "
+                        "--enable-unified-memory or run without PD."
+                    )
+                bundle = self._init_unified_mamba_swa_pools(
+                    max_num_reqs=sizes.max_running_requests,
+                    full_max_total_num_tokens=sizes.full_max_total_num_tokens,
+                    swa_max_total_num_tokens=sizes.swa_max_total_num_tokens,
+                    unified_total_bytes=sizes.unified_total_bytes,
+                )
+            elif self.mambaish_config is not None:
                 if pd_enabled and not self.use_mla_backend:
                     raise ValueError(
                         "--enable-unified-memory with PD disaggregation "
@@ -423,7 +446,7 @@ class KVCacheConfigurator:
                     max_total_num_tokens=sizes.max_total_num_tokens,
                     unified_total_bytes=sizes.unified_total_bytes,
                 )
-            elif self.is_hybrid_swa and not is_deepseek_v4(self.model_config.hf_config):
+            elif self.is_hybrid_swa and not is_dsv4:
                 if pd_enabled:
                     raise ValueError(
                         "--enable-unified-memory with PD disaggregation does "
@@ -654,6 +677,119 @@ class KVCacheConfigurator:
         )
         return bundle
 
+    def _init_unified_mamba_swa_pools(
+        self,
+        *,
+        max_num_reqs: int,
+        full_max_total_num_tokens: Optional[int],
+        swa_max_total_num_tokens: Optional[int],
+        unified_total_bytes: Optional[int] = None,
+    ) -> UnifiedPoolBundle:
+        """TRI-pool stack for models that are BOTH mambaish and hybrid-SWA
+        (Inkling-class): full KV + SWA KV + mamba/conv state in one buffer,
+        chain [mamba(up) | swa(float) | full(down)]."""
+        from sglang.srt.mem_cache.unified_memory_pool import (
+            init_unified_mamba_swa_pools,
+        )
+
+        config = self.mambaish_config
+        assert config is not None and self.is_hybrid_swa
+        assert self.page_size >= 1, f"page_size must be >= 1, got {self.page_size}"
+        assert (
+            not self.use_mla_backend
+        ), "unified tri-pool does not support an MLA full side yet"
+        # Mirror the non-shared path's extra_max_context_len computation.
+        extra_max_context_len = 4
+        if get_spec().speculative_num_draft_tokens is not None:
+            extra_max_context_len += get_spec().speculative_num_draft_tokens
+
+        head_num = self.model_config.get_num_kv_heads(
+            get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+        )
+        head_dim = self.model_config.head_dim
+        if self.is_hybrid_swa_compress:
+            # Asymmetric full/SWA head geometry (Inkling): SWA dims from the
+            # hf text config, same as the 2-pool SWA wrapper.
+            v_head_dim = self.model_config.hf_text_config.v_head_dim
+            swa_head_num = max(
+                1,
+                self.model_config.hf_text_config.swa_num_key_value_heads
+                // get_parallel().attn_tp_size,
+            )
+            swa_head_dim = self.model_config.hf_text_config.swa_head_dim
+            swa_v_head_dim = self.model_config.hf_text_config.swa_v_head_dim
+        else:
+            v_head_dim = head_dim
+            swa_head_num = head_num
+            swa_head_dim = head_dim
+            swa_v_head_dim = head_dim
+
+        # From the sglang ModelConfig WRAPPER, never the HF config's
+        # full_attention_layer_ids: that property feeds the conv/attention
+        # pairing, not the KV-lifetime split, and returns ALL layers.
+        swa_attention_layer_ids = [
+            i
+            for i in self.model_config.swa_attention_layer_ids
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+        full_attention_layer_ids = [
+            i
+            for i in self.model_config.full_attention_layer_ids
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+        n_local_layers = self.layer_info.end_layer - self.layer_info.start_layer
+        assert (
+            len(full_attention_layer_ids) + len(swa_attention_layer_ids)
+            == n_local_layers
+        ), (
+            "tri-pool KV split must cover every local attention layer exactly "
+            f"once: full={len(full_attention_layer_ids)} + "
+            f"swa={len(swa_attention_layer_ids)} != {n_local_layers} layers in "
+            f"[{self.layer_info.start_layer}, {self.layer_info.end_layer}) — "
+            "the ModelConfig full/swa split is wrong for this architecture"
+        )
+        mamba_layer_ids = [
+            i
+            for i in config.mamba2_cache_params.layers
+            if self.layer_info.start_layer <= i < self.layer_info.end_layer
+        ]
+
+        return init_unified_mamba_swa_pools(
+            device=self.device,
+            kv_cache_dtype=self.kv_cache_dtype,
+            head_num=head_num,
+            head_dim=head_dim,
+            v_head_dim=v_head_dim,
+            swa_head_num=swa_head_num,
+            swa_head_dim=swa_head_dim,
+            swa_v_head_dim=swa_v_head_dim,
+            page_size=self.page_size,
+            start_layer=self.layer_info.start_layer,
+            end_layer=self.layer_info.end_layer,
+            swa_attention_layer_ids=swa_attention_layer_ids,
+            full_attention_layer_ids=full_attention_layer_ids,
+            mamba_layer_ids=mamba_layer_ids,
+            mamba2_cache_params=config.mamba2_cache_params,
+            full_max_total_num_tokens=full_max_total_num_tokens,
+            swa_max_total_num_tokens=swa_max_total_num_tokens,
+            max_mamba_cache_size=get_schedule().max_mamba_cache_size,
+            model_context_len=self.model_config.context_len,
+            extra_max_context_len=extra_max_context_len,
+            max_num_reqs=max_num_reqs,
+            enable_memory_saver=get_exec().features.enable_memory_saver,
+            enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
+            disable_overlap_schedule=get_schedule().disable_overlap_schedule,
+            need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
+            speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
+            forward_stream=self.forward_stream,
+            lazy_compaction=_should_enable_lazy_compaction(),
+            # Draft workers keep the token-count byte sum (spec is asserted
+            # off under unified; belt only).
+            unified_total_bytes=(None if self.is_draft_worker else unified_total_bytes),
+            # bs=1 feasibility floor input (context len is already passed).
+            sliding_window_size=self.model_config.sliding_window_size,
+        )
+
     def _init_unified_swa_pools(
         self,
         *,

python/sglang/srt/mem_cache/unified_memory_pool.py

@@ -11,13 +11,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""UnifiedKVPool — one physical `uint8` byte buffer shared by 2 sub-pools.
+"""UnifiedKVPool -- one physical `uint8` byte buffer shared by N sub-pools.
 
-Two `MultiEndedAllocator`s grow from opposite ends; eager-compacting `free`
-keeps each pool's byte range hole-free. Layout is envelope-major (a slot's data
-for all its layers in one contiguous byte envelope) so a freed slot vacates a
-region the peer can grow into. Everything above the allocator stores virtual
-slot IDs; the allocator owns the per-sub-pool virtual<->physical tables and
+Two END `MultiEndedAllocator`s grow inward from opposite ends; optional
+"float" MIDDLE pools live between their frontiers (chain order
+`[up end, floats..., down end]`). Eager- or lazy-compacting `free` keeps each
+pool's byte range reclaimable. Layout is envelope-major (a slot's data for all
+its layers in one contiguous byte envelope) so a freed slot vacates a region a
+neighbor can grow into. Everything above the allocator stores virtual slot
+IDs; the allocator owns the per-sub-pool virtual<->physical tables and
 compaction only mutates those (no reference rewriting).
 """
 
@@ -26,7 +28,7 @@ from __future__ import annotations
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Dict, List, NamedTuple, Optional, Tuple
+from typing import ClassVar, Dict, List, NamedTuple, Optional, Tuple
 
 import torch
 from torch.profiler import record_function
@@ -70,17 +72,28 @@ def _store_dtype_for(kv_cache_dtype: torch.dtype) -> torch.dtype:
 
 @dataclass(frozen=True, kw_only=True)
 class SubPoolSpec(ABC):
-    """Abstract per-slot layout of one sub-pool in a `UnifiedKVPool`."""
+    """Abstract per-slot layout of one sub-pool in a `UnifiedKVPool`.
+
+    ``grow_direction`` places the sub-pool in the buffer's chain: the two
+    ``"up"``/``"down"`` END pools own the buffer's two ends and grow inward;
+    ``"float"`` middles live between the ends' frontiers (a float's position is
+    carried entirely by its allocator's two watermarks — every view spans the
+    whole buffer, so relocation never rebuilds views).
+    """
+
+    # Grow directions this subclass accepts. Scratch-class specs (e.g. the
+    # spec-decode band) narrow this to ("float",).
+    _allowed_grow_directions: ClassVar[Tuple[str, ...]] = ("up", "down", "float")
 
     name: str
     layer_num: int
-    grow_direction: str  # "up" | "down"
+    grow_direction: str  # "up" | "down" | "float"
 
     def __post_init__(self):
-        assert self.grow_direction in (
-            "up",
-            "down",
-        ), f"grow_direction must be 'up' or 'down'; got {self.grow_direction!r}"
+        assert self.grow_direction in self._allowed_grow_directions, (
+            f"{type(self).__name__}.grow_direction must be one of "
+            f"{self._allowed_grow_directions}; got {self.grow_direction!r}"
+        )
         assert self.layer_num > 0, f"layer_num must be positive; got {self.layer_num}"
 
     @abstractmethod
@@ -278,9 +291,11 @@ def _reserved_floor_bytes(sub_pool_specs: List[SubPoolSpec], page_size: int) ->
 
 
 class UnifiedKVPool:
-    """One physical `uint8` byte buffer shared by 2 sub-pools, each exposing
+    """One physical `uint8` byte buffer shared by N sub-pools, each exposing
     per-layer views over its own byte range (contiguous per layer for KV,
-    strided for the Mamba state). Allocators keep byte ranges disjoint; no usage tracking here.
+    strided for the Mamba state). Two END pools (one grow-up, one grow-down)
+    own the buffer's ends; optional "float" MIDDLE pools live between their
+    frontiers. Allocators keep byte ranges disjoint; no usage tracking here.
     """
 
     def __init__(
@@ -293,21 +308,33 @@ class UnifiedKVPool:
         page_size: int = 1,
     ):
         assert page_size >= 1, f"page_size must be >= 1; got {page_size}"
-        assert len(sub_pool_specs) == 2, (
-            f"UnifiedKVPool currently supports exactly 2 sub-pools; got "
-            f"{len(sub_pool_specs)} (N>2 is not yet implemented)"
-        )
+        assert (
+            len(sub_pool_specs) >= 2
+        ), f"UnifiedKVPool needs >= 2 sub-pools; got {len(sub_pool_specs)}"
         names = [s.name for s in sub_pool_specs]
-        assert len(set(names)) == 2, f"sub-pool names must be unique; got {names}"
-        directions = sorted(s.grow_direction for s in sub_pool_specs)
-        assert directions == ["down", "up"], (
-            f"UnifiedKVPool needs one grow-up and one grow-down sub-pool; "
-            f"got {directions}"
+        assert len(set(names)) == len(
+            names
+        ), f"sub-pool names must be unique; got {names}"
+        # Per-spec direction validity already ran in each spec's __post_init__.
+        up_specs = [s for s in sub_pool_specs if s.grow_direction == "up"]
+        down_specs = [s for s in sub_pool_specs if s.grow_direction == "down"]
+        float_specs = [s for s in sub_pool_specs if s.grow_direction == "float"]
+        assert len(up_specs) == 1 and len(down_specs) == 1, (
+            f"UnifiedKVPool needs exactly one grow-up and one grow-down END "
+            f"sub-pool; got directions "
+            f"{[s.grow_direction for s in sub_pool_specs]}"
         )
 
         self.device = device
         self.total_bytes = total_bytes
-        self.sub_pool_specs = sub_pool_specs
+        # Canonical chain order, low byte end -> high: grow-up end, float
+        # middles (input order preserved), grow-down end. The allocators'
+        # neighbour wiring follows it; all other access is by-name.
+        self.sub_pool_specs: List[SubPoolSpec] = [
+            up_specs[0],
+            *float_specs,
+            down_specs[0],
+        ]
         self._page_size = page_size
         self._specs_by_name: Dict[str, SubPoolSpec] = {
             s.name: s for s in sub_pool_specs
@@ -348,9 +375,9 @@ class UnifiedKVPool:
         # For a page-aware sub-pool the slot-0 write touches layer blocks spread
         # across the WHOLE page-0 envelope (up to page_size * entry_bytes), not
         # just one slot envelope — reserve the max of both.
-        reserved_floor = _reserved_floor_bytes(sub_pool_specs, page_size)
+        reserved_floor = _reserved_floor_bytes(self.sub_pool_specs, page_size)
 
-        for spec in sub_pool_specs:
+        for spec in self.sub_pool_specs:
             entry_bytes = spec.entry_bytes()
             max_slots = total_bytes // entry_bytes
             min_slot_index = (reserved_floor + entry_bytes - 1) // entry_bytes  # ceil
@@ -390,9 +417,9 @@ class UnifiedKVPool:
             "%d sub-pool(s)",
             total_bytes / GB,
             total_bytes,
-            len(sub_pool_specs),
+            len(self.sub_pool_specs),
         )
-        for s in sub_pool_specs:
+        for s in self.sub_pool_specs:
             logger.info(
                 "[unified-memory-pool]   sub-pool %r: kind=%s, layer_num=%d, grow=%s, "
                 "entry_bytes=%d, max_slots=%d, min_slot_index=%d (slots [0,%d) reserved)",
@@ -843,9 +870,11 @@ class UnifiedMambaPool(MambaPool):
     # Inherited MambaPool state ops (copy_from/clear_slots/get_cpu_copy/load_cpu_copy)
     # take PHYSICAL slot ids; callers translate via the slot allocator first.
 
-    def _copy_from_physical(self, src_index: torch.Tensor, dst_index: torch.Tensor):
-        # Physical-slot copy used by the allocator's `_compact_pending`.
-        MambaPool.copy_from(self, src_index, dst_index)
+    def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
+        # Cross-pool physical-move contract, implemented by every pool the
+        # MultiEndedAllocator wraps. Ids are PHYSICAL slots; `MambaPool.copy_from`
+        # takes (src, dst), hence the swap.
+        MambaPool.copy_from(self, src_loc, tgt_loc)
 
     # -- PD state transfer (StateType.MAMBA) --
     # The transfer item is the whole per-slot envelope, addressed as
@@ -1776,3 +1805,222 @@ def init_unified_swa_pools(
         token_to_kv_pool=token_to_kv_pool,
         token_to_kv_pool_allocator=allocator,
     )
+
+
+def init_unified_mamba_swa_pools(
+    *,
+    device: str,
+    kv_cache_dtype: torch.dtype,
+    head_num: int,
+    head_dim: int,
+    v_head_dim: int,
+    swa_head_num: int,
+    swa_head_dim: int,
+    swa_v_head_dim: int,
+    page_size: int,
+    start_layer: int,
+    end_layer: int,
+    swa_attention_layer_ids: List[int],
+    full_attention_layer_ids: List[int],
+    mamba_layer_ids: List[int],
+    mamba2_cache_params,
+    full_max_total_num_tokens: int,
+    swa_max_total_num_tokens: int,
+    max_mamba_cache_size: int,
+    model_context_len: int,
+    extra_max_context_len: int,
+    max_num_reqs: int,
+    enable_memory_saver: bool,
+    enable_mamba_extra_buffer: bool,
+    disable_overlap_schedule: bool,
+    need_sort: bool,
+    speculative_num_draft_tokens: Optional[int] = None,
+    forward_stream: Optional[torch.cuda.Stream] = None,
+    lazy_compaction: bool = False,
+    unified_total_bytes: Optional[int] = None,
+    sliding_window_size: Optional[int] = None,
+) -> UnifiedPoolBundle:
+    """Build the TRI-pool unified-memory-pool stack for models with full KV +
+    SWA KV + mamba/conv state (Inkling-class: `mambaish_config` AND
+    `is_hybrid_swa` simultaneously — Inkling's SConv state is conv-only but
+    rides the mamba machinery, so "mamba" here == the conv state pool).
+
+    Chain: ``[mamba (up END) | swa (FLOAT) | full (down END)]``. The KV side
+    is a `UnifiedSWAKVPool` (per-layer full/swa routing, asymmetric head
+    geometry supported); the state side is a `UnifiedHybridReqToTokenPool`
+    whose `mamba_pool` the model reads directly (sconv:
+    `req_to_token_pool.mamba2_layer_cache(layer).conv[...]` with
+    `translate_mamba_indices` for v->p).
+
+    Sizing inputs are the same token counts the 2-pool factories take (ratio-
+    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 (
+        UnifiedMambaSWATokenToKVPoolAllocator,
+    )
+
+    assert page_size >= 1, f"page_size must be >= 1, got {page_size}"
+    assert (
+        len(full_attention_layer_ids) > 0
+    ), "tri-pool with zero full-attention layers is degenerate"
+    assert (
+        len(swa_attention_layer_ids) > 0
+    ), "tri-pool with zero SWA-attention layers is degenerate"
+    assert len(mamba_layer_ids) > 0, "tri-pool with zero state layers is degenerate"
+
+    store_dtype = _store_dtype_for(kv_cache_dtype)
+    # mamba/conv at the LOWEST bytes, full KV at the HIGHEST, SWA floating
+    # between: ends never relocate, and SWA's window-capped span is cheapest to move.
+    full_spec = MHASubPoolSpec(
+        name="full",
+        layer_num=len(full_attention_layer_ids),
+        head_num=head_num,
+        head_dim=head_dim,
+        v_head_dim=v_head_dim,
+        store_dtype=store_dtype,
+        grow_direction="down",
+    )
+    swa_spec = MHASubPoolSpec(
+        name="swa",
+        layer_num=len(swa_attention_layer_ids),
+        head_num=swa_head_num,
+        head_dim=swa_head_dim,
+        v_head_dim=swa_v_head_dim,
+        store_dtype=store_dtype,
+        grow_direction="float",
+    )
+    cp = mamba2_cache_params
+    mamba_spec = MambaSubPoolSpec(
+        name="mamba",
+        layer_num=len(mamba_layer_ids),
+        conv_state_shapes=tuple(tuple(int(x) for x in s) for s in cp.shape.conv),
+        conv_dtype=cp.dtype.conv,
+        temporal_state_shape=tuple(int(x) for x in cp.shape.temporal),
+        temporal_dtype=cp.dtype.temporal,
+        grow_direction="up",
+    )
+    if unified_total_bytes is not None:
+        # PROFILED byte budget for the token side (captured pre-ratio-floor);
+        # the state pool's bytes ride on top. The token counts stay boot
+        # labels / conserve caps -- the runtime split floats.
+        total_bytes = (
+            unified_total_bytes + max_mamba_cache_size * mamba_spec.entry_bytes()
+        )
+    else:
+        total_bytes = (
+            full_max_total_num_tokens * full_spec.entry_bytes()
+            + swa_max_total_num_tokens * swa_spec.entry_bytes()
+            + max_mamba_cache_size * mamba_spec.entry_bytes()
+        )
+    # bs=1 floor: ONE sliding window of swa KV (+ a page of slack, clamped to
+    # the context) + the state slots one running request locks (1 active + 2
+    # radix checkpoints, a FLOOR not headroom) + the slot-0 sink. The
+    # full-attention side is not charged: `max_req_len` already clamps to the pool.
+    swa_bs1_tokens = (
+        min(model_context_len, sliding_window_size + page_size)
+        if sliding_window_size is not None
+        else model_context_len
+    )
+    _check_bs1_feasibility_floor(
+        total_bytes=total_bytes,
+        floor_terms=[
+            ("swa_window_kv", swa_bs1_tokens * swa_spec.entry_bytes()),
+            ("bs1_state_slots", 3 * mamba_spec.entry_bytes()),
+            (
+                "sink",
+                _reserved_floor_bytes([full_spec, swa_spec, mamba_spec], page_size),
+            ),
+        ],
+        factory="init_unified_mamba_swa_pools",
+    )
+    shared_pool = UnifiedKVPool(
+        total_bytes=total_bytes,
+        sub_pool_specs=[full_spec, swa_spec, mamba_spec],
+        device=device,
+        enable_memory_saver=enable_memory_saver,
+        page_size=page_size,
+    )
+    token_to_kv_pool = UnifiedSWAKVPool(
+        unified_buffer=shared_pool,
+        swa_attention_layer_ids=swa_attention_layer_ids,
+        full_attention_layer_ids=full_attention_layer_ids,
+        page_size=page_size,
+        start_layer=start_layer,
+        end_layer=end_layer,
+        enable_memory_saver=enable_memory_saver,
+    )
+    req_to_token_pool = UnifiedHybridReqToTokenPool(
+        unified_buffer=shared_pool,
+        mamba_sub_pool_name="mamba",
+        size=max_num_reqs,
+        mamba_spec_state_size=max_num_reqs,
+        max_context_len=model_context_len + extra_max_context_len,
+        device=device,
+        enable_memory_saver=enable_memory_saver,
+        cache_params=mamba2_cache_params,
+        mamba_layer_ids=mamba_layer_ids,
+        enable_mamba_extra_buffer=enable_mamba_extra_buffer,
+        speculative_num_draft_tokens=speculative_num_draft_tokens,
+        enable_overlap_schedule=not disable_overlap_schedule,
+        start_layer=start_layer,
+    )
+    allocator = UnifiedMambaSWATokenToKVPoolAllocator(
+        unified_buffer=shared_pool,
+        kvcache=token_to_kv_pool,
+        mamba_kvcache=req_to_token_pool.mamba_pool,
+        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,
+    )
+    # Wrap the composite's mamba end in the slot allocator (PHYSICAL view) the
+    # radix MambaComponent / model-side sconv reads consume.
+    mamba_slot_allocator = UnifiedMambaSlotAllocator(
+        allocator.mamba_allocator,
+        max_size=req_to_token_pool._shared_mamba_size,
+        device=device,
+    )
+    req_to_token_pool.mamba_allocator = mamba_slot_allocator
+
+    logger.info(
+        "[unified-memory-pool] ============================================================"
+    )
+    logger.info(
+        "[unified-memory-pool] UNIFIED MEMORY POOL ENABLED -- path=SWA+Mamba tri-pool"
+    )
+    logger.info(
+        "[unified-memory-pool]   full_layers=%d, swa_layers=%d, state_layers=%d, "
+        "head_num=%d/%d, head_dim=%d/%d, page_size=%d",
+        len(full_attention_layer_ids),
+        len(swa_attention_layer_ids),
+        len(mamba_layer_ids),
+        head_num,
+        swa_head_num,
+        head_dim,
+        swa_head_dim,
+        page_size,
+    )
+    logger.info(
+        "[unified-memory-pool]   total_bytes=%d (=%.2f GB), full_max=%d, swa_max=%d, "
+        "max_mamba_cache_size=%d, max_num_reqs=%d, joint_available=%d",
+        total_bytes,
+        total_bytes / GB,
+        full_max_total_num_tokens,
+        swa_max_total_num_tokens,
+        max_mamba_cache_size,
+        max_num_reqs,
+        allocator.available_size(),
+    )
+    logger.info(
+        "[unified-memory-pool] ============================================================"
+    )
+    return UnifiedPoolBundle(
+        unified_memory_pool=shared_pool,
+        token_to_kv_pool=token_to_kv_pool,
+        token_to_kv_pool_allocator=allocator,
+        req_to_token_pool=req_to_token_pool,
+    )