From 1183f608288b5bfe988c88f941ec5f699ff832b8 Mon Sep 17 00:00:00 2001 From: Dan Sedano <47065799+thedandano@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:56:40 -0700 Subject: [PATCH] Use efficient attention memory estimate when --use-flash-attention is set memory_required() only checks xformers_enabled() and pytorch_attention_flash_attention() when picking between the efficient (area * dtype_size * 0.01 * factor) and conservative (area * 0.15 * factor) estimate formulas. With --use-flash-attention the conservative formula is used even though flash attention's memory footprint is comparable to SDPA's, inflating the estimate 7.5x at bf16. On AMD this bites whenever ENABLE_PYTORCH_ATTENTION is not auto-enabled (e.g. torch ROCm wheels that ship without aotriton.images, like 2.12.0+rocm7.14.0): a 12.5 GB fp8 model at 2560x1440 on a 16 GB card gets a 20.5 GB estimate, falls onto the MIN_WEIGHT_MEMORY_RATIO floor, and streams 7 GB of weights over PCIe every step. With this change the same workflow fully loads and runs 2.7x faster. Verified no behavior change for large-latent workloads (Wan 2.2 I2V) that legitimately need the floor. Fixes #15585 --- comfy/model_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index 90cab7ac0..aec298216 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -423,7 +423,7 @@ class BaseModel(torch.nn.Module): if len(shape) > 0: input_shapes += shape - if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention(): + if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention() or comfy.model_management.flash_attention_enabled(): dtype = self.get_dtype_inference() #TODO: this needs to be tweaked area = sum(map(lambda input_shape: input_shape[0] * math.prod(input_shape[2:]), input_shapes))