From 600340d91db76caeb4f74e8bbdf53fd440ad9a52 Mon Sep 17 00:00:00 2001 From: EnragedAntelope Date: Fri, 5 Dec 2025 14:30:37 -0500 Subject: [PATCH] feat: Refactor Flash Attention 2 compatibility to optimization check, improve device mapping for model loading, and enhance attention mode selection logic. --- nodes.py | 60 +++++++++++++++++++++++++++++--------------------------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/nodes.py b/nodes.py index f803b61..c88743e 100644 --- a/nodes.py +++ b/nodes.py @@ -44,13 +44,13 @@ def _check_quantization_support(): def _get_gpu_info(): - """Get GPU compute capability and check Flash Attention 2 compatibility. + """Get GPU compute capability and check Flash Attention 2 optimization status. - Flash Attention 2 supports SM_80-SM_90 (Ampere, Ada, Hopper). - Blackwell (RTX 50-series) is SM_120 and NOT supported by Flash Attention 2. + Returns info about whether FA2 is likely OPTIMAL for this GPU (hint, not blocker). + FA2 may still work on newer GPUs even if not fully optimized yet. """ if not torch.cuda.is_available(): - return None, None, False + return None, None, True # Assume compatible if we can't check try: device = torch.device("cuda:0") @@ -59,21 +59,21 @@ def _get_gpu_info(): compute_cap = f"{major}.{minor}" gpu_name = props.name - # Flash Attention 2 supports SM 80-90 (Ampere, Ada, Hopper) - # SM 120+ is Blackwell - NOT supported by FA2 - fa2_compatible = 80 <= (major * 10 + minor) <= 90 + # FA2 is optimized for SM 80-90 (Ampere, Ada, Hopper) + # Newer architectures (SM 120+ Blackwell) may not have optimized kernels yet + # But this could change with future FA2 updates + sm_version = major * 10 + minor + fa2_optimized = 80 <= sm_version <= 90 - return gpu_name, compute_cap, fa2_compatible + return gpu_name, compute_cap, fa2_optimized except Exception: - return None, None, False + return None, None, True # Assume compatible if we can't check # Get GPU info at module load time -_GPU_NAME, _GPU_COMPUTE_CAP, _FA2_COMPATIBLE = _get_gpu_info() +_GPU_NAME, _GPU_COMPUTE_CAP, _FA2_OPTIMIZED = _get_gpu_info() if _GPU_NAME: print(f"[SCG_LocalVLM] GPU: {_GPU_NAME} (SM {_GPU_COMPUTE_CAP})") - if not _FA2_COMPATIBLE: - print(f"[SCG_LocalVLM] NOTE: Flash Attention 2 not optimized for this GPU. SDPA recommended.") @@ -427,7 +427,6 @@ class QwenVL: } if quantization == "4bit": - # Quantization requires device_map for bitsandbytes quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=compute_dtype, @@ -435,7 +434,8 @@ class QwenVL: bnb_4bit_use_double_quant=True, ) load_kwargs["quantization_config"] = quantization_config - load_kwargs["device_map"] = "auto" # Let Accelerate handle quantized loading + # Use explicit GPU placement to avoid Accelerate memory management overhead + load_kwargs["device_map"] = {"": 0} if torch.cuda.device_count() == 1 else "auto" print(f"[SCG_LocalVLM] Loading with 4-bit quantization") elif quantization == "8bit": @@ -443,35 +443,41 @@ class QwenVL: load_in_8bit=True, ) load_kwargs["quantization_config"] = quantization_config - load_kwargs["device_map"] = "auto" + load_kwargs["device_map"] = {"": 0} if torch.cuda.device_count() == 1 else "auto" print(f"[SCG_LocalVLM] Loading with 8-bit quantization") else: - # Non-quantized: NO device_map - use native PyTorch loading for speed - # Model will load to CPU first, then we move to GPU - print("[SCG_LocalVLM] Loading without quantization (native PyTorch)") + # Non-quantized: load directly to GPU for maximum speed + load_kwargs["device_map"] = {"": "cuda:0"} + print("[SCG_LocalVLM] Loading without quantization") - # Handle attention mode - simplified - if attention_mode == "flash_attention_2" and _FA2_COMPATIBLE: + # Handle attention mode + if attention_mode == "flash_attention_2": try: import flash_attn load_kwargs["attn_implementation"] = "flash_attention_2" attention_used = "flash_attention_2" + if not _FA2_OPTIMIZED: + print(f"[SCG_LocalVLM] Using Flash Attention 2 (may not be optimal for {_GPU_NAME or 'this GPU'})") + else: + print("[SCG_LocalVLM] Using Flash Attention 2") except ImportError: + print("[SCG_LocalVLM] flash-attn not installed, using SDPA instead") load_kwargs["attn_implementation"] = "sdpa" attention_used = "sdpa" - elif attention_mode == "sdpa" or (attention_mode == "flash_attention_2" and not _FA2_COMPATIBLE): + elif attention_mode == "sdpa": load_kwargs["attn_implementation"] = "sdpa" attention_used = "sdpa" + print("[SCG_LocalVLM] Using SDPA attention") elif attention_mode == "eager": load_kwargs["attn_implementation"] = "eager" attention_used = "eager" + print("[SCG_LocalVLM] Using eager attention") else: - # auto - use sdpa as default - load_kwargs["attn_implementation"] = "sdpa" + # auto - use sdpa as default (best compatibility + good performance) + load_kwargs["attn_implementation"] = "sdpa" attention_used = "sdpa" - - print(f"[SCG_LocalVLM] Using {attention_used} attention") + print("[SCG_LocalVLM] Using SDPA attention (auto)") # Load the model model_class = _get_model_class(model, is_vl=True) @@ -491,10 +497,6 @@ class QwenVL: load_time = time.time() - load_start - # For non-quantized models: move to GPU using native PyTorch - if quantization == "none": - self.model = self.model.cuda() - # Set model to eval mode self.model.eval()