[None][feat] DSv4: sparse MLA attention backend

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py

@@ -14,6 +14,21 @@
 # limitations under the License.
 
 from .cache_manager import DeepseekV4CacheManager
-from .deepseek_v4 import DeepseekV4AttentionType
+from .deepseek_v4 import (
+    DeepseekV4AttentionType,
+    DeepSeekV4MetadataParams,
+    DeepSeekV4Params,
+    DeepseekV4TrtllmAttention,
+    make_deepseek_v4_sparse_metadata_params,
+    make_deepseek_v4_sparse_params,
+)
 
-__all__ = ["DeepseekV4AttentionType", "DeepseekV4CacheManager"]
+__all__ = [
+    "DeepSeekV4MetadataParams",
+    "DeepSeekV4Params",
+    "DeepseekV4AttentionType",
+    "DeepseekV4CacheManager",
+    "DeepseekV4TrtllmAttention",
+    "make_deepseek_v4_sparse_metadata_params",
+    "make_deepseek_v4_sparse_params",
+]

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py

@@ -138,11 +138,21 @@ def __init__(
         mapping: Mapping,
         dtype: DataType = DataType.BF16,
         compressor_dtype: DataType = DataType.FLOAT,
-        sparse_attn_config: DeepSeekV4SparseAttentionConfig,
+        sparse_attn_config: Optional[DeepSeekV4SparseAttentionConfig] = None,
         max_input_len: Optional[int] = None,
         max_num_tokens: Optional[int] = None,
         **kwargs,
     ) -> None:
+        if sparse_attn_config is None:
+            sparse_attn_config = kwargs.pop("sparse_attention_config", None)
+        if sparse_attn_config is None and kwargs.get("model_config") is not None:
+            sparse_attn_config = kwargs["model_config"].sparse_attention_config
+        if sparse_attn_config is None:
+            raise ValueError(
+                "sparse_attn_config or sparse_attention_config is required "
+                "for DeepseekV4CacheManager"
+            )
+
         # DeepSeek-V4 specific attributes initialization
         assert kv_cache_type == CacheTypeCpp.SELFKONLY, "DeepSeek-V4 only supports SELFKONLY"
         assert num_kv_heads == 1, "DeepSeek-V4 only supports num_kv_heads == 1"