diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index e8c8816ba4c8..81feb03a25a7 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.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 typing import Dict, List, Optional, Tuple import numpy as np @@ -337,14 +351,28 @@ def _build_layer_ptrs( layer_offsets: Dict[int, int], overlapping_layers: List[int], slot: int, + layer_slot0_addresses: Optional[Dict[int, int]] = None, ) -> np.ndarray: - """Build per-layer pointers for a given pool (conv or ssm) and slot.""" + """Build per-layer pointers for a given pool (conv or SSM) and slot. + + V1 stores states layer-major, so its layer base is derived from the + Mamba-local layer offset. V2 stores buffers inside coalesced slot-major + pools; its manager-provided slot-0 addresses preserve the per-layer + offsets within that physical slot. + """ ptrs = [] + slot_stride_bytes = pool.slot_stride_bytes + assert slot_stride_bytes is not None 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 * slot_stride_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 +458,29 @@ 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 + ) 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, ) 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, ) src_region = SpecRegion( diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 92f0a9c2f4ad..ddf876bd7dbd 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.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 typing import Dict, List import numpy as np @@ -21,7 +35,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 @@ -73,10 +90,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) @@ -110,7 +129,8 @@ def _build_layer_group_for_mamba( ) # 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() @@ -132,6 +152,70 @@ def _build_layer_group_for_mamba( ) +def _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_slot_stride_bytes = _slot_stride_bytes(first_conv_state) + ssm_slot_stride_bytes = _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]) + + # V2 coalesces equal-size buffers into slot-major physical pools. The + # SHARED tensor bases include each layer/role's offset within slot 0, while + # stride(0) is the distance to the same buffer in the next physical slot. + # Preserve both pieces: V1's layer-major ``layer * num_slots`` formula does + # not describe this layout. + 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() + _, 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, + slot_stride_bytes=conv_slot_stride_bytes, + ), + ssm_states=PhysicalPool( + base_address=int(first_ssm_state.data_ptr()), + slot_bytes=ssm_slot_bytes, + num_slots=num_slots, + slot_stride_bytes=ssm_slot_stride_bytes, + ), + 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, + ) + + 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 +423,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 +484,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) diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 81514c06d6a0..28537ff9a5bb 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.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 dataclasses import dataclass, field @@ -48,12 +62,24 @@ class PhysicalPool: base_address: int # uint64 slot_bytes: int num_slots: int + # Distance between the starts of adjacent slots. Most pools are densely + # packed, so the stride defaults to the transferable payload size. V2 + # Mamba views point into a coalesced physical slot whose stride can be + # larger than the state payload described by ``slot_bytes``. + slot_stride_bytes: Optional[int] = None + + def __post_init__(self) -> None: + if self.slot_stride_bytes is None: + self.slot_stride_bytes = self.slot_bytes + if self.slot_stride_bytes < self.slot_bytes: + raise ValueError("slot_stride_bytes must be greater than or equal to slot_bytes") def to_dict(self) -> dict: return { "base_address": int(self.base_address), "slot_bytes": int(self.slot_bytes), "num_slots": int(self.num_slots), + "slot_stride_bytes": int(self.slot_stride_bytes), } @staticmethod @@ -62,6 +88,11 @@ def from_dict(data: dict) -> "PhysicalPool": base_address=int(data["base_address"]), slot_bytes=int(data["slot_bytes"]), num_slots=int(data["num_slots"]), + slot_stride_bytes=( + int(data["slot_stride_bytes"]) + if data.get("slot_stride_bytes") is not None + else None + ), ) @@ -204,6 +235,11 @@ class MambaLayerGroup(LayerGroup): ssm_states: Optional[PhysicalPool] = None conv_section_bytes: Optional[List[int]] = None ssm_bytes_per_head: Optional[int] = None + # V2 pools are slot-major and may coalesce several layer/role buffers into + # one physical slot. These are the manager-provided buffer offsets within + # slot 0; they cannot be derived with V1's layer-major pointer formula. + conv_layer_slot0_addresses: Optional[Dict[int, int]] = None + ssm_layer_slot0_addresses: Optional[Dict[int, int]] = None def to_dict(self) -> dict: return { @@ -213,12 +249,24 @@ 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, } @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 +276,14 @@ 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, ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 21c4d98bd2aa..874911f81fd9 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/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 @@ -10,7 +24,7 @@ def get_pool_bytes(pool: PhysicalPool) -> int: - """Total bytes across all slots in this pool.""" + """Total transferable payload bytes across all slots in this pool.""" return pool.slot_bytes * pool.num_slots @@ -18,7 +32,8 @@ def get_slot_address(pool: PhysicalPool, slot_id: int) -> int: """Base address of *slot_id*.""" if slot_id >= pool.num_slots: raise ValueError(f"slot_id {slot_id} >= num_slots {pool.num_slots}") - return pool.base_address + slot_id * pool.slot_bytes + assert pool.slot_stride_bytes is not None + return pool.base_address + slot_id * pool.slot_stride_bytes # ------------------------------------------------------------------------- @@ -117,9 +132,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 diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 4fe7fa28e50f..3bcc7ec72d6f 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/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,10 @@ 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): + if 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( diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 5c2e60c4706a..2938c88507b1 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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 hybrid-state transfers.""" + return "PYTHON" + @staticmethod def lora_config(model_dir: str): """Nemotron-H-specific LoRA configuration. diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 6fc43129a353..4bd9de60c284 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/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 Python 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) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 0d91b4ebad74..87925dbebe06 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -18,7 +18,7 @@ import copy import os from types import SimpleNamespace -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Literal, Optional import torch @@ -992,6 +992,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 hybrid-state transfers.""" + return "PYTHON" + def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 032e8e1881fa..23b59ff2f1f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -88,6 +88,22 @@ def _non_hybrid_kv_cache_manager_cls(config, kv_cache_config: KvCacheConfig): return KVCacheManagerV2 if needs_v2 else KVCacheManager +def _resolve_disagg_transceiver_route( + cache_transceiver_config: Optional[CacheTransceiverConfig], +) -> tuple[Optional[str], Optional[str]]: + """Return the effective backend and runtime used for manager routing.""" + if cache_transceiver_config is None: + return None, None + + backend, _ = cache_transceiver_config._resolve_default_backend() + runtime = cache_transceiver_config.transceiver_runtime + if runtime == "auto": + # Model loading normally resolves ``auto``. Paths that skip model + # defaults use the global C++ fallback, matching transceiver creation. + runtime = None + return backend, runtime + + def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -101,8 +117,14 @@ def get_kv_cache_manager_cls( unified-pool default. V1 is the default for hybrid Mamba models. V2 is selected only by an - explicit ``kv_cache_config.use_kv_cache_manager_v2=True`` and is rejected - for disaggregated serving until its transceiver/page-table adapter exists. + explicit ``kv_cache_config.use_kv_cache_manager_v2=True``. In + disaggregated serving, V2 additionally requires the Python transceiver + with the NIXL backend. Unsupported explicit V2 routes fail rather than + falling back to a different manager. + + Env-var overrides: + * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager in aggregated serving. + * ``TLLM_MAMBA_MANAGER_PREFERENCE`` — explicit manager preference. """ config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config @@ -131,12 +153,6 @@ def get_kv_cache_manager_cls( or state_config.additional_snapshot_offsets_from_end) use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True - if is_disagg and use_v2: - raise ValueError( - "KV cache manager V2 for hybrid Mamba models is not " - "supported with disaggregated serving. Set " - "use_kv_cache_manager_v2=False or 'auto' to use a V1 " - "Mamba cache manager.") if has_additional_snapshots and not use_v2: raise ValueError("Mamba additional snapshot offsets require " "use_kv_cache_manager_v2=True; V1 supports only " @@ -151,16 +167,29 @@ def get_kv_cache_manager_cls( # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. if is_disagg: - if kv_cache_config.enable_block_reuse: + backend, runtime = _resolve_disagg_transceiver_route( + cache_transceiver_config) + if use_v2: + if runtime != "PYTHON" or backend != "NIXL": + raise ValueError( + "KV cache manager V2 for hybrid Mamba disaggregated " + "serving requires transceiver_runtime='PYTHON' with " + "backend='NIXL'.") + else: + if (kv_cache_config.enable_block_reuse and runtime == "PYTHON"): + raise ValueError( + "Hybrid Mamba disaggregated serving with block reuse " + "and transceiver_runtime='PYTHON' requires " + "use_kv_cache_manager_v2=True.") + if kv_cache_config.enable_block_reuse: + return CppMambaHybridCacheManager + if runtime == "PYTHON" and backend == "NIXL": + logger.info("Python transceiver detected; using " + "MixedMambaHybridCacheManager for hybrid model") + return MixedMambaHybridCacheManager return CppMambaHybridCacheManager - if (cache_transceiver_config is not None and - cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.info("Python transceiver detected; using " - "MixedMambaHybridCacheManager for hybrid model") - return MixedMambaHybridCacheManager - return CppMambaHybridCacheManager - if use_py_mamba_cache_manager(): + if use_py_mamba_cache_manager() and not is_disagg: if use_v2: raise ValueError( "TRTLLM_USE_PY_MAMBA=1 conflicts with explicit " @@ -1913,6 +1942,8 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats + if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + manager_extra_kwargs["is_disagg"] = is_disagg if is_mla(config): kv_cache_manager = kv_cache_manager_cls( @@ -2016,7 +2047,9 @@ def _create_kv_cache_manager( and mamba_params.mamba_ssm_cache_dtype == torch.float16) mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - if not issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + mamba_manager_extra_kwargs["conv_state_layout"] = "x_b_c" + else: mamba_manager_extra_kwargs["model_type"] = "nemotron_hybrid" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters @@ -2115,7 +2148,9 @@ def _create_kv_cache_manager( ("ENABLED" if use_replay else "DISABLED")) mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - if not issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + if issubclass(kv_cache_manager_cls, V2MambaHybridCacheManager): + mamba_manager_extra_kwargs["conv_state_layout"] = "q_k_v" + else: mamba_manager_extra_kwargs["model_type"] = "qwen3_next" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 348d47a9659c..9a2ed0ef52fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -16,7 +16,8 @@ from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, - MixedMambaHybridCacheManager) + MixedMambaHybridCacheManager, + V2MambaHybridCacheManager) from .resource_manager import KVCacheManager CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver @@ -157,14 +158,33 @@ def create_kv_cache_transceiver( "UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server " "hangs or lower-than-expected performance.") - # Select transceiver implementation based on transceiver_runtime + # Select transceiver implementation based on transceiver_runtime. # transceiver_runtime == None or "CPP" -> use C++ transceiver (default) - # transceiver_runtime == "PYTHON" -> use Python transceiver - if cache_transceiver_config.transceiver_runtime == "PYTHON": - # Python transceiver currently only supports NIXL and DEFAULT backend - if cache_transceiver_config.backend not in ("DEFAULT", "NIXL"): + # transceiver_runtime == "PYTHON" -> use Python transceiver. + # + # V2MambaHybridCacheManager is backed by the Python KVCacheManagerV2 core, + # not the C++ BaseKVCacheManager binding required by CacheTransceiverCpp. + is_v2_mamba_hybrid = isinstance(mamba_cache_manager, + V2MambaHybridCacheManager) + use_python_transceiver = ( + cache_transceiver_config.transceiver_runtime == "PYTHON") + + if is_v2_mamba_hybrid and not use_python_transceiver: + raise ValueError( + "V2MambaHybridCacheManager requires transceiver_runtime='PYTHON' " + "with backend='NIXL'; it cannot use the C++ transceiver.") + + if use_python_transceiver: + if isinstance(mamba_cache_manager, CppMambaHybridCacheManager): + raise ValueError( + "transceiver_runtime='PYTHON' cannot drive " + "CppMambaHybridCacheManager (C++ pool backed). Use " + "transceiver_runtime='CPP', or select the V2 manager " + "with use_kv_cache_manager_v2=True.") + # DEFAULT has already been resolved above, so Python must see NIXL. + if cache_transceiver_config.backend != "NIXL": raise ValueError( - f"Python transceiver currently only supports NIXL or DEFAULT backend, " + f"Python transceiver currently only supports the NIXL backend, " f"got {cache_transceiver_config.backend}. " f"Please use transceiver_runtime='CPP' for MPI, UCX, or MOONCAKE backends." ) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 96550e5fe53e..b3de90d0d779 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -17,8 +17,8 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import (TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, - Tuple, Union) +from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, + Optional, Tuple, Union) import torch import triton @@ -2466,8 +2466,12 @@ def __init__( is_draft: bool = False, use_replay_state_update: bool = False, mamba_ssm_stochastic_rounding: bool = False, + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", **kwargs, ) -> None: + if conv_state_layout not in ("x_b_c", "q_k_v"): + raise ValueError( + f"Unsupported convolution state layout: {conv_state_layout!r}") total_layers = len(mamba_layer_mask) if layer_mask is None: full_attention_layer_mask = [False] * total_layers @@ -2518,18 +2522,42 @@ def __init__( if self.local_num_mamba_layers > 0: tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 d_inner = mamba_head_dim * mamba_num_heads - conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state + grouped_state_dim = mamba_n_groups * mamba_d_state + conv_dim = d_inner + 2 * grouped_state_dim nheads = mamba_num_heads assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if kwargs.get("is_disagg", + False) and grouped_state_dim % tp_size != 0: + raise ValueError( + "Disaggregated Mamba transfer requires each convolution " + "state section to be divisible by tp_size") if use_replay_state_update: assert mamba_n_groups % tp_size == 0, \ "replay state update requires mamba_n_groups divisible by tp_size" self._n_groups_per_rank = mamba_n_groups // tp_size + d_inner_local = d_inner // tp_size + grouped_state_dim_local = grouped_state_dim // tp_size conv_dim = conv_dim // tp_size nheads = nheads // tp_size self.conv_state_shape = [conv_dim, mamba_d_conv - 1] self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + # TP-mismatch disaggregated transfers must split the flat + # convolution state at its true semantic boundaries. Mamba2 stores + # [x | B | C], while GDN stores [Q | K | V]. The large section is + # therefore first for Mamba2 and last for GDN. + if conv_state_layout == "x_b_c": + self.conv_section_dims = [ + d_inner_local, + grouped_state_dim_local, + grouped_state_dim_local, + ] + else: + self.conv_section_dims = [ + grouped_state_dim_local, + grouped_state_dim_local, + d_inner_local, + ] self.ssm_count = math.prod(self.ssm_state_shape) self.conv_count = math.prod(self.conv_state_shape) self.ssm_bytes = self.ssm_count * self.ssm_state_dtype.itemsize @@ -2541,6 +2569,7 @@ def __init__( self._n_groups_per_rank = 0 self.conv_state_shape = [] self.ssm_state_shape = [] + self.conv_section_dims = [] self.ssm_count = 0 self.conv_count = 0 self.ssm_bytes = 0 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 08ceedf61054..a8dd45606309 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -680,16 +680,6 @@ def drafting_loop_wrapper(model): config = model_engine.model.model_config.pretrained_config max_num_seq_slots = getattr(model_engine, "max_num_seq_slots", max_batch_size * getattr(mapping, "pp_size", 1)) - if is_hybrid_linear(config) and kv_cache_config.enable_block_reuse and ( - cache_transceiver_config is not None - and cache_transceiver_config.backend is not None - and cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.warning( - "Disabling block reuse for MambaHybridCacheManager-based models when disagg + Python transceiver enabled" - ) - kv_cache_config.enable_block_reuse = False - _set_model_engines_cache_reuse([model_engine, draft_model_engine], - False) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 @@ -907,17 +897,16 @@ def drafting_loop_wrapper(model): if is_disagg and is_hybrid: # NOTE: TRTLLM_USE_PY_MAMBA is an agg-mode-only override and has - # no effect in disagg. The disagg manager choice is driven solely - # by transceiver_runtime: PYTHON => PythonMambaCacheManager, - # otherwise CppMambaHybridCacheManager (unified pool, default). + # no effect in disagg. The disagg manager choice is driven by + # get_kv_cache_manager_cls and cache_transceiver_config. if os.environ.get("TRTLLM_USE_PY_MAMBA", "0") == "1": logger.warning( "TRTLLM_USE_PY_MAMBA is ignored in disaggregated serving; " - "use cache_transceiver_config.transceiver_runtime='PYTHON' " - "to select PythonMambaCacheManager.") + "configure transceiver_runtime='PYTHON' with backend='NIXL' " + "to select MixedMambaHybridCacheManager.") else: logger.info("Disaggregated serving with hybrid model detected. " - "Using CppMambaHybridCacheManager.") + "Using the configured Mamba cache manager.") # Get draft config for one-engine speculative decoding if available draft_config = getattr(model_engine.model, 'draft_config', None) diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 3d9a3f83dcc3..f5bdc6de75e4 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -23,6 +23,10 @@ from tensorrt_llm._torch.pyexecutor import py_executor as executor_module from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + CppMambaHybridCacheManager, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig @@ -519,6 +523,59 @@ def test_flag_unset_preserves_python_transceiver(monkeypatch): constructor.assert_called_once() +def test_python_nixl_transceiver_accepts_v2_mamba_manager(monkeypatch): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + expected = object() + constructor = Mock(return_value=expected) + fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) + manager = object.__new__(V2MambaHybridCacheManager) + + result = transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + assert result is expected + constructor.assert_called_once() + + +@pytest.mark.parametrize("runtime", [None, "CPP", "auto"]) +def test_cpp_runtime_rejects_v2_mamba_manager(runtime): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(V2MambaHybridCacheManager) + + with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + +def test_python_runtime_rejects_cpp_mamba_manager(): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + manager = object.__new__(CppMambaHybridCacheManager) + + with pytest.raises(ValueError, match="cannot drive CppMambaHybridCacheManager"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + +@pytest.mark.parametrize("runtime", [None, "CPP"]) +def test_cpp_runtime_keeps_cpp_mamba_manager(monkeypatch, runtime): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(CppMambaHybridCacheManager) + expected = object() + constructor = Mock(return_value=expected) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + result = transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + assert result is expected + constructor.assert_called_once() + + def test_flag_unset_preserves_libfabric_selection(monkeypatch): monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, "LIBFABRIC") config = CacheTransceiverConfig(backend="NIXL") diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index e28046cd8bfb..db6f8aaa1b2d 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -9,6 +9,9 @@ import pytest import torch +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager +from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup +from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -46,12 +49,16 @@ from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, KvCacheConfig, MambaStateConfig, MTPDecodingConfig, TorchLlmArgs, ) -from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto +from tensorrt_llm.llmapi.llm_utils import ( + _resolve_kv_cache_manager_v2_auto, + _resolve_transceiver_runtime_auto, +) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import GpuCacheTierConfig, LayerId from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager @@ -240,6 +247,7 @@ def __init__(self, *args, **kwargs): assert captured_cpp["model_type"] == "qwen3_next" assert captured_v2["use_replay_state_update"] is False assert "model_type" not in captured_v2 + assert captured_v2["conv_state_layout"] == "q_k_v" def test_hybrid_cache_manager_factory_rejects_cpp_preference_with_explicit_v2( @@ -357,6 +365,93 @@ def test_hybrid_cache_manager_factory_routes_explicit_snapshots_to_v2( ) +@pytest.mark.parametrize("backend", ["NIXL", "DEFAULT"]) +def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, backend): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime="PYTHON" + ), + ) + is V2MambaHybridCacheManager + ) + + +def test_hybrid_cache_manager_factory_rejects_python_v1_disagg_reuse(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + with pytest.raises(ValueError, match="requires use_kv_cache_manager_v2=True"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=False, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + ), + ) + + +@pytest.mark.parametrize( + ("backend", "runtime", "backend_env"), + [ + ("DEFAULT", "PYTHON", "TRTLLM_USE_UCX_KVCACHE"), + ("UCX", "PYTHON", None), + ("NIXL", "auto", None), + ("NIXL", None, None), + ("NIXL", "CPP", None), + ("UCX", None, None), + ], +) +def test_hybrid_cache_manager_factory_rejects_unsupported_v2_disagg_route( + monkeypatch, backend, runtime, backend_env +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + if backend_env is not None: + monkeypatch.setenv(backend_env, "1") + + with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime + ), + ) + + @pytest.mark.parametrize( ("env_name", "env_value", "expected_error"), [ @@ -413,35 +508,52 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2) use_kv_cache_manager_v2=use_v2, ), is_disagg=True, - cache_transceiver_config=SimpleNamespace(transceiver_runtime="PYTHON"), + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON" + ), ) is MixedMambaHybridCacheManager ) -@pytest.mark.parametrize("enable_block_reuse", [False, True]) -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -def test_hybrid_cache_manager_factory_rejects_v2_disagg( - monkeypatch, enable_block_reuse, transceiver_runtime -): - monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) - monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - transceiver_config = ( - None - if transceiver_runtime is None - else SimpleNamespace(transceiver_runtime=transceiver_runtime) - ) +def test_hybrid_models_resolve_auto_to_python_transceiver(monkeypatch): + from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM + from tensorrt_llm._torch.models.modeling_qwen3_5 import Qwen3_5VLModel + from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM - with pytest.raises(ValueError, match="V2.*not supported with disaggregated"): - get_kv_cache_manager_cls( - _hybrid_model_config(), - KvCacheConfig( - enable_block_reuse=enable_block_reuse, - use_kv_cache_manager_v2=True, - ), - is_disagg=True, - cache_transceiver_config=transceiver_config, + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + + for model_cls in (NemotronHForCausalLM, Qwen3NextForCausalLM, Qwen3_5VLModel): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), ) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): + manager = object.__new__(V2MambaHybridCacheManager) + manager.local_num_mamba_layers = 0 + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._kv_cache_manager = manager + transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) + transceiver._page_table = SimpleNamespace(layer_groups=[]) + request = SimpleNamespace( + is_generation_only_request=lambda: False, + prompt_len=0, + py_request_id=123, + ) + + kv_slice = transceiver._create_kv_slice(request) + + assert kv_slice.mamba_state_index is None @pytest.mark.parametrize( @@ -1453,6 +1565,7 @@ def _build_v2_hybrid_with_mamba_layer( enable_attention_dp=False, enable_swa_scratch_reuse=False, dtype=DataType.HALF, + conv_state_layout="x_b_c", ): """Construct a real V2MambaHybridCacheManager.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -1496,6 +1609,7 @@ def _build_v2_hybrid_with_mamba_layer( vocab_size=1024, use_replay_state_update=use_replay_state_update, dtype=dtype, + conv_state_layout=conv_state_layout, ) @@ -1841,6 +1955,47 @@ def test_v2_hybrid_swa_scratch_keeps_ssm_placeholder_rows(): mgr.shutdown() +@skip_no_cuda +def test_v2_hybrid_disagg_page_table_preserves_lifecycle_indices(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4, num_mamba_layers=2) + try: + page_table = build_page_table_from_manager(mgr) + + assert len(page_table.layer_groups) == mgr.impl._storage.num_life_cycles + assert isinstance(page_table.layer_groups[0], MambaLayerGroup) + assert isinstance(page_table.layer_groups[1], AttentionLayerGroup) + + requests = mgr.add_dummy_requests([123], token_nums=[64], is_gen=False) + assert len(requests) == 1 + attention_blocks = list( + mgr.kv_cache_map[123].get_aggregated_page_indices(1, valid_only=True) + ) + assert attention_blocks + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_disagg_page_table_uses_qwen3_next_conv_sections(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + conv_state_layout="q_k_v", + ) + try: + page_table = build_page_table_from_manager(mgr) + mamba_group = page_table.layer_groups[0] + + assert isinstance(mamba_group, MambaLayerGroup) + d_conv_m1 = mgr.conv_state_shape[1] + conv_elem_size = mgr.all_conv_states[0].element_size() + assert mamba_group.conv_section_bytes == [ + dim * d_conv_m1 * conv_elem_size for dim in mgr.conv_section_dims + ] + assert mgr.conv_section_dims == [8, 8, 32] + finally: + mgr.shutdown() + + @skip_no_cuda def test_cpp_hybrid_replay_buffers_size_by_tokens_per_gen_step(): spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 6fc2f68fef68..e31180ce18ec 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.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. + import numpy as np import pytest @@ -13,6 +27,7 @@ get_num_layer_groups, get_num_layers, get_physical_pool, + get_slot_address, get_unique_layers, ) from tensorrt_llm._torch.pyexecutor.resource_manager import ( @@ -235,8 +250,18 @@ def test_layer_group_meta_serialization(): def test_mamba_layer_group_serialization(): from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup, PhysicalPool - conv_pool = PhysicalPool(base_address=1000, slot_bytes=128, num_slots=10) - ssm_pool = PhysicalPool(base_address=8000, slot_bytes=256, num_slots=8) + conv_pool = PhysicalPool( + base_address=1000, + slot_bytes=128, + num_slots=10, + slot_stride_bytes=512, + ) + ssm_pool = PhysicalPool( + base_address=8000, + slot_bytes=256, + num_slots=8, + slot_stride_bytes=1024, + ) mlg = MambaLayerGroup( pool_group_idx=1, mamba_layer_offsets={10: 0, 11: 1, 12: 2}, @@ -244,11 +269,17 @@ def test_mamba_layer_group_serialization(): ssm_states=ssm_pool, conv_section_bytes=[512, 256, 256], ssm_bytes_per_head=128, + conv_layer_slot0_addresses={10: 1000, 11: 2000, 12: 3000}, + ssm_layer_slot0_addresses={10: 8000, 11: 9000, 12: 10000}, ) d = mlg.to_dict() assert d["mamba_layer_offsets"] == {10: 0, 11: 1, 12: 2} assert d["conv_section_bytes"] == [512, 256, 256] + assert d["conv_layer_slot0_addresses"] == {10: 1000, 11: 2000, 12: 3000} + assert d["ssm_layer_slot0_addresses"] == {10: 8000, 11: 9000, 12: 10000} + assert d["conv_states"]["slot_stride_bytes"] == 512 + assert d["ssm_states"]["slot_stride_bytes"] == 1024 from tensorrt_llm._torch.disaggregation.resource.page import LayerGroup @@ -258,11 +289,102 @@ def test_mamba_layer_group_serialization(): assert restored.conv_states.base_address == 1000 assert restored.conv_states.slot_bytes == 128 assert restored.conv_states.num_slots == 10 + assert restored.conv_states.slot_stride_bytes == 512 + assert get_slot_address(restored.conv_states, 3) == 1000 + 3 * 512 assert restored.ssm_states.base_address == 8000 assert restored.ssm_states.slot_bytes == 256 assert restored.ssm_states.num_slots == 8 + assert restored.ssm_states.slot_stride_bytes == 1024 assert restored.conv_section_bytes == [512, 256, 256] assert restored.ssm_bytes_per_head == 128 + assert restored.conv_layer_slot0_addresses == {10: 1000, 11: 2000, 12: 3000} + assert restored.ssm_layer_slot0_addresses == {10: 8000, 11: 9000, 12: 10000} + + legacy_pool = PhysicalPool.from_dict({"base_address": 1000, "slot_bytes": 128, "num_slots": 10}) + assert legacy_pool.slot_stride_bytes == legacy_pool.slot_bytes + + +def test_v2_mamba_registration_uses_coalesced_physical_pool(): + from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MambaLayerGroup, + PhysicalPool, + PhysicalPoolGroup, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + state_bytes = 64 + num_layers = 2 + num_slots = 8 + # Equal-sized SSM and convolution states share one interleaved V2 pool. + physical_slot_bytes = state_bytes * num_layers * 2 + physical_pool = PhysicalPool( + base_address=1000, + slot_bytes=physical_slot_bytes, + num_slots=num_slots, + ) + mamba_group = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=PhysicalPool( + base_address=1000 + state_bytes, + slot_bytes=state_bytes, + num_slots=num_slots, + slot_stride_bytes=physical_slot_bytes, + ), + ssm_states=PhysicalPool( + base_address=1000, + slot_bytes=state_bytes, + num_slots=num_slots, + slot_stride_bytes=physical_slot_bytes, + ), + conv_layer_slot0_addresses={ + 1: 1000 + state_bytes, + 2: 1000 + state_bytes * 3, + }, + ssm_layer_slot0_addresses={ + 1: 1000, + 2: 1000 + state_bytes * 2, + }, + ) + page_table = KVCachePageTable( + tokens_per_block=16, + layer_groups=[mamba_group], + pool_groups=[PhysicalPoolGroup(pools=[physical_pool])], + ) + + assert get_unique_pool_memory_descs(page_table, device_id=3) == [ + (1000, physical_slot_bytes * num_slots, 3, "kv_cache_memory_pool0") + ] + + +def test_legacy_mamba_registration_uses_layer_major_pools(): + from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MambaLayerGroup, + PhysicalPool, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + num_layers = 3 + conv_pool = PhysicalPool(base_address=1000, slot_bytes=128, num_slots=10) + ssm_pool = PhysicalPool(base_address=8000, slot_bytes=256, num_slots=8) + mamba_group = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={10: 0, 11: 1, 12: 2}, + conv_states=conv_pool, + ssm_states=ssm_pool, + ) + page_table = KVCachePageTable( + tokens_per_block=16, + layer_groups=[mamba_group], + pool_groups=[], + ) + + assert get_unique_pool_memory_descs(page_table, device_id=3) == [ + (1000, num_layers * conv_pool.num_slots * conv_pool.slot_bytes, 3, "kv_cache_memory_pool0"), + (8000, num_layers * ssm_pool.num_slots * ssm_pool.slot_bytes, 3, "kv_cache_memory_pool1"), + ] def test_mixed_page_table_serialization(): diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 0cf0b3209899..0b5c3716c95c 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -16,6 +16,7 @@ import uuid from typing import Dict, List +import numpy as np import pytest import torch @@ -23,13 +24,19 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native import rank_info +from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer +from tensorrt_llm._torch.disaggregation.resource import page from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, LlmRequestType, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MixedMambaHybridCacheManager, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp @@ -183,8 +190,14 @@ def _init(rank): return results -def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=False): - """Create MixedMambaHybridCacheManagers for all TP ranks (PP=1). +def _create_managers( + tp, + max_batch_size=MAX_BATCH_SIZE, + enable_attention_dp=False, + use_v2=False, + conv_state_layout="x_b_c", +): + """Create Mamba hybrid cache managers for all TP ranks (PP=1). Layer 0 is a dummy attention layer required by page table infrastructure. Layers 1..NUM_MAMBA_LAYERS are mamba layers under test. @@ -194,7 +207,16 @@ def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=Fals mapping = Mapping( world_size=tp, rank=rank, tp_size=tp, pp_size=1, enable_attention_dp=enable_attention_dp ) - mgr = MixedMambaHybridCacheManager( + manager_cls = V2MambaHybridCacheManager if use_v2 else MixedMambaHybridCacheManager + manager_kwargs = ( + { + "is_disagg": True, + "conv_state_layout": conv_state_layout, + } + if use_v2 + else {} + ) + mgr = manager_cls( mamba_d_state=MAMBA_D_STATE, mamba_d_conv=MAMBA_D_CONV, mamba_num_heads=MAMBA_NUM_HEADS, @@ -220,19 +242,169 @@ def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=Fals max_batch_size=max_batch_size, mapping=mapping, dtype=DataType.FLOAT, + **manager_kwargs, ) managers.append(mgr) return managers +def _mamba_layer_ids(manager): + if isinstance(manager, V2MambaHybridCacheManager): + return manager.mamba_layer_offsets + return manager._impl.mamba_layer_offsets + + +def _mamba_state_slot(manager, request_id): + if isinstance(manager, V2MambaHybridCacheManager): + return manager.get_state_indices([request_id])[0] + return manager.mamba_cache_index[request_id] + + +def _zero_mamba_states(manager): + for layer_idx in _mamba_layer_ids(manager): + manager.get_conv_states(layer_idx).zero_() + manager.get_ssm_states(layer_idx).zero_() + + +def test_mamba_policy_layer_major_v1_ptrs(): + pool = page.PhysicalPool(base_address=100, slot_bytes=10, num_slots=8) + + ptrs = peer.MambaPolicy._build_layer_ptrs( + pool=pool, + layer_offsets={1: 0, 2: 1}, + overlapping_layers=[1, 2], + slot=3, + ) + + np.testing.assert_array_equal(ptrs, [130, 210]) + + +def test_mamba_policy_slot_major_layer_ptrs(): + """V2 Mamba state tensors step by physical slots, not V1 layers.""" + self_mlg = page.MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=page.PhysicalPool( + base_address=1000, + slot_bytes=10, + num_slots=8, + slot_stride_bytes=20, + ), + ssm_states=page.PhysicalPool( + base_address=3000, + slot_bytes=20, + num_slots=8, + slot_stride_bytes=40, + ), + conv_section_bytes=[10], + ssm_bytes_per_head=10, + conv_layer_slot0_addresses={1: 1000, 2: 1010}, + ssm_layer_slot0_addresses={1: 3000, 2: 3020}, + ) + peer_mlg = page.MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=page.PhysicalPool( + base_address=5000, + slot_bytes=10, + num_slots=8, + slot_stride_bytes=20, + ), + ssm_states=page.PhysicalPool( + base_address=7000, + slot_bytes=20, + num_slots=8, + slot_stride_bytes=40, + ), + conv_section_bytes=[10], + ssm_bytes_per_head=10, + conv_layer_slot0_addresses={1: 5000, 2: 5010}, + ssm_layer_slot0_addresses={1: 7000, 2: 7020}, + ) + self_ri = rank_info.RankInfo( + instance_name="self", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=bytes(), + ) + peer_ri = rank_info.RankInfo( + instance_name="peer", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=bytes(), + ) + + src_frags, dst_frags, kv_sizes = peer.MambaPolicy.build_mamba_frags( + self_mlg=self_mlg, + peer_mlg=peer_mlg, + src_slot=3, + dst_slot=5, + self_ri=self_ri, + peer_ri=peer_ri, + ) + + assert src_frags == [1060, 1070, 3120, 3140] + assert dst_frags == [5100, 5110, 7200, 7220] + assert kv_sizes == [10, 10, 20, 20] + + +def test_mamba_policy_slot_major_interleaved_role_ptrs(): + """Layer/role offsets remain inside each coalesced V2 physical slot.""" + state_bytes = 64 + physical_slot_bytes = 4 * state_bytes + conv_pool = page.PhysicalPool( + base_address=1000 + state_bytes, + slot_bytes=state_bytes, + num_slots=8, + slot_stride_bytes=physical_slot_bytes, + ) + + ptrs = peer.MambaPolicy._build_layer_ptrs( + pool=conv_pool, + layer_offsets={1: 0, 2: 1}, + overlapping_layers=[1, 2], + slot=3, + layer_slot0_addresses={ + 1: 1000 + state_bytes, + 2: 1000 + 3 * state_bytes, + }, + ) + + np.testing.assert_array_equal( + ptrs, + [ + 1000 + state_bytes + 3 * physical_slot_bytes, + 1000 + 3 * state_bytes + 3 * physical_slot_bytes, + ], + ) + + # --------------------------------------------------------------------------- # Ground truth: generate, shard, write, compute expected, read actual # --------------------------------------------------------------------------- -def _full_conv_section_dims() -> List[int]: - """Full (unsharded) first-dim sizes: [x(d_inner) | B(ng*ds) | C(ng*ds)].""" +def _full_conv_section_dims(conv_state_layout="x_b_c") -> List[int]: + """Full first-dimension sizes in the model's convolution-state order.""" d_inner = MAMBA_HEAD_DIM * MAMBA_NUM_HEADS ng_ds = MAMBA_N_GROUPS * MAMBA_D_STATE - return [d_inner, ng_ds, ng_ds] + if conv_state_layout == "x_b_c": + return [d_inner, ng_ds, ng_ds] + if conv_state_layout == "q_k_v": + return [ng_ds, ng_ds, d_inner] + raise ValueError(f"Unsupported convolution state layout: {conv_state_layout!r}") def _generate_ground_truth(num_requests: int, seed: int = 12345): @@ -270,29 +442,48 @@ def _shard_ssm(full_ssm: torch.Tensor, tp: int, tp_rank: int) -> torch.Tensor: return full_ssm[tp_rank * n : (tp_rank + 1) * n].clone() -def _shard_conv(full_conv: torch.Tensor, tp: int, tp_rank: int) -> torch.Tensor: - """Shard conv per-section along dim 0: [x | B | C] each independently.""" +def _shard_conv( + full_conv: torch.Tensor, + tp: int, + tp_rank: int, + conv_state_layout="x_b_c", +) -> torch.Tensor: + """Shard each semantic convolution-state section independently.""" parts = [] offset = 0 - for sec_dim in _full_conv_section_dims(): + for sec_dim in _full_conv_section_dims(conv_state_layout): n = sec_dim // tp parts.append(full_conv[offset + tp_rank * n : offset + (tp_rank + 1) * n]) offset += sec_dim return torch.cat(parts, dim=0).clone() -def _write_ground_truth_to_ctx(managers, tp, ground_truth, request_ids): +def _write_ground_truth_to_ctx( + managers, + tp, + ground_truth, + request_ids, + conv_state_layout="x_b_c", +): """Write sharded ground truth into ctx managers' allocated mamba slots.""" for rank, mgr in enumerate(managers): for req_idx, rid in enumerate(request_ids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: + slot = _mamba_state_slot(mgr, rid) + for layer_idx in _mamba_layer_ids(mgr): full = ground_truth[req_idx][layer_idx] mgr.get_ssm_states(layer_idx)[slot] = _shard_ssm(full["ssm"], tp, rank) - mgr.get_conv_states(layer_idx)[slot] = _shard_conv(full["conv"], tp, rank) - - -def _compute_expected(ground_truth, gen_managers, gen_tp, gen_request_ids) -> Dict: + mgr.get_conv_states(layer_idx)[slot] = _shard_conv( + full["conv"], tp, rank, conv_state_layout + ) + + +def _compute_expected( + ground_truth, + gen_managers, + gen_tp, + gen_request_ids, + conv_state_layout="x_b_c", +) -> Dict: """Compute expected mamba states BEFORE transfer. Returns: {(gen_rank, req_idx, layer_idx): {"conv": Tensor, "ssm": Tensor}} @@ -300,11 +491,11 @@ def _compute_expected(ground_truth, gen_managers, gen_tp, gen_request_ids) -> Di expected = {} for gen_rank, mgr in enumerate(gen_managers): for req_idx in range(len(gen_request_ids)): - for layer_idx in mgr._impl.mamba_layer_offsets: + for layer_idx in _mamba_layer_ids(mgr): full = ground_truth[req_idx][layer_idx] expected[(gen_rank, req_idx, layer_idx)] = { "ssm": _shard_ssm(full["ssm"], gen_tp, gen_rank), - "conv": _shard_conv(full["conv"], gen_tp, gen_rank), + "conv": _shard_conv(full["conv"], gen_tp, gen_rank, conv_state_layout), } return expected @@ -317,8 +508,8 @@ def _read_actual(gen_managers, gen_request_ids) -> Dict: actual = {} for gen_rank, mgr in enumerate(gen_managers): for req_idx, rid in enumerate(gen_request_ids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: + slot = _mamba_state_slot(mgr, rid) + for layer_idx in _mamba_layer_ids(mgr): actual[(gen_rank, req_idx, layer_idx)] = { "conv": mgr.get_conv_states(layer_idx)[slot].cpu().clone(), "ssm": mgr.get_ssm_states(layer_idx)[slot].cpu().clone(), @@ -358,14 +549,26 @@ def test_mamba_disagg_attention_dp_dummy_with_batch_size_one(): mgr.shutdown() -def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): +def run_mamba_transfer_test( + ctx_tp: int, + gen_tp: int, + use_v2: bool = False, + conv_state_layout: str = "x_b_c", +): """Test mamba transfer: ctx_tp -> gen_tp (PP=1, no DP).""" # -- 1. Create managers, zero mamba caches -- - ctx_mgrs = _create_managers(ctx_tp) - gen_mgrs = _create_managers(gen_tp) + ctx_mgrs = _create_managers( + ctx_tp, + use_v2=use_v2, + conv_state_layout=conv_state_layout, + ) + gen_mgrs = _create_managers( + gen_tp, + use_v2=use_v2, + conv_state_layout=conv_state_layout, + ) for mgr in ctx_mgrs + gen_mgrs: - mgr._impl.mamba_cache.conv.zero_() - mgr._impl.mamba_cache.temporal.zero_() + _zero_mamba_states(mgr) # -- 2. Create transceivers -- config = CacheTransceiverConfig( @@ -419,14 +622,29 @@ def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): # -- 4. Allocate slots -- ctx_batch = ScheduledRequests() - ctx_batch.reset_context_requests(ctx_reqs) - for mgr in ctx_mgrs: - mgr.prepare_resources(ctx_batch) + if use_v2: + ctx_batch.context_requests_last_chunk = ctx_reqs + for mgr in ctx_mgrs: + for req in ctx_reqs: + assert mgr.prepare_context(req) + assert mgr.resize_context(req, req.context_chunk_size) + mgr.prepare_resources(ctx_batch) + else: + ctx_batch.reset_context_requests(ctx_reqs) + for mgr in ctx_mgrs: + mgr.prepare_resources(ctx_batch) gen_batch = ScheduledRequests() - gen_batch.reset_context_requests(gen_reqs) - for mgr in gen_mgrs: - mgr.prepare_resources(gen_batch) + if use_v2: + gen_batch.context_requests_last_chunk = gen_reqs + for mgr in gen_mgrs: + for req in gen_reqs: + assert mgr.prepare_disagg_gen_init(req) + mgr.prepare_resources(gen_batch) + else: + gen_batch.reset_context_requests(gen_reqs) + for mgr in gen_mgrs: + mgr.prepare_resources(gen_batch) for req in ctx_reqs + gen_reqs: req.context_current_position = req.prompt_len @@ -438,10 +656,22 @@ def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): # -- 5. Ground truth -> shard -> write to ctx -- ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) - _write_ground_truth_to_ctx(ctx_mgrs, ctx_tp, ground_truth, ctx_rids) + _write_ground_truth_to_ctx( + ctx_mgrs, + ctx_tp, + ground_truth, + ctx_rids, + conv_state_layout, + ) # -- 6. Compute expected BEFORE transfer -- - expected = _compute_expected(ground_truth, gen_mgrs, gen_tp, gen_rids) + expected = _compute_expected( + ground_truth, + gen_mgrs, + gen_tp, + gen_rids, + conv_state_layout, + ) # -- 7. Transfer -- for rank in range(gen_tp): @@ -499,3 +729,23 @@ def test_mamba_transfer(ctx_tp, gen_tp): print(f"\nMamba transfer test: ctx_tp={ctx_tp} -> gen_tp={gen_tp}") run_mamba_transfer_test(ctx_tp, gen_tp) print("PASSED") + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,gen_tp,conv_state_layout", + [ + (2, 2, "x_b_c"), + (2, 4, "q_k_v"), + (4, 2, "x_b_c"), + ], + ids=["same_tp_xbc", "expand_tp_qkv", "contract_tp_xbc"], +) +def test_v2_mamba_transfer(ctx_tp, gen_tp, conv_state_layout): + """Transfer slot-major V2 Mamba states through Python/NIXL.""" + run_mamba_transfer_test( + ctx_tp, + gen_tp, + use_v2=True, + conv_state_layout=conv_state_layout, + )