[None][feat] Support V2 Mamba disaggregated serving

tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py

@@ -337,14 +337,24 @@ def _build_layer_ptrs(
         layer_offsets: Dict[int, int],
         overlapping_layers: List[int],
         slot: int,
+        layer_slot0_addresses: Optional[Dict[int, int]] = None,
+        physical_slot_stride_bytes: Optional[int] = None,
     ) -> np.ndarray:
         """Build per-layer pointers for a given pool (conv or ssm) and slot."""
         ptrs = []
         for glid in overlapping_layers:
-            lid = layer_offsets[glid]
-            ptrs.append(
-                pool.base_address + lid * pool.num_slots * pool.slot_bytes + slot * pool.slot_bytes
-            )
+            if layer_slot0_addresses is not None:
+                ptrs.append(
+                    layer_slot0_addresses[glid]
+                    + slot * (physical_slot_stride_bytes or pool.slot_bytes)
+                )
+            else:
+                lid = layer_offsets[glid]
+                ptrs.append(
+                    pool.base_address
+                    + lid * pool.num_slots * pool.slot_bytes
+                    + slot * pool.slot_bytes
+                )
         return np.array(ptrs, dtype=np.int64)
 
     @staticmethod
@@ -430,11 +440,41 @@ def build_mamba_frags(
             (self_mlg.conv_states, peer_mlg.conv_states, True),
             (self_mlg.ssm_states, peer_mlg.ssm_states, False),
         ]:
