Add Flex Attention and expose torch.compile modes by apersomany · Pull Request #15870 · Comfy-Org/ComfyUI

apersomany · 2026-08-24T20:52:59Z

Summary

This PR adds an opt-in PyTorch Flex Attention backend and exposes torch.compile backends and modes through TorchCompileModel.

The changes are complementary: Flex Attention benefits substantially from selecting an appropriate compile mode, especially max-autotune-no-cudagraphs on ROCm.

Changes

Flex Attention

Add --use-flex-attention.
Register Flex Attention as an attention backend.
Support masks, scaling, GQA, and existing reshape modes.
Use the raw Flex Attention implementation while an outer torch.compile graph is being traced.
Fall back to PyTorch attention for unsupported backend runtime errors while propagating out-of-memory errors.
Preserve existing attention behavior when Flex Attention is not selected.

TorchCompileModel

Expose all backends returned by torch.compiler.list_backends().
Add an advanced compile mode input supporting values such as default, max-autotune, and max-autotune-no-cudagraphs.
Preserve the existing guard_filter_fn behavior when no compile mode is selected.
Omit guard_filter_fn when a mode is selected because PyTorch rejects combining that option with mode.

Benchmark evidence

The attention shapes below were derived from an actual Anima Turbo sampling pass. For each input size, the largest observed attention call was replayed by the microbenchmark.

Measurements used FP16 on AMD Radeon Graphics with ROCm 7.2 and PyTorch 2.13.0+rocm7.2. default means Inductor compilation without an explicit compile mode. Values are steady-state medians; compilation and autotuning startup are excluded from timed evaluations.

512²

Q: [1, 16, 1024, 128]
K: [1, 16, 1024, 128]
V: [1, 16, 1024, 128]
Attention FLOPs: 8.59 GFLOPs

Attention backend	Default	Max autotune	Max autotune, no cudagraphs
PyTorch (AOTriton Flash)	0.38 ms (22.61 TFLOP/s)	0.45 ms (19.27 TFLOP/s)	0.38 ms (22.71 TFLOP/s)
Flex (Triton 3.7.1)	0.42 ms (20.59 TFLOP/s)	0.38 ms (22.38 TFLOP/s)	0.30 ms (28.82 TFLOP/s)
Flex (Triton 3.8.0 + local patches)	0.54 ms (15.89 TFLOP/s)	0.31 ms (27.77 TFLOP/s)	0.25 ms (34.27 TFLOP/s)

1024²

Q: [1, 16, 4096, 128]
K: [1, 16, 4096, 128]
V: [1, 16, 4096, 128]
Attention FLOPs: 137.44 GFLOPs

Attention backend	Default	Max autotune	Max autotune, no cudagraphs
PyTorch (AOTriton Flash)	4.21 ms (32.66 TFLOP/s)	4.51 ms (30.45 TFLOP/s)	4.29 ms (32.00 TFLOP/s)
Flex (Triton 3.7.1)	6.24 ms (22.03 TFLOP/s)	3.77 ms (36.48 TFLOP/s)	3.30 ms (41.64 TFLOP/s)
Flex (Triton 3.8.0 + local patches)	7.80 ms (17.63 TFLOP/s)	3.01 ms (45.59 TFLOP/s)	2.59 ms (53.02 TFLOP/s)

2048²

Q: [1, 16, 16384, 128]
K: [1, 16, 16384, 128]
V: [1, 16, 16384, 128]
Attention FLOPs: 2.20 TFLOPs

Attention backend	Default	Max autotune	Max autotune, no cudagraphs
PyTorch (AOTriton Flash)	59.46 ms (36.98 TFLOP/s)	60.69 ms (36.24 TFLOP/s)	59.66 ms (36.86 TFLOP/s)
Flex (Triton 3.7.1)	97.16 ms (22.63 TFLOP/s)	49.31 ms (44.59 TFLOP/s)	47.67 ms (46.13 TFLOP/s)
Flex (Triton 3.8.0 + local patches)	114.05 ms (19.28 TFLOP/s)	40.04 ms (54.93 TFLOP/s)	38.70 ms (56.82 TFLOP/s)

These results suggest that max-autotune-no-cudagraphs may be a better default for Flex Attention compilation on ROCm, and possibly CUDA. No default change is proposed here: this needs broader testing across NVIDIA and AMD hardware, driver versions, PyTorch/Triton versions, model families, masks, GQA, sequence lengths, and compilation cold-start costs.

The Triton 3.8.0 local-patch results are supporting evidence only; this PR does not add a Triton dependency.

Appendix: Reproducing the benchmarks

The benchmark implementation is available at bench.py @ 73abe92.

The aggregated results and raw JSON outputs are available in results/flex-attention-sdpa @ 252aed9.

Common arguments:

/path/to/anima-turbo.safetensors --device cuda:0 --dtype fp16 --context-length 64 --target sdpa --attention BACKEND --compile inductor --json

The size-specific arguments were:

512²:  --width 512  --height 512  --steps 16 --warmup-steps 6
1024²: --width 1024 --height 1024 --steps 8  --warmup-steps 3
2048²: --width 2048 --height 2048 --steps 8  --warmup-steps 3

BACKEND was pytorch or flex. The compile-mode variants used:

default:                         omit --compile-mode
max-autotune:                    --compile-mode max-autotune
max-autotune-no-cudagraphs:      --compile-mode max-autotune-no-cudagraphs

Compilation runs were launched with CXX=clang++ in a nix shell nixpkgs#clang environment. The patched Flex results used Triton 3.8.0 with the local patches; the baseline Flex results used Triton 3.7.1.

Testing

Existing full-model sampling behavior remains unchanged when Flex Attention is not selected.
Flex Attention was tested at 512², 1024², and 2048² attention shapes.
Compile modes default, max-autotune, and max-autotune-no-cudagraphs were exercised.
Changed Python files pass syntax compilation.

github-actions · 2026-08-24T20:53:14Z

✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
_{Posted by the CLA Assistant Lite bot.}

apersomany · 2026-08-24T21:06:02Z

I have read and agree to the Contributor License Agreement

coderabbitai · 2026-08-24T21:23:14Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 6c507e7c-3f1f-4b41-901d-97d0d77f248b

📥 Commits

Reviewing files that changed from the base of the PR and between eb8cad7 and 99e2986.

📒 Files selected for processing (4)

comfy/cli_args.py
comfy/ldm/modules/attention.py
comfy/model_management.py
comfy_extras/nodes_torch_compile.py

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

📜 Recent review details

🧰 Additional context used

📓 Path-based instructions (7)

Community-contributed extra nodes. Focus on:

⚙️ CodeRabbit configuration file

Files:

comfy_extras/nodes_torch_compile.py

Core ML/diffusion engine. Focus on:

⚙️ CodeRabbit configuration file

Files:

comfy/model_management.py
comfy/cli_args.py
comfy/ldm/modules/attention.py

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

⚙️ CodeRabbit configuration file

Files:

comfy/model_management.py
comfy_extras/nodes_torch_compile.py
comfy/cli_args.py
comfy/ldm/modules/attention.py

Treat legacy combo, `io.Combo`, and `io.DynamicCombo` values affecting filesystem access as untrusted; revalidate them at load/save boundaries with `folder_paths`, containment checks, or fixed allowlists.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

comfy/model_management.py
comfy_extras/nodes_torch_compile.py
comfy/cli_args.py
comfy/ldm/modules/attention.py

Keep state and capability flags on the object that owns the behavior. Prefer explicit parent-owned attributes over probing child objects with `getattr`; use child checks only when the child owns the delegated behavior.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

comfy/model_management.py
comfy_extras/nodes_torch_compile.py
comfy/cli_args.py
comfy/ldm/modules/attention.py

Keep changes small, direct, and limited to the narrowest necessary code path and smallest number of files.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

comfy/model_management.py
comfy_extras/nodes_torch_compile.py
comfy/cli_args.py
comfy/ldm/modules/attention.py

Keep warning and info messages short and actionable, remove noisy or misleading logging, and make documentation edits concise, factual, and tied to changed behavior.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

comfy/model_management.py
comfy_extras/nodes_torch_compile.py
comfy/cli_args.py
comfy/ldm/modules/attention.py

🔇 Additional comments (9)

comfy_extras/nodes_torch_compile.py (4)

1-1: LGTM!

20-28: LGTM!

35-43: LGTM!

19-19: 🩺 Stability & Availability

Declare or guard the minimum PyTorch version.

torch.compiler.list_backends() requires PyTorch 2.1 or later, but the repository does not declare a PyTorch minimum.

comfy/cli_args.py (1)

150-150: LGTM!
comfy/ldm/modules/attention.py (3)
56-66: LGTM!

916-918: LGTM!

Also applies to: 953-954, 964-965

589-593: 🩺 Stability & Availability

Do not flag this fallback.

comfy/ldm/modules/attention.py is unchanged by this PR. The comment targets pre-existing code.
			> Likely an incorrect or invalid review comment.
comfy/model_management.py (1)

1675-1676: LGTM!

📝 Walkthrough

Walkthrough

Adds the --use-flex-attention flag and integrates PyTorch Flex Attention with mask handling, GQA, optional compilation, backend registration, selection, and fallback behavior. Adds flex_attention_enabled() for configuration checks. Updates TorchCompileModel to list PyTorch compiler backends dynamically and accept an optional backend-specific mode during execution.

Merge Risk: 🔵 Low · up to 99e29

The PR adds opt-in attention and compile-mode behavior while preserving existing defaults. Its compile backend discovery requires PyTorch 2.1 or newer without declaring or guarding that requirement, so older installations could fail at runtime; the PR is mergeable with explicit owner awareness.

🚥 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 10 functions across 4 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 summarizes both primary changes: adding Flex Attention and exposing torch.compile modes.
Description check	✅ Passed	The description directly explains the Flex Attention backend, compile modes, benchmarks, and testing included in the changeset.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

comfy-legal added a commit to Comfy-Org/comfy-cla that referenced this pull request Aug 24, 2026

@apersomany has signed the CLA in Comfy-Org/ComfyUI#15870

017b17b

apersomany added 2 commits August 25, 2026 06:18

Add Flex Attention backend

fa24b4b

Expose compile modes in TorchCompileModel

99e2986

apersomany force-pushed the flex-attention-compile-modes branch from bdd2f06 to 99e2986 Compare August 24, 2026 21:18

apersomany marked this pull request as ready for review August 24, 2026 21:19

apersomany requested review from Kosinkadink, alexisrolland, comfyanonymous, guill, kijai and rattus128 as code owners August 24, 2026 21:19

coderabbitai Bot approved these changes Aug 24, 2026

View reviewed changes

Navigation Menu

Uh oh!

Add Flex Attention and expose torch.compile modes - #15870

Add Flex Attention and expose torch.compile modes#15870
apersomany wants to merge 2 commits into
Comfy-Org:masterfrom
apersomany:flex-attention-compile-modes

apersomany commented Aug 24, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 24, 2026 •

edited

Loading

Uh oh!

apersomany commented Aug 24, 2026

Uh oh!

coderabbitai Bot commented Aug 24, 2026

Walkthrough

❌ Failed checks (1 warning)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

apersomany commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Flex Attention

TorchCompileModel

Benchmark evidence

512²

1024²

2048²

Appendix: Reproducing the benchmarks

Testing

Uh oh!

github-actions Bot commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

apersomany commented Aug 24, 2026

Uh oh!

coderabbitai Bot commented Aug 24, 2026

Walkthrough

❌ Failed checks (1 warning)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

apersomany commented Aug 24, 2026 •

edited

Loading

github-actions Bot commented Aug 24, 2026 •

edited

Loading