feat: AMD aiter Flash Attention via source-level SDPA rewrite in compiler.py

unsloth_zoo/compiler.py

@@ -373,6 +373,122 @@ def replace_with_grouped_query_attention(module, source):
 pass
 
 
+
+
+def replace_sdpa_with_amd_aiter(source):
+    """
+    For AMD ROCm with amd-aiter installed: replace scaled_dot_product_attention
+    calls with amd-aiter Flash Attention in the compiled source.
+
+    Activation requirements (ALL must hold):
+      1. get_amd_attention_implementation() == "amd_aiter"
+         — requires ROCm >= 7.0, aiter installed, gfx942 or gfx950 arch
+      2. rest_args is exactly ", is_causal=True" (optional trailing comma/whitespace)
+         — any other argument (positional or keyword) leaves the call unchanged
+      3. Call is not immediately followed by a method chain, subscript, or operator
+         — .transpose(), * scale, [0], etc. are rejected (negative lookahead)
+      4. At runtime, _aiter_ok additionally checks:
+         - q/k seq lengths equal (causal mask alignment differs between SDPA and aiter)
+         - input dtype is float16 or bfloat16 (aiter does not support float32)
+         - no input has requires_grad=True (aiter asserts return_lse for backward;
+           we omit return_lse, so the fast path is inference-only)
+      5. aiter call wrapped in try/except with SDPA fallback for unsupported head dims
+         or JIT failures
+
+    Known limitation: the rewrite currently only matches user-code SDPA calls with
+    the bare is_causal=True form. Unsloth's own compiled SDPA paths emit attn_mask=
+    or is_causal=<variable> which are correctly rejected by guard 2. Wiring this
+    to Unsloth's SDPA shim requires a separate follow-up PR.
+
+    No-op on NVIDIA and on ROCm without amd-aiter, ROCm < 7.0, or unsupported arch.
+    """
+    # Import inside function — avoids compile-time NameError in caller scope
+    from unsloth_zoo.device_type import get_amd_attention_implementation
+    if get_amd_attention_implementation() != "amd_aiter":
+        return source
+
+    import re
+
+    # Pattern:  out = <prefix.>scaled_dot_product_attention(q, k, v, ...)
+    # Excludes: disable_compile_scaled_dot_product_attention (opt-out shim)
+    # Uses ((?:[^)]|\([^)]*\))*) to handle one level of nested parens in args
+    sdpa_call_pattern = (
+        # Negative lookbehind: reject attribute assignments like self.attn_output = ...
+        # which would leave "self." dangling before the generated if/else block.
+        r"([ \t]*)(?<![.\w])([A-Za-z_][A-Za-z0-9_]*)[ \t]*=[ \t]*"
+        r"(?:(?:torch\.nn\.functional|F|nn\.functional)\.)?scaled_dot_product_attention"
+        r"\("
+        r"[ \t\n]*([A-Za-z_][A-Za-z0-9_]*)[ \t]*,"
+        r"[ \t\n]*([A-Za-z_][A-Za-z0-9_]*)[ \t]*,"
+        r"[ \t\n]*([A-Za-z_][A-Za-z0-9_]*)[ \t]*"
+        r"((?:[^)]|\([^)]*\))*)\)"
+        r"(?![ \t]*[^\s#])"  # reject any non-ws/non-comment suffix (* scale, [0], .method)
+    )
+
+    def aiter_replacement(m):
+        indent  = m.group(1)
+        out_var = m.group(2)
+        q_var   = m.group(3)
+        k_var   = m.group(4)
+        v_var   = m.group(5)
+        rest_args = m.group(6)
+
+        # Only bare is_causal=True (optional trailing comma/whitespace) is rewritten.
+        # Any other argument — positional or keyword — is left unchanged;
+        # positional args like (q, k, v, None, 0.1, ...) carry semantics aiter ignores.
+        rest_args_stripped = rest_args.strip().lstrip(",").strip().rstrip(",").strip()
+        if not re.fullmatch(r"is_causal\s*=\s*True", rest_args_stripped):
+            return m.group(0)
+
+        # rest_args is ", is_causal=True" here (fullmatch guarantees it).
+        # Reconstruct with leading comma for the SDPA fallback.
+        _kw = rest_args.strip().lstrip(",").strip().rstrip(",").strip()
+        rest_args_formatted = (", " + _kw) if _kw else ""
+
+        # Generated code is exec()'d in isolation — all symbols must be imported inside.
+        lines = [
+            f"{indent}# AMD aiter Flash Attention (MFMA + causal tile skip)",
+            f"{indent}# Requires: pip install amd-aiter  (ROCm >= 7.0)",
+            f"{indent}from unsloth_zoo.device_type import get_amd_flash_attn_func as _get_aiter_fn",
+            f"{indent}_aiter_fn = _get_aiter_fn()",
+            f"{indent}_aiter_ok = (",
+            f"{indent}    _aiter_fn is not None",
+            f"{indent}    and {q_var}.shape[-2] == {k_var}.shape[-2]",  # equal seq lengths
+            f"{indent}    and {q_var}.dtype in (torch.float16, torch.bfloat16)",  # dtype guard
+            # Inference-only: aiter.flash_attn_func asserts return_lse=True when
+            # any input has requires_grad=True (aiter/ops/mha.py:2502,2536). We
+            # omit return_lse, so calling it during training raises AssertionError.
+            # Must check all three of q/k/v: a LoRA targeting k_proj/v_proj but
+            # not q_proj leaves k and v with grad while q.requires_grad is False.
+            f"{indent}    and not ({q_var}.requires_grad or {k_var}.requires_grad or {v_var}.requires_grad)",
+            f"{indent})",
+            f"{indent}if _aiter_ok:",
+            f"{indent}    _aiter_q = {q_var}.transpose(1, 2)",
+            f"{indent}    _aiter_k = {k_var}.transpose(1, 2)",
+            f"{indent}    _aiter_v = {v_var}.transpose(1, 2)",
+            # _call_aiter_safe returns None on failure; avoids try/except in
+            # generated code which breaks torch.compile(fullgraph=True).
+            f"{indent}    def _call_aiter_safe(_fn, _q, _k, _v):",
+            f"{indent}        try: return _fn(_q, _k, _v, causal=True)",
+            f"{indent}        except Exception: return None",
+            f"{indent}    _aiter_out = _call_aiter_safe(_aiter_fn, _aiter_q, _aiter_k, _aiter_v)",
+            f"{indent}    if _aiter_out is not None:",
+            f"{indent}        {out_var} = _aiter_out.transpose(1, 2)",
+            f"{indent}    else:",
+            f"{indent}        {out_var} = torch.nn.functional.scaled_dot_product_attention(",
+            f"{indent}            {q_var}, {k_var}, {v_var}{rest_args_formatted})",
+            # else: _aiter_ok is False — fall back to plain SDPA directly
+            f"{indent}else:",
+            f"{indent}    {out_var} = torch.nn.functional.scaled_dot_product_attention(",
+            f"{indent}        {q_var}, {k_var}, {v_var}{rest_args_formatted})",
+        ]
+        return "\n".join(lines)
+
+    new_source, n = re.subn(sdpa_call_pattern, aiter_replacement, source, flags=re.DOTALL)
+    return new_source
+
+
+
 def _get_compile_folder(use_tempfile=False):
     global UNSLOTH_COMPILE_LOCATION
     global UNSLOTH_COMPILE_USE_TEMP
