From 88aa07951a44148b43dda620922e015815bfd4b0 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 16 Jun 2026 14:42:42 +0800 Subject: [PATCH] [None][fix] Keep hybrid KV routing for Skip Softmax Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 5ae9ce95fbd0..f5cc763bec3b 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -88,16 +88,26 @@ def get_kv_cache_manager_cls( * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager with PythonMambaCacheManager. """ config = model_config.pretrained_config - sparse_attention_config = model_config.sparse_attention_config - if sparse_attention_config is not None: - return get_sparse_attn_kv_cache_manager(sparse_attention_config) - elif is_hybrid_linear(config): + sparse_attn_config = model_config.sparse_attention_config + sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) + if is_hybrid_linear(config): # Degenerate case: model is flagged as hybrid but the config has zero - # mamba layers. Fall through to the standard non-hybrid manager. + # mamba layers. Fall through to the standard non-hybrid routes. if model_config.get_num_mamba_layers() == 0: logger.info("Hybrid linear model has 0 mamba layers; using " - "KVCacheManager without mamba caching") + "KV cache manager without mamba caching") + if sparse_attn_config is not None: + return get_sparse_attn_kv_cache_manager(sparse_attn_config) return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) + + if (sparse_attn_config is not None + and sparse_attn_algorithm != "skip_softmax"): + raise ValueError( + f"Sparse attention algorithm {sparse_attn_algorithm!r} is not " + "supported with hybrid Mamba / linear-attention models.") + + # Skip Softmax only changes attention kernels. Hybrid models still + # need a Mamba-capable cache manager for recurrent state. if kv_cache_config.enable_block_reuse: return CppMambaHybridCacheManager if use_cpp_mamba_cache_manager() or use_py_mamba_cache_manager(): @@ -128,6 +138,8 @@ def get_kv_cache_manager_cls( f"Expected 'CPP' or 'MIXED'. Using default {default_cls.__name__}." ) return default_cls + elif sparse_attn_config is not None: + return get_sparse_attn_kv_cache_manager(sparse_attn_config) else: return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config)