diff --git a/comfy/ldm/ace/ace_step15.py b/comfy/ldm/ace/ace_step15.py index 2ca2d26c4..02182c49f 100644 --- a/comfy/ldm/ace/ace_step15.py +++ b/comfy/ldm/ace/ace_step15.py @@ -217,10 +217,7 @@ class AceStepAttention(nn.Module): cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - n_rep = self.num_heads // self.num_kv_heads - if n_rep > 1: - key_states = key_states.repeat_interleave(n_rep, dim=1) - value_states = value_states.repeat_interleave(n_rep, dim=1) + gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {} attn_bias = None if self.sliding_window is not None and not self.is_cross_attention: @@ -244,7 +241,7 @@ class AceStepAttention(nn.Module): else: attn_bias = window_bias - attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False) + attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False, **gqa_kwargs) attn_output = self.o_proj(attn_output) return attn_output diff --git a/comfy/ldm/audio/dit.py b/comfy/ldm/audio/dit.py index c28be5b49..b0759a240 100644 --- a/comfy/ldm/audio/dit.py +++ b/comfy/ldm/audio/dit.py @@ -425,19 +425,16 @@ class Attention(nn.Module): if n == 1 and causal: causal = False - if h != kv_h: - # Repeat interleave kv_heads to match q_heads - heads_per_kv_head = h // kv_h - k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v)) + gqa_kwargs = {"enable_gqa": True} if h != kv_h else {} if self.differential: q, q_diff = q.unbind(dim=1) k, k_diff = k.unbind(dim=1) - out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options) - out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options) + out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs) + out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs) out = out - out_diff else: - out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options) + out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs) out = self.to_out(out) diff --git a/comfy/ldm/boogu/model.py b/comfy/ldm/boogu/model.py index 966f3c583..ca88bdeb1 100644 --- a/comfy/ldm/boogu/model.py +++ b/comfy/ldm/boogu/model.py @@ -74,11 +74,8 @@ class BooguDoubleStreamProcessor(nn.Module): key = key.transpose(1, 2) value = value.transpose(1, 2) - if attn.kv_heads < attn.heads: - key = key.repeat_interleave(attn.heads // attn.kv_heads, dim=1) - value = value.repeat_interleave(attn.heads // attn.kv_heads, dim=1) - - hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options) + gqa_kwargs = {"enable_gqa": True} if attn.kv_heads < attn.heads else {} + hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs) # Split back to instruction/image, apply per-stream output projections, recombine. instruct_hidden_states = self.instruct_out(hidden_states[:, :L_instruct]) diff --git a/comfy/ldm/omnigen/omnigen2.py b/comfy/ldm/omnigen/omnigen2.py index b8da4cf39..d18a9f6d0 100644 --- a/comfy/ldm/omnigen/omnigen2.py +++ b/comfy/ldm/omnigen/omnigen2.py @@ -141,11 +141,8 @@ class Attention(nn.Module): key = key.transpose(1, 2) value = value.transpose(1, 2) - if self.kv_heads < self.heads: - key = key.repeat_interleave(self.heads // self.kv_heads, dim=1) - value = value.repeat_interleave(self.heads // self.kv_heads, dim=1) - - hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options) + gqa_kwargs = {"enable_gqa": True} if self.kv_heads < self.heads else {} + hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs) hidden_states = self.to_out[0](hidden_states) return hidden_states diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index e9f38a9a2..ce5cc1645 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -550,10 +550,8 @@ class Attention(nn.Module): xv = xv[:, :, -sliding_window:] attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None - xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - - output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True) + gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {} + output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs) return self.o_proj(output), present_key_value class MLP(nn.Module): diff --git a/comfy/text_encoders/qwen35.py b/comfy/text_encoders/qwen35.py index 71a17990f..304a4357f 100644 --- a/comfy/text_encoders/qwen35.py +++ b/comfy/text_encoders/qwen35.py @@ -366,12 +366,8 @@ class GatedAttention(nn.Module): xv = torch.cat((past_value[:, :, :index], xv), dim=2) present_key_value = (xk, xv, index + num_tokens) - # Expand KV heads for GQA - if self.num_heads != self.num_kv_heads: - xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - - output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True) + gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {} + output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs) output = output * gate.sigmoid() return self.o_proj(output), present_key_value