[None][refactor] Unify sparse attention framework with clean backend interfaces

docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md

@@ -220,7 +220,7 @@ Figure 4 illustrates the prediction implementation within TensorRT LLM. To suppo
 
 **Auxiliary memory management.** Managing the paged KT cache presented another challenge. `RocketKVCacheManager` inherits from `KVCacheManager` and extends it with a dedicated `BlockManager` for the auxiliary KT cache at the Python level. The main KV cache and the KT cache share block IDs for each request, so that the lifecycle of KT cache blocks is automatically tied to the corresponding KV cache blocks. The `BlockManager` handles slot allocation and deallocation for the KT cache independently, while `RocketKVCacheManager` overrides methods such as `get_cache_bytes_per_token` and `prepare_resources` to ensure that memory sizing accounts for the extra KT cache footprint and that the correct KT cache pointers are passed to prediction kernels at each step. This design keeps the integration lightweight and easy to iterate on, though it inherits the limitations of Python-level management—namely, no automatic support for KV cache reuse or disaggregated serving. 
 
-The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/rocket.py`.
+The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/rocket/`.
 
 ### DeepSeek Sparse Attention (DSA)
 
@@ -261,7 +261,7 @@ As with RocketKV, a dedicated metadata class `DSATrtllmAttentionMetadata` is def
 
 **Auxiliary memory management.** DSA requires an auxiliary **indexer K cache** to store the low-rank K projections for reuse across decoding steps. `DSAKVCacheManager` inherits from `KVCacheManager`, but unlike RocketKV's Python-level KT cache management, DSA's indexer K cache is integrated directly into the C++ `KVCacheManager`. This design enables compatibility with advanced features such as KV cache reuse, chunked prefill, and disaggregated serving—features that would be difficult to support with a Python-level manager. 
 
-The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`.
+The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/dsa/`.
 
 For a comprehensive description of DSA kernel optimizations, precision strategies, feature support (MTP, disaggregated serving, Wide-EP), and benchmark results, please refer to the dedicated blog post: [Optimizing DeepSeek-V3.2 on NVIDIA Blackwell GPUs](blog15_Optimizing_DeepSeek_V32_on_NVIDIA_Blackwell_GPUs.md).
 

docs/source/developer-guide/sparse-attention-development-guide.md

@@ -137,19 +137,26 @@ KV eviction are tracked as future work.
 subclasses override:
 
 ```python
-sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict(q, k, metadata, **kwargs)
-sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict(q, k, metadata, **kwargs)
+sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict(q, k, metadata, forward_args)
+sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict(q, k, metadata, forward_args)
 ```
 
+`hooks.py` writes these results to `SparseRuntimeParams`. SkipSoftmax writes
+its thresholds to the same runtime interface consumed by `AttentionOp`.
+
 Different KV heads are allowed to emit different sparse index sets; Q
 heads that map to the same KV head share the KV head's sparse pattern.
 
 Algorithm implementations live under
 `tensorrt_llm/_torch/attention_backend/sparse/`:
 
-- `rocket.py`, `dsa.py` — concrete algorithms.
-- `kernel.py` — custom Triton kernels (importance scoring, Top-K).
-- `utils.py` — dispatch helpers.
+- `rocket/` — RocketKV backend, metadata, cache manager, parameters, and kernels.
+- `dsa/` — DSA backend, indexer, metadata, cache manager, parameters, custom ops, and kernels.
+- `deepseek_v4/` — DeepSeek-V4 backend, indexer, metadata, cache manager,
+  parameters, module hooks, and index conversion kernels.
+- `skip_softmax/` — SkipSoftmax parameter parsing and runtime scheduler.
+- `hooks.py` — module hooks and common backend prediction orchestration.
+- `registry.py` — backend, metadata, and cache-manager dispatch helpers.
 
 ### AttentionOp behavior
 
