ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix xformers Blackwell guard: broader coverage and root cause docs (#4338)

Browse source 
* Remove outdated xformers Blackwell version guard

The guard at _utils.py:976-989 blocked xformers 0.0.32.post2 on
Blackwell/RTX 50x/Jetson GPUs (SM 10.0/11.0/12.0) due to a FA3
dispatch bug that caused CUDA errors (issue #1329).

This is no longer needed because:

1. xformers fixed the FA3 dispatch in 0.0.33.post2 by capping it
   at SM <= 9.0, so FA3 is never attempted on Blackwell. The FA2
   backend works correctly via PTX forward compatibility.

2. The only blocked version (0.0.32.post2) was built for torch 2.8.0
   and cannot load on torch 2.9+ due to ABI mismatch, so the guard
   never actually triggers for any current user.

3. The existing _register_extensions() check plus the except Exception
   fallback already handle broken xformers installs gracefully by
   falling back to SDPA.

Verified on NVIDIA RTX PRO 6000 Blackwell (SM 12.0) with both
pre-built wheels (0.0.33.post2) and source builds -- all attention
tests pass with exact numerical match vs SDPA.

* Update xformers Blackwell guard with root cause and broader coverage

Changes to the xformers version guard for Blackwell/RTX 50x/Jetson GPUs:

1. Broaden version check from `in (0.0.32.post2,)` to `<= 0.0.32.post2`
   to cover all versions with the broken FA3 dispatch, not just one.

2. Add `DEVICE_TYPE == "cuda"` guard to avoid calling
   `get_device_capability()` on non-CUDA devices (XPU, etc.).

3. Document the root cause: xformers <= 0.0.32.post2 used
   `capability >= (9, 0)` in the FA3 dispatch, which matched
   Blackwell SM 12.0 and attempted sm_90a Hopper kernels on it.
   Fixed upstream in 0.0.33 with `<= (9, 0)`.

4. Update error message to include the installed version, mention
   the fix (upgrade to >= 0.0.33), and keep the build-from-source
   fallback. The raise is caught by `except Exception` which shows
   the message when UNSLOTH_ENABLE_LOGGING is set and falls back
   to SDPA.

Verified on NVIDIA RTX PRO 6000 Blackwell (SM 12.0):
- xformers 0.0.33.post2 pre-built wheel: works (FA2 via PTX)
- xformers source build: works (FA2 native)
- Both have exact numerical match vs SDPA

---------

Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-03-16 21:30:03 -07:00 • committed by GitHub
commit dc2879a048
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 18 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

31
unsloth/models/_utils.py
View file

@ -973,20 +973,25 @@ except:
try:
from xformers import __version__ as xformers_version
# [TODO] Xformers does NOT work on RTX 50x (12), B200 (10), Jetson (11)
# Xformers <= 0.0.32.post2 has a broken FA3 dispatch on Blackwell/RTX 50x GPUs.
# The FA3 check used `capability >= (9, 0)` which matches SM 10.0/11.0/12.0,
# causing sm_90a kernels to be attempted on non-Hopper GPUs (CUDA error in
# flash_fwd_launch_template.h:188). Fixed in 0.0.33 with `<= (9, 0)`.
# See https://github.com/facebookresearch/xformers/issues/1329
# CUDA error (/workspace/xfrm2/third_party/flash-attention/hopper/flash_fwd_launch_template.h:188)
major_version, minor_version = torch.cuda.get_device_capability()
if (f"{major_version}.{minor_version}" in ("10.0", "11.0", "12.0")) and (
Version(xformers_version) in (Version("0.0.32.post2"),)
):
raise NotImplementedError(
"Unsloth: Xformers does not work in RTX 50X, Blackwell GPUs as of yet. Please build from source via\n"
"```\n"
"pip install ninja\n"
"pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers\n"
"```\n"
)
if DEVICE_TYPE == "cuda":
major_version, minor_version = torch.cuda.get_device_capability()
if (f"{major_version}.{minor_version}" in ("10.0", "11.0", "12.0")) and (
Version(xformers_version) <= Version("0.0.32.post2")
):
raise NotImplementedError(
f"Unsloth: Xformers {xformers_version} has a broken FA3 dispatch on "
f"SM {major_version}.{minor_version} GPUs. Please upgrade to >= 0.0.33 or build from source via\n"
"```\n"
"pip install ninja\n"
"pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers\n"
"```\n"
)
# Temporarily disable 0.0.27 and higher - inference issues
if False: # Version(xformers_version) >= Version("0.0.27"):
raise ImportError(