@@ -3809,6 +3925,15 @@ def _def_pos(name):
                 disabled_scaled_dot_product_attention_modules.append(module)
             pass
         pass
+        # AMD ROCm: replace SDPA with amd-aiter Flash Attention if available.
+        # Note: this fires on the model's compiled source after Unsloth's SDPA
+        # rewriting. The Unsloth-rewritten SDPA calls carry attn_mask= or
+        # is_causal=is_causal (variable, not literal True) and will not match
+        # the aiter guards. User-added model code with bare
+        # scaled_dot_product_attention(q, k, v, is_causal=True) will match.
+        # A future PR can also rewrite the Unsloth SDPA shim to emit
+        # is_causal=True literally where provably safe.
+        new_source = replace_sdpa_with_amd_aiter(new_source)
         scaled_dot_product_attention_modules[module] = new_source
     pass
 

unsloth_zoo/device_type.py

@@ -14,6 +14,8 @@
 
 __all__ = [
     "is_hip",
+    "get_amd_attention_implementation",
+    "get_amd_flash_attn_func",
     "get_device_type",
     "DEVICE_TYPE",
     "DEVICE_TYPE_TORCH",
@@ -213,6 +215,138 @@ def is_hip():
     return bool(getattr(getattr(torch, "version", None), "hip", None))
 pass
 
+
+@functools.cache
+def _detect_gfx_arch():
+    """Return the GPU architecture string (e.g. 'gfx942') or None if undetectable.
+
+    Uses the PyTorch-active device first so that HIP_VISIBLE_DEVICES and multi-arch
+    hosts are handled correctly. Subprocess probes (rocminfo, hipconfig) scan
+    system-wide output and may return a different arch than the active device.
+    """
+    import re as _re
+    # First: query the currently active PyTorch device — always returns the right
+    # arch for the GPU PyTorch is using, respects HIP_VISIBLE_DEVICES.
+    try:
+        if torch.cuda.is_available():
+            dev = torch.cuda.current_device()
+            arch = torch.cuda.get_device_properties(dev).gcnArchName
+            # gcnArchName may include suffixes e.g. "gfx942:sramecc+:xnack-"
+            m = _re.match(r"gfx[0-9a-f]+", arch)
+            if m:
+                return m.group(0)
+    except Exception:
+        pass
+    # Subprocess fallbacks for unusual builds where PyTorch cannot query the device.
+    import subprocess
+    try:
+        r = subprocess.run(["rocminfo"], capture_output=True, text=True, timeout=10)
+        m = _re.search(r"gfx[0-9a-f]+", r.stdout)
+        if m:
+            return m.group(0)
+    except Exception:
+        pass
+    try:
+        r = subprocess.run(["hipconfig", "--full"], capture_output=True, text=True, timeout=10)
+        m = _re.search(r"gfx[0-9a-f]+", r.stdout)
+        if m:
+            return m.group(0)
+    except Exception:
+        pass
+    return None
+pass
+
+
+@functools.cache
+def get_amd_attention_implementation():
+    """
+    Return the best available attention implementation for AMD ROCm.
+
+    Priority:
+    1. "amd_aiter" — AMD aiter (pip install amd-aiter, ROCm >= 7.0 required)
+    2. "sdpa"      — PyTorch SDPA via MIOpen (always available on ROCm)
+
+    Returns: "amd_aiter" | "sdpa"
+    """
+    if not is_hip():
+        return "sdpa"
+
+    # Gate on ROCm >= 7.0 (amd-aiter has hard ABI dep on libamdhip64.so.7).
+    # Prefer torch.version.hip — always present in ROCm PyTorch wheels and not
+    # affected by rocm-smi which reports the kernel driver version, not the
+    # ROCm runtime version (the two can differ on wheel-only installs).
+    try:
+        _hip_ver = getattr(torch.version, "hip", None)
+        if _hip_ver is not None:
+            _hip_major = int(str(_hip_ver).split(".")[0])
+            if _hip_major < 7:
+                return "sdpa"
+        else:
+            # Fallback for unusual builds without torch.version.hip
+            rocm_version = _detect_rocm_major_minor()
+            if rocm_version is not None:
+                if int(rocm_version.split(".")[0]) < 7:
+                    return "sdpa"
+    except Exception:
+        return "sdpa"
+
+    # aiter's CK/ASM kernels are CDNA-only: gfx942 (MI300X/MI325X) and gfx950
+    # (MI355X). flash_attn_func imports on all archs but only runs on these two.
+    try:
+        _gfx = _detect_gfx_arch()
+        _AITER_SUPPORTED_ARCHS = ("gfx942", "gfx950")  # CDNA3, CDNA4
+        if _gfx not in _AITER_SUPPORTED_ARCHS:
+            return "sdpa"
+    except Exception:
+        return "sdpa"
+
+    # Check for amd-aiter (AMD AI Tensor Engine for ROCm)
+    try:
+        import importlib.util  # must import submodule explicitly (not bare importlib)
+        if importlib.util.find_spec("aiter") is not None:
+            import aiter as _aiter
+            # Validate: AMD AI tensor engine exposes the functional flash_attn_func API.
+            # FlashAttnFunc (class API) requires 13+ positional args and cannot be
+            # wrapped safely; only accept the simpler flash_attn_func functional API.
+            if hasattr(_aiter, "flash_attn_func"):
+                return "amd_aiter"
+    except Exception:
+        pass
+
+    return "sdpa"
+
+
+@functools.cache
+def get_amd_flash_attn_func():
+    """
+    Return the amd-aiter flash attention function, or None if unavailable.
+
+    Checks both naming conventions across amd-aiter versions:
+    - flash_attn_func (newer, lowercase)
+    - FlashAttnFunc   (older, class-based callable)
+
+    Call signature: func(q, k, v, causal=True)
+    Shapes: q/k/v = (batch, seqlen, nheads, headdim), float16 or bfloat16
+    Returns: output tensor, same shape as q
+    """
+    if get_amd_attention_implementation() != "amd_aiter":
+        return None
+    try:
+        import aiter as _aiter
+        if hasattr(_aiter, "flash_attn_func"):
+            return _aiter.flash_attn_func
+        # FlashAttnFunc is a torch.autograd.Function whose .apply() requires 13+
+        # positional arguments (dropout_p, softmax_scale, causal, window_size,
+        # bias, alibi_slopes, deterministic, return_lse, return_softmax,
+        # is_grad_enabled, ...) — see ROCm/aiter:aiter/ops/mha.py.
+        # We cannot safely wrap it without knowing the required defaults for the
+        # installed aiter version.  Return None so callers fall back to SDPA.
+        # (FlashAttnFunc environments should expose flash_attn_func in aiter >= 0.7)
+    except Exception:
+        pass
+    return None
+
+
 @functools.cache
 def get_device_type():
     if _IS_MLX: