diff --git a/pyproject.toml b/pyproject.toml index ffdac15..b8c7112 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,15 @@ classifiers = [ "Operating System :: OS Independent", ] +[project.optional-dependencies] +# Selection stays a runtime decision via RING_FLASH_ATTN_BACKEND; these extras only +# document which distribution provides each backend. Neither is on PyPI as of this +# writing -- both are built from a flash-attention checkout: +# fa3: `pip install -e /hopper` (needs CUDA_HOME; sm90a by default) +# fa4: `pip install -e /flash_attn/cute` +fa3 = ["flash-attn-3"] +fa4 = ["flash-attn-4"] + [project.urls] Homepage = "https://github.com/zhuzilin/ring-flash-attention" Issues = "https://github.com/zhuzilin/ring-flash-attention/issues" diff --git a/ring_flash_attn/flash_attn_backend.py b/ring_flash_attn/flash_attn_backend.py new file mode 100644 index 0000000..8527c52 --- /dev/null +++ b/ring_flash_attn/flash_attn_backend.py @@ -0,0 +1,1064 @@ +"""Backend seam between the ring attention variants and the flash-attn kernels. + +The ring variants only need four primitives: a block forward that returns +``(out, lse)`` and a block backward that writes into caller-supplied ``dq``/``dk``/``dv``, +each in a dense and a varlen flavour. This module normalises those four onto the +FA2 keyword set so the variant modules contain no version- or backend-specific +branching. + +Three backends are available: + +``fa2`` + ``flash_attn.flash_attn_interface`` (FlashAttention 2.x). The default. Routes + through ``torch.ops.flash_attn.*`` where the installed versions register it, so + the dense path stays traceable under ``torch.compile``. + +``fa3`` + ``flash_attn_3`` (FlashAttention 3). Opt in with ``RING_FLASH_ATTN_BACKEND=fa3``. + Keeps FA2's LSE layout, in-place ``dq``/``dk``/``dv`` contract, and ``-1`` window + sentinel, and registers real ``torch.library`` custom ops, so it traces under + ``torch.compile``. No dropout and no ALiBi. Dense and varlen share one op per + direction. Its backward renames ``causal`` to ``is_causal`` and orders + ``cu_seqlens_*`` before ``dq``/``dk``/``dv``, so calls here are keyword-only. + +``fa4`` + ``flash_attn.cute.interface`` (FlashAttention 4, distributed as ``flash-attn-4``). + Opt in with ``RING_FLASH_ATTN_BACKEND=fa4``. No dropout, no ALiBi, and no + ``torch.library`` registration upstream, so it is not traceable under + ``torch.compile``. Its backward additionally requires compute capability 9.0+. +""" + +import importlib +import importlib.metadata +import importlib.util +import os +from pathlib import Path +from typing import Optional, Tuple + +import torch + +from .utils import get_default_args + +__all__ = [ + "BACKEND", + "CAPABILITIES", + "Capabilities", + "available_backends", + "check_variant_supported", + "fa_backward", + "fa_forward", + "fa_varlen_backward", + "fa_varlen_forward", + "is_fa3_available", + "is_fa4_available", +] + +_ENV_VAR = "RING_FLASH_ATTN_BACKEND" +_VALID_BACKENDS = ("fa2", "fa3", "fa4") +_FA4_INSTALL_HINT = ( + "install it with `pip install -e /flash_attn/cute`" +) +_FA3_INSTALL_HINT = ( + "build it with `pip install -e /hopper` (needs CUDA_HOME " + "and compiles for sm90a by default)" +) + +# FA3's interface module has moved between releases. Tried in this order; the one that +# resolves is recorded in FA3_MODULE so errors can name it. +# +# `hopper.flash_attn_interface` is only reachable if someone puts a source checkout on +# the path *and* has a built flash_attn_3._C from elsewhere -- the vendored hopper/ trees +# that appear in site-packages have no extension, so is_fa3_available() rejects them +# before this list is consulted. Kept as a last resort rather than a supported layout. +_FA3_MODULE_CANDIDATES = ( + "flash_attn_3.flash_attn_interface", + "flash_attn_interface", + "hopper.flash_attn_interface", +) + + +def _flash_attn_search_locations() -> Tuple[str, ...]: + """Locate the ``flash_attn`` package directories without importing it. + + Deliberately avoids ``find_spec("flash_attn.cute.interface")``: resolving a + *submodule* spec imports the parent package, which is precisely what fails when + the FA2 extension module is broken. ``find_spec("flash_attn")`` only locates the + package and does not execute its ``__init__.py``. + """ + try: + spec = importlib.util.find_spec("flash_attn") + except Exception: + # A broken parent package can make even locating it fail. + return () + if spec is None or not spec.submodule_search_locations: + return () + return tuple(spec.submodule_search_locations) + + +def _fa4_distribution_installed() -> bool: + """Whether a flash-attn-4 distribution is registered with the installer. + + FA4's ``pyproject.toml`` names the distribution ``flash-attn-4`` while its + ``cute/__init__.py`` looks up ``version("fa4")``, so both spellings are checked. + Name lookup is normalised, so ``flash_attn_4`` matches ``flash-attn-4``. + """ + for name in ("flash-attn-4", "fa4"): + try: + importlib.metadata.distribution(name) + return True + except Exception: + continue + return False + + +def is_fa4_available() -> bool: + """Return whether flash-attn-4 looks importable, without importing it. + + A real import is unsuitable here -- FA4 pulls in ``nvidia-cutlass-dsl`` and + ``quack-kernels`` and is slow to load -- so this works off packaging metadata and + the filesystem. + + Note the presence of ``flash_attn/cute/`` alone is *not* sufficient: flash-attn + 2.8.x bundles an older, unrelated ``cute`` module (its own ``__version__`` is + "0.1.0") whose ``_flash_attn_fwd`` has a different signature. Selecting it as FA4 + would fail at call time with a confusing TypeError. FA4 proper is distinguished by + a registered distribution, or failing that by ``flash_fwd_combine.py``, which the + bundled 2.8.x copy does not ship. + """ + # Metadata first: an editable install (`pip install -e .../flash_attn/cute`) routes + # through a path hook and never materialises a `cute/` directory next to flash_attn, + # so a filesystem-only probe misses it. + if _fa4_distribution_installed(): + return True + # Fall back to the filesystem for a vendored or PYTHONPATH checkout with no metadata. + return any( + (Path(location) / "cute" / "interface.py").is_file() + and (Path(location) / "cute" / "flash_fwd_combine.py").is_file() + for location in _flash_attn_search_locations() + ) + + +def is_fa3_available() -> bool: + """Return whether flash-attn-3 looks importable, without importing its interface. + + Probes the compiled extension ``flash_attn_3._C`` rather than the interface module, + because the interface alone is not enough: several environments here carry a raw + ``hopper/`` directory copied into site-packages with no extension built, so an + interface-only probe reports available and then dies at the first kernel call. + + ``find_spec("flash_attn_3._C")`` is safe in a way the FA4 equivalent was not -- + ``flash_attn_3/__init__.py`` does not import the interface, so resolving this + submodule does not pull in the kernels. We deliberately do *not* probe + ``flash_attn_3.flash_attn_interface``, which would execute that module's + ``import flash_attn_3._C`` at module scope. + """ + try: + return importlib.util.find_spec("flash_attn_3._C") is not None + except Exception: + # A broken or partially-installed parent package can make even locating it fail. + return False + + +def available_backends() -> Tuple[str, ...]: + """Backends usable in this environment, for tests to skip on. + + ``fa2`` is reported only if it genuinely imports -- an installed-but-broken + flash-attn (e.g. built against a different torch) is not "available". That import + is free when fa2 is already the active backend. + + ``fa4`` reflects installation, not a trial import, because importing it pulls in + ``nvidia-cutlass-dsl`` and ``quack-kernels``. Note FA4 shares the ``flash_attn`` + package namespace, so a broken flash-attn 2.x makes FA4 unimportable too even + when it is installed. + """ + found = [] + try: + importlib.import_module("flash_attn.flash_attn_interface") + found.append("fa2") + except Exception: + pass + if is_fa3_available(): + found.append("fa3") + if is_fa4_available(): + found.append("fa4") + return tuple(found) + + +class Capabilities: + """What the active backend actually supports. + + Every flag gates a real call-time check in this module. Properties that no ring + code can act on (such as ``torch.compile`` traceability) are documented rather + than expressed here, so nothing in this record can be mistaken for a runtime + guard that does not exist. + """ + + def __init__( + self, + name: str, + supports_dropout: bool, + supports_alibi: bool, + supports_softcap: bool, + supports_window: bool, + supports_backward: bool, + ): + self.name = name + self.supports_dropout = supports_dropout + self.supports_alibi = supports_alibi + self.supports_softcap = supports_softcap + self.supports_window = supports_window + self.supports_backward = supports_backward + + def __repr__(self) -> str: + return f"Capabilities(name={self.name!r})" + + +def _select_backend() -> str: + backend = os.environ.get(_ENV_VAR, "fa2").strip().lower() + if backend not in _VALID_BACKENDS: + raise ValueError( + f"{_ENV_VAR}={backend!r} is not a known backend. " + f"Expected one of {', '.join(_VALID_BACKENDS)}." + ) + return backend + + +BACKEND = _select_backend() + + +def _set_window(params: dict, window_size) -> None: + """Apply the window in whichever arity the resolved signature uses. + + Both FA2 and FA3 have shipped releases taking either a ``window_size`` tuple or the + two ints ``window_size_left``/``window_size_right``. ``params`` comes from signature + introspection, so the key-presence check is meaningful rather than a guess. + + Shared by the fa2 and fa3 adapters. Not used by fa4, whose schema needs ``None`` + rather than the ``-1`` sentinel -- see ``_split_window`` in that branch. + """ + if "window_size" in params: + params["window_size"] = window_size + else: + params["window_size_left"] = window_size[0] + params["window_size_right"] = window_size[1] + + +# --------------------------------------------------------------------------- +# fa2 +# --------------------------------------------------------------------------- + +if BACKEND == "fa2": + try: + import flash_attn + from flash_attn.flash_attn_interface import ( + _flash_attn_backward, + _flash_attn_forward, + _flash_attn_varlen_backward, + _flash_attn_varlen_forward, + ) + except Exception as e: + # Catch broadly: a flash_attn built against a different torch fails at + # extension-load time as ImportError (undefined symbol), OSError, or + # RuntimeError depending on the mismatch -- not ModuleNotFoundError. + _alternates = [] + if is_fa3_available(): + _alternates.append( + f"flash-attn-3 is installed -- set {_ENV_VAR}=fa3 to use it. It is a " + "separate `flash_attn_3` package, so it is unaffected by whatever is " + "wrong with flash-attn 2.x here." + ) + if is_fa4_available(): + _alternates.append( + f"flash-attn-4 is installed -- set {_ENV_VAR}=fa4 to use it. Note it " + "installs into the same `flash_attn` package namespace, so if the " + "failure above came from `flash_attn/__init__.py` itself, the broken " + "flash-attn 2.x must be repaired or uninstalled first." + ) + if _alternates: + extra = " ".join(_alternates) + else: + extra = ( + f"No other backend was found either -- for flash-attn-3 {_FA3_INSTALL_HINT}, " + f"or for flash-attn-4 {_FA4_INSTALL_HINT}." + ) + raise ImportError( + f"ring_flash_attn could not load the flash_attn (FA2) backend: {e}. {extra}" + ) from e + + # FA >= 2.7 registers the python entry points as torch library custom ops, which + # keeps them traceable under torch.compile. Preserved verbatim from the per-module + # gates this seam replaces, including the asymmetry that the backward gate does not + # consult flash_attn.__version__. + if torch.__version__ >= "2.4.0" and flash_attn.__version__ >= "2.7.0": + _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward + else: + _wrapped_flash_attn_forward = _flash_attn_forward + + if torch.__version__ >= "2.4.0": + _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward + else: + _wrapped_flash_attn_backward = _flash_attn_backward + + CAPABILITIES = Capabilities( + name="fa2", + supports_dropout=True, + supports_alibi=True, + supports_softcap=True, + supports_window=True, + supports_backward=True, + ) + + def _unpack_forward(outputs): + """FA <= 2.6 returns 8 values, FA >= 2.7 returns 4. We want out and lse.""" + if len(outputs) == 8: + return outputs[0], outputs[5] + assert len(outputs) == 4 + return outputs[0], outputs[1] + + def fa_forward( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + use_custom_op=True, + ): + fn = _wrapped_flash_attn_forward if use_custom_op else _flash_attn_forward + params = get_default_args(_flash_attn_forward).copy() + params.update( + { + "q": q, + "k": k, + "v": v, + "dropout_p": dropout_p, + "softmax_scale": softmax_scale, + "causal": causal, + "softcap": softcap, + "alibi_slopes": alibi_slopes, + "return_softmax": True and dropout_p > 0, + } + ) + _set_window(params, window_size) + return _unpack_forward(fn(**params)) + + def fa_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + deterministic=False, + use_custom_op=True, + ): + fn = _wrapped_flash_attn_backward if use_custom_op else _flash_attn_backward + params = get_default_args(_flash_attn_backward).copy() + params.update( + { + "dout": dout, + "q": q, + "k": k, + "v": v, + "out": out, + "softmax_lse": softmax_lse, + "dq": dq, + "dk": dk, + "dv": dv, + "dropout_p": dropout_p, + "softmax_scale": softmax_scale, + "causal": causal, + "softcap": softcap, + "alibi_slopes": alibi_slopes, + "deterministic": deterministic, + } + ) + _set_window(params, window_size) + fn(**params) + + def fa_varlen_forward( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + ): + params = get_default_args(_flash_attn_varlen_forward).copy() + params.update( + { + "q": q, + "k": k, + "v": v, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + "dropout_p": dropout_p, + "softmax_scale": softmax_scale, + "causal": causal, + "softcap": softcap, + "alibi_slopes": alibi_slopes, + "return_softmax": True and dropout_p > 0, + } + ) + _set_window(params, window_size) + return _unpack_forward(_flash_attn_varlen_forward(**params)) + + def fa_varlen_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + deterministic=False, + ): + params = get_default_args(_flash_attn_varlen_backward).copy() + params.update( + { + "dout": dout, + "q": q, + "k": k, + "v": v, + "out": out, + "softmax_lse": softmax_lse, + "dq": dq, + "dk": dk, + "dv": dv, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + "dropout_p": dropout_p, + "softmax_scale": softmax_scale, + "causal": causal, + "softcap": softcap, + "alibi_slopes": alibi_slopes, + "deterministic": deterministic, + } + ) + _set_window(params, window_size) + _flash_attn_varlen_backward(**params) + + +# --------------------------------------------------------------------------- +# fa3 +# --------------------------------------------------------------------------- + +elif BACKEND == "fa3": + if not is_fa3_available(): + raise ImportError( + f"{_ENV_VAR}=fa3 was requested but the flash_attn_3._C extension was not " + f"found. A bare `hopper/` source directory is not enough -- the CUDA " + f"extension must be built; {_FA3_INSTALL_HINT}." + ) + + FA3_MODULE = None + _fa3_errors = [] + for _candidate in _FA3_MODULE_CANDIDATES: + try: + _fa3 = importlib.import_module(_candidate) + FA3_MODULE = _candidate + break + except Exception as e: # noqa: PERF203 - we want the reason for each candidate + _fa3_errors.append(f"{_candidate}: {type(e).__name__}: {e}") + if FA3_MODULE is None: + raise ImportError( + "ring_flash_attn found flash_attn_3._C but could not import the flash-attn-3 " + "interface module. Tried:\n " + "\n ".join(_fa3_errors) + f"\n{_FA3_INSTALL_HINT}." + ) + + _flash_attn_forward = _fa3._flash_attn_forward + _flash_attn_backward = _fa3._flash_attn_backward + + # FA3 can compile out features ring attention depends on; those surface only as + # TORCH_CHECK failures deep in C++, so report them here instead. + # + # The config module is generated at build time (hopper/setup.py:89) and sits beside + # the interface, so look for it next to whichever module resolved above before + # falling back to the top-level name that setup.py's py_modules installs. + # + # Note the keys are FLASHATTENTION_DISABLE_* -- no underscore after FLASH -- while + # the env vars that set them are FLASH_ATTENTION_DISABLE_*. Using the env-var + # spelling here silently disables this whole check. + _fa3_disable_flags = { + "FLASHATTENTION_DISABLE_VARLEN": "the varlen ring paths", + "FLASHATTENTION_DISABLE_LOCAL": "sliding-window attention", + "FLASHATTENTION_DISABLE_BACKWARD": "training (backward)", + } + _fa3_config_candidates = [] + if "." in FA3_MODULE: + _fa3_config_candidates.append(FA3_MODULE.rsplit(".", 1)[0] + ".flash_attn_config") + _fa3_config_candidates.append("flash_attn_config") + + try: + _flags = None + for _cand in _fa3_config_candidates: + try: + _flags = importlib.import_module(_cand).CONFIG["build_flags"] + break + except Exception: + continue + # Values are real booleans (hopper/setup.py:51 computes them as `== "TRUE"`), + # so plain truthiness is correct here. + if isinstance(_flags, dict): + for _flag, _why in _fa3_disable_flags.items(): + if _flags.get(_flag): + raise ImportError( + f"this flash-attn-3 build was compiled with {_flag}=TRUE, which " + f"disables {_why}. Rebuild without it, or use {_ENV_VAR}=fa2." + ) + except ImportError: + raise + except Exception: + # The config module is optional and its shape is a build-time detail; an + # unrecognised layout means "no opinion", never a failed import. + pass + + CAPABILITIES = Capabilities( + name="fa3", + supports_dropout=False, + supports_alibi=False, + supports_softcap=True, + supports_window=True, + supports_backward=True, + ) + + # FA3 unifies dense and varlen into ONE op per direction -- there is no + # _flash_attn_varlen_forward to import. The varlen wrappers below call the same two + # ops, just supplying cu_seqlens_*/max_seqlen_*. + # + # Two shapes of drift are absorbed by introspecting the real signature: releases + # differ on window arity (`window_size` tuple vs `window_size_left`/`_right` ints) + # and on the backward causal kwarg (`causal` vs `is_causal`). get_default_args + # handles the CustomOpDef via its `_init_fn` fallback and is cached per function, + # so this costs nothing per ring step. + _FA3_BWD_CAUSAL_KEY = ( + "is_causal" if "is_causal" in get_default_args(_flash_attn_backward) else "causal" + ) + + def _fa3_forward_params( + q, k, v, softmax_scale, causal, window_size, softcap, dropout_p, alibi_slopes + ): + _check_unsupported(dropout_p, alibi_slopes) + params = get_default_args(_flash_attn_forward).copy() + params.update( + { + "q": q, + "k": k, + "v": v, + "softmax_scale": softmax_scale, + "causal": causal, + "softcap": softcap, + } + ) + _set_window(params, window_size) + return params + + def _fa3_backward_params( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + softmax_scale, + causal, + window_size, + softcap, + deterministic, + dropout_p, + alibi_slopes, + ): + _check_unsupported(dropout_p, alibi_slopes) + params = get_default_args(_flash_attn_backward).copy() + params.update( + { + "dout": dout, + "q": q, + "k": k, + "v": v, + "out": out, + "softmax_lse": softmax_lse, + # Must always be supplied: they are Optional, and the C++ silently + # allocates throwaway buffers when omitted, computing the gradient and + # discarding it with no error. + "dq": dq, + "dk": dk, + "dv": dv, + "softmax_scale": softmax_scale, + _FA3_BWD_CAUSAL_KEY: causal, + "softcap": softcap, + "deterministic": deterministic, + } + ) + _set_window(params, window_size) + return params + + def fa_forward( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + use_custom_op=True, # accepted and ignored: the FA3 symbols *are* the custom ops + ): + params = _fa3_forward_params( + q, k, v, softmax_scale, causal, window_size, softcap, dropout_p, alibi_slopes + ) + outputs = _flash_attn_forward(**params) + # (out, softmax_lse, out_accum, softmax_lse_accum); the accum pair is empty + # unless num_splits > 1, which we never set. + # + # `out_` is deliberately left None: FA3's fake impl raises on a preallocated + # output under tracing, so writing straight into a ring buffer here would cost + # torch.compile support. + return outputs[0], outputs[1] + + def fa_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + deterministic=False, + use_custom_op=True, # accepted and ignored: the FA3 symbols *are* the custom ops + ): + params = _fa3_backward_params( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + softmax_scale, + causal, + window_size, + softcap, + deterministic, + dropout_p, + alibi_slopes, + ) + # Writes into dq/dk/dv in place (mutates_args); returns softmax_d, unused here. + _flash_attn_backward(**params) + + def fa_varlen_forward( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + ): + params = _fa3_forward_params( + q, k, v, softmax_scale, causal, window_size, softcap, dropout_p, alibi_slopes + ) + params.update( + { + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + } + ) + outputs = _flash_attn_forward(**params) + return outputs[0], outputs[1] + + def fa_varlen_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + deterministic=False, + ): + params = _fa3_backward_params( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + softmax_scale, + causal, + window_size, + softcap, + deterministic, + dropout_p, + alibi_slopes, + ) + params.update( + { + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": max_seqlen_k, + } + ) + _flash_attn_backward(**params) + + +# --------------------------------------------------------------------------- +# fa4 +# --------------------------------------------------------------------------- + +elif BACKEND == "fa4": + if not is_fa4_available(): + raise ImportError( + f"{_ENV_VAR}=fa4 was requested but flash_attn.cute was not found; " + f"{_FA4_INSTALL_HINT}." + ) + try: + from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd + except Exception as e: + raise ImportError( + f"ring_flash_attn could not load the flash-attn-4 backend: {e}. " + "flash-attn-4 lives under the `flash_attn` package namespace, so a broken " + "flash-attn 2.x install in the same environment breaks it too -- if the " + "error above names flash_attn_2_cuda or flash_attn/__init__.py, repair or " + "uninstall flash-attn 2.x rather than reinstalling FA4. FA4 also requires " + f"nvidia-cutlass-dsl and quack-kernels; {_FA4_INSTALL_HINT}." + ) from e + + CAPABILITIES = Capabilities( + name="fa4", + supports_dropout=False, + supports_alibi=False, + supports_softcap=True, + supports_window=True, + supports_backward=True, # gated per-device, see _check_backward_arch + ) + + # Forward accepts sm_80 and up; backward is sm_90 and up. + _MIN_BACKWARD_MAJOR = 9 + + def _check_backward_arch(device) -> None: + major, minor = torch.cuda.get_device_capability(device) + if major < _MIN_BACKWARD_MAJOR: + raise RuntimeError( + f"flash-attn-4 backward requires compute capability " + f"{_MIN_BACKWARD_MAJOR}.0+, but this device is sm_{major}{minor}. " + "FA4 forward works on sm_80+; only the backward kernel is gated. " + "Use RING_FLASH_ATTN_BACKEND=fa2 to train on this device." + ) + + def _split_window(window_size) -> Tuple[Optional[int], Optional[int]]: + """Map FA2's -1 sentinel onto FA4's ``None``. + + FA4 reads these as literal bounds, so passing -1 through would be taken as a + one-column window rather than "unbounded". Its own ``(-1, -1)`` special case + only fires when *both* sides are negative, which is not general enough. + """ + left, right = window_size + return ( + None if left is None or left < 0 else left, + None if right is None or right < 0 else right, + ) + + def fa_forward( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + use_custom_op=True, # accepted and ignored: FA4 registers no torch ops + ): + _check_unsupported(dropout_p, alibi_slopes) + window_left, window_right = _split_window(window_size) + out, lse = _flash_attn_fwd( + q, + k, + v, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap, + window_size_left=window_left, + window_size_right=window_right, + # Required: FA4 returns lse=None unless asked, and autograd.Function.forward + # runs with grad disabled so the requires_grad shortcut never fires. + return_lse=True, + ) + return out, lse + + def fa_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + deterministic=False, + use_custom_op=True, # accepted and ignored: FA4 registers no torch ops + ): + _check_unsupported(dropout_p, alibi_slopes) + _check_backward_arch(q.device) + window_left, window_right = _split_window(window_size) + out_dq, out_dk, out_dv = _flash_attn_bwd( + q, + k, + v, + out, + dout, + softmax_lse, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap, + window_size_left=window_left, + window_size_right=window_right, + deterministic=deterministic, + dq=dq, + dk=dk, + dv=dv, + ) + _copy_back(dq, dk, dv, out_dq, out_dk, out_dv) + + def fa_varlen_forward( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + ): + _check_unsupported(dropout_p, alibi_slopes) + window_left, window_right = _split_window(window_size) + out, lse = _flash_attn_fwd( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap, + window_size_left=window_left, + window_size_right=window_right, + return_lse=True, + ) + return out, lse + + def fa_varlen_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq, + dk, + dv, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=0.0, + alibi_slopes=None, + deterministic=False, + ): + _check_unsupported(dropout_p, alibi_slopes) + _check_backward_arch(q.device) + window_left, window_right = _split_window(window_size) + out_dq, out_dk, out_dv = _flash_attn_bwd( + q, + k, + v, + out, + dout, + softmax_lse, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap, + window_size_left=window_left, + window_size_right=window_right, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + deterministic=deterministic, + dq=dq, + dk=dk, + dv=dv, + ) + _copy_back(dq, dk, dv, out_dq, out_dk, out_dv) + + def _copy_back(dq, dk, dv, out_dq, out_dk, out_dv) -> None: + """Honour FA2's write-into-the-caller's-buffer contract. + + FA4 writes into the tensors it is handed, so these copies are normally no-ops. + They matter only if a future FA4 version reallocates internally -- callers here + hold views into ring buffers (see llama3's ``dkv_buffer[0][local_k_slice]``) and + would silently lose the gradient. + """ + for dst, src in ((dq, out_dq), (dk, out_dk), (dv, out_dv)): + if dst is not None and src is not None and dst.data_ptr() != src.data_ptr(): + dst.copy_(src) + + +# --------------------------------------------------------------------------- +# Backend-agnostic guards (shared by all adapters, driven by CAPABILITIES) +# --------------------------------------------------------------------------- + + +def _check_unsupported(dropout_p, alibi_slopes) -> None: + """Fail loudly rather than silently dropping an argument the backend cannot honour. + + Reads CAPABILITIES so the flags are load-bearing rather than descriptive. Under fa2 + every flag is True and this is a no-op. + """ + if dropout_p and not CAPABILITIES.supports_dropout: + raise NotImplementedError( + f"the {CAPABILITIES.name} backend has no dropout support; got " + f"dropout_p={dropout_p}. Use {_ENV_VAR}=fa2 if you need it." + ) + if alibi_slopes is not None and not CAPABILITIES.supports_alibi: + raise NotImplementedError( + f"the {CAPABILITIES.name} backend has no ALiBi support; got a non-None " + f"alibi_slopes. Use {_ENV_VAR}=fa2 if you need it." + ) + + +# Ring variants that are in scope for each non-default backend. Everything else is +# routed through this seam but was never exercised against those kernels, so it is +# refused rather than silently run with unverified numerics. +# +# This marks *scope*, not proven correctness -- adding a name asserts only that the +# variant is intended to work there. fa2 is absent because it is the reference backend +# and every variant is in scope. +_VALIDATED_VARIANTS = { + "fa3": frozenset( + { + "llama3_flash_attn_varlen", + "llama_fwd_ring_bwd_flash_attn", + "ring_flash_attn_backward", + } + ), + "fa4": frozenset( + { + "llama3_flash_attn_varlen", + "llama_fwd_ring_bwd_flash_attn", + "ring_flash_attn_backward", + } + ), +} + +_BACKEND_DISPLAY_NAME = {"fa2": "flash-attn 2.x", "fa3": "flash-attn-3", "fa4": "flash-attn-4"} + + +def check_variant_supported(variant: str) -> None: + """Refuse ring variants that are out of scope for the active backend. + + Called by the variant modules deliberately left out of the fa3/fa4 scope. Under + fa2 (the default) this never fires, since fa2 has no entry in _VALIDATED_VARIANTS. + """ + in_scope = _VALIDATED_VARIANTS.get(BACKEND) + if in_scope is not None and variant not in in_scope: + raise NotImplementedError( + f"'{variant}' is not in scope for the " + f"{_BACKEND_DISPLAY_NAME.get(BACKEND, BACKEND)} backend, so it is refused " + f"rather than run with unverified numerics. In scope under {BACKEND}: " + f"{', '.join(sorted(in_scope))}. " + f"Unset {_ENV_VAR} (or set it to fa2) to use this variant." + ) diff --git a/ring_flash_attn/llama3_flash_attn_varlen.py b/ring_flash_attn/llama3_flash_attn_varlen.py index bbd5909..427a7dd 100644 --- a/ring_flash_attn/llama3_flash_attn_varlen.py +++ b/ring_flash_attn/llama3_flash_attn_varlen.py @@ -1,11 +1,8 @@ import torch import torch.distributed as dist -from flash_attn.flash_attn_interface import ( - _flash_attn_varlen_forward, - _flash_attn_varlen_backward, -) +from .flash_attn_backend import fa_varlen_forward, fa_varlen_backward import logging -from .utils import get_default_args, AllGatherComm as Comm +from .utils import AllGatherComm as Comm def llama3_flash_attn_prepare_cu_seqlens( @@ -121,37 +118,21 @@ def llama3_flash_attn_varlen_forward( v_i = kv_buffer[1][local_k_slice] # logging.debug(f"fwd i {i} k_ishape {k_i.shape} q.shape {q.shape} kv_buffer[0] {kv_buffer[0].shape} local_k_slice {local_k_slice}") - # params = get_default_args(_flash_attn_varlen_forward).copy() - params = { - "q": q_i, - "k": k_i, - "v": v_i, - "cu_seqlens_q": cu_seqlens_q, - "cu_seqlens_k": cu_seqlens_k, - "max_seqlen_q": max_seqlen_q, - "max_seqlen_k": max_seqlen_k, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "softcap": softcap, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - outputs = _flash_attn_varlen_forward(**params) - if len(outputs) == 8: - out, _, _, _, _, lse, _, _ = outputs - else: - assert len(outputs) == 4 - out, lse, _, _ = outputs + out, lse = fa_varlen_forward( + q_i, + k_i, + v_i, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + ) out_list.append(out) lse_list.append(lse) @@ -251,38 +232,28 @@ def llama3_flash_attn_varlen_backward( dv_i = dkv_buffer[1][local_k_slice] # logging.debug(f"bwd i {i} q_slice {q_slice} k_ishape {k_i.shape} dv_i.shape {dv_i.shape} q.shape {q.shape}") - # params = get_default_args(_flash_attn_varlen_backward).copy() - params = { - "dout": dout_i, - "q": q_i, - "k": k_i, - "v": v_i, - "out": out_i, - "softmax_lse": lse_i, - "dq": dq_i, - "dk": dk_i, - "dv": dv_i, - "cu_seqlens_q": cu_seqlens_q, - "cu_seqlens_k": cu_seqlens_k, - "max_seqlen_q": max_seqlen_q, - "max_seqlen_k": max_seqlen_k, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "softcap": softcap, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - _flash_attn_varlen_backward(**params) + fa_varlen_backward( + dout_i, + q_i, + k_i, + v_i, + out_i, + lse_i, + dq_i, + dk_i, + dv_i, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + deterministic=deterministic, + ) if heads_k_stride != nheads_k: # reduce_scatter needs contiguous buffer diff --git a/ring_flash_attn/llama_fwd_ring_bwd_flash_attn.py b/ring_flash_attn/llama_fwd_ring_bwd_flash_attn.py index 99ccb38..5032cdc 100644 --- a/ring_flash_attn/llama_fwd_ring_bwd_flash_attn.py +++ b/ring_flash_attn/llama_fwd_ring_bwd_flash_attn.py @@ -1,28 +1,23 @@ import torch import torch.distributed as dist -from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward +from .flash_attn_backend import fa_forward, fa_backward from .ring_flash_attn import ring_flash_attn_backward from einops import rearrange from typing import Optional, Tuple from .utils import ( - get_default_args, AllGatherComm as Comm, ReduceScatterHandleManager, ) import logging import torch.distributed._tensor as distp_tensor -import flash_attn import os -if torch.__version__ >= "2.4.0" and flash_attn.__version__ >= "2.7.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward +# Kept under their historical names: these are the seam the kernel-level backend +# selection happens behind, and test_llama_bwd_stride_flags.py patches +# `_wrapped_flash_attn_backward` on this module by name. Call sites below pass +# keyword arguments only, which that test's stub relies on to identify dq/dk/dv. +_wrapped_flash_attn_forward = fa_forward +_wrapped_flash_attn_backward = fa_backward def llama_flash_attn_forward( process_group: dist.ProcessGroup, @@ -207,7 +202,6 @@ def llama_flash_attn_forward( k_i = rearrange(current_kv_buffer[0].contiguous(), "w b s hs dh -> b (w s) hs dh") v_i = rearrange(current_kv_buffer[1].contiguous(), "w b s hs dh -> b (w s) hs dh") - # params = get_default_args(_flash_attn_varlen_forward).copy() params = { "q": q_i, "k": k_i, @@ -215,23 +209,15 @@ def llama_flash_attn_forward( "dropout_p": dropout_p, "softmax_scale": softmax_scale, "causal": causal, # 'step' was not defined in this scope - "window_size_left": window_size[0], - "window_size_right": window_size[1], + "window_size": window_size, "softcap": softcap, "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, } - # logging.debug(f"fwd i {i} k_ishape {k_i.shape} s{k_i[0,:3,0,:2]} e{k_i[0,-3:,0,:2]} q_i.shape {q_i.shape} params {params}") + # logging.debug(f"fwd i {i} k_ishape {k_i.shape} s{k_i[0,:3,0,:2]} e{k_i[0,-3:,0,:2]} q_i.shape {q_i.shape} params {params}") # process_id = os.getpid() # if not os.path.exists('./logging/k_buffer_{}.pt'.format(process_id)): # torch.save(k_i.detach(), './logging/k_buffer_{}.pt'.format(process_id)) - # out, _, _, _, _, lse, _, _ = _flash_attn_varlen_forward(**params) - outputs = _wrapped_flash_attn_forward(**params) - if len(outputs) == 8: - out, _, _, _, _, lse, _, _ = outputs - else: - assert len(outputs) == 4 - out, lse, _, _ = outputs + out, lse = _wrapped_flash_attn_forward(**params) out_list.append(out) lse_list.append(lse) @@ -503,8 +489,7 @@ def scatter_in_shape(width): "dropout_p": dropout_p, "softmax_scale": softmax_scale, "causal": causal, - "window_size_left": window_size[0], - "window_size_right": window_size[1], + "window_size": window_size, "softcap": softcap, "alibi_slopes": alibi_slopes, "deterministic": deterministic, @@ -1057,19 +1042,12 @@ def forward( "dropout_p": dropout_p, "softmax_scale": softmax_scale, "causal": causal, - "window_size_left": window_size[0], - "window_size_right": window_size[1], + "window_size": window_size, "softcap": 0.0, "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, } - outputs = _wrapped_flash_attn_forward(**params) - if len(outputs) == 8: - out, _, _, _, _, softmax_lse, _, _ = outputs - else: - assert len(outputs) == 4 - out, softmax_lse, _, _ = outputs + out, softmax_lse = _wrapped_flash_attn_forward(**params) ctx.bwd_event_sync = False else: # out shape (batch, seq, heads, head_dim) @@ -1128,8 +1106,7 @@ def backward(ctx, dout, *args): "dropout_p": ctx.dropout_p, "softmax_scale": ctx.softmax_scale, "causal": ctx.causal, - "window_size_left": ctx.window_size[0], - "window_size_right": ctx.window_size[1], + "window_size": ctx.window_size, "softcap": 0.0, "alibi_slopes": ctx.alibi_slopes, "deterministic": ctx.deterministic, @@ -1231,19 +1208,12 @@ def forward( "dropout_p": dropout_p, "softmax_scale": softmax_scale, "causal": causal, - "window_size_left": window_size[0], - "window_size_right": window_size[1], + "window_size": window_size, "softcap": 0.0, "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, } - outputs = _wrapped_flash_attn_forward(**params) - if len(outputs) == 8: - out, _, _, _, _, softmax_lse, _, _ = outputs - else: - assert len(outputs) == 4 - out, softmax_lse, _, _ = outputs + out, softmax_lse = _wrapped_flash_attn_forward(**params) ctx.bwd_event_sync = False else: # out shape (batch, seq, heads, head_dim) diff --git a/ring_flash_attn/ring_flash_attn.py b/ring_flash_attn/ring_flash_attn.py index 626e21a..10e257c 100644 --- a/ring_flash_attn/ring_flash_attn.py +++ b/ring_flash_attn/ring_flash_attn.py @@ -1,20 +1,10 @@ import torch import torch.distributed as dist -from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward -from .utils import RingComm, update_out_and_lse, get_default_args +from .flash_attn_backend import fa_forward, fa_backward, check_variant_supported +from .utils import RingComm, update_out_and_lse import logging import gc -import flash_attn -if torch.__version__ >= "2.4.0" and flash_attn.__version__ >= "2.7.0": - _wrapped_flash_attn_forward = torch.ops.flash_attn._flash_attn_forward -else: - _wrapped_flash_attn_forward = _flash_attn_forward - -if torch.__version__ >= "2.4.0": - _wrapped_flash_attn_backward = torch.ops.flash_attn._flash_attn_backward -else: - _wrapped_flash_attn_backward = _flash_attn_backward def ring_flash_attn_forward( process_group, @@ -29,6 +19,7 @@ def ring_flash_attn_forward( alibi_slopes=None, deterministic=False, ): + check_variant_supported("ring_flash_attn_forward") comm = RingComm(process_group) out = None @@ -41,35 +32,17 @@ def ring_flash_attn_forward( next_k, next_v = comm.send_recv_kv(k, v) if not causal or step <= comm.rank: - params = get_default_args(_flash_attn_forward).copy() - params.update( - { - "q": q, - "k": k, - "v": v, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal and step == 0, - "softcap": softcap, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } + block_out, block_lse = fa_forward( + q, + k, + v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal and step == 0, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - outputs = _wrapped_flash_attn_forward(**params) - if len(outputs) == 8: - block_out, _, _, _, _, block_lse, _, _ = outputs - else: - assert len(outputs) == 4 - block_out, block_lse, _, _ = outputs out, lse = update_out_and_lse(out, lse, block_out, block_lse) if step + 1 != comm.world_size: @@ -118,37 +91,24 @@ def ring_flash_attn_backward( if step <= kv_comm.rank or not causal: bwd_causal = causal and step == 0 - params = get_default_args(_flash_attn_backward).copy() - params.update( - { - "dout": dout, - "q": q, - "k": k, - "v": v, - "out": out, - "softmax_lse": softmax_lse, - "dq": block_dq_buffer, - "dk": block_dk_buffer, - "dv": block_dv_buffer, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": bwd_causal, - "softcap": softcap, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } + fa_backward( + dout, + q, + k, + v, + out, + softmax_lse, + block_dq_buffer, + block_dk_buffer, + block_dv_buffer, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=bwd_causal, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + deterministic=deterministic, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - # logging.debug(f"q {params['q'].shape} k {params['k'].shape} v {params['v'].shape} dout {params['dout'].shape} softmax_lse {params['softmax_lse'].shape}") - _wrapped_flash_attn_backward(**params) if dq is None: dq = block_dq_buffer.to(torch.float32) diff --git a/ring_flash_attn/ring_flash_attn_varlen.py b/ring_flash_attn/ring_flash_attn_varlen.py index f7f47b1..0837d93 100644 --- a/ring_flash_attn/ring_flash_attn_varlen.py +++ b/ring_flash_attn/ring_flash_attn_varlen.py @@ -1,13 +1,13 @@ import torch import torch.distributed as dist -from flash_attn.flash_attn_interface import ( - _flash_attn_varlen_forward, - _flash_attn_varlen_backward, +from .flash_attn_backend import ( + fa_varlen_forward, + fa_varlen_backward, + check_variant_supported, ) from .utils import ( RingComm, update_out_and_lse, - get_default_args, ) try: @@ -37,6 +37,7 @@ def ring_flash_attn_varlen_forward( alibi_slopes=None, deterministic=False, ): + check_variant_supported("ring_flash_attn_varlen") comm = RingComm(process_group) out = None @@ -48,40 +49,21 @@ def ring_flash_attn_varlen_forward( if step + 1 != comm.world_size: next_k, next_v = comm.send_recv_kv(k, v) if not causal or step <= comm.rank: - params = get_default_args(_flash_attn_varlen_forward).copy() - params.update( - { - "q": q, - "k": k, - "v": v, - "cu_seqlens_q": cu_seqlens, - "cu_seqlens_k": cu_seqlens, - "max_seqlen_q": max_seqlen, - "max_seqlen_k": max_seqlen, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal and step == 0, - "softcap": softcap, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } + block_out, block_lse = fa_varlen_forward( + q, + k, + v, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal and step == 0, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - - outputs = _flash_attn_varlen_forward(**params) - if len(outputs) == 8: - block_out, _, _, _, _, block_lse, _, _ = outputs - else: - assert len(outputs) == 4 - block_out, block_lse, _, _ = outputs if block_lse.dim() == 3: old_lse = True block_lse = flatten_varlen_lse( @@ -120,6 +102,7 @@ def ring_flash_attn_varlen_backward( alibi_slopes=None, deterministic=False, ): + check_variant_supported("ring_flash_attn_varlen") kv_comm = RingComm(process_group) d_kv_comm = RingComm(process_group) dq, dk, dv = None, None, None @@ -138,40 +121,28 @@ def ring_flash_attn_varlen_backward( if step <= kv_comm.rank or not causal: bwd_causal = causal and step == 0 - params = get_default_args(_flash_attn_varlen_backward).copy() - params.update( - { - "dout": dout, - "q": q, - "k": k, - "v": v, - "out": out, - "softmax_lse": softmax_lse, - "dq": block_dq_buffer, - "dk": block_dk_buffer, - "dv": block_dv_buffer, - "cu_seqlens_q": cu_seqlens, - "cu_seqlens_k": cu_seqlens, - "max_seqlen_q": max_seqlen, - "max_seqlen_k": max_seqlen, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": bwd_causal, - "softcap": softcap, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } + fa_varlen_backward( + dout, + q, + k, + v, + out, + softmax_lse, + block_dq_buffer, + block_dk_buffer, + block_dv_buffer, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=bwd_causal, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + deterministic=deterministic, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - _flash_attn_varlen_backward(**params) if dq is None: dq = block_dq_buffer.to(torch.float32) diff --git a/ring_flash_attn/stripe_flash_attn.py b/ring_flash_attn/stripe_flash_attn.py index de1dd07..8f21760 100644 --- a/ring_flash_attn/stripe_flash_attn.py +++ b/ring_flash_attn/stripe_flash_attn.py @@ -1,7 +1,7 @@ import torch import torch.distributed as dist -from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward -from .utils import RingComm, update_out_and_lse, get_default_args +from .flash_attn_backend import fa_forward, fa_backward, check_variant_supported +from .utils import RingComm, update_out_and_lse def stripe_flash_attn_forward( @@ -16,6 +16,7 @@ def stripe_flash_attn_forward( alibi_slopes=None, deterministic=False, ): + check_variant_supported("stripe_flash_attn") assert ( causal ), "stripe flash attn only supports causal attention, if not causal, use ring flash attn instead" @@ -30,64 +31,32 @@ def stripe_flash_attn_forward( if step + 1 != comm.world_size: next_k, next_v = comm.send_recv_kv(k, v) - params = get_default_args(_flash_attn_forward).copy() + # use_custom_op=False keeps this on the plain python entry point, as before. if step <= comm.rank: - params.update( - { - "q": q, - "k": k, - "v": v, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } + block_out, block_lse = fa_forward( + q, + k, + v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + use_custom_op=False, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - outputs = _flash_attn_forward(**params) - if len(outputs) == 8: - block_out, _, _, _, _, block_lse, _, _ = outputs - else: - assert len(outputs) == 4 - block_out, block_lse, _, _ = outputs out, lse = update_out_and_lse(out, lse, block_out, block_lse) else: - params.update( - { - "q": q[:, 1:], - "k": k[:, :-1], - "v": v[:, :-1], - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } + block_out, block_lse = fa_forward( + q[:, 1:], + k[:, :-1], + v[:, :-1], + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + use_custom_op=False, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - outputs = _flash_attn_forward(**params) - if len(outputs) == 8: - block_out, _, _, _, _, block_lse, _, _ = outputs - else: - assert len(outputs) == 4 - block_out, block_lse, _, _ = outputs out, lse = update_out_and_lse( out, lse, block_out, block_lse, slice_=(slice(None), slice(1, None)) ) @@ -119,6 +88,7 @@ def stripe_flash_attn_backward( assert ( causal ), "stripe flash attn only supports causal attention, if not causal, ring flash attn instead" + check_variant_supported("stripe_flash_attn") kv_comm = RingComm(process_group) d_kv_comm = RingComm(process_group) dq, dk, dv = None, None, None @@ -135,68 +105,48 @@ def stripe_flash_attn_backward( shift_causal = step > kv_comm.rank softmax_lse_1 = None - params = get_default_args(_flash_attn_backward).copy() + # use_custom_op=False keeps this on the plain python entry point, as before. if not shift_causal: - params.update( - { - "dout": dout, - "q": q, - "k": k, - "v": v, - "out": out, - "softmax_lse": softmax_lse, - "dq": block_dq_buffer, - "dk": block_dk_buffer, - "dv": block_dv_buffer, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } + fa_backward( + dout, + q, + k, + v, + out, + softmax_lse, + block_dq_buffer, + block_dk_buffer, + block_dv_buffer, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + deterministic=deterministic, + use_custom_op=False, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - _flash_attn_backward(**params) else: if softmax_lse_1 is None: # lazy init, since the last rank does not need softmax_lse_1 softmax_lse_1 = softmax_lse[:, :, 1:].contiguous() - params.update( - { - "dout": dout[:, 1:], - "q": q[:, 1:], - "k": k[:, :-1], - "v": v[:, :-1], - "out": out[:, 1:], - "softmax_lse": softmax_lse_1, - "dq": block_dq_buffer[:, 1:], - "dk": block_dk_buffer[:, :-1], - "dv": block_dv_buffer[:, :-1], - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } + fa_backward( + dout[:, 1:], + q[:, 1:], + k[:, :-1], + v[:, :-1], + out[:, 1:], + softmax_lse_1, + block_dq_buffer[:, 1:], + block_dk_buffer[:, :-1], + block_dv_buffer[:, :-1], + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + deterministic=deterministic, + use_custom_op=False, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - _flash_attn_backward(**params) if dq is None: dq = block_dq_buffer.to(torch.float32) diff --git a/ring_flash_attn/zigzag_ring_flash_attn.py b/ring_flash_attn/zigzag_ring_flash_attn.py index a0bc36d..1cd2aa6 100644 --- a/ring_flash_attn/zigzag_ring_flash_attn.py +++ b/ring_flash_attn/zigzag_ring_flash_attn.py @@ -1,7 +1,7 @@ import torch import torch.distributed as dist -from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward -from .utils import RingComm, update_out_and_lse, get_default_args +from .flash_attn_backend import fa_forward, fa_backward, check_variant_supported +from .utils import RingComm, update_out_and_lse def zigzag_ring_flash_attn_forward( @@ -16,6 +16,7 @@ def zigzag_ring_flash_attn_forward( alibi_slopes=None, deterministic=False, ): + check_variant_supported("zigzag_ring_flash_attn") assert causal == True, "zigzag ring is meaningless for causal=False" comm = RingComm(process_group) @@ -27,35 +28,18 @@ def zigzag_ring_flash_attn_forward( next_k, next_v = None, None def forward(q, k, v, causal): - params = get_default_args(_flash_attn_forward).copy() - params.update( - { - "q": q, - "k": k, - "v": v, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } + # use_custom_op=False keeps this on the plain python entry point, as before. + return fa_forward( + q, + k, + v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + use_custom_op=False, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - outputs = _flash_attn_forward(**params) - if len(outputs) == 8: - block_out, _, _, _, _, block_lse, _, _ = outputs - else: - assert len(outputs) == 4 - block_out, block_lse, _, _ = outputs - return block_out, block_lse for step in range(comm.world_size): if step + 1 != comm.world_size: @@ -103,6 +87,7 @@ def zigzag_ring_flash_attn_backward( alibi_slopes=None, deterministic=False, ): + check_variant_supported("zigzag_ring_flash_attn") assert causal == True, "zigzag ring is meaningless for causal=False" kv_comm = RingComm(process_group) d_kv_comm = RingComm(process_group) @@ -125,35 +110,24 @@ def zigzag_ring_flash_attn_backward( def backward(dout, q, k, v, out, softmax_lse, causal): seqlen_q = q.shape[1] seqlen_kv = k.shape[1] - params = get_default_args(_flash_attn_backward).copy() - params.update( - { - "dout": dout, - "q": q, - "k": k, - "v": v, - "out": out, - "softmax_lse": softmax_lse, - "dq": dq_buffer[:, :seqlen_q], - "dk": dk_buffer[:, :seqlen_kv], - "dv": dv_buffer[:, :seqlen_kv], - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } + fa_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq_buffer[:, :seqlen_q], + dk_buffer[:, :seqlen_kv], + dv_buffer[:, :seqlen_kv], + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + deterministic=deterministic, + use_custom_op=False, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - _flash_attn_backward(**params) for step in range(kv_comm.world_size): if step + 1 != kv_comm.world_size: diff --git a/ring_flash_attn/zigzag_ring_flash_attn_varlen.py b/ring_flash_attn/zigzag_ring_flash_attn_varlen.py index 7f1be94..75731a8 100644 --- a/ring_flash_attn/zigzag_ring_flash_attn_varlen.py +++ b/ring_flash_attn/zigzag_ring_flash_attn_varlen.py @@ -1,12 +1,12 @@ import torch -from flash_attn.flash_attn_interface import ( - _flash_attn_varlen_forward, - _flash_attn_varlen_backward, +from .flash_attn_backend import ( + fa_varlen_forward, + fa_varlen_backward, + check_variant_supported, ) from .utils import ( RingComm, update_out_and_lse, - get_default_args, ) try: @@ -88,6 +88,7 @@ def zigzag_ring_flash_attn_varlen_forward( deterministic=False, ): assert causal == True, "zigzag ring is meaningless for causal=False" + check_variant_supported("zigzag_ring_flash_attn_varlen") comm = RingComm(process_group) block_seq_len = q.shape[0] // 2 @@ -107,40 +108,21 @@ def forward(q, k, v, causal): cu_seqlens_kv = half_cu_seqlens if seqlen_kv == block_seq_len else cu_seqlens max_seqlen_kv = half_max_seqlen if seqlen_kv == block_seq_len else max_seqlen - params = get_default_args(_flash_attn_varlen_forward).copy() - params.update( - { - "q": q, - "k": k, - "v": v, - # the first half and the second half are the same - "cu_seqlens_q": cu_seqlens_q, - "cu_seqlens_k": cu_seqlens_kv, - "max_seqlen_q": max_seqlen_q, - "max_seqlen_k": max_seqlen_kv, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "return_softmax": True and dropout_p > 0, - } + return fa_varlen_forward( + q, + k, + v, + # the first half and the second half are the same + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - outputs = _flash_attn_varlen_forward(**params) - if len(outputs) == 8: - block_out, _, _, _, _, block_lse, _, _ = outputs - else: - assert len(outputs) == 4 - block_out, block_lse, _, _ = outputs - return block_out, block_lse old_lse = False for step in range(comm.world_size): @@ -211,6 +193,7 @@ def zigzag_ring_flash_attn_varlen_backward( deterministic=False, ): assert causal == True, "zigzag ring is meaningless for causal=False" + check_variant_supported("zigzag_ring_flash_attn_varlen") kv_comm = RingComm(process_group) d_kv_comm = RingComm(process_group) dq, dk, dv = None, None, None @@ -239,40 +222,28 @@ def backward(dout, q, k, v, out, softmax_lse, causal): max_seqlen_q = half_max_seqlen if seqlen_q == block_seq_len else max_seqlen cu_seqlens_kv = half_cu_seqlens if seqlen_kv == block_seq_len else cu_seqlens max_seqlen_kv = half_max_seqlen if seqlen_kv == block_seq_len else max_seqlen - params = get_default_args(_flash_attn_varlen_backward).copy() - params.update( - { - "dout": dout, - "q": q, - "k": k, - "v": v, - "out": out, - "softmax_lse": softmax_lse, - "dq": dq_buffer[:seqlen_q], - "dk": dk_buffer[:seqlen_kv], - "dv": dv_buffer[:seqlen_kv], - # the first half and the second half are the same - "cu_seqlens_q": cu_seqlens_q, - "cu_seqlens_k": cu_seqlens_kv, - "max_seqlen_q": max_seqlen_q, - "max_seqlen_k": max_seqlen_kv, - "dropout_p": dropout_p, - "softmax_scale": softmax_scale, - "causal": causal, - "alibi_slopes": alibi_slopes, - "deterministic": deterministic, - } + fa_varlen_backward( + dout, + q, + k, + v, + out, + softmax_lse, + dq_buffer[:seqlen_q], + dk_buffer[:seqlen_kv], + dv_buffer[:seqlen_kv], + # the first half and the second half are the same + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + deterministic=deterministic, ) - if "window_size" in params: - params.update({"window_size": window_size}) - else: - params.update( - { - "window_size_left": window_size[0], - "window_size_right": window_size[1], - } - ) - _flash_attn_varlen_backward(**params) for step in range(kv_comm.world_size): if step + 1 != kv_comm.world_size: diff --git a/test/test.sh b/test/test.sh index 94f7d8a..d746f5b 100755 --- a/test/test.sh +++ b/test/test.sh @@ -18,12 +18,16 @@ tests=( ) cpu_tests=( - test.test_reduce_scatter_handle_manager - test.test_llama_bwd_stride_flags + test_reduce_scatter_handle_manager + test_llama_bwd_stride_flags + test_fa3_arg_binding ) +# Run as bare module names with test/ on the path. The previous `test.` form +# resolved to the CPython stdlib `test` package (there is no test/__init__.py here), +# so these never actually ran. for test in "${cpu_tests[@]}"; do - python -m unittest "$test" + PYTHONPATH=.:test python -m unittest "$test" done for test in "${tests[@]}"; do diff --git a/test/test_fa3_arg_binding.py b/test/test_fa3_arg_binding.py new file mode 100644 index 0000000..47e60e1 --- /dev/null +++ b/test/test_fa3_arg_binding.py @@ -0,0 +1,447 @@ +"""Bind the fa3 adapter's kwargs against the real flash-attn-3 signatures, without a GPU. + +flash-attn-3 cannot be imported on a machine with no `flash_attn_3._C`, and its kernels +need sm90-class hardware. But the highest-risk errors in the adapter are pure +argument-binding mistakes, and those are checkable against the signature alone: + + * backward renames `causal` to `is_causal` -- a straight copy of the FA2 call raises + * backward puts `cu_seqlens_*`/`max_seqlen_*` BEFORE `dq`/`dk`/`dv`, where FA2 puts + `dq`/`dk`/`dv` immediately after `softmax_lse`, so a positional call misbinds + * `dq`/`dk`/`dv` are Optional and the C++ silently allocates throwaways when omitted, + computing the gradient and discarding it with no error + * forward takes no `dropout_p`/`alibi_slopes`/`return_softmax` + +The signatures are loaded from a flash-attention checkout's hopper/flash_attn_interface.py +with `flash_attn_3._C` stubbed, since `_init_fn` is the undecorated python function and its +signature survives without any kernel. + +Set FLASH_ATTENTION_HOPPER to point at a checkout's hopper/ directory; the test skips if +it is not found. +""" + +import inspect +import os +import sys +import types +import unittest +from pathlib import Path + +_DEFAULT_HOPPER = Path.home() / "projects/repos/flash-attention/hopper" + + +def _hopper_dir(): + return Path(os.environ.get("FLASH_ATTENTION_HOPPER", _DEFAULT_HOPPER)) + + +def _load_fa3_interface(): + """Import hopper/flash_attn_interface.py with the CUDA extension stubbed out.""" + import importlib.util + + path = _hopper_dir() / "flash_attn_interface.py" + if not path.is_file(): + raise unittest.SkipTest(f"no flash-attn-3 interface at {path}") + + import torch + + # The module does `import flash_attn_3._C` at module scope purely to register ops. + stub_pkg = types.ModuleType("flash_attn_3") + stub_pkg.__path__ = [] + stub_ext = types.ModuleType("flash_attn_3._C") + saved = {k: sys.modules.get(k) for k in ("flash_attn_3", "flash_attn_3._C")} + sys.modules["flash_attn_3"] = stub_pkg + sys.modules["flash_attn_3._C"] = stub_ext + + # torch.ops.flash_attn_3 resolves lazily, so binding it needs no real library. + try: + spec = importlib.util.spec_from_file_location("_fa3_iface_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + except Exception as e: # pragma: no cover - environment dependent + raise unittest.SkipTest(f"could not load flash-attn-3 interface: {e!r}") + finally: + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + return module + + +def _signature_of(op): + """Real signature of a torch.library custom op (or a plain function).""" + # inspect.signature on a CustomOpDef only reports (*args, **kwargs). + return inspect.signature(getattr(op, "_init_fn", op)) + + +class TestFA3ArgBinding(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.iface = _load_fa3_interface() + cls.fwd_sig = _signature_of(cls.iface._flash_attn_forward) + cls.bwd_sig = _signature_of(cls.iface._flash_attn_backward) + + # -- forward --------------------------------------------------------------- + + def test_forward_rejects_fa2_only_kwargs(self): + """dropout_p / alibi_slopes / return_softmax must not be passed to FA3.""" + for absent in ("dropout_p", "alibi_slopes", "return_softmax"): + self.assertNotIn( + absent, + self.fwd_sig.parameters, + f"{absent} unexpectedly present; the adapter assumes FA3 has no such arg", + ) + + def test_forward_window_is_two_ints(self): + self.assertIn("window_size_left", self.fwd_sig.parameters) + self.assertIn("window_size_right", self.fwd_sig.parameters) + self.assertNotIn("window_size", self.fwd_sig.parameters) + # -1 is the native "unbounded" sentinel; FA4's None mapping must not leak here. + self.assertEqual(self.fwd_sig.parameters["window_size_left"].default, -1) + self.assertEqual(self.fwd_sig.parameters["window_size_right"].default, -1) + + def test_forward_binds_adapter_kwargs(self): + params = { + "q": object(), + "k": object(), + "v": object(), + "softmax_scale": 0.125, + "causal": True, + "softcap": 0.0, + "window_size_left": -1, + "window_size_right": -1, + "cu_seqlens_q": object(), + "cu_seqlens_k": object(), + "max_seqlen_q": 16, + "max_seqlen_k": 32, + } + bound = self.fwd_sig.bind(**params) + self.assertEqual(bound.arguments["causal"], True) + self.assertEqual(bound.arguments["max_seqlen_q"], 16) + + # -- backward -------------------------------------------------------------- + + def test_backward_uses_is_causal_not_causal(self): + self.assertIn("is_causal", self.bwd_sig.parameters) + self.assertNotIn("causal", self.bwd_sig.parameters) + + def test_backward_dqdkdv_bind_to_grad_params_not_cu_seqlens(self): + """The reordering trap: cu_seqlens_*/max_seqlen_* sit before dq/dk/dv in FA3.""" + dq, dk, dv = object(), object(), object() + params = { + "dout": object(), + "q": object(), + "k": object(), + "v": object(), + "out": object(), + "softmax_lse": object(), + "dq": dq, + "dk": dk, + "dv": dv, + "softmax_scale": 0.125, + "is_causal": True, + "softcap": 0.0, + "deterministic": False, + "window_size_left": -1, + "window_size_right": -1, + } + bound = self.bwd_sig.bind(**params) + # Identity, not just presence -- a positional misbind can still bind cleanly. + self.assertIs(bound.arguments["dq"], dq) + self.assertIs(bound.arguments["dk"], dk) + self.assertIs(bound.arguments["dv"], dv) + for grad_name, grad in (("dq", dq), ("dk", dk), ("dv", dv)): + for other in ("cu_seqlens_q", "cu_seqlens_k", "sequed_q", "sequed_k"): + self.assertIsNot( + bound.arguments.get(other), + grad, + f"{grad_name} misbound to {other}", + ) + + def test_backward_grad_params_are_optional_so_must_be_passed(self): + """Documents why the adapter always passes dq/dk/dv: omitting them is silent.""" + for name in ("dq", "dk", "dv"): + self.assertIs( + self.bwd_sig.parameters[name].default, + None, + f"{name} is Optional; the C++ allocates a throwaway and discards the " + "gradient without error when it is omitted", + ) + + def test_backward_seqused_typo_is_still_present(self): + """Upstream misspells seqused_* as sequed_* in backward only. + + The adapter never passes seqused, but if this typo is ever fixed upstream a + future change that does pass it would silently no-op. Fail loudly instead. + """ + self.assertIn("sequed_q", self.bwd_sig.parameters) + self.assertIn("sequed_k", self.bwd_sig.parameters) + self.assertIn("seqused_q", self.fwd_sig.parameters) + + +class TestFA3AdapterParams(unittest.TestCase): + """Bind the params the adapter actually builds against the real FA3 signatures. + + The tests above check the signature; these check *our* call against it, which is the + thing that can actually be wrong. The seam is imported with RING_FLASH_ATTN_BACKEND=fa3 + while flash_attn_3 is stubbed to serve the checkout's real ops, so the adapter + configures itself exactly as it would on Hopper -- including the introspection that + picks `is_causal` over `causal`. + """ + + @classmethod + def setUpClass(cls): + import importlib + + import importlib.machinery + + iface = _load_fa3_interface() + + # The seam's probe uses importlib.util.find_spec, which consults sys.modules and + # reads __spec__ -- a bare ModuleType has __spec__ = None and would look absent. + pkg = types.ModuleType("flash_attn_3") + pkg.__path__ = [] + pkg.__spec__ = importlib.machinery.ModuleSpec( + "flash_attn_3", loader=None, is_package=True + ) + ext = types.ModuleType("flash_attn_3._C") + ext.__spec__ = importlib.machinery.ModuleSpec("flash_attn_3._C", loader=None) + cls._saved_modules = { + k: sys.modules.get(k) + for k in ( + "flash_attn_3", + "flash_attn_3._C", + "flash_attn_3.flash_attn_interface", + "ring_flash_attn.flash_attn_backend", + ) + } + sys.modules["flash_attn_3"] = pkg + sys.modules["flash_attn_3._C"] = ext + sys.modules["flash_attn_3.flash_attn_interface"] = iface + pkg.flash_attn_interface = iface + + cls._saved_env = os.environ.get("RING_FLASH_ATTN_BACKEND") + os.environ["RING_FLASH_ATTN_BACKEND"] = "fa3" + sys.modules.pop("ring_flash_attn.flash_attn_backend", None) + try: + cls.seam = importlib.import_module("ring_flash_attn.flash_attn_backend") + except Exception as e: # pragma: no cover - environment dependent + cls._restore() + raise unittest.SkipTest(f"could not load seam under fa3: {e!r}") + cls.fwd_sig = _signature_of(iface._flash_attn_forward) + cls.bwd_sig = _signature_of(iface._flash_attn_backward) + + @classmethod + def _restore(cls): + if cls._saved_env is None: + os.environ.pop("RING_FLASH_ATTN_BACKEND", None) + else: + os.environ["RING_FLASH_ATTN_BACKEND"] = cls._saved_env + for k, v in cls._saved_modules.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + + @classmethod + def tearDownClass(cls): + cls._restore() + + def test_backend_selected_and_capabilities(self): + self.assertEqual(self.seam.BACKEND, "fa3") + self.assertFalse(self.seam.CAPABILITIES.supports_dropout) + self.assertFalse(self.seam.CAPABILITIES.supports_alibi) + + def test_adapter_picked_is_causal(self): + self.assertEqual(self.seam._FA3_BWD_CAUSAL_KEY, "is_causal") + + def test_adapter_forward_params_bind(self): + params = self.seam._fa3_forward_params( + q="Q", k="K", v="V", + softmax_scale=0.125, causal=True, window_size=(-1, -1), + softcap=0.0, dropout_p=0.0, alibi_slopes=None, + ) + bound = self.fwd_sig.bind(**params) # must not raise + self.assertEqual(bound.arguments["causal"], True) + self.assertEqual(bound.arguments["window_size_left"], -1) + self.assertNotIn("dropout_p", params) + self.assertNotIn("alibi_slopes", params) + + def test_adapter_backward_params_bind_and_place_grads(self): + dq, dk, dv = object(), object(), object() + params = self.seam._fa3_backward_params( + dout="DO", q="Q", k="K", v="V", out="O", softmax_lse="LSE", + dq=dq, dk=dk, dv=dv, + softmax_scale=0.125, causal=True, window_size=(-1, -1), + softcap=0.0, deterministic=False, dropout_p=0.0, alibi_slopes=None, + ) + bound = self.bwd_sig.bind(**params) # must not raise + self.assertIs(bound.arguments["dq"], dq) + self.assertIs(bound.arguments["dk"], dk) + self.assertIs(bound.arguments["dv"], dv) + self.assertEqual(bound.arguments["is_causal"], True) + self.assertNotIn("causal", params) + + def test_adapter_varlen_backward_params_bind(self): + dq, dk, dv = object(), object(), object() + params = self.seam._fa3_backward_params( + dout="DO", q="Q", k="K", v="V", out="O", softmax_lse="LSE", + dq=dq, dk=dk, dv=dv, + softmax_scale=0.125, causal=True, window_size=(-1, -1), + softcap=0.0, deterministic=False, dropout_p=0.0, alibi_slopes=None, + ) + params.update({ + "cu_seqlens_q": "CQ", "cu_seqlens_k": "CK", + "max_seqlen_q": 16, "max_seqlen_k": 32, + }) + bound = self.bwd_sig.bind(**params) + self.assertIs(bound.arguments["dq"], dq) + self.assertEqual(bound.arguments["cu_seqlens_q"], "CQ") + self.assertEqual(bound.arguments["max_seqlen_k"], 32) + + def test_adapter_rejects_dropout_and_alibi(self): + with self.assertRaises(NotImplementedError): + self.seam._fa3_forward_params( + q="Q", k="K", v="V", softmax_scale=0.125, causal=True, + window_size=(-1, -1), softcap=0.0, dropout_p=0.1, alibi_slopes=None, + ) + with self.assertRaises(NotImplementedError): + self.seam._fa3_forward_params( + q="Q", k="K", v="V", softmax_scale=0.125, causal=True, + window_size=(-1, -1), softcap=0.0, dropout_p=0.0, alibi_slopes="SLOPES", + ) + + def _assert_extras_are_declared_defaults(self, params, explicit, sig): + """Every arg we pass but do not set must equal its own declared default. + + get_default_args makes the adapter pass all ~34 forward args explicitly, of + which we set only a handful. That is a semantic no-op *only* if the rest carry + the function's own defaults -- the same property that settled the equivalent + question for FA2. _get_default_args force-sets `softcap` and pads non-defaulted + params with None, so this is worth asserting rather than assuming. + """ + for name, value in params.items(): + if name in explicit: + continue + declared = sig.parameters[name].default + if declared is inspect.Parameter.empty: + # Required params get None-padded by _get_default_args; we must be + # setting all of those explicitly, or the call is malformed. + self.fail(f"{name} is required but not set explicitly (got {value!r})") + self.assertEqual( + value, declared, f"{name} passed as {value!r}, declared default {declared!r}" + ) + + def test_forward_extras_equal_declared_defaults(self): + params = self.seam._fa3_forward_params( + q="Q", k="K", v="V", + softmax_scale=0.125, causal=True, window_size=(-1, -1), + softcap=0.0, dropout_p=0.0, alibi_slopes=None, + ) + explicit = { + "q", "k", "v", "softmax_scale", "causal", "softcap", + "window_size_left", "window_size_right", + } + self._assert_extras_are_declared_defaults(params, explicit, self.fwd_sig) + + def test_backward_extras_equal_declared_defaults(self): + params = self.seam._fa3_backward_params( + dout="DO", q="Q", k="K", v="V", out="O", softmax_lse="LSE", + dq="DQ", dk="DK", dv="DV", + softmax_scale=0.125, causal=True, window_size=(-1, -1), + softcap=0.0, deterministic=False, dropout_p=0.0, alibi_slopes=None, + ) + explicit = { + "dout", "q", "k", "v", "out", "softmax_lse", "dq", "dk", "dv", + "softmax_scale", "is_causal", "softcap", "deterministic", + "window_size_left", "window_size_right", + } + self._assert_extras_are_declared_defaults(params, explicit, self.bwd_sig) + + def test_out_of_scope_variant_refused_under_fa3(self): + self.seam.check_variant_supported("llama3_flash_attn_varlen") # in scope + self.seam.check_variant_supported("ring_flash_attn_backward") # in scope + with self.assertRaises(NotImplementedError): + self.seam.check_variant_supported("zigzag_ring_flash_attn") + + +class TestFA3BuildFlagGuard(unittest.TestCase): + """The guard must fire on a build that compiled out something ring attention needs. + + Regression test for a real bug: the guard originally used the env-var spelling + FLASH_ATTENTION_DISABLE_* while the generated config keys are FLASHATTENTION_DISABLE_* + (no underscore after FLASH), so every lookup missed and the check was silently dead. + """ + + def _load_seam_with_config(self, build_flags): + import importlib + import importlib.machinery + + iface = _load_fa3_interface() + + pkg = types.ModuleType("flash_attn_3") + pkg.__path__ = [] + pkg.__spec__ = importlib.machinery.ModuleSpec( + "flash_attn_3", loader=None, is_package=True + ) + ext = types.ModuleType("flash_attn_3._C") + ext.__spec__ = importlib.machinery.ModuleSpec("flash_attn_3._C", loader=None) + cfg = types.ModuleType("flash_attn_3.flash_attn_config") + cfg.CONFIG = {"build_flags": build_flags} + + names = ( + "flash_attn_3", + "flash_attn_3._C", + "flash_attn_3.flash_attn_interface", + "flash_attn_3.flash_attn_config", + "ring_flash_attn.flash_attn_backend", + ) + saved = {k: sys.modules.get(k) for k in names} + saved_env = os.environ.get("RING_FLASH_ATTN_BACKEND") + sys.modules["flash_attn_3"] = pkg + sys.modules["flash_attn_3._C"] = ext + sys.modules["flash_attn_3.flash_attn_interface"] = iface + sys.modules["flash_attn_3.flash_attn_config"] = cfg + os.environ["RING_FLASH_ATTN_BACKEND"] = "fa3" + sys.modules.pop("ring_flash_attn.flash_attn_backend", None) + try: + return importlib.import_module("ring_flash_attn.flash_attn_backend") + finally: + if saved_env is None: + os.environ.pop("RING_FLASH_ATTN_BACKEND", None) + else: + os.environ["RING_FLASH_ATTN_BACKEND"] = saved_env + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + + def test_guard_fires_on_disabled_varlen(self): + with self.assertRaises(ImportError) as cm: + self._load_seam_with_config({"FLASHATTENTION_DISABLE_VARLEN": True}) + self.assertIn("FLASHATTENTION_DISABLE_VARLEN", str(cm.exception)) + + def test_guard_fires_on_disabled_backward(self): + with self.assertRaises(ImportError) as cm: + self._load_seam_with_config({"FLASHATTENTION_DISABLE_BACKWARD": True}) + self.assertIn("training", str(cm.exception)) + + def test_guard_silent_on_a_healthy_build(self): + seam = self._load_seam_with_config( + { + "FLASHATTENTION_DISABLE_VARLEN": False, + "FLASHATTENTION_DISABLE_LOCAL": False, + "FLASHATTENTION_DISABLE_BACKWARD": False, + } + ) + self.assertEqual(seam.BACKEND, "fa3") + + def test_guard_tolerates_unexpected_config_shape(self): + """An unrecognised layout means 'no opinion', never a failed import.""" + for weird in ([], "TRUE", None, {"SOMETHING_ELSE": True}): + seam = self._load_seam_with_config(weird) + self.assertEqual(seam.BACKEND, "fa3") + + +if __name__ == "__main__": + unittest.main()