Comfy-Org / Comfy-Org/ComfyUI

--use-flash-attention is ignored by the memory estimator, causing a 7.5x overestimate → forced partial load / PCIe thrashing (AMD, torch wheels without aotriton)

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

Description

## Custom Node Testing

- [x] Reproduced with the estimator math directly in core code paths; custom nodes not involved

## Expected Behavior

With `--use-flash-attention`, `BaseModel.memory_required()` should use the
memory-efficient attention estimate. Flash attention's memory footprint is
comparable to pytorch SDPA's, and SDPA does get the efficient estimate.

## Actual Behavior

`BaseModel.memory_required()` (comfy/model_base.py:425) selects the estimate
formula with:

```python
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
# efficient: area * dtype_size * 0.01 * memory_usage_factor
else:
# conservative: area * 0.15 * memory_usage_factor <- 7.5x larger at bf16
```

`flash_attention_enabled()` is never consulted. A user running
`--use-flash-attention` only gets the efficient estimate if
`ENABLE_PYTORCH_ATTENTION` also happens to be on. On AMD that depends on
`aotriton_supported()` (comfy/model_management.py:493), which is a directory
listing of `torch/lib/aotriton.images` inside the wheel.

Secondary bug: when that directory does not exist, `os.listdir` raises
`FileNotFoundError`. The broad `try:` wrapping the AMD block
(model_management.py:480) swallows it in a bare `except: pass`, which also
aborts the `SUPPORT_FP8_OPS` detection below it. `ENABLE_PYTORCH_ATTENTION`
silently stays False with no log line, and the estimator flips to the
conservative formula.

Wheel packaging context: torch 2.13.0+rocm7.2 ships `aotriton.images` (it
contains `amd-gfx110x`, which matches gfx1101). torch 2.12.0+rocm7.14.0 ships
`libaotriton_v2.so.0.11.2` in the same `torch/lib/` but not the kernel-image
directory, so the library is present while the directory check fails. The
images are not relocated elsewhere in the modular ROCm SDK (no aotriton entry
in any `rocm_sdk_*` module manifest, including `rocm_sdk_device_gfx1101`), and
`libaotriton_v2` resolves `aotriton.images` relative to its own location with
no env override. ComfyUI's check looks in the same place the library itself
would; the files are simply not in the wheel. That part is AMD's packaging
problem, but ComfyUI turns it into a silent 7.5x estimate inflation.

## Consequences (measured, RX 7800 XT 16 GB / gfx1101, torch 2.12.0+rocm7.14.0)

Krea2 fp8 (12.5 GB), 2560x1440 txt2img, latent (1, 16, 1, 180, 320), area 57600,
`memory_usage_factor = 2.2`, bf16:

| | estimate formula | predicted | observed (logs) |
|---|---|---|---|
| torch wheel with aotriton (2.13.0+rocm7.2) | efficient | 2534 MB | min_req ≈4.1-4.5 GB: normal path, 1.6 GB offloaded |
| torch wheel without aotriton (2.12.0+rocm7.14) | conservative | 19008 MB | min_req 20544 MB: estimate exceeds VRAM |

Observed min_req runs about 1.5 GB above the raw formula in both rows. That is
an additive overhead from additional models attached to the conditioning, and
it does not affect which branch dominates.

When the estimate exceeds free VRAM, `load_models_gpu` falls through to the
`MIN_WEIGHT_MEMORY_RATIO` floor (0.4 default for non-NVIDIA; 0.0 on NVIDIA).
Only ≈6.1 GB of the 12.5 GB model stays resident and ≈7 GB streams over PCIe
on every sampling step:

```
loaded partially; 6140.00 MB usable, 5468.00 MB loaded, 7062.31 MB offloaded, lowvram patches: 206
```

Host RSS hit 27.8 GB on a 30 GB box (the offloaded weights sit in RAM), heavy
swap pressure, roughly 15 minutes per image.

With a one-line change adding `flash_attention_enabled()` to the check:

```
loaded completely; 12939.60 MB usable, 12532.86 MB loaded, full load: True
```

