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
* security: fix five vulnerabilities (GHSA-779p-m5rp-r4h4)
- CVE-2026-56670: force download of SVG/XML responses on /view to prevent stored XSS
- CVE-2026-56671: contain /experiment/models/preview reads within the model folder
- CVE-2026-56672: stop inline rendering of uploaded /userdata/{file} content
- CVE-2026-56673: prevent path traversal in get_annotated_filepath (LoadImage /prompt input)
- CVE-2026-56674: reject opaque/null Origin to close the CSRF middleware bypass
Adds regression tests under tests-unit/security_test/ covering all five.
* security: address review feedback on GHSA-779p fixes
- Fix Windows CI failure in test_get_annotated_filepath: compare against
os.path.abspath(...) to match the intentional abspath normalization added
by the traversal hardening (abspath prepends the drive letter on Windows).
- origin_check: narrow the bare `except:` in is_loopback() to ValueError so
genuine interrupts aren't swallowed (review nit).
- origin_check: guard .port access in is_cross_origin_forbidden() so a
malformed/out-of-range port (e.g. Origin: http://127.0.0.1:99999) fails
closed with a 403 instead of surfacing an uncaught 500 in the middleware.
- server /view: escape backslash/quote in the Content-Disposition filename
(RFC 6266 quoted-string) so a filename containing a double quote can't
malform the response header.
* security: address CodeRabbit review feedback on GHSA-779p tests
- test #3: guard the symlink-escape test with a try/except skip so it no
longer errors on Windows CI where os.symlink needs elevated privileges /
Developer Mode (mirrors the guard in the sibling test #2).
- test #5: refresh the stale module docstring to describe the actual /view
gating (view_image closure calling folder_paths.is_dangerous_content_type,
the normalising check) instead of the bypassable raw set-membership test.
* revert(security): drop CVE-2026-56674 Origin: null CSRF change
Per maintainer review, the reported CSRF is already mitigated by the pre-existing
Sec-Fetch-Site: cross-site check for current browsers, and the null-origin
rejection risked breaking legitimate sandboxed-iframe embeds. Restores
origin_only_middleware and is_loopback in server.py to their prior state
(the Sec-Fetch-Site check is retained) and removes utils/origin_check.py and its
regression test. The other four GHSA-779p fixes are unaffected.
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>
Scott Yewell
Scott Yewell
Silver
comfyanonymous
Barish Ozbay
Alexander Piskun
Alexis Rolland
Daxiong (Lin)
Robin Huang
Matt Miller
comfyanonymous
Comfy Org PR Bot
pythongosssss