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

Use gqa on models that repeat k and v.
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Details
Python Linting / Run Pylint (push) Has been cancelled
Details

Browse Source
This commit is contained in:
comfyanonymous 2026-07-04 21:57:30 -04:00
parent 2131e630e0
commit ec8c2e7315
6 changed files with 14 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
comfy/ldm/ace/ace_step15.py
View File

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

11
comfy/ldm/audio/dit.py
View File

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

7
comfy/ldm/boogu/model.py
View File

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

7
comfy/ldm/omnigen/omnigen2.py
View File

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

6
comfy/text_encoders/llama.py
View File

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

8
comfy/text_encoders/qwen35.py
View File

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