[None][feat] Add cuDNN attention backend

docs/source/models/visual-generation.md

@@ -12,9 +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 (`VANILLA`), TRT-LLM kernels (`TRTLLM`), TRT-LLM CuTe DSL kernels (`CUTEDSL`, Blackwell-class GPUs), and Flash Attention 4 (`FA4`).
+- Pluggable attention backends: PyTorch SDPA (`VANILLA`), TRT-LLM kernels (`TRTLLM`), TRT-LLM CuTe DSL kernels (`CUTEDSL`, Blackwell-class GPUs), Flash Attention 4 (`FA4`), and cuDNN fused SDPA (`CUDNN`).
 - 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).
+- Quantized attention support: `QK16PV8` to quantize Bmm2 on `CUTEDSL`, `SAGE` to run SageAttention on `TRTLLM`, and per-tensor FP8 / block-scaled MXFP8 on `CUDNN` (requires Blackwell SM100).
 - Sparse attention support: see [VisualGen Sparse Attention](../visual-gen/features/sparse-attention.md).
 - Multi-GPU parallelism (CFG parallel, Ulysses sequence parallel, Tensor parallelism).
 - **Step caching** — two runtime caching backends (**TeaCache** and **Cache-DiT**) that skip transformer computation on steps where the step-to-step change is small.
@@ -169,9 +169,10 @@ By default, `strict=True` raises when adapter tensors cannot be matched, have un
 
 ### 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` and are mutually exclusive with each other.
+In addition to linear-layer quantization, VisualGen exposes several **attention-level** quantization presets that operate inside the attention kernel. They are configured through `AttentionConfig.quant_attention_config` 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`.
+- **FP8 / MXFP8** (`CUDNN` backend): Runs cuDNN's fused FP8 SDPA. `qk_dtype='fp8'` uses one scale per tensor; `qk_dtype='mxfp8'` with `v_dtype='mxfp8'` uses MXFP8 block scaling. Both require Blackwell-class GPUs and `head_dim in {32, 64, 96, 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.
 
 
@@ -213,6 +214,23 @@ args = VisualGenArgs(
 )
 ```
 
+Python API for cuDNN MXFP8:
+
+```python
+from tensorrt_llm import VisualGenArgs
+
+args = VisualGenArgs(
+    model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    attention_config={
+        "backend": "CUDNN",
+        "quant_attention_config": {
+            "qk_dtype": "mxfp8",
+            "v_dtype": "mxfp8",
+        },
+    },
+)
+```
+
 ### CUDA Graphs
 
 VisualGen CUDA graphs capture transformer forward calls during denoising and replay them for later steps with compatible inputs. See [VisualGen CUDA Graphs](../visual-gen/features/cuda-graph.md) for capture scope, graph keys, and sparse-attention phase behavior.

requirements.txt

@@ -95,6 +95,7 @@ cuda-core
 llist
 cuda-tile>=1.0.1
 nvidia-cuda-tileiras>=13.1,<13.2
+nvidia-cudnn-frontend==1.27.0
 etcd-sdk-python==0.0.7
 # etcd-sdk-python imports google.protobuf but omits it from its package metadata.
 protobuf>=5.27.2

tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py

@@ -20,6 +20,7 @@
 simplified metadata that doesn't require KV caching.
 """
 
+from .cudnn import CuDNNAttention
 from .cute_dsl import (
     VSA_TILE_SIZE,
     CuTeDSLAttention,
@@ -42,6 +43,7 @@
     "AttentionTensorLayout",
     "get_visual_gen_attention_backend",
     "create_attention",
+    "CuDNNAttention",
     "CuTeDSLAttention",
     "VSAAttention",
     "FlashAttn4Attention",