Declare hybrid Mamba indexer host pool

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -1131,7 +1131,7 @@ def build_deepseek_v4_hicache_stack(
 def build_hybrid_mamba_stack(
     *,
     params: CacheInitParams,
-    kv_pool: Any,
+    decls: tuple[HostPoolDecl, ...],
     mamba_pool: Any,
     full_layer_mapping: dict[int, int],
     mamba_layer_mapping: dict[int, int],
@@ -1145,18 +1145,23 @@ def build_hybrid_mamba_stack(
     model_name: Optional[str] = None,
     storage_backend_extra_config: Optional[dict] = None,
     enable_storage_metrics: bool = False,
-) -> tuple[HostPoolGroup, HybridCacheController]:
+) -> tuple[HostPoolGroup, HybridCacheController, HostPoolGroupConfig]:
+    """KV plus every pool the hybrid pool declares (e.g. a sparse indexer),
+    then the Mamba state pool, which keeps its own host path."""
     transfer_layer_id_max = (
         max(full_layer_mapping.keys() | mamba_layer_mapping.keys()) + 1
     )
     mamba_allocator = params.req_to_token_pool.mamba_allocator
-    from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
-
     qsa_draft_pools = params.mtp_draft_device_pools if qsa_pool is not None else ()
-    mtp_draft_device_pools = tuple(
-        pool.full_kv_pool if isinstance(pool, HybridLinearKVPool) else pool
-        for pool in params.mtp_draft_device_pools
+    config = prepare_host_pool_config(
+        decls=decls,
+        full_layer_mapping=full_layer_mapping,
+        transfer_layer_id_max=transfer_layer_id_max,
+        transfer_page_size=params.page_size,
+        packed_draft_device_pools=params.mtp_draft_device_pools,
     )
+    root = config.pools[0]
+    kv_pool = root.decl.device_pool
     kv_host_size, mamba_host_size = None, 0
     if get_memory().hicache_size > 0:
         kv_host_size, mamba_host_size = _split_hicache_size(
@@ -1175,15 +1180,8 @@ def build_hybrid_mamba_stack(
         page_size=params.page_size,
         use_mla=use_mla,
         host_size=kv_host_size,
-        mtp_draft_device_pools=mtp_draft_device_pools,
+        mtp_draft_device_pools=root.packed_draft_device_pools,
     )
-    if mtp_draft_device_pools:
-        full_layer_mapping = with_packed_draft_layer_mapping(
-            full_layer_mapping,
-            transfer_layer_start=transfer_layer_id_max,
-            target_device_layer_num=kv_pool.layer_num,
-            draft_layer_num=len(mtp_draft_device_pools),
-        )
     # MambaPoolHost only supports page_first_direct; the global layout may be
     # page_first_kv_split (e.g. MLA + KDA hybrid on NPU). The Mamba/KDA state
     # pool has no separate K/V buffers, so kv_split does not apply; override
@@ -1200,16 +1198,7 @@ def build_hybrid_mamba_stack(
         allocator_type=_get_allocator_type(),
         layout=mamba_layout,
     )
-    entries = [
-        build_pool_entry(
-            name=PoolName.KV,
-            host_pool=kv_host_pool,
-            device_pool=kv_pool,
-            layer_mapping=full_layer_mapping,
-            transfer_layer_id_max=transfer_layer_id_max + len(mtp_draft_device_pools),
-            is_anchor=True,
-            packed_draft_device_pools=mtp_draft_device_pools,
-        ),
+    entries = _build_pool_entries(config=config, root_host_pool=kv_host_pool) + [
         build_pool_entry(
             name=PoolName.MAMBA,
             host_pool=mamba_host_pool,
@@ -1220,7 +1209,7 @@ def build_hybrid_mamba_stack(
             device_evict_fn=device_mamba_evict_fn,
             device_alloc_fn=mamba_allocator.alloc,
             device_free_fn=mamba_allocator.free,
-        ),
+        )
     ]
     if qsa_pool is not None:
         entries.append(
@@ -1233,8 +1222,9 @@ def build_hybrid_mamba_stack(
                     mtp_draft_device_pools=qsa_draft_pools,
                 ),
                 device_pool=qsa_pool,
-                layer_mapping=full_layer_mapping,
-                transfer_layer_id_max=transfer_layer_id_max + len(qsa_draft_pools),
+                # The KV config carries the packed draft tail layers.
+                layer_mapping=root.layer_mapping,
+                transfer_layer_id_max=root.transfer_layer_id_max,
                 packed_draft_device_pools=qsa_draft_pools,
             )
         )
@@ -1258,7 +1248,7 @@ def build_hybrid_mamba_stack(
         enable_storage_metrics=enable_storage_metrics,
         host_memory_mode=get_memory().hicache_host_memory_mode,
     )
-    return host_pool_group, cache_controller
+    return host_pool_group, cache_controller, config
 
 
 def build_hybrid_mamba_swa_stack(
@@ -1390,13 +1380,20 @@ def build_host_pool_group(
 ) -> HostPoolGroup:
     """Allocate host pools and transfer entries from prepared configs."""
     root = config.pools[0]
-    host_pools = {
-        root.decl.pool_name: build_kv_host_pool(
-            kv_pool=root.decl.device_pool,
-            page_size=config.transfer_page_size,
-            mtp_draft_device_pools=root.packed_draft_device_pools,
-        )
-    }
+    root_host_pool = build_kv_host_pool(
+        kv_pool=root.decl.device_pool,
+        page_size=config.transfer_page_size,
+        mtp_draft_device_pools=root.packed_draft_device_pools,
+    )
+    return HostPoolGroup(
+        _build_pool_entries(config=config, root_host_pool=root_host_pool)
+    )
+
+
+def _build_pool_entries(
+    *, config: HostPoolGroupConfig, root_host_pool: Any
+) -> list[PoolEntry]:
+    host_pools = {config.pools[0].decl.pool_name: root_host_pool}
     entries = []
     for pool_config in config.pools:
         decl = pool_config.decl
@@ -1418,7 +1415,7 @@ def build_host_pool_group(
                 packed_draft_device_pools=pool_config.packed_draft_device_pools,
             )
         )
-    return HostPoolGroup(entries)
+    return entries
 
 
 def _build_mha_mla_host_pool(
@@ -1832,9 +1829,9 @@ def build(
         mamba_layer_mapping = _stage_local_layer_mapping(
             params.req_to_token_pool.mamba_map, kvcache.start_layer
         )
-        host_pool_group, cache_controller = build_hybrid_mamba_stack(
+        host_pool_group, cache_controller, config = build_hybrid_mamba_stack(
             params=params,
-            kv_pool=kvcache.full_kv_pool,
+            decls=kvcache.host_pool_decls(),
             mamba_pool=params.req_to_token_pool.mamba_pool,
             full_layer_mapping=full_layer_mapping,
             mamba_layer_mapping=mamba_layer_mapping,
@@ -1856,13 +1853,21 @@ def build(
                 ComponentType.FULL: host_pool_group.get_pool(PoolName.KV),
                 ComponentType.MAMBA: host_pool_group.get_pool(PoolName.MAMBA),
             },
-            sidecars=(
+            sidecars=[
+                c.decl.sidecar_spec() for c in config.pools if not c.decl.is_primary
+            ]
+            + (
                 [SidecarPoolSpec(PoolName.INDEXER, PoolName.KV)]
                 if qsa_pool is not None
                 else []
             ),
+            pool_declarations=tuple(c.decl for c in config.pools),
             register_req_to_token_counter=True,
-            pools_desc="KV + MAMBA + INDEXER" if qsa_pool is not None else "KV + MAMBA",
+            pools_desc=" + ".join(
+                [c.decl.pool_name.value.upper() for c in config.pools]
+                + ["MAMBA"]
+                + (["INDEXER"] if qsa_pool is not None else [])
+            ),
         )
 
 
@@ -2223,7 +2228,7 @@ def _select_strategy(kvcache: Any, components: set[ComponentType]) -> StackStrat
 
 # Strategies that assemble from host_pool_decls(), so every declared pool has an
 # entry by construction. A miss here is a bug, not an unsupported combination.
-_DECLARATION_VERIFIED_STRATEGIES: tuple[type, ...] = (_DsaStrategy,)
+_DECLARATION_VERIFIED_STRATEGIES: tuple[type, ...] = (_DsaStrategy, _MambaStrategy)
 
 
 def _check_declared_pools_present(

test/registered/unit/mem_cache/test_hybrid_pool_assembler.py

@@ -176,13 +176,14 @@ def test_strategies_pass_stage_local_maps_without_changing_device_maps(self):
                         tp_cache_group=None,
                         pp_cache_group=None,
                     )
+                    controller = SimpleNamespace(transfer_layer_id_max=4)
+                    built = (
+                        (MagicMock(), controller, SimpleNamespace(pools=()))
+                        if strategy_cls is _MambaStrategy
+                        else (MagicMock(), controller)
+                    )
                     with patch.object(
-                        hybrid_pool_assembler,
-                        builder_name,
-                        return_value=(
-                            MagicMock(),
-                            SimpleNamespace(transfer_layer_id_max=4),
-                        ),
+                        hybrid_pool_assembler, builder_name, return_value=built
                     ) as build_stack:
                         result = strategy_cls().build(
                             cache=SimpleNamespace(page_size=1),
@@ -732,6 +733,113 @@ def test_packed_draft_with_a_different_kv_row_is_rejected(self):
         self._build(pool, (_dsa_pool_stub(layer_num=1),))
 
 
+class TestHybridMambaDeclaredIndexer(CustomTestCase):
+    """hybrid Mamba + DSA reuses the target's indexer declaration: the KV/Mamba
+    stack gains an INDEXER entry whose mirror and layer mapping cover only the
+    DSA layers that own index buffers."""
+
+    def test_stack_declares_indexer_and_skips_empty_layers(self):
+        from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
+
+        kv_pool = _dsa_pool_stub(layer_num=3)
+        kv_pool.skip_topk_layers = [False, True, False]  # device layer 1: no buffer
+        mamba_pool = SimpleNamespace(layer_num=2, size=8)
+        # transfer layers 0,2,4 are DSA (device 0,1,2), 1,3 are Mamba
+        full_mapping = {0: 0, 2: 1, 4: 2}
+        mamba_mapping = {1: 0, 3: 1}
+        params = SimpleNamespace(
+            page_size=64,
+            mtp_draft_device_pools=(),
+            token_to_kv_pool_allocator=None,
+            tp_cache_group=None,
+            attn_cp_cache_group=None,
+            attn_tp_cache_group=None,
+            pp_cache_group=None,
+            req_to_token_pool=SimpleNamespace(
+                mamba_allocator=SimpleNamespace(
+                    alloc=lambda n: None, free=lambda x: None
+                )
+            ),
+        )
+        real_indexer_host = pool_host_dsa.DSAIndexerPoolHost
+
+        def dummy_kv_host(**kwargs):
+            return MLATokenToKVPoolHost(
+                kwargs["kv_pool"],
+                host_to_device_ratio=2,
+                host_size=0,
+                page_size=kwargs["page_size"],
+                layout="page_first",
+                pin_memory=False,
+                is_dummy=True,
+                override_kv_cache_dim=None,
+                mtp_draft_device_pools=kwargs["mtp_draft_device_pools"],
+            )
+
+        def dummy_indexer_host(
+            decl, anchor_host, *, allocator_type, packed_draft_device_pools=()
+        ):
+            return real_indexer_host(
+                decl=decl,
+                anchor_host=anchor_host,
+                packed_draft_device_pools=packed_draft_device_pools,
+                allocator_type=allocator_type,
+                pin_memory=False,
+                is_dummy=True,
+            )
+
+        with (
+            patch.object(hybrid_pool_assembler, "build_kv_host_pool", dummy_kv_host),
+            patch.object(pool_host_dsa, "DSAIndexerPoolHost", dummy_indexer_host),
+            patch.object(
+                hybrid_pool_assembler,
+                "MambaPoolHost",
+                return_value=SimpleNamespace(layer_num=2, can_use_write_back_jit=False),
+            ),
+            patch.object(hybrid_pool_assembler, "HybridCacheController", MagicMock()),
+            patch.object(
+                hybrid_pool_assembler, "_get_allocator_type", return_value="default"
+            ),
+            patch.object(
+                hybrid_pool_assembler,
+                "get_memory",
+                return_value=SimpleNamespace(
+                    hicache_size=0,
+                    hicache_ratio=2,
+                    hicache_mem_layout="page_first",
+                    hicache_write_policy="write_through",
+                    hicache_io_backend="kernel",
+                    hicache_host_memory_mode="cache",
+                ),
+            ),
+        ):
+            group, _, config = hybrid_pool_assembler.build_hybrid_mamba_stack(
+                params=params,
+                decls=kv_pool.host_pool_decls(),
+                mamba_pool=mamba_pool,
+                full_layer_mapping=full_mapping,
+                mamba_layer_mapping=mamba_mapping,
+                load_cache_event=None,
+                storage_backend=None,
+                use_mla=True,
+            )
+
+        names = [e.name for e in group.entries]
+        self.assertEqual(names, [PoolName.KV, PoolName.INDEXER, PoolName.MAMBA])
+        indexer = group.entry_map[PoolName.INDEXER]
+        kv = group.entry_map[PoolName.KV]
+        # KV still maps every DSA transfer layer. The indexer drops device layer 1.
+        self.assertEqual([kv.layer_mapper(t) for t in range(5)], [0, None, 1, None, 2])
+        self.assertEqual(
+            [indexer.layer_mapper(t) for t in range(5)], [0, None, None, None, 2]
+        )
+        self.assertEqual(indexer.host_pool.layer_num, 2)
+        self.assertEqual(indexer.host_pool._host_layer_index(2), 1)
+        self.assertEqual(
+            [c.decl.pool_name for c in config.pools], [PoolName.KV, PoolName.INDEXER]
+        )
+
+
 class TestDeclaredPoolPlanning(CustomTestCase):
     """Sidecar indices resolve from one primary source in HostPoolGroup, so the
     planner must reject self-references and sidecar chains up front."""