From 55760cdeb58cdf283472f44fc0cc0ecd1bb633e8 Mon Sep 17 00:00:00 2001
From: Scott Yewell <scott@discretioncapital.com>
Date: Wed, 8 Jul 2026 16:31:12 -0400
Subject: [PATCH] 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 <noreply@anthropic.com>
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):