@@ -231,7 +238,7 @@ Create a new backend class inheriting from `TrtllmAttention` (or
 `tensorrt_llm/_torch/attention_backend/sparse/`. Override one or both
 prediction methods.
 
-**`sparse_kv_predict(self, q, k, metadata, **kwargs)`**
+**`sparse_kv_predict(self, q, k, metadata, forward_args)`**
 
 - **Behavior**: return the indices of tokens to retain in the KV cache.
 - **Outputs**:
@@ -245,7 +252,7 @@ prediction methods.
   in-place gather (`updateSparseKvCacheAfterFmha`) is safe. The sort
   cost buys compatibility with chunked prefill and similar features.
 
-**`sparse_attn_predict(self, q, k, metadata, **kwargs)`**
+**`sparse_attn_predict(self, q, k, metadata, forward_args)`**
 
 - **Behavior**: return the sparse indices used by the generation-phase
   attention computation.
@@ -285,7 +292,7 @@ If the algorithm needs extra tensors beyond the main KV cache:
 ### 4. Registration and dispatch
 
 - Register the new config + backend in
-  `tensorrt_llm/_torch/attention_backend/sparse/utils.py` and
+  `tensorrt_llm/_torch/attention_backend/sparse/registry.py` and
   `tensorrt_llm/_torch/pyexecutor/_util.py` so the runtime routes
   requests to your backend when the config is present.
 - If your algorithm exposes new C++ parameters, plumb them through
@@ -318,5 +325,3 @@ for the kernel-side specifics.
   eviction as a compromise that keeps KV cache flexibility manageable.
 - **Unified auxiliary memory management** — let custom auxiliary pools
   inherit KV-cache features (reuse, offloading) by default.
-- **Code refactoring** — as more algorithms land, unify the
-  framework-level scaffolding for maintainability.

docs/source/torch/adding_custom_kernels.md

@@ -290,13 +290,13 @@ The adapter and env helpers live in [`tensorrt_llm/_torch/custom_ops/cutedsl_mat
 
 - **The kernel.** [`cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h) declares `invokeIndexerKCacheScatter(...)`; [`cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu) defines the `__global__` kernel and the host launcher. Picked up automatically by the `*.cu` glob in `cpp/tensorrt_llm/kernels/CMakeLists.txt`.  
 - **The binding.** [`cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp) validates inputs, fetches the stream, and calls the launcher; registers the op via `TORCH_LIBRARY_FRAGMENT(trtllm, m)` and `TORCH_LIBRARY_IMPL(trtllm, CUDA, m)`. A one-line addition to [`cpp/tensorrt_llm/thop/CMakeLists.txt`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/CMakeLists.txt) wires it into `th_common` (this directory is not globbed).  
-- **The integration.** In [`tensorrt_llm/_torch/attention_backend/sparse/dsa.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/attention_backend/sparse/dsa.py), `_update_k_cache` calls `torch.ops.trtllm.indexer_k_cache_scatter_op(...)` in place of the prior Python scatter loop.  
+- **The integration.** In [`tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py), `_update_k_cache` calls `torch.ops.trtllm.indexer_k_cache_scatter_op(...)` in place of the prior Python scatter loop.
 - **The tests.** [`tests/unittest/_torch/attention/sparse/test_dsa_indexer.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py) — `test_indexer_k_cache_scatter_custom_op` runs the CUDA op against a PyTorch reference on two layers of the same cache pool and asserts the outputs match.
 
 The full call chain is three lines:
 
 ```py
-# tensorrt_llm/_torch/attention_backend/sparse/dsa.py
+# tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
 torch.ops.trtllm.indexer_k_cache_scatter_op(
     k_fp8, k_scale, k_cache,
     metadata.slot_mapping_fp8, metadata.slot_mapping_scale,

tensorrt_llm/_torch/attention_backend/fmha/fallback.py

@@ -32,7 +32,7 @@
 # transitively reads; entries here are exempt.
 _THOP_EXCLUDED_FIELDS: frozenset = frozenset(
     {
-        "topk_indices",  # DSA-only
+        "sparse_backend_args",  # consumed by sparse prediction before the attention op
         "attention_mask_data",  # custom-mask code path
         "out_scale_sf",  # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path
         "skip_mla_rope_generation",  # handled in ``TrtllmAttention.forward`` for the test-only MLA path
@@ -60,7 +60,7 @@ def forward(
         attn = self.attn
 
         # Every kwarg sources from ``attn`` / ``metadata`` / ``forward_args``
-        # (with ``forward_args.sparse_prediction`` for sparse-attn inputs),
+        # (with ``forward_args.sparse_runtime_params`` for sparse inputs),
         # or a literal allowlisted in ``_THOP_LITERALS``.
         # ``test_attention_op_sync.py`` enforces this statically.
         thop.attention(
@@ -141,8 +141,6 @@ def forward(
             cross_kv=forward_args.cross_kv,
             relative_attention_bias=forward_args.relative_attention_bias,
             relative_attention_max_distance=forward_args.relative_attention_max_distance,
-            skip_softmax_threshold_scale_factor_prefill=forward_args.skip_softmax_kernel_params.threshold_scale_factor_prefill,
-            skip_softmax_threshold_scale_factor_decode=forward_args.skip_softmax_kernel_params.threshold_scale_factor_decode,
             # --- Module config (TrtllmAttention) ---
             rotary_inv_freq=attn.rotary_inv_freq,
             rotary_cos_sin=attn.rotary_cos_sin,
@@ -171,12 +169,22 @@ def forward(
             rope_append=attn.rope_append,
             attention_chunk_size=attn.attention_chunk_size,
             skip_softmax_stat=attn.skip_softmax_stat,
-            # --- Sparse-specific (AttentionForwardArgs.sparse_prediction) ---
-            sparse_kv_indices=forward_args.sparse_prediction.sparse_kv_indices,
-            sparse_kv_offsets=forward_args.sparse_prediction.sparse_kv_offsets,
-            sparse_attn_indices=forward_args.sparse_prediction.sparse_attn_indices,
-            sparse_attn_offsets=forward_args.sparse_prediction.sparse_attn_offsets,
-            sparse_attn_indices_block_size=forward_args.sparse_prediction.sparse_attn_indices_block_size,
-            sparse_mla_topk_lens=forward_args.sparse_prediction.sparse_mla_topk_lens,
-            compressed_kv_cache_pool_ptr=forward_args.sparse_prediction.compressed_kv_cache_pool_ptr,
+            # --- Sparse runtime parameters ---
+            sparse_kv_indices=forward_args.sparse_runtime_params.sparse_kv_indices,
+            sparse_kv_offsets=forward_args.sparse_runtime_params.sparse_kv_offsets,
+            sparse_attn_indices=forward_args.sparse_runtime_params.sparse_attn_indices,
+            sparse_attn_offsets=forward_args.sparse_runtime_params.sparse_attn_offsets,
+            sparse_attn_indices_block_size=(
+                forward_args.sparse_runtime_params.sparse_attn_indices_block_size
+            ),
+            sparse_mla_topk_lens=forward_args.sparse_runtime_params.sparse_mla_topk_lens,
+            compressed_kv_cache_pool_ptr=(
+                forward_args.sparse_runtime_params.compressed_kv_cache_pool_ptr
+            ),
+            skip_softmax_threshold_scale_factor_prefill=(
+                forward_args.sparse_runtime_params.threshold_scale_factor_prefill
+            ),
+            skip_softmax_threshold_scale_factor_decode=(
+                forward_args.sparse_runtime_params.threshold_scale_factor_decode
+            ),
         )

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

@@ -620,8 +620,8 @@ def _is_supported_with_reason(
             return False, "sage attention."
         if meta.helix_position_offsets is not None:
             return False, "helix parallelism."
-        sparse_kv_indices = fwd.sparse_prediction.sparse_kv_indices
-        sparse_attn_indices = fwd.sparse_prediction.sparse_attn_indices
+        sparse_kv_indices = fwd.sparse_runtime_params.sparse_kv_indices
+        sparse_attn_indices = fwd.sparse_runtime_params.sparse_attn_indices
         if (
             (sparse_kv_indices is not None and sparse_kv_indices.numel() > 0)
             or (sparse_attn_indices is not None and sparse_attn_indices.numel() > 0)

tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py

@@ -5,7 +5,7 @@
 MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and
 participates in the standard TrtllmAttention.forward dispatch loop. The
 owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the
-per-query KV blocks and publishes them on forward_args.sparse_prediction;
+per-query KV blocks and publishes them on forward_args.sparse_runtime_params;
 this class attends over them.
 """
 
@@ -173,7 +173,7 @@ class MsaSparseGqaFmha(Fmha):
     """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel.
 
     Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's
-    selected KV block indices on forward_args.sparse_prediction.sparse_attn_indices
+    selected KV block indices on forward_args.sparse_runtime_params.sparse_attn_indices
     and attend those blocks; dense layers leave the indices None and attend the
     full page table.
 
@@ -216,7 +216,7 @@ def forward(
         # Sparse layers attend the per-query top-k blocks with the sparse plan;
         # dense layers leave the indices None and attend the full page table
         # with the dense plan.
-        kv_block_indexes = forward_args.sparse_prediction.sparse_attn_indices
+        kv_block_indexes = forward_args.sparse_runtime_params.sparse_attn_indices
         plan = (
             metadata.msa_decode_gqa_plan
             if kv_block_indexes is not None

tensorrt_llm/_torch/attention_backend/interface.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import copy
 import weakref
 from collections import namedtuple
@@ -27,7 +30,8 @@
 from ..pyexecutor.resource_manager import KVCacheManager
 from ..pyexecutor.trace_log_utils import log_tensor_size
 from ..utils import get_model_extra_attrs
-from .sparse.params import SkipSoftmaxKernelParams, SparseMetadataParams
+from .sparse.params import (SparseBackendForwardArgs, SparseMetadataParams,
+                            SparseRuntimeParams)
 
 try:
     # Transformers v5
@@ -886,27 +890,6 @@ class CustomAttentionMask(str, Enum):
 AttentionMask = Union[PredefinedAttentionMask, CustomAttentionMask]
 
 
-@dataclass(kw_only=True, slots=True)
-class SparsePrediction:
-    """Sparse KV / attention indices predicted by the framework backends.
-
-    RocketKV and DSA produce these from ``sparse_kv_predict`` /
-    ``sparse_attn_predict``, telling the attention op which KV tokens to keep
-    and which blocks to attend to. Backends that don't predict leave
-    ``AttentionForwardArgs.sparse_prediction`` at its default-constructed value
-    (all-``None`` / ``0`` fields).
-    """
-    sparse_kv_indices: Optional[torch.Tensor] = None
-    sparse_kv_offsets: Optional[torch.Tensor] = None
-    sparse_attn_indices: Optional[torch.Tensor] = None
-    sparse_attn_offsets: Optional[torch.Tensor] = None
-    sparse_attn_indices_block_size: int = 0
-    # DeepSeek-V4 sparse-MLA only: per-token compressed top-k lengths and the
-    # base pointer of the compressed KV cache pool (compress_ratio > 1).
-    sparse_mla_topk_lens: Optional[torch.Tensor] = None
-    compressed_kv_cache_pool_ptr: Optional[int] = None
-
-
 @dataclass(kw_only=True, slots=True)
 class AttentionForwardArgs:
     """Per-forward optional arguments for attention backends."""
@@ -960,17 +943,14 @@ class AttentionForwardArgs:
     sage_attn_num_elts_per_blk_v: int = 0
     sage_attn_qk_int8: bool = False
 
-    topk_indices: Optional[torch.Tensor] = None
-
     is_fused_qkv: bool = False
     update_kv_cache: bool = True
     # Optional normalized diffusion timestep for timestep-varying sparse attention.
     timestep: Optional[torch.Tensor] = None
 
-    sparse_prediction: SparsePrediction = field(
-        default_factory=SparsePrediction)
-    skip_softmax_kernel_params: SkipSoftmaxKernelParams = field(
-        default_factory=SkipSoftmaxKernelParams)
+    sparse_backend_args: Optional[SparseBackendForwardArgs] = None
+    sparse_runtime_params: SparseRuntimeParams = field(
+        default_factory=SparseRuntimeParams)
 
     @property
     def mask_type(self) -> int:

tensorrt_llm/_torch/attention_backend/sparse/__init__.py

@@ -1,7 +1,10 @@
-from .utils import (get_flashinfer_sparse_attn_attention_backend,
-                    get_sparse_attn_kv_cache_manager,
-                    get_trtllm_sparse_attn_attention_backend,
-                    get_vanilla_sparse_attn_attention_backend)
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .registry import (get_flashinfer_sparse_attn_attention_backend,
+                       get_sparse_attn_kv_cache_manager,
+                       get_trtllm_sparse_attn_attention_backend,
+                       get_vanilla_sparse_attn_attention_backend)
 
 __all__ = [
     "get_sparse_attn_kv_cache_manager",

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

@@ -13,22 +13,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .backend import DeepseekV4TrtllmAttention
 from .cache_manager import DeepseekV4CacheManager
-from .deepseek_v4 import (
-    DeepseekV4AttentionType,
-    DeepSeekV4MetadataParams,
-    DeepSeekV4Params,
-    DeepseekV4TrtllmAttention,
-    make_deepseek_v4_sparse_metadata_params,
-    make_deepseek_v4_sparse_params,
-)
+from .indexer import DeepseekV4Indexer
+from .metadata import DeepseekV4TrtllmAttentionMetadata
+from .params import DeepseekV4AttentionType, DeepSeekV4MetadataParams, DeepSeekV4Params
 
 __all__ = [
     "DeepSeekV4MetadataParams",
     "DeepSeekV4Params",
     "DeepseekV4AttentionType",
     "DeepseekV4CacheManager",
+    "DeepseekV4Indexer",
     "DeepseekV4TrtllmAttention",
-    "make_deepseek_v4_sparse_metadata_params",
-    "make_deepseek_v4_sparse_params",
+    "DeepseekV4TrtllmAttentionMetadata",
 ]