[TRTLLM-15585][feat] Wire SkipSoftmax sparse attention into the CuTeDSL backend

docs/source/visual-gen/features/sparse-attention.md

@@ -90,7 +90,7 @@ User configuration is supplied through Python or YAML and controls how the check
 
 `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers.
 
-Skip Softmax Attention only works with the **TRTLLM** attention backend in VisualGen. Set `attention_config.backend` to `TRTLLM` when enabling it.
+Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA is the only CUTEDSL sparse-attention algorithm that is mutually exclusive with quantized attention.
 
 #### Python API
 
@@ -130,6 +130,19 @@ args = VisualGenArgs(
 )
 ```
 
+```python
+# CUTEDSL backend:
+args = VisualGenArgs(
+    model="<path_or_hf_id>",
+    attention_config=AttentionConfig(
+        backend="CUTEDSL",
+        sparse_attention_config=SkipSoftmaxAttentionConfig(
+            threshold_scale_factor=5000.0,
+        ),
+    ),
+)
+```
+
 #### YAML
 
 ```yaml
@@ -151,6 +164,15 @@ attention_config:
     disabled_until_timestep: 0.6
 ```
 
+```yaml
+# CUTEDSL backend:
+attention_config:
+  backend: CUTEDSL
+  sparse_attention_config:
+    algorithm: skip_softmax
+    threshold_scale_factor: 5000.0
+```
+
 ### CUDA Graphs
 
 `disabled_until_timestep` creates two sparse-attention phases when it is set: the high-timestep disabled phase and the enabled phase after the cutoff. VisualGen includes that phase in CUDA graph keys so graph capture does not reuse a graph across different Skip Softmax Attention settings. See [VisualGen CUDA Graphs](cuda-graph.md) for the general capture and replay design.

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py

@@ -13,18 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-CuTe DSL (NVIDIA kernels) Dense FMHA Backend for Visual Generation Models
+CuTe DSL (NVIDIA kernels) FMHA Backend for Visual Generation Models
 
-JIT-compiles the dense FMHA kernel and caches the compiled artifact for each kernel configuration.
-Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA sparse path uses
-VSAAttention from vsa.py instead.
+JIT-compiles dense or SkipSoftmax FMHA and caches the compiled artifact for each kernel
+configuration. Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA
+sparse path uses VSAAttention from vsa.py instead.
 """
 
 import math
 from typing import Any, NamedTuple, Tuple
 
 import torch
 
+from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxParams
 from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.args import QuantAttentionConfig
 
@@ -64,6 +65,34 @@
 # ============================================================================
 
 
+def _resolve_skip_softmax_threshold_scale_factor(
+    threshold_scale_factor: float | None,
+    sparse_params: SkipSoftmaxParams | None,
+    timestep: Any,
+) -> float | None:
+    """Resolve the active CuTeDSL threshold for the current denoising phase."""
+    if sparse_params is not None:
+        if timestep is None and sparse_params.scheduler.disabled_until_timestep is not None:
+            # Fail-open: a missing timestep resolves to the full (unthrottled)
+            # threshold, i.e. skip-softmax runs during the high-noise steps
+            # `disabled_until_timestep` exists to protect. This is silent
+            # elsewhere (a quality regression, not an error), so surface it
+            # once per process instead of only in this function's return value.
+            logger.warning_once(
+                "SkipSoftmax scheduler has disabled_until_timestep="
+                f"{sparse_params.scheduler.disabled_until_timestep} configured, but no "
+                "`timestep` was passed to the CuTeDSL attention forward call. Skip-softmax "
+                "will run unthrottled (as if past the cutoff) until `timestep` is threaded "
+                "through.",
+                key="cute_dsl_skip_softmax_missing_timestep",
+            )
+        runtime_params = sparse_params.scheduler.get_runtime_params(timestep=timestep)
+        threshold_scale_factor = runtime_params.threshold_scale_factor_prefill
+    if threshold_scale_factor is None or threshold_scale_factor <= 0.0:
+        return None
+    return threshold_scale_factor
+
+
 def _check_cute_runtime_available() -> None:
     if _cute_dsl_import_error is None:
         return
@@ -239,6 +268,8 @@ def cute_dsl_fmha_fwd(
     scale_o: float | torch.Tensor = 1.0,
     is_persistent: bool = True,
     skip_softmax_threshold_scale_factor: float | None = None,
+    sparse_params: SkipSoftmaxParams | None = None,
+    timestep: Any = None,
     qk_sf_vec: int = 0,
     q_sf: torch.Tensor | None = None,
     k_sf: torch.Tensor | None = None,
@@ -252,7 +283,14 @@ def cute_dsl_fmha_fwd(
     When `qk_sf_vec` is non-zero, dispatches to the block-scaled kernel class:
     32 selects MXFP8 (Q/K stored as FP8 e4m3, SFs as Float8E8M0FNU uint8 storage);
     16 selects NVFP4 (Q/K stored as packed FP4 in torch.uint8, SFs as Float8E4M3FN in uint8 storage).
+    When `sparse_params` is set, its timestep-aware scheduler overrides the direct threshold.
     """
