From 7be129913f5061427b95433080f62b8bfe315805 Mon Sep 17 00:00:00 2001 From: Scott Yewell Date: Wed, 8 Jul 2026 15:02:40 -0400 Subject: [PATCH 1/3] Fix MPS attention corruption for large sequences (Wan2.2 long-video) 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 Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW --- comfy/ldm/modules/attention.py | 171 +++++++++++++++++++++++++++++---- comfy/model_management.py | 2 +- 2 files changed, 152 insertions(+), 21 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 2411aff5c..d3528b2ec 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -81,6 +81,65 @@ def get_attn_precision(attn_precision, current_dtype): return FORCE_UPCAST_ATTENTION_DTYPE[current_dtype] return attn_precision +_DIAG_MAX_ATTN_ELEMENTS = {"split": 0, "pytorch": 0, "basic": 0, "sub_quad": 0} + +def _diag_log_attn_size(tag, b_heads, seq_q, seq_k, extra=""): + total = b_heads * seq_q * seq_k + if total > _DIAG_MAX_ATTN_ELEMENTS[tag]: + _DIAG_MAX_ATTN_ELEMENTS[tag] = total + logging.warning( + f"[ATTN-DIAG:{tag}] new max attn matrix size: b*heads={b_heads} seq_q={seq_q} seq_k={seq_k} " + f"total_elements={total} (2^31={2**31}) ratio={total / (2**31):.3f} {extra}" + ) + +# 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 instead +# of raising, regardless of how much unified memory is free. Shared by every +# attention implementation that needs to force chunking to stay under that limit. +MPS_MAX_ATTN_ELEMENTS = 2 ** 30 + +def _mps_forced_attn_chunk_steps(elements_full, steps, tag): + if elements_full <= MPS_MAX_ATTN_ELEMENTS: + return steps + mps_steps = 2 ** math.ceil(math.log(elements_full / MPS_MAX_ATTN_ELEMENTS, 2)) + if mps_steps > steps: + # Quieted to debug for now -- this fires on every call once the guard engages + # (not deduplicated like _diag_log_attn_size), which floods the log during a + # long render. Bump back to logging.warning if you need it visible again. + logging.debug( + f"[MPS-ATTN-FIX:{tag}] forcing attention chunking steps={mps_steps} (was {steps}) to stay " + f"under MPS's ~2^31-element indexing limit (full attn matrix would be {elements_full} elements)" + ) + return mps_steps + return steps + +def _mps_cap_subquad_chunk_sizes(batch_x_heads, q_tokens, k_tokens, query_chunk_size, kv_chunk_size, tag): + effective_kv = kv_chunk_size if kv_chunk_size is not None else max(1, int(math.sqrt(k_tokens))) + elements_full = batch_x_heads * min(query_chunk_size, q_tokens) * effective_kv + if elements_full <= MPS_MAX_ATTN_ELEMENTS: + return query_chunk_size, kv_chunk_size + + # Cap kv_chunk_size too (only if the caller had set one explicitly) in case + # batch_x_heads*kv_chunk_size alone already exceeds the ceiling -- unrealistic + # for current model shapes, but keeps the "regardless of free memory" guarantee + # symmetric with the other MPS attention fixes rather than relying on query + # capping alone. + max_kv_for_ceiling = max(1, MPS_MAX_ATTN_ELEMENTS // batch_x_heads) + new_kv = min(effective_kv, max_kv_for_ceiling) + if kv_chunk_size is not None and new_kv < kv_chunk_size: + kv_chunk_size = new_kv + effective_kv = new_kv + + new_query_chunk_size = max(1, MPS_MAX_ATTN_ELEMENTS // (batch_x_heads * effective_kv)) + if new_query_chunk_size < query_chunk_size: + logging.debug( + f"[MPS-ATTN-FIX:{tag}] capping query_chunk_size={new_query_chunk_size} (was {query_chunk_size}) " + f"kv_chunk_size={kv_chunk_size} to stay under MPS's ~2^31-element indexing limit " + f"(full attn matrix would be {elements_full} elements)" + ) + query_chunk_size = new_query_chunk_size + return query_chunk_size, kv_chunk_size + def exists(val): return val is not None @@ -205,32 +264,51 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v)) - # force cast to fp32 to avoid overflowing - if attn_precision == torch.float32: - sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale - else: - sim = einsum('b i d, b j d -> b i j', q, k) * scale - - del q, k + steps = 1 + if q.device.type == "mps": + elements_full = q.shape[0] * q.shape[1] * k.shape[1] + steps = _mps_forced_attn_chunk_steps(elements_full, steps, tag="basic") + slice_size = math.ceil(q.shape[1] / steps) + is_bool_mask = False if exists(mask): if mask.dtype == torch.bool: + is_bool_mask = True mask = rearrange(mask, 'b ... -> b (...)') #TODO: check if this bool part matches pytorch attention - max_neg_value = -torch.finfo(sim.dtype).max + max_neg_value = -torch.finfo(torch.float32 if attn_precision == torch.float32 else q.dtype).max mask = repeat(mask, 'b j -> (b h) () j', h=h) - sim.masked_fill_(~mask, max_neg_value) else: if len(mask.shape) == 2: bs = 1 else: bs = mask.shape[0] mask = mask.reshape(bs, -1, mask.shape[-2], mask.shape[-1]).expand(b, heads, -1, -1).reshape(-1, mask.shape[-2], mask.shape[-1]) - sim.add_(mask) - # attention, what we cannot get enough of - sim = sim.softmax(dim=-1) + out = torch.empty((q.shape[0], q.shape[1], v.shape[2]), dtype=v.dtype, device=q.device) + for i in range(0, q.shape[1], slice_size): + end = min(i + slice_size, q.shape[1]) + # force cast to fp32 to avoid overflowing + if attn_precision == torch.float32: + sim = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * scale + else: + sim = einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale - out = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v) + if exists(mask): + if is_bool_mask: + sim.masked_fill_(~mask, max_neg_value) + else: + if mask.shape[1] == 1: + sim += mask + else: + sim += mask[:, i:end] + + # attention, what we cannot get enough of + sim = sim.softmax(dim=-1) + + out[:, i:end] = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v) + _diag_log_attn_size("basic", q.shape[0], end - i, k.shape[1], extra=f"steps={steps} b={b} heads={heads} full_seq_q={q.shape[1]}") + + del q, k if skip_output_reshape: out = ( @@ -298,6 +376,15 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None, if query_chunk_size is None: query_chunk_size = 512 + if query.device.type == "mps": + query_chunk_size, kv_chunk_size = _mps_cap_subquad_chunk_sizes( + batch_x_heads, q_tokens, k_tokens, query_chunk_size, kv_chunk_size, tag="sub_quad") + _diag_log_attn_size( + "sub_quad", batch_x_heads, min(query_chunk_size, q_tokens), + kv_chunk_size if kv_chunk_size is not None else int(math.sqrt(k_tokens)), + extra=f"q_tokens={q_tokens} k_tokens={k_tokens}" + ) + if mask is not None: if len(mask.shape) == 2: bs = 1 @@ -370,6 +457,12 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB " # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}") + # See _mps_forced_attn_chunk_steps for why this is needed on MPS, independent + # of the memory-based steps calculation above. + if q.device.type == "mps": + elements_full = q.shape[0] * q.shape[1] * k.shape[1] + steps = _mps_forced_attn_chunk_steps(elements_full, steps, tag="split") + if steps > 64: max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). ' @@ -387,7 +480,8 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape cleared_cache = False while True: try: - slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1] + slice_size = min(q.shape[1], math.ceil(q.shape[1] / steps)) + _diag_log_attn_size("split", q.shape[0], slice_size, k.shape[1], extra=f"steps={steps} b={b} heads={heads} full_seq_q={q.shape[1]}") for i in range(0, q.shape[1], slice_size): end = i + slice_size if upcast: @@ -405,11 +499,16 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape else: s1 += mask[:, i:end] - s2 = s1.softmax(dim=-1).to(v.dtype) + s2 = s1.softmax(dim=-1) + if not upcast: + s2 = s2.to(v.dtype) del s1 first_op_done = True - r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v) + if upcast: + r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v.float()).to(r1.dtype) + else: + r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v) del s2 break except Exception as e: @@ -521,6 +620,8 @@ else: @wrap_attn def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + attn_precision = get_attn_precision(attn_precision, q.dtype) + if skip_reshape: b, _, _, dim_head = q.shape else: @@ -537,15 +638,42 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha if mask.ndim == 3: mask = mask.unsqueeze(1) + do_upcast = attn_precision == torch.float32 + if do_upcast: + orig_dtype = q.dtype + q, k, v = q.float(), k.float(), v.float() + if mask is not None and mask.dtype != torch.bool: + mask = mask.float() + sdpa_keys = ("scale", "enable_gqa") sdpa_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys} if SDP_BATCH_LIMIT >= b: - out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra) - if not skip_output_reshape: - out = ( - out.transpose(1, 2).reshape(b, -1, heads * dim_head) + # See _mps_forced_attn_chunk_steps: on MPS a single attention matrix over + # ~2^31 elements silently corrupts, so chunk over seq_q to stay under that. + # steps stays 1 (one iteration over the full sequence, identical to the + # previous unchunked call) everywhere this doesn't apply. + steps = 1 + if q.device.type == "mps": + elements_full = q.shape[0] * q.shape[1] * q.shape[2] * k.shape[2] + steps = _mps_forced_attn_chunk_steps(elements_full, steps, tag="pytorch") + slice_size = math.ceil(q.shape[2] / steps) + + raw_out = torch.empty((b, q.shape[1], q.shape[2], dim_head), dtype=q.dtype, layout=q.layout, device=q.device) + for i in range(0, q.shape[2], slice_size): + end = i + slice_size + mask_chunk = mask + if mask is not None and mask.shape[-2] != 1: + mask_chunk = mask[..., i:end, :] + _diag_log_attn_size("pytorch", q.shape[0] * q.shape[1], min(end, q.shape[2]) - i, k.shape[2], extra=f"steps={steps} b={b} heads={q.shape[1]} full_seq_q={q.shape[2]} dtype={q.dtype}") + raw_out[:, :, i:end] = comfy.ops.scaled_dot_product_attention( + q[:, :, i:end], k, v, attn_mask=mask_chunk, dropout_p=0.0, is_causal=False, **sdpa_extra ) + + if skip_output_reshape: + out = raw_out + else: + out = raw_out.transpose(1, 2).reshape(b, -1, heads * dim_head) else: out = torch.empty((b, q.shape[2], heads * dim_head), dtype=q.dtype, layout=q.layout, device=q.device) for i in range(0, b, SDP_BATCH_LIMIT): @@ -561,6 +689,9 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha attn_mask=m, dropout_p=0.0, is_causal=False, **sdpa_extra ).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head) + + if do_upcast: + out = out.to(orig_dtype) return out @wrap_attn diff --git a/comfy/model_management.py b/comfy/model_management.py index b15d08ba1..715893a7b 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1646,7 +1646,7 @@ def force_upcast_attention_dtype(): upcast = True if upcast: - return {torch.float16: torch.float32} + return {torch.float16: torch.float32, torch.bfloat16: torch.float32} else: return None From 7b7d48faa2989bb2b2f6929fde7fc3b10f1e4b48 Mon Sep 17 00:00:00 2001 From: Scott Yewell Date: Wed, 8 Jul 2026 15:44:09 -0400 Subject: [PATCH 2/3] Remove diagnostic logging scaffolding from the MPS attention fix 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 Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW --- comfy/ldm/modules/attention.py | 50 ++++++---------------------------- 1 file changed, 8 insertions(+), 42 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index d3528b2ec..18efdf148 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -81,39 +81,19 @@ def get_attn_precision(attn_precision, current_dtype): return FORCE_UPCAST_ATTENTION_DTYPE[current_dtype] return attn_precision -_DIAG_MAX_ATTN_ELEMENTS = {"split": 0, "pytorch": 0, "basic": 0, "sub_quad": 0} - -def _diag_log_attn_size(tag, b_heads, seq_q, seq_k, extra=""): - total = b_heads * seq_q * seq_k - if total > _DIAG_MAX_ATTN_ELEMENTS[tag]: - _DIAG_MAX_ATTN_ELEMENTS[tag] = total - logging.warning( - f"[ATTN-DIAG:{tag}] new max attn matrix size: b*heads={b_heads} seq_q={seq_q} seq_k={seq_k} " - f"total_elements={total} (2^31={2**31}) ratio={total / (2**31):.3f} {extra}" - ) - # 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 instead # of raising, regardless of how much unified memory is free. Shared by every # attention implementation that needs to force chunking to stay under that limit. MPS_MAX_ATTN_ELEMENTS = 2 ** 30 -def _mps_forced_attn_chunk_steps(elements_full, steps, tag): +def _mps_forced_attn_chunk_steps(elements_full, steps): if elements_full <= MPS_MAX_ATTN_ELEMENTS: return steps mps_steps = 2 ** math.ceil(math.log(elements_full / MPS_MAX_ATTN_ELEMENTS, 2)) - if mps_steps > steps: - # Quieted to debug for now -- this fires on every call once the guard engages - # (not deduplicated like _diag_log_attn_size), which floods the log during a - # long render. Bump back to logging.warning if you need it visible again. - logging.debug( - f"[MPS-ATTN-FIX:{tag}] forcing attention chunking steps={mps_steps} (was {steps}) to stay " - f"under MPS's ~2^31-element indexing limit (full attn matrix would be {elements_full} elements)" - ) - return mps_steps - return steps + return max(mps_steps, steps) -def _mps_cap_subquad_chunk_sizes(batch_x_heads, q_tokens, k_tokens, query_chunk_size, kv_chunk_size, tag): +def _mps_cap_subquad_chunk_sizes(batch_x_heads, q_tokens, k_tokens, query_chunk_size, kv_chunk_size): effective_kv = kv_chunk_size if kv_chunk_size is not None else max(1, int(math.sqrt(k_tokens))) elements_full = batch_x_heads * min(query_chunk_size, q_tokens) * effective_kv if elements_full <= MPS_MAX_ATTN_ELEMENTS: @@ -131,13 +111,7 @@ def _mps_cap_subquad_chunk_sizes(batch_x_heads, q_tokens, k_tokens, query_chunk_ effective_kv = new_kv new_query_chunk_size = max(1, MPS_MAX_ATTN_ELEMENTS // (batch_x_heads * effective_kv)) - if new_query_chunk_size < query_chunk_size: - logging.debug( - f"[MPS-ATTN-FIX:{tag}] capping query_chunk_size={new_query_chunk_size} (was {query_chunk_size}) " - f"kv_chunk_size={kv_chunk_size} to stay under MPS's ~2^31-element indexing limit " - f"(full attn matrix would be {elements_full} elements)" - ) - query_chunk_size = new_query_chunk_size + query_chunk_size = min(query_chunk_size, new_query_chunk_size) return query_chunk_size, kv_chunk_size def exists(val): @@ -267,7 +241,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape steps = 1 if q.device.type == "mps": elements_full = q.shape[0] * q.shape[1] * k.shape[1] - steps = _mps_forced_attn_chunk_steps(elements_full, steps, tag="basic") + steps = _mps_forced_attn_chunk_steps(elements_full, steps) slice_size = math.ceil(q.shape[1] / steps) is_bool_mask = False @@ -306,7 +280,6 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape sim = sim.softmax(dim=-1) out[:, i:end] = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v) - _diag_log_attn_size("basic", q.shape[0], end - i, k.shape[1], extra=f"steps={steps} b={b} heads={heads} full_seq_q={q.shape[1]}") del q, k @@ -378,12 +351,7 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None, if query.device.type == "mps": query_chunk_size, kv_chunk_size = _mps_cap_subquad_chunk_sizes( - batch_x_heads, q_tokens, k_tokens, query_chunk_size, kv_chunk_size, tag="sub_quad") - _diag_log_attn_size( - "sub_quad", batch_x_heads, min(query_chunk_size, q_tokens), - kv_chunk_size if kv_chunk_size is not None else int(math.sqrt(k_tokens)), - extra=f"q_tokens={q_tokens} k_tokens={k_tokens}" - ) + batch_x_heads, q_tokens, k_tokens, query_chunk_size, kv_chunk_size) if mask is not None: if len(mask.shape) == 2: @@ -461,7 +429,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape # of the memory-based steps calculation above. if q.device.type == "mps": elements_full = q.shape[0] * q.shape[1] * k.shape[1] - steps = _mps_forced_attn_chunk_steps(elements_full, steps, tag="split") + steps = _mps_forced_attn_chunk_steps(elements_full, steps) if steps > 64: max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 @@ -481,7 +449,6 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape while True: try: slice_size = min(q.shape[1], math.ceil(q.shape[1] / steps)) - _diag_log_attn_size("split", q.shape[0], slice_size, k.shape[1], extra=f"steps={steps} b={b} heads={heads} full_seq_q={q.shape[1]}") for i in range(0, q.shape[1], slice_size): end = i + slice_size if upcast: @@ -656,7 +623,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha steps = 1 if q.device.type == "mps": elements_full = q.shape[0] * q.shape[1] * q.shape[2] * k.shape[2] - steps = _mps_forced_attn_chunk_steps(elements_full, steps, tag="pytorch") + steps = _mps_forced_attn_chunk_steps(elements_full, steps) slice_size = math.ceil(q.shape[2] / steps) raw_out = torch.empty((b, q.shape[1], q.shape[2], dim_head), dtype=q.dtype, layout=q.layout, device=q.device) @@ -665,7 +632,6 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha mask_chunk = mask if mask is not None and mask.shape[-2] != 1: mask_chunk = mask[..., i:end, :] - _diag_log_attn_size("pytorch", q.shape[0] * q.shape[1], min(end, q.shape[2]) - i, k.shape[2], extra=f"steps={steps} b={b} heads={q.shape[1]} full_seq_q={q.shape[2]} dtype={q.dtype}") raw_out[:, :, i:end] = comfy.ops.scaled_dot_product_attention( q[:, :, i:end], k, v, attn_mask=mask_chunk, dropout_p=0.0, is_causal=False, **sdpa_extra ) From 55760cdeb58cdf283472f44fc0cc0ecd1bb633e8 Mon Sep 17 00:00:00 2001 From: Scott Yewell Date: Wed, 8 Jul 2026 16:31:12 -0400 Subject: [PATCH 3/3] attention_pytorch: skip the chunk buffer/copy when steps==1 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 Claude-Session: https://claude.ai/code/session_01HAMqHgFBD6e9wjU8k8U6uW --- comfy/ldm/modules/attention.py | 35 ++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 18efdf148..439436179 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -618,28 +618,31 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha if SDP_BATCH_LIMIT >= b: # See _mps_forced_attn_chunk_steps: on MPS a single attention matrix over # ~2^31 elements silently corrupts, so chunk over seq_q to stay under that. - # steps stays 1 (one iteration over the full sequence, identical to the - # previous unchunked call) everywhere this doesn't apply. steps = 1 if q.device.type == "mps": elements_full = q.shape[0] * q.shape[1] * q.shape[2] * k.shape[2] steps = _mps_forced_attn_chunk_steps(elements_full, steps) - slice_size = math.ceil(q.shape[2] / steps) - raw_out = torch.empty((b, q.shape[1], q.shape[2], dim_head), dtype=q.dtype, layout=q.layout, device=q.device) - for i in range(0, q.shape[2], slice_size): - end = i + slice_size - mask_chunk = mask - if mask is not None and mask.shape[-2] != 1: - mask_chunk = mask[..., i:end, :] - raw_out[:, :, i:end] = comfy.ops.scaled_dot_product_attention( - q[:, :, i:end], k, v, attn_mask=mask_chunk, dropout_p=0.0, is_causal=False, **sdpa_extra - ) - - if skip_output_reshape: - out = raw_out + if steps == 1: + out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra) + if not skip_output_reshape: + out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) else: - out = raw_out.transpose(1, 2).reshape(b, -1, heads * dim_head) + slice_size = math.ceil(q.shape[2] / steps) + raw_out = torch.empty((b, q.shape[1], q.shape[2], dim_head), dtype=q.dtype, layout=q.layout, device=q.device) + for i in range(0, q.shape[2], slice_size): + end = i + slice_size + mask_chunk = mask + if mask is not None and mask.shape[-2] != 1: + mask_chunk = mask[..., i:end, :] + raw_out[:, :, i:end] = comfy.ops.scaled_dot_product_attention( + q[:, :, i:end], k, v, attn_mask=mask_chunk, dropout_p=0.0, is_causal=False, **sdpa_extra + ) + + if skip_output_reshape: + out = raw_out + else: + out = raw_out.transpose(1, 2).reshape(b, -1, heads * dim_head) else: out = torch.empty((b, q.shape[2], heads * dim_head), dtype=q.dtype, layout=q.layout, device=q.device) for i in range(0, b, SDP_BATCH_LIMIT):