[None][fix] Keep hybrid KV routing for Skip Softmax

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)