Patch Flash Attention KJ (PatchFlashAttentionKJ) - ComfyUI-KJNodes

Patch Flash Attention KJ(PatchFlashAttentionKJ)

Nodo experimental para parchear la atención y usar flash attention, sin la alternativa silenciosa de SDPA que hace ComfyUI por defecto. Parchea la atención del modelo que pasa a través de este nodo; para deshabilitar, omita o desconecte este nodo. Requiere que la biblioteca flash_attn esté instalada.

Descripción

Nodo experimental para parchear el mecanismo de atención de un modelo de difusión para usar Flash Attention (a través de la biblioteca flash_attn). A diferencia de la alternativa silenciosa de SDPA predeterminada de ComfyUI, este nodo aplica flash attention directamente. Para deshabilitar el parche, omita o desconecte este nodo. Requiere que el paquete de Python flash_attn esté instalado (ejecute pip install flash-attn).

Entradas

Nombre	Tipo	Requerido	Por defecto	Descripción
`model`	MODEL	Sí	–	El modelo de difusión a parchear. Las capas de atención serán reemplazadas con operaciones de flash attention.
`allow_compile`	BOOLEANO	No	Falso	Si es `Verdadero`, permite que los kernels de flash attention sean compilados (ej., usando `torch.compile` de PyTorch) para posibles ganancias de rendimiento. Esto puede aumentar el tiempo de inicio y el uso de memoria. Habilítelo solo si tiene un entorno compatible.

Salidas

Nombre	Tipo	Descripción
`MODEL`	MODEL	El modelo parcheado con flash attention habilitado.

Notas de uso

Este nodo es experimental y puede no funcionar con todas las arquitecturas de modelo o configuraciones de CUDA.
La biblioteca flash_attn debe estar instalada en su entorno. Instrucciones de instalación: pip install flash-attn (puede requerir una compilación de torch compatible con CUDA).
Si encuentra errores, intente omitir el nodo para volver al mecanismo de atención predeterminado.
La opción allow_compile está desactivada por defecto; actívela solo si comprende las implicaciones (primera inferencia más larga, posible inestabilidad).