+    skip_softmax_threshold_scale_factor = _resolve_skip_softmax_threshold_scale_factor(
+        skip_softmax_threshold_scale_factor,
+        sparse_params,
+        timestep,
+    )
+
     _check_cute_runtime_available()
     _validate_inputs(q, k, v, o)
     if qk_sf_vec != 0:
@@ -583,7 +621,13 @@ def __init__(
         num_kv_heads: int | None = None,
         dtype: torch.dtype | None = None,
         quant_attention_config: QuantAttentionConfig | None = None,
+        # Legacy static threshold, superseded by `sparse_params`'s timestep-aware
+        # scheduler for every in-tree construction path (`create_attention` never
+        # forwards this). Kept only for direct-construction debug/testing use
+        # (e.g. unit tests that want a fixed threshold without a scheduler);
+        # mutually exclusive with `sparse_params` below.
         skip_softmax_threshold_scale: float | None = None,
+        sparse_params: SkipSoftmaxParams | None = None,
         **kwargs,
     ):
         self.layer_idx = layer_idx
@@ -592,7 +636,10 @@ def __init__(
         self.num_kv_heads = num_kv_heads or num_heads
         self.dtype = dtype
         self.quant_attention_config = quant_attention_config
+        if skip_softmax_threshold_scale is not None and sparse_params is not None:
+            raise ValueError("Set either skip_softmax_threshold_scale or sparse_params, not both.")
         self.skip_softmax_threshold_scale = skip_softmax_threshold_scale
+        self.sparse_params = sparse_params
         self.scale = 1.0 / math.sqrt(head_dim)
 
         # CuTe DSL expects [B, S, H, D] format
@@ -727,6 +774,8 @@ def _fwd(
             scale_v_channels=scale_v_channels,
             scale_o=kwargs.get("scale_o", 1.0),
             skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale,
+            sparse_params=self.sparse_params,
+            timestep=kwargs.get("timestep"),
             qk_sf_vec=qk_sf_vec,
             q_sf=q_sf,
             k_sf=k_sf,

tensorrt_llm/_torch/visual_gen/attention_backend/utils.py

@@ -49,8 +49,8 @@ def get_visual_gen_attention_backend(
                     Better performance but requires fused QKV
         - "FA4": Flash Attention 4; provides higher speedup on Blackwell GPUs (sm100)
                  Requires flash-attn package with cute interface
-        - "CUTEDSL": CuTe DSL kernels. create_attention selects dense FMHA or VSA from
-                      AttentionConfig.sparse_attention_config.
+        - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA
+                      from AttentionConfig.sparse_attention_config.
     """
     # Lazy imports to avoid circular dependency
     from .cute_dsl import CuTeDSLAttention

tensorrt_llm/_torch/visual_gen/models/modeling.py

@@ -77,7 +77,9 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None:
         if not isinstance(sparse_config, SkipSoftmaxAttentionConfig):
             return
 
-        disabled_until_timestep = sparse_config.disabled_until_timestep
+        disabled_until_timestep = sparse_config.resolve_disabled_until_timestep(
+            pretrained_config=self.model_config.pretrained_config,
+        )
         if disabled_until_timestep is None:
             return
 

tensorrt_llm/_torch/visual_gen/modules/attention.py

@@ -220,10 +220,14 @@ def __init__(
             backend_num_heads = self.local_num_attention_heads
             backend_num_kv_heads = self.local_num_key_value_heads
 
-        # Resolve sparse attention params for TRTLLM backend
+        # Lower the shared SkipSoftmax user/checkpoint config for each backend
+        # whose kernel consumes SkipSoftmaxParams.
         sparse_params = None
         ss_cfg = config.attention.sparse_attention_config
-        if isinstance(ss_cfg, SkipSoftmaxAttentionConfig) and backend_name == "TRTLLM":
+        if isinstance(ss_cfg, SkipSoftmaxAttentionConfig) and backend_name in (
+            "TRTLLM",
+            "CUTEDSL",
+        ):
             sparse_params = ss_cfg.to_sparse_params(
                 module_name=self.module_name,
                 pretrained_config=config.pretrained_config,

tensorrt_llm/visual_gen/args.py

@@ -116,7 +116,7 @@ class AttentionConfig(StrictBaseModel):
         status="prototype",
         description=(
             "Sparse attention recipe. Discriminated by algorithm: "
-            "skip_softmax (TRTLLM backend) or VSA (CUTEDSL backend)."
+            "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)."
         ),
     )
 
@@ -177,29 +177,34 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig":
             return self
 
         algo = self.sparse_attention_config.algorithm
-        required_backend = {"skip_softmax": "TRTLLM", "vsa": "CUTEDSL"}.get(algo)
-        if required_backend is None:
+        supported_backends = {
+            "skip_softmax": ("TRTLLM", "CUTEDSL"),
+            "vsa": ("CUTEDSL",),
+        }.get(algo)
+        if supported_backends is None:
             return self
 
-        if self.backend != required_backend:
+        if self.backend not in supported_backends:
             raise ValueError(
                 f"sparse_attention_config with algorithm='{algo}' requires "
-                f"backend='{required_backend}', got backend='{self.backend}'. "
-                f"Either set backend='{required_backend}' or remove "
+                f"backend in {supported_backends}, got backend='{self.backend}'. "
+                f"Either select a supported backend or remove "
                 f"sparse_attention_config."
             )
         return self
 
     @model_validator(mode="after")
     def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig":
-        # quant_attention_config and sparse_attention_config are mutually exclusive.
+        # VSA replaces the dense CuTeDSL path and cannot compose with quantized
+        # attention. SkipSoftmax is part of that dense path and can compose.
         if (
             self.backend == "CUTEDSL"
             and self.quant_attention_config is not None
             and self.sparse_attention_config is not None
+            and self.sparse_attention_config.algorithm == "vsa"
         ):
             raise ValueError(
-                "CUTEDSL backend: quant_attention_config and "
+                "CUTEDSL backend: quant_attention_config and VSA "
                 "sparse_attention_config are mutually exclusive (the "
                 "CuTeDSLAttention dispatcher selects either the dense path "
                 "or the sparse VSA path, not both)."

tensorrt_llm/visual_gen/sparse_attention.py

@@ -72,7 +72,6 @@ def to_sparse_params(self, **kwargs):
         from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import (
             SkipSoftmaxParams,
             SkipSoftmaxScheduler,
-            skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config,
         )
 
         module_name = kwargs.get("module_name", None)
@@ -81,13 +80,9 @@ def to_sparse_params(self, **kwargs):
         if module_name is not None and self._is_disabled(module_name, ckpt_sparse_attention_config):
             return None
 
-        disabled_until_timestep = self.disabled_until_timestep
-        if disabled_until_timestep is None:
-            disabled_until_timestep = (
-                skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config(
-                    ckpt_sparse_attention_config
-                )
-            )
+        disabled_until_timestep = self.resolve_disabled_until_timestep(
+            checkpoint_config=ckpt_sparse_attention_config,
+        )
 
         threshold_scale_factor = self.resolve_threshold_scale_factor(ckpt_sparse_attention_config)
         if threshold_scale_factor is None:
@@ -104,6 +99,33 @@ def to_sparse_params(self, **kwargs):
             return None
         return SkipSoftmaxParams(scheduler=scheduler)
 
+    def resolve_disabled_until_timestep(
+        self,
+        *,
+        checkpoint_config: Optional[Dict[str, Any]] = None,
+        pretrained_config: Any = None,
+    ) -> Optional[float]:
+        """Resolve the user override or checkpoint-provided timestep cutoff.
+
+        Exactly one of ``checkpoint_config`` (a raw checkpoint config dict) or
+        ``pretrained_config`` (the model's pretrained config object/dict) is
+        expected per call site; both default to ``None`` so a mistyped keyword
+        raises ``TypeError`` here instead of silently resolving to ``None``.
+        """
+        if self.disabled_until_timestep is not None:
+            return self.disabled_until_timestep
+
+        from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import (
+            skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config,
+        )
+
+        ckpt_sparse_attention_config = self._ckpt_sparse_attention_config_from_kwargs(
+            {"checkpoint_config": checkpoint_config, "pretrained_config": pretrained_config}
+        )
+        return skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config(
+            ckpt_sparse_attention_config
+        )
+
     def _is_disabled(
         self,
         module_name: str,

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -43,6 +43,7 @@ l0_cpu:
   - unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py
   - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py
   - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py
+  - unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py
   - unittest/_torch/visual_gen/test_attention_integration.py
   - unittest/_torch/visual_gen/test_cache_dit.py
   - unittest/_torch/visual_gen/test_flux_infer.py