From 2131e630e0a12e24c7d80a5be7e2e0fd7e663e52 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sat, 4 Jul 2026 20:53:17 -0400 Subject: [PATCH] Fix. --- comfy/ldm/modules/attention.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index b4cdcc8ed..2411aff5c 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -643,9 +643,6 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape skip_output_reshape=skip_output_reshape, **kwargs ) - if kwargs.get("enable_gqa", False): - k, v = _repeat_kv_for_gqa(k, v, H, -3) - q_s, k_s, v_s = q, k, v N = q.shape[2] dim_head = D else: @@ -671,7 +668,13 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape **kwargs ) - if not skip_reshape: + if skip_reshape: + q_s = q + if kwargs.get("enable_gqa", False): + k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3) + else: + k_s, v_s = k, v + else: q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False)) q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s)) B, H, L, D = q_s.shape