[None][feat] Add DSA Vanilla sparse attention

tensorrt_llm/_torch/attention/backends/sparse/dsa/__init__.py

@@ -23,6 +23,7 @@
 )
 from .metadata import DSAtrtllmAttentionMetadata, build_req_idx_per_token
 from .params import DSABackendForwardArgs, DSAMetadataParams, DSAParams
+from .vanilla_backend import DSAVanillaAttention, DSAVanillaIndexer
 
 __all__ = [
     "HAS_FAST_HADAMARD",
@@ -33,6 +34,8 @@
     "DSAParams",
     "DSATrtllmAttention",
     "DSAtrtllmAttentionMetadata",
+    "DSAVanillaAttention",
+    "DSAVanillaIndexer",
     "Indexer",
     "IndexerParams",
     "IndexerPrefillChunkMetadata",

tensorrt_llm/_torch/attention/backends/sparse/registry.py

@@ -75,6 +75,10 @@ def get_vanilla_sparse_attn_attention_backend(
 
     if sparse_params.algorithm == "rocket":
         return RocketVanillaAttention
+    elif sparse_params.algorithm == "dsa":
+        from .dsa import DSAVanillaAttention
+
+        return DSAVanillaAttention
     elif sparse_params.algorithm == "minimax_m3":
         return _resolve_minimax_m3_backend_cls(sparse_params)
     else:

tensorrt_llm/_torch/attention/backends/vanilla.py

@@ -804,6 +804,10 @@ def forward(self,
                 raise ValueError("Vanilla MLA requires a KV cache manager.")
             if forward_args.latent_cache is None:
                 raise ValueError("Vanilla MLA requires latent_cache.")
+            if self.sparse_params is not None:
+                raise NotImplementedError(
+                    f"{self.sparse_params.algorithm} requires its specialized "
+                    "Vanilla attention backend")
             if forward_args.attention_input_type == AttentionInputType.context_only:
                 assert k is not None and v is not None
                 return self._mla_forward_context(q, k, v, metadata,

tests/unittest/_torch/attention/backend_capability.py

@@ -20,7 +20,7 @@
 #   fp4_kv         - NVFP4 KV cache (Blackwell only)
 #   sliding_window - sliding-window attention via attention_window_size
 #   no_cache       - ragged/prefill forward with kv_cache_manager=None
-#   sparse         - sparse-attention forward plumbing (degenerate regime here)
+#   sparse         - sparse-attention forward plumbing
 #   mla            - multi-head latent attention
 #   cross_attn     - cross-attention (encoder-decoder)
 #   kv_layouts     - supported paged-cache block layouts ("NHD" / "HND")
@@ -60,7 +60,7 @@
         fp4_kv=False,
         sliding_window=True,
         no_cache=True,
-        sparse=False,
+        sparse=True,
         mla=True,
         cross_attn=True,
         kv_layouts=("NHD",),  # reads the NHD get_buffers view
@@ -86,7 +86,7 @@ def required_features(case) -> set:
         feats.add("sliding_window")
     if getattr(case, "cache", "paged") == "none":
         feats.add("no_cache")
-    if getattr(case, "sparse", "off") != "off":
+    if getattr(case, "sparse_attention_config", None) is not None:
         feats.add("sparse")
     if getattr(case, "is_mla", False):
         feats.add("mla")
@@ -114,6 +114,17 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
         if not caps.get(feat, False):
             return f"{backend} does not support feature '{feat}'"
 
+    sparse_config = getattr(case, "sparse_attention_config", None)
+    if sparse_config is not None:
+        algorithm = sparse_config.algorithm
+        if backend == "TRTLLM" and algorithm == "dsa":
+            # DSA selected-attention runs the trtllm-gen DynamicTokenSparse FMHA
+            # kernels, which only ship for Blackwell (sm_100+). On Hopper (sm90)
+            # MLA generation falls back to the dense FlashMLA kernel, which has no
+            # sparse path, so top-k selection is silently ignored.
+            if sm < 100:
+                return f"TRTLLM DSA requires sm>=100/Blackwell (have sm{sm})"
+
     # KV-cache block layout: a case may request a specific layout (NHD/HND). A
     # backend that cannot store the cache that way is skipped (e.g. TRTLLM is
     # head-major HND only). The Vanilla golden always runs in its native NHD and