Flex Attention (torch.compile)

Flex Attention (torch.compile) is the rare ComfyUI node that asks for nothing and sometimes gives you something. One input, one output, zero dependencies, and it swaps your model's attention path for a fused, torch-compiled PyTorch kernel that runs on any GPU torch.compile supports. The whole pitch is in the README's one comparison: Flash Attention 4 needs Blackwell hardware, Flex Attention doesn't. If you're on a 30-series or 40-series card, this is the "new attention kernel" trick you can actually run.

Where it goes and what it does

It's a patch node, not a generator. Drop it between your model/LoRA loader and the KSampler, wire model in and the output MODEL into the sampler's model input, and every attention call in that model gets rerouted through PyTorch's flex_attention - compiled via torch.compile on the first generation. It's model-agnostic: feed it SD1.5, SDXL, or Flux and it patches whatever comes in.

Under the hood it does the same thing every optimized-attention override does in ComfyUI: it clones your model, then sets transformer_options["optimized_attention_override"] so Comfy's attention wrapper calls the new function. The function reshapes q/k/v into (B, H, N, D) head-major layout, calls the compiled kernel, and reshapes back. Two details worth knowing from the source:

• First call compiles. Expect a one-time warmup delay of a few seconds to a couple of minutes depending on your machine. It's not a hang, and the compiled kernel is cached to disk, so it's roughly a once-per-machine cost.
• Masks silently fall back. If attention is called with a mask (some regional-prompting, AnimateDiff, and inpainting setups do this), the node hands off to plain PyTorch SDPA instead. That's the correct behavior - flex_attention wants masks expressed as block masks, and this node doesn't translate them - but it means you can't benchmark it with a masked workflow and conclude anything.

Install

Painless, because the author shipped it with literally no Python dependencies (dependencies = [] in pyproject, no requirements.txt - a rarity in this ecosystem, where half your installs are dependency roulette). Two routes:

• ComfyUI Manager - search "Flex Attention" and install.
• Manual - cd ComfyUI/custom_nodes && git clone https://github.com/sorryhyun/comfyui-flex-attention, then restart ComfyUI.

The real requirement is PyTorch 2.5+ - flex_attention only exists there - plus a CUDA GPU that torch.compile will work with. No model files to download, no config to set.

Common issues

• Import error on old PyTorch. The node imports torch.nn.attention.flex_attention at runtime. On PyTorch older than 2.5 that import fails outright. Update PyTorch (or the whole ComfyUI portable install) and it goes away.
• Errors mentioning torch/_inductor/kernel/flex_attention.py or TritonTemplate. This is torch's own Triton stack tripping over an incompatible torch/triton pairing, not this node's code - the traceback doesn't even reach your custom_nodes folder. Fix is to update or re-align PyTorch and Triton. It's a known rough edge, most often reported on Windows.
• "Why isn't it any faster?" Be honest about expectations. Attention is rarely the whole story - token encoding, VAE, and sampler step count dominate wall time, and torch.compile gives the biggest wins on heavier transformer models like Flux and Wan, not on small SD1.5 UNets. It's a free-ish experiment worth A/B-testing once (the node clones the model, so you can bypass it to compare), but if your goal is a big speedup, the popular lever is still SageAttention-style kernels or full-model compile patches like KJNodes' TorchCompile nodes.

The version here is minimal on purpose - a single node, no knobs. That's its charm and its limit. For a five-minute install that might shave a chunk off your Flux iterations without touching anything else, it's a fine thing to wire in.