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.
This commit is contained in:
Dan Sedano 2026-08-13 12:49:01 -07:00
parent 1183f60828
commit 2bc2e8d73e
1 changed files with 63 additions and 0 deletions

View File

@ -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