From ba18408acbd7a5ff0d0a8d078385dc13065d1bfe Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Tue, 19 May 2026 23:41:44 -0700 Subject: [PATCH] [None][feature] Add env variables to help debugging mamba modules Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 8 +++- tensorrt_llm/_torch/pyexecutor/_util.py | 38 ++++++++++++++----- .../_torch/pyexecutor/mamba_cache_manager.py | 27 ++++++++++++- .../_torch/pyexecutor/py_executor_creator.py | 11 ++++++ 4 files changed, 72 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index cba1fd0996de..1f8a2af500a0 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 53238799fd5d..6506e9c445f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index da682bebc858..ff942210b303 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index fa33e0ce8643..de2fac2b6330 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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. "