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.
* 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
* 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
* fix attention OOM in xformers
* allow passing attention mask in flux attention
* allow an attn_mask in flux
* attn masks can be done using replace patches instead of a separate dict
* fix return types
* fix return order
* enumerate
* patch the right keys
* arg names
* fix a silly bug
* fix xformers masks
* replace match with if, elif, else
* mask with image_ref_size
* remove unused import
* remove unused import 2
* fix pytorch/xformers attention
This corrects a weird inconsistency with skip_reshape.
It also allows masks of various shapes to be passed, which will be
automtically expanded (in a memory-efficient way) to a size that is
compatible with xformers or pytorch sdpa respectively.
* fix mask shapes
Scott
comfyanonymous
Scott Yewell
John Pollock
Jukka Seppänen
rattus
mengqin
blepping
Jedrzej Kosinski
contentis
Kohaku-Blueleaf
Raphael Walker
comfyanonymous
FeepingCreature
Dr.Lt.Data
Chenlei Hu