From e1cec33fb7c986b26f4cf0f46f00e5420c824195 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:49:23 -0700 Subject: [PATCH 1/3] [nvbugs/6550132][fix] Don't propose KV cache manager V2 under a V1 Mamba override Hybrid Mamba models default `use_kv_cache_manager_v2` to True for the hybrid state layout. `TRTLLM_USE_PY_MAMBA=1` and `TLLM_MAMBA_MANAGER_PREFERENCE` in (CPP, MIXED) select a mutually exclusive V1-route manager, and `get_kv_cache_manager_cls` rejects that pair. A user who set only the env var left the setting at "auto", which resolved to the model default of True, so the factory raised about an "explicit" V2 the user never selected -- inside the MPI worker, surfacing as "RuntimeError: Executor worker returned error". Suppress the default at the producer instead of correcting it afterwards: the hybrid `get_model_defaults` now propose V2 only when no V1-route override is active, so the contradicted value is never proposed and nothing is downgraded. The predicate lives next to `use_py_mamba_cache_manager()`, which already owns these env vars, and returns False for disaggregated serving because both overrides are aggregated-mode-only. An explicit `use_kv_cache_manager_v2=True` still raises -- that is a genuine user contradiction, which finally makes the "explicit" wording truthful. `get_model_defaults` is called with None (perf-sanity transceiver precheck) and with a plain dict (config-database test) as well as with `TorchLlmArgs`, so the predicate reads `llm_args` via `getattr`; both loose forms now get the correct answer rather than the precheck's silent "assuming V1" fallback, and are covered by a new regression test. Removes the now-stale waiver for the aggregated H100 parametrization. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/models/modeling_nemotron_h.py | 7 +- .../_torch/models/modeling_qwen3_next.py | 10 +- .../_torch/pyexecutor/mamba_cache_manager.py | 24 ++++- tests/integration/test_lists/waives.txt | 1 - .../executor/test_mamba_cache_manager.py | 100 +++++++++++++++++- 5 files changed, 133 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 8516a3dbdb1e..c527a4890db9 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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, } } diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 15047b48a4d5..6b3a9f55f61c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/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, } } diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index a670074dd23b..dd5b4e484b58 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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,28 @@ def use_py_mamba_cache_manager() -> bool: return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' +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. + + ``llm_args`` is read defensively: :meth:`get_model_defaults` is also called + with ``None`` (perf-sanity transceiver precheck) and with a plain dict + (``tests/unittest/llmapi/test_config_database.py``). + """ + preference = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', '').upper() + if not (use_py_mamba_cache_manager() or preference in ('CPP', 'MIXED')): + return False + + # Both overrides are agg-mode-only; in disagg the manager follows the + # transceiver configuration instead. + transceiver_config = getattr(llm_args, 'cache_transceiver_config', None) + return getattr(transceiver_config, 'backend', None) is None + + class ReplayStateUpdateMetadata(NamedTuple): """Shared tensors and fixed sizes for replay state updates.""" prev_num_accepted_tokens: torch.Tensor diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 647729a6c0da..e9ffeba62b9f 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index dda9c28b93c9..2b3bb4222557 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -515,11 +515,20 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2) ) -def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): +def _hybrid_model_classes(): + """The hybrid classes that each define their own ``get_model_defaults``. + + ``Qwen3_5VLModel`` delegates to ``Qwen3NextForCausalLM``; the Qwen3.5 text + classes inherit it verbatim, so they add no coverage here. + """ 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 test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): for env_var in ( "TRTLLM_USE_NIXL_KVCACHE", "TRTLLM_USE_UCX_KVCACHE", @@ -528,7 +537,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 +551,93 @@ 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) + + for model_cls in _hybrid_model_classes(): + # 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), + ) + model_defaults = model_cls.get_model_defaults(llm_args) + assert model_defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is False + + 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_and_unknown_values(monkeypatch): + """The override is agg-mode-only, and an unrecognized value is not an override.""" + 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"), + ) + for model_cls in _hybrid_model_classes(): + defaults = model_cls.get_model_defaults(disagg_args) + assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "V2") + llm_args = TorchLlmArgs(model="/tmp/dummy_model") + for model_cls in _hybrid_model_classes(): + defaults = model_cls.get_model_defaults(llm_args) + assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is 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) + for model_cls in _hybrid_model_classes(): + defaults = model_cls.get_model_defaults(loose_llm_args) + assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + + monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1") + for model_cls in _hybrid_model_classes(): + defaults = model_cls.get_model_defaults(loose_llm_args) + assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is 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 From 2968450263e0d812cadcf622f162fc309a16b0f1 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 3 Aug 2026 12:38:20 -0700 Subject: [PATCH 2/3] [nvbugs/6550132][fix] Scope the Mamba manager preference override to aggregated serving TLLM_MAMBA_MANAGER_PREFERENCE selects a compatibility (V1) Mamba manager, so like TRTLLM_USE_PY_MAMBA it only applies in aggregated serving -- in disagg the manager follows the transceiver configuration. It was missing that guard, and get_kv_cache_manager_cls's is_disagg block does not return on a supported V2 route, so a hybrid model whose auto-resolved default is V2 fell through to the preference branch and raised about an 'explicit' V2 the user never set. Give the preference branch the same 'and not is_disagg' guard the sibling env var already has. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 12 ++- .../_torch/pyexecutor/mamba_cache_manager.py | 25 +++-- .../executor/test_mamba_cache_manager.py | 102 +++++++++++------- 3 files changed, 91 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7b4ffa4d4b2f..e27d0cd2c69d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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( diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index dd5b4e484b58..0c98be1bc8bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -221,6 +221,17 @@ 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. @@ -229,16 +240,14 @@ def mamba_manager_override_forces_v1(llm_args: 'TorchLlmArgs') -> bool: default while one is active — an unset setting would resolve to V2 and collide with the route the override already selected. - ``llm_args`` is read defensively: :meth:`get_model_defaults` is also called - with ``None`` (perf-sanity transceiver precheck) and with a plain dict - (``tests/unittest/llmapi/test_config_database.py``). + Both overrides are agg-mode-only, so this returns False once a disagg + transceiver backend is configured. ``llm_args`` is read via ``getattr`` + because :meth:`get_model_defaults` is also called with ``None`` and with a + plain dict in-tree. """ - preference = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', '').upper() - if not (use_py_mamba_cache_manager() or preference in ('CPP', 'MIXED')): + if not (use_py_mamba_cache_manager() + or mamba_manager_preference() in ('CPP', 'MIXED')): return False - - # Both overrides are agg-mode-only; in disagg the manager follows the - # transceiver configuration instead. transceiver_config = getattr(llm_args, 'cache_transceiver_config', None) return getattr(transceiver_config, 'backend', None) is None diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 2b3bb4222557..ffe8ce70c478 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -515,11 +515,35 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2) ) +@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 each define their own ``get_model_defaults``. + """The hybrid classes that define their own ``get_model_defaults``. - ``Qwen3_5VLModel`` delegates to ``Qwen3NextForCausalLM``; the Qwen3.5 text - classes inherit it verbatim, so they add no coverage here. + 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 @@ -528,6 +552,14 @@ def _hybrid_model_classes(): 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", @@ -565,45 +597,47 @@ def test_hybrid_defaults_do_not_propose_v2_under_v1_override(monkeypatch, env_na monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) monkeypatch.setenv(env_name, env_value) - for model_cls in _hybrid_model_classes(): - # 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), - ) - model_defaults = model_cls.get_model_defaults(llm_args) - assert model_defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is False + # 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} - 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, - ) + # 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_and_unknown_values(monkeypatch): - """The override is agg-mode-only, and an unrecognized value is not an override.""" +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"), ) - for model_cls in _hybrid_model_classes(): - defaults = model_cls.get_model_defaults(disagg_args) - assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + 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") - llm_args = TorchLlmArgs(model="/tmp/dummy_model") - for model_cls in _hybrid_model_classes(): - defaults = model_cls.get_model_defaults(llm_args) - assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + + assert _proposed_use_v2(TorchLlmArgs(model="/tmp/dummy_model")) == {True} @pytest.mark.parametrize("loose_llm_args", [None, {}], ids=["none", "dict"]) @@ -616,14 +650,10 @@ def test_hybrid_defaults_accept_loose_llm_args(monkeypatch, loose_llm_args): """ monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - for model_cls in _hybrid_model_classes(): - defaults = model_cls.get_model_defaults(loose_llm_args) - assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + assert _proposed_use_v2(loose_llm_args) == {True} monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1") - for model_cls in _hybrid_model_classes(): - defaults = model_cls.get_model_defaults(loose_llm_args) - assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is False + assert _proposed_use_v2(loose_llm_args) == {False} def test_explicit_v2_still_conflicts_with_py_mamba_override(monkeypatch): From 1f92b32e8ae9e0e533f5d63b259cf8bf49c7d7cf Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:06:57 +0800 Subject: [PATCH 3/3] [nvbugs/6550132][fix] Handle serialized disagg config in Mamba routing (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/pyexecutor/mamba_cache_manager.py | 14 +++++++++----- .../_torch/executor/test_mamba_cache_manager.py | 13 +++++++++++++ 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 0c98be1bc8bb..0ddf49188a99 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -241,15 +241,19 @@ def mamba_manager_override_forces_v1(llm_args: 'TorchLlmArgs') -> bool: 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`` is read via ``getattr`` - because :meth:`get_model_defaults` is also called with ``None`` and with a - plain dict in-tree. + 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 = getattr(llm_args, 'cache_transceiver_config', None) - return getattr(transceiver_config, 'backend', None) is None + 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): diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index ffe8ce70c478..79894954de09 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -632,6 +632,19 @@ def test_v1_override_is_ignored_for_disagg(monkeypatch): 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)