Патч Flash Attention KJ (PatchFlashAttentionKJ) - ComfyUI-KJNodes

Патч Flash Attention KJ(PatchFlashAttentionKJ)

Экспериментальный узел для патчинга внимания с использованием flash attention, без тихого отката к SDPA, как это делает ComfyUI по умолчанию. Патчит внимание модели, проходящей через этот узел; для отключения патча обойдите или отключите этот узел. Требует установки библиотеки flash_attn.

Описание

Экспериментальный узел для патчинга механизма внимания диффузионной модели с использованием Flash Attention (через библиотеку flash_attn). В отличие от стандартного отката к SDPA в ComfyUI, этот узел напрямую применяет flash attention. Для отключения патча обойдите или отключите этот узел. Требует установки пакета Python flash-attn (выполните pip install flash-attn).

Входы

Имя	Тип	Обязательно	По умолчанию	Описание
`model`	MODEL	Да	–	Диффузионная модель для патчинга. Слои внимания будут заменены на операции flash attention.
`allow_compile`	БУЛЕВ	Нет	ложь	Если `истина`, разрешает компиляцию ядер flash attention (например, с помощью `torch.compile` в PyTorch) для потенциального повышения производительности. Это может увеличить время запуска и использование памяти. Включайте только при совместимом окружении.

Выходы

Имя	Тип	Описание
`MODEL`	MODEL	Модель с включённым flash attention.

Примечания по использованию

Этот узел является экспериментальным и может не работать со всеми архитектурами моделей или конфигурациями CUDA.
Библиотека flash_attn должна быть установлена в вашем окружении. Команда установки: pip install flash-attn (может потребоваться сборка torch с поддержкой CUDA).
Если возникают ошибки, попробуйте обойти узел, чтобы вернуться к стандартному механизму внимания.
Опция allow_compile по умолчанию отключена; включайте её только если понимаете последствия (более долгий первый вывод, возможная нестабильность).