NVIDIA · VALLIS-NERIA · Jul 21, 2026 · Jul 21, 2026 · Jul 21, 2026 · Jul 21, 2026
@@ -28,7 +28,7 @@ unset or when the safety sanitizer rejects the runtime value.
 
 ### `TorchLlmArgs`
 
-269 captured fields.
+270 captured fields.
 
 | Captured key | Annotation | Kind | Converter | Allowed values |
 |--------------|------------|------|-----------|----------------|
@@ -43,6 +43,7 @@ unset or when the safety sanitizer rejects the runtime value.
 | `attention_dp_config.kv_cache_routing_load_balance_weight` | `<class 'float'>` | `value` |  |  |
 | `attention_dp_config.kv_cache_routing_match_rate_threshold` | `<class 'float'>` | `value` |  |  |
 | `attention_dp_config.kv_cache_routing_max_sessions` | `<class 'int'>` | `value` |  |  |
+| `attention_dp_config.kv_cache_routing_new_conv_placement` | `Literal['round_robin', 'least_queued']` | `categorical` |  | `round_robin`, `least_queued` |
 | `attention_dp_config.timeout_iters` | `<class 'int'>` | `value` |  |  |
 | `attn_backend` | `<class 'str'>` | `categorical` | allowlist | `VANILLA`, `TRTLLM`, `FLASHINFER`, `FLASHINFER_STAR_ATTENTION` |
 | `backend` | `Literal['pytorch']` | `categorical` |  | `pytorch` |
@@ -123,7 +124,7 @@ unset or when the safety sanitizer rejects the runtime value.
 | `kv_cache_config.mamba_ssm_cache_dtype` | `Literal['auto', 'float16', 'bfloat16', 'float32']` | `categorical` |  | `auto`, `float16`, `bfloat16`, `float32` |
 | `kv_cache_config.mamba_ssm_philox_rounds` | `<class 'int'>` | `value` |  |  |
 | `kv_cache_config.mamba_ssm_stochastic_rounding` | `<class 'bool'>` | `value` |  |  |
-| `kv_cache_config.mamba_state_cache_interval` | `<class 'int'>` | `value` |  |  |
+| `kv_cache_config.mamba_state_config.periodic_snapshot_interval` | `<class 'int'>` | `value` |  |  |
 | `kv_cache_config.max_attention_window` | `Optional[List[int]]` | `value` |  |  |
 | `kv_cache_config.max_gpu_total_bytes` | `<class 'int'>` | `value` |  |  |
 | `kv_cache_config.max_tokens` | `Optional[int]` | `value` |  |  |

@@ -74,6 +74,49 @@ scheduler_config:
   enable_prefix_aware_scheduling: false
 ```
 
+### Mamba Snapshot Boundaries
+
+Hybrid Mamba models must retain the recurrent Mamba state together with the
+attention KV prefix. Snapshot policy is grouped under
+`kv_cache_config.mamba_state_config`. `periodic_snapshot_interval` controls
+periodic boundaries. They are disabled by default; set the interval to a
+positive value to enable them. The deprecated
+`kv_cache_config.mamba_state_cache_interval` alias remains accepted for
+compatibility and is copied to the nested field during validation. New code and
+configuration files should use the nested field. The prototype
+`additional_snapshot_offsets_from_start` and
+`additional_snapshot_offsets_from_end` options add fixed boundaries. Start
+offsets count tokens from the beginning of the prompt. End offsets count
+backward from the prompt end, and an end offset of `0` selects the final
+prompt boundary. The `per_conversation` block reuse policy disables periodic
+Mamba snapshots, so configure one or more explicit stable boundaries (usually
+an end offset of `0`) when using it with a hybrid Mamba model. For example:
+
+```yaml
+kv_cache_config:
+  enable_block_reuse: true
+  use_kv_cache_manager_v2: true
+  avg_seq_len: 2048
+  mamba_state_config:
+    periodic_snapshot_interval: 0
+    additional_snapshot_offsets_from_start: [128]
+    additional_snapshot_offsets_from_end: [0, 32]
+```
+
+This retains snapshots after the first 128 tokens, at the end of the prompt,
+and before the final 32 prompt tokens. Positions outside a particular prompt
+are ignored. Set `avg_seq_len` to the workload's average total sequence length
+so V2 can size the attention KV and Mamba state pools in the right proportion.
+If neither `avg_seq_len` nor an explicit `pool_ratio` is configured, hybrid
+Mamba models warn and fall back to half of `max_seq_len`, which can produce a
+suboptimal pool split. Exact explicit boundaries currently require
+`MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid
+Mamba models select V2 by default when
+`use_kv_cache_manager_v2: auto`; set it to `false` to select the V1 C++
+compatibility manager. In disaggregated serving, V2 Mamba requires the Python
+NIXL transceiver (`transceiver_runtime: PYTHON`); V1 routes support periodic
+snapshots only.
+
 ### KV Cache Salting for Secure Reuse
 
 KV cache salting provides a security mechanism to control which requests can reuse cached KV states. When a `cache_salt` parameter is provided with a request, the KV cache system will only allow reuse of cached blocks given the same cache salt value. This prevents potential security issues such as prompt theft attacks, where malicious users might try to infer information from cached states of other users' requests.

