From c35a622acdabf99f60bba737f07977fef1ef97f4 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Mon, 13 Jul 2026 12:52:28 -0700
Subject: [PATCH] Fix hidream o1 regression. (#14923)

---
 comfy/ldm/hidream_o1/attention.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/comfy/ldm/hidream_o1/attention.py b/comfy/ldm/hidream_o1/attention.py
index 1b68f1771..afb2be9b8 100644
--- a/comfy/ldm/hidream_o1/attention.py
+++ b/comfy/ldm/hidream_o1/attention.py
@@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None):
     The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
     """
 
-    def two_pass_attention(q, k, v, heads, **kwargs):
+    def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs):
         B, H, T, D = q.shape
 
         if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
-            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
+            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
         elif ar_len >= T:
-            out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
+            out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
         elif ar_len <= 0:
-            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
+            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
         else:
             out_ar = comfy.ops.scaled_dot_product_attention(
                 q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
-                attn_mask=None, dropout_p=0.0, is_causal=True,
+                attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
             )
             out_gen = optimized_attention(
                 q[:, :, ar_len:], k, v, heads,
                 mask=None, skip_reshape=True, skip_output_reshape=True,
-                transformer_options=transformer_options,
+                transformer_options=transformer_options, enable_gqa=enable_gqa,
             )
             out = torch.cat([out_ar, out_gen], dim=2)