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,
+    )

test/registered/models_e2e/test_inkling_unified.py

@@ -0,0 +1,240 @@
+"""Inkling under ``--enable-unified-memory`` -- the first TRI-pool model.
+
+Boots the shrunken ``thinkingmachines/Inkling`` checkpoint (``test`` revision)
+with the unified memory pool: one byte buffer, chain
+``[mamba/conv (up END) | swa (FLOAT) | full (down END)]``. Inkling is the only
+in-tree model that is BOTH mambaish (conv-only SConv state riding the mamba
+machinery) and hybrid-SWA, so booting AT ALL proves the tri routing branch
+(the 2-pool branches would either mis-store SWA KV at full lifetime — the
+pre-tri hazard — or fail loud).
+
+Guards (undertrained checkpoint — code-path correctness, not answer quality):
+  - tri boot + generation through the Triton backend (unified forces triton;
+    Inkling's fa4 default is NOT unified-compatible);
+  - decode/prefill KV consistency via the input-vs-output logprobs match
+    (catches wrong-slot reads through the v2p translate on any of the three
+    pools);
+  - multi-turn prefix reuse: a repeated prefix must reproduce identical
+    logprobs (radix reuse + SWA tombstone recycling + conv COW);
+  - a long-generation turn that slides past the SWA window (exercises
+    out-of-window free_swa -> float holes -> in-place reuse during decode).
+
+An optional env-gated parity class re-runs fixed prompts on the STATIC pools
+and compares logprobs (``INKLING_UNIFIED_PARITY=1``) — the eval-host lane;
+kept out of per-commit CI to bound cost.
+
+    python -m pytest test/registered/models/test_inkling_unified.py -v
+"""
+
+import os
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+
+# Aliased so pytest does not collect the imported `test_`-prefixed helper.
+from sglang.test.kl_test_utils import (
+    test_input_output_logprobs_match_helper as assert_logprobs_match,
+)
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large")
+
+_MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling")
+_MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test")
+
+
+def _unified_args():
+    """Server args for the tri-pool boot. Mirrors test_inkling.py's fixture
+    minus the multimodal/parser surface (KV-path focus), plus the unified
+    flags. The ratios still feed boot sizing until the byte configurator
+    lands; the runtime split floats regardless."""
+    args = [
+        "--trust-remote-code",
+        "--enable-unified-memory",
+        # Unified requires the Triton strided page-major read/write paths.
+        "--attention-backend",
+        "triton",
+        "--page-size",
+        "128",
+        "--mamba-radix-cache-strategy",
+        "extra_buffer",
+        # Inkling defaults to a FULL prefill graph, which unified rejects at
+        # boot: the prefill graph runner bypasses the virtual->physical rebind.
+        "--cuda-graph-backend-prefill",
+        "disabled",
+        "--swa-full-tokens-ratio",
+        "0.1",
+        "--mamba-full-memory-ratio",
+        "0.1",
+        "--mem-fraction-static",
+        "0.5",
+    ]
+    if _MODEL_REVISION:
+        args += ["--revision", _MODEL_REVISION]
+    return args
+
+
+def _static_args():
+    args = [
+        "--trust-remote-code",
+        "--attention-backend",
+        "triton",
+        "--page-size",
+        "128",
+        "--mamba-radix-cache-strategy",
+        "extra_buffer",
+        # Inkling defaults to a FULL prefill graph, which unified rejects at
+        # boot: the prefill graph runner bypasses the virtual->physical rebind.
+        "--cuda-graph-backend-prefill",
+        "disabled",
+        "--swa-full-tokens-ratio",
+        "0.1",
+        "--mamba-full-memory-ratio",
+        "0.1",
+        "--mem-fraction-static",
+        "0.5",
+    ]
+    if _MODEL_REVISION:
+        args += ["--revision", _MODEL_REVISION]
+    return args
+
+
+_PARITY_PROMPTS = [
+    "The capital of France is",
+    "1 + 2 + 3 + 4 + 5 =",
+    "List three prime numbers:",
+]
+
+
+def _greedy_generate(base_url, text, max_new_tokens=32, logprobs=False):
+    payload = {
+        "text": text,
+        "sampling_params": {"temperature": 0.0, "max_new_tokens": max_new_tokens},
+    }
+    if logprobs:
+        payload["return_logprob"] = True
+        payload["logprob_start_len"] = 0
+    resp = requests.post(f"{base_url}/generate", json=payload, timeout=120)
+    assert resp.status_code == 200, resp.text
+    return resp.json()
+
+
+class TestInklingUnifiedTriPool(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = _MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=_unified_args(),
+            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if getattr(cls, "process", None) is not None:
+            kill_process_tree(cls.process.pid)
+
+    def test_generation_basic(self):
+        """Booting IS the tri-routing gate; every prompt must complete."""
+        for prompt in _PARITY_PROMPTS:
+            data = _greedy_generate(self.base_url, prompt, max_new_tokens=16)
+            self.assertIn("text", data, data)
+            self.assertGreater(len(data["text"].strip()), 0, data)
+
+    def test_input_output_logprobs_match(self):
+        """Prefill-vs-decode KV consistency through all three sub-pools'
+        translates (wrong-slot reads surface as logprob mismatches)."""
+        assert_logprobs_match(
+            self.base_url,
+            {self.model: {"kl_div": 1e-2}},
+            self.model,
+            max_samples=4,
+            max_new_tokens=256,
+            trust_remote_code=True,
+        )
+
+    def test_repeated_prefix_reproduces_logprobs(self):
+        """Multi-turn prefix reuse: radix hit + conv COW + swa recycling must
+        not change the numerics of a greedy re-run."""
+        prompt = (
+            "In a quiet village by the sea, a clockmaker kept a ledger of "
+            "every tide. One morning the ledger read:"
+        )
+        first = _greedy_generate(
+            self.base_url, prompt, max_new_tokens=24, logprobs=True
+        )
+        second = _greedy_generate(
+            self.base_url, prompt, max_new_tokens=24, logprobs=True
+        )
+        self.assertEqual(first["text"], second["text"])
+        lp1 = [t[0] for t in first["meta_info"]["output_token_logprobs"]]
+        lp2 = [t[0] for t in second["meta_info"]["output_token_logprobs"]]
+        for a, b in zip(lp1, lp2):
+            self.assertAlmostEqual(a, b, places=3)
+
+    def test_long_decode_slides_past_swa_window(self):
+        """A generation long enough to age tokens out of the SWA window
+        exercises free_swa -> float holes -> in-place reuse mid-decode."""
+        data = _greedy_generate(
+            self.base_url,
+            "Write an unbroken story about a lighthouse: ",
+            max_new_tokens=512,
+        )
+        self.assertGreater(len(data["text"].strip()), 0, data)
+
+
+@unittest.skipUnless(
+    os.environ.get("INKLING_UNIFIED_PARITY") == "1",
+    "eval-host lane: set INKLING_UNIFIED_PARITY=1 (two sequential server boots)",
+)
+class TestInklingUnifiedVsStaticParity(CustomTestCase):
+    """Greedy logprob parity: unified tri-pool vs static pools, same prompts.
+    Two sequential boots -- the strongest wrong-slot tripwire short of GSM8K."""
+
+    @classmethod
+    def _collect(cls, other_args):
+        proc = popen_launch_server(
+            _MODEL_PATH,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+        )
+        try:
+            out = []
+            for p in _PARITY_PROMPTS:
+                data = _greedy_generate(
+                    DEFAULT_URL_FOR_TEST, p, max_new_tokens=32, logprobs=True
+                )
+                out.append(
+                    (
+                        data["text"],
+                        [t[0] for t in data["meta_info"]["output_token_logprobs"]],
+                    )
+                )
+            return out
+        finally:
+            kill_process_tree(proc.pid)
+
+    def test_parity(self):
+        static = self._collect(_static_args())
+        unified = self._collect(_unified_args())
+        for (s_text, s_lp), (u_text, u_lp) in zip(static, unified):
+            self.assertEqual(s_text, u_text)
+            for a, b in zip(s_lp, u_lp):
+                self.assertAlmostEqual(a, b, places=2)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_unified_memory_move_gate.py

@@ -148,9 +148,20 @@ class TestGatedPeerHolesAreNotSchedulable(CustomTestCase):
             self.entry_bytes_per_page = 512
             self.disagg_move_gate = gate
 
+        def _is_frontier_transparent(self):
+            return False
+
     class _Owner:
+        """Stands in for a grow-up END pool: the credit walks the chain from
+        `_growth_side_neighbor()`, so the stub must expose what that walk reads,
+        not the pre-chain `_peer` slot it used to."""
+
         def __init__(self, peer):
-            self._peer = peer
+            self.grow_direction = "up"
+            self.high_peer = peer
+            self.low_peer = None
+
+        _growth_side_neighbor = MultiEndedAllocator._growth_side_neighbor
 
     def _credit(self, gate):
         peer = self._Peer(gate)

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -30,6 +30,7 @@ import unittest
 import torch
 
 from sglang.srt.mem_cache.multi_ended_allocator import (
+    FloatMultiEndedAllocator,
     MultiEndedAllocator,
     UnifiedMambaTokenToKVPoolAllocator,
     UnifiedSWATokenToKVPoolAllocator,
@@ -1825,8 +1826,8 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
         full_kv.attach_allocator = lambda allocator: None
         mamba_kv = _FakeKVCache(pool.max_slots("mamba"))
         mamba_kv.attach_allocator = lambda allocator: None
-        # _copy_from_physical for the mamba sub-pool (kept un-translated).
-        mamba_kv._copy_from_physical = lambda src, dst: None
+        # The physical-move contract for the mamba sub-pool (un-translated);
+        # _FakeKVCache already provides move_kv_cache, so nothing extra.
 
         class _FakeHybridLinearKVPool:
             full_kv_pool = full_kv
@@ -2822,5 +2823,481 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase):
         self.assertTrue(bool((got < 2**31).all().item()))
 
 
+class _ChainStub:
+    """Duck-typed chain member for frontier-walk tests: only the attributes the
+    walk itself touches. Lets the tests position an (opaque|transparent) middle
+    at exact byte coordinates without the full allocator machinery (the real
+    float allocator lands in a later commit)."""
+
+    def __init__(self, *, low_byte: int, high_byte: int, transparent: bool):
+        self._low_byte = low_byte
+        self._high_byte = high_byte
+        self.transparent = transparent
+        self.low_peer = None
+        self.high_peer = None
+        self.lazy_compaction = False
+        self._free_phys_pages = torch.empty(0, dtype=torch.int64)
+        self.entry_bytes_per_page = 1
+        self.sub_pool_name = "stub"
+        self.grow_direction = "float"
+
+    def _is_frontier_transparent(self):
+        return self.transparent
+
+    def _byte_low_frontier(self):
+        return self._low_byte
+
+    def _byte_high_frontier(self):
+        return self._high_byte
+
+
+class TestChainFrontierWalk(unittest.TestCase):
+    """N-pool chain walk: 2-pool byte-identity to the old single-peer formulas,
+    transparent-middle skipping, and growth-side-neighbor credit routing.
+
+    Guarded failure modes: (a) a rewrite of the walk silently changes the
+    2-pool gap math (golden identity); (b) an empty/parked middle walls off
+    free space it does not occupy; (c) drainable-hole credit reads the wrong
+    chain member.
+    """
+
+    def _build_pair(self):
+        full = _make_mha_spec("full", "up", layer_num=2)
+        mamba = _make_mamba_spec("mamba", "down", layer_num=2)
+        total = full.entry_bytes() * 64 + mamba.entry_bytes() * 16
+        pool = UnifiedKVPool(
+            total_bytes=total,
+            sub_pool_specs=[full, mamba],
+            device=_DEV,
+            enable_memory_saver=False,
+        )
+        fa = MultiEndedAllocator(
+            kvcache=_FakeKVCache(pool.max_slots("full")),
+            unified_buffer=pool,
+            sub_pool_name="full",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        ma = MultiEndedAllocator(
+            kvcache=_FakeKVCache(pool.max_slots("mamba")),
+            unified_buffer=pool,
+            sub_pool_name="mamba",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        fa.bind_peer(ma)
+        ma.bind_peer(fa)
+        return pool, fa, ma
+
+    def test_bind_peer_mirrors_growth_side(self):
+        _, fa, ma = self._build_pair()
+        self.assertIs(fa.high_peer, ma)  # grow-up's neighbor sits above
+        self.assertIsNone(fa.low_peer)
+        self.assertIs(ma.low_peer, fa)  # grow-down's neighbor sits below
+        self.assertIsNone(ma.high_peer)
+
+    def test_bind_peer_rejects_float_members(self):
+        _, fa, ma = self._build_pair()
+        stub = _ChainStub(low_byte=0, high_byte=0, transparent=True)
+        with self.assertRaisesRegex(AssertionError, "END-pool-only"):
+            fa.bind_peer(stub)
+        fa.grow_direction = "float"
+        try:
+            with self.assertRaisesRegex(AssertionError, "END-pool-only"):
+                fa.bind_peer(ma)
+        finally:
+            fa.grow_direction = "up"
+
+    def test_two_pool_gap_equals_old_single_peer_formula(self):
+        # Golden identity: with no middles, the chain walk must reproduce the
+        # pre-chain closed form  gap_up = peer_low - my_high ;
+        # gap_down = my_low - peer_high  at every allocation state.
+        _, fa, ma = self._build_pair()
+        for n_full, n_mamba in ((0, 0), (8, 0), (8, 4), (32, 16)):
+            fa.clear()
+            ma.clear()
+            if n_full:
+                self.assertIsNotNone(fa.alloc(n_full))
+            if n_mamba:
+                self.assertIsNotNone(ma.alloc(n_mamba))
+            self.assertEqual(
+                fa._current_gap_bytes(),
+                max(0, ma._byte_low_frontier() - fa._byte_high_frontier()),
+            )
+            # Old down-side closed form: my_low - peer_high (symmetric band).
+            self.assertEqual(
+                ma._current_gap_bytes(),
+                max(0, ma._byte_low_frontier() - fa._byte_high_frontier()),
+            )
+
+    def test_transparent_middle_is_skipped(self):
+        pool, fa, ma = self._build_pair()
+        mid_lo = fa.entry_bytes_per_page * 8
+        mid_hi = fa.entry_bytes_per_page * 12
+        stub = _ChainStub(low_byte=mid_lo, high_byte=mid_hi, transparent=True)
+        fa.bind_high_peer(stub)
+        stub.low_peer = fa
+        stub.high_peer = ma
+        ma.bind_low_peer(stub)
+
+        # Transparent: both ends see straight through to each other.
+        self.assertEqual(
+            fa._current_gap_bytes(),
+            ma._byte_low_frontier() - fa._byte_high_frontier(),
+        )
+        self.assertEqual(
+            ma._current_gap_bytes(),
+            ma._byte_low_frontier() - fa._byte_high_frontier(),
+        )
+
+        # Opaque: each end's gap stops at the middle's near frontier.
+        stub.transparent = False
+        self.assertEqual(fa._current_gap_bytes(), mid_lo - fa._byte_high_frontier())
+        self.assertEqual(ma._current_gap_bytes(), ma._byte_low_frontier() - mid_hi)
+
+    def test_multi_hop_walk_stops_at_first_opaque(self):
+        _, fa, ma = self._build_pair()
+        t1 = _ChainStub(low_byte=100, high_byte=100, transparent=True)
+        t2 = _ChainStub(low_byte=200, high_byte=260, transparent=False)
+        fa.bind_high_peer(t1)
+        t1.low_peer = fa
+        t1.high_peer = t2
+        t2.low_peer = t1
+        t2.high_peer = ma
+        ma.bind_low_peer(t2)
+        self.assertEqual(fa._current_gap_bytes(), 200 - fa._byte_high_frontier())
+        self.assertIs(fa._growth_side_neighbor(), t2)
+        t2.transparent = True
+        self.assertIs(fa._growth_side_neighbor(), ma)
+        self.assertEqual(
+            fa._current_gap_bytes(),
+            ma._byte_low_frontier() - fa._byte_high_frontier(),
+        )
+
+    def test_drainable_credit_reads_walked_neighbor(self):
+        _, fa, ma = self._build_pair()
+        stub = _ChainStub(low_byte=64, high_byte=128, transparent=True)
+        fa.bind_high_peer(stub)
+        stub.low_peer = fa
+        stub.high_peer = ma
+        ma.bind_low_peer(stub)
+
+        # Walked-through to the far end: its holes are credited iff lazy.
+        self.assertEqual(fa._peer_drainable_hole_bytes(), 0)  # ma not lazy
+        ma.lazy_compaction = True
+        ma._free_phys_pages = torch.arange(3, dtype=torch.int64)
+        self.assertEqual(fa._peer_drainable_hole_bytes(), 3 * ma.entry_bytes_per_page)
+        # Opaque non-lazy middle blocks the far end's credit.
+        stub.transparent = False
+        self.assertEqual(fa._peer_drainable_hole_bytes(), 0)
+
+
+class TestFloatMultiEndedAllocator(unittest.TestCase):
+    """Holes-first float middle: midpoint placement, in-place hole recycling,
+    larger-gap boundary extension, boundary absorption + park-on-empty
+    transparency, and the on-demand data movers (`make_room` boundary
+    relocation / leapfrog, `compact_holes` ordered pack) with data-move
+    verification through the fake KV marker.
+
+    Guarded failure modes: a float that copies on steady-state churn, walls
+    off free space when empty, extends toward the tighter gap, moves more
+    than min(L_live, G) pages, loses data across relocation, or corrupts
+    v2p/p2v span/hole bookkeeping.
+    """
+
+    def _build_tri(self, n_state=8, n_float=32, n_full=32):
+        state = _make_mamba_spec("state", "up", layer_num=2)
+        fl = _make_mha_spec("swa", "float", layer_num=2)
+        full = _make_mha_spec("full", "down", layer_num=2)
+        total = (
+            state.entry_bytes() * n_state
+            + fl.entry_bytes() * n_float
+            + full.entry_bytes() * n_full
+        )
+        pool = UnifiedKVPool(
+            total_bytes=total,
+            sub_pool_specs=[full, fl, state],
+            device=_DEV,
+            enable_memory_saver=False,
+        )
+        sa = MultiEndedAllocator(
+            kvcache=_FakeKVCache(pool.max_slots("state")),
+            unified_buffer=pool,
+            sub_pool_name="state",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        fkv = _FakeKVCache(pool.max_slots("swa"))
+        fla = FloatMultiEndedAllocator(
+            kvcache=fkv,
+            unified_buffer=pool,
+            sub_pool_name="swa",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        dkv = _FakeKVCache(pool.max_slots("full"))
+        da = MultiEndedAllocator(
+            kvcache=dkv,
+            unified_buffer=pool,
+            sub_pool_name="full",
+            device=_DEV,
+            is_id_owner=True,
+        )
+        # Chain wiring state <-> float <-> full.
+        sa.bind_high_peer(fla)
+        fla.bind_low_peer(sa)
+        fla.bind_high_peer(da)
+        da.bind_low_peer(fla)
+        return pool, sa, fla, da, fkv
+
+    def _stamp(self, alloc, kv, v):
+        kv.buf[alloc.virtual_to_physical[v]] = v
+
+    def _interior_block(self, fla, blocks):
+        """The allocated block whose physical pages touch neither span
+        boundary (robust to the extend-direction policy)."""
+        for v in blocks:
+            pages = set(int(x) for x in fla.virtual_to_physical[v].tolist())
+            if fla.low_wm_page not in pages and (fla.high_wm_page - 1) not in pages:
+                return v
+        raise AssertionError("no interior block in layout")
+
+    def _check_float_state(self, fla, kv):
+        holes = set(int(x) for x in fla._free_phys_pages.tolist())
+        span = range(fla.low_wm_page, fla.high_wm_page)
+        live = [p for p in span if p not in holes]
+        self.assertEqual(fla._live_pages(), len(live))
+        for h in holes:
+            self.assertTrue(fla.low_wm_page < h < fla.high_wm_page - 1)
+        for p in live:
+            v = int(fla.physical_to_virtual[p].item())
+            self.assertNotEqual(v, -1, f"live page {p} unbound")
+            self.assertEqual(int(fla.virtual_to_physical[v].item()), p)
+            self.assertEqual(int(kv.buf[p].item()), v, f"data lost at {p}")
+
+    def test_midpoint_initial_placement(self):
+        _, _, fla, _, _ = self._build_tri()
+        self.assertTrue(fla._is_frontier_transparent())
+        v = fla.alloc(4)
+        self.assertIsNotNone(v)
+        lo, hi = fla._region_bounds_pages()
+        self.assertEqual(fla.low_wm_page, lo + (hi - lo - 4) // 2)
+        self.assertEqual(fla._span_pages(), 4)
+        # Gap on BOTH sides.
+        gap_low, gap_high = fla._gap_pages()
+        self.assertGreater(gap_low, 0)
+        self.assertGreater(gap_high, 0)
+
+    def test_holes_first_reuse_is_zero_copy(self):
+        _, _, fla, _, kv = self._build_tri()
+        va = fla.alloc(2)
+        vb = fla.alloc(2)
+        vc = fla.alloc(2)
+        for v in (va, vb, vc):
+            self._stamp(fla, kv, v)
+        span_before = (fla.low_wm_page, fla.high_wm_page)
+        fla.free(self._interior_block(fla, (va, vb, vc)))  # interior -> holes
+        self.assertEqual(fla._hole_pages(), 2)
+        self.assertEqual((fla.low_wm_page, fla.high_wm_page), span_before)
+        vd = fla.alloc(2)  # must recycle the holes in place
+        self._stamp(fla, kv, vd)
+        self.assertEqual(fla._hole_pages(), 0)
+        self.assertEqual((fla.low_wm_page, fla.high_wm_page), span_before)
+        self.assertEqual(len(fla._inverse_history), 0)  # zero copies
+        self._check_float_state(fla, kv)
+
+    def test_boundary_free_absorbed_at_the_deferred_point(self):
+        """Boundary holes shrink the span ZERO-COPY -- but the shrink is
+        DEFERRED out of `free`, which must stay host-sync-free (deciding how
+        far to walk needs the hole set on the host). `free` records the holes;
+        `_absorb_span_boundary_holes` (per-step flush / shortfall ladder) reclaims
+        the span. Skipping it is only ever conservative."""
+        _, _, fla, _, kv = self._build_tri()
+        va = fla.alloc(2)
+        vb = fla.alloc(2)  # extends one side; frees at that edge absorb
+        self._stamp(fla, kv, va)
+        self._stamp(fla, kv, vb)
+        span = fla._span_pages()
+        fla.free(vb)
+        # Deferred: span still claims the freed edge, the pages are holes.
+        self.assertEqual(fla._hole_pages(), 2)
+        self.assertEqual(fla._span_pages(), span)
+        self.assertEqual(fla._live_pages(), 2)  # exact regardless
+        absorbed = fla._flush(urgent=False)
+        self.assertEqual(absorbed, 2)
+        self.assertEqual(fla._hole_pages(), 0)
+        self.assertEqual(fla._span_pages(), span - 2)
+        self.assertEqual(len(fla._inverse_history), 0)  # zero copies
+        self._check_float_state(fla, kv)
+
+    def test_free_is_host_sync_free(self):
+        """The per-decode-step property: no D2H anywhere in the float's free
+        (the base's lazy free earns this by deferring absorption to `_flush`;
+        the float now follows the same model)."""
+        from unittest import mock
+
+        _, _, fla, _, kv = self._build_tri()
+        v = fla.alloc(4)
+        self._stamp(fla, kv, v)
+        with mock.patch.object(
+            torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H")
+        ), mock.patch.object(
+            torch.Tensor, "item", side_effect=AssertionError("item = D2H")
+        ), mock.patch.object(
+            torch, "unique", side_effect=AssertionError("unique = host sync")
+        ):
+            fla.free(v[:2], _pages=v[:2])
+
+    def test_deferred_absorption_reaches_the_same_state_as_eager(self):
+        """Derived property: deferring must not change WHERE the span lands,
+        only when. Two floats, identical ops; one absorbs after every free,
+        one only at the end."""
+        _, _, f1, _, kv1 = self._build_tri()
+        _, _, f2, _, kv2 = self._build_tri()
+        for f, kv, absorb_each in ((f1, kv1, True), (f2, kv2, False)):
+            blocks = [f.alloc(2) for _ in range(3)]
+            for v in blocks:
+                self._stamp(f, kv, v)
+            for v in (blocks[2], blocks[0]):
+                f.free(v)
+                if absorb_each:
+                    f._flush(urgent=False)
+        f2._flush(urgent=False)
+        self.assertEqual(f1.low_wm_page, f2.low_wm_page)
+        self.assertEqual(f1.high_wm_page, f2.high_wm_page)
+        self.assertEqual(f1._hole_pages(), f2._hole_pages())
+        self.assertEqual(f1.available_size(), f2.available_size())
+
+    def test_park_on_empty_restores_transparency(self):
+        _, sa, fla, da, _ = self._build_tri()
+        base_gap = da._current_gap_bytes()
+        v = fla.alloc(4)
+        self.assertLess(da._current_gap_bytes(), base_gap)  # float blocks
+        fla.free(v)
+        self.assertTrue(fla._is_frontier_transparent())
+        self.assertEqual(fla._hole_pages(), 0)
+        self.assertEqual(da._current_gap_bytes(), base_gap)  # sees through again
+        self.assertEqual(sa._current_gap_bytes(), base_gap)
+
+    def test_extends_toward_larger_gap(self):
+        _, _, fla, da, kv = self._build_tri()
+        v = fla.alloc(4)
+        self._stamp(fla, kv, v)
+        # Consume most of the high gap with the full end; low gap now larger.
+        self.assertIsNotNone(da.alloc(24))
+        gap_low, gap_high = fla._gap_pages()
+        self.assertGreater(gap_low, gap_high)
+        lo_before = fla.low_wm_page
+        hi_before = fla.high_wm_page
+        v2 = fla.alloc(2)
+        self.assertIsNotNone(v2)
+        self._stamp(fla, kv, v2)
+        self.assertEqual(fla.low_wm_page, lo_before - 2)  # grew low side
+        self.assertEqual(fla.high_wm_page, hi_before)
+        self._check_float_state(fla, kv)
+
+    def test_available_is_max_gap_plus_holes(self):
+        _, _, fla, da, _ = self._build_tri()
+        va = fla.alloc(2)
+        vb = fla.alloc(2)
+        vc = fla.alloc(2)
+        fla.free(self._interior_block(fla, (va, vb, vc)))
+        self.assertEqual(fla._hole_pages(), 2)
+        gap_low, gap_high = fla._gap_pages()
+        self.assertEqual(fla.available_size(), max(gap_low, gap_high) + 2)
+        del da
+
+    def test_make_room_boundary_relocation(self):
+        _, _, fla, da, kv = self._build_tri()
+        v = fla.alloc(6)
+        self._stamp(fla, kv, v)
+        epp = fla.entry_bytes_per_page
+        _, gap_high = fla._gap_pages()
+        ask = (gap_high + 3) * epp  # 3 pages beyond the current high gap
+        opened = fla.make_room(side="high", min_bytes=ask)
+        self.assertGreaterEqual(opened, ask)
+        # Cost min(L_live, G): moved exactly the 3 boundary pages, not 6.
+        moved = sum(int(s.numel()) for s, _, _ in fla._inverse_history)
+        self.assertEqual(moved, 3)
+        self._check_float_state(fla, kv)
+        # The opened space is real: the full end can now take it.
+        self.assertGreaterEqual(da.available_size(), 3)
+
+    def test_make_room_leapfrog_cost_bounded_by_live(self):
+        _, _, fla, _, kv = self._build_tri()
+        v = fla.alloc(2)  # tiny live mass
+        self._stamp(fla, kv, v)
+        epp = fla.entry_bytes_per_page
+        _, gap_high = fla._gap_pages()
+        ask = (gap_high + 10) * epp  # demand >> live
+        opened = fla.make_room(side="high", min_bytes=ask)
+        self.assertGreaterEqual(opened, ask)
+        moved = sum(int(s.numel()) for s, _, _ in fla._inverse_history)
+        self.assertEqual(moved, 2)  # min(L_live, G) == L_live
+        self._check_float_state(fla, kv)
+
+    def test_make_room_impossible_leaves_state_unchanged(self):
+        _, _, fla, _, kv = self._build_tri(n_float=8)
+        v = fla.alloc(6)
+        self._stamp(fla, kv, v)
+        lo, hi = fla._region_bounds_pages()
+        epp = fla.entry_bytes_per_page
+        snapshot = (
+            fla.low_wm_page,
+            fla.high_wm_page,
+            fla._hole_pages(),
+            len(fla._inverse_history),
+        )
+        opened = fla.make_room(side="high", min_bytes=(hi - lo) * epp)
+        self.assertLess(opened, (hi - lo) * epp)
+        self.assertEqual(
+            snapshot,
+            (
+                fla.low_wm_page,
+                fla.high_wm_page,
+                fla._hole_pages(),
+                len(fla._inverse_history),
+            ),
+        )
+        self._check_float_state(fla, kv)
+
+    def test_make_room_uses_far_holes_before_far_gap(self):
+        _, _, fla, _, kv = self._build_tri()
+        va = fla.alloc(2)
+        vb = fla.alloc(2)
+        vc = fla.alloc(2)
+        for v in (va, vb, vc):
+            self._stamp(fla, kv, v)
+        lo_before = fla.low_wm_page
+        fla.free(self._interior_block(fla, (va, vb, vc)))  # 2 interior holes
+        self.assertEqual(fla._hole_pages(), 2)
+        epp = fla.entry_bytes_per_page
+        _, gap_high = fla._gap_pages()
+        opened = fla.make_room(side="high", min_bytes=(gap_high + 2) * epp)
+        self.assertGreaterEqual(opened, (gap_high + 2) * epp)
+        # The two holes absorbed the two moved pages: low side untouched.
+        self.assertEqual(fla.low_wm_page, lo_before)
+        self.assertEqual(fla._hole_pages(), 0)
+        self._check_float_state(fla, kv)
+
+    def test_compact_holes_ordered_pack(self):
+        _, _, fla, _, kv = self._build_tri()
+        vs = [fla.alloc(2) for _ in range(4)]
+        for v in vs:
+            self._stamp(fla, kv, v)
+        fla.free(vs[1])  # interleaved holes
+        span_before = fla._span_pages()
+        moved = fla.compact_holes(retreat_side="high")
+        self.assertEqual(fla._hole_pages(), 0)
+        self.assertEqual(fla._span_pages(), span_before - 2)
+        self.assertGreater(moved, 0)
+        self._check_float_state(fla, kv)
+
+    def test_bind_peer_raises_on_float(self):
+        _, sa, fla, _, _ = self._build_tri()
+        with self.assertRaises(AssertionError):
+            fla.bind_peer(sa)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_capacity_memo.py

@@ -0,0 +1,299 @@
+# 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.
+# ==============================================================================
+"""Epoch-memoized capacity views on the allocator chain (2-pool subset).
+
+The capacity views (`available_size` / `schedulable_available_size` per band,
+plus the composite joint view) are pure functions of a handful of
+CPU-resident fields across the chain; schedulers read them O(queue) times
+between mutations. `_CapacityField` descriptors bump `_capacity_epoch` on
+every rebind, so the memos invalidate by construction.
+
+The failure mode being guarded: a memo serving a STALE value after a mutation
+the epoch machinery missed — either a new mutation site writing a field the
+descriptors don't cover, or an in-place write that bypasses `__set__`. Stale
+capacity is silent over-/under-admission, not a crash. Hence:
+
+  * every mutation kind is followed by memo == fresh-recompute assertions;
+  * a randomized op-sequence property test (seeded) catches interactions no
+    hand-written sequence covers;
+  * a deliberate descriptor-bypassing write must be caught by the IDLE check
+    (`verify_byte_accounting`) — readers cannot detect it, the battery must.
+
+The tri-pool cases (float span fields, three-band joint view) join at the
+tri phase; this file pins the 2-pool chain form they build on.
+
+    python -m pytest test/registered/unit/mem_cache/test_unified_capacity_memo.py -v
+"""
+
+import random
+import unittest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
+
+
+def _build(lazy: bool):
+    # Function-scope import: the fixture is a TestCase subclass, and a
+    # module-scope binding would make pytest collect its tests AGAIN here.
+    from test_multi_ended_allocator import (
+        TestUnifiedSWATokenToKVPoolAllocator as _SwaFixture,
+    )
+
+    inst = _SwaFixture([m for m in dir(_SwaFixture) if m.startswith("test_")][0])
+    pool, allocator, kvcache = inst._build()
+    allocator.full_attn_allocator.lazy_compaction = lazy
+    allocator.swa_attn_allocator.lazy_compaction = lazy
+    allocator.lazy_compaction = lazy
+    return inst, allocator, kvcache
+
+
+class TestCapacityMemoCoherence(unittest.TestCase):
+    def _assert_memos_fresh(self, allocator):
+        """Every memoized capacity view must equal a fresh recompute."""
+        self.assertEqual(
+            allocator.available_size(), allocator._compute_available_size()
+        )
+        for band in (
+            allocator.full_attn_allocator,
+            allocator.swa_attn_allocator,
+        ):
+            self.assertEqual(
+                band.available_size(),
+                band._available_tokens(),
+                msg=f"stale available_size memo on {band.sub_pool_name!r}",
+            )
+            self.assertEqual(
+                band.schedulable_available_size(),
+                band._available_tokens(
+                    extra_gap_bytes=band._peer_drainable_hole_bytes()
+                ),
+                msg=f"stale schedulable memo on {band.sub_pool_name!r}",
+            )
+        self.assertEqual(allocator.verify_byte_accounting(), [])
+
+    def test_memos_track_every_mutation_kind(self):
+        for lazy in (False, True):
+            with self.subTest(lazy_compaction=lazy):
+                inst, allocator, kvcache = _build(lazy)
+                self._assert_memos_fresh(allocator)
+
+                v1 = inst._alloc(allocator, kvcache, 8)  # both-side bind
+                self.assertIsNotNone(v1)
+                self._assert_memos_fresh(allocator)
+
+                allocator.free_swa(v1[2:6])  # swa-side tombstones
+                self._assert_memos_fresh(allocator)
+
+                inst._free(allocator, kvcache, v1)  # both-side free
+                self._assert_memos_fresh(allocator)
+
+                v2 = inst._alloc(allocator, kvcache, 4)
+                self.assertIsNotNone(v2)
+                allocator.free_group_begin()  # grouped free path
+                allocator.free(v2)
+                allocator.free_group_end()
+                self._assert_memos_fresh(allocator)
+
+                if lazy:
+                    allocator.full_attn_allocator._flush(urgent=True)
+                    self._assert_memos_fresh(allocator)
+
+                allocator.clear()
+                self._assert_memos_fresh(allocator)
+
+    def test_random_op_sequence_value_identity(self):
+        """Property: after ANY mutation sequence, memoized views equal fresh
+        recomputes. Seeded (deterministic) — the sequences cover interleavings
+        (alloc / partial swa free / grouped free / flush / clear) that no
+        hand-written case enumerates."""
+        rng = random.Random(0xC0FFEE)
+        for lazy in (False, True):
+            with self.subTest(lazy_compaction=lazy):
+                inst, allocator, kvcache = _build(lazy)
+                live = []
+                for step in range(60):
+                    op = rng.choice(("alloc", "free", "free_swa", "flush", "clear"))
+                    if op == "alloc":
+                        v = inst._alloc(allocator, kvcache, rng.choice((1, 2, 4)))
+                        if v is not None:
+                            live.append(v)
+                    elif op == "free" and live:
+                        inst._free(allocator, kvcache, live.pop())
+                    elif op == "free_swa" and live:
+                        v = live[-1]
+                        if v.numel() > 1:
+                            allocator.free_swa(v[: v.numel() // 2])
+                    elif op == "flush":
+                        allocator.full_attn_allocator._flush(urgent=True)
+                        allocator.swa_attn_allocator._flush(urgent=True)
+                    elif op == "clear":
+                        allocator.clear()
+                        live.clear()
+                    self._assert_memos_fresh(allocator)
+
+    def test_bypassing_write_is_caught_by_the_idle_check(self):
+        inst, allocator, kvcache = _build(lazy=False)
+        v = inst._alloc(allocator, kvcache, 8)
+        self.assertIsNotNone(v)
+        fa = allocator.full_attn_allocator
+        # Prime every memo at the current epoch.
+        allocator.available_size()
+        fa.available_size()
+        fa.schedulable_available_size()
+        # Mutate capacity state BYPASSING the _CapacityField descriptor -- the
+        # epoch does not move, so the memos go stale undetectably for readers...
+        fa.__dict__["watermark_physical"] = fa.watermark_physical + 2
+        # ...but the idle-time coherence check must flag it.
+        violations = allocator.verify_byte_accounting()
+        self.assertTrue(
+            any("stale" in msg for msg in violations),
+            msg=f"bypassing write not caught: {violations}",
+        )
+
+    def test_joint_memo_invalidates_on_swa_only_mutation(self):
+        """The joint view depends on the swa end's frontier through the chain
+        walk; an swa-side-only mutation (tombstoning) must invalidate the
+        composite memo even though the full side never moved."""
+        inst, allocator, kvcache = _build(lazy=True)
+        v = inst._alloc(allocator, kvcache, 8)
+        self.assertIsNotNone(v)
+        before = allocator.available_size()
+        allocator.free_swa(v[:4])  # swa band only
+        after = allocator.available_size()
+        self.assertEqual(after, allocator._compute_available_size())
+        self.assertGreaterEqual(after, before)  # holes only ever add room
+
+    def test_float_only_span_move_invalidates_every_memo(self):
+        """A hole-free float alloc rebinds NO free-list and has no watermark --
+        the span fields are its ONLY capacity state. If they are not
+        `_CapacityField` descriptors, the float's own memo AND both
+        neighbours' (the span flips transparency, walling off their gaps)
+        keep serving pre-move values.
+
+        Driven on a hand-wired end+float+end chain (the composite arrives
+        with the tri phase); the float is exercised alone so no end-pool
+        descriptor write can mask a missing span bump."""
+        from test_multi_ended_allocator import TestFloatMultiEndedAllocator
+
+        inst = TestFloatMultiEndedAllocator(
+            [m for m in dir(TestFloatMultiEndedAllocator) if m.startswith("test_")][0]
+        )
+        _pool, sa, fla, da, _kv = inst._build_tri()
+        self.assertEqual(fla._hole_pages(), 0)  # hole-free extension path
+        self.assertTrue(fla._is_frontier_transparent())
+
+        # Prime every memo while the float is empty/transparent.
+        float_cached = fla.available_size()
+        low_end_cached = sa.available_size()
+        high_end_cached = da.available_size()
+
+        v = fla.alloc(4)  # float-only mutation: span move, no end-pool write
+        self.assertIsNotNone(v)
+        self.assertFalse(fla._is_frontier_transparent())  # span now opaque
+
+        self.assertEqual(fla.available_size(), fla._available_tokens())
+        self.assertEqual(sa.available_size(), sa._available_tokens())
+        self.assertEqual(da.available_size(), da._available_tokens())
+        # The opaque midpoint span must actually reduce what the neighbours
+        # see, i.e. the memos above were not merely re-serving primed values.
+        self.assertLess(sa.available_size(), low_end_cached)
+        self.assertLess(da.available_size(), high_end_cached)
+        self.assertLessEqual(fla.available_size(), float_cached)
+
+    def test_bind_rewiring_bumps_the_epoch(self):
+        """Rewiring changes what the chain walks see; a memo primed before a
+        re-bind must not survive it."""
+        inst, allocator, kvcache = _build(lazy=False)
+        fa = allocator.full_attn_allocator
+        e0 = fa._chain_capacity_epoch()
+        fa.bind_peer(allocator.swa_attn_allocator)  # re-bind (same peer)
+        self.assertGreater(fa._chain_capacity_epoch(), e0)
+
+
+class TestTriCapacityMemoCoherence(unittest.TestCase):
+    """Tri-composite twins of the 2-pool cases: the joint view walks THREE
+    bands (mamba end, swa float, full end), so a mutation on ANY of them must
+    invalidate the composite memo — including the two mutations only the tri
+    has: a mamba-end state draw and a float span move behind the composite."""
+
+    def _build_tri(self, lazy=False):
+        from test_unified_tri_pool import TestUnifiedTriPool
+
+        inst = TestUnifiedTriPool(
+            [m for m in dir(TestUnifiedTriPool) if m.startswith("test_")][0]
+        )
+        pool, allocator, kvcache, mamba_kv = inst._build(lazy_compaction=lazy)
+        return inst, allocator
+
+    def _assert_memos_fresh(self, allocator):
+        self.assertEqual(
+            allocator.available_size(), allocator._compute_available_size()
+        )
+        for band in (
+            allocator.full_attn_allocator,
+            allocator.swa_attn_allocator,
+            allocator.mamba_allocator,
+        ):
+            self.assertEqual(
+                band.available_size(),
+                band._available_tokens(),
+                msg=f"stale available_size memo on {band.sub_pool_name!r}",
+            )
+        self.assertEqual(allocator.verify_byte_accounting(), [])
+
+    def test_memos_track_tri_mutation_kinds(self):
+        for lazy in (False, True):
+            with self.subTest(lazy_compaction=lazy):
+                inst, allocator = self._build_tri(lazy)
+                ma = allocator.mamba_allocator
+                self._assert_memos_fresh(allocator)
+
+                v1 = allocator.alloc(8)  # composite alloc (full + swa bind)
+                self.assertIsNotNone(v1)
+                self._assert_memos_fresh(allocator)
+
+                s1 = ma.alloc(2)  # mamba end alloc
+                self.assertIsNotNone(s1)
+                self._assert_memos_fresh(allocator)
+
+                allocator.free_swa(v1[2:6])  # interior float holes
+                self._assert_memos_fresh(allocator)
+
+                allocator.free(v1)  # both-side free
+                self._assert_memos_fresh(allocator)
+
+                ma.free(s1)  # mamba free
+                self._assert_memos_fresh(allocator)
+
+                allocator.clear()
+                ma.clear()
+                self._assert_memos_fresh(allocator)
+
+    def test_joint_memo_invalidates_on_mamba_only_mutation(self):
+        """The joint view depends on the mamba end's frontier through the
+        chain walk; a mamba-only mutation must invalidate the composite memo
+        even though neither KV side moved."""
+        inst, allocator = self._build_tri()
+        ma = allocator.mamba_allocator
+        before = allocator.available_size()
+        slots = ma.alloc(4)
+        self.assertIsNotNone(slots)
+        after = allocator.available_size()
+        self.assertEqual(after, allocator._compute_available_size())
+        self.assertLessEqual(after, before)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_free_no_host_sync.py

@@ -65,12 +65,55 @@ def _paged_allocator(lazy: bool):
 # 1. tombstone scatters
 # --------------------------------------------------------------------------
 
+_TABLES = {"virtual_to_physical", "physical_to_virtual"}
+
+# Methods that MUST tombstone through index_fill_. Explicit, because "this
+# method writes a tombstone" is a design fact per method, not something a scan
+# can infer -- but `test_every_allocator_free_path_is_listed` below fails if a
+# new allocator arrives with its own free path and is not added here.
 _TOMBSTONE_METHODS = [
     (mea.MultiEndedAllocator, "_free_lazy"),
     (mea.MultiEndedAllocator, "free"),
     (mea.MultiEndedAllocator, "_commit_move_batch"),
+    (mea.FloatMultiEndedAllocator, "free"),
+    (mea.FloatMultiEndedAllocator, "make_room"),
+    (mea.FloatMultiEndedAllocator, "_relocate_to_positions"),
 ]
-_TABLES = {"virtual_to_physical", "physical_to_virtual"}
+
+
+def _allocators_in_module():
+    """Every allocator class DEFINED in multi_ended_allocator (not imported)."""
+    return sorted(
+        (
+            c
+            for c in vars(mea).values()
+            if isinstance(c, type)
+            and c.__module__ == mea.__name__
+            and "Allocator" in c.__name__
+        ),
+        key=lambda c: c.__name__,
+    )
+
+
+def _table_touching_methods():
+    """Every own method of every allocator whose source names a page table.
+
+    DISCOVERY, not a list: a hardcoded list stops guarding the moment a new
+    allocator class arrives with its own free path -- which is what happened
+    when FloatMultiEndedAllocator was added and inherited no coverage.
+    """
+    out = []
+    for cls in _allocators_in_module():
+        for name, fn in vars(cls).items():
+            if not inspect.isfunction(fn):
+                continue
+            try:
+                src = inspect.getsource(fn)
+            except OSError:
+                continue
+            if any(f".{t}[" in src for t in _TABLES):
+                out.append((cls, name))
+    return sorted(out, key=lambda pair: (pair[0].__name__, pair[1]))
 
 
 def _scalar_index_assignments(fn):
@@ -102,13 +145,22 @@ def _scalar_index_assignments(fn):
                 continue
             if isinstance(tgt.slice, ast.Slice):
                 continue
+            # A CONSTANT integer index (`t[0] = 0`, `t[-1] = -1`) is a
+            # single-element sentinel write, not the tensor-index tombstone this
+            # guard exists to find, and the only such writes are in `clear()`.
+            if _is_scalar_literal(tgt.slice):
+                continue
             bad.append(ast.unparse(node))
     return bad
 
 
 class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
     def test_no_scalar_index_assignment(self):
-        for cls, name in _TOMBSTONE_METHODS:
+        discovered = _table_touching_methods()
+        self.assertGreaterEqual(
+            len(discovered), len(_TOMBSTONE_METHODS), "discovery scan went blind"
+        )
+        for cls, name in discovered:
             with self.subTest(method=f"{cls.__name__}.{name}"):
                 bad = _scalar_index_assignments(getattr(cls, name))
                 self.assertEqual(
@@ -132,6 +184,41 @@ class TestTombstonesDoNotCrossTheBus(unittest.TestCase):
 
         self.assertEqual(len(_scalar_index_assignments(_offender)), 1)
 
+    def test_every_allocator_free_path_is_listed(self):
+        """The positive list must name every allocator that owns a free path.
+
+        REGRESSION: the list used to hold three MultiEndedAllocator methods, so
+        adding FloatMultiEndedAllocator with its own `free` silently dropped that
+        free path out of coverage -- and it shipped a scalar tombstone. Fail here
+        instead, loudly, the next time an allocator arrives.
+        """
+        listed = {(cls.__name__, name) for cls, name in _TOMBSTONE_METHODS}
+        for cls in _allocators_in_module():
+            for name, fn in vars(cls).items():
+                if not inspect.isfunction(fn):
+                    continue
+                try:
+                    src = inspect.getsource(fn)
+                except OSError:
+                    continue
+                # WRITES a page table -- either correctly (index_fill_) or in the
+                # banned scalar form the scan below catches. A method that only
+                # READS a table has nothing to tombstone.
+                if not (
+                    any(f"{t}.index_fill_" in src for t in _TABLES)
+                    or _scalar_index_assignments(fn)
+                ):
+                    continue
+                self.assertIn(
+                    (cls.__name__, name),
+                    listed,
+                    msg=(
+                        f"{cls.__name__}.{name} writes a page table but is not in "
+                        f"_TOMBSTONE_METHODS, so the index_fill_ guard does not "
+                        f"cover it. Add it."
+                    ),
+                )
+
     def test_free_paths_actually_use_index_fill(self):
         """Positive form, so deleting the scatter entirely cannot pass."""
         for cls, name in _TOMBSTONE_METHODS:
@@ -302,5 +389,110 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
                 self.assertIn("free_page_reps_group", inspect.getsource(cls))
 
 
+class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
+    """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice
+    with host-int, page-aligned bounds — the same shape `free_segment` was
+    built for. `free_swa(..., start_pos=)` must therefore reach the swa side
+    with caller-derived page ids: no `torch.unique` (data-dependent shape =
+    host sync) and no stale-slot `.item()` on the per-step path.
+
+    Poisoning the ops is the decisive form (a textual guard can be fooled).
+    """
+
+    PS = 4
+
+    def _swa_composite(self, lazy=True):
+        from test_multi_ended_allocator import _FakeKVCache, _make_mha_spec
+
+        from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
+
+        full = _make_mha_spec("full", "up", layer_num=4)
+        swa = _make_mha_spec("swa", "down", layer_num=2)
+        total = 64 * full.entry_bytes() + 64 * swa.entry_bytes()
+        pool = UnifiedKVPool(
+            total_bytes=total,
+            sub_pool_specs=[full, swa],
+            device="cpu",
+            enable_memory_saver=False,
+            page_size=self.PS,
+        )
+
+        class _KV:
+            def __init__(self, p):
+                self.full_kv_pool = _FakeKVCache(p.max_slots("full"))
+                self.swa_kv_pool = _FakeKVCache(p.max_slots("swa"))
+
+            def attach_allocators(self, **kwargs):
+                pass
+
+        return mea.UnifiedSWATokenToKVPoolAllocator(
+            unified_buffer=pool,
+            kvcache=_KV(pool),
+            device="cpu",
+            full_max_total_num_tokens=64,
+            swa_max_total_num_tokens=64,
+            page_size=self.PS,
+            need_sort=False,
+            forward_stream=None,
+            lazy_compaction=lazy,
+        )
+
+    def test_ratchet_shape_free_swa_never_syncs(self):
+        """Aligned bounds (the ratchet guarantees them at ps>1): no unique,
+        no item -- on the lazy production config."""
+        alloc = self._swa_composite(lazy=True)
+        v = alloc.alloc(8 * self.PS)
+        self.assertIsNotNone(v)
+        with mock.patch.object(
+            torch, "unique", side_effect=AssertionError("unique = host sync")
+        ), mock.patch.object(
+            torch.Tensor, "item", side_effect=AssertionError("item = host sync")
+        ):
+            alloc.free_swa(v[: 4 * self.PS], start_pos=0)
+            alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS)
+
+    def test_unaligned_start_pos_still_no_sync(self):
+        """`_page_reps_pieces` covers a misaligned start with a second piece;
+        the sync-free property must not depend on alignment."""
+        alloc = self._swa_composite(lazy=True)
+        v = alloc.alloc(8 * self.PS)
+        with mock.patch.object(
+            torch, "unique", side_effect=AssertionError("unique = host sync")
+        ):
+            alloc.free_swa(v[1 : 5 * self.PS], start_pos=1)
+
+    def test_start_pos_path_matches_the_fallback_end_state(self):
+        """Derived property: the stride-rep path and the dedup fallback must
+        leave IDENTICAL allocator state (v2p tombstones, capacity)."""
+        for lazy in (True, False):
+            with self.subTest(lazy=lazy):
+                a1 = self._swa_composite(lazy=lazy)
+                a2 = self._swa_composite(lazy=lazy)
+                v1 = a1.alloc(6 * self.PS)
+                v2 = a2.alloc(6 * self.PS)
+                self.assertTrue(torch.equal(v1, v2))
+                a1.free_swa(v1[: 4 * self.PS], start_pos=0)
+                a2.free_swa(v2[: 4 * self.PS])  # fallback (radix shape)
+                self.assertTrue(
+                    torch.equal(
+                        a1.swa_attn_allocator.virtual_to_physical,
+                        a2.swa_attn_allocator.virtual_to_physical,
+                    )
+                )
+                self.assertEqual(a1.available_size(), a2.available_size())
+                self.assertEqual(
+                    a1.swa_attn_allocator.schedulable_available_size(),
+                    a2.swa_attn_allocator.schedulable_available_size(),
+                )
+
+    def test_double_ratchet_is_filtered_not_crashed(self):
+        """Freeing an already-tombstoned range again must no-op through the
+        liveness filter (radix eviction and the ratchet can overlap)."""
+        alloc = self._swa_composite(lazy=True)
+        v = alloc.alloc(4 * self.PS)
+        alloc.free_swa(v, start_pos=0)
+        alloc.free_swa(v, start_pos=0)  # all tombstoned -> filtered to empty
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_npool_sweep.py

@@ -0,0 +1,271 @@
+# 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.
+# ==============================================================================
+"""N-sub-pool construction sweep for ``UnifiedKVPool``.
+
+The pool accepts N sub-pool specs: exactly one grow-up END, exactly one
+grow-down END, and >= 0 "float" MIDDLE pools between their frontiers. These
+tests pin the constructor contract the N-pool chain machinery builds on:
+
+  - canonical chain order ``[up end, floats (input order), down end]`` —
+    input list order is irrelevant (2-pool configs stay byte-identical);
+  - by-name geometry (``max_slots = total_bytes // entry_bytes``,
+    ``min_slot_index`` past the shared reserved floor) independent of N;
+  - the reserved slot-0 sink covers EVERY sub-pool's page-0 dummy-write
+    envelope, floats included (mamba stays page_size=1);
+  - validation: unique names, exactly one up + one down, >= 2 specs, and
+    per-spec ``_allowed_grow_directions`` narrowing;
+  - float sub-pool views build and round-trip like end-pool views (all views
+    span the whole buffer at anchor 0; keeping the bands disjoint is the
+    allocators' job).
+
+Pure CPU geometry — no allocator, no GPU.
+
+    python -m pytest test/registered/unit/mem_cache/test_unified_npool_sweep.py -v
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.unified_memory_pool import (
+    MambaSubPoolSpec,
+    MHASubPoolSpec,
+    MLASubPoolSpec,
+    UnifiedKVPool,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+# Plain unittest.TestCase, importing only ci_register -- the deliberate
+# hermetic convention of the pool-geometry tests in this directory (see
+# test_multi_ended_allocator.py): no heavy sglang.test.test_utils import
+# chain, so the suite runs in a lean torch-only environment.
+register_cpu_ci(est_time=30, suite="base-a-test-cpu")
+
+_DEV = "cpu"
+
+
+def _mha(
+    name: str,
+    grow_direction: str,
+    *,
+    layer_num: int = 2,
+    head_num: int = 2,
+    head_dim: int = 8,
+) -> MHASubPoolSpec:
+    return MHASubPoolSpec(
+        name=name,
+        layer_num=layer_num,
+        grow_direction=grow_direction,
+        head_num=head_num,
+        head_dim=head_dim,
+        store_dtype=torch.bfloat16,
+    )
+
+
+def _mla(name: str, grow_direction: str, *, layer_num: int = 2) -> MLASubPoolSpec:
+    return MLASubPoolSpec(
+        name=name,
+        layer_num=layer_num,
+        grow_direction=grow_direction,
+        kv_lora_rank=16,
+        qk_rope_head_dim=8,
+        store_dtype=torch.bfloat16,
+    )
+
+
+def _mamba(name: str, grow_direction: str, *, layer_num: int = 2) -> MambaSubPoolSpec:
+    return MambaSubPoolSpec(
+        name=name,
+        layer_num=layer_num,
+        grow_direction=grow_direction,
+        conv_state_shapes=((4, 6),),
+        conv_dtype=torch.bfloat16,
+        temporal_state_shape=(2, 4, 4),
+        temporal_dtype=torch.float32,
+    )
+
+
+def _make_pool(specs, *, total_bytes: int = 1 << 20, page_size: int = 1):
+    return UnifiedKVPool(
+        total_bytes=total_bytes,
+        sub_pool_specs=specs,
+        device=_DEV,
+        enable_memory_saver=False,
+        page_size=page_size,
+    )
+
+
+def _chain_names(pool: UnifiedKVPool):
+    return [s.name for s in pool.sub_pool_specs]
+
+
+class TestNPoolCanonicalOrder(unittest.TestCase):
+    def test_two_pool_input_order_irrelevant(self):
+        for specs in (
+            [_mha("full", "down"), _mamba("mamba", "up")],
+            [_mamba("mamba", "up"), _mha("full", "down")],
+        ):
+            pool = _make_pool(specs)
+            self.assertEqual(_chain_names(pool), ["mamba", "full"])
+
+    def test_three_pool_float_in_the_middle(self):
+        for specs in (
+            [_mha("full", "down"), _mha("swa", "float"), _mamba("conv", "up")],
+            [_mha("swa", "float"), _mamba("conv", "up"), _mha("full", "down")],
+            [_mamba("conv", "up"), _mha("full", "down"), _mha("swa", "float")],
+        ):
+            pool = _make_pool(specs)
+            self.assertEqual(_chain_names(pool), ["conv", "swa", "full"])
+
+    def test_four_pool_float_input_order_preserved(self):
+        pool = _make_pool(
+            [
+                _mha("full", "down"),
+                _mha("f1", "float"),
+                _mamba("state", "up"),
+                _mha("f0", "float", layer_num=1),
+            ]
+        )
+        # Ends canonical; floats keep INPUT order between them.
+        self.assertEqual(_chain_names(pool), ["state", "f1", "f0", "full"])
+
+    def test_by_name_geometry_independent_of_n(self):
+        two = _make_pool([_mha("full", "down"), _mamba("mamba", "up")])
+        three = _make_pool(
+            [_mha("full", "down"), _mha("swa", "float"), _mamba("mamba", "up")]
+        )
+        for name in ("full", "mamba"):
+            self.assertEqual(
+                two.max_slots(name),
+                two.total_bytes // two.spec(name).entry_bytes(),
+            )
+            self.assertEqual(two.max_slots(name), three.max_slots(name))
+        for pool in (two, three):
+            for s in pool.sub_pool_specs:
+                self.assertEqual(pool.anchor_bytes(s.name), 0)
+
+
+class TestNPoolValidation(unittest.TestCase):
+    def test_duplicate_names_rejected(self):
+        with self.assertRaisesRegex(AssertionError, "unique"):
+            _make_pool([_mha("x", "down"), _mamba("x", "up")])
+
+    def test_fewer_than_two_specs_rejected(self):
+        with self.assertRaisesRegex(AssertionError, ">= 2 sub-pools"):
+            _make_pool([_mha("full", "down")])
+
+    def test_two_ups_rejected(self):
+        with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
+            _make_pool([_mha("a", "up"), _mamba("b", "up")])
+
+    def test_missing_down_end_rejected(self):
+        with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
+            _make_pool([_mha("a", "up"), _mha("b", "float")])
+
+    def test_missing_up_end_rejected(self):
+        with self.assertRaisesRegex(AssertionError, "exactly one grow-up"):
+            _make_pool([_mha("a", "down"), _mha("b", "float"), _mha("c", "float")])
+
+    def test_bogus_direction_rejected_at_spec_level(self):
+        with self.assertRaisesRegex(AssertionError, "grow_direction"):
+            _mha("a", "sideways")
+
+    def test_float_accepted_on_all_cache_spec_kinds(self):
+        # Every cache-class spec kind may float (the chain decides placement).
+        pool = _make_pool(
+            [
+                _mamba("state", "up"),
+                _mha("f_mha", "float"),
+                _mla("f_mla", "float", layer_num=1),
+                _mamba("f_mamba", "float", layer_num=1),
+                _mha("full", "down"),
+            ]
+        )
+        self.assertEqual(
+            _chain_names(pool), ["state", "f_mha", "f_mla", "f_mamba", "full"]
+        )
+
+
+class TestReservedFloorWithFloats(unittest.TestCase):
+    def test_float_page_envelope_extends_the_sink(self):
+        # The float MHA has the largest page-0 envelope; every pool's
+        # min_slot_index must clear it (mamba is page_size=1 and excluded from
+        # the page-aware term, but still must clear the byte floor).
+        page_size = 4
+        big_float = _mha("swa", "float", layer_num=8, head_num=4, head_dim=32)
+        specs = [_mamba("state", "up"), big_float, _mha("full", "down")]
+        pool = _make_pool(specs, total_bytes=1 << 22, page_size=page_size)
+        floor = max(
+            max(s.entry_bytes() for s in specs),
+            page_size * big_float.entry_bytes(),
+            page_size * specs[2].entry_bytes(),
+        )
+        for s in specs:
+            e = s.entry_bytes()
+            self.assertEqual(pool.min_slot_index(s.name), (floor + e - 1) // e)
+
+    def test_too_small_buffer_fails_loud(self):
+        # 2048 B with page_size=16 and 128 B/entry MHA specs: the page-0 sink
+        # (16*128 = 2048 B) consumes the whole buffer -> min_slot_index ==
+        # max_slots for the MHA pools -> no allocatable slot -> loud error.
+        with self.assertRaisesRegex(RuntimeError, "no room"):
+            _make_pool(
+                [_mamba("state", "up"), _mha("swa", "float"), _mha("full", "down")],
+                total_bytes=2048,
+                page_size=16,
+            )
+
+
+class TestFloatViews(unittest.TestCase):
+    def test_float_mha_views_shape_and_roundtrip(self):
+        page_size = 2
+        spec = _mha("swa", "float", layer_num=3, head_num=2, head_dim=8)
+        pool = _make_pool(
+            [_mamba("state", "up"), spec, _mha("full", "down")],
+            total_bytes=1 << 20,
+            page_size=page_size,
+        )
+        k_views, v_views = pool.mha_views_for("swa")
+        self.assertEqual(len(k_views), spec.layer_num)
+        self.assertEqual(len(v_views), spec.layer_num)
+        num_pages = pool.max_slots("swa") // page_size
+        blocks = 2 * spec.layer_num  # K at block 2l, V at 2l+1
+        n_rows = num_pages * blocks * page_size
+        for k in (*k_views, *v_views):
+            # Stock 3-D per-layer MHA signature; the row index is the
+            # kernel-facing id, each view's storage_offset folding in its block
+            # origin (see `build_mha_views`).
+            self.assertEqual(tuple(k.shape), (n_rows, spec.head_num, spec.head_dim))
+        # Round-trip: a float view is a real strided window into _raw.
+        slot = pool.min_slot_index("swa")
+        row = (slot // page_size) * (page_size * blocks) + slot % page_size
+        pattern = (
+            torch.arange(spec.head_num * spec.head_dim, dtype=torch.float32)
+            .reshape(spec.head_num, spec.head_dim)
+            .to(torch.bfloat16)
+        )
+        k_views[1][row] = pattern
+        torch.testing.assert_close(k_views[1][row], pattern)
+
+    def test_float_mamba_views_zero_visible(self):
+        pool = _make_pool(
+            [_mamba("state", "up"), _mamba("fstate", "float"), _mha("full", "down")]
+        )
+        conv_views, temporal = pool.mamba_views_for("fstate")
+        self.assertTrue(all(v.eq(0).all() for v in conv_views))
+        self.assertTrue(temporal.eq(0).all())
+
+
+if __name__ == "__main__":
+    unittest.main()