From 0a7b20f408043de3597347cf9c87e0467c1457ee Mon Sep 17 00:00:00 2001 From: Serge Zaigraeff Date: Tue, 31 Mar 2026 12:01:49 +0300 Subject: [PATCH] Add Flash Attention 2 capability check for older GPUs --- app/core/transcriber.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/app/core/transcriber.py b/app/core/transcriber.py index b38e909..f0d4276 100644 --- a/app/core/transcriber.py +++ b/app/core/transcriber.py @@ -143,13 +143,23 @@ class WhisperTranscriber: 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: - if self.device.type == "cuda": + if use_flash_attn: model_kwargs["attn_implementation"] = "flash_attention_2" self.model = WhisperForConditionalGeneration.from_pretrained( self.model_path, **model_kwargs ).to(self.device) - if self.device.type == "cuda": + if use_flash_attn: logger.info("Используется Flash Attention 2") except Exception as e: logger.warning("Не удалось загрузить модель с Flash Attention: %s", e)