Addresses CodeRabbit feedback on #14838: the chunking loop always
allocated a temporary output buffer and copied SDPA's result into it,
even in the common case (any non-MPS device, or MPS below the size
threshold) where steps==1 and the loop only ever runs once. Adds a
fast path that calls scaled_dot_product_attention directly and
reshapes its result, matching the original pre-fix code for that case
-- no extra allocation or copy. The chunked path (steps>1) is
unchanged.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW
Keep the fix itself (size-based chunking/capping) but drop the
_diag_log_attn_size helper, the debug-level [MPS-ATTN-FIX:...]
messages, and the now-unused tag parameters that only existed to
support that logging -- a minimal correctness fix without added
observability scaffolding.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW
MPS uses 32-bit indexing internally for many ops; a single attention
matrix (b*heads*seq_q*seq_k) at or above ~2^31 elements silently
corrupts output instead of raising, regardless of free unified memory.
A 61-frame/832x640 Wan2.2 case measured 7.68B elements, 3.6x over the
limit.
Fixes all four MPS-reachable attention implementations in
comfy/ldm/modules/attention.py, not just the ones a given launch
config happens to select:
- model_management.py: add bf16 to FORCE_UPCAST_ATTENTION_DTYPE
(was fp16-only).
- attention_split: extend fp32 upcast to the second (weights x V)
matmul, not just Q.Kt; add size-based diagnostic logging
(_diag_log_attn_size); force chunking via a new shared
_mps_forced_attn_chunk_steps() helper once the attention matrix
exceeds 2^30 elements, independent of the existing memory-pressure
based steps calculation (which never triggers on a 512GB unified
memory machine); fix a slice_size divisibility bug that silently
fell back to computing the whole unchunked sequence.
- attention_pytorch: port the same MPS chunking guard to the native
SDPA path (previously unguarded and unsafe for large sequences on
MPS), chunking over seq_q with a preallocated output buffer; zero
behavior change when chunking doesn't engage. Also add
attn_precision handling (get_attn_precision + float32 upcast of
q/k/v/mask), which this path previously ignored entirely.
- attention_basic: reached via optimized_attention_for_device's
small_input=True path whenever pytorch attention isn't enabled --
a live path on MPS for text/image encoder attention (CLIP, T5,
Llama, Gemma, Qwen-VL, DINO, BiRefNet, RT-DETR, etc). Had zero
chunking of any kind; introduces a chunking loop using the same
shared helper, preserving both bool-mask (masked_fill_) and
float-mask (additive) handling.
- attention_sub_quad: the actual default MPS attention path (selected
whenever no --use-*-attention flag is passed at all). Its
free-memory-based query_chunk_size/kv_chunk_size selection has the
same class of gap attention_split had -- on a 512GB unified-memory
machine it always picks the largest candidate and disables
kv-chunking entirely. Adds a new _mps_cap_subquad_chunk_sizes()
helper that caps both values against the same MPS ceiling, entirely
in the caller so the third-party MIT-licensed
sub_quadratic_attention.py delegate stays untouched.
- Quiet the MPS-ATTN-FIX log to debug level (fires on every call once
a guard engages, not deduplicated like the diagnostic logging).
Verified: standalone numerical tests confirm chunked/capped output
matches unchunked baseline (fp16-rounding-level max abs diff) for all
four functions, across no-mask/float-mask/bool-mask and fp32-upcast
variants, including the attention_sub_quad boundary case that forces
the delegate's real kv-chunked branch instead of its fast path. Shape
parity confirmed against the real 61-frame/832x640 repro case: all
four functions independently compute consistent forced chunk sizes
for the identical b*heads=40, seq=13860 shape.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW
* initial gemma4 support
* parity with reference implementation
outputs can 100% match transformers with same sdpa flags, checkpoint this and then optimize
* Cleanup, video fixes
* cleanup, enable fused rms norm by default
* update comment
* Cleanup
* Update sd.py
* Various fixes
* Add fp8 scaled embedding support
* small fixes
* Translate think tokens
* Fix image encoder attention mask type
So it works with basic attention
* Handle thinking tokens different only for Gemma4
* Code cleanup
* Update nodes_textgen.py
* Use embed scale class instead of buffer
Slight difference to HF, but technically more accurate and simpler code
* Default to fused rms_norm
* Update gemma4.py
Pytorch only filters for OOMs in its own allocators however there are
paths that can OOM on allocators made outside the pytorch allocators.
These manifest as an AllocatorError as pytorch does not have universal
error translation to its OOM type on exception. Handle it. A log I have
for this also shows a double report of the error async, so call the
async discarder to cleanup and make these OOMs look like OOMs.
* causal_video_ae: Remove attention ResNet
This attention_head_dim argument does not exist on this constructor so
this is dead code. Remove as generic attention mid VAE conflicts with
temporal roll.
* ltx-vae: consoldate causal/non-causal code paths
* ltx-vae: add cache rolling adder
* ltx-vae: use cached adder for resnet
* ltx-vae: Implement rolling VAE
Implement a temporal rolling VAE for the LTX2 VAE.
Usually when doing temporal rolling VAEs you can just chunk on time relying
on causality and cache behind you as you go. The LTX VAE is however
non-causal.
So go whole hog and implement per layer run ahead and backpressure between
the decoder layers using recursive state beween the layers.
Operations are ammended with temporal_cache_state{} which they can use to
hold any state then need for partial execution. Convolutions cache their
inputs behind the up to N-1 frames, and skip connections need to cache the
mismatch between convolution input and output that happens due to missing
future (non-causal) input.
Each call to run_up() processes a layer accross a range on input that
may or may not be complete. It goes depth first to process as much as
possible to try and digest frames to the final output ASAP. If layers run
out of input due to convolution losses, they simply return without action
effectively applying back-pressure to the earlier layers. As the earlier
layers do more work and caller deeper, the partial states are reconciled
and output continues to digest depth first as much as possible.
Chunking is done using a size quota rather than a fixed frame length and
any layer can initiate chunking, and multiple layers can chunk at different
granulatiries. This remove the old limitation of always having to process
1 latent frame to entirety and having to hold 8 full decoded frames as
the VRAM peak.
* Add support for sage attention 3 in comfyui, enable via new cli arg
--use-sage-attiention3
* Fix some bugs found in PR review. The N dimension at which Sage
Attention 3 takes effect is reduced to 1024 (although the improvement is
not significant at this scale).
* Remove the Sage Attention3 switch, but retain the attention function
registration.
* Fix a ruff check issue in attention.py
* hunyuan upsampler: rework imports
Remove the transitive import of VideoConv3d and Resnet and takes these
from actual implementation source.
* model: remove unused give_pre_end
According to git grep, this is not used now, and was not used in the
initial commit that introduced it (see below).
This semantic is difficult to implement temporal roll VAE for (and would
defeat the purpose). Rather than implement the complex if, just delete
the unused feature.
(venv) rattus@rattus-box2:~/ComfyUI$ git log --oneline
220afe33 (HEAD) Initial commit.
(venv) rattus@rattus-box2:~/ComfyUI$ git grep give_pre
comfy/ldm/modules/diffusionmodules/model.py: resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
comfy/ldm/modules/diffusionmodules/model.py: self.give_pre_end = give_pre_end
comfy/ldm/modules/diffusionmodules/model.py: if self.give_pre_end:
(venv) rattus@rattus-box2:~/ComfyUI$ git co origin/master
Previous HEAD position was 220afe33 Initial commit.
HEAD is now at 9d8a8179 Enable async offloading by default on Nvidia. (#10953)
(venv) rattus@rattus-box2:~/ComfyUI$ git grep give_pre
comfy/ldm/modules/diffusionmodules/model.py: resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
comfy/ldm/modules/diffusionmodules/model.py: self.give_pre_end = give_pre_end
comfy/ldm/modules/diffusionmodules/model.py: if self.give_pre_end:
* move refiner VAE temporal roller to core
Move the carrying conv op to the common VAE code and give it a better
name. Roll the carry implementation logic for Resnet into the base
class and scrap the Hunyuan specific subclass.
* model: Add temporal roll to main VAE decoder
If there are no attention layers, its a standard resnet and VideoConv3d
is asked for, substitute in the temporal rolloing VAE algorithm. This
reduces VAE usage by the temporal dimension (can be huge VRAM savings).
* model: Add temporal roll to main VAE encoder
If there are no attention layers, its a standard resnet and VideoConv3d
is asked for, substitute in the temporal rolling VAE algorithm. This
reduces VAE usage by the temporal dimension (can be huge VRAM savings).
* Looking into a @wrap_attn decorator to look for 'optimized_attention_override' entry in transformer_options
* Created logging code for this branch so that it can be used to track down all the code paths where transformer_options would need to be added
* Fix memory usage issue with inspect
* Made WAN attention receive transformer_options, test node added to wan to test out attention override later
* Added **kwargs to all attention functions so transformer_options could potentially be passed through
* Make sure wrap_attn doesn't make itself recurse infinitely, attempt to load SageAttention and FlashAttention if not enabled so that they can be marked as available or not, create registry for available attention
* Turn off attention logging for now, make AttentionOverrideTestNode have a dropdown with available attention (this is a test node only)
* Make flux work with optimized_attention_override
* Add logs to verify optimized_attention_override is passed all the way into attention function
* Make Qwen work with optimized_attention_override
* Made hidream work with optimized_attention_override
* Made wan patches_replace work with optimized_attention_override
* Made SD3 work with optimized_attention_override
* Made HunyuanVideo work with optimized_attention_override
* Made Mochi work with optimized_attention_override
* Made LTX work with optimized_attention_override
* Made StableAudio work with optimized_attention_override
* Made optimized_attention_override work with ACE Step
* Made Hunyuan3D work with optimized_attention_override
* Make CosmosPredict2 work with optimized_attention_override
* Made CosmosVideo work with optimized_attention_override
* Made Omnigen 2 work with optimized_attention_override
* Made StableCascade work with optimized_attention_override
* Made AuraFlow work with optimized_attention_override
* Made Lumina work with optimized_attention_override
* Made Chroma work with optimized_attention_override
* Made SVD work with optimized_attention_override
* Fix WanI2VCrossAttention so that it expects to receive transformer_options
* Fixed Wan2.1 Fun Camera transformer_options passthrough
* Fixed WAN 2.1 VACE transformer_options passthrough
* Add optimized to get_attention_function
* Disable attention logs for now
* Remove attention logging code
* Remove _register_core_attention_functions, as we wouldn't want someone to call that, just in case
* Satisfy ruff
* Remove AttentionOverrideTest node, that's something to cook up for later
Scott
comfyanonymous
John Pollock
Scott Yewell
Jukka Seppänen
omahs
rattus
mengqin
blepping
Jedrzej Kosinski
contentis
chaObserv
Kohaku-Blueleaf
Raphael Walker
comfyanonymous
FeepingCreature
Dr.Lt.Data