[None][feature] Add env variables to help debugging mamba modules.

tensorrt_llm/_torch/model_config.py

@@ -45,15 +45,19 @@ def _unified_kv_pool_includes_mamba(
 
       * disaggregated serving forces the C++ mamba manager
         (``TRTLLM_USE_CPP_MAMBA=1`` enables the same path locally), or
+      * ``TRTLLM_USE_PY_MAMBA=1`` forces the Python mamba manager locally
+        (agg-mode override), or
       * one-model speculative decoding splits mamba and attention into
         separate caches.
 
     Single source of truth for the binding-side layer-counting decision; do
     not duplicate the predicate at call sites.
     """
-    use_disagg = is_disagg or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1'
+    use_split_pool = is_disagg \
+        or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \
+        or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1'
     use_spec = spec_config is not None
-    return not (use_disagg or use_spec)
+    return not (use_split_pool or use_spec)
 
 
 @contextlib.contextmanager

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -41,7 +41,8 @@
 from .mamba_cache_manager import (BaseMambaCacheManager,
                                   CppMambaHybridCacheManager,
                                   MixedMambaHybridCacheManager,
-                                  use_cpp_mamba_cache_manager)
+                                  use_cpp_mamba_cache_manager,
+                                  use_py_mamba_cache_manager)
 from .model_engine import PyTorchModelEngine
 from .py_executor import PyExecutor
 from .resource_manager import (KVCacheManager, KVCacheManagerV2,
@@ -75,9 +76,14 @@ def get_kv_cache_manager_cls(model_config: ModelConfig,
     """Resolve the concrete KV cache manager class for ``model_config``.
 
     For hybrid mamba models the choice between ``Mixed`` (separate pools,
-    needed for disagg / TRTLLM_USE_CPP_MAMBA) and ``Cpp`` (unified pool with
-    block reuse) is made here. Callers that don't care about disagg can omit
-    ``is_disagg`` and get the unified-pool default.
+    needed for disagg / TRTLLM_USE_CPP_MAMBA / TRTLLM_USE_PY_MAMBA) and
+    ``Cpp`` (unified pool with block reuse) is made here. Callers that don't
+    care about disagg can omit ``is_disagg`` and get the unified-pool default.
+
+    Env-var overrides (agg mode only — disagg picks its inner impl via
+    ``cache_transceiver_config.transceiver_runtime``):
+      * ``TRTLLM_USE_CPP_MAMBA=1`` — Mixed manager with CppMambaCacheManager.
+      * ``TRTLLM_USE_PY_MAMBA=1``  — Mixed manager with PythonMambaCacheManager.
     """
     config = model_config.pretrained_config
     sparse_attn_config = model_config.sparse_attention_config
@@ -90,7 +96,8 @@ def get_kv_cache_manager_cls(model_config: ModelConfig,
             logger.info("Hybrid linear model has 0 mamba layers; using "
                         "KVCacheManager without mamba caching")
             return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config)
-        if is_disagg or use_cpp_mamba_cache_manager():
+        if is_disagg or use_cpp_mamba_cache_manager(
+        ) or use_py_mamba_cache_manager():
             return MixedMambaHybridCacheManager
         return CppMambaHybridCacheManager
     else:
@@ -236,14 +243,16 @@ def _get_model_kv_cache_manager_cls(self, model_engine: PyTorchModelEngine):
                     "event buffer max size > 0, or cache transceiver. Falling back to KVCacheManager."
                 )
                 cls = KVCacheManager
-        # The V1-route hybrid mamba managers (disagg via TRTLLM_USE_CPP_MAMBA,
-        # or one-model speculative decoding) keep mamba state in a separate
-        # cache that doesn't honor block reuse. Warn at the routing site so
-        # users see the warning where the decision is actually made.
+        # The V1-route hybrid mamba managers (disagg, TRTLLM_USE_CPP_MAMBA,
+        # TRTLLM_USE_PY_MAMBA, or one-model speculative decoding) keep mamba
+        # state in a separate cache that doesn't honor block reuse. Warn at
+        # the routing site so users see the warning where the decision is
+        # actually made.
         if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \
                 and self._kv_cache_config.enable_block_reuse:
             uses_v1_mamba_route = self._is_disagg \
                 or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \
+                or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \
                 or self._speculative_config is not None
             if uses_v1_mamba_route:
                 logger.warning(
@@ -1265,6 +1274,17 @@ def _create_kv_cache_manager(
                         "using legacy MTP path for stochastic rounding support")
             use_replay = False
 
+        # Use replay algorithm for mamba (default is on).
+        enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY',
+                                                '1') == '0'
+        if enforce_disable_replay:
+            logger.info(
+                "Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0")
+            use_replay = False
+        else:
+            logger.info(
+                "Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1")
+
         kv_cache_manager = kv_cache_manager_cls(
             # mamba cache parameters
             mamba_params.state_size,

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -60,7 +60,32 @@ def use_cpp_mamba_cache_manager() -> bool:
     Returns True if TRTLLM_USE_CPP_MAMBA='1' is set, False otherwise.
     By default, PythonMambaCacheManager is used.
     """
-    return os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1'
+    cpp = os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1'
+    py = os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1'
+    if cpp and py:
+        raise ValueError(
+            "TRTLLM_USE_CPP_MAMBA=1 and TRTLLM_USE_PY_MAMBA=1 are mutually "
+            "exclusive; unset one of them.")
+    return cpp
+
+
+def use_py_mamba_cache_manager() -> bool:
+    """Check if PythonMambaCacheManager should be forced (agg mode override).
+
+    Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise.
+
+    Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager
+    with PythonMambaCacheManager inside instead of the default unified-pool
+    CppMambaHybridCacheManager. Disagg mode is unaffected — it already picks
+    PythonMambaCacheManager when transceiver_runtime='PYTHON'.
+    """
+    cpp = os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1'
+    py = os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1'
+    if cpp and py:
+        raise ValueError(
+            "TRTLLM_USE_CPP_MAMBA=1 and TRTLLM_USE_PY_MAMBA=1 are mutually "
+            "exclusive; unset one of them.")
+    return py
 
 
 class BaseMambaCacheManager(ABC):

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -829,6 +829,17 @@ def drafting_loop_wrapper(model):
         is_hybrid = is_hybrid_linear(config)
 
         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 CppMambaCacheManager. Clear the var here so the
+            # mutual-exclusion check in use_cpp_mamba_cache_manager() does
+            # not fire after we force TRTLLM_USE_CPP_MAMBA=1 below.
+            if os.environ.pop("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.")
             if cache_transceiver_config.transceiver_runtime != "PYTHON" or os.environ.get(
                     "TRTLLM_USE_CPP_MAMBA") == "1":
                 logger.info("Disaggregated serving with hybrid model detected. "