[https://nvbugs/6550132][fix] Preserve disaggregated Mamba cache routing

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -52,6 +52,7 @@
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
 from ..peft.lora.layer import LoraLayer, LoraModuleType
+from ..pyexecutor.mamba_cache_manager import mamba_manager_override_forces_v1
 from ..speculative import SpecMetadata
 from ..utils import AuxStreamType, EventType, Fp4QuantizedTensor
 from .modeling_deepseekv3 import DeepseekV3MTPHead
@@ -985,13 +986,15 @@ def load_weights(self,
     def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         """Model-specific defaults for NemotronH.
 
-        Uses KV cache manager V2 for the hybrid state layout. Block reuse
+        Uses KV cache manager V2 for the hybrid state layout, unless an env
+        override has pinned the mutually exclusive V1 Mamba route. Block reuse
         remains opt-in because it also requires a Mamba snapshot policy.
         """
+        propose_v2 = not mamba_manager_override_forces_v1(llm_args)
         return {
             "kv_cache_config": {
                 "enable_block_reuse": False,
-                "use_kv_cache_manager_v2": True,
+                "use_kv_cache_manager_v2": propose_v2,
             }
         }
 

tensorrt_llm/_torch/models/modeling_qwen3_next.py

@@ -40,6 +40,8 @@
 from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata
 from tensorrt_llm._torch.pyexecutor.config_utils import \
     get_qwen3_hybrid_layer_types
+from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import \
+    mamba_manager_override_forces_v1
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
@@ -1058,13 +1060,15 @@ def __init__(
     def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict:
         """Use V2 for the hybrid state layout.
 
-        Block reuse remains opt-in because it also requires a recurrent-state
-        snapshot policy.
+        Yields to an env override that pins the mutually exclusive V1 Mamba
+        route. Block reuse remains opt-in because it also requires a
+        recurrent-state snapshot policy.
         """
+        propose_v2 = not mamba_manager_override_forces_v1(llm_args)
         return {
             "kv_cache_config": {
                 "enable_block_reuse": False,
-                "use_kv_cache_manager_v2": True,
+                "use_kv_cache_manager_v2": propose_v2,
             }
         }
 

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -60,6 +60,7 @@
                                   CppMambaHybridCacheManager,
                                   MambaHybridCacheManagerV2,
                                   MixedMambaHybridCacheManager,
+                                  mamba_manager_preference,
                                   use_py_mamba_cache_manager)
 from .model_engine import PyTorchModelEngine
 from .py_executor import PyExecutor
@@ -122,7 +123,9 @@ def get_kv_cache_manager_cls(
     the Python transceiver with the NIXL backend. Unsupported V2 routes fail
     rather than falling back to a different manager.
 
-    Env-var overrides:
+    Env-var overrides. Both select a compatibility manager, so both are
+    consulted only in aggregated serving; in disaggregated serving the manager
+    follows the transceiver configuration instead.
       * ``TRTLLM_USE_PY_MAMBA=1``  — Mixed manager in aggregated serving.
       * ``TLLM_MAMBA_MANAGER_PREFERENCE`` — explicit manager preference.
     """
@@ -166,6 +169,8 @@ def get_kv_cache_manager_cls(
                         "KV cache manager V2 for hybrid Mamba disaggregated "
                         "serving requires transceiver_runtime='PYTHON' with "
                         "backend='NIXL'.")
+                # Falls through to the shared V2 route below, which still owns
+                # the remaining V2 compatibility checks.
             else:
                 if (kv_cache_config.enable_block_reuse and runtime == "PYTHON"):
                     raise ValueError(
@@ -194,9 +199,8 @@ def get_kv_cache_manager_cls(
             logger.info(
                 "Using MixedMambaHybridCacheManager for hybrid mamba model")
             return MixedMambaHybridCacheManager
-        env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None)
-        if env_override is not None:
-            env_override = env_override.upper()
+        env_override = mamba_manager_preference()
+        if env_override is not None and not is_disagg:
             if env_override == 'MIXED':
                 if use_v2:
                     raise ValueError(

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -26,7 +26,7 @@
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
-    from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
+    from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, TorchLlmArgs
 
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import (
     BlockReusePolicy, KVCacheManagerV2, Role)
@@ -221,6 +221,41 @@ def use_py_mamba_cache_manager() -> bool:
     return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1'
 
 
+def mamba_manager_preference() -> Optional[str]:
+    """Normalized value of the TLLM_MAMBA_MANAGER_PREFERENCE override.
+
+    Upper-cased so callers compare against 'CPP' / 'MIXED' consistently;
+    None when unset. Unrecognized values are returned as-is so the consumer
+    can warn about them.
+    """
+    preference = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None)
+    return preference.upper() if preference is not None else None
+
+
+def mamba_manager_override_forces_v1(llm_args: 'TorchLlmArgs') -> bool:
+    """Whether an env override pins the V1 Mamba route for aggregated serving.
+
+    ``get_kv_cache_manager_cls`` rejects these overrides alongside V2, so a
+    hybrid model must not *propose* ``use_kv_cache_manager_v2=True`` as a
+    default while one is active — an unset setting would resolve to V2 and
+    collide with the route the override already selected.
+
+    Both overrides are agg-mode-only, so this returns False once a disagg
+    transceiver backend is configured. ``llm_args`` may be ``None``, a
+    ``TorchLlmArgs``, or its nested-dict serialization.
+    """
+    if not (use_py_mamba_cache_manager()
+            or mamba_manager_preference() in ('CPP', 'MIXED')):
+        return False
+    transceiver_config = (llm_args.get('cache_transceiver_config')
+                          if isinstance(llm_args, dict) else getattr(
+                              llm_args, 'cache_transceiver_config', None))
+    backend = (transceiver_config.get('backend')
+               if isinstance(transceiver_config, dict) else getattr(
+                   transceiver_config, 'backend', None))
+    return backend is None
+
+
 class ReplayStateUpdateMetadata(NamedTuple):
     """Shared tensors and fixed sizes for replay state updates."""
     prev_num_accepted_tokens: torch.Tensor

tests/integration/test_lists/waives.txt

@@ -84,7 +84,6 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torc
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=False] SKIP (https://nvbugs/6427411)
 accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_fp4_tp2pp2[torch_compile=True-enable_gemm_allreduce_fusion=True] SKIP (https://nvbugs/6211441)
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] SKIP (https://nvbugs/6159132)
-accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-python_mamba_cache] SKIP (https://nvbugs/6550132)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2] SKIP (https://nvbugs/6427411)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] SKIP (https://nvbugs/6535767)
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_parallelism[ADP2_PP2] SKIP (https://nvbugs/6427411)

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -515,11 +515,52 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2)
     )
 
 
-def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
+@pytest.mark.parametrize("env_value", ["CPP", "MIXED"])
+def test_disagg_v2_route_ignores_agg_only_mamba_preference(monkeypatch, env_value):
+    """Both env overrides select a compatibility manager, so both are agg-only.
+
+    ``TRTLLM_USE_PY_MAMBA`` was already scoped that way, but
+    ``TLLM_MAMBA_MANAGER_PREFERENCE`` used to be consulted in disagg too, where
+    it could only ever reject the V2 route the transceiver had already selected.
+    """
+    monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)
+    monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", env_value)
+
+    assert (
+        get_kv_cache_manager_cls(
+            _hybrid_model_config(),
+            KvCacheConfig(enable_block_reuse=False, use_kv_cache_manager_v2=True),
+            is_disagg=True,
+            cache_transceiver_config=CacheTransceiverConfig(
+                backend="NIXL", transceiver_runtime="PYTHON"
+            ),
+        )
+        is MambaHybridCacheManagerV2
+    )
+
+
+def _hybrid_model_classes():
+    """The hybrid classes that define their own ``get_model_defaults``.
+
+    The Qwen3.5 text classes inherit ``Qwen3NextForCausalLM``'s verbatim, so
+    they add no coverage; ``Qwen3_5VLModel`` delegates to it explicitly.
+    """
     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
 
+    return (NemotronHForCausalLM, Qwen3NextForCausalLM, Qwen3_5VLModel)
+
+
+def _proposed_use_v2(llm_args):
+    """The ``use_kv_cache_manager_v2`` every hybrid class proposes, as a set."""
+    return {
+        model_cls.get_model_defaults(llm_args)["kv_cache_config"]["use_kv_cache_manager_v2"]
+        for model_cls in _hybrid_model_classes()
+    }
+
+
+def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
     for env_var in (
         "TRTLLM_USE_NIXL_KVCACHE",
         "TRTLLM_USE_UCX_KVCACHE",
@@ -528,7 +569,7 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
     ):
         monkeypatch.delenv(env_var, raising=False)
 
-    for model_cls in (NemotronHForCausalLM, Qwen3NextForCausalLM, Qwen3_5VLModel):
+    for model_cls in _hybrid_model_classes():
         llm_args = TorchLlmArgs(
             model="/tmp/dummy_model",
             cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"),
@@ -542,6 +583,104 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
         assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON"
 
 
+@pytest.mark.parametrize(
+    ("env_name", "env_value"),
+    [
+        ("TRTLLM_USE_PY_MAMBA", "1"),
+        ("TLLM_MAMBA_MANAGER_PREFERENCE", "MIXED"),
+        ("TLLM_MAMBA_MANAGER_PREFERENCE", "CPP"),
+    ],
+)
+def test_hybrid_defaults_do_not_propose_v2_under_v1_override(monkeypatch, env_name, env_value):
+    """A V1-route override must suppress the V2 default, not collide with it."""
+    monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)
+    monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
+    monkeypatch.setenv(env_name, env_value)
+
+    # Block reuse stays off: the Mixed manager the override selects rejects it.
+    llm_args = TorchLlmArgs(
+        model="/tmp/dummy_model",
+        kv_cache_config=KvCacheConfig(enable_block_reuse=False),
+    )
+    assert _proposed_use_v2(llm_args) == {False}
+
+    # All three propose the same dict (asserted above), so resolve one of them
+    # the rest of the way through to the factory.
+    representative_cls, *_ = _hybrid_model_classes()
+    model_defaults = representative_cls.get_model_defaults(llm_args)
+    apply_model_defaults_to_llm_args(llm_args, model_defaults)
+    assert (
+        _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto")
+        is False
+    )
+    # The factory previously raised ValueError for this pair.
+    assert get_kv_cache_manager_cls(_hybrid_model_config(), llm_args.kv_cache_config) in (
+        MixedMambaHybridCacheManager,
+        CppMambaHybridCacheManager,
+    )
+
+
+def test_v1_override_is_ignored_for_disagg(monkeypatch):
+    """The override is agg-mode-only, so disagg keeps proposing V2."""
+    monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
+    monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1")
+
+    disagg_args = TorchLlmArgs(
+        model="/tmp/dummy_model",
+        cache_transceiver_config=CacheTransceiverConfig(backend="NIXL"),
+    )
+    assert _proposed_use_v2(disagg_args) == {True}
+
+
+def test_v1_override_detects_disagg_in_model_dump(monkeypatch):
+    """A serialized transceiver config must still identify disagg serving."""
+    monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1")
+    monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
+    disagg_args = TorchLlmArgs(
+        model="/tmp/dummy_model",
+        cache_transceiver_config=CacheTransceiverConfig(backend="NIXL"),
+    ).model_dump()
+
+    assert isinstance(disagg_args["cache_transceiver_config"], dict)
+    assert _proposed_use_v2(disagg_args) == {True}
+
+
+def test_unrecognized_preference_is_not_a_v1_override(monkeypatch):
+    """Only 'CPP' and 'MIXED' select a V1 manager; anything else is ignored."""
+    monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)
+    monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "V2")
+
+    assert _proposed_use_v2(TorchLlmArgs(model="/tmp/dummy_model")) == {True}
+
+
+@pytest.mark.parametrize("loose_llm_args", [None, {}], ids=["none", "dict"])
+def test_hybrid_defaults_accept_loose_llm_args(monkeypatch, loose_llm_args):
+    """``get_model_defaults`` is called with None and with a plain dict in-tree.
+
+    The perf-sanity transceiver precheck passes ``None`` and
+    ``tests/unittest/llmapi/test_config_database.py`` passes a ``model_dump()``
+    dict, so the override predicate must not require a real ``TorchLlmArgs``.
+    """
+    monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)
+    monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
+    assert _proposed_use_v2(loose_llm_args) == {True}
+
+    monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1")
+    assert _proposed_use_v2(loose_llm_args) == {False}
+
+
+def test_explicit_v2_still_conflicts_with_py_mamba_override(monkeypatch):
+    """An explicit True is a real user contradiction and must keep raising."""
+    monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
+    monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1")
+
+    with pytest.raises(ValueError, match="conflicts with explicit"):
+        get_kv_cache_manager_cls(
+            _hybrid_model_config(),
+            KvCacheConfig(enable_block_reuse=False, use_kv_cache_manager_v2=True),
+        )
+
+
 def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank():
     manager = object.__new__(MambaHybridCacheManagerV2)
     manager.local_num_mamba_layers = 0