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

unsloth_zoo/compiler.py

@@ -373,6 +373,109 @@ 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.
+
+    Safety requirements (all must hold for the aiter path to activate):
+      1. get_amd_attention_implementation() == "amd_aiter"
+      2. is_causal=True as a literal (not a complex expression)
+      3. No extra SDPA args that aiter does not support:
+         attn_mask, dropout_p, scale, enable_gqa
+      4. q/k seq lengths equal (non-square causal has different bias alignment)
+      5. Input dtype is float16 or bfloat16 (aiter does not support float32)
+
+    No-op on NVIDIA and on ROCm without amd-aiter or ROCm < 7.0.
+    """
+    # 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 = (
+        r"([ \t]*)([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"((?:[^)]|\([^)]*\))*)\)"
+    )
+
+    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 rewrite when rest_args is exactly ", is_causal=True" (with optional
+        # whitespace). aiter accepts only (q, k, v, causal=True) — any additional
+        # argument (positional or keyword) may carry semantics the kernel ignores
+        # (e.g. positional dropout_p, attn_mask, scale, enable_gqa). Keyword-name
+        # checking alone is insufficient: positional calls like
+        # scaled_dot_product_attention(q, k, v, None, 0.1, is_causal=True) contain
+        # no keyword named "dropout_p" yet silently drop the positional dropout value.
+        # Accepting only the bare causal-only form guarantees correctness.
+        # Strip leading/trailing whitespace, leading comma, and optional trailing
+        # comma (Black/multiline formatting). Then fullmatch the bare keyword.
+        rest_args_stripped = rest_args.strip().lstrip(",").strip().rstrip(",").strip()
+        if not re.fullmatch(r"is_causal\s*=\s*True", rest_args_stripped):
+            # rest_args has more than just is_causal=True — leave call unchanged
+            return m.group(0)
+
+        # Clean leading newlines from rest_args
+        rest_args_clean = rest_args.lstrip("\r\n").strip()
+        # rest_args already starts with a comma when non-empty; do NOT add another
+        # rest_args is ", is_causal=True" at this point — pass it through to the
+        # SDPA fallback so the original call semantics are preserved.
+        rest_args_clean = rest_args.lstrip("\r\n").strip()
+        rest_args_formatted = rest_args_clean  # already starts with comma
+
+        # Generate the replacement block.
+        # All symbols used at runtime are imported inside the generated code so
+        # that the exec()'d module namespace is self-contained (no NameError).
+        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
+            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)",
+            # Wrap in try/except: unsupported head dim, arch mismatch, or JIT build
+            # failure inside flash_attn_func must fail-soft to SDPA, not crash training.
+            f"{indent}    try:",
+            f"{indent}        {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal=True).transpose(1, 2)",
+            f"{indent}    except Exception:",
+            f"{indent}        {out_var} = torch.nn.functional.scaled_dot_product_attention(",
+            f"{indent}            {q_var}, {k_var}, {v_var}{rest_args_formatted})",
+            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 +3912,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,135 @@ 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."""
+    import subprocess, re as _re
+    try:
+        # rocminfo is the most reliable source
+        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:
+        # hipconfig fallback
+        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
+    # Final fallback: query torch directly — available in all ROCm wheel installs
+    # even when rocminfo and hipconfig are absent (e.g. runtime-only containers).
+    try:
+        if torch.cuda.is_available():
+            arch = torch.cuda.get_device_properties(0).gcnArchName
+            # gcnArchName may be e.g. "gfx942:sramecc+:xnack-" — extract gfxNNNN
+            import re as _re2
+            m = _re2.match(r"gfx[0-9a-f]+", arch)
+            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"
+
+    # Gate on GPU architecture — aiter's CK/ASM kernels are CDNA-only.
+    # flash_attn_func is always importable on any arch where aiter installs,
+    # but the underlying kernels only work on gfx942 (MI300X/MI325X) and
+    # gfx950 (MI355X). Consumer RDNA (gfx1100, gfx1200, etc.) is not supported.
+    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"  # fail safe: unknown arch → use 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: