`optimized_attention_for_device(small_input=True)` falls back to `attention_basic`, OOMing on GPUs without aotriton SDPA kernels

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.

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:

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 <Picture i> 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 <Picture 1> .. <Picture 8>.
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

     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,

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.