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
Add a --enable-asset-hashing CLI flag (action=store_true, default False)
and plumb it into the two asset-seeder call sites in main.py that
previously hardcoded compute_hashes=True (the startup scan and the
post-job output enqueue). Local runs now skip blake3 hashing unless the
user opts in, avoiding the startup/per-output cost on large models
directories while keeping hashing available for asset-portability
features.
Co-authored-by: Alexis Rolland <alexisrolland@hotmail.com>
* main: implement --vram-headroom
Implement --vram-headroom for dynamic vram as a hybrid debug/diagnostic
option that can be used for people who still report shared VRAM spills.
They can trial and error the setting to maintain a bit more headroom to
avoid shared VRAM spills.
* main: implement --reserve-vram
Implement --reserve-vram as extra headroom on the simple method which
is semantically as close as possible to the stated functionality and
formet behaviour of non-dynamic VRAM.
Add this option for users who know they have so much ram they want
to pin everything or have a pagefile that outruns their disk speed.
The removes the RAM pressure caps completely and pins behind the
primary model load forcing all models to be permanently comitted
to RAM.
Some custom nodes .to weights completely out of load context which
can wreak havoc if its for a model that is not active. Detect this
condition and just let it fall-through to the non-dynamic loader
straight up.
Some custom nodes try to set this true globally. It messes with dynamic
VRAM with one-off spikes that can OOM but this is also very high risk
for windows where such allocations might get serviced by shared memory
fallback.
Trump it.
cleanup_models_gc can be called once per load_models_gpu via
free_memory, which in turn can de-activate an active model via
this reset_cast_buffers.
cleanup_models_gc() could also come via obscure garbage collector
paths so limit reset_cast_buffers to the post-node callsite instead.
* mm: split off registration helper to doer and headroom calc
* pinned_memory: implement registration comfy side
Move away from Aimdo buffer registrations which seem fraught with
danger and do it comfy side. Just start with the basic move.
* pinned_memory: do registrations as portable memory
* pinned_memory: discard async errors on registration fail
Like the good ol days.
* pinned_memory: implement abs shortfall retry
If pinned registration happens to fail despite the previous budget
ensures, consider the allocation shortfall, ensure it again, and
try again. This allows comfy pins to interoperate with other software
that might be doing substantive pinning.
* fix (MultiGPU): prevent freeze on manual abort when using MultiGPU CFG Split
Problem:
Upon manual abort application hangs indefinitely.
`InterruptProcessingException` inherits from `BaseException` and bypasses MultiGPU's worker error handling block so thread dies silently, leaving the main thread waiting forever for `result_q.get()`
Fix:
Catch `comfy.model_management.InterruptProcessingException` instead of `Exception` so it's caught and passed back via `result_q` to unblock the main thread when manual abort signal fires.
* oops
Scott Yewell
Scott Yewell
Silver
comfyanonymous
Matt Miller
Jukka Seppänen
Barish Ozbay
rattus
Jedrzej Kosinski
Kohaku-Blueleaf
Talmaj
kelseyee
Alexis Rolland
John Pollock
Quasar of Mikus
person4268