diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index ee5185ccaa42..2f54a29eef45 100644 --- a/docs/source/models/visual-generation.md +++ b/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. diff --git a/requirements.txt b/requirements.txt index abcaef09250a..664737e25204 100644 --- a/requirements.txt +++ b/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 diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 6b5d9b538d8b..6d8f44416c30 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/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", diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py new file mode 100644 index 000000000000..85e5f44f0bdf --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py @@ -0,0 +1,777 @@ +# 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. +""" +cuDNN SDPA backend for visual generation models. + +Three recipes, selected by ``quant_attention_config`` (see +``tensorrt_llm.visual_gen.args.AttentionConfig``): + +============ ========================================= ================== +Recipe ``quant_attention_config`` cuDNN node +============ ========================================= ================== +``no_quant`` ``None`` ``sdpa`` +``fp8`` ``qk_dtype='fp8'``, ``v_dtype='fp8'`` ``sdpa_fp8`` +``mxfp8`` ``qk_dtype='mxfp8'``, ``v_dtype='mxfp8'`` ``sdpa_mxfp8`` +============ ========================================= ================== + +Layout: NHD ``[B, S, H, D]``. cuDNN takes explicit per-tensor strides, so the +NHD buffers are described to the graph as ``[B, H, S, D]`` with BSHD strides and +never transposed into a real HND copy. +""" + +import math +import threading +from dataclasses import dataclass, field +from typing import Any, ClassVar, Dict, Optional, Tuple + +import cudnn +import torch + +from tensorrt_llm.logger import logger +from tensorrt_llm.visual_gen.args import QuantAttentionConfig + +from ...attention_backend.interface import PredefinedAttentionMask +from .interface import AttentionBackend, AttentionTensorLayout + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def _pad_up(x: int, multiple: int) -> int: + return _ceil_div(x, multiple) * multiple + + +def _row_major_stride(*dims: int) -> list: + """Row-major (contiguous) strides for ``dims``.""" + acc = 1 + strides = [] + for dim in reversed(dims): + strides.append(acc) + acc *= dim + return list(reversed(strides)) + + +def _torch_to_cudnn_dtype(dtype: torch.dtype) -> Any: + mapping = { + torch.float16: cudnn.data_type.HALF, + torch.bfloat16: cudnn.data_type.BFLOAT16, + torch.float32: cudnn.data_type.FLOAT, + torch.float8_e4m3fn: cudnn.data_type.FP8_E4M3, + } + if dtype not in mapping: + raise ValueError(f"No cudnn.data_type mapping for torch dtype {dtype}.") + return mapping[dtype] + + +# ============================================================================ +# Quantization helpers +# ============================================================================ + + +def _quantize_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize ``x`` to FP8 e4m3 with a single scale. + + The fused amax+quantize op computes the scale and the FP8 data in one pass and + returns the descale on device, so the recipe stays free of host synchronization. + + Returns: + x_q: FP8 tensor with the same shape as ``x``. + descale: ``[1, 1, 1, 1]`` float32 device tensor with ``x ~= x_q * descale``. + """ + # The op requires contiguous input. + x_q, descale = torch.ops.trtllm.quantize_e4m3_per_tensor(x.contiguous()) + return x_q, descale.float().reshape(1, 1, 1, 1) + + +def _quantize_mxfp8_qk(x_bhsd: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a Q or K tensor to MXFP8, blocking along the head dimension. + + cuDNN wants ``descale_q``/``descale_k`` as ``[B, H, S_padded, D_scale]`` with + ``S_padded`` a multiple of 128, ``D_scale = ceil(D / 32)`` padded to a multiple + of 4, ``stride[3] == 1`` and ``F8_128x4`` reordering, which is what + ``torch.ops.trtllm.mxfp8_quantize(..., is_sf_swizzled_layout=True)`` emits for a + ``[B * H * S_padded, D]`` matrix. S is padded to a multiple of 128 so that the + 128-row scale-factor tiles align with ``(b, h)`` boundaries and the flat + scale-factor buffer can be viewed as ``[B, H, S_padded, D_scale]``. + + The quantized data is returned as a view into the S-padded buffer; cuDNN accepts + a strided Q/K as long as the head-dim stride is 1. + + Returns: + x_q: ``[B, H, S, D]`` float8_e4m3fn view into a ``[B, H, S_padded, D]`` buffer. + x_sf: ``[B, H, S_padded, D_scale]`` uint8 E8M0 scale factors. + """ + if x_bhsd.dim() != 4: + raise ValueError(f"_quantize_mxfp8_qk expects [B, H, S, D]; got {tuple(x_bhsd.shape)}.") + b, h, s, d = x_bhsd.shape + if d % 32 != 0: + raise ValueError(f"head_dim={d} must be a multiple of the MXFP8 block size 32.") + + s_pad = _pad_up(s, 128) + if s_pad != s: + x_padded = x_bhsd.new_zeros(b, h, s_pad, d) + x_padded[:, :, :s, :] = x_bhsd + else: + x_padded = x_bhsd.contiguous() + + x_q_2d, x_sf_1d = torch.ops.trtllm.mxfp8_quantize(x_padded.reshape(b * h * s_pad, d), True, 32) + d_scale = _pad_up(d // 32, 4) + x_q = x_q_2d.view(b, h, s_pad, d)[:, :, :s, :] + return x_q, x_sf_1d.view(b, h, s_pad, d_scale) + + +def _quantize_mxfp8_v(x_bhsd: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize a V tensor to MXFP8, blocking along the *sequence* dimension. + + The second attention GEMM contracts over S, so V's scale factors block along S + rather than D. cuDNN wants ``descale_v`` as ``[B, H, S_scale, D_padded]`` with + ``S_scale = ceil(S_kv / 32)`` padded to a multiple of 4, ``stride[2] == 1`` and + ``F8_128x4`` reordering. Quantizing the transposed tensor + ``[B, H, D_padded, S_padded]`` produces that buffer; the returned view carries + the transposed (S-scale contiguous) strides cuDNN asks for. + + Returns: + x_q: ``[B, H, S, D]`` float8_e4m3fn (unpadded). + x_sf: ``[B, H, S_scale, D_padded]`` uint8 E8M0 scale factors, ``stride[2] == 1``. + """ + if x_bhsd.dim() != 4: + raise ValueError(f"_quantize_mxfp8_v expects [B, H, S, D]; got {tuple(x_bhsd.shape)}.") + b, h, s, d = x_bhsd.shape + + s_pad = _pad_up(s, 128) + d_pad = _pad_up(d, 128) + # [B, H, D_padded, S_padded]. Zero padding does not affect the per-block amax. + x_t = x_bhsd.new_zeros(b, h, d_pad, s_pad) + x_t[:, :, :d, :s] = x_bhsd.transpose(2, 3) + + x_q_2d, x_sf_1d = torch.ops.trtllm.mxfp8_quantize(x_t.reshape(b * h * d_pad, s_pad), True, 32) + s_scale = _pad_up(s_pad // 32, 4) + x_q = x_q_2d.view(b, h, d_pad, s_pad)[:, :, :d, :s].permute(0, 1, 3, 2).contiguous() + # [B, H, D_padded, S_scale] -> [B, H, S_scale, D_padded] (stride[2] == 1). + x_sf = x_sf_1d.view(b, h, d_pad, s_scale).permute(0, 1, 3, 2) + return x_q, x_sf + + +# ============================================================================ +# Graph geometry +# ============================================================================ + + +@dataclass +class _CuDNNGraphBundle: + """A built cuDNN graph plus the tensor handles needed to bind buffers.""" + + graph: Any + workspace_size: int + inputs: Dict[str, Any] = field(default_factory=dict) + outputs: Dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class _CuDNNProblemShape: + """Problem geometry shared by all three recipes. + + Every tensor's strides are part of the geometry, and of the graph cache key: + the FP8 and unquantized recipes hand cuDNN BSHD-strided views of NHD buffers, + and the MXFP8 recipe passes Q/K as views into sequence-padded buffers. + """ + + b: int + h_q: int + h_kv: int + s_q: int + s_kv: int + d_qk: int + d_v: int + q_strides: Tuple[int, ...] + k_strides: Tuple[int, ...] + v_strides: Tuple[int, ...] + o_strides: Tuple[int, ...] + + +# ============================================================================ +# VisualGen AttentionBackend class +# ============================================================================ + + +class CuDNNAttention(AttentionBackend): + """cuDNN SDPA backend for visual generation. + + Runs unquantized (bf16/fp16), per-tensor FP8, or block-scaled MXFP8 attention + through cuDNN's fused ``sdpa`` / ``sdpa_fp8`` / ``sdpa_mxfp8`` nodes. The recipe + comes from ``quant_attention_config``; ``None`` means unquantized. + + The compiled-graph cache is process-wide and shared by every instance. + cuDNN handles and cached graphs are isolated by CUDA device. + """ + + _cudnn_lib_version = None + _cudnn_handles: ClassVar[Dict[int, Any]] = {} + _graph_cache: ClassVar[Dict[Tuple, _CuDNNGraphBundle]] = {} + _scales_cache: ClassVar[Dict[Tuple[int, str], Any]] = {} + _cache_lock: ClassVar[threading.Lock] = threading.Lock() + + 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, + **kwargs, + ): + 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 or torch.bfloat16 + self.quant_attention_config = quant_attention_config + self.recipe = self.resolve_recipe(quant_attention_config) + self.check_library_feature(self.recipe) + self.scale = 1.0 / math.sqrt(head_dim) + + # Always use NHD via strides to avoid transpose. + self._preferred_layout = AttentionTensorLayout.NHD + + @staticmethod + def resolve_recipe(quant_attention_config: Optional[QuantAttentionConfig]) -> str: + """Map the validated public recipe onto a cuDNN SDPA node.""" + if quant_attention_config is None: + return "no_quant" + qk_dtype, v_dtype = quant_attention_config.qk_dtype, quant_attention_config.v_dtype + if qk_dtype != v_dtype: + raise ValueError( + f"cuDNN backend requires qk_dtype == v_dtype; got qk_dtype={qk_dtype!r}, " + f"v_dtype={v_dtype!r}. cuDNN's FP8 and MXFP8 SDPA nodes quantize both GEMMs " + "with the same element format." + ) + if qk_dtype in ("fp8", "mxfp8"): + return qk_dtype + raise ValueError( + f"cuDNN backend does not support qk_dtype={qk_dtype!r}; supported recipes are " + "unquantized (quant_attention_config=None), 'fp8' and 'mxfp8'." + ) + + # ------------------------------------------------------------------ + # cuDNN handle and compiled-graph cache + # ------------------------------------------------------------------ + + @classmethod + def _get_lib_version(cls): + if cls._cudnn_lib_version is None: + cls._cudnn_lib_version = cudnn.backend_version() + torch_cudnn_lib_version = torch.backends.cudnn.version() + if cls._cudnn_lib_version != torch_cudnn_lib_version: + logger.critical( + "PyTorch and cuDNN Frontend loaded different cuDNN backends: " + f"PyTorch:v{torch_cudnn_lib_version} != cuDNN-FE:v{cls._cudnn_lib_version}. " + ) + return cls._cudnn_lib_version + + @classmethod + def _get_handle(cls, device: torch.device) -> Any: + device_index = device.index if device.index is not None else torch.cuda.current_device() + with cls._cache_lock: + handle = cls._cudnn_handles.get(device_index) + if handle is None: + with torch.cuda.device(device_index): + handle = cudnn.create_handle() + cls._cudnn_handles[device_index] = handle + return handle + + @classmethod + def check_hardware_compatibility(cls, device: torch.device, recipe: str = "no_quant") -> None: + compute_capability = torch.cuda.get_device_capability(device) + gpu_arch = f"sm_{compute_capability[0]}{compute_capability[1]}a" + if gpu_arch not in ("sm_100a", "sm_103a") and recipe != "no_quant": + raise ImportError("cuDNN quantized attention requires NVIDIA Blackwell-class GPU.") + + @classmethod + def check_library_feature(cls, recipe: str = "no_quant") -> None: + # Check if the cuDNN library supports the requested functionality. + if cls._get_lib_version() < 90100: + raise ImportError("cuDNN attention backend requires cuDNN library v9.1.0 or later.") + if cls._get_lib_version() < 92100 and recipe == "mxfp8": + raise ImportError("cuDNN MXFP8 attention requires cuDNN library v9.21.0 or later.") + + @classmethod + def clear_graph_cache(cls) -> None: + """Drop every compiled cuDNN graph (used by tests).""" + with cls._cache_lock: + cls._graph_cache.clear() + cls._scales_cache.clear() + + @staticmethod + def _build_graph( + recipe: str, + shape: _CuDNNProblemShape, + is_causal: bool, + sm_scale: float, + out_dtype: torch.dtype, + with_lse: bool, + ) -> _CuDNNGraphBundle: + s = shape + out_cudnn_dtype = _torch_to_cudnn_dtype(out_dtype) + fp8 = cudnn.data_type.FP8_E4M3 + e8m0 = cudnn.data_type.FP8_E8M0 + f32 = cudnn.data_type.FLOAT + + io_dtype = fp8 if recipe in ("fp8", "mxfp8") else out_cudnn_dtype + graph = cudnn.pygraph( + io_data_type=io_dtype, + intermediate_data_type=f32, + compute_data_type=f32, + name=f"visual_gen_sdpa_{recipe}", + ) + + def _tensor(name: str, dims: Tuple[int, ...], dtype: Any, **kwargs) -> Any: + return graph.tensor( + name=name, + dim=list(dims), + stride=_row_major_stride(*dims), + data_type=dtype, + **kwargs, + ) + + q_t = graph.tensor( + name="q", dim=[s.b, s.h_q, s.s_q, s.d_qk], stride=list(s.q_strides), data_type=io_dtype + ) + k_t = graph.tensor( + name="k", + dim=[s.b, s.h_kv, s.s_kv, s.d_qk], + stride=list(s.k_strides), + data_type=io_dtype, + ) + v_t = graph.tensor( + name="v", dim=[s.b, s.h_kv, s.s_kv, s.d_v], stride=list(s.v_strides), data_type=io_dtype + ) + inputs: Dict[str, Any] = {"q": q_t, "k": k_t, "v": v_t} + amax_s_t = None + + if recipe == "no_quant": + o_t, stats_t = graph.sdpa( + q=q_t, + k=k_t, + v=v_t, + attn_scale=sm_scale, + use_causal_mask=is_causal, + generate_stats=with_lse, + ) + amax_o_t = None + elif recipe == "fp8": + # Per-tensor descales for Q/K/V plus the FP8-quantized softmax output S. + for name in ("descale_q", "descale_k", "descale_v", "descale_s", "scale_s", "scale_o"): + inputs[name] = _tensor(name, (1, 1, 1, 1), f32) + o_t, stats_t, amax_s_t, amax_o_t = graph.sdpa_fp8( + q=q_t, + k=k_t, + v=v_t, + descale_q=inputs["descale_q"], + descale_k=inputs["descale_k"], + descale_v=inputs["descale_v"], + descale_s=inputs["descale_s"], + scale_s=inputs["scale_s"], + scale_o=inputs["scale_o"], + attn_scale=sm_scale, + use_causal_mask=is_causal, + generate_stats=with_lse, + ) + elif recipe == "mxfp8": + s_q_pad = _pad_up(s.s_q, 128) + s_kv_pad = _pad_up(s.s_kv, 128) + qk_d_scale = _pad_up(_ceil_div(s.d_qk, 32), 4) + v_s_scale = _pad_up(s_kv_pad // 32, 4) + v_d_pad = _pad_up(s.d_v, 128) + reorder = {"reordering_type": cudnn.tensor_reordering.F8_128x4} + + inputs["descale_q"] = _tensor( + "descale_q", (s.b, s.h_q, s_q_pad, qk_d_scale), e8m0, **reorder + ) + inputs["descale_k"] = _tensor( + "descale_k", (s.b, s.h_kv, s_kv_pad, qk_d_scale), e8m0, **reorder + ) + # descale_v blocks along S, so its S-scale dimension is the contiguous one. + inputs["descale_v"] = graph.tensor( + name="descale_v", + dim=[s.b, s.h_kv, v_s_scale, v_d_pad], + stride=[s.h_kv * v_s_scale * v_d_pad, v_s_scale * v_d_pad, 1, v_s_scale], + data_type=e8m0, + **reorder, + ) + o_t, stats_t, amax_o_t = graph.sdpa_mxfp8( + q=q_t, + k=k_t, + v=v_t, + descale_q=inputs["descale_q"], + descale_k=inputs["descale_k"], + descale_v=inputs["descale_v"], + attn_scale=sm_scale, + use_causal_mask=is_causal, + generate_stats=with_lse, + ) + else: + raise ValueError(f"Unknown cuDNN SDPA recipe {recipe!r}.") + + out_dims = (s.b, s.h_q, s.s_q, s.d_v) + o_t.set_output(True).set_dim(list(out_dims)).set_stride(list(s.o_strides)).set_data_type( + out_cudnn_dtype + ) + outputs: Dict[str, Any] = {"o": o_t} + + if with_lse: + # Stats follows HN1 layout, unlike input tensors. + stats_dims = (s.b, s.h_q, s.s_q, 1) + stats_t.set_output(True).set_dim(list(stats_dims)).set_stride( + _row_major_stride(*stats_dims) + ).set_data_type(f32) + outputs["stats"] = stats_t + # The quantized nodes emit amax(S) / amax(O) for calibration; inference binds + # scratch buffers for them. + for name, tensor in (("amax_s", amax_s_t), ("amax_o", amax_o_t)): + if tensor is not None: + tensor.set_output(True).set_dim([1, 1, 1, 1]).set_stride( + [1, 1, 1, 1] + ).set_data_type(f32) + outputs[name] = tensor + + graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + return _CuDNNGraphBundle( + graph=graph, workspace_size=graph.get_workspace_size(), inputs=inputs, outputs=outputs + ) + + @classmethod + @torch.compiler.disable + def _get_or_build_graph( + cls, + recipe: str, + shape: _CuDNNProblemShape, + is_causal: bool, + sm_scale: float, + out_dtype: torch.dtype, + with_lse: bool, + device: torch.device, + ) -> _CuDNNGraphBundle: + device_index = device.index if device.index is not None else torch.cuda.current_device() + key = (device_index, recipe, shape, is_causal, sm_scale, out_dtype, with_lse) + with cls._cache_lock: + bundle = cls._graph_cache.get(key) + if bundle is None: + cls.check_hardware_compatibility(device, recipe) + logger.debug( + f"[CuDNNAttention] building graph on cuda:{device_index}: " + f"recipe={recipe} {shape} causal={is_causal}" + ) + with torch.cuda.device(device_index): + bundle = cls._build_graph( + recipe, shape, is_causal, sm_scale, out_dtype, with_lse + ) + cls._graph_cache[key] = bundle + return bundle + + @classmethod + @torch.compiler.disable + def _execute_graph( + cls, + bundle: _CuDNNGraphBundle, + tensor_map: Dict[Any, torch.Tensor], + device: torch.device, + ) -> None: + with torch.cuda.device(device): + handle = cls._get_handle(device) + cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) + workspace = torch.empty(bundle.workspace_size, dtype=torch.uint8, device=device) + bundle.graph.execute(tensor_map, workspace, handle=handle) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + @classmethod + @torch.compiler.disable + def _softmax_scales(cls, recipe: str, device: torch.device) -> Any: + """Get the per-device cache for recipe-determined constants. + + - recipe="fp8": (scale_s, descale_s, scale_o) for the FP8 recipe + Softmax output lies in [0, 1], so it is scaled by 448 to fill the FP8 range before Bmm2. + """ + key = (device.index if device.index is not None else torch.cuda.current_device(), recipe) + with cls._cache_lock: + scales = cls._scales_cache.get(key) + if scales is None: + if recipe == "fp8": + scale_s = torch.full((1, 1, 1, 1), 448.0, dtype=torch.float32, device=device) + scale_o = torch.ones(1, 1, 1, 1, dtype=torch.float32, device=device) + scales = (scale_s, scale_s.reciprocal(), scale_o) + cls._scales_cache[key] = scales + else: + # "no_quant" and "mxfp8" don't need device-side constants. + pass + return scales + + @staticmethod + @torch.compiler.disable + def _is_fused_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> bool: + """Whether Q, K, and V are gap-free slices from one fused-QKV buffer. + + At TRTLLM/VisualGen, Attention.get_qkv() at QKVMode.FUSE_QKV produces this type of buffer, + where one qkv_proj output is split along the last dim. + """ + if q.dim() != 4 or k.shape != v.shape or q.shape[3] != k.shape[3]: + return False + if q.dtype != k.dtype or q.dtype != v.dtype: + return False + storage_ptr = q.untyped_storage().data_ptr() + if any(x.untyped_storage().data_ptr() != storage_ptr for x in (k, v)): + return False + + b, s, h_q, d = q.shape + h_kv = k.shape[2] + q_dim, kv_dim = h_q * d, h_kv * d + total = q_dim + 2 * kv_dim + + # Heads packed within each slice, and one shared row pitch across all three. + for x, h in ((q, h_q), (k, h_kv), (v, h_kv)): + if x.shape[:3] != (b, s, h): + return False + if x.stride() != (s * total, total, d, 1): + return False + + base = q.storage_offset() + return k.storage_offset() == base + q_dim and v.storage_offset() == base + q_dim + kv_dim + + @staticmethod + @torch.compiler.disable + def _as_fused_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Zero-copy [B, S, q_dim + 2 * kv_dim] view of the parent QKV buffer.""" + b, s, h_q, d = q.shape + total = (h_q + 2 * k.shape[2]) * d + return torch.as_strided(q, (b, s, total), (s * total, total, 1), q.storage_offset()) + + def _validate_inputs(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + for name, tensor in (("q", q), ("k", k), ("v", v)): + if tensor.dim() != 4: + raise ValueError( + f"cuDNN backend expects a 4D [B, S, H, D] {name}; got {tuple(tensor.shape)}." + ) + if k.shape[:3] != v.shape[:3]: + raise ValueError(f"K/V shape mismatch: {tuple(k.shape)} vs {tuple(v.shape)}.") + if q.shape[0] != k.shape[0]: + raise ValueError(f"Batch size mismatch: q={q.shape[0]} vs k={k.shape[0]}.") + if q.shape[3] != k.shape[3]: + raise ValueError(f"Q/K head_dim mismatch: {q.shape[3]} vs {k.shape[3]}.") + if q.shape[3] != self.head_dim: + raise ValueError( + f"cuDNN backend was configured with head_dim={self.head_dim}, " + f"but received head_dim={q.shape[3]}." + ) + if q.shape[2] != self.num_heads: + raise ValueError( + f"cuDNN backend was configured with num_heads={self.num_heads}, " + f"but received num_heads={q.shape[2]}." + ) + if q.shape[2] % k.shape[2] != 0: + raise ValueError( + f"num_heads={q.shape[2]} must be a multiple of num_kv_heads={k.shape[2]} for GQA." + ) + # cuDNN's FP8 and MXFP8 SDPA engines support head_dim <= 128. + if self.recipe != "no_quant" and max(q.shape[3], v.shape[3]) > 128: + raise ValueError( + f"cuDNN quantized SDPA supports head_dim <= 128; got qk={q.shape[3]}, " + f"v={v.shape[3]}. Drop quant_attention_config to run unquantized." + ) + if self.recipe == "mxfp8" and q.shape[3] % 32 != 0: + raise ValueError( + f"cuDNN MXFP8 requires head_dim to be a multiple of 32; got {q.shape[3]}." + ) + + def _run( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + is_causal: bool, + with_lse: bool, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + self._validate_inputs(q, k, v) + b, s_q, h_q, d_qk = q.shape + _, s_kv, h_kv, d_v = v.shape + device = q.device + out_dtype = self.dtype if self.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 + q, k, v = q.to(out_dtype), k.to(out_dtype), v.to(out_dtype) + + # Q/K/V arrive as NHD [B, S, H, D]. cuDNN wants a [B, H, S, D] *logical* tensor, + # but takes the strides explicitly, so `.transpose(1, 2)` stays a view and the + # BSHD bytes are fed to the kernel untouched -- no transpose copy anywhere. + buffers: Dict[str, torch.Tensor] = {} + if self.recipe == "no_quant": + buffers.update( + q=q.contiguous().transpose(1, 2), + k=k.contiguous().transpose(1, 2), + v=v.contiguous().transpose(1, 2), + ) + elif self.recipe == "fp8": + if d_qk == d_v and self._is_fused_qkv(q, k, v): + # One amax+cast over the packed buffer for fused-QKV, mirroring TransformerEngine's + # handling of same-layout buffers. + qkv_q, descale = _quantize_fp8(self._as_fused_qkv(q, k, v)) + q_q, k_q, v_q = ( + t.unflatten(-1, (h, d_qk)) + for t, h in zip( + qkv_q.split([h_q * d_qk, h_kv * d_qk, h_kv * d_qk], dim=-1), + (h_q, h_kv, h_kv), + ) + ) + descale_q = descale_k = descale_v = descale + else: + q_q, descale_q = _quantize_fp8(q) + k_q, descale_k = _quantize_fp8(k) + v_q, descale_v = _quantize_fp8(v) + scale_s, descale_s, scale_o = self._softmax_scales(self.recipe, device) + buffers.update( + q=q_q.transpose(1, 2), + k=k_q.transpose(1, 2), + v=v_q.transpose(1, 2), + descale_q=descale_q, + descale_k=descale_k, + descale_v=descale_v, + descale_s=descale_s, + scale_s=scale_s, + scale_o=scale_o, + ) + else: # mxfp8 + # The MXFP8 quantizers pad along S and emit their own HND buffers, so they + # take the [B, H, S, D] view and do the layout change as part of the pack. + q_q, descale_q = _quantize_mxfp8_qk(q.transpose(1, 2)) + k_q, descale_k = _quantize_mxfp8_qk(k.transpose(1, 2)) + v_q, descale_v = _quantize_mxfp8_v(v.transpose(1, 2)) + buffers.update( + q=q_q, k=k_q, v=v_q, descale_q=descale_q, descale_k=descale_k, descale_v=descale_v + ) + + shape = _CuDNNProblemShape( + b=b, + h_q=h_q, + h_kv=h_kv, + s_q=s_q, + s_kv=s_kv, + d_qk=d_qk, + d_v=d_v, + q_strides=tuple(buffers["q"].stride()), + k_strides=tuple(buffers["k"].stride()), + v_strides=tuple(buffers["v"].stride()), + # O is written straight into an NHD buffer, likewise via strides. + o_strides=(s_q * h_q * d_v, d_v, h_q * d_v, 1), + ) + bundle = self._get_or_build_graph( + self.recipe, + shape, + is_causal=is_causal, + sm_scale=self.scale, + out_dtype=out_dtype, + with_lse=with_lse, + device=device, + ) + + output = torch.empty(b, s_q, h_q, d_v, dtype=out_dtype, device=device) + tensor_map = {bundle.inputs[name]: tensor for name, tensor in buffers.items()} + tensor_map[bundle.outputs["o"]] = output.transpose(1, 2) + + stats: Optional[torch.Tensor] = None + if with_lse: + stats = torch.empty(b, h_q, s_q, 1, dtype=torch.float32, device=device) + tensor_map[bundle.outputs["stats"]] = stats + for amax_name in ("amax_s", "amax_o"): + if amax_name in bundle.outputs: + tensor_map[bundle.outputs[amax_name]] = torch.empty( + 1, 1, 1, 1, dtype=torch.float32, device=device + ) + + self._execute_graph(bundle, tensor_map, device) + + # Stats is packed [B, H, S, 1]; other backends expose LSE as [B, S, H]. + lse = None if stats is None else stats.squeeze(-1).transpose(1, 2).contiguous() + return output, lse + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + key_padding_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """Run attention. + + Args: + q: Query tensor ``[B, S_q, H, D]``. + k: Key tensor ``[B, S_kv, H_kv, D]``. + v: Value tensor ``[B, S_kv, H_kv, D_v]``. + attention_mask: ``CAUSAL`` or ``FULL``. + key_padding_mask: Not supported by this backend. + + Returns: + Output tensor ``[B, S_q, H, D_v]``. + """ + output, _ = self._run( + q, k, v, is_causal=self._resolve_mask(attention_mask, key_padding_mask), with_lse=False + ) + return output + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + key_padding_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Same as :meth:`forward`, additionally returning the softmax log-sum-exp. + + Returns: + output: ``[B, S_q, H, D_v]`` + lse: ``[B, S_q, H]`` float32 + """ + output, lse = self._run( + q, k, v, is_causal=self._resolve_mask(attention_mask, key_padding_mask), with_lse=True + ) + assert lse is not None, "cuDNN graph was built with stats but returned none." + return output, lse + + @staticmethod + def _resolve_mask( + attention_mask: PredefinedAttentionMask, key_padding_mask: Optional[torch.Tensor] + ) -> bool: + if key_padding_mask is not None: + raise NotImplementedError( + "cuDNN backend does not support key_padding_mask; use the VANILLA backend." + ) + return attention_mask == PredefinedAttentionMask.CAUSAL + + @classmethod + def support_lse(cls) -> bool: + return True + + @classmethod + def support_fused_qkv(cls) -> bool: + return False + + @property + def preferred_layout(self) -> AttentionTensorLayout: + return self._preferred_layout diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 12108ad84ece..3e3e4dbd4e41 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -37,7 +37,7 @@ def get_visual_gen_attention_backend( Get diffusion attention backend class by name. Args: - backend_name: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL") + backend_name: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL", "CUDNN") Returns: Diffusion attention backend class @@ -51,8 +51,11 @@ def get_visual_gen_attention_backend( Requires flash-attn package with cute interface - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA from AttentionConfig.sparse_attention_config. + - "CUDNN": cuDNN fused SDPA. Unquantized by default; quant_attention_config + selects per-tensor FP8 or block-scaled MXFP8 (Blackwell). """ # Lazy imports to avoid circular dependency + from .cudnn import CuDNNAttention from .cute_dsl import CuTeDSLAttention from .flash_attn4 import FlashAttn4Attention from .trtllm import TrtllmAttention @@ -68,6 +71,8 @@ def get_visual_gen_attention_backend( return FlashAttn4Attention elif backend_name == "CUTEDSL": return CuTeDSLAttention + elif backend_name == "CUDNN": + return CuDNNAttention else: # Default to VANILLA for maximum compatibility return VanillaAttention @@ -94,7 +99,7 @@ def create_attention( internally, simplifying the forward() call. Args: - backend: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL") + backend: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL", "CUDNN") layer_idx: Layer index in the model num_heads: Number of attention heads head_dim: Dimension per head @@ -116,7 +121,8 @@ def create_attention( """ attn_cls = get_visual_gen_attention_backend(backend) - # Forward the validated quantization recipe to TRTLLM or the dense CuTe DSL FMHA backend. + # Forward the validated quantization recipe to the TRTLLM, dense CuTe DSL FMHA, + # or cuDNN backend. if attention_config is not None and attention_config.quant_attention_config is not None: kwargs["quant_attention_config"] = attention_config.quant_attention_config if backend.upper() == "TRTLLM": diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index d314c92b24ae..e8243bd935a2 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -44,7 +44,7 @@ class QuantAttentionConfig(StrictBaseModel): - """Attention quantization recipe (TRTLLM / CUTEDSL backends). + """Attention quantization recipe (TRTLLM / CUTEDSL / CUDNN backends). Specifies Q/K and V quantization formats and their optional block sizes. @@ -60,10 +60,13 @@ class QuantAttentionConfig(StrictBaseModel): "integer and floating-point element formats; mxfp8 and nvfp4 are block-scaled formats." ), ) - v_dtype: Literal["fp8"] = Field( + v_dtype: Literal["fp8", "mxfp8"] = Field( "fp8", status="prototype", - description="V quantization dtype. The current kernels always load V in FP8 (e4m3).", + description=( + "V quantization format. fp8 is the 8-bit floating-point element format; mxfp8 is " + "the block-scaled format." + ), ) q_block_size: int = Field( 0, @@ -97,16 +100,16 @@ class QuantAttentionConfig(StrictBaseModel): class AttentionConfig(StrictBaseModel): """Configuration for Attention layers.""" - backend: Literal["VANILLA", "TRTLLM", "FA4", "CUTEDSL"] = Field( + backend: Literal["VANILLA", "TRTLLM", "FA4", "CUTEDSL", "CUDNN"] = Field( "VANILLA", status="prototype", - description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, FA4, CUTEDSL", + description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, FA4, CUTEDSL, CUDNN", ) quant_attention_config: Optional[QuantAttentionConfig] = Field( None, status="prototype", description=( - "Quantized-attention recipe (TRTLLM / CUTEDSL backends). " + "Quantized-attention recipe (TRTLLM / CUTEDSL / CUDNN backends). " "Set to a QuantAttentionConfig instance to enable quantized " "attention; leave as None to disable." ), @@ -130,6 +133,11 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("fp8", "fp8", (1, 1, 1)), ("fp8", "fp8", (1, 4, 1)), } + # cuDNN fused SDPA quantizes both GEMMs with the same element format. + CUDNN_RECIPES = { + ("fp8", "fp8", (0, 0, 0)), + ("mxfp8", "mxfp8", (0, 0, 0)), + } CUTEDSL_RECIPES = { ("bf16", "fp8", (0, 0, 0)), ("mxfp8", "fp8", (0, 0, 0)), @@ -163,9 +171,18 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": f"(qk_dtype, v_dtype, (q_block, k_block, v_block)): " f"{sorted(CUTEDSL_RECIPES)}." ) + elif self.backend == "CUDNN": + if recipe not in CUDNN_RECIPES: + raise ValueError( + f"Unsupported quant_attention_config={self.quant_attention_config!r} " + f"for backend='CUDNN'. Supported recipes " + f"(qk_dtype, v_dtype, (q_block, k_block, v_block)): " + f"{sorted(CUDNN_RECIPES)}. Omit quant_attention_config to run " + f"unquantized attention." + ) else: raise ValueError( - f"quant_attention_config requires backend in ('TRTLLM', 'CUTEDSL'), " + f"quant_attention_config requires backend in ('TRTLLM', 'CUTEDSL', 'CUDNN'), " f"got backend='{self.backend}'. Either change backend or " f"remove quant_attention_config." ) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8d47a2e993e5..60119deaca09 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -241,6 +241,7 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py + - unittest/_torch/visual_gen/test_attention_cudnn.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_fa4_key_padding_mask.py - unittest/_torch/visual_gen/test_attention_perf.py diff --git a/tests/unittest/_torch/visual_gen/test_attention_cudnn.py b/tests/unittest/_torch/visual_gen/test_attention_cudnn.py new file mode 100644 index 000000000000..75737855a676 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_cudnn.py @@ -0,0 +1,243 @@ +# 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. + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask +from tensorrt_llm._torch.utils import unswizzle_sf +from tensorrt_llm._torch.visual_gen.attention_backend import create_attention +from tensorrt_llm._torch.visual_gen.attention_backend.cudnn import ( + CuDNNAttention, + _quantize_mxfp8_qk, + _quantize_mxfp8_v, +) +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig + +# Recipe name -> quant_attention_config accepted by AttentionConfig. +RECIPES = { + "no_quant": None, + "fp8": QuantAttentionConfig(qk_dtype="fp8", v_dtype="fp8"), + "mxfp8": QuantAttentionConfig(qk_dtype="mxfp8", v_dtype="mxfp8"), +} + +# (name, batch, num_heads, num_kv_heads, seq_len_q, seq_len_kv, head_dim) +SHAPES = [ + ("mha", 2, 8, 8, 512, 512, 128), + ("gqa", 1, 16, 4, 1024, 1024, 128), + ("mqa", 1, 16, 1, 256, 256, 64), + ("cross_gqa", 1, 8, 2, 512, 333, 128), +] + +# Cosine similarity against an FP32 reference. The quantized recipes are bounded by +# the FP8 P*V GEMM that cuDNN performs internally on both quantized paths. +MIN_COSINE = {"no_quant": 0.9999, "fp8": 0.995, "mxfp8": 0.995} + + +def _require_cudnn(recipe: str) -> None: + if not torch.cuda.is_available(): + pytest.skip("cuDNN attention backend requires CUDA.") + if recipe != "no_quant" and torch.cuda.get_device_capability()[0] < 10: + pytest.skip("cuDNN FP8/MXFP8 SDPA requires a Blackwell-class GPU (sm100+).") + + +def _make_qkv(batch, num_heads, num_kv_heads, seq_q, seq_kv, head_dim, device): + torch.manual_seed(0) + q = torch.randn(batch, seq_q, num_heads, head_dim, device=device, dtype=torch.bfloat16) + k = torch.randn(batch, seq_kv, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16) + v = torch.randn(batch, seq_kv, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16) + return q, k, v + + +def _reference(q, k, v, is_causal): + """FP32 SDPA reference plus its log-sum-exp.""" + q, k, v = (t.transpose(1, 2) for t in (q, k, v)) + num_heads, num_kv_heads = q.shape[1], k.shape[1] + out = F.scaled_dot_product_attention( + q.float(), k.float(), v.float(), is_causal=is_causal, enable_gqa=num_heads != num_kv_heads + ) + k_rep = k.float().repeat_interleave(num_heads // num_kv_heads, dim=1) + logits = (q.float() @ k_rep.transpose(-1, -2)) * q.shape[-1] ** -0.5 + if is_causal: + seq_q, seq_kv = q.shape[2], k.shape[2] + causal_mask = torch.ones(seq_q, seq_kv, device=q.device, dtype=torch.bool) + logits = logits.masked_fill(causal_mask.triu(seq_kv - seq_q + 1), float("-inf")) + return out.transpose(1, 2), torch.logsumexp(logits, dim=-1).transpose(1, 2) + + +@pytest.mark.parametrize("recipe", list(RECIPES)) +@pytest.mark.parametrize("shape", SHAPES, ids=[s[0] for s in SHAPES]) +@pytest.mark.parametrize("is_causal", [False, True]) +def test_cudnn_attention(recipe, shape, is_causal): + """Output and LSE match an FP32 SDPA reference for every recipe and mask.""" + _require_cudnn(recipe) + _, batch, num_heads, num_kv_heads, seq_q, seq_kv, head_dim = shape + if is_causal and seq_q != seq_kv: + pytest.skip("Causal masking is only meaningful for self attention.") + + device = torch.device("cuda") + q, k, v = _make_qkv(batch, num_heads, num_kv_heads, seq_q, seq_kv, head_dim, device) + attention = CuDNNAttention( + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + dtype=torch.bfloat16, + quant_attention_config=RECIPES[recipe], + ) + mask = PredefinedAttentionMask.CAUSAL if is_causal else PredefinedAttentionMask.FULL + output, lse = attention.forward_with_lse(q, k, v, attention_mask=mask) + + ref_out, ref_lse = _reference(q, k, v, is_causal) + assert output.shape == (batch, seq_q, num_heads, head_dim) + assert output.dtype == torch.bfloat16 + assert lse.shape == (batch, seq_q, num_heads) + + cosine = F.cosine_similarity(output.float().flatten(), ref_out.flatten(), dim=0).item() + assert cosine > MIN_COSINE[recipe], f"{recipe}: cosine similarity {cosine} too low" + # LSE is computed in FP32 by cuDNN on all recipes; only the quantized logits differ. + torch.testing.assert_close(lse, ref_lse, atol=0.5 if recipe != "no_quant" else 1e-2, rtol=0.0) + + +@pytest.mark.parametrize("seq_len", [128, 1000]) +@pytest.mark.parametrize("head_dim", [32, 64, 128]) +def test_mxfp8_scale_factor_layout_roundtrip(seq_len, head_dim): + """Dequantizing with cuDNN's documented layout reproduces the input tensor. + + Guards the two scale-factor layouts the MXFP8 recipe relies on: Q/K block along + the head dim (``[B, H, S_padded, D_scale]``, ``stride[3] == 1``) and V blocks + along the sequence dim (``[B, H, S_scale, D_padded]``, ``stride[2] == 1``). + Both come out of ``torch.ops.trtllm.mxfp8_quantize``'s swizzled (``F8_128x4``) + buffer, so a layout regression shows up here rather than as a silent accuracy + loss inside cuDNN. + """ + _require_cudnn("mxfp8") + device = torch.device("cuda") + batch, num_heads = 2, 3 + torch.manual_seed(0) + + def dequant_e8m0(scale_factors): + return torch.exp2(scale_factors.float() - 127.0) + + # Block-aligned dynamic range: a wrong scale mapping cannot average out. + channel_gain = 10.0 ** (torch.arange(head_dim, device=device) // 32 % 7 - 3) + x_qk = ( + torch.randn(batch, num_heads, seq_len, head_dim, device=device) * channel_gain + ).bfloat16() + x_q, scale_factors = _quantize_mxfp8_qk(x_qk) + seq_padded, d_scale = scale_factors.shape[2], scale_factors.shape[3] + linear_sf = unswizzle_sf( + scale_factors.reshape(-1), batch * num_heads * seq_padded, d_scale * 32, 32 + ).view(batch, num_heads, seq_padded, d_scale)[:, :, :seq_len, : head_dim // 32] + dequantized = x_q.float() * dequant_e8m0(linear_sf).repeat_interleave(32, dim=-1) + rel_err = ((dequantized - x_qk.float()).norm() / x_qk.float().norm()).item() + assert rel_err < 0.05, f"Q/K MXFP8 round-trip error {rel_err} exceeds e4m3 block noise" + + token_gain = (10.0 ** (torch.arange(seq_len, device=device) // 32 % 7 - 3)).view(seq_len, 1) + x_v = (torch.randn(batch, num_heads, seq_len, head_dim, device=device) * token_gain).bfloat16() + v_q, v_sf = _quantize_mxfp8_v(x_v) + s_scale, d_padded = v_sf.shape[2], v_sf.shape[3] + assert v_sf.stride(2) == 1, "cuDNN requires descale_v to have a contiguous S-scale dimension" + linear_v_sf = unswizzle_sf( + v_sf.permute(0, 1, 3, 2).reshape(-1), batch * num_heads * d_padded, s_scale * 32, 32 + ).view(batch, num_heads, d_padded, s_scale)[:, :, :head_dim, :] + v_scale = dequant_e8m0(linear_v_sf).repeat_interleave(32, dim=-1)[..., :seq_len] + dequantized_v = v_q.float() * v_scale.permute(0, 1, 3, 2) + rel_err_v = ((dequantized_v - x_v.float()).norm() / x_v.float().norm()).item() + assert rel_err_v < 0.05, f"V MXFP8 round-trip error {rel_err_v} exceeds e4m3 block noise" + + +def test_cudnn_graph_cache_reuses_plans(): + """Repeated calls with the same geometry reuse one compiled graph.""" + _require_cudnn("no_quant") + device = torch.device("cuda") + q, k, v = _make_qkv(1, 4, 4, 256, 256, 64, device) + attention = CuDNNAttention(num_heads=4, head_dim=64, dtype=torch.bfloat16) + + CuDNNAttention.clear_graph_cache() + attention.forward(q, k, v) + assert len(CuDNNAttention._graph_cache) == 1 + attention.forward(q, k, v) + assert len(CuDNNAttention._graph_cache) == 1 + # A different mask needs its own plan. + attention.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL) + assert len(CuDNNAttention._graph_cache) == 2 + + +def test_cudnn_backend_wires_validated_recipes(): + """create_attention("CUDNN") wires the validated recipe into the backend.""" + _require_cudnn("mxfp8") + attention = create_attention( + backend="CUDNN", + layer_idx=0, + num_heads=8, + head_dim=128, + num_kv_heads=8, + dtype=torch.bfloat16, + attention_config=AttentionConfig(backend="CUDNN", quant_attention_config=RECIPES["mxfp8"]), + ) + assert isinstance(attention, CuDNNAttention) + assert attention.recipe == "mxfp8" + assert attention.preferred_layout == AttentionTensorLayout.NHD + assert attention.support_lse() and not attention.support_fused_qkv() + + +@pytest.mark.parametrize("num_kv_heads", [8, 2], ids=["mha", "gqa"]) +def test_cudnn_fp8_attention_with_fused_qkv(num_kv_heads): + """Packed Q/K/V are quantized through the shared-scale path; separate ones are not. + + The shared scale shifts the output slightly, so the two paths are compared at the + suite's FP8 tolerance rather than for equality. + """ + _require_cudnn("fp8") + device = torch.device("cuda") + batch, num_heads, seq_len, head_dim = 1, 8, 256, 64 + q_dim, kv_dim = num_heads * head_dim, num_kv_heads * head_dim + + # Q/K/V as views into one buffer, the way get_qkv splits under FUSE_QKV. + torch.manual_seed(0) + qkv = torch.randn(batch, seq_len, q_dim + 2 * kv_dim, device=device, dtype=torch.bfloat16) + q, k, v = ( + x.view(batch, seq_len, -1, head_dim) for x in qkv.split([q_dim, kv_dim, kv_dim], dim=-1) + ) + + assert CuDNNAttention._is_fused_qkv(q, k, v) + fused_qkv = CuDNNAttention._as_fused_qkv(q, k, v) + assert fused_qkv.data_ptr() == qkv.data_ptr(), "fused view must alias, not copy" + torch.testing.assert_close(fused_qkv, qkv, atol=0.0, rtol=0.0) + + # Cloning breaks the shared storage, so equal values take the per-tensor path. + separate = tuple(x.contiguous() for x in (q, k, v)) + assert not CuDNNAttention._is_fused_qkv(*separate) + + attention = CuDNNAttention( + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + dtype=torch.bfloat16, + quant_attention_config=RECIPES["fp8"], + ) + fused_out = attention.forward(q, k, v) + separate_out = attention.forward(*separate) + ref_out, _ = _reference(q, k, v, is_causal=False) + + for what, a, b in ( + ("fused vs reference", fused_out.float(), ref_out), + ("separate vs reference", separate_out.float(), ref_out), + ("fused vs separate", fused_out.float(), separate_out.float()), + ): + cosine = F.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() + assert cosine > MIN_COSINE["fp8"], f"{what}: cosine similarity {cosine} too low" diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index a372db9641f2..67fd60c1e18f 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -18,6 +18,7 @@ # ============================================================================ # Flash Attention 4 availability # ============================================================================ +from tensorrt_llm._torch.visual_gen.attention_backend.cudnn import CuDNNAttention from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd from tensorrt_llm._torch.visual_gen.attention_backend.parallel import ( @@ -166,11 +167,23 @@ def create_model_config( return config -def _require_attention_backend(attn_backend: str) -> None: +def _require_attention_backend( + attn_backend: str, + quant_attention_config: "QuantAttentionConfig | None" = None, +) -> None: if attn_backend == "FA4" and not _flash_attn4_available: pytest.fail("FlashAttention 4 backend is required for FA4 attention test") if attn_backend == "CUTEDSL" and not _cute_dsl_available: pytest.fail("CuTe DSL backend is required for CUTEDSL attention test") + if attn_backend == "CUDNN": + recipe = CuDNNAttention.resolve_recipe(quant_attention_config) + try: + if not torch.cuda.is_available(): + raise ImportError("CUDA not available") + CuDNNAttention.check_hardware_compatibility(torch.device("cuda"), recipe) + CuDNNAttention.check_library_feature(recipe) + except ImportError as e: + pytest.skip(f"cuDNN detected hardware/library incompatibility: {e}") if attn_backend == "CUTEDSL": compute_capability = torch.cuda.get_device_capability() gpu_arch = f"sm_{compute_capability[0]}{compute_capability[1]}a" @@ -349,13 +362,16 @@ def test_cross_attention_with_sage_config_falls_back_to_vanilla(self): ("FA4", None), ("CUTEDSL", None), ("CUTEDSL", QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8")), + ("CUDNN", None), + ("CUDNN", QuantAttentionConfig(qk_dtype="fp8", v_dtype="fp8")), + ("CUDNN", QuantAttentionConfig(qk_dtype="mxfp8", v_dtype="mxfp8")), ], ) def test_self_attention_equivalence( head_dim: int, attn_backend: str, quant_attention_config: "QuantAttentionConfig | None" ): """Test that integrated self-attention produces same output as naive.""" - _require_attention_backend(attn_backend) + _require_attention_backend(attn_backend, quant_attention_config) print("\n" + "=" * 60) print("Testing Self-Attention Equivalence") @@ -409,7 +425,14 @@ def test_self_attention_equivalence( # Compare (using looser tolerance for bf16) max_diff = (out_naive - out_integrated).abs().max().item() mean_diff = (out_naive - out_integrated).abs().mean().item() - tol = 1e-2 if quant_attention_config is None else 2e-2 + if quant_attention_config is None: + tol = 1e-2 + elif quant_attention_config.qk_dtype == "bf16": + # V-only quantization (QK16PV8): Bmm1 still runs in BF16. + tol = 2e-2 + else: + # Q/K quantized as well (CUDNN FP8 / MXFP8): both GEMMs carry FP8 noise. + tol = 4e-2 is_close = torch.allclose(out_naive, out_integrated, rtol=tol, atol=tol) print("\nResults:") @@ -548,7 +571,7 @@ def test_cross_attention_equivalence( head_dim: int, attn_backend: str, quant_attention_config: "QuantAttentionConfig | None" ): """Test that integrated cross-attention produces same output as naive.""" - _require_attention_backend(attn_backend) + _require_attention_backend(attn_backend, quant_attention_config) print("\n" + "=" * 60) print("Testing Cross-Attention Equivalence") @@ -637,6 +660,8 @@ def test_cross_attention_equivalence( ("FA4", None), ("CUTEDSL", None), ("CUTEDSL", QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8")), + ("CUDNN", None), + ("CUDNN", QuantAttentionConfig(qk_dtype="mxfp8", v_dtype="mxfp8")), ], ) def test_fast_cross_attention_wan_shapes( @@ -649,7 +674,7 @@ def test_fast_cross_attention_wan_shapes( quant_attention_config: "QuantAttentionConfig | None", ): """Test fast cross-attention correctness at Wan-realistic shapes.""" - _require_attention_backend(attn_backend) + _require_attention_backend(attn_backend, quant_attention_config) hidden_size = num_heads * head_dim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 74667c3d714f..c1bc5183265a 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -162,6 +162,23 @@ def test_blockscaled_qk_dtype_rejected_on_trtllm(self): quant_attention_config=QuantAttentionConfig(qk_dtype="nvfp4"), ) + def test_supported_quant_config_cudnn_fp8(self): + attention = AttentionConfig( + backend="CUDNN", + quant_attention_config=QuantAttentionConfig(qk_dtype="fp8", v_dtype="fp8"), + ) + + assert attention.quant_attention_config is not None + + def test_supported_quant_config_cudnn_mxfp8(self): + attention = AttentionConfig( + backend="CUDNN", + quant_attention_config=QuantAttentionConfig(qk_dtype="mxfp8", v_dtype="mxfp8"), + ) + + assert attention.quant_attention_config is not None + assert attention.quant_attention_config.v_dtype == "mxfp8" + def test_sage_qk_block_size_rejected_on_cute(self): with pytest.raises(ValidationError, match="Unsupported quant_attention_config"): AttentionConfig(