{megapaks} Add FlashAttention (again)

#144

Author: YanWenKun

cu126-megapak/Dockerfile

@@ -296,17 +296,12 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 RUN --mount=type=cache,target=/root/.cache/pip \
     # Update UV
     pip install -U uv \
+    # Compile-install FlashAttention (version paired with xFormers)
+    && pip install flash-attn==2.8.3 \
     # Nunchaku version needs to sync with PyTorch version
     && pip install \
 https://github.com/nunchaku-tech/nunchaku/releases/download/v1.0.2/nunchaku-1.0.2+torch2.9-cp312-cp312-linux_x86_64.whl
 
-# Notes on FlashAttention:
-# <xformers 0.0.32.post2> requires <flash-attn [2.7.1, 2.8.2]>,
-# but flash-attn 2.8.2 does not have binary wheel for PyTorch 2.8.
-# Now using xformers embbeded flash-attn.
-# May need to suppress some custom nodes that explicit depends on flash-attn.
-# Wait for xFormers to update.
-
 ################################################################################
 # Bundle ComfyUI in the image
 

cu128-megapak/Dockerfile

@@ -296,6 +296,8 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 RUN --mount=type=cache,target=/root/.cache/pip \
     # Update UV
     pip install -U uv \
+    # Compile-install FlashAttention (version paired with xFormers)
+    && pip install flash-attn==2.8.3 \
     # Nunchaku version needs to sync with PyTorch version
     && pip install \
 https://github.com/nunchaku-tech/nunchaku/releases/download/v1.0.2/nunchaku-1.0.2+torch2.9-cp312-cp312-linux_x86_64.whl