Comfy-Org / Comfy-Org/ComfyUI

Option not to use quantized attention in Cross Attention

Open
#10,079 0 comments 1 reaction 0 assignees View on GitHub
Feature
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Feature Idea

From the code at the attention class https://github.com/comfyanonymous/ComfyUI/blob/555f902fc1ed20e98201f9102172f0fc190c2c42/comfy/ldm/modules/attention.py#L675 I get that Cross Attention uses `optimized_attention`, which can be, for example, sage attention (especially, the new sage attention 3), meaning the cross attention has been quantized.

Cross attention (image-text attention) because of its nature is the cornerstone of prompt-following in non-MMDiT models (example: Wan2.1/2.2). While self attention (image-image attention, seq_lenseq_len in complexity) is the bottleneck and it is justified to lower the precision, cross attention (image-text attention, seq_lentext_len in complexity) takes much less resources (tens of thousands in image tokens vs. max 512 tokens for text) and it is not really justified to quantize it at the cost of prompt adherence.

The solution is to propose is to add a CLI option like `--cross-attn-type flash`, to set after the quantization option like `--use-sage-attention`

### Existing Solutions

_No response_

### Other

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.