From 2ac6336b575da3b540abe9ae999b0047a3f8d4ef Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Sun, 1 Nov 2026 09:14:02 -0800 Subject: [PATCH] [feat]: add opt-in cuDNN SDPA for MiniMax-H3 dense attention The H3 dense self-attention call always ran on FlashAttention. On GB200, cuDNN's fused SDPA engine is a few percent faster in-model at 37k-109k tokens, but the ATen op has no shape metadata: it always reads inputs as [B, H, S, D], so handing it contiguous [B, S, H, D] is silently misread. Add FASTVIDEO_MINIMAX_H3_DENSE_ATTN=flash|cudnn (default flash) and route the dense, unmasked, inference-only call through a helper that presents BHSD strided views. Training calls, compiled regions, head sizes >128, non-contiguous inputs, VSA blocks and NVFP4 stay on their existing paths. FA4_* tuning and the default behaviour are unchanged. --- docs/contributing/env_vars.md | 1 + docs/inference/optimizations.md | 18 +++ fastvideo/attention/backends/flash_attn.py | 14 ++ fastvideo/attention/utils/cudnn_sdpa.py | 86 +++++++++++ fastvideo/envs.py | 7 + fastvideo/models/dits/minimax_h3.py | 5 + .../tests/attention/test_cudnn_sdpa_dense.py | 142 ++++++++++++++++++ 7 files changed, 273 insertions(+) create mode 100644 fastvideo/attention/utils/cudnn_sdpa.py create mode 100644 fastvideo/tests/attention/test_cudnn_sdpa_dense.py diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index e2cd6646a..2fcbae2e4 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -194,6 +194,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | | `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | | `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | +| `FASTVIDEO_MINIMAX_H3_DENSE_ATTN` | one of flash, cudnn | `flash` | attention | Kernel for MiniMax-H3 dense DiT self-attention: flash keeps the selected FlashAttention backend; with FLASH_ATTN, cudnn sends the no-mask, no-grad dense call through cuDNN's fused SDPA. cuDNN has a different reduction order, so this is an inference-only opt-in. | | `FASTVIDEO_VSA_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel; grad-tracking forwards need its autograd operator, else Triton. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. | | `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | | `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index d8a4d91d4..2ad6edfe0 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -114,6 +114,24 @@ the Preview checkpoint's sparse VSA blocks. Packed-varlen changes floating-point reduction order relative to fixed-length FA4, so treat it as a speed/quality evaluation option rather than an exact-parity mode. +### cuDNN SDPA for H3 dense attention (opt-in) + +With the `FLASH_ATTN` backend, MiniMax-H3's dense, unmasked, inference-only DiT +self-attention can use cuDNN's fused SDPA engine instead of FlashAttention: + +```bash +export FASTVIDEO_MINIMAX_H3_DENSE_ATTN=cudnn +``` + +It keeps the same precision contract (bf16 operands, fp32 softmax/accumulation, +bf16 output) but uses a different reduction order, so treat it as a +speed/quality option rather than an exact-parity mode. Training calls, compiled +regions, head sizes above 128, non-contiguous inputs, and the sparse VSA blocks +stay on their existing paths. Standalone kernel timings at 37.7k and 109k tokens +are within measurement noise of FA4; an in-model A/B on one GB200 measured a +2.2% faster DiT forward. Requires CUDA and a torch build with +`aten::_scaled_dot_product_cudnn_attention` (cuDNN 9.x). + ### FP4 Flash Attention 4 (Blackwell only) **`FLASH_ATTN`** with **`--nvfp4_fa4`** diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index eeb02fb27..993b5b4a0 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -6,6 +6,7 @@ import torch.nn.functional as F from dataclasses import dataclass import fastvideo.envs as envs +from fastvideo.attention.utils.cudnn_sdpa import cudnn_sdpa_bshd, cudnn_sdpa_supported from fastvideo.attention.utils.flash_attn_default import ( fa_version, flash_attn_func_compilable, @@ -290,6 +291,13 @@ class FlashAttentionImpl(AttentionImpl): # identical to before the knob; kwargs win over the # FASTVIDEO_FA4_PV_MODE env bridge (same pattern as nvfp4_fa4 above). self.fa4_pv_mode = _resolve_fa4_pv_mode(extra_impl_args) + # MiniMax-H3 dense attention may opt into cuDNN's fused SDPA. It is an + # explicit impl argument (threaded from FASTVIDEO_MINIMAX_H3_DENSE_ATTN + # by the H3 model) so no other model can inherit it by accident. + self.cudnn_sdpa = bool(extra_impl_args.get("cudnn_sdpa", False)) + if self.cudnn_sdpa and self.nvfp4_fa4: + raise ValueError("cudnn_sdpa and nvfp4_fa4 select different Q/K precision; enable only one of " + "FASTVIDEO_MINIMAX_H3_DENSE_ATTN=cudnn and FASTVIDEO_NVFP4_FA4=1") if self.nvfp4_fa4: cap = torch.cuda.get_device_capability() assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}") @@ -484,6 +492,12 @@ class FlashAttentionImpl(AttentionImpl): elif self.nvfp4_fa4: output = self._forward_nvfp4(query, key, value) + elif self.cudnn_sdpa and cudnn_sdpa_supported(query, key, value): + # H3's dense, unmasked, inference-only route. The helper presents + # the BSHD inputs to cuDNN as BHSD strided views (contiguous BSHD + # is silently misread by the ATen op). + output = cudnn_sdpa_bshd(query, key, value, scale=self.softmax_scale, causal=self.causal) + elif (self.fa4_packed_varlen and fa_version == "4" and not torch.is_grad_enabled() and query.shape[0] == 1 and query.shape[1] == key.shape[1] == value.shape[1]): # FA4's packed-varlen entry point is materially faster for H3's diff --git a/fastvideo/attention/utils/cudnn_sdpa.py b/fastvideo/attention/utils/cudnn_sdpa.py new file mode 100644 index 000000000..9af72430c --- /dev/null +++ b/fastvideo/attention/utils/cudnn_sdpa.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""cuDNN fused scaled-dot-product attention for dense ``[B, S, H, D]`` tensors. + +Drop-in alternative to ``flash_attn_func`` for the no-mask, inference-only dense +self-attention call. It keeps the same precision contract (bf16/fp16 operands, +fp32 softmax/accumulation, bf16/fp16 output) but uses cuDNN's attention engine, +which has a different reduction order than FlashAttention. + +The ATen entry point has no shape metadata: it always reads its inputs as +``[B, H, S, D]``. Handing it the caller's contiguous ``[B, S, H, D]`` tensors is +silently misread -- when ``H`` and ``S`` swap into the same output shape the +call even returns a correctly-shaped tensor -- so every call here goes through +``transpose(1, 2)`` to present a BHSD *strided view*, exactly like the +documented PyTorch usage. Never pass contiguous BSHD to the op directly. +""" + +from __future__ import annotations + +import torch + +# cuDNN's fused attention kernels take fp16/bf16 only; the H3 DiT runs bf16. +_CUDNN_SDPA_DTYPES = (torch.float16, torch.bfloat16) +# cuDNN 9.x supports head sizes up to 128 on the fused SDPA path. Wider heads +# stay on FlashAttention instead of risking an engine-selection failure. +_CUDNN_SDPA_MAX_HEAD_DIM = 128 + +# Older torch builds may not register the op at all; keep the module importable +# (the guard then reports unsupported and the backend stays on FlashAttention). +_ATEN_CUDNN_SDPA = getattr(torch.ops.aten, "_scaled_dot_product_cudnn_attention", None) +if _ATEN_CUDNN_SDPA is not None: + _ATEN_CUDNN_SDPA = _ATEN_CUDNN_SDPA.default + + +def cudnn_sdpa_supported( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> bool: + """Whether the dense cuDNN route can take this Q/K/V triple unchanged. + + Grad-enabled calls stay on the training path, and calls inside a compiled + region fall through to the backend's traceable FlashAttention route so the + opt-in never changes what ``torch.compile`` can capture. + """ + if _ATEN_CUDNN_SDPA is None: + return False + if torch.is_grad_enabled() or torch.compiler.is_compiling(): + return False + if not (query.is_cuda and key.is_cuda and value.is_cuda): + return False + if query.dtype not in _CUDNN_SDPA_DTYPES or key.dtype != query.dtype or value.dtype != query.dtype: + return False + if query.ndim != 4 or key.shape != query.shape or value.shape != query.shape: + return False + if query.shape[-1] > _CUDNN_SDPA_MAX_HEAD_DIM: + return False + # The BHSD view below only describes the intended (B, S, H, D) layout when + # the caller's heads are head-dim contiguous; require that explicitly. + return query.is_contiguous() and key.is_contiguous() and value.is_contiguous() + + +def cudnn_sdpa_bshd( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + scale: float, + causal: bool = False, +) -> torch.Tensor: + """Dense scaled-dot-product attention on contiguous ``[B, S, H, D]`` inputs.""" + if _ATEN_CUDNN_SDPA is None: + raise RuntimeError("torch build does not provide aten::_scaled_dot_product_cudnn_attention") + # compute_log_sumexp=False: inference needs only the forward result. + output = _ATEN_CUDNN_SDPA( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + None, + False, + 0.0, + causal, + False, + scale=scale, + )[0] + # cuDNN returns [B, H, S, D]; restore the caller's [B, S, H, D] contract. + return output.transpose(1, 2).contiguous() diff --git a/fastvideo/envs.py b/fastvideo/envs.py index dc8a99563..9414af2ef 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -346,6 +346,13 @@ FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN = EnvBool( category="attention", doc="MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the " "floating-point reduction order, so it is an inference-only opt-in.") +FASTVIDEO_MINIMAX_H3_DENSE_ATTN = EnvChoice( + "flash", + choices=("flash", "cudnn"), + category="attention", + doc="Kernel for MiniMax-H3 dense DiT self-attention: flash keeps the selected FlashAttention backend; with " + "FLASH_ATTN, cudnn sends the no-mask, no-grad dense call through cuDNN's fused SDPA. cuDNN has a different " + "reduction order, so this is an inference-only opt-in.") FASTVIDEO_VSA_PTX_BLACKWELL = EnvBool( False, category="attention", diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index faa70ce60..31df2e0be 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -205,6 +205,7 @@ class MiniMaxH3Attention(nn.Module): prefix: str, fuse_qknorm_rope: bool = False, fa4_packed_varlen: bool = False, + cudnn_sdpa: bool = False, exact_rope: bool = False, vsa_fp4: dict[str, Any] | None = None, ) -> None: @@ -261,6 +262,7 @@ class MiniMaxH3Attention(nn.Module): supported_attention_backends=supported_attention_backends, prefix=prefix, fa4_packed_varlen=fa4_packed_varlen, + cudnn_sdpa=cudnn_sdpa, ) # Opt-in inference route: VSA-H3 selection on the block-sparse FP4 # kernel (see minimax_h3_vsa_fp4); grad and compile keep the generic path. @@ -656,6 +658,7 @@ class MiniMaxH3TransformerBlock(nn.Module): fuse_qknorm_rope: bool = False, fuse_swiglu: bool = False, fa4_packed_varlen: bool = False, + cudnn_sdpa: bool = False, exact_kernels: frozenset[str] = frozenset(), vsa_fp4: dict[str, Any] | None = None, ) -> None: @@ -671,6 +674,7 @@ class MiniMaxH3TransformerBlock(nn.Module): prefix=f"{prefix}.attn", fuse_qknorm_rope=fuse_qknorm_rope, fa4_packed_varlen=fa4_packed_varlen, + cudnn_sdpa=cudnn_sdpa, exact_rope="rope" in exact_kernels, vsa_fp4=vsa_fp4, ) @@ -933,6 +937,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions, fuse_swiglu="swiglu" in self.enabled_fusions, fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(), + cudnn_sdpa=envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() == "cudnn", exact_kernels=self.enabled_exact_kernels, vsa_fp4=config.vsa_fp4, ) for index in range(arch.num_layers) diff --git a/fastvideo/tests/attention/test_cudnn_sdpa_dense.py b/fastvideo/tests/attention/test_cudnn_sdpa_dense.py new file mode 100644 index 000000000..d3c369bbe --- /dev/null +++ b/fastvideo/tests/attention/test_cudnn_sdpa_dense.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-H3's opt-in cuDNN SDPA route for dense self-attention. + +The ATen cuDNN attention op carries no shape metadata and always reads its +inputs as ``[B, H, S, D]``; a contiguous ``[B, S, H, D]`` tensor is silently +misread. These tests pin the BHSD-strided-view contract and the opt-in flag. +""" + +import math + +import pytest +import torch + +from fastvideo import envs +from fastvideo.attention.utils.cudnn_sdpa import cudnn_sdpa_bshd, cudnn_sdpa_supported + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuDNN SDPA requires CUDA") + + +def _qkv(seq_len: int, *, heads: int = 8, head_dim: int = 64, dtype=torch.bfloat16): + generator = torch.Generator(device="cuda").manual_seed(0) + return [ + torch.randn(1, seq_len, heads, head_dim, device="cuda", dtype=dtype, generator=generator) for _ in range(3) + ] + + +def _rel_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + return ((actual.float() - expected.float()).norm() / expected.float().norm()).item() + + +def test_cudnn_route_presents_bhsd_strided_views(): + """Contiguous BSHD is silently misread; the BHSD view is the supported form.""" + query, key, value = _qkv(257) + scale = 1.0 / math.sqrt(query.shape[-1]) + with torch.no_grad(): + output = cudnn_sdpa_bshd(query, key, value, scale=scale) + # The trap: hand the op the caller's contiguous BSHD tensors directly. + misread = torch.ops.aten._scaled_dot_product_cudnn_attention.default( + query, key, value, None, False, 0.0, False, False, scale=scale)[0] + + # Same shape survives the misread, which is why it is silent. + assert misread.shape == output.shape + assert _rel_l2(misread, output) > 1.0 + assert output.is_contiguous() + assert output.dtype is torch.bfloat16 + assert torch.isfinite(output).all() + + +def test_cudnn_route_is_bit_identical_to_cudnn_dispatched_sdpa(): + query, key, value = _qkv(2048) + scale = 1.0 / math.sqrt(query.shape[-1]) + previous = torch.backends.cuda.cudnn_sdp_enabled() + torch.backends.cuda.enable_cudnn_sdp(True) + try: + with torch.no_grad(): + reference = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), scale=scale).transpose(1, 2) + finally: + torch.backends.cuda.enable_cudnn_sdp(previous) + + with torch.no_grad(): + output = cudnn_sdpa_bshd(query, key, value, scale=scale) + assert torch.equal(output, reference.contiguous()) + + +def test_cudnn_route_tracks_eager_math_sdpa(): + """Same precision contract (bf16 in, fp32 softmax, bf16 out), different reduction order.""" + query, key, value = _qkv(257) + scale = 1.0 / math.sqrt(query.shape[-1]) + with torch.no_grad(): + output = cudnn_sdpa_bshd(query, key, value, scale=scale) + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + reference = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), scale=scale).transpose(1, 2) + assert _rel_l2(output, reference) < 5e-3 + assert (output.float() - reference.float()).abs().max().item() < 0.02 + + +def test_cudnn_route_guard_rejects_training_unsupported_dtypes_and_shapes(): + query, key, value = _qkv(64, heads=4, head_dim=128) + with torch.no_grad(): + assert cudnn_sdpa_supported(query, key, value) + assert not cudnn_sdpa_supported(query.float(), key.float(), value.float()) + assert not cudnn_sdpa_supported(query.transpose(1, 2).transpose(1, 2).transpose(1, 2), key, value) + wide, _, _ = _qkv(64, heads=2, head_dim=256) + assert not cudnn_sdpa_supported(wide, wide, wide) + # A grad-enabled training call stays on the FlashAttention path. + assert torch.is_grad_enabled() + assert not cudnn_sdpa_supported(query, key, value) + + +def test_dense_attention_env_default_is_flash_and_accepts_cudnn(): + assert envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() == "flash" + with envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.override("cudnn"): + assert envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() == "cudnn" + with pytest.raises(ValueError): + with envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.override("cutlass"): + envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() + + +def test_flash_attn_impl_routes_dense_calls_to_cudnn(): + pytest.importorskip("flash_attn") + from fastvideo.attention.backends.flash_attn import FlashAttentionImpl + + query, key, value = _qkv(257) + scale = 1.0 / math.sqrt(query.shape[-1]) + impl = FlashAttentionImpl(num_heads=query.shape[2], + head_size=query.shape[3], + causal=False, + softmax_scale=scale, + cudnn_sdpa=True) + assert impl.cudnn_sdpa + with torch.no_grad(): + routed = impl.forward(query, key, value, None) + expected = cudnn_sdpa_bshd(query, key, value, scale=scale) + assert torch.equal(routed, expected) + + # The cuDNN opt-in wins over FA4 packed-varlen when both are configured. + packed = FlashAttentionImpl(num_heads=query.shape[2], + head_size=query.shape[3], + causal=False, + softmax_scale=scale, + cudnn_sdpa=True, + fa4_packed_varlen=True) + with torch.no_grad(): + assert torch.equal(packed.forward(query, key, value, None), expected) + + # Default construction keeps the existing FlashAttention behavior. + default = FlashAttentionImpl(num_heads=query.shape[2], + head_size=query.shape[3], + causal=False, + softmax_scale=scale) + assert not default.cudnn_sdpa + + +def test_flash_attn_impl_rejects_cudnn_and_nvfp4_together(): + pytest.importorskip("flash_attn") + from fastvideo.attention.backends.flash_attn import FlashAttentionImpl + + with pytest.raises(ValueError, match="cudnn_sdpa and nvfp4_fa4"): + FlashAttentionImpl(num_heads=4, head_size=64, causal=False, softmax_scale=0.125, cudnn_sdpa=True, + nvfp4_fa4=True)