[TRTLLM-12436][feat] visual_gen: add CuTe DSL attention via exported binaries

.gitattributes

@@ -16,3 +16,4 @@ docs/source/blogs/media/tech_blog10_full_strategy_performance.png filter=lfs dif
 docs/source/blogs/media/tech_blog10_context_wait_performance.png  filter=lfs diff=lfs merge=lfs -text
 cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cubin/kernelMetaInfo_cubin.cpp filter=lfs diff=lfs merge=lfs -text
 cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cubin/xqa_kernel_cubin.cpp filter=lfs diff=lfs merge=lfs -text
+tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/*/*/*.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -57,6 +57,7 @@ tensorrt_llm/flash_mla/
 tensorrt_llm/flash_mla_cpp_tllm.*.so
 tensorrt_llm/flash_mla_cpp_tllm.pyi
 tensorrt_llm/runtime/kv_cache_manager_v2/**/*.so
+!tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/**/*.so
 **/*__mypyc*.so
 tensorrt_llm/scripts
 *docs/cpp_docs*

docs/source/models/visual-generation.md

@@ -12,8 +12,9 @@ Visual generation models based on diffusion transformers (DiT) have become the s
 TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion models, with a pipeline architecture separate from the LLM inference path. Key capabilities include:
 
 - A shared pipeline abstraction covering the denoising loop, guidance strategies, and component loading.
-- Pluggable attention backends (PyTorch SDPA and TRT-LLM optimized kernels).
+- Pluggable attention backends: PyTorch SDPA (`VANILLA`), TRT-LLM kernels (`TRTLLM`), TRT-LLM CuTe DSL kernels (`CUTEDSL`, Blackwell-class GPUs), and Flash Attention 4 (`FA4`).
 - Quantization support (dynamic and static) using the [ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer) configuration format.
+- Quantized attention support: `QK16PV8` to quantize Bmm2 on `CUTEDSL`, `SAGE` to run SageAttention on `TRTLLM` (requires Blackwell SM100).
 - Multi-GPU parallelism (CFG parallel, Ulysses sequence parallel).
 - **TeaCache** — a runtime caching optimization that skips transformer steps when timestep embeddings change slowly.
 - `trtllm-serve` integration with OpenAI-compatible API endpoints for image and video generation.
@@ -107,6 +108,52 @@ args = VisualGenArgs(
 )
 ```
 
+### Quantized Attention
+
+In addition to linear-layer quantization, VisualGen exposes two **attention-level** quantization presets that operate inside the attention kernel. They are configured through `AttentionConfig.quant_attention_config` (or the `--quant_attention_mode` flag in the example scripts) and are mutually exclusive with each other.
+
+- **QK16PV8** (`CUTEDSL` backend): Keeps Q & K in BF16 and quantizes only V to FP8 (E4M3, per-tensor), thus Bmm1 will be carried out in BF16 with Bmm2 in FP8. Targets Blackwell-class GPUs (`sm_100a` / `sm_103a`) with `head_dim = 128`.
+- **SAGE** (`TRTLLM` backend): Quantizes Q, K, and V with per-block scaling factors. Q/K are stored as INT8 or FP8 (e4m3) and V as FP8 (e4m3); block sizes are tunable per axis (typically `(q, k, v) = (1, 4, 1)` for Wan-1.3B and `(1, 16, 1)` for larger Wan / FLUX checkpoints). Supported recipes are validated at runtime.
+
+
+Python API for SageAttention:
+
+```python
+from tensorrt_llm import VisualGenArgs
+
+args = VisualGenArgs(
+    model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    attention_config={
+        "backend": "TRTLLM",
+        "quant_attention_config": {
+            "qk_dtype": "int8",
+            "q_block_size": 1,
+            "k_block_size": 16,
+            "v_block_size": 1,
+        },
+    },
+)
+```
+
+Python API for QK16PV8:
+
+```python
+from tensorrt_llm import VisualGenArgs
+
+args = VisualGenArgs(
+    model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    attention_config={
+        "backend": "CUTEDSL",
+        "quant_attention_config": {
+            "qk_dtype": "bf16",
+            "q_block_size": 0,
+            "k_block_size": 0,
+            "v_block_size": 0,
+        },
+    },
+)
+```
+
 ### TeaCache
 
 TeaCache caches transformer outputs when timestep embeddings change slowly between denoising steps, skipping redundant computation. Enable via `VisualGenArgs.cache_config` (YAML or programmatic):
@@ -126,7 +173,7 @@ The `teacache_thresh` parameter controls the similarity threshold. Cache-DiT is
 - **CFG Parallelism** (`--cfg_size 2`): Splits positive/negative guidance prompts across GPUs.
 - **Ulysses Parallelism** (`--ulysses_size N`): Splits the sequence dimension across GPUs for longer sequences.
 - **Parallel VAE** (`--parallel_vae_size N`): Shards the final VAE decode along a spatial axis across GPUs, useful to reduce VAE latency and improve GPU utilization (Constraint: `parallel_vae_size ≤ world_size`). Currently only supported for WAN models.
-- **Attention Parallel**: There are 2 methods supported to run attention parallel. Both of these methods require the attention backend to support LSE (only FA4 currently) - 
+- **Attention Parallel**: There are 2 methods supported to run attention parallel. Both of these methods require the attention backend to support LSE (`FA4` and `CUTEDSL`) - 
     - **Attention2D Parallelism** (`--attn2d_row_size N`, `--attn2d_col_size M`): Shards the sequence axis across a 2D `N x M` device mesh, all-gathering Q along rows and K/V along columns so each rank computes a sub-block of the attention matrix (total CP degree = `N * M`; not currently combinable with Ulysses).
     - **Ring Attention Parallelism** (`--ring_size N`): Shards the sequence axis across a 1D ring of `N` ranks and streams K/V blocks around the ring so each rank computes its attention output without materializing the full K/V (mutually exclusive with Attention2D).
 ## Developer Guide

setup.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -171,6 +171,7 @@ def has_ext_modules(self):
     # Include CUDA source for fused MoE align extension so runtime JIT can find it in wheels
     '_torch/auto_deploy/custom_ops/fused_moe/moe_align_kernel.cu',
     '_torch/auto_deploy/custom_ops/fused_moe/triton_fused_moe_configs/*',
+    '_torch/visual_gen/cute_dsl_kernels/blackwell/attention/cubins/**/*.so',
     'usage/schemas/*.json',
 ]
 

tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py

@@ -20,6 +20,7 @@
 simplified metadata that doesn't require KV caching.
 """
 
+from .cute_dsl import CuTeDSLAttention
 from .flash_attn4 import FlashAttn4Attention
 from .interface import AttentionBackend, AttentionTensorLayout
 from .parallel import Attention2DAttention, RingAttention, UlyssesAttention
@@ -33,6 +34,7 @@
     "AttentionTensorLayout",
     "get_visual_gen_attention_backend",
     "create_attention",
+    "CuTeDSLAttention",
     "FlashAttn4Attention",
     "TrtllmAttention",
     "TrtllmAttentionMetadata",

tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl.py

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+CuTe DSL (NVIDIA kernels) Backend for Visual Generation Models
+
+Uses pre-compiled cubins derived from CUTLASS CuTe DSL FMHA.
+Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+
+from tensorrt_llm.visual_gen.args import QuantAttentionConfig
+
+from ...attention_backend.interface import PredefinedAttentionMask
+from .interface import AttentionBackend, AttentionTensorLayout
+
+_cute_dsl_import_error = None
+try:
+    import tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.attention as cute_dsl
+    from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.attention.fmha import (
+        _cute_runtime_import_error,
+    )
+
+    if _cute_runtime_import_error is not None:
+        raise ImportError(_cute_runtime_import_error)
+except (ImportError, OSError) as e:
+    cute_dsl = None
+    _cute_dsl_import_error = e
+
+
+class CuTeDSLAttention(AttentionBackend):
+    """
+    CuTe DSL (NVIDIA kernels) backend for diffusion models.
+
+    Uses pre-compiled cubin kernels (head_dim=128 only).
+    """
+
+    def __init__(
+        self,
+        layer_idx: int = 0,
+        num_heads: int = 8,
+        head_dim: int = 64,
+        num_kv_heads: Optional[int] = None,
+        dtype: Optional[torch.dtype] = None,
+        quant_attention_config: Optional[QuantAttentionConfig] = None,
+        skip_softmax_threshold_scale: Optional[float] = None,
+        **kwargs,
+    ):
+        # Only head_dim=128 cubins are packaged.
+        if head_dim != 128:
+            raise ValueError(f"CUTEDSL cubins require head_dim=128, got head_dim={head_dim}.")
+        self.layer_idx = layer_idx
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.dtype = dtype
+        self.quant_attention_config = quant_attention_config
+        self.skip_softmax_threshold_scale = skip_softmax_threshold_scale
+        self.scale = 1.0 / math.sqrt(head_dim)
+
+        # CuTe DSL expects [B, S, H, D] format
+        self._preferred_layout = AttentionTensorLayout.NHD
+
+    def _prepare_inputs(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attention_mask: PredefinedAttentionMask,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool, torch.dtype]:
+        """Cast inputs to CuTeDSL-compatible dtype and resolve causal flag."""
+        if _cute_dsl_import_error is not None:
+            raise ImportError(
+                f"CuTe DSL kernels are not available. Import error: {_cute_dsl_import_error}"
+            ) from _cute_dsl_import_error
+
+        is_causal = attention_mask == PredefinedAttentionMask.CAUSAL
+
+        # Packaged cubins support float16 and bfloat16 only.
+        origin_dtype = q.dtype
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+        return q, k, v, is_causal, origin_dtype
+
+    # cute_dsl.cute_dsl_fmha_fwd is already decorated with @torch.compiler.disable
+    # Allow torch.compile to fuse preceding linear/norm with quantization of V / seq-preprocess
+    def _fwd(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        is_causal: bool,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len_q, num_heads, _ = q.shape
+        _, seq_len_kv, _, value_head_dim = v.shape
+        out = torch.empty(
+            batch_size,
+            seq_len_q,
+            num_heads,
+            value_head_dim,
+            dtype=q.dtype,
+            device=q.device,
+        )
+        lse = torch.empty(
+            batch_size,
+            seq_len_q,
+            num_heads,
+            dtype=torch.float32,
+            device=q.device,
+        )
+
+        # Options that instructs quantization of V
+        scale_v = kwargs.get("scale_v", 1.0)
+        if self.quant_attention_config is not None:
+            v_qscale = 448.0 / v.abs().amax().clamp(min=1e-3)
+            v = (v * v_qscale).to(torch.float8_e4m3fn)
+            scale_v = scale_v / v_qscale
+
+        # Sequence preproc.
+        qo_indptr_host = [i * seq_len_q for i in range(batch_size + 1)]
+        qo_indptr = torch.tensor(qo_indptr_host).to(device=q.device, dtype=torch.int32)
+        kv_indptr_host = [i * seq_len_kv for i in range(batch_size + 1)]
+        kv_indptr = torch.tensor(kv_indptr_host).to(device=q.device, dtype=torch.int32)
+
+        # Skip softmax.
+        skip_softmax_threshold_scale = self.skip_softmax_threshold_scale
+        if skip_softmax_threshold_scale is not None and skip_softmax_threshold_scale <= 0.0:
+            skip_softmax_threshold_scale = None
+
+        cute_dsl.cute_dsl_fmha_fwd(
+            q.flatten(0, 1).contiguous(),
+            k.flatten(0, 1).contiguous(),
+            v.flatten(0, 1).contiguous(),
+            out.flatten(0, 1),
+            qo_indptr=qo_indptr,
+            kv_indptr=kv_indptr,
+            is_causal=is_causal,
+            sm_scale=self.scale,
+            lse=lse.flatten(0, 1).contiguous(),
+            scale_q=kwargs.get("scale_q", 1.0),
+            scale_k=kwargs.get("scale_k", 1.0),
+            scale_v=scale_v,
+            scale_o=kwargs.get("scale_o", 1.0),
+            max_qo_len=seq_len_q,
+            max_kv_len=seq_len_kv,
+            is_persistent=False,
+            skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale,
+        )
+        return out, lse
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Forward pass using CuTe DSL (NVIDIA kernels).
+
+        Dimensions are derived from tensor shapes (NHD layout: ``[B, S, H, D]``).
+
+        Args:
+            q: Query tensor [batch_size, seq_len, num_heads, head_dim]
+            k: Key tensor [batch_size, seq_len_kv, num_kv_heads, head_dim]
+            v: Value tensor [batch_size, seq_len_kv, num_kv_heads, head_dim]
+            attention_mask: Attention mask type (CAUSAL or FULL)
+
+        Returns:
+            Output tensor [batch_size, seq_len, num_heads, head_dim]
+        """
+        output, _ = self.forward_with_lse(q, k, v, attention_mask=attention_mask, **kwargs)
+        return output
+
+    def forward_with_lse(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass returning both output and log-sum-exp (LSE).
+
+        Returns:
+            output: [batch_size, seq_len, num_heads, head_dim]
+            lse:    [batch_size, num_heads, seq_len] - log-sum-exp per query position,
+                    always in float32. Used for numerically stable combination of
+                    partial attention results in Attention2D parallelism.
+        """
+        q, k, v, is_causal, origin_dtype = self._prepare_inputs(q, k, v, attention_mask)
+        output, lse = self._fwd(q, k, v, is_causal, **kwargs)
+        if output.dtype != origin_dtype:
+            output = output.to(origin_dtype)
+        return output, lse.transpose(1, 2)
+
+    @classmethod
+    def support_lse(cls) -> bool:
+        return True
+
+    @property
+    def preferred_layout(self) -> AttentionTensorLayout:
+        """Return the preferred tensor layout for this backend."""
+        return self._preferred_layout
+
+    @classmethod
+    def support_fused_qkv(cls) -> bool:
+        return False

tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py

@@ -16,7 +16,7 @@
 Parallelism Wrappers
 
 Wraps any attention backend with a parallelism strategy. Not a standalone
-backend — compose around a real backend (VANILLA/TRTLLM/FA4).
+backend — compose around a real backend (VANILLA/TRTLLM/FA4/CUTEDSL).
 
 """
 
@@ -256,8 +256,8 @@ class Attention2DAttention(AttentionBackend):
     Supported inner backends
     ------------------------
     The inner backend must support LSE output (``support_lse() -> True``) — required
-    for the reduce-scatter combine step.  Currently only the FA4 backend
-    (``FlashAttn4Attention``) meets this requirement.
+    for the reduce-scatter combine step.  Currently the FA4 and CUTEDSL
+    backends meet this requirement.
 
     Note: ``AttentionTensorLayout.NHD`` and ``AttentionTensorLayout.HND`` are both
     handled transparently; transposition is applied before the inner forward and