wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-07 15:10:50 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Support gqa on all attention backends, drop support for pytorch 2.4

Browse Source
This commit is contained in:
comfyanonymous 2026-07-04 20:28:39 -04:00
parent 7f287b705e
commit 501b808481
3 changed files with 95 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@ -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:

156
comfy/ldm/modules/attention.py
View File

@ -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

6
comfy/text_encoders/gpt_oss.py
View File

@ -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]