feat: AMD aiter Flash Attention via source-level SDPA rewrite in compiler.py - #920
feat: AMD aiter Flash Attention via source-level SDPA rewrite in compiler.py#920zhihuidu-amd wants to merge 12 commits into
Conversation
…er.py)
Properly integrates amd-aiter Flash Attention by hooking into unsloth-zoo's
existing source-rewriting pipeline in compiler.py. No-op on NVIDIA or
AMD ROCm without amd-aiter / ROCm < 7.0.
## Approach
unsloth-zoo compiles attention forwards by rewriting their source code as
strings (inspect.getsource() + regex transforms). This PR adds a new
source-level transform that replaces scaled_dot_product_attention calls
with the amd-aiter kernel when available.
The transform runs in the existing SDPA processing loop in
unsloth_compile_transformers(), after the existing causal mask and GQA
rewrites. It is completely invisible to NVIDIA/CUDA — get_amd_attention_
implementation() returns 'sdpa' immediately when not is_hip().
## Changes
### unsloth_zoo/device_type.py
- get_amd_attention_implementation(): detects amd-aiter with ROCm >= 7.0 gate
and importlib.util (not bare importlib, which fails Python 3.13+)
- get_amd_flash_attn_func(): returns flash_attn_func/FlashAttnFunc or None
(handles both naming conventions across amd-aiter versions)
### unsloth_zoo/compiler.py
- replace_sdpa_with_amd_aiter(source): new source-rewriting function.
Matches scaled_dot_product_attention calls with regex, rewrites them as:
_aiter_q = q.transpose(1,2) # (B,H,S,D) — aiter convention
_aiter_k = k.transpose(1,2)
_aiter_v = v.transpose(1,2)
_aiter_fn = get_amd_flash_attn_func()
if _aiter_fn is not None:
attn_output = _aiter_fn(_aiter_q,_aiter_k,_aiter_v, causal=...).transpose(1,2)
else:
attn_output = torch.nn.functional.scaled_dot_product_attention(...)
The transpose converts transformers (B,H,S,D) to aiter (B,S,H,D) convention.
Falls back to standard SDPA if aiter is unavailable at runtime.
- Hooked into the scaled_dot_product_attention_modules processing loop.
## Expected performance (MI300X/MI325X gfx942, ROCm >= 7.0)
Prefill seq=2048, H=32, D=128:
amd-aiter full MFMA: 51.5 TFLOPS (+5.9x vs PyTorch SDPA)
amd-aiter + causal tile skip: 71.3 TFLOPS (+8.2x, causal only)
PyTorch SDPA (current): ~5.7 TFLOPS
## Requirements
- pip install amd-aiter (https://github.com/ROCm/aiter)
- ROCm >= 7.0 (hard ABI requirement of amd-aiter)
- Automatic activation — no user code changes required
## Testing
- NVIDIA: no change — all new code gated behind is_hip() check
- AMD without amd-aiter: no change — falls back to SDPA in the rewritten source
- AMD with amd-aiter (ROCm >= 7.0): aiter kernel activated at compile time
Note: cannot run live benchmark on Alola (ROCm 6.2); performance numbers
from AMD's validated benchmarks on MI300X gfx942 (KB entries).
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for AMD ROCm by replacing PyTorch's scaled_dot_product_attention (SDPA) with AMD's amd-aiter Flash Attention implementation during compilation when available. The review feedback highlights critical compile-time and runtime NameError bugs due to missing imports of the helper functions, potential regex failures with nested parentheses, and incorrect routing of non-causal attention to the causal-only fast path. A refactored replacement function is suggested to address these issues.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if get_amd_attention_implementation() != "amd_aiter": | ||
| return source | ||
|
|
||
| # Pattern: capture the SDPA call and all its arguments | ||
| # Matches: attn_output = torch.nn.functional.scaled_dot_product_attention( | ||
| # query_states, key_states, value_states, | ||
| # attn_mask=..., dropout_p=..., is_causal=... | ||
| # ) | ||
| sdpa_call_pattern = ( | ||
| r"([ ]*)([A-Za-z_][A-Za-z0-9_]*)[ ]*=[ ]*" | ||
| r"(?:[A-Za-z_.]*scaled_dot_product_attention)" | ||
| r"\(" | ||
| r"[ | ||
| ]*([A-Za-z_][A-Za-z0-9_]*)[ ]*," # query | ||
| r"[ | ||
| ]*([A-Za-z_][A-Za-z0-9_]*)[ ]*," # key | ||
| r"[ | ||
| ]*([A-Za-z_][A-Za-z0-9_]*)[ ]*" # value | ||
| r"([^)]*)\)" # remaining args | ||
| ) | ||
|
|
||
| import re | ||
|
|
||
| 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) | ||
|
|
||
| # Extract is_causal from rest_args (default False if not found) | ||
| causal_match = re.search(r"is_causal\s*=\s*(True|False|[A-Za-z_][A-Za-z0-9_]*)", rest_args) | ||
| is_causal = causal_match.group(1) if causal_match else "False" | ||
|
|
||
| # Build aiter call: | ||
| # aiter expects (B, S, H, D) but transformers uses (B, H, S, D) -> transpose 1,2 | ||
| return ( | ||
| f"{indent}# AMD aiter Flash Attention (MFMA + fast softmax + causal tile skip) | ||
| " | ||
| f"{indent}# Requires: pip install amd-aiter (ROCm >= 7.0) | ||
| " | ||
| f"{indent}_aiter_q = {q_var}.transpose(1, 2) # (B,H,S,D) -> (B,S,H,D) | ||
| " | ||
| f"{indent}_aiter_k = {k_var}.transpose(1, 2) | ||
| " | ||
| f"{indent}_aiter_v = {v_var}.transpose(1, 2) | ||
| " | ||
| f"{indent}_aiter_fn = get_amd_flash_attn_func() | ||
| " | ||
| f"{indent}if _aiter_fn is not None: | ||
| " | ||
| f"{indent} {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal={is_causal}).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})" | ||
| ) | ||
|
|
||
| new_source, n = re.subn(sdpa_call_pattern, aiter_replacement, source, flags=re.DOTALL) | ||
| if n > 0: | ||
| pass # Successfully replaced SDPA with aiter | ||
| return new_source |
There was a problem hiding this comment.
There are two critical NameError bugs, some robustness issues, and compliance concerns with repository rules in this implementation:
- Compile-time NameError: get_amd_attention_implementation is called on line 398 but is never imported in unsloth_zoo/compiler.py.
- Runtime NameError: get_amd_flash_attn_func is used in the generated code on line 446, but it is not imported in the compiled module, which will fail when the compiled module is loaded and executed.
- Regex Robustness & Line Structure Preservation: The regex pattern ([^)]) for rest_args will fail if any argument contains nested parentheses. We can use ((?:[^)]|([^)]))*)) to safely support one level of nested parentheses. Additionally, when patching code via regex, we must ensure replacement strings do not introduce unintended blank lines by checking if captured suffixes (like rest_args) already contain leading newlines.
- Causal Fast-Path Routing: When optimizing attention mechanisms by routing to fast paths (like aiter FlashAttention) that assume causal masking, we must ensure that bidirectional attention layers (where is_causal is False and attention_mask is None) are not incorrectly forced to be causal. We should explicitly verify that the call is causal before routing to the causal-only fast path.
We can resolve all of these issues cleanly by importing get_amd_attention_implementation at the start of the function, importing get_amd_flash_attn_func directly inside the generated code block, checking if is_causal is True before routing to _aiter_fn, cleaning up leading newlines in rest_args, and updating the regex patterns.
from unsloth_zoo.device_type import get_amd_attention_implementation
if get_amd_attention_implementation() != "amd_aiter":
return source
sdpa_call_pattern = (
r"([ \t]*)([A-Za-z_][A-Za-z0-9_]*)[ \t]*=[ \t]*"
r"(?:[A-Za-z_.]*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"((?:[^)]|\([^)]*\))*)\)"
)
import re
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)
causal_match = re.search(r"is_causal\s*=\s*([^,]+)", rest_args)
is_causal = causal_match.group(1).strip() if causal_match else "False"
rest_args_clean = rest_args.lstrip('\r\n')
if rest_args_clean != rest_args:
rest_args_formatted = ", " + rest_args_clean.strip() if rest_args_clean.strip() else ""
else:
rest_args_formatted = rest_args
return (
f"{indent}# AMD aiter Flash Attention (MFMA + fast softmax + causal tile skip)\n"
f"{indent}# Requires: pip install amd-aiter (ROCm >= 7.0)\n"
f"{indent}from unsloth_zoo.device_type import get_amd_flash_attn_func\n"
f"{indent}_aiter_fn = get_amd_flash_attn_func()\n"
f"{indent}if _aiter_fn is not None and {is_causal}:\n"
f"{indent} _aiter_q = {q_var}.transpose(1, 2)\n"
f"{indent} _aiter_k = {k_var}.transpose(1, 2)\n"
f"{indent} _aiter_v = {v_var}.transpose(1, 2)\n"
f"{indent} {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal=True).transpose(1, 2)\n"
f"{indent}else:\n"
f"{indent} {out_var} = torch.nn.functional.scaled_dot_product_attention(\n"
f"{indent} {q_var}, {k_var}, {v_var}{rest_args_formatted})"
)
new_source, n = re.subn(sdpa_call_pattern, aiter_replacement, source, flags=re.DOTALL)
if n > 0:
pass
return new_sourceReferences
- When patching code via regex, ensure replacement strings do not introduce unintended blank lines by checking if captured suffixes already contain leading newlines. Verify the fix using tests that check for line structure preservation and, if applicable, live round-trip checks against the upstream source.
- When optimizing attention mechanisms by routing to fast paths (like FlashAttention) that assume causal masking, ensure that bidirectional attention layers (where is_causal is False and attention_mask is None) are not incorrectly forced to be causal. Explicitly verify that the call is causal or has an explicit mask before routing to a causal-only fast path.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 57d4cf4238
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| r"([ ]*)([A-Za-z_][A-Za-z0-9_]*)[ ]*=[ ]*" | ||
| r"(?:[A-Za-z_.]*scaled_dot_product_attention)" | ||
| r"\(" | ||
| r"[ |
There was a problem hiding this comment.
This raw string is split by a physical newline, so compiler.py is not parseable (python -m py_compile unsloth_zoo/compiler.py stops here with SyntaxError: unterminated string literal). Any path that imports or uses the compiler now fails before the AMD/ROCm gate can no-op; the same escaping issue appears in the f-strings that build the replacement body below.
Useful? React with 👍 / 👎.
| Only activates on AMD ROCm when get_amd_attention_implementation() == "amd_aiter". | ||
| No-op on NVIDIA and on ROCm without amd-aiter or ROCm < 7.0. | ||
| """ | ||
| if get_amd_attention_implementation() != "amd_aiter": |
There was a problem hiding this comment.
replace_sdpa_with_amd_aiter() is called for each SDPA module, but compiler.py only imports DEVICE_TYPE from the package and never imports get_amd_attention_implementation. Once the syntax error is fixed, any SDPA-capable model that enters this loop will raise NameError here even on NVIDIA/ROCm setups where this should be a no-op.
Useful? React with 👍 / 👎.
| " | ||
| f"{indent}_aiter_v = {v_var}.transpose(1, 2) | ||
| " | ||
| f"{indent}_aiter_fn = get_amd_flash_attn_func() |
There was a problem hiding this comment.
The replacement emits get_amd_flash_attn_func() into the compiled cache file, but create_new_function() builds that file with fixed imports plus symbols from the Transformers modeling module; it does not import this helper from unsloth_zoo.device_type. On ROCm+aiter, the rewritten forward will therefore fail with NameError at runtime as soon as it reaches the aiter branch.
Useful? React with 👍 / 👎.
| " | ||
| f"{indent}if _aiter_fn is not None: | ||
| " | ||
| f"{indent} {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal={is_causal}).transpose(1, 2) |
There was a problem hiding this comment.
The aiter branch uses only the extracted is_causal flag and discards the rest of the original SDPA arguments. On AMD systems with aiter, any attention call that carries a padding/sliding attn_mask, training dropout_p, custom scale, or enable_gqa will silently compute different attention than the original SDPA call instead of falling back to the SDPA path that still receives rest_args.
Useful? React with 👍 / 👎.
| # ) | ||
| sdpa_call_pattern = ( | ||
| r"([ ]*)([A-Za-z_][A-Za-z0-9_]*)[ ]*=[ ]*" | ||
| r"(?:[A-Za-z_.]*scaled_dot_product_attention)" |
There was a problem hiding this comment.
This pattern matches any callable name ending in scaled_dot_product_attention, so on ROCm+aiter it also rewrites the disable_compile_scaled_dot_product_attention(...) calls inserted just above when the compiler decided an SDPA module was unsafe. That undoes the existing opt-out for those modules and reintroduces the SDPA/aiter call the preceding branch was trying to disable.
Useful? React with 👍 / 👎.
| " | ||
| f"{indent}if _aiter_fn is not None: | ||
| " | ||
| f"{indent} {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal={is_causal}).transpose(1, 2) |
There was a problem hiding this comment.
This branch passes q/k/v to aiter unconditionally, but the repo has forced-float32 attention patches that deliberately call SDPA with float32 tensors (for example unsloth_zoo/temporary_patches/qwen3_moe_float32.py:153-155 and :218-226) to avoid fp16 overflows, while the new helper documents aiter inputs as float16 or bfloat16. On AMD+aiter with UNSLOTH_FORCE_FLOAT32=1, the rewrite sends those float32 tensors into aiter instead of the SDPA fallback, causing unsupported-dtype failures or bypassing the stability fix.
Useful? React with 👍 / 👎.
| if hasattr(_aiter, "flash_attn_func"): | ||
| return _aiter.flash_attn_func | ||
| if hasattr(_aiter, "FlashAttnFunc"): | ||
| return _aiter.FlashAttnFunc |
There was a problem hiding this comment.
When an installed aiter build exposes only FlashAttnFunc, this returns the autograd Function class itself. The rewritten forward later calls the returned object as _aiter_fn(q, k, v, causal=...), but that class-based API needs .apply(...) with the full positional argument list, so those environments crash instead of using aiter or falling back to SDPA.
Useful? React with 👍 / 👎.
| rest_args = m.group(6) | ||
|
|
||
| # Extract is_causal from rest_args (default False if not found) | ||
| causal_match = re.search(r"is_causal\s*=\s*(True|False|[A-Za-z_][A-Za-z0-9_]*)", rest_args) |
There was a problem hiding this comment.
This regex captures only a literal or a single identifier after is_causal=, so common SDPA calls such as is_causal=self.is_causal and attention_mask is None and q_len > 1 are rewritten with causal=self. On AMD+aiter that makes the aiter kernel see the module object as a truthy causal flag instead of the original boolean expression, producing causal attention in cases where the original SDPA call would not.
Useful? React with 👍 / 👎.
| ]*([A-Za-z_][A-Za-z0-9_]*)[ ]*," # key | ||
| r"[ | ||
| ]*([A-Za-z_][A-Za-z0-9_]*)[ ]*" # value | ||
| r"([^)]*)\)" # remaining args |
There was a problem hiding this comment.
The rest_args group stops at the first ), not the closing parenthesis of scaled_dot_product_attention. On AMD+aiter, any SDPA call with a nested expression in its keywords, for example scale=getattr(self, "scaling", None), is only partially replaced and leaves the remaining arguments after the generated fallback, yielding an invalid compiled forward instead of a working SDPA fallback.
Useful? React with 👍 / 👎.
1. Compile-time NameError: add 'from unsloth_zoo.device_type import
get_amd_attention_implementation' at top of replace_sdpa_with_amd_aiter().
Previously called at module scope where the import was absent.
2. Runtime NameError: add 'from unsloth_zoo.device_type import
get_amd_flash_attn_func' to the generated code string itself.
The compiled module runs in its own namespace; the outer import
is not visible there.
3. Regex robustness: update rest_args pattern from ([^)]*) to
((?:[^)]|\([^)]*\))*) to handle one level of nested parentheses
without mismatching. Also strip leading \r\n from rest_args before
formatting to avoid blank lines in generated code.
4. Causal fast-path routing: change 'if _aiter_fn is not None:' to
'if _aiter_fn is not None and {is_causal}:'. amd-aiter flash_attn
assumes causal masking; bidirectional attention (is_causal=False,
attn_mask=None) must not be routed to the causal-only fast path.
|
All 4 issues from the Gemini Code Assist review addressed in commit Fix 1 (compile-time NameError): Added Fix 2 (runtime NameError): Added Fix 3 (regex robustness + blank lines): Updated rest_args pattern from Fix 4 (causal fast-path routing): Changed |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f64a7594c0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| # Fix 3: clean leading newlines from rest_args to avoid blank lines | ||
| rest_args_clean = rest_args.lstrip("\r\n") | ||
| rest_args_formatted = (", " + rest_args_clean.strip()) if rest_args_clean.strip() else "" |
There was a problem hiding this comment.
When rest_args is non-empty it already includes the comma that follows the value tensor, because the regex captures everything after v_var. Prepending another ", " here makes the generated fallback parse as q, k, v, , attn_mask=...; on AMD+aiter, ordinary HF SDPA calls with attn_mask/dropout_p/is_causal therefore produce an unparseable standalone attention module before either branch can run, so the fast-attention patch fails for those modules.
Useful? React with 👍 / 👎.
| f"{indent} _aiter_q = {q_var}.transpose(1, 2)\n" | ||
| f"{indent} _aiter_k = {k_var}.transpose(1, 2)\n" | ||
| f"{indent} _aiter_v = {v_var}.transpose(1, 2)\n" | ||
| f"{indent} {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal=True).transpose(1, 2)\n" |
There was a problem hiding this comment.
On AMD+aiter, this hardcoded causal=True changes the mask semantics whenever the original SDPA call has is_causal=True with non-square attention, such as chunked/cached decoding where q_len != key_len: PyTorch SDPA documents non-square causal bias as upper-left, while AITER flash_attn_func documents bottom-right alignment. In that context the rewritten branch can attend to a different set of keys than the original model, so it should either prove equal sequence lengths before using aiter or preserve the SDPA path.
Useful? React with 👍 / 👎.
Addresses all issues raised in the ChatGPT Codex review of PR unslothai#920. device_type.py: - Fix P1: FlashAttnFunc is a class requiring .apply(q,k,v,causal); wrap in lambda so all callers always receive a plain callable with signature func(q, k, v, causal=True). compiler.py: - Fix P1: Restrict is_causal match to literal True only (not complex expressions like `self.is_causal and mask is None`). Complex expressions cannot be safely evaluated at rewrite time; leave those calls unchanged. - Fix P1: Guard against unsupported SDPA args: attn_mask, dropout_p, scale, enable_gqa. If any are present, leave the call unchanged (SDPA handles them). - Fix P1: Guard against float32 inputs — aiter only supports float16/bfloat16. Added dtype check in _aiter_ok before dispatch. - Fix P1: Guard against non-square Q/K sequence lengths (chunked/cached decode). SDPA and aiter have different causal bias alignment for non-square inputs; only dispatch to aiter when q.shape[-2] == k.shape[-2]. - Fix P2: Tighten SDPA pattern to explicit prefixes only (torch.nn.functional., F., nn.functional.) — prevents accidentally matching disable_compile_scaled_dot_product_attention opt-out shim. - Fix P2: rest_args already leads with a comma; remove the extra ', ' that would produce `q, k, v, , attn_mask=...` in the SDPA fallback. - Rename import alias to _get_aiter_fn to avoid shadowing _aiter_fn variable. All existing fixes from f64a759 (Gemini review) are preserved.
|
All Codex P1/P2 review issues addressed in commit Summary of fixes: device_type.py
compiler.py
All five fixes from |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 30710aba98
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # Class-based API: must call via .apply(); wrap in lambda so callers | ||
| # always receive a plain callable with signature (q, k, v, causal=True) | ||
| _cls = _aiter.FlashAttnFunc | ||
| return lambda q, k, v, causal=True: _cls.apply(q, k, v, causal) |
There was a problem hiding this comment.
When an installed aiter build exposes only FlashAttnFunc, the new lambda still calls the autograd class API with just (q, k, v, causal). The current wrapper in ROCm/aiter invokes FlashAttnFunc.apply with the full positional list (dropout_p, softmax_scale, causal, window_size, etc.; see https://github.com/ROCm/aiter/blob/main/aiter/ops/mha.py), so in these class-only environments this sends the causal bool as dropout_p and leaves required arguments missing, causing a TypeError on the AMD fast path instead of falling back to SDPA. Fresh evidence versus the prior comment is that the code now wraps the class in a lambda, but the lambda still omits those required apply arguments.
Useful? React with 👍 / 👎.
Four issues raised in Codex review of PR unslothai#943: 1. aiter APIs not present in this tree: add pytest.mark.skipif guards on all TestAmdAiterDetection, TestAmdFlashAttnFunc, and TestReplaceSDPAWithAmdAiter classes. They skip gracefully when get_amd_attention_implementation / get_amd_flash_attn_func / replace_sdpa_with_amd_aiter are absent (pending PR unslothai#920). Avoids AttributeError/ImportError in CI on this branch. 2. --rocm-live flag: remove non-existent CLI option from docstring. Live hardware validation is now documented as plain 'pytest -v' with UNSLOTH_IS_PRESENT=1 env var, which actually works. 3. DEVICE_TYPE_TORCH assertion too weak: test_device_type_torch_is_cuda_on_hip now asserts DEVICE_TYPE_TORCH == 'cuda' exactly when running on real ROCm hardware (DEVICE_TYPE == 'hip'). Also adds source-inspection test that the hip->cuda alias transformation is present in device_type.py. 4. SM cap test duplicates production logic: TestSmCapGuard now inspects the real vllm_utils.py source to confirm 'sm_cap = 0' and 'is_hip()' are present (guards against silent removal). The mock test simulates the production guard directly and asserts get_device_capability() is never called when is_hip() returns True.
Codex P1 on PR unslothai#920: the FlashAttnFunc.apply() call in device_type.py was passing only (q, k, v, causal) but the actual aiter API signature requires 13+ positional arguments: dropout_p, softmax_scale, causal, window_size, bias, alibi_slopes, deterministic, return_lse, return_softmax, is_grad_enabled, and more (see ROCm/aiter:aiter/ops/mha.py:2485-2506). The previous lambda sent 'causal' as 'dropout_p', left 9 required args missing, and raised TypeError on the AMD fast path. Fix: remove the FlashAttnFunc wrapper entirely. If only FlashAttnFunc is exposed (no flash_attn_func), return None from get_amd_flash_attn_func() so callers fall back to SDPA rather than crashing. Also tighten get_amd_attention_implementation(): only return 'amd_aiter' when flash_attn_func (functional API) is present, not FlashAttnFunc (class API). This avoids detecting 'amd_aiter' and then immediately returning None from get_amd_flash_attn_func(). aiter >= 0.7 exposes flash_attn_func; environments with only FlashAttnFunc are treated as SDPA.
|
P1 fixed in commit Checked the actual Fix: Removed the Also: Tightened |
Addressing four points from human review on RX 9060 XT (gfx1200, ROCm 7.13): 1. Threshold fix (reviewer's exact suggestion, verified on hardware): Changed 'suggested <= batch_size' to 'suggested < 2 * batch_size' so the advisory is silent for small increments like batch 4->5 that are inconsistent with the 4->16 benchmark citation. 256 GB / 240 GB free at batch=4 still suggests 16; 40 GB free at batch=4 goes silent (suggests only 5). 2. Wire check_amd_vram_utilization into training_utils.py: Called immediately after the debug_info print in the Unsloth trainer loop, where batch_size and seq_len are both known. No behavior change on NVIDIA or CPU (is_hip() gate fires first). 3. Wire get_recommended_attn_implementation into empty_model.py: Used as the attn_implementation argument to AutoModelForCausalLM.from_config, replacing the hardcoded 'eager'. On AMD: 'sdpa' (MIOpen fused kernel, +37% prefill). On NVIDIA/CPU: returns None, falls back to 'eager' (no change). Downstream compiler SDPA support check (_supports_sdpa / ALL_ATTENTION_FUNCTIONS) handles models that don't support SDPA (Pixtral, Mistral 3). 4. Clarify naming vs get_amd_attention_implementation (PR unslothai#920): Updated docstring to note the difference: get_recommended_attn_implementation() is a top-level attn_implementation selector returning 'sdpa'|None; get_amd_attention_implementation() is an AMD-internal dispatch selector returning 'sdpa'|'amd_aiter' and is not intended as a top-level selector.
|
thanks for this, the AMD attention gap is worth closing. i set the branch up locally and tested on windows with torch 2.11.0+rocm7.13 on gfx1200 (RX 9060 XT), and WSL ubuntu-24.04 one thing that's right and easy to miss: the 1. generated code is a syntax error (blocking)
rest_args_formatted = (", " + rest_args_clean) if rest_args_clean else ""the comment above it is correct, else:
out = torch.nn.functional.scaled_dot_product_attention(
q, k, v, , is_causal=True)
2. the rewrite can never fire where it's called (blocking)
with the aiter path forced on, running all 32 the this is also why issue 1 hasn't shown up yet, the feature is currently inert. 3. no GPU arch gate
aiter's own support table lists RDNA4 only as gfx1201 (Radeon AI PRO R9700, experimental) 4. no fallback if the aiter call raisesgenerated block at if _aiter_ok:
out = _aiter_fn(q_t, k_t, v_t, causal=True).transpose(1, 2)
else:
out = torch.nn.functional.scaled_dot_product_attention(...)every guard is a pre-call predicate and the fallback only triggers on happy to re-test once these are sorted, the isolated env is still up on my end. |
…ssues
Tested on gfx1200 (RX 9060 XT) + gfx942 (MI325X), ROCm 7.2/7.13:
1. Double comma syntax error (compiler.py:438):
rest_args already starts with a comma when non-empty (e.g. ', is_causal=True').
The old code prepended another ', ' producing 'q, k, v, , is_causal=True' which
fails ast.parse and causes exec() to reject every rewritten module.
Fix: rest_args_formatted = rest_args_clean (keep leading comma as-is).
2. Fail-soft fallback on aiter call (compiler.py):
The generated if _aiter_ok: block had no exception handler. An unsupported head
dim, arch mismatch, or JIT build failure inside flash_attn_func would propagate
out of model forward and crash mid-training.
Fix: wrap the _aiter_fn() call in try/except with SDPA as the except handler.
3. GPU arch gate (device_type.py):
hasattr(aiter, 'flash_attn_func') is True on any arch where aiter installs —
the attribute check says nothing about hardware support. aiter's CK/ASM kernels
are CDNA-only; consumer RDNA (gfx1200, gfx1100) is not supported and not listed
in aiter's support table.
Fix: add _detect_gfx_arch() helper (rocminfo -> hipconfig fallback) and gate
get_amd_attention_implementation() on _AITER_SUPPORTED_ARCHS = ('gfx942', 'gfx950').
Unknown arch falls back to SDPA (fail safe).
4. Document call site limitation (compiler.py):
The existing Unsloth SDPA rewriting emits attn_mask= or is_causal=is_causal
(a variable, not literal True) — both of which the aiter guards correctly reject.
Added a comment explaining this so future contributors know the rewrite currently
applies only to user-code SDPA calls with bare is_causal=True literals, and
pointing to the path for a follow-up fix.
|
All 4 issues fixed in commit Fix 1 — Double comma syntax error (blocking): Fix 2 — Fail-soft fallback on aiter call: Fix 3 — GPU arch gate: Fix 4 — Call site limitation documented: The |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8494840b8d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| unsupported = ("attn_mask", "dropout_p", "scale", "enable_gqa") | ||
| if any(kw in rest_args for kw in unsupported): |
There was a problem hiding this comment.
Fresh evidence versus the prior rest-args issue is that this new guard only rejects unsupported keyword names. On AMD+aiter, a valid training call such as scaled_dot_product_attention(q, k, v, None, dropout, is_causal=True) contains no dropout_p keyword, so it passes this check and the generated aiter branch calls the kernel with only causal=True, silently dropping the original dropout while the fallback SDPA call would still apply it. Reject any positional tail other than the causal literal, or parse the SDPA arguments before rewriting.
Useful? React with 👍 / 👎.
…ue only
Codex review: the keyword-name guard (attn_mask, dropout_p, scale, enable_gqa)
only checks keyword argument names. A positional call like:
scaled_dot_product_attention(q, k, v, None, 0.1, is_causal=True)
contains no keyword named 'dropout_p', passes the guard, and the generated aiter
branch calls flash_attn_func(q, k, v, causal=True) — silently dropping the
positional dropout value that SDPA would have applied.
Fix: replace the two-step guard with a single fullmatch check:
re.fullmatch(r'is_causal\s*=\s*True', rest_args.strip().lstrip(',').strip())
Only rewrite when rest_args is exactly ', is_causal=True' (whitespace-tolerant).
Any other content — positional or keyword — leaves the call unchanged.
Verified on 8 test cases:
', is_causal=True' → rewrite ✅
', is_causal = True' → rewrite ✅ (spaces)
', None, 0.1, is_causal=True' → skip ✅ (positional dropout)
', attn_mask=mask, is_causal=True' → skip ✅ (keyword attn_mask)
', is_causal=True, dropout_p=0.0' → skip ✅ (keyword dropout_p)
'' → skip ✅ (no is_causal)
', is_causal=False' → skip ✅ (not causal)
', scale=0.5, is_causal=True' → skip ✅ (keyword scale)
|
Fixed in commit Replaced the two-step keyword-name guard with a single rest_args_stripped = rest_args.strip().lstrip(",").strip()
if not re.fullmatch(r"is_causal\s*=\s*True", rest_args_stripped):
return m.group(0)Only rewrites when Verified against 8 test cases including positional-dropout form, keyword attn_mask, keyword dropout_p, empty args, is_causal=False, and keyword scale — all correct. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3605a61e18
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| rocm_version = _detect_rocm_major_minor() | ||
| if rocm_version is not None: | ||
| major = int(rocm_version.split(".")[0]) | ||
| if major < 7: | ||
| return "sdpa" |
There was a problem hiding this comment.
In wheel-only ROCm 7 PyTorch installs that lack /opt/rocm/hipcc but do have rocm-smi, _detect_rocm_major_minor() can fall through to rocm-smi --showdriverversion; the ROCm SMI docs describe that flag as printing the AMDGPU module/kernel driver version, not the ROCm runtime version. A 6.x kernel driver string therefore makes this new gate return sdpa before checking aiter, even when torch.version.hip is 7.x, so the advertised AMD aiter path is disabled on supported wheel-based installs.
Useful? React with 👍 / 👎.
| # no keyword named "dropout_p" yet silently drop the positional dropout value. | ||
| # Accepting only the bare causal-only form guarantees correctness. | ||
| rest_args_stripped = rest_args.strip().lstrip(",").strip() | ||
| if not re.fullmatch(r"is_causal\s*=\s*True", rest_args_stripped): |
There was a problem hiding this comment.
For AMD+aiter, a safe causal-only call formatted in the common multiline style as scaled_dot_product_attention(q, k, v, is_causal=True,) is left unchanged because stripping only the leading comma leaves is_causal=True,, which this fullmatch rejects. That means many otherwise eligible HF/Black-formatted SDPA calls never take the new aiter fast path even though they satisfy the intended “bare causal-only” guard; allowing an optional trailing comma would preserve the safety check without disabling the feature for this format.
Useful? React with 👍 / 👎.
| return m.group(0) | ||
| except Exception: | ||
| pass | ||
| return None |
There was a problem hiding this comment.
In ROCm runtime or PyTorch-wheel containers that expose an MI300/MI355 GPU to torch but do not include the developer utilities rocminfo or hipconfig, both subprocess probes fail and _detect_gfx_arch() returns None; the architecture gate then treats the supported card as unsupported and returns sdpa, so the new aiter rewrite is never enabled. Querying the active torch device's gcnArchName before returning None would keep the fast path available in these supported runtime-only deployments.
Useful? React with 👍 / 👎.
…obe, trailing comma Three issues from Codex review of commit 3605a61: 1. ROCm version gate uses torch.version.hip first (device_type.py): _detect_rocm_major_minor() can fall through to 'rocm-smi --showdriverversion' which reports the AMDGPU kernel driver version, not the ROCm runtime version. A 6.x kernel driver string would gate out aiter even on ROCm 7 wheel installs. Fix: check torch.version.hip first (always present in ROCm PyTorch, reliable). Fall back to _detect_rocm_major_minor() only when torch.version.hip is absent. 2. _detect_gfx_arch() falls back to torch.cuda.get_device_properties (device_type.py): In runtime-only containers where rocminfo and hipconfig are absent, both subprocess probes failed and the function returned None, causing the arch gate to treat supported cards (gfx942, gfx950) as unsupported. Fix: add torch.cuda.get_device_properties(0).gcnArchName as final fallback. gcnArchName is always available when torch+ROCm is installed (no dev tools needed). 3. Accept trailing comma on causal-only SDPA calls (compiler.py): Black/multiline-formatted calls like scaled_dot_product_attention(q, k, v, is_causal=True,) were left unchanged because 'is_causal=True,' failed fullmatch. Fix: strip both leading AND trailing comma+whitespace before fullmatch: rest_args.strip().lstrip(',').strip().rstrip(',').strip() Verified on 8 cases including spaces+trailing comma (Black format).
|
All 3 issues addressed in commit Fix 1 — ROCm version gate uses Fix 2 — Fix 3 — Trailing comma: Strip both leading AND trailing comma+whitespace before |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a33eafacf
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # 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)", |
There was a problem hiding this comment.
When training on AMD with aiter, the attention inputs generally require gradients, but aiter.flash_attn_func's autograd path asserts that return_lse is enabled when gradients are active (upstream FlashAttnFunc.forward). This generated call omits return_lse=True, so each training forward that passes _aiter_ok first tries aiter, raises, and then recomputes with SDPA in the except, making the new fast path exception-driven double work instead of an optimization; either guard it to inference-only or request and unwrap the LSE-returning training form.
Useful? React with 👍 / 👎.
| 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"((?:[^)]|\([^)]*\))*)\)" |
There was a problem hiding this comment.
When AMD+aiter sees an assignment that immediately chains the SDPA result, such as attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2), this matcher stops at the SDPA closing paren and re.sub leaves the .transpose(...) suffix after the generated multi-line if/else block. That produces invalid compiled source instead of preserving a valid attention implementation; require an end-of-statement after the call or include the suffix in the rewrite.
Useful? React with 👍 / 👎.
…#920) Two issues from Codex review of commit 5a33eaf: 1. LSE assertion during training (compiler.py): aiter.flash_attn_func's autograd path (FlashAttnFunc.forward) asserts return_lse=True when any input requires_grad=True. The generated call omits return_lse, so every training forward would try aiter, raise AssertionError, then recompute with SDPA in the except branch — exception-driven double work, not a fast path. Fix: add 'and not {q_var}.requires_grad' to _aiter_ok. The fast path now fires only during inference (no_grad / eval mode). Training always uses SDPA via the else branch, zero overhead. 2. Chained SDPA calls leave dangling suffix (compiler.py): 'out = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2)' — the regex stops at the SDPA closing paren, re.sub replaces only that part, and '.transpose(1, 2)' is left dangling after the generated if/else block, producing invalid Python that exec() rejects. Fix: add negative lookahead r'(?![ \t]*\.)' at the end of the pattern. Any SDPA call immediately followed by a method chain is left unchanged. Verified: plain call matches; .transpose() and .view() chains are rejected.
|
Both issues addressed in commit Fix 1 — Inference-only guard: Added Fix 2 — Chained SDPA calls: Added negative lookahead |
|
re-tested on e20f0a4. all four blocking items are fixed, verified by running them rather first a correction to the record: i only tested gfx1200. i don't have an MI300 or MI355 verified fixed
on the two codex P2s from 5a33eafLSE guard: correct. i checked this against aiter source and your reasoning holds. one gap: your guard checks only chained calls: partially. the rewrites to (trimmed): the
issue 2 is still openstill 0 rewrites across all 32 the inference-only guard narrows it further
that's the thing worth settling before merge. this is ~250 lines that provably never smaller
codex's last pass was on 5a33eaf, so e20f0a4 hasn't been reviewed by anything yet, and |
Four concrete fixes from hardware review (gfx1200, e20f0a4): 1. requires_grad check covers all three of q, k, v: aiter/ops/mha.py:2502 checks 'any(x.requires_grad for x in [q, k, v])'. The old guard checked only q.requires_grad. A LoRA targeting k_proj/v_proj but not q_proj leaves k and v with grad while q.requires_grad is False — the guard passed but the assert fired. try/except caught it, causing exception-driven double work. Now guards all three. 2. Lookahead broadened from dot-only to any non-whitespace/non-comment: The old '(?![ \t]*\.)' rejected .transpose() but not '* scale' or '[0]'. Both forms parse and produce no exec error, but the two branches compute different numbers — the scalar multiply or subscript lands on the SDPA fallback branch only, causing silent divergence. New lookahead: '(?![ \t]*[^\s#])' rejects any non-ws/non-comment suffix. Trailing comments and clean EOL still match; .method, * op, [idx] all reject. Verified on: plain call, trailing comment, .transpose, * scale, [0]. 3. Removed duplicate rest_args_clean assignment — leftover from double-comma fix. 4. Updated docstring to match current guards: removed stale attn_mask/dropout_p rejection rule (now handled by the fullmatch guard), added arch gate, added inference-only guard, documented the known call-site limitation clearly.
|
All four concrete fixes addressed in commit Fix 1 — requires_grad covers q, k, v: Changed to Fix 2 — Lookahead broadened: Changed from Fix 3 — Duplicate rest_args_clean removed: Cleaned up the leftover from the double-comma fix. Fix 4 — Docstring updated: Removed the stale attn_mask/dropout_p rejection rule (now handled by the fullmatch guard), added the arch gate and inference-only guard, and explicitly documented the known call-site limitation. On the open question about the call site: you're right that ~250 lines provably don't execute in Unsloth's own compiled paths today. The honest description is that this PR delivers the infrastructure (correct guards, fail-soft, arch detection) for user-code SDPA calls with bare |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e63e24c1b2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| f"{indent} try:", | ||
| f"{indent} {out_var} = _aiter_fn(_aiter_q, _aiter_k, _aiter_v, causal=True).transpose(1, 2)", | ||
| f"{indent} except Exception:", |
There was a problem hiding this comment.
When the SDPA attention class is actually compiled (disable=False, for example with sdpa_dynamic_compile=False), this generated try/except becomes part of a @torch.compile(fullgraph=...) forward. TorchDynamo does not support tracing Python exception handling in fullgraph mode, so an otherwise eligible AMD+aiter attention module can fail on the first forward instead of falling back to SDPA. Keep the fail-soft fallback out of the compiled region, or skip this rewrite when the SDPA forward is compiled.
Useful? React with 👍 / 👎.
| 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) |
There was a problem hiding this comment.
On mixed-architecture ROCm hosts, or when HIP_VISIBLE_DEVICES exposes a different GPU than physical device 0, this returns the first gfx... token from system-wide rocminfo output before checking PyTorch's active device. Since the later gate only accepts gfx942/gfx950, a supported visible MI300/MI355 can be forced to SDPA if another agent is printed first, or an unsupported visible GPU can repeatedly try aiter because a supported agent was printed first; prefer torch.cuda.get_device_properties() for the current visible device before parsing global CLI output.
Useful? React with 👍 / 👎.
|
verified all four on e63e24c.
re-ran the earlier checks against the new head for regressions: arch gate still returns that clears everything i raised. agreed the scope question is @danielhanchen 's call, and codex's last pass was on 5a33eaf, so it has now skipped both e20f0a4 and e63e24c. those |
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep them coming! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
… and arch detection Two issues from Codex review of e63e24c: 1. try/except incompatible with torch.compile(fullgraph=True) (compiler.py): TorchDynamo cannot trace Python exception handling in fullgraph mode. A generated try/except inside a compiled attention forward would fail on the first compiled call rather than falling back to SDPA. Fix: replace try/except with a nested _call_aiter_safe() helper that returns None on exception. The generated code uses if/else on the return value — no exception handling in the hot path, compatible with fullgraph. 2. rocminfo returns first gfx* token from system-wide output (device_type.py): On mixed-arch ROCm hosts or with HIP_VISIBLE_DEVICES, the first token in rocminfo output may not be the active device. This could gate in or out the wrong architecture. Fix: move torch.cuda.get_device_properties(current_device()).gcnArchName to FIRST — it always reflects the device PyTorch is actually using. rocminfo/hipconfig are demoted to subprocess fallbacks for environments where torch cannot query device properties (unusual builds).
|
Both Codex issues addressed in commit Fix 1 — compile-safe fallback: Replaced try/except in generated code with a nested Fix 2 — arch detection uses active device first: |
…nslothai#944) The docstring referenced get_amd_attention_implementation() from PR unslothai#920, which is not yet merged. If unslothai#944 lands first, main would carry a docstring citing a function not in the tree. Cross-references are cheaper pointing backwards — let whichever lands second carry the disambiguation note. The preceding line already tells callers what they need: None means no AMD-specific override, keep your own default.
|
Codex came back clean on e63e24c ( |
|
Separate thread from correctness, same note I left on #944: comment volume relative to code. Against the files' own baselines, The clearest example is compiler.py:448-456, six comment lines wrapped around three lines # 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 (fullmatch guarantees this).
# Pass it through to the SDPA fallback as-is.
rest_args_formatted = rest_args.strip().lstrip(",").strip().rstrip(",").strip()
# Reconstruct with leading comma for the SDPA fallback call
rest_args_formatted = (", " + rest_args_formatted) if rest_args_formatted else ""Two real problems hiding in there, not just verbosity. The 26-line docstring on Same as on #944, treat the above as representative and do one general pass over the whole
Those three are the reason someone won't undo your guards in six months. Most of the rest Also: ccddae0 landed after my last pass, so I'll re-verify against it separately. |
LeoBorcherding: added lines were running at ~2x the repo's comment density. compiler.py: - Remove dead rest_args_clean assignment (never read after double-comma fix) - Remove contradictory 'do NOT add another' comment (code did add one) - Trim verbose 9-line block before fullmatch to 3 lines stating the rule - Remove 'leave call unchanged' after return — code speaks for itself - Shorten 'Generate the replacement block' to one line - Shorten _call_aiter_safe comment to 2 lines device_type.py: - Shorten subprocess fallback comment (1 line, not 3) - Shorten CDNA-only arch gate comment (2 lines, keep the non-obvious fact) - Remove inline 'fail safe: unknown arch' on return Kept the non-obvious comments that protect correctness: why q/k seq lengths must match (SDPA top-left vs aiter bottom-right causal alignment), why requires_grad triggers the inference-only guard (return_lse assertion in aiter/ops/mha.py), and why the arch allowlist is CDNA-only.
|
Done in commit compiler.py: removed dead device_type.py: shortened subprocess fallback comment (3→1 lines), shortened CDNA-only arch gate comment (4→2 lines, kept the non-obvious fact), removed inline 'fail safe: unknown arch' comment. Kept the three non-obvious ones: q/k seq length guard (SDPA top-left vs aiter bottom-right causal alignment), inference-only guard (return_lse assertion in mha.py:2502/2536), and CDNA-only arch allowlist. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: add32da42c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # 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]*" |
There was a problem hiding this comment.
In AMD+aiter builds, this pattern is not anchored to a statement boundary, so it also matches the attn_output = substring in attribute assignments such as self.attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True). re.sub leaves the self. prefix in place and inserts the generated multi-line if block after it, making the compiled cache module syntactically invalid instead of falling back to SDPA; require a line/start-of-statement anchor or a negative lookbehind for ./identifier characters.
Useful? React with 👍 / 👎.
Codex: 'self.attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True)' matched because the pattern captured 'attn_output' as the LHS variable but left 'self.' dangling before the generated if/else block, producing invalid syntax. Fix: add negative lookbehind (?<![.\w]) before the LHS variable group. Rejects any assignment where the variable is preceded by a dot or word character (attribute access like self.x, module.x). Plain local variable assignments are unchanged. Verified: 'out = F.sdpa(...)' matches, 'self.out = F.sdpa(...)' and 'module.out = F.sdpa(...)' are rejected.
|
Fixed in commit Added negative lookbehind Verified on four cases: |
|
Already fixed in commit Added |
|
Already fixed — the bot is reviewing The keyword-name guard ( |
|
@codex review Head is now add32da. Your last pass was e63e24c, so this covers two commits: ccddae0 Verified locally on gfx1200 / ROCm 7.13: both files compile, the arch gate returns sdpa on |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
2 participants
Summary
Reopening with proper implementation after closing #916.
The key insight: unsloth-zoo compiles attention by rewriting source code as strings
(
inspect.getsource()+ regex transforms incompiler.py). The correct integrationpoint is to add a new source-level transform in that same pipeline — not to hook
wrappers around module forwards.
How it works
replace_sdpa_with_amd_aiter(source)is a new source-rewriting function thatruns in the existing
scaled_dot_product_attention_modulesprocessing loop inunsloth_compile_transformers().It matches
scaled_dot_product_attention(query, key, value, ...)calls andrewrites them as:
The transpose converts transformers convention
(B, H, S, D)to aiter's(B, S, H, D). Falls back to standard SDPA if aiter is unavailable at runtime.Changes
unsloth_zoo/device_type.pyget_amd_attention_implementation(): detects amd-aiter with ROCm ≥7.0 gate,uses
import importlib.utilexplicitly (fixes Python 3.13+ issue from feat: AMD aiter Flash Attention dispatch for ROCm MI300X/MI325X (amd-aiter) #916)get_amd_flash_attn_func(): returnsflash_attn_funcorFlashAttnFunc(both naming conventions)unsloth_zoo/compiler.pyreplace_sdpa_with_amd_aiter(source): new source-rewriting function, hookedinto the existing SDPA processing loop. No-op on NVIDIA or ROCm < 7.0.
Expected Performance (MI300X/MI325X, ROCm ≥7.0)
Testing
is_hip()returns False, all code gatedget_amd_attention_implementation()returns"sdpa"Cannot benchmark live (clusters run ROCm 6.2); numbers from AMD's validated
benchmarks on MI300X gfx942.
Install
pip install amd-aiter # https://github.com/ROCm/aiter — ROCm >= 7.0Closes (reopens) #916.