From e443eb4180347b3bf3701543e227915224be8f7d Mon Sep 17 00:00:00 2001 From: eric-tc-wong Date: Fri, 31 Jul 2026 20:08:08 +0800 Subject: [PATCH 1/3] Add a backend seam so ring attention can target FA2 or flash-attn-4 The flash_attn dependency was not abstracted: 7 modules imported it directly and 20 call sites each re-implemented the same three compatibility branches inline (the 8-vs-4 forward tuple unpack, the window_size arity split, and the torch.ops.flash_attn version gate). Introduce ring_flash_attn/flash_attn_backend.py, which normalises the four primitives the ring variants actually need -- fa_forward / fa_backward and their varlen flavours -- onto the FA2 keyword set, and route every call site through it. Backend selection is RING_FLASH_ATTN_BACKEND, defaulting to fa2, so existing behaviour is unchanged unless a caller opts in. FA4 notes: its cute interface returns a 2-tuple, needs an explicit return_lse=True (autograd.Function.forward runs with grad disabled, so the requires_grad shortcut never fires), takes no dropout/ALiBi, maps the -1 window sentinel to None, and registers no torch.library ops -- so it is not traceable under torch.compile. Its backward additionally requires sm90+. The loader raises rather than falling back between backends: flash-attn-4 installs into the flash_attn namespace, so a broken flash-attn 2.x makes FA4 unimportable too, and a silent fallback would change numerics and hardware requirements underneath the caller. Variants outside FA4 scope are refused via check_variant_supported rather than run with unverified numerics. ring_flash_attn_backward is in the validated set because both llama hybrid entry points route into it. Verified against FA2: CPU tests plus 11 GPU tests and 7 compile-leg runs at nproc_per_node 1, all metrics bit-identical to a stashed baseline except dq diff, which is nondeterministic under FA2's atomic dq accumulation (shown via repeat runs on both sides, and by confirming the seam's only added kwargs are parameters passed at their own declared defaults). Co-Authored-By: Claude Opus 5 --- pyproject.toml | 7 + ring_flash_attn/flash_attn_backend.py | 681 ++++++++++++++++++ ring_flash_attn/llama3_flash_attn_varlen.py | 107 +-- .../llama_fwd_ring_bwd_flash_attn.py | 62 +- ring_flash_attn/ring_flash_attn.py | 100 +-- ring_flash_attn/ring_flash_attn_varlen.py | 111 ++- ring_flash_attn/stripe_flash_attn.py | 170 ++--- ring_flash_attn/zigzag_ring_flash_attn.py | 90 +-- .../zigzag_ring_flash_attn_varlen.py | 111 ++- 9 files changed, 947 insertions(+), 492 deletions(-) create mode 100644 ring_flash_attn/flash_attn_backend.py diff --git a/pyproject.toml b/pyproject.toml index ffdac15..b7249a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,13 @@ classifiers = [ "Operating System :: OS Independent", ] +[project.optional-dependencies] +# Selection stays a runtime decision via RING_FLASH_ATTN_BACKEND; this extra only +# documents the package that provides the `flash_attn.cute` backend. flash-attn-4 is +# not on PyPI as of this writing -- install it from a flash-attention checkout with +# `pip install -e /flash_attn/cute`. +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..0da8af3 --- /dev/null +++ b/ring_flash_attn/flash_attn_backend.py @@ -0,0 +1,681 @@ +"""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. + +Two 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``. + +``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_fa4_available", +] + +_ENV_VAR = "RING_FLASH_ATTN_BACKEND" +_VALID_BACKENDS = ("fa2", "fa4") +_FA4_INSTALL_HINT = ( + "install it with `pip install -e /flash_attn/cute`" +) + + +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 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_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() + + +# --------------------------------------------------------------------------- +# 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. + if is_fa4_available(): + extra = ( + "flash-attn-4 is installed, though -- set " + f"{_ENV_VAR}=fa4 to use it instead. Note that flash-attn-4 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 before either backend can load." + ) + else: + extra = f"flash-attn-4 was not found either; {_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 _set_window(params: dict, window_size) -> None: + """Apply the window in whichever arity the installed FA2 signature uses. + + FA < 2.7 takes a ``window_size`` tuple; FA >= 2.7 takes two ints. ``params`` + comes from signature introspection, so the key presence check is meaningful. + """ + 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] + + 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) + + +# --------------------------------------------------------------------------- +# 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 both 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 whose numerics have actually been validated against FA4. Everything +# else is routed through this seam but never checked against the cute kernels, so it +# is refused rather than silently run. See the plan's scope section. +_FA4_VALIDATED_VARIANTS = frozenset( + { + "llama3_flash_attn_varlen", + "llama_fwd_ring_bwd_flash_attn", + "ring_flash_attn_backward", + } +) + + +def check_variant_supported(variant: str) -> None: + """Refuse ring variants that have not been validated against the active backend. + + Called by the variant modules that were deliberately left out of FA4 scope. Under + fa2 (the default) this never fires. + """ + if BACKEND == "fa4" and variant not in _FA4_VALIDATED_VARIANTS: + raise NotImplementedError( + f"'{variant}' has not been validated against the flash-attn-4 backend, so " + f"it is refused rather than run with unverified numerics. Validated under " + f"fa4: {', '.join(sorted(_FA4_VALIDATED_VARIANTS))}. " + 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: From 937c703e29688d3ff0f73bfebb785003a477b280 Mon Sep 17 00:00:00 2001 From: eric-tc-wong Date: Fri, 31 Jul 2026 20:17:15 +0800 Subject: [PATCH 2/3] Add a flash-attn-3 backend FA3 ships as a separate flash_attn_3 package with its own interface. Add it as a third backend behind the existing seam, selected with RING_FLASH_ATTN_BACKEND=fa3. FA2 remains the default and is unchanged. FA3 is a much closer fit for ring attention than FA4: * LSE layout matches FA2 exactly -- (batch, nheads, seqlen) dense and (nheads, total_q) varlen -- so update_out_and_lse needs no transpose * dq/dk/dv are written in place (mutates_args), FA3's only mode, so there is no copy-back shim * the -1 window sentinel is native; mapping it to None as FA4 needs would fail FA3's torch.library schema, which types these as plain int * both ops are real torch.library custom ops with register_fake, so unlike FA4 they are traceable under torch.compile Three ways the signatures differ from FA2, each of which fails quietly rather than loudly if copied: * backward renames causal to is_causal * backward puts cu_seqlens_*/max_seqlen_* BEFORE dq/dk/dv, where FA2 puts dq/dk/dv straight after softmax_lse, so a positional call misbinds * dq/dk/dv are Optional and the C++ allocates throwaways when omitted, computing the gradient and discarding it with no error The adapter therefore calls keyword-only and always passes dq/dk/dv. Window arity and the causal kwarg name are resolved by introspecting the real signature via get_default_args, whose _init_fn fallback already handles CustomOpDef -- older FA3 releases (still present as vendored hopper/ trees in several envs here) use plain functions, a window_size tuple, and causal. Varlen is unified upstream: one op per direction, varlen selected by passing cu_seqlens_q. The four seam entry points are kept, with the varlen pair as thin wrappers. is_fa3_available probes flash_attn_3._C rather than the interface module, because a bare hopper/ source tree with no built extension otherwise looks available and dies at the first kernel call. _FA4_VALIDATED_VARIANTS becomes a per-backend dict; fa3 shares fa4's llama-path scope. _set_window moves to module scope, now shared by fa2 and fa3. FA3 cannot run on this machine (no flash_attn_3._C, and stock builds gencode sm90a), so it ships unrun. test_fa3_arg_binding covers what is checkable without a GPU: it loads the seam under fa3 against a checkout's real ops with the extension stubbed, and binds the adapter's own params dicts against the true signatures, asserting on BoundArguments identity rather than just that bind succeeded. FA2 verified unchanged: CPU tests plus 11 GPU tests and 7 compile-leg runs. Co-Authored-By: Claude Opus 5 --- pyproject.toml | 10 +- ring_flash_attn/flash_attn_backend.py | 415 +++++++++++++++++++++++--- test/test_fa3_arg_binding.py | 321 ++++++++++++++++++++ 3 files changed, 705 insertions(+), 41 deletions(-) create mode 100644 test/test_fa3_arg_binding.py diff --git a/pyproject.toml b/pyproject.toml index b7249a7..b8c7112 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,10 +18,12 @@ classifiers = [ ] [project.optional-dependencies] -# Selection stays a runtime decision via RING_FLASH_ATTN_BACKEND; this extra only -# documents the package that provides the `flash_attn.cute` backend. flash-attn-4 is -# not on PyPI as of this writing -- install it from a flash-attention checkout with -# `pip install -e /flash_attn/cute`. +# 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] diff --git a/ring_flash_attn/flash_attn_backend.py b/ring_flash_attn/flash_attn_backend.py index 0da8af3..a25bec9 100644 --- a/ring_flash_attn/flash_attn_backend.py +++ b/ring_flash_attn/flash_attn_backend.py @@ -41,14 +41,27 @@ "fa_forward", "fa_varlen_backward", "fa_varlen_forward", + "is_fa3_available", "is_fa4_available", ] _ENV_VAR = "RING_FLASH_ATTN_BACKEND" -_VALID_BACKENDS = ("fa2", "fa4") +_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. +_FA3_MODULE_CANDIDATES = ( + "flash_attn_3.flash_attn_interface", + "flash_attn_interface", + "hopper.flash_attn_interface", +) def _flash_attn_search_locations() -> Tuple[str, ...]: @@ -112,6 +125,27 @@ def is_fa4_available() -> bool: ) +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. @@ -130,6 +164,8 @@ def available_backends() -> Tuple[str, ...]: found.append("fa2") except Exception: pass + if is_fa3_available(): + found.append("fa3") if is_fa4_available(): found.append("fa4") return tuple(found) @@ -177,6 +213,23 @@ def _select_backend() -> str: 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 # --------------------------------------------------------------------------- @@ -194,16 +247,27 @@ def _select_backend() -> str: # 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(): - extra = ( - "flash-attn-4 is installed, though -- set " - f"{_ENV_VAR}=fa4 to use it instead. Note that flash-attn-4 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 before either backend can load." + _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"flash-attn-4 was not found either; {_FA4_INSTALL_HINT}." + 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 @@ -231,18 +295,6 @@ def _select_backend() -> str: supports_backward=True, ) - def _set_window(params: dict, window_size) -> None: - """Apply the window in whichever arity the installed FA2 signature uses. - - FA < 2.7 takes a ``window_size`` tuple; FA >= 2.7 takes two ints. ``params`` - comes from signature introspection, so the key presence check is meaningful. - """ - 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] - def _unpack_forward(outputs): """FA <= 2.6 returns 8 values, FA >= 2.7 returns 4. We want out and lse.""" if len(outputs) == 8: @@ -409,6 +461,278 @@ def fa_varlen_backward( _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. These surface only as + # TORCH_CHECK failures deep in C++, so surface them here instead. + try: + from flash_attn_config import CONFIG as _FA3_CONFIG + except Exception: + _FA3_CONFIG = None + if _FA3_CONFIG is not None: + _flags = _FA3_CONFIG.get("build_flags", {}) or {} + for _flag, _why in ( + ("FLASH_ATTENTION_DISABLE_VARLEN", "the varlen ring paths"), + ("FLASH_ATTENTION_DISABLE_LOCAL", "sliding-window attention"), + ("FLASH_ATTENTION_DISABLE_BACKWARD", "training (backward)"), + ): + 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." + ) + + 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. + 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 # --------------------------------------------------------------------------- @@ -654,28 +978,45 @@ def _check_unsupported(dropout_p, alibi_slopes) -> None: ) -# Ring variants whose numerics have actually been validated against FA4. Everything -# else is routed through this seam but never checked against the cute kernels, so it -# is refused rather than silently run. See the plan's scope section. -_FA4_VALIDATED_VARIANTS = frozenset( - { - "llama3_flash_attn_varlen", - "llama_fwd_ring_bwd_flash_attn", - "ring_flash_attn_backward", - } -) +# 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 have not been validated against the active backend. + """Refuse ring variants that are out of scope for the active backend. - Called by the variant modules that were deliberately left out of FA4 scope. Under - fa2 (the default) this never fires. + 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. """ - if BACKEND == "fa4" and variant not in _FA4_VALIDATED_VARIANTS: + in_scope = _VALIDATED_VARIANTS.get(BACKEND) + if in_scope is not None and variant not in in_scope: raise NotImplementedError( - f"'{variant}' has not been validated against the flash-attn-4 backend, so " - f"it is refused rather than run with unverified numerics. Validated under " - f"fa4: {', '.join(sorted(_FA4_VALIDATED_VARIANTS))}. " + 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/test/test_fa3_arg_binding.py b/test/test_fa3_arg_binding.py new file mode 100644 index 0000000..77a6076 --- /dev/null +++ b/test/test_fa3_arg_binding.py @@ -0,0 +1,321 @@ +"""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 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") + + +if __name__ == "__main__": + unittest.main() From f86553e8b1cdc5e617257cdf5c2c84f6ae8cc35e Mon Sep 17 00:00:00 2001 From: eric-tc-wong Date: Fri, 31 Jul 2026 20:22:34 +0800 Subject: [PATCH 3/3] Fix the fa3 build-flag guard, and make the CPU tests actually run The build-flag guard was dead. It looked up FLASH_ATTENTION_DISABLE_*, the spelling of the env vars that *set* the flags, but the generated config keys are FLASHATTENTION_DISABLE_* with no underscore after FLASH (hopper/setup.py:91), so every lookup missed and no build was ever rejected. Values are real booleans (setup.py:51 computes them as `== "TRUE"`), so plain truthiness is right once the keys are. The guard also only looked for a top-level `flash_attn_config`, which is what hopper/setup.py's py_modules installs -- but the layout that exposes flash_attn_3.flash_attn_interface keeps its config alongside, so the guard was a no-op on exactly the install this backend targets. Derive the config module from whichever interface module resolved, then fall back to the top-level name, and treat any unrecognised shape as "no opinion" rather than a failed import. Tests cover all of it: the guard fires on DISABLE_VARLEN and DISABLE_BACKWARD, stays quiet on a healthy build, and tolerates a list/str/None/unknown-keys config. Also assert that every kwarg the adapter passes but does not explicitly set equals its own declared default -- the property that makes get_default_args' pass-everything approach a semantic no-op, previously verified for FA2 but not FA3. test.sh's cpu_tests ran `python -m unittest test.`, which resolved to the CPython stdlib `test` package because there is no test/__init__.py here; all three entries errored on import and the block never tested anything. Run them as bare module names with test/ on the path, and add the fa3 test so it is not orphaned. Co-Authored-By: Claude Opus 5 --- ring_flash_attn/flash_attn_backend.py | 78 ++++++++++++---- test/test.sh | 10 +- test/test_fa3_arg_binding.py | 126 ++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 21 deletions(-) diff --git a/ring_flash_attn/flash_attn_backend.py b/ring_flash_attn/flash_attn_backend.py index a25bec9..8527c52 100644 --- a/ring_flash_attn/flash_attn_backend.py +++ b/ring_flash_attn/flash_attn_backend.py @@ -6,13 +6,21 @@ FA2 keyword set so the variant modules contain no version- or backend-specific branching. -Two backends are available: +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 @@ -57,6 +65,11 @@ # 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", @@ -491,24 +504,49 @@ def fa_varlen_backward( _flash_attn_forward = _fa3._flash_attn_forward _flash_attn_backward = _fa3._flash_attn_backward - # FA3 can compile out features ring attention depends on. These surface only as - # TORCH_CHECK failures deep in C++, so surface them here instead. + # 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: - from flash_attn_config import CONFIG as _FA3_CONFIG + _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: - _FA3_CONFIG = None - if _FA3_CONFIG is not None: - _flags = _FA3_CONFIG.get("build_flags", {}) or {} - for _flag, _why in ( - ("FLASH_ATTENTION_DISABLE_VARLEN", "the varlen ring paths"), - ("FLASH_ATTENTION_DISABLE_LOCAL", "sliding-window attention"), - ("FLASH_ATTENTION_DISABLE_BACKWARD", "training (backward)"), - ): - 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." - ) + # 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", @@ -611,6 +649,10 @@ def fa_forward( 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( @@ -956,7 +998,7 @@ def _copy_back(dq, dk, dv, out_dq, out_dk, out_dv) -> None: # --------------------------------------------------------------------------- -# Backend-agnostic guards (shared by both adapters, driven by CAPABILITIES) +# Backend-agnostic guards (shared by all adapters, driven by CAPABILITIES) # --------------------------------------------------------------------------- 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 index 77a6076..47e60e1 100644 --- a/test/test_fa3_arg_binding.py +++ b/test/test_fa3_arg_binding.py @@ -310,6 +310,53 @@ def test_adapter_rejects_dropout_and_alibi(self): 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 @@ -317,5 +364,84 @@ def test_out_of_scope_variant_refused_under_fa3(self): 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()