From 501b808481c91b1481c70ce22d2427569ca37b75 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 4 Jul 2026 20:28:39 -0400 Subject: [PATCH 1/3] Support gqa on all attention backends, drop support for pytorch 2.4 --- README.md | 2 +- comfy/ldm/modules/attention.py | 156 ++++++++++++++++++++------------- comfy/text_encoders/gpt_oss.py | 6 +- 3 files changed, 95 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index bcec86377..14c8d2cb2 100644 --- a/README.md +++ b/README.md @@ -229,7 +229,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12 -torch 2.4 and above is supported but some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. +torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it. ### Instructions: diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 55360535a..b4cdcc8ed 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -1,5 +1,6 @@ import math import sys +import inspect import torch import torch.nn.functional as F @@ -14,16 +15,16 @@ from .sub_quadratic_attention import efficient_dot_product_attention from comfy import model_management -TORCH_HAS_GQA = model_management.torch_version_numeric >= (2, 5) - if model_management.xformers_enabled(): import xformers import xformers.ops SAGE_ATTENTION_IS_AVAILABLE = False +SAGE_ATTENTION_SUPPORTS_MASK = False try: from sageattention import sageattn SAGE_ATTENTION_IS_AVAILABLE = True + SAGE_ATTENTION_SUPPORTS_MASK = "attn_mask" in inspect.signature(sageattn).parameters except ImportError as e: if model_management.sage_attention_enabled(): if e.name == "sageattention": @@ -89,6 +90,44 @@ def default(val, d): return val return d +def _gqa_repeat_factor(query_heads, key_heads, value_heads): + if key_heads != value_heads: + raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}") + if query_heads == key_heads: + return 1 + if query_heads % key_heads != 0: + raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}") + return query_heads // key_heads + +def _repeat_kv_for_gqa(k, v, query_heads, head_dim): + n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim]) + if n_rep > 1: + k = k.repeat_interleave(n_rep, dim=head_dim) + v = v.repeat_interleave(n_rep, dim=head_dim) + return k, v + +def _heads_from_dim(tensor, dim_head, name): + inner_dim = tensor.shape[-1] + if inner_dim % dim_head != 0: + raise ValueError(f"{name} inner dimension {inner_dim} is not divisible by head dimension {dim_head}") + return inner_dim // dim_head + +def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_kv=True): + q = q.unsqueeze(3).reshape(b, -1, heads, dim_head) + if enable_gqa: + key_heads = _heads_from_dim(k, dim_head, "Key") + value_heads = _heads_from_dim(v, dim_head, "Value") + else: + key_heads = heads + value_heads = heads + k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head) + v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head) + if enable_gqa: + _gqa_repeat_factor(heads, key_heads, value_heads) + if expand_kv: + k, v = _repeat_kv_for_gqa(k, v, heads, -2) + return q, k, v + # feedforward class GEGLU(nn.Module): @@ -152,28 +191,19 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape b, _, dim_head = q.shape dim_head //= heads - if kwargs.get("enable_gqa", False) and q.shape[-3] != k.shape[-3]: - n_rep = q.shape[-3] // k.shape[-3] - k = k.repeat_interleave(n_rep, dim=-3) - v = v.repeat_interleave(n_rep, dim=-3) - scale = kwargs.get("scale", dim_head ** -0.5) h = heads if skip_reshape: - q, k, v = map( + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3) + q, k, v = map( lambda t: t.reshape(b * heads, -1, dim_head), (q, k, v), ) else: - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, -1, heads, dim_head) - .permute(0, 2, 1, 3) - .reshape(b * heads, -1, dim_head) - .contiguous(), - (q, k, v), - ) + 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: @@ -231,13 +261,16 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None, query = query * (kwargs["scale"] * dim_head ** 0.5) if skip_reshape: + if kwargs.get("enable_gqa", False): + key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3) query = query.reshape(b * heads, -1, dim_head) value = value.reshape(b * heads, -1, dim_head) key = key.reshape(b * heads, -1, dim_head).movedim(1, 2) else: - query = query.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) - value = value.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) - key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1) + query, key, value = _reshape_qkv_to_heads(query, key, value, b, heads, dim_head, kwargs.get("enable_gqa", False)) + query = query.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) + value = value.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) + key = key.permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1) dtype = query.dtype @@ -304,19 +337,15 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape scale = kwargs.get("scale", dim_head ** -0.5) if skip_reshape: - q, k, v = map( + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3) + q, k, v = map( lambda t: t.reshape(b * heads, -1, dim_head), (q, k, v), ) else: - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, -1, heads, dim_head) - .permute(0, 2, 1, 3) - .reshape(b * heads, -1, dim_head) - .contiguous(), - (q, k, v), - ) + 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)) r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) @@ -438,7 +467,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh disabled_xformers = True if disabled_xformers: - return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs) + return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs) if skip_reshape: # b h k d -> b k h d @@ -446,13 +475,12 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh lambda t: t.permute(0, 2, 1, 3), (q, k, v), ) + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2) # actually do the reshaping else: dim_head //= heads - q, k, v = map( - lambda t: t.reshape(b, -1, heads, dim_head), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) if mask is not None: # add a singleton batch dimension @@ -474,7 +502,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh mask = mask_out[..., :mask.shape[-1]] mask = mask.expand(b, heads, -1, -1) - out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask) + out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, scale=kwargs.get("scale", None)) if skip_output_reshape: out = out.permute(0, 2, 1, 3) @@ -498,10 +526,8 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha else: b, _, dim_head = q.shape dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False) + q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) if mask is not None: # add a batch dimension if there isn't already one @@ -511,9 +537,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha if mask.ndim == 3: mask = mask.unsqueeze(1) - # Pass through extra SDPA kwargs (scale, enable_gqa) if provided - # enable_gqa requires PyTorch 2.5+; older versions use manual KV expansion above - sdpa_keys = ("scale", "enable_gqa") if TORCH_HAS_GQA else ("scale",) + 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: @@ -541,20 +565,19 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha @wrap_attn def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): - if kwargs.get("low_precision_attention", True) is False: + if kwargs.get("low_precision_attention", True) is False or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK): return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs) exception_fallback = False if skip_reshape: b, _, _, dim_head = q.shape tensor_layout = "HND" + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3) else: b, _, dim_head = q.shape dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) tensor_layout = "NHD" if mask is not None: @@ -565,8 +588,12 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= if mask.ndim == 3: mask = mask.unsqueeze(1) + sage_kwargs = {"is_causal": False, "tensor_layout": tensor_layout, "sm_scale": kwargs.get("scale", None), "smooth_k": False} + if mask is not None: + sage_kwargs["attn_mask"] = mask + try: - out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout) + out = sageattn(q, k, v, **sage_kwargs) except Exception as e: logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e)) exception_fallback = True @@ -616,6 +643,8 @@ 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 @@ -643,10 +672,8 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape ) if not skip_reshape: - q_s, k_s, v_s = map( - lambda t: t.view(B, -1, heads, dim_head).permute(0, 2, 1, 3).contiguous(), - (q, k, v), - ) + 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 try: @@ -662,7 +689,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape q, k, v, heads, mask=mask, attn_precision=attn_precision, - skip_reshape=False, + skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs ) @@ -681,19 +708,20 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape try: @torch.library.custom_op("flash_attention::flash_attn", mutates_args=()) def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor: - return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal) + dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor: + softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale + return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal, softmax_scale=softmax_scale_arg) @flash_attn_wrapper.register_fake - def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False): + def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False, softmax_scale=-1.0): # Output shape is the same as q return q.new_empty(q.shape) except AttributeError as error: FLASH_ATTN_ERROR = error def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor: + dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor: assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}" @wrap_attn @@ -703,10 +731,8 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape else: b, _, dim_head = q.shape dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False) + q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) if mask is not None: # add a batch dimension if there isn't already one @@ -725,10 +751,16 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape v.transpose(1, 2), dropout_p=0.0, causal=False, + softmax_scale=kwargs.get("scale", -1.0), ).transpose(1, 2) except Exception as e: logging.warning(f"Flash Attention failed, using default SDPA: {e}") - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) + sdpa_extra = {} + if kwargs.get("enable_gqa", False): + sdpa_extra["enable_gqa"] = True + if "scale" in kwargs: + sdpa_extra["scale"] = kwargs["scale"] + out = torch.nn.functional.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) @@ -1209,5 +1241,3 @@ class SpatialVideoTransformer(SpatialTransformer): x = self.proj_out(x) out = x + x_in return out - - diff --git a/comfy/text_encoders/gpt_oss.py b/comfy/text_encoders/gpt_oss.py index d596ef9a0..066796b6a 100644 --- a/comfy/text_encoders/gpt_oss.py +++ b/comfy/text_encoders/gpt_oss.py @@ -12,7 +12,7 @@ import torch.nn.functional as F import comfy.ops from comfy import sd1_clip -from comfy.ldm.modules.attention import TORCH_HAS_GQA, optimized_attention_for_device +from comfy.ldm.modules.attention import optimized_attention_for_device from comfy.text_encoders.llama import RMSNorm, apply_rope @@ -110,10 +110,6 @@ def _attention_with_sinks(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sin putting the sink logit in the mask at that column. """ - if num_kv_groups > 1 and not TORCH_HAS_GQA: - k = k.repeat_interleave(num_kv_groups, dim=1) - v = v.repeat_interleave(num_kv_groups, dim=1) - B, _, S_q, D = q.shape H_kv = k.shape[1] S_kv = k.shape[-2] 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 2/3] 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 From ec8c2e731549d57bc329f73dc63a57a2720cfdca Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 4 Jul 2026 21:57:30 -0400 Subject: [PATCH 3/3] Use gqa on models that repeat k and v. --- comfy/ldm/ace/ace_step15.py | 7 ++----- comfy/ldm/audio/dit.py | 11 ++++------- comfy/ldm/boogu/model.py | 7 ++----- comfy/ldm/omnigen/omnigen2.py | 7 ++----- comfy/text_encoders/llama.py | 6 ++---- comfy/text_encoders/qwen35.py | 8 ++------ 6 files changed, 14 insertions(+), 32 deletions(-) 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