From 5a02ed4f0f5db71bdc2d664eef0198ac64f11363 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 18 Feb 2026 08:58:25 +0000 Subject: [PATCH] Disable flex attention on Blackwell+ GPUs (sm_120+) at startup --- studio/backend/main.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index 3d9e2a0f9a..e027998dcf 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1,6 +1,7 @@ """ Main FastAPI application for Unsloth UI Backend """ +import os import secrets import shutil from contextlib import asynccontextmanager @@ -15,7 +16,7 @@ from datetime import datetime # Import routers from routes import training_router, models_router, inference_router, datasets_router, auth_router, export_router from auth import storage -from utils.hardware import detect_hardware +from utils.hardware import detect_hardware, get_device, DeviceType import utils.hardware.hardware as _hw_module UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache" @@ -27,6 +28,18 @@ async def lifespan(app: FastAPI): # Detect hardware first — sets DEVICE global used everywhere detect_hardware() + # Disable flex attention on Blackwell+ GPUs (sm_120 and above) + if get_device() == DeviceType.CUDA: + import torch + props = torch.cuda.get_device_properties(0) + sm_version = props.major * 10 + props.minor + if sm_version >= 120: + os.environ["UNSLOTH_FLEX_ATTENTION"] = "0" + import logging + logging.getLogger(__name__).info( + f"GPU sm_{sm_version} detected — setting UNSLOTH_FLEX_ATTENTION=0" + ) + if not storage.is_initialized(): setup_token = secrets.token_urlsafe(32) storage.save_setup_token(setup_token)