Fix small_input attention fallback OOMing without pytorch attention

[chelsealong]

Fixes #16160

Problem

optimized_attention_for_device(small_input=True) only offered a two-way
choice: attention_pytorch when pytorch attention (SDPA) is enabled, or
attention_basic otherwise. attention_basic materializes the full
N x N score matrix, which is fine for typical short text-encoder
sequences but becomes the largest allocation in the run once N grows —
e.g. MiniMax H3's <Picture i> reference-image tokens are appended to the
same sequence CLIP encodes with small_input=True
(comfy/text_encoders/llama.py:757), so encoding several reference images
pushes N into the thousands.

On GPUs where pytorch_attention_enabled() is False — notably AMD
RDNA2-and-older architectures (gfx1030, gfx1031, gfx1035, gfx1010,
gfx1011, gfx1012, gfx906, gfx900, gfx803), which ship no aotriton
SDPA kernels — there was no way to avoid this path. --use-pytorch-cross-attention
makes it worse (falls through to PyTorch's math SDPA backend, which uses
~5x more memory than attention_basic at the same shape per the issue's
own benchmark), and --use-split-cross-attention / --use-quad-cross-attention
don't help either since the small_input=True branch is hardcoded and
bypasses the module-level optimized_attention rebinding those flags use.

Fix

Change the small_input=True, no-pytorch-attention fallback from
attention_basic to attention_sub_quad, which:

• has an identical call signature (mask, skip_reshape,
  skip_output_reshape, enable_gqa via **kwargs),
• is already ComfyUI's own hardware-agnostic, chunked default fallback
  used elsewhere in this file (e.g. for device == cpu),
• per the issue's measurements, drops peak allocation for the reported
  repro from 10.25 GiB (OOM) to ~1.3 GiB.

attention_split was considered and rejected (by the issue reporter) as
the alternative: its chunk size falls back to no chunking at all when the
sequence length isn't evenly divisible by the step count, which can still
OOM.

Test plan

Added tests-unit/comfy_test/optimized_attention_for_device_test.py,
asserting that optimized_attention_for_device(..., small_input=True)
returns attention_sub_quad (not attention_basic) when
pytorch_attention_enabled() is False.

This sandbox (and CI's unit-test job, which installs CPU-only torch)
has no CUDA device, and importing comfy.ldm.modules.attention pulls
in comfy.model_management, which calls torch.cuda.current_device()
at module import time. The test file guards against this the same way
tests-unit/comfy_test/gemma4_template_test.py already does — setting
comfy.cli_args.args.cpu = True before the import when CUDA isn't
available — so it collects and passes standalone, with no dependency
on import order relative to other test files.

Ran the exact command from .github/workflows/test-unit.yml in
isolation, both before and after the source fix:

$ python -m pytest tests-unit/comfy_test/optimized_attention_for_device_test.py -v
# before the fix (comfy/ldm/modules/attention.py reverted via `git checkout HEAD~2 -- ...`):
FAILED tests-unit/comfy_test/optimized_attention_for_device_test.py::test_small_input_without_pytorch_attention_uses_chunked_fallback
AssertionError: assert <function attention_basic at 0x7f815d159120> is attention_sub_quad
1 failed in 2.58s

# after the fix:
tests-unit/comfy_test/optimized_attention_for_device_test.py::test_small_input_without_pytorch_attention_uses_chunked_fallback PASSED
1 passed in 2.57s

Also ran ruff check on both changed files: all checks passed.

🤖 Generated with Claude Code

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Advanced

Run ID: 195b96e3-fda8-49b4-94e4-6c6f1ad27d89

📥 Commits

Reviewing files that changed from the base of the PR and between a7b1d39 and 47eee96.

📒 Files selected for processing (2)

• comfy/ldm/modules/attention.py
• tests-unit/comfy_test/optimized_attention_for_device_test.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📜 Recent review details

⏰ Context from checks skipped due to timeout. (8)

• GitHub Check: test (macos-latest)
• GitHub Check: test (ubuntu-latest)
• GitHub Check: test (ubuntu-latest)
• GitHub Check: test (windows-2022)
• GitHub Check: test (macos-latest)
• GitHub Check: test (windows-latest)
• GitHub Check: Run Pylint
• GitHub Check: test

🧰 Additional context used

📓 Path-based instructions (3)

Core ML/diffusion engine.

⚙️ CodeRabbit configuration file

Files:

• comfy/ldm/modules/attention.py

IMPORTANT: Only comment on issues directly introduced by this PR's code changes.

⚙️ CodeRabbit configuration file

Files:

• tests-unit/comfy_test/optimized_attention_for_device_test.py
• comfy/ldm/modules/attention.py

Documentation and README edits should be concise, factual, and tied to the changed behavior.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

• tests-unit/comfy_test/optimized_attention_for_device_test.py
• comfy/ldm/modules/attention.py

🔇 Additional comments (2)

comfy/ldm/modules/attention.py (1)

911-911: LGTM!

tests-unit/comfy_test/optimized_attention_for_device_test.py (1)

1-22: LGTM!

---

📝 Walkthrough

Walkthrough

The small-input branch of optimized_attention_for_device now returns attention_sub_quad when PyTorch attention is unavailable. A unit test patches the attention capability check and verifies that CPU selection returns attention_sub_quad rather than attention_basic.

Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 47eee

The change prevents full attention score materialization for the affected fallback path while preserving targeted regression coverage, so it is ready to merge after normal checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the small-input attention fallback change and the OOM problem it addresses.
Description check	✅ Passed	The description directly explains the OOM cause, the fallback change, affected hardware, and the test results.
Linked Issues check	✅ Passed	The PR implements issue #16160 by selecting chunked attention_sub_quad instead of attention_basic when PyTorch attention is unavailable, and adds coverage for the required fallback behavior.
Out of Scope Changes check	✅ Passed	The source change and CPU-safe unit test directly support the linked issue and stated objectives. No unrelated code changes are described.

• Fix all pre-merge checks with AI

---

_{Comment @coderabbitai help to get the list of available commands.}