From dc2879a048f2c8853554aa0c3280f1b0e6afbc17 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 21:30:03 -0700 Subject: [PATCH] Fix xformers Blackwell guard: broader coverage and root cause docs (#4338) * 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 --- unsloth/models/_utils.py | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index cbaebcc7ac..44ebc0a445 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -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(