+            self_layer_slot0_addresses = (
+                self_mlg.conv_layer_slot0_addresses
+                if is_conv
+                else self_mlg.ssm_layer_slot0_addresses
+            )
+            peer_layer_slot0_addresses = (
+                peer_mlg.conv_layer_slot0_addresses
+                if is_conv
+                else peer_mlg.ssm_layer_slot0_addresses
+            )
+            self_physical_slot_stride_bytes = (
+                self_mlg.conv_physical_slot_stride_bytes
+                if is_conv
+                else self_mlg.ssm_physical_slot_stride_bytes
+            )
+            peer_physical_slot_stride_bytes = (
+                peer_mlg.conv_physical_slot_stride_bytes
+                if is_conv
+                else peer_mlg.ssm_physical_slot_stride_bytes
+            )
             src_ptrs = MambaPolicy._build_layer_ptrs(
-                self_pool, self_mlg.mamba_layer_offsets, overlapping_layers, src_slot
+                self_pool,
+                self_mlg.mamba_layer_offsets,
+                overlapping_layers,
+                src_slot,
+                self_layer_slot0_addresses,
+                self_physical_slot_stride_bytes,
             )
             dst_ptrs = MambaPolicy._build_layer_ptrs(
-                peer_pool, peer_mlg.mamba_layer_offsets, overlapping_layers, dst_slot
+                peer_pool,
+                peer_mlg.mamba_layer_offsets,
+                overlapping_layers,
+                dst_slot,
+                peer_layer_slot0_addresses,
+                peer_physical_slot_stride_bytes,
             )
 
             src_region = SpecRegion(

tensorrt_llm/_torch/disaggregation/native/rank_info.py

@@ -59,6 +59,7 @@ def from_kv_cache_manager(
         m = kv_cache_manager.mapping
         kvm = kv_cache_manager
         enable_attention_dp = m.enable_attention_dp
+        kv_heads_per_rank = next((h for h in kvm.num_kv_heads_per_layer if h > 0), 0)
         return cls(
             instance_name=instance_name,
             instance_rank=m.rank,
@@ -77,7 +78,7 @@ def from_kv_cache_manager(
             self_endpoint="",
             transfer_engine_info=bytes(),
             attention=AttentionInfo(
-                kv_heads_per_rank=kvm.num_kv_heads_per_layer[0],
+                kv_heads_per_rank=kv_heads_per_rank,
                 tokens_per_block=kvm.tokens_per_block,
                 dims_per_head=kvm.head_dim,
                 element_bytes=get_size_in_bytes(1, kvm.dtype),

tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py

@@ -21,7 +21,10 @@
     PoolView,
 )
 from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool
-from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager
+from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
+    MambaHybridCacheManager,
+    V2MambaHybridCacheManager,
+)
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm._utils import get_size_in_bytes, nvtx_range
 from tensorrt_llm.bindings import DataType
@@ -132,6 +135,65 @@ def _build_layer_group_for_mamba(
     )
 
 
+def _physical_slot_stride_bytes(tensor) -> int:
+    return int(tensor.stride(0) * tensor.element_size())
+
+
+def _build_layer_group_for_v2_mamba(
+    manager: V2MambaHybridCacheManager, pool_group_idx: int
+) -> MambaLayerGroup:
+    mamba_layer_offsets = {
+        int(global_layer_id): int(local_layer_id)
+        for global_layer_id, local_layer_id in manager.mamba_layer_offsets.items()
+    }
+
+    first_conv_state = manager.all_conv_states[0]
+    first_ssm_state = manager.all_ssm_states[0]
+    conv_physical_slot_stride_bytes = _physical_slot_stride_bytes(first_conv_state)
+    ssm_physical_slot_stride_bytes = _physical_slot_stride_bytes(first_ssm_state)
+    conv_slot_bytes = int(first_conv_state[0].numel() * first_conv_state.element_size())
+    ssm_slot_bytes = int(first_ssm_state[0].numel() * first_ssm_state.element_size())
+    num_slots = int(first_ssm_state.shape[0])
+
+    conv_layer_slot0_addresses = {
+        int(global_layer_id): int(manager.all_conv_states[offset].data_ptr())
+        for global_layer_id, offset in mamba_layer_offsets.items()
+    }
+    ssm_layer_slot0_addresses = {
+        int(global_layer_id): int(manager.all_ssm_states[offset].data_ptr())
+        for global_layer_id, offset in mamba_layer_offsets.items()
+    }
+
+    d_conv_m1 = manager.conv_state_shape[1]
+    conv_elem_size = first_conv_state.element_size()
+    nheads, head_dim, d_state = manager.ssm_state_shape
+    conv_section_bytes = [dim * d_conv_m1 * conv_elem_size for dim in manager.conv_section_dims]
+
+    ssm_elem_size = first_ssm_state.element_size()
+    ssm_bytes_per_head = head_dim * d_state * ssm_elem_size
+
+    return MambaLayerGroup(
+        pool_group_idx=pool_group_idx,
+        mamba_layer_offsets=mamba_layer_offsets,
+        conv_states=PhysicalPool(
+            base_address=int(first_conv_state.data_ptr()),
+            slot_bytes=conv_slot_bytes,
+            num_slots=num_slots,
+        ),
+        ssm_states=PhysicalPool(
+            base_address=int(first_ssm_state.data_ptr()),
+            slot_bytes=ssm_slot_bytes,
+            num_slots=num_slots,
+        ),
+        conv_section_bytes=conv_section_bytes,
+        ssm_bytes_per_head=ssm_bytes_per_head,
+        conv_layer_slot0_addresses=conv_layer_slot0_addresses,
+        ssm_layer_slot0_addresses=ssm_layer_slot0_addresses,
+        conv_physical_slot_stride_bytes=conv_physical_slot_stride_bytes,
+        ssm_physical_slot_stride_bytes=ssm_physical_slot_stride_bytes,
+    )
+
+
 def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable:
     """Build a KVCachePageTable from a KVCacheManager (V1)."""
     if kv_cache_manager.dtype == DataType.NVFP4:
@@ -339,6 +401,14 @@ def _window_size_for_layer(internal_layer_id: int):
         for variant in pg_desc.slot_desc.variants:
             layer_group_id = int(variant.layer_group_id)
             all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id])
+            if isinstance(manager, V2MambaHybridCacheManager) and any(
+                manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids
+            ):
+                layer_groups_by_id[layer_group_id] = _build_layer_group_for_v2_mamba(
+                    manager, storage_pg_to_list_idx[storage_pg_idx]
+                )
+                continue
+
             all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id)
 
             local_layers = [
@@ -392,7 +462,9 @@ def _window_size_for_layer(internal_layer_id: int):
             raise ValueError(f"Missing V2 layer group descriptor for layer group {layer_group_id}")
         layer_groups.append(layer_group)
 
-    if isinstance(manager, MambaHybridCacheManager):
+    if isinstance(manager, MambaHybridCacheManager) and not isinstance(
+        manager, V2MambaHybridCacheManager
+    ):
         mamba_layer_group_idx = len(pool_groups)
         mamba_layer_group = _build_layer_group_for_mamba(manager, mamba_layer_group_idx)
         layer_groups.append(mamba_layer_group)

tensorrt_llm/_torch/disaggregation/resource/page.py

@@ -204,6 +204,10 @@ class MambaLayerGroup(LayerGroup):
     ssm_states: Optional[PhysicalPool] = None
     conv_section_bytes: Optional[List[int]] = None
     ssm_bytes_per_head: Optional[int] = None
+    conv_layer_slot0_addresses: Optional[Dict[int, int]] = None
+    ssm_layer_slot0_addresses: Optional[Dict[int, int]] = None
+    conv_physical_slot_stride_bytes: Optional[int] = None
+    ssm_physical_slot_stride_bytes: Optional[int] = None
 
     def to_dict(self) -> dict:
         return {
@@ -213,12 +217,26 @@ def to_dict(self) -> dict:
             "ssm_states": self.ssm_states.to_dict(),
             "conv_section_bytes": self.conv_section_bytes,
             "ssm_bytes_per_head": self.ssm_bytes_per_head,
+            "conv_layer_slot0_addresses": {
+                int(k): int(v) for k, v in (self.conv_layer_slot0_addresses or {}).items()
+            }
+            if self.conv_layer_slot0_addresses is not None
+            else None,
+            "ssm_layer_slot0_addresses": {
+                int(k): int(v) for k, v in (self.ssm_layer_slot0_addresses or {}).items()
+            }
+            if self.ssm_layer_slot0_addresses is not None
+            else None,
+            "conv_physical_slot_stride_bytes": self.conv_physical_slot_stride_bytes,
+            "ssm_physical_slot_stride_bytes": self.ssm_physical_slot_stride_bytes,
         }
 
     @classmethod
     def from_dict(cls, data: dict) -> "MambaLayerGroup":
         conv_section_bytes = data.get("conv_section_bytes")
         ssm_bytes_per_head = data.get("ssm_bytes_per_head")
+        conv_layer_slot0_addresses = data.get("conv_layer_slot0_addresses")
+        ssm_layer_slot0_addresses = data.get("ssm_layer_slot0_addresses")
         return cls(
             pool_group_idx=int(data["pool_group_idx"]),
             mamba_layer_offsets={int(k): int(v) for k, v in data["mamba_layer_offsets"].items()},
@@ -228,6 +246,24 @@ def from_dict(cls, data: dict) -> "MambaLayerGroup":
             if conv_section_bytes is not None
             else None,
             ssm_bytes_per_head=int(ssm_bytes_per_head) if ssm_bytes_per_head is not None else None,
+            conv_layer_slot0_addresses={
+                int(k): int(v) for k, v in conv_layer_slot0_addresses.items()
+            }
+            if conv_layer_slot0_addresses is not None
+            else None,
+            ssm_layer_slot0_addresses={int(k): int(v) for k, v in ssm_layer_slot0_addresses.items()}
+            if ssm_layer_slot0_addresses is not None
+            else None,
+            conv_physical_slot_stride_bytes=(
+                int(data["conv_physical_slot_stride_bytes"])
+                if data.get("conv_physical_slot_stride_bytes") is not None
+                else None
+            ),
+            ssm_physical_slot_stride_bytes=(
+                int(data["ssm_physical_slot_stride_bytes"])
+                if data.get("ssm_physical_slot_stride_bytes") is not None
+                else None
+            ),
         )
 
 

tensorrt_llm/_torch/disaggregation/resource/utils.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
 from __future__ import annotations
 
 from typing import Dict, List, Set
@@ -117,9 +131,22 @@ def get_unique_pool_memory_descs(
     pool_counter = 0
     for lg_idx, lg in enumerate(page_table.layer_groups):
         if isinstance(lg, MambaLayerGroup):
-            num_mamba_layers = len(lg.mamba_layer_offsets)
-            for pool in [lg.conv_states, lg.ssm_states]:
-                pool_size = num_mamba_layers * pool.num_slots * pool.slot_bytes
+            is_v2_layout = (
+                lg.conv_layer_slot0_addresses is not None
+                or lg.ssm_layer_slot0_addresses is not None
+            )
+            if is_v2_layout:
+                pools_and_sizes = [
+                    (pool, get_pool_bytes(pool))
+                    for pool in page_table.pool_groups[int(lg.pool_group_idx)].pools
+                ]
+            else:
+                num_mamba_layers = len(lg.mamba_layer_offsets)
+                pools_and_sizes = [
+                    (pool, num_mamba_layers * pool.num_slots * pool.slot_bytes)
+                    for pool in [lg.conv_states, lg.ssm_states]
+                ]
+            for pool, pool_size in pools_and_sizes:
                 pool_key = (pool.base_address, pool_size)
                 if pool_key not in unique_pools:
                     unique_pools[pool_key] = pool_counter

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -30,7 +30,10 @@
 from tensorrt_llm._torch.distributed.communicator import Distributed
 from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
-from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager
+from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
+    MambaHybridCacheManager,
+    V2MambaHybridCacheManager,
+)
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm._utils import nvtx_range
 from tensorrt_llm.bindings import LlmRequestState
@@ -137,7 +140,9 @@ def _init_sync_policy(self):
     def _exchange_rank_info(self):
         endpoints = cast(list, self._dist.allgather(self._transfer_worker.sender_endpoint))
         layer_num = len(self._kv_cache_manager.pp_layers)
-        if isinstance(self._kv_cache_manager, MambaHybridCacheManager):
+        if isinstance(self._kv_cache_manager, MambaHybridCacheManager) and not isinstance(
+            self._kv_cache_manager, V2MambaHybridCacheManager
+        ):
             layer_num += len(self._kv_cache_manager._impl.mamba_layer_offsets)
         layer_num_per_pp = cast(list, getattr(self._dist, "pp_allgather")(layer_num))
         self._transfer_worker.populate_instance_and_rank_info(
@@ -230,7 +235,12 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
             groups.append(block_ids)
 
         mamba_state_index = None
-        if isinstance(self._kv_cache_manager, MambaHybridCacheManager):
+        if (
+            isinstance(self._kv_cache_manager, V2MambaHybridCacheManager)
+            and self._kv_cache_manager.local_num_mamba_layers > 0
+        ):
+            mamba_state_index = self._kv_cache_manager.get_state_indices([req.py_request_id])[0]
+        elif isinstance(self._kv_cache_manager, MambaHybridCacheManager):
             mamba_state_index = self._kv_cache_manager.mamba_cache_index[req.py_request_id]
 
         return KVSlice(

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -17,7 +17,7 @@
 import re
 from contextlib import contextmanager
 from dataclasses import replace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import torch
 
@@ -991,6 +991,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         # is supported for Mamba/SSM-based models
         return {"kv_cache_config": {"enable_block_reuse": False}}
 
+    @classmethod
+    def get_preferred_transceiver_runtime(cls,
+                                          pretrained_config: object
+                                          | None = None) -> Literal["PYTHON"]:
+        """Use the Python transceiver for V2 hybrid-state transfers."""
+        return "PYTHON"
+
     @staticmethod
     def lora_config(model_dir: str):
         """Nemotron-H-specific LoRA configuration.

tensorrt_llm/_torch/models/modeling_qwen3_5.py

@@ -15,7 +15,7 @@
 
 import re
 from types import SimpleNamespace
-from typing import Dict, List
+from typing import Dict, List, Literal
 
 import torch
 from transformers import PretrainedConfig
@@ -674,6 +674,13 @@ def get_model_defaults(cls, llm_args):
         # would silently fall back to the global default (block reuse on).
         return Qwen3NextForCausalLM.get_model_defaults(llm_args)
 
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls, pretrained_config: object | None = None
+    ) -> Literal["PYTHON"]:
+        """Match the hybrid text decoder's V2 disaggregated route."""
+        return "PYTHON"
+
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs):
         kwargs["vision_model_class"] = Qwen3VisionModel
         kwargs["disable_fuse_rope"] = kwargs.get("disable_fuse_rope", False)