[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.
2026-11-01 09:14:02 -08:00
parent f73735f491
commit 2ac6336b57
7 changed files with 273 additions and 0 deletions
@@ -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.                                                                                                                                                                                                                                                                                                                      |
@@ -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`**
@@ -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
@@ -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()
@@ -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",
@@ -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)
@@ -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)