Dymas/whisper-api-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Flash Attention 2 capability check for older GPUs

Browse Source
This commit is contained in:
Serge Zaigraeff
2026-03-31 12:01:49 +03:00
parent 2c6f7aee8e
commit 0a7b20f408
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/core/transcriber.py
+12 -2
View File
@@ -143,13 +143,23 @@ class WhisperTranscriber:
use_safetensors=True, use_safetensors=True,
) )
use_flash_attn = False
if self.device.type == "cuda":
# Flash Attention 2 требует архитектуру Ampere или новее (compute capability >= 8.0)
capability = torch.cuda.get_device_capability(self.device.index)
if capability[0] >= 8:
use_flash_attn = True
logger.info("GPU поддерживает Flash Attention 2 (compute capability: %d.%d)", *capability)
else:
logger.info("GPU не поддерживает Flash Attention 2 (compute capability: %d.%d), используется стандартный режим", *capability)
try: try:
if self.device.type == "cuda": if use_flash_attn:
model_kwargs["attn_implementation"] = "flash_attention_2" model_kwargs["attn_implementation"] = "flash_attention_2"
self.model = WhisperForConditionalGeneration.from_pretrained( self.model = WhisperForConditionalGeneration.from_pretrained(
self.model_path, **model_kwargs self.model_path, **model_kwargs
).to(self.device) ).to(self.device)
if self.device.type == "cuda": if use_flash_attn:
logger.info("Используется Flash Attention 2") logger.info("Используется Flash Attention 2")
except Exception as e: except Exception as e:
logger.warning("Не удалось загрузить модель с Flash Attention: %s", e) logger.warning("Не удалось загрузить модель с Flash Attention: %s", e)