Patch Sage Attention DM

SageAttention is the community's favorite free attention speedup - it trades the standard attention path for a memory-efficient, roughly 8-bit one that can cut several GB of peak VRAM on long video gens and noticeably speed up sampling. This node is that trick packaged for ComfyUI, with a twist you need to know before you add it: it does not use ComfyUI's model patching system. It's a global, runtime patch. Disabling it means running the node again with sage_attention set to disabled - you can't just bypass the node, because the patch is applied on each execution and reverted on cleanup around whatever model you wired it to.

The author is upfront about this in the node's own description: "experimental." Read the rest of this page before you build a workflow around it.

How it works

The node takes a MODEL, clones it, and hooks two callbacks: ON_PRE_RUN swaps ComfyUI's attention override to the SageAttention function you picked, and ON_CLEANUP restores the original (falling back to Flash-Attention if you have it). That means the effect is scoped to the execution of that model, but while it runs, every attention call in the graph that shares the override goes through Sage. It's a different mechanism from the per-workflow SageAttention nodes in packs like KJNodes, which is why it behaves differently.

The inputs

• model - the MODEL to run under the patch.
• sage_attention - the picker. disabled reverts everything. auto picks a sensible implementation. Then the explicit ones: sageattn_qk_int8_pv_fp16_cuda and a Triton variant, int8/fp8 CUDA pairs, sageattn3, and sageattn3_per_block_mean for Blackwell hardware. The cuda++ variant is the optimized fp8 path.
• allow_compile (optional) - lets it torch.compile the SageAttention function. Only meaningful with sageattn 2.2.0+, and compile adds startup overhead.

The two traps

1. It isn't installed for you. This pack's requirements.txt is just torch, psutil, nvidia-ml-py - no sageattention, no Triton. You have to install those yourself, and the KB's Wan writeup is blunt about Windows: Triton requires a manual install and one community member effectively single-handedly maintains Windows builds. On Linux it's usually pip install sageattention and you're off.

2. It can genuinely break models. The known bad case is Z-Image Base, where SageAttn produces severe artifacts - patchy/blurry output, matrix-code lines, sometimes fully black frames, confirmed across 3090/4090/5090 (Z-Image Turbo was reported fine). Don't assume Sage is safe for whatever you're sampling; if output goes weird, the first thing to try is re-running with disabled.

There's a nice touch in the pack: it installs a startup guard that quietly swallows the repeated "unsupported head_dim=160" SageAttention errors some models hit, so you don't drown in console spam.

Installing it

cd ComfyUI/custom_nodes
git clone https://github.com/ussoewwin/ComfyUI-DistorchMemoryManager.git
cd ComfyUI-DistorchMemoryManager
pip install -r requirements.txt
# SageAttention is NOT in requirements - install it yourself:
pip install sageattention

Or use ComfyUI Manager, searching "Distorch" or "VRAM Manager". Restart ComfyUI; the node appears under Memory. Start with auto, on a cheap quick test render, and check the console logs - the node helpfully prints which SageAttention version and implementation it picked up, which tells you fast whether the install actually took.