File size: 2,447 Bytes
8c8128e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """Forward-only FlashAttention-4 CuTe runtime for SM100-family GPUs."""
from __future__ import annotations
import os
import sys
import cutlass
from cutlass import cute
def _version_tuple(value: str) -> tuple[int, int]:
fields = value.split(".")
return tuple(int(field) for field in fields[:2]) # type: ignore[return-value]
_DSL_VERSION = _version_tuple(str(getattr(cutlass, "__version__", "0.0")))
if not ((4, 4) <= _DSL_VERSION < (4, 7)):
raise RuntimeError(
"fa4-cute-runtime requires nvidia-cutlass-dsl 4.4.x, 4.5.x, or 4.6.x; "
f"found {getattr(cutlass, '__version__', 'unknown')}"
)
os.environ.setdefault(
"CUTE_DSL_ARCH", "sm_101a" if _DSL_VERSION >= (4, 5) else "sm_110a"
)
os.environ.setdefault("FLASH_ATTENTION_ARCH", "sm_100a")
# CUTLASS DSL 4.6 promoted these public types out of ``cute.core``. The
# vendored FA4 sources use the 4.4/4.5 annotation paths; restoring the aliases
# keeps those annotations importable without changing generated kernels.
if _DSL_VERSION >= (4, 6):
if not hasattr(cute.core, "ThrMma"):
cute.core.ThrMma = cute.ThrMma
if not hasattr(cute.core, "ThrCopy"):
cute.core.ThrCopy = cute.ThrCopy
if not hasattr(cute, "make_fragment"):
cute.make_fragment = cute.make_rmem_tensor
# Kernel Hub imports a noarch variant under a content-derived module name after
# flattening this package into the variant root. Preserve the public package
# alias used by the vendored FA4 absolute imports.
sys.modules.setdefault("fa4_cute_runtime", sys.modules[__name__])
# kernel-builder copies only the directory matching ``general.name`` for a
# torch-noarch package. The private vendor and quack subset therefore live
# inside this module and use package-qualified imports.
from .flashrt_fa4.cute import flash_attn_func, flash_attn_varlen_func # noqa: E402
from .flashrt_fa4.cute.interface_fwd_sm100 import _flash_attn_fwd # noqa: E402
def forward_static(
q,
k,
v,
out,
*,
softmax_scale=None,
causal=False,
pack_gqa=None,
seqused_k=None,
):
"""Run forward attention into a caller-owned output tensor."""
result, _ = _flash_attn_fwd(
q,
k,
v,
softmax_scale=softmax_scale,
causal=causal,
pack_gqa=pack_gqa,
seqused_k=seqused_k,
out=out,
)
return result
__all__ = ["flash_attn_func", "flash_attn_varlen_func", "forward_static"]
|