Add FA4 monkey-patch path for low-precision attention

Adds FlashAttention 4 backend support for the monkey-patch SDPA path.
FA4 supports both Hopper (SM 9.x) and Blackwell (SM 10.x) hardware
with the flash_attn.cute.interface package.

Key additions:
- FP8_FA4 enum value in AttentionBackend
- _is_blackwell(), _is_fa4_available() hardware/library checks
- FA4 dispatch in apply_low_precision_attention
- fp8_fa4/ directory with fp8_fa4_sdpa entry point
- FA4 backend config in test suite with eager probe guard
- RoPE fusion placeholder (fuse_rope=True raises NotImplementedError)

ghstack-source-id: 05d86df
Pull-Request: pytorch#3960

Author: howardzhang-cv

test/prototype/attention/test_fp8_attention.py

@@ -5,11 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 
 """
-Tests for FP8 low-precision attention (FA3 backend).
+Tests for FP8 low-precision attention (FA3 and FA4 backends).
 
-Tests are gated on Hopper (SM 9.x) with flash-attn installed.
-When the backend is not available on the current hardware, tests are
-automatically skipped.
+Tests are parametrized over available backends. On Hopper (SM 9.x) with
+flash-attn installed, FA3 tests run. On Hopper or Blackwell (SM 10.x)
+with flash_attn.cute.interface installed, FA4 tests run. Backends that
+are not available on the current hardware are automatically skipped.
 """
 
 import unittest
@@ -42,7 +43,9 @@
     apply_low_precision_attention,
 )
 from torchao.prototype.attention.utils import (
+    _is_blackwell,
     _is_fa3_available,
+    _is_fa4_available,
     _is_hopper,
 )
 
@@ -55,10 +58,10 @@ class BackendConfig:
     """Configuration for a single backend under test."""
 
     name: str
-    flash_impl: str  # "FA3"
+    flash_impl: str  # "FA3" or "FA4"
     attention_backend: AttentionBackend
     sdpa_fn: Callable  # fp8_fa3_sdpa
-    rope_sdpa_fn: Callable  # fp8_fa3_rope_sdpa
+    rope_sdpa_fn: Callable  # fp8_fa3_rope_sdpa, or None if not yet available
     available_eager: bool  # Can run direct sdpa calls
     available_compiled: bool  # Can run via apply_low_precision_attention
     skip_msg: str
@@ -121,6 +124,38 @@ def _build_backend_configs() -> List[BackendConfig]:
         )
     )
 
+    # FA4: Hopper or Blackwell
+    fa4_available = (
+        _TORCH_VERSION_AT_LEAST_2_11
+        and (_is_hopper() or _is_blackwell())
+        and _is_fa4_available()
+    )
+    if fa4_available:
+        from torchao.prototype.attention.fp8_fa4.attention import fp8_fa4_sdpa
+
+        sdpa_fn = fp8_fa4_sdpa
+        eager_ok = _probe_eager_quantized_sdpa(sdpa_fn, "FA4")
+    else:
+        sdpa_fn = None
+        eager_ok = False
+
+    configs.append(
+        BackendConfig(
+            name="FA4",
+            flash_impl="FA4",
+            attention_backend=AttentionBackend.FP8_FA4,
+            sdpa_fn=sdpa_fn,
+            rope_sdpa_fn=None,  # FA4 rope not yet available
+            available_eager=eager_ok,
+            available_compiled=eager_ok,
+            skip_msg=(
+                "FP8 FA4 requires Hopper (SM 9.x) or Blackwell (SM 10.x), "
+                "flash-attn with FA4 support installed, "
+                "and PyTorch with flash activation APIs"
+            ),
+        )
+    )
+
     return configs
 
 
@@ -288,6 +323,8 @@ def test_rope_sdpa_accuracy(self, shape, dtype):
             )
 
         for backend in _EAGER_BACKENDS:
+            if backend.rope_sdpa_fn is None:
+                continue  # Backend doesn't support fused RoPE yet
             self._activate(backend)
             with torch.no_grad():
                 out_fp8 = backend.rope_sdpa_fn(q, k, v, cos, sin, is_causal=False)

torchao/prototype/attention/api.py

@@ -97,4 +97,9 @@ def apply_low_precision_attention(
 
         return setup_fp8_fa3(model, config)
 
+    if backend == AttentionBackend.FP8_FA4:
+        from torchao.prototype.attention.fp8_fa4.setup import setup_fp8_fa4
+
+        return setup_fp8_fa4(model, config)
+
     raise ValueError(f"Unknown backend: {backend}")

torchao/prototype/attention/config.py

@@ -23,6 +23,9 @@ class AttentionBackend(str, Enum):
     FP8_FA3 = "FP8_FA3"
     """FlashAttention 3 via PyTorch core. Requires SM90+ (Hopper)."""
 
+    FP8_FA4 = "FP8_FA4"
+    """FlashAttention 4 via PyTorch core. Requires SM90+ (Hopper) or SM100+ (Blackwell)."""
+
 
 @dataclass
 class LowPrecisionAttentionConfig:

torchao/prototype/attention/fp8_fa4/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FP8 attention implementation using FA4 backend.
+
+Use apply_low_precision_attention() from torchao.prototype.attention as the public API.
+For lower-level access, use fp8_fa4_sdpa() directly.
+"""
+
+from torchao.prototype.attention.fp8_fa4.attention import fp8_fa4_sdpa
+from torchao.prototype.attention.quantization import _fp8_sdpa_quantize
+
+__all__ = [
+    "fp8_fa4_sdpa",
+    "_fp8_sdpa_quantize",
+]

torchao/prototype/attention/fp8_fa4/attention.py

@@ -0,0 +1,44 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FP8 scaled dot-product attention using FA4 backend.
+
+This is a thin wrapper around the shared implementation in
+``shared_utils/attention.py``.  It exists so that the FA4 backend has
+a named entry point (``fp8_fa4_sdpa``) and backend-specific error messages.
+
+.. important::
+
+    When using this function directly (not through
+    ``apply_low_precision_attention``), you **must** activate the FA4
+    flash attention implementation yourself::
+
+        from torch.nn.attention import (
+            activate_flash_attention_impl,
+            restore_flash_attention_impl,
+        )
+
+        activate_flash_attention_impl("FA4")
+        try:
+            out = fp8_fa4_sdpa(q, k, v, is_causal=True)
+        finally:
+            restore_flash_attention_impl()
+
+    The high-level ``apply_low_precision_attention`` API handles this
+    automatically.
+"""
+
+from functools import partial
+
+from torchao.prototype.attention.shared_utils.attention import (
+    _fp8_sdpa,
+)
+
+fp8_fa4_sdpa = partial(_fp8_sdpa, backend_name="FA4")
+fp8_fa4_sdpa.__doc__ = _fp8_sdpa.__doc__
+fp8_fa4_sdpa.__name__ = "fp8_fa4_sdpa"
+fp8_fa4_sdpa.__qualname__ = "fp8_fa4_sdpa"

torchao/prototype/attention/fp8_fa4/setup.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FP8 FA4 backend setup.
+
+Thin wrapper around the shared ``setup_fp8_backend``, binding the FA4
+attention function.
+"""
+
+import torch.nn as nn
+
+from torchao.prototype.attention.config import LowPrecisionAttentionConfig
+from torchao.prototype.attention.shared_utils.setup import setup_fp8_backend
+
+
+def _compile_not_available(model, config):
+    raise NotImplementedError(
+        "FA4 RoPE fusion (fuse_rope=True) is not yet available. "
+        "Use fuse_rope=False (default) for the monkey-patch path."
+    )
+
+
+def setup_fp8_fa4(
+    model: nn.Module,
+    config: LowPrecisionAttentionConfig,
+) -> nn.Module:
+    """Set up FP8 FA4 attention on *model* and wrap it."""
+    from torchao.prototype.attention.fp8_fa4.attention import fp8_fa4_sdpa
+
+    return setup_fp8_backend(
+        model,
+        config,
+        flash_impl_name="FA4",
+        sdpa_fn=fp8_fa4_sdpa,
+        compile_fn=_compile_not_available,
+    )

torchao/prototype/attention/utils.py

@@ -25,6 +25,16 @@ def _is_hopper() -> bool:
     return major == 9
 
 
+def _is_blackwell() -> bool:
+    """
+    Check if the current CUDA device is Blackwell (SM 10.x).
+    """
+    if not torch.cuda.is_available():
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major == 10
+
+
 def _is_fa3_available() -> bool:
     """
     Check if the flash attention 3 library (flash_attn_interface) is installed.
@@ -36,6 +46,17 @@ def _is_fa3_available() -> bool:
         return False
 
 
+def _is_fa4_available() -> bool:
+    """
+    Check if the flash attention 4 library (flash_attn.cute.interface) is installed.
+    """
+    try:
+        importlib.import_module("flash_attn.cute.interface")
+        return True
+    except ModuleNotFoundError:
+        return False
+
+
 def _get_available_backend() -> AttentionBackend:
     """
     Get the best available backend for current hardware.
@@ -51,10 +72,18 @@ def _get_available_backend() -> AttentionBackend:
 
     capability = torch.cuda.get_device_capability()
 
+    # FA4 on Blackwell (SM 10.x) with flash_attn.cute.interface
+    if _is_blackwell() and _is_fa4_available():
+        return AttentionBackend.FP8_FA4
+
     # FA3 requires exactly Hopper (SM 9.x) and flash_attn_interface
     if _is_hopper() and _is_fa3_available():
         return AttentionBackend.FP8_FA3
 
+    # FA4 also supports Hopper (SM 9.x) with flash_attn.cute.interface
+    if _is_hopper() and _is_fa4_available():
+        return AttentionBackend.FP8_FA4
+
     raise RuntimeError(f"No compatible backend for SM{capability[0]}{capability[1]}.")
 
 
@@ -84,5 +113,22 @@ def _check_backend_available(backend: AttentionBackend) -> None:
                 "FP8_FA3 backend requires the flash-attn package with FA3 support. "
             )
 
+    elif backend == AttentionBackend.FP8_FA4:
+        if not torch.cuda.is_available():
+            raise RuntimeError("FP8_FA4 backend requires CUDA.")
+
+        if not (_is_hopper() or _is_blackwell()):
+            capability = torch.cuda.get_device_capability()
+            raise RuntimeError(
+                f"FP8_FA4 backend requires Hopper (SM 9.x) or Blackwell (SM 10.x). "
+                f"Current device: SM{capability[0]}{capability[1]}. "
+            )
+
+        if not _is_fa4_available():
+            raise RuntimeError(
+                "FP8_FA4 backend requires the flash-attn package with FA4 support "
+                "(flash_attn.cute.interface). "
+            )
+
     else:
         raise ValueError(f"Unknown backend: {backend}")