xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-14 10:30:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

wan: Optimize qkv intermediates on attention

Browse Source 
As commented. The former logic computed independent pieces of QKV in
parallel which help more inference intermediates in VRAM spiking
VRAM usage. Fully roping Q and garbage collecting the intermediates
before touching K reduces the peak inference VRAM usage.
This commit is contained in:
Rattus 2025-09-16 21:03:44 +10:00
parent 8922c21c9e
commit 39eb74cabe
1 changed files with 13 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
comfy/ldm/wan/model.py
View File

@ -8,7 +8,7 @@ from einops import rearrange
from comfy.ldm.modules.attention import optimized_attention from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.flux.layers import EmbedND from comfy.ldm.flux.layers import EmbedND
from comfy.ldm.flux.math import apply_rope from comfy.ldm.flux.math import apply_rope1
import comfy.ldm.common_dit import comfy.ldm.common_dit
import comfy.model_management import comfy.model_management
import comfy.patcher_extension import comfy.patcher_extension
@ -60,20 +60,24 @@ class WanSelfAttention(nn.Module):
""" """
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
# query, key, value function def qkv_fn_q(x):
def qkv_fn(x):
q = self.norm_q(self.q(x)).view(b, s, n, d) q = self.norm_q(self.q(x)).view(b, s, n, d)
k = self.norm_k(self.k(x)).view(b, s, n, d) return apply_rope1(q, freqs)
v = self.v(x).view(b, s, n * d)
return q, k, v
q, k, v = qkv_fn(x) def qkv_fn_k(x):
q, k = apply_rope(q, k, freqs) k = self.norm_k(self.k(x)).view(b, s, n, d)
return apply_rope1(k, freqs)
#These two are VRAM hogs, so we want to do all of q computation and
#have pytorch garbage collect the intermediates on the sub function
#return before we touch k
q = qkv_fn_q(x)
k = qkv_fn_k(x)
x = optimized_attention( x = optimized_attention(
q.view(b, s, n * d), q.view(b, s, n * d),
k.view(b, s, n * d), k.view(b, s, n * d),
v, self.v(x).view(b, s, n * d),
heads=self.num_heads, heads=self.num_heads,
transformer_options=transformer_options, transformer_options=transformer_options,
) )