CU128 with 5090 CUDA error caused by xformers

[wx2058]

After using the yanwk/comfyui-boot:cu128-slim image, an error occurred during use. The error message is as follows:

model weight dtype torch.float16, manual cast: None
model_type EPS
Using xformers attention in VAE
Using xformers attention in VAE
VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
Requested to load SDXLClipModel
FETCH ComfyRegistry Data: 5/95
loaded completely 30383.8625 1560.802734375 True
Loads SAM model: /root/ComfyUI/models/sams/sam_vit_b_01ec64.pth (device:AUTO)
Requested to load SDXL
loaded completely 28650.93463973999 4897.0483474731445 True
0%| | 0/30 [00:00<?, ?it/s]CUDA error (/__w/xformers/xformers/third_party/flash-attention/hopper/flash_fwd_launch_template.h:188): invalid argument
[INFO] Running set-proxy script...
[INFO] Continue without proxy.

then container was restarted with driver version 580.76.05 and image yanwk/comfyui-boot:cu128-slim，How can I fix this error?

[pedroleiteluizalabs]

same error here...

[YanWenKun]

Hi,

I guess there's something wrong with xformers/third_party/flash-attention/hopper since you're running on Blackwell.

You can either:

1. Run with xFormers disabled:

docker run \

...

  -e CLI_ARGS="--use-pytorch-cross-attention --fast" \
  yanwk/comfyui-boot:cu128-slim

2. Try the new cu129 images:

https://github.com/YanWenKun/ComfyUI-Docker/tree/main/cu129-slim
https://github.com/YanWenKun/ComfyUI-Docker/tree/main/cu129-megapak

[DarthGibblet]

Just wanted to chime in here with another data point. Unfortunately upgrading to the cu129 branch didn't fix the error, but moving away from xformers with --use-pytorch-cross-attention did. Hope this helps anybody else having the same issue!

[YanWenKun]

It's a confirmed bug for Blackwell GPUs, the solution is to disable xFormers as showed above. Or wait for xFormers to update.

I've updated docs, now all cu128 and later images will suggest to disable xFormers by default.