`optimized_attention_for_device(small_input=True)` falls back to `attention_basic`, OOMing on GPUs without aotriton SDPA kernels
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Custom Node Testing
- [ ] I have tried disabling custom nodes and the issue persists
Not ticked, and I want to be straight about why rather than claim a test I did not run: I have already applied the local fix below, so the crash no longer reproduces on my install. What I can say precisely is that **every node in the reproduction graph is core** (verified against a live `/object_info` dump — `LoadImage`, `ResizeImagesByLongerEdge` from `comfy_extras.nodes_dataset`, `CLIPLoader`, `VAELoader`, `MiniMaxH3ReferenceToVideo`, `PreviewAny`), and **the entire traceback is inside core files** — `comfy_extras/nodes_minimax_h3.py` → `comfy/sd1_clip.py` → `comfy/text_encoders/*` → `comfy/ldm/modules/attention.py`. No custom node appears anywhere in the call path.
The log below does contain a few `[MultiGPU ...]` lines from a device-placement custom node. Those are incidental — it only pins which GPU a loader targets and does not appear in the traceback. Happy to reproduce on a clean install if that would help.
### Expected Behavior
Encoding a MiniMax H3 `ref2va` prompt that references several images should succeed, or at worst fall back to a memory-efficient attention implementation. ComfyUI already ships two that are hardware-agnostic and chunked (`attention_sub_quad`, `attention_split`).
### Actual Behavior
It OOMs in the text encoder before sampling starts, trying to allocate a single **10.25 GiB** tensor.
`optimized_attention_for_device()` offers only a two-way choice when `small_input=True`, and the fallback materializes the full N×N score matrix:
```python
def optimized_attention_for_device(device, mask=False, small_input=False):
if small_input:
if model_management.pytorch_attention_enabled():
return attention_pytorch
else:
return attention_basic # <-- materializes N x N
```
`comfy/text_encoders/llama.py:757` requests this path with `small_input=True`. That assumption holds for a plain text prompt, but MiniMax H3's `` reference images are tokenized into the **same** sequence (`comfy/text_encoders/minimax.py:157`), so N grows with the number of references and the quadratic matrix becomes the largest allocation in the run.
On any GPU where `pytorch_attention_enabled()` is False there is currently no supported way to avoid this path — see Other for why `--use-pytorch-cross-attention` and the split/quad flags do not help.
### Steps to Reproduce
Core-only graph. Image content and prompt text are irrelevant — only the resulting sequence length matters, so any 8 images and any prompt reproduce it:
```
CLIPLoader (MiniMax H3 text encoder, type "minimax")
VAELoader (H3 video VAE)
VAELoader (H3 audio VAE)
8x [ LoadImage -> ResizeImagesByLongerEdge (longer_edge = 768) ]
-> MiniMaxH3ReferenceToVideo.ref_images.ref_image_0..7
(width 1344, height 768, length 175, ref_image_size "match")
-> PreviewAny
```
1. Build the graph above on an AMD RDNA2 GPU (gfx1030).
2. Write a prompt referencing `` .. ``.
3. Queue it.
It OOMs during `MiniMaxH3ReferenceToVideo`, before sampling. The same graph with **2** reference images completes normally — this scales with reference count, not prompt length. Each 768 px reference contributes roughly 1,344 vision tokens, so 8 references add ~10.7k tokens and the score matrix is quadratic in that total.
### Debug Logs
Model filenames are generalized below; they are ordinary public MiniMax H3 weights.
```
got prompt
VAE load device: cuda:1, offload device: cpu, dtype: torch.float16
Requested to load MiniMaxH3VideoVAE
loaded completely; 17439.22 MB usable, 4966.19 MB loaded, full load: True
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
gguf qtypes: Q4_K (390), F32 (433), Q6_K (50), Q5_K (27), F16 (2)
!!! Exception during processing !!! CUDA out of memory. Tried to allocate 10.25 GiB.
GPU 0 has a total capacity of 29.98 GiB of which 7.65 GiB is free.
Of the allocated memory 20.79 GiB is allocated by PyTorch, and 1.08 GiB is reserved
by PyTorch but unallocated.
Traceback (most recent call last):
File "ComfyUI/execution.py", line 545, in execute
File "ComfyUI/execution.py", line 344, in get_output_data
File "ComfyUI/execution.py", line 318, in _async_map_node_over_list
File "ComfyUI/execution.py", line 306, in process_inputs
result = f(**inputs)
File "ComfyUI/comfy_api/latest/_io.py", line 1990, in EXECUTE_NORMALIZED
File "ComfyUI/comfy_extras/nodes_minimax_h3.py", line 277, in execute
cond = clip.encode_from_tokens_scheduled(tokens)
File "ComfyUI/comfy/sd.py", line 336, in encode_from_tokens_scheduled
File "ComfyUI/comfy/sd.py", line 405, in encode_from_tokens
File "ComfyUI/comfy/sd1_clip.py", line 743, in encode_token_weights
File "ComfyUI/comfy/text_encoders/minimax.py", line 110, in encode_token_weights
File "ComfyUI/comfy/sd1_clip.py", line 45, in encode_token_weights
File "ComfyUI/comfy/sd1_clip.py", line 306, in encode
File "ComfyUI/comfy/sd1_clip.py", line 279, in forward
File "ComfyUI/comfy/text_encoders/minimax.py", line 96, in forward
File "ComfyUI/comfy/text_encoders/qwen3vl.py", line 100, in forward
File "ComfyUI/comfy/text_encoders/llama.py", line 782, in forward
File "ComfyUI/comfy/text_encoders/llama.py", line 602, in forward
File "ComfyUI/comfy/text_encoders/llama.py", line 565, in forward
output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask,
skip_reshape=True, **gqa_kwargs)
File "ComfyUI/comfy/ldm/modules/attention.py", line 160, in wrapper
return func(*args, **kwargs)
File "ComfyUI/comfy/ldm/modules/attention.py", line 194, in attention_basic
sim = einsum('b i d, b j d -> b i j', q, k) * scale
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 10.25 GiB.
Memory summary (device 0):
| CUDA OOMs: 1 | cudaMalloc retries: 1 |
| Allocated memory | 21288 MiB | 23975 MiB |
| GPU reserved memory | 22394 MiB | 27740 MiB |
```
After the one-line change below, the identical graph:
```
Requested to load MiniMaxH3VideoVAE
loaded completely; 29458.80 MB usable, 4966.19 MB loaded, full load: True
Requested to load MiniMaxH3TEModel_
loaded completely; 29458.80 MB usable, 15393.81 MB loaded, full load: True
Prompt executed in 158.43 seconds
```
### Other
**Environment**
| | |
|---|---|
| GPU | 2 x AMD Radeon Pro V620 (Navi 21, **gfx1030**, RDNA2), 32 GB GDDR6 each, `[1002:73a1]` |
| CPU | AMD EPYC 7452 (32c / 64t, Zen 2) |
| Host | Proxmox VE 9.1.7, kernel 6.17.13-2-pve, 128 GB DDR4 ECC |
| Guest | Ubuntu 24.04.4 LTS, kernel 6.8.0, 24 vCPU, 76 GB RAM; both GPUs via `vfio-pci` |
| Stack | ROCm 7.2, torch 2.11.0+rocm7.2, Python 3.13.12, ComfyUI 0.31.0 |
Passthrough is not relevant — the path is selected purely by `pytorch_attention_enabled()`.
**Not V620-specific.** `gfx1030` is Navi 21, so this equally affects the RX 6800 / 6800 XT / 6900 XT and the Radeon Pro W6800, and more broadly every architecture in ComfyUI's own `AMD_RDNA2_AND_OLDER_ARCH` list (`gfx1030, gfx1031, gfx1035, gfx1010, gfx1011, gfx1012, gfx906, gfx900, gfx803`), none of which ship aotriton kernels.
**Why the pytorch path is unavailable, and why forcing it is worse**
`model_management.py` only auto-enables `ENABLE_PYTORCH_ATTENTION` on AMD for architectures with aotriton kernels:
```
arch: gfx1030
aotriton images: ['gfx11xx', 'gfx120x', 'gfx90a', 'gfx942', 'gfx950']
aotriton_supported(gfx1030) = False
```
Forcing `--use-pytorch-cross-attention` makes it worse. Direct backend probe, B=1 H=8 N=4096 D=128 fp16:
```
FLASH FAIL No available kernel. ("Flash attention was not compiled for
current AMD GPU architecture ... gfx1030")
EFFICIENT FAIL No available kernel. ("Mem Efficient attention was not compiled
for current AMD GPU architecture")
MATH OK peak +1248.0 MiB
EINSUM OK peak + 256.0 MiB <- attention_basic
```
SDPA falls through to the math backend, which used **~4.9x more** memory than `attention_basic` at identical shape. `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` (already set unconditionally in `main.py`) changes nothing.
`--use-split-cross-attention` / `--use-quad-cross-attention` also do not affect this path — they only rebind the module-level `optimized_attention`, while the `small_input=True` branch is hardcoded.
**Suggested fix**
```diff
if small_input:
if model_management.pytorch_attention_enabled():
return attention_pytorch
else:
- return attention_basic
+ return attention_sub_quad
```
`attention_sub_quad` has an identical signature, handles `mask` / `skip_reshape` / `enable_gqa` the same way, and is already ComfyUI's own default fallback for the main model elsewhere in the same file.
`attention_split` is the other candidate but looks less safe: its slicing is gated on divisibility,
```python
slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
```
so for a sequence length not divisible by `steps` it silently does no chunking at all, OOMs, doubles `steps`, and can never become divisible.
**Measured result**
| | peak allocation |
|---|---|
| `attention_basic` (current) | **10.25 GiB** (OOM) |
| `attention_sub_quad` (patched) | **~1.3 GiB** |
Measured by sampling free VRAM at 1 Hz across a full encode: idle baseline with the encoder resident was 15.97 GiB free, minimum observed during the run 14.68 GiB. Two runs completed with `status_str: success`. Had the quadratic path still been active, free memory would have had to drop to about 4.7 GiB.
Happy to open a PR if this direction is acceptable.
Contributor guide
Research direction
Start in comfy/ldm/modules/attention.py at optimized_attention_for_device() and compare the small_input fallback with attention_sub_quad; then inspect the call from comfy/text_encoders/llama.py:757. Reproduce the MiniMax H3 graph with eight reference images on hardware without aotriton kernels, and consider the issue done when encoding completes without the quadratic OOM while the existing attention behavior remains intact.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100