@@ -1004,6 +1004,11 @@ def update_spec_dec_param(
         self.use_spec_decoding = self.is_spec_decoding_enabled
         self.is_spec_dec_tree = is_spec_dec_tree
         self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree
+        # A hybrid model's first executed attention layer can have a nonzero
+        # cache-local index because recurrent layers precede it.  Do not rely
+        # on the C++ ``layer_idx == 0`` fallback to rebuild the target mask:
+        # the dynamic draft loop clears that mask before the next target step.
+        self.force_prepare_spec_dec_tree_mask = is_spec_dec_dynamic_tree
         # Forward static tree length to FMHA kernel selection.
         self.max_total_draft_tokens = max_total_draft_tokens
 

@@ -587,6 +587,11 @@ def head_match(self, peer_ri: RankInfo) -> tuple[bool, bool]:
         return head_match, is_dup_head
 
     def duplicate_head_factors(self, peer_ri: RankInfo) -> tuple[int, int]:
+        # Head duplication only applies to attention KV.  In particular, do
+        # not divide 0 / 0 for an attention-free hybrid PP stage; MambaPolicy
+        # computes its TP mapping separately.
+        if self._ri.attention.kv_heads_per_rank == 0 or peer_ri.attention.kv_heads_per_rank == 0:
+            return 1, 1
         factor_self, factor_peer = self._head_factors(peer_ri)
         dup_head = max(1, factor_self // factor_peer)
         peer_dup_head = max(1, factor_peer // factor_self)

@@ -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 typing import Dict, List, Optional, Tuple
 
 import numpy as np
@@ -338,13 +352,17 @@ def _build_layer_ptrs(
         overlapping_layers: List[int],
         slot: int,
     ) -> 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
-            )
+        """Build per-layer pointers from a pool's affine layer/slot layout."""
+        slot_stride_bytes = pool.slot_stride_bytes
+        layer_stride_bytes = pool.layer_stride_bytes
+        assert slot_stride_bytes is not None
+        assert layer_stride_bytes is not None
+        ptrs = [
+            pool.base_address
+            + layer_offsets[global_layer_id] * layer_stride_bytes
+            + slot * slot_stride_bytes
+            for global_layer_id in overlapping_layers
+        ]
         return np.array(ptrs, dtype=np.int64)
 
     @staticmethod
@@ -431,10 +449,16 @@ def build_mamba_frags(
             (self_mlg.ssm_states, peer_mlg.ssm_states, False),
         ]:
             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,
             )
             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,
             )
 
             src_region = SpecRegion(

@@ -74,6 +74,10 @@ def from_kv_cache_manager(
         m = kv_cache_manager.mapping
         kvm = kv_cache_manager
         enable_attention_dp = m.enable_attention_dp
+        # Keep AttentionInfo on attention-free PP stages so it can still carry
+        # the attention-DP topology used by Mamba transfers.  A zero head count
+        # means that this rank has no local attention cache; AttentionPolicy
+        # must not perform head-ratio arithmetic for such ranks.
         kv_heads_per_rank = next((h for h in kvm.num_kv_heads_per_layer if h > 0), 0)
         # Eight is the smallest element count guaranteed to occupy whole bytes
         # for every supported sub-byte cache dtype (including NVFP4).

@@ -14,9 +14,10 @@
 # limitations under the License.
 
 from collections import defaultdict
-from typing import Dict, List
+from typing import Dict, List, Sequence
 
 import numpy as np
+import torch
 
 from tensorrt_llm._torch.disaggregation.base.region import (
     DataLayout,
@@ -41,7 +42,10 @@
     get_physical_pool,
 )
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
-from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager
+from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import (
+    MambaHybridCacheManager,
+    MambaHybridCacheManagerV2,
+)
 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
@@ -106,10 +110,12 @@ def extract(
 
         base_ptr = pool.base_address
         block_size = pool.slot_bytes
+        block_stride = pool.slot_stride_bytes
+        assert block_stride is not None
 
         # KV cache: filter out invalid block_ids (BAD_PAGE_INDEX = -1)
         valid = region_ids >= 0
-        ptrs = base_ptr + block_size * region_ids[valid]
+        ptrs = base_ptr + block_stride * region_ids[valid]
         memory = MemRegionGroup(ptrs=ptrs, bytes_per_region=block_size)
         return SpecRegion(memory=memory)
 
@@ -134,16 +140,19 @@ def _build_layer_group_for_mamba(
         base_address=conv_state.data_ptr(),
         slot_bytes=conv_state.stride(1) * conv_state.element_size(),
         num_slots=conv_state.shape[1],
+        layer_stride_bytes=conv_state.stride(0) * conv_state.element_size(),
     )
 
     ssm_pool = PhysicalPool(
         base_address=ssm_state.data_ptr(),
         slot_bytes=ssm_state.stride(1) * ssm_state.element_size(),
         num_slots=ssm_state.shape[1],
+        layer_stride_bytes=ssm_state.stride(0) * ssm_state.element_size(),
     )
 
     # Per-section bytes for conv_state and per-head bytes for ssm_state.
-    # conv_state layout: [x: d_inner/tp | B: ng*ds/tp | C: ng*ds/tp] x (d_conv-1)
+    # The section ordering is supplied by the cache manager because Mamba2
+    # uses [x | B | C], while GDN uses [Q | K | V].
     # ssm_state layout: (nheads/tp, head_dim, d_state)
     d_conv_m1 = conv_state.shape[3]
     conv_elem_size = conv_state.element_size()
@@ -165,6 +174,91 @@ def _build_layer_group_for_mamba(
     )
 
 
+def _slot_stride_bytes(tensor: torch.Tensor) -> int:
+    return int(tensor.stride(0) * tensor.element_size())
+
+
+def _build_v2_mamba_state_pool(states: Sequence[torch.Tensor]) -> PhysicalPool:
+    """Describe affine layer/slot addressing for one V2 Mamba state role."""
+    if not states:
+        raise ValueError("V2 Mamba state pool requires at least one layer")
+
+    first_state = states[0]
+    base_address = int(first_state.data_ptr())
+    num_slots = int(first_state.shape[0])
+    slot_bytes = int(first_state[0].numel() * first_state.element_size())
+    slot_stride_bytes = _slot_stride_bytes(first_state)
+
+    num_layers = len(states)
+    if slot_stride_bytes % num_layers != 0:
+        raise ValueError("V2 Mamba physical slot must divide evenly across layers")
+    # Each role appears once per layer in its size-class pool. Equal-size SSM
+    # and convolution states share that pool and are interleaved, so their
+    # layer stride includes both role payloads.
+    layer_stride_bytes = slot_stride_bytes // num_layers
+
+    for layer_offset, state in enumerate(states):
+        state_slot_bytes = int(state[0].numel() * state.element_size())
+        if (
+            int(state.shape[0]) != num_slots
+            or state_slot_bytes != slot_bytes
+            or _slot_stride_bytes(state) != slot_stride_bytes
+        ):
+            raise ValueError("V2 Mamba state tensors must share one slot layout per role")
+        expected_address = base_address + layer_offset * layer_stride_bytes
+        if int(state.data_ptr()) != expected_address:
+            raise ValueError("V2 Mamba state tensors must have a uniform layer stride per role")
+
+    return PhysicalPool(
+        base_address=base_address,
+        slot_bytes=slot_bytes,
+        num_slots=num_slots,
+        slot_stride_bytes=slot_stride_bytes,
+        layer_stride_bytes=layer_stride_bytes,
+    )
+
+
+def _build_layer_group_for_v2_mamba(
+    manager: MambaHybridCacheManagerV2, 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()
+    }
+
+    expected_offsets = list(range(len(mamba_layer_offsets)))
+    if sorted(mamba_layer_offsets.values()) != expected_offsets:
+        raise ValueError("V2 Mamba layer offsets must be dense")
+    if len(manager.all_conv_states) != len(expected_offsets) or len(manager.all_ssm_states) != len(
+        expected_offsets
+    ):
+        raise ValueError("V2 Mamba state tensors must match the layer-offset table")
+
+    first_conv_state = manager.all_conv_states[0]
+    first_ssm_state = manager.all_ssm_states[0]
+    conv_pool = _build_v2_mamba_state_pool(manager.all_conv_states)
+    ssm_pool = _build_v2_mamba_state_pool(manager.all_ssm_states)
+    if conv_pool.num_slots != ssm_pool.num_slots:
+        raise ValueError("V2 Mamba convolution and SSM states must have the same number of slots")
+
+    d_conv_m1 = manager.conv_state_shape[1]
+    conv_elem_size = first_conv_state.element_size()
+    _, 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=conv_pool,
+        ssm_states=ssm_pool,
+        conv_section_bytes=conv_section_bytes,
+        ssm_bytes_per_head=ssm_bytes_per_head,
+    )
+
+
 def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable:
     """Build a KVCachePageTable from a KVCacheManager (V1)."""
     if kv_cache_manager.dtype == DataType.NVFP4:
@@ -419,6 +513,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, MambaHybridCacheManagerV2) 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 = [
@@ -519,7 +621,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, MambaHybridCacheManagerV2
+    ):
         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)