328 s per image (≈2.7x faster), zero swap. Wan 2.2 I2V, a large-latent
workload, was verified unaffected by the change: its post-patch efficient
estimate (≈12.4 GB at latent (1,16,21,86,152), factor dim/2222 = 2.304) still
leaves `free - min_req` (≈3 GB) below the 0.4 floor (≈6.25 GB), so the same
floor governs the load before and after. We observed the identical
`loaded partially; 6248 MB usable` both ways, with zero OOM.

## Proposed fix

```python
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention() or comfy.model_management.flash_attention_enabled():
```

The same reasoning probably applies to `sage_attention_enabled()`; we did not
measure that case.

For the secondary bug, `aotriton_supported()` could catch `FileNotFoundError`
and return False explicitly, and log one line saying pytorch attention was not
auto-enabled because the wheel lacks aotriton images. As it stands, the
failure is invisible: the only symptom is that model loads get slower and
partial, several layers of cause away.

## Steps to Reproduce

1. AMD RDNA3 (gfx1101), any torch ROCm wheel whose `torch/lib/` has no
`aotriton.images` directory (e.g. 2.12.0+rocm7.14.0 from
repo.amd.com/rocm/whl-multi-arch)
2. Launch with `--use-flash-attention`
3. Run any model and resolution where the conservative estimate exceeds free
VRAM but the efficient one does not (e.g. a 12.5 GB fp8 checkpoint at
2560x1440 on 16 GB). Observe `loaded partially` plus per-step
host-to-device streaming.
4. Confirm branch selection (from the ComfyUI root, same venv):
`python -c "import comfy.model_management as mm; print(mm.ENABLE_PYTORCH_ATTENTION, mm.pytorch_attention_flash_attention(), mm.xformers_enabled())"`
prints `False False False`. All three disjuncts of the estimator condition
are off while flash attention is the active runtime backend.

## Debug Logs

```
[INFO] Total VRAM 16368 MB, total RAM 31288 MB
[INFO] pytorch version: 2.12.0+rocm7.14.0
[INFO] AMD arch: gfx1101
[INFO] ROCm version: (7, 14)
[INFO] Set vram state to: NORMAL_VRAM
[INFO] Using Flash Attention
[DETAIL] Sampler: model=Krea2 latent_shapes=[(1, 16, 1, 180, 320)]
[INFO] Requested to load Krea2
[DETAIL] Non dynamic memory free called! memory_required=55929143367.2 pins_required=13141695832 ram_required=0
[DETAIL] Non dynamic memory free called! memory_required=21541945344.0 pins_required=0 ram_required=0
[INFO] 0 models unloaded.
[INFO] loaded partially; 6140.00 MB usable, 5468.00 MB loaded, 7062.31 MB offloaded, 672.00 MB buffer reserved, lowvram patches: 206
```

After the one-line fix, same workflow:

```
[INFO] loaded completely; 12939.60 MB usable, 12532.86 MB loaded, full load: True
[INFO] Prompt executed in 328.66 seconds
```

## Other

Searched existing issues.
[#15356](https://github.com/Comfy-Org/ComfyUI/issues/15356) concerns the LTXAV
estimate value (video-only vs nested AV latents), not attention-backend flag
selection; [#15443](https://github.com/Comfy-Org/ComfyUI/issues/15443) and
[#15481](https://github.com/Comfy-Org/ComfyUI/issues/15481) are eviction and
RAM-cache complaints. None covers this.

Flash attention on this box is not a wheel: flash-attn 2.8.4 is built from
source (Dao-AILab/flash-attention @ 2409214a, aiter submodule) as a uv path
dependency, running the Triton AMD backend
(`FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE`). So `--use-flash-attention` is a
supported, working configuration on ROCm; the estimator just never asks about
it.

What OS: Linux (Fedora 44)
GPU: AMD RX 7800 XT 16 GB (gfx1101), headless compute
Launch flags: `--use-flash-attention --disable-pinned-memory --reserve-vram 0 --cache-none --fast-disk --enable-triton-backend --async-offload`

Contributor guide

Open the contributing guide

Research direction

Start in comfy/model_base.py around memory_required() and inspect the attention checks in comfy/model_management.py, including flash_attention_enabled() and aotriton_supported(). Reproduce the reported branch selection with the provided Python command on an AMD setup, then verify that --use-flash-attention selects the efficient estimate and avoids the reported partial-load behavior; handle the missing aotriton directory only if that secondary scope is included.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.