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 HiDream01-image support
* Cleanup nodes
* Cleaner handling of empty placeholder models
* Remove snap_to_predefined, prefer tooltip for the trained resolutions
* Add model and block wrappers
* Fix shift tooltip
* Add node to work around the patch tile issue
Experimental, runs multiple passes with the patch grid offset and blends with various different methods.
* Qwen35 vision rotary_pos_emb cast fix
* Fix embedding layout type
* Some small optimizations
* Cleanup, don't need this fallback
* Prefix KV cache, cleanup
Bit of speed, reduce redundant code
* Get rid of redundant custom sampler, refactor noise scaling
Our existing lcm sampler is mathematically same, just added the missing options to it instead and a node to control them. Refactored the noise scaling and fix it for the stochastic samplers, add a generic node to control the initial noise scale.
* Update nodes_hidream_o1.py
* Fix some cache validation cases
* Keep existing sampling params
* Remove redundant video vision path
* Replace some numpy ops with torch
* Fx RoPE index for batch size > 1
* Prefer torch preprocessing
* Rename block_type to be compatible with existing patch nodes
* Fixes and tweaks
* initial WanDancer support
* nodes_wandancer: Add list form of chunker.
Create an alternate list form of the node so the chunk gens can be
trivially looped by the comfy executor.
* Closer match to original soxr resampling
* Remove librosa node
* Cleanup
---------
Co-authored-by: Rattus <rattus128@gmail.com>
* 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
* mm: Use Aimdo raw allocator for cast buffers
pytorch manages allocation of growing buffers on streams poorly. Pyt
has no windows support for the expandable segments allocator (which is
the right tool for this job), while also segmenting the memory by
stream such that it can be generally re-used. So kick the problem to
aimdo which can just grow a virtual region thats freed per stream.
* plan
* ops: move cpu handler up to the caller
* ops: split up prefetch from weight prep block prefetching API
Split up the casting and weight formating/lora stuff in prep for
arbitrary prefetch support.
* ops: implement block prefetching API
allow a model to construct a prefetch list and operate it for increased
async offload.
* ltxv2: Implement block prefetching
* Implement lora async offload
Implement async offload of loras.
Scott Yewell
Scott Yewell
comfyanonymous
Jukka Seppänen
Barish Ozbay
Talmaj
Alexis Rolland
John Pollock
person4268
Ivan Zorin
Jedrzej Kosinski
Cezarijus Kivylius
omahs
rattus