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 1/3] 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)) From 2bc2e8d73e6686e7ee54e84c7a0416708c4df1da Mon Sep 17 00:00:00 2001 From: Dan Sedano <47065799+thedandano@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:49:01 -0700 Subject: [PATCH 2/3] Add unit tests for the attention-dependent memory estimate Covers which attention checks select the efficient formula in BaseModel.memory_required, including the new flash_attention_enabled() path, and pins the 7.5x bf16 ratio between the conservative and efficient estimates so a retune of either formula fails the test and gets updated consciously. --- tests-unit/comfy_test/memory_required_test.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 tests-unit/comfy_test/memory_required_test.py diff --git a/tests-unit/comfy_test/memory_required_test.py b/tests-unit/comfy_test/memory_required_test.py new file mode 100644 index 000000000..76e2d45a3 --- /dev/null +++ b/tests-unit/comfy_test/memory_required_test.py @@ -0,0 +1,63 @@ +import math + +import torch + +import comfy.model_base +import comfy.model_management + + +class _StubModel: + """Minimal stand-in for BaseModel: just the attributes memory_required reads.""" + memory_usage_factor_conds = () + memory_usage_shape_process = {} + memory_usage_factor = 2.0 + + def get_dtype_inference(self): + return torch.bfloat16 + + +INPUT_SHAPE = (1, 16, 1, 180, 320) +AREA = INPUT_SHAPE[0] * math.prod(INPUT_SHAPE[2:]) +DTYPE_SIZE = 2 # bf16 +EFFICIENT = AREA * DTYPE_SIZE * 0.01 * _StubModel.memory_usage_factor * (1024 * 1024) +CONSERVATIVE = AREA * 0.15 * _StubModel.memory_usage_factor * (1024 * 1024) + + +def _patch_attention(monkeypatch, xformers=False, pytorch_flash=False, flash=False): + monkeypatch.setattr(comfy.model_management, "xformers_enabled", lambda: xformers) + monkeypatch.setattr(comfy.model_management, "pytorch_attention_flash_attention", lambda: pytorch_flash) + monkeypatch.setattr(comfy.model_management, "flash_attention_enabled", lambda: flash) + + +def _estimate(): + return comfy.model_base.BaseModel.memory_required(_StubModel(), INPUT_SHAPE) + + +def test_no_efficient_attention_uses_conservative_estimate(monkeypatch): + _patch_attention(monkeypatch) + assert _estimate() == CONSERVATIVE + + +def test_pytorch_flash_attention_uses_efficient_estimate(monkeypatch): + _patch_attention(monkeypatch, pytorch_flash=True) + assert _estimate() == EFFICIENT + + +def test_flash_attention_flag_uses_efficient_estimate(monkeypatch): + # --use-flash-attention must select the efficient estimate even when + # pytorch attention was not auto enabled (e.g. torch builds without + # working aotriton), otherwise the estimate is 7.5x too large. + _patch_attention(monkeypatch, flash=True) + assert _estimate() == EFFICIENT + + +def test_conservative_estimate_is_7_5x_efficient_at_bf16(monkeypatch): + # Documents the size of the gap between the two formulas: at bf16 the + # conservative path asks for 7.5x more working memory than the efficient + # one for the same shapes. If either formula is retuned, this ratio (and + # the impact of picking the wrong branch) changes; update it consciously. + _patch_attention(monkeypatch, flash=True) + efficient = _estimate() + _patch_attention(monkeypatch) + conservative = _estimate() + assert conservative == 7.5 * efficient From b5b5c40fd7660abb5173fd0da1eb0be5a4cec1c5 Mon Sep 17 00:00:00 2001 From: Dan Sedano <47065799+thedandano@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:15:52 -0700 Subject: [PATCH 3/3] Add xformers-only estimate test --- tests-unit/comfy_test/memory_required_test.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests-unit/comfy_test/memory_required_test.py b/tests-unit/comfy_test/memory_required_test.py index 76e2d45a3..37ad8e498 100644 --- a/tests-unit/comfy_test/memory_required_test.py +++ b/tests-unit/comfy_test/memory_required_test.py @@ -43,6 +43,11 @@ def test_pytorch_flash_attention_uses_efficient_estimate(monkeypatch): assert _estimate() == EFFICIENT +def test_xformers_uses_efficient_estimate(monkeypatch): + _patch_attention(monkeypatch, xformers=True) + assert _estimate() == EFFICIENT + + def test_flash_attention_flag_uses_efficient_estimate(monkeypatch): # --use-flash-attention must select the efficient estimate even when # pytorch attention was not auto enabled (e.g. torch builds without