From e77fb313d325bca87449ddc81b994dd61d9a6bb6 Mon Sep 17 00:00:00 2001 From: Haoming Date: Sat, 4 Oct 2025 02:07:54 +0800 Subject: [PATCH] update Attention https://github.com/comfyanonymous/ComfyUI/blob/v0.3.62/comfy/ldm/modules/attention.py --- backend/args.py | 56 +++++++++ backend/attention.py | 223 ++++++++++++++++++++++++----------- backend/memory_management.py | 8 +- modules/cmd_args.py | 3 - 4 files changed, 215 insertions(+), 75 deletions(-) diff --git a/backend/args.py b/backend/args.py index 6e02d392..3fd62de3 100644 --- a/backend/args.py +++ b/backend/args.py @@ -1,4 +1,25 @@ import argparse +import enum + + +class EnumAction(argparse.Action): + """Argparse `action` for handling Enum""" + + def __init__(self, **kwargs): + enum_type = kwargs.pop("type", None) + assert issubclass(enum_type, enum.Enum) + + choices = tuple(e.value for e in enum_type) + kwargs.setdefault("choices", choices) + kwargs.setdefault("metavar", f"[{','.join(list(choices))}]") + + super(EnumAction, self).__init__(**kwargs) + self._enum = enum_type + + def __call__(self, parser, namespace, values, option_string=None): + value = self._enum(values) + setattr(namespace, self.dest, value) + parser = argparse.ArgumentParser() @@ -37,6 +58,10 @@ upcast = parser.add_mutually_exclusive_group() upcast.add_argument("--force-upcast-attention", action="store_true") upcast.add_argument("--disable-attention-upcast", action="store_true") +parser.add_argument("--xformers", action="store_true", help="install xformers for cross attention") +parser.add_argument("--sage", action="store_true", help="install sageattention") +parser.add_argument("--flash", action="store_true", help="install flash_attn") + parser.add_argument("--disable-xformers", action="store_true") parser.add_argument("--disable-sage", action="store_true") parser.add_argument("--disable-flash", action="store_true") @@ -65,6 +90,37 @@ parser.add_argument("--fast-fp16", action="store_true") parser.add_argument("--mmap-torch-files", action="store_true") parser.add_argument("--disable-mmap", action="store_true") + +class SageAttentionFuncs(enum.Enum): + auto = "auto" + fp16_triton = "fp16_triton" + fp16_cuda = "fp16_cuda" + fp8_cuda = "fp8_cuda" + + +class Sage_quantization_backend(enum.Enum): + cuda = "cuda" + triton = "triton" + + +class Sage_qk_quant_gran(enum.Enum): + per_warp = "per_warp" + per_thread = "per_thread" + + +class Sage_pv_accum_dtype(enum.Enum): + fp16 = "fp16" + fp32 = "fp32" + fp16fp32 = "fp16+fp32" + fp32fp32 = "fp32+fp32" + + +parser.add_argument("--sage2-function", type=SageAttentionFuncs, default=SageAttentionFuncs.auto, action=EnumAction) +parser.add_argument("--sage-quantization-backend", type=Sage_quantization_backend, default=Sage_quantization_backend.triton, action=EnumAction) +parser.add_argument("--sage-quant-gran", type=Sage_qk_quant_gran, default=Sage_qk_quant_gran.per_thread, action=EnumAction) +parser.add_argument("--sage-accum-dtype", type=Sage_pv_accum_dtype, default=Sage_pv_accum_dtype.fp32, action=EnumAction) + + args, _ = parser.parse_known_args() dynamic_args = dict(embedding_dir="./embeddings", emphasis_name="original") diff --git a/backend/attention.py b/backend/attention.py index 8024ff2c..beb4fb0e 100644 --- a/backend/attention.py +++ b/backend/attention.py @@ -4,7 +4,7 @@ import einops import torch from backend import memory_management -from backend.args import args +from backend.args import args, SageAttentionFuncs from modules.errors import display_once if memory_management.xformers_enabled(): @@ -40,9 +40,6 @@ if memory_management.flash_enabled(): return q.new_empty(q.shape) -FORCE_UPCAST_ATTENTION_DTYPE = memory_management.force_upcast_attention_dtype() - - def get_xformers_flash_attention_op(q, k, v): try: flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp @@ -55,11 +52,14 @@ def get_xformers_flash_attention_op(q, k, v): return None -def get_attn_precision(attn_precision=torch.float32): - if args.disable_attention_upcast: +FORCE_UPCAST_ATTENTION_DTYPE = memory_management.force_upcast_attention_dtype() + + +def get_attn_precision(attn_precision, current_dtype): + if args.dont_upcast_attention: return None if FORCE_UPCAST_ATTENTION_DTYPE is not None: - return FORCE_UPCAST_ATTENTION_DTYPE + return FORCE_UPCAST_ATTENTION_DTYPE.get(current_dtype, attn_precision) return attn_precision @@ -67,11 +67,23 @@ def exists(val): return val is not None +def default(val, d): + if exists(val): + return val + return d + + +if memory_management.is_nvidia(): + SDP_BATCH_LIMIT = 2**15 +else: + SDP_BATCH_LIMIT = 2**31 + + # ========== Diffusion ========== # -def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): - attn_precision = get_attn_precision(attn_precision) +def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + attn_precision = get_attn_precision(attn_precision, q.dtype) if skip_reshape: b, _, _, dim_head = q.shape @@ -116,12 +128,15 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape sim = sim.softmax(dim=-1) out = torch.einsum("b i j, b j d -> b i d", sim.to(v.dtype), v) - out = out.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head) - return out + + if skip_output_reshape: + return out.unsqueeze(0).reshape(b, heads, -1, dim_head) + else: + return out.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head) -def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): - attn_precision = get_attn_precision(attn_precision) +def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + attn_precision = get_attn_precision(attn_precision, q.dtype) if skip_reshape: b, _, _, dim_head = q.shape @@ -131,7 +146,6 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape scale = dim_head**-0.5 - h = heads if skip_reshape: q, k, v = map( lambda t: t.reshape(b * heads, -1, dim_head), @@ -162,8 +176,6 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape if mem_required > mem_free_total: steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2))) - # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB " - # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}") if steps > 64: max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 @@ -176,7 +188,6 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape bs = mask.shape[0] mask = mask.reshape(bs, -1, mask.shape[-2], mask.shape[-1]).expand(b, heads, -1, -1).reshape(-1, mask.shape[-2], mask.shape[-1]) - # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size) first_op_done = False cleared_cache = False while True: @@ -194,7 +205,10 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape if len(mask.shape) == 2: s1 += mask[i:end] else: - s1 += mask[:, i:end] + if mask.shape[1] == 1: + s1 += mask + else: + s1 += mask[:, i:end] s2 = s1.softmax(dim=-1).to(v.dtype) del s1 @@ -208,59 +222,69 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape memory_management.soft_empty_cache(True) if cleared_cache == False: cleared_cache = True - print("out of memory error, emptying cache and trying again") + print(f"[Out of Memory Error] emptying cache and trying again...") continue steps *= 2 if steps > 64: raise e - print("out of memory error, increasing steps and trying again {}".format(steps)) + print(f"[Out of Memory Error] increasing steps and trying again {steps}...") else: raise e del q, k, v - r1 = r1.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head) - return r1 - - -def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): - if skip_reshape: - b, _, _, dim_head = q.shape + if skip_output_reshape: + return r1.unsqueeze(0).reshape(b, heads, -1, dim_head) else: - b, _, dim_head = q.shape - dim_head //= heads + return r1.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head) + + +def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + b = q.shape[0] + dim_head = q.shape[-1] + disabled_xformers = False if BROKEN_XFORMERS and b * heads > 65535: - return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape) + disabled_xformers = True + + if not disabled_xformers: + disabled_xformers = torch.jit.is_tracing() or torch.jit.is_scripting() + + if disabled_xformers: + return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs) if skip_reshape: q, k, v = map( - lambda t: t.reshape(b * heads, -1, dim_head), + lambda t: t.permute(0, 2, 1, 3), (q, k, v), ) else: + dim_head //= heads q, k, v = map( lambda t: t.reshape(b, -1, heads, dim_head), (q, k, v), ) if mask is not None: - pad = 8 - q.shape[1] % 8 - mask_out = torch.empty([q.shape[0], q.shape[1], q.shape[1] + pad], dtype=q.dtype, device=q.device) - mask_out[:, :, : mask.shape[-1]] = mask - mask = mask_out[:, :, : mask.shape[-1]] + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(1) + pad = 8 - mask.shape[-1] % 8 + mask_out = torch.empty([mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device) + mask_out[..., : mask.shape[-1]] = mask + 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) - if skip_reshape: - out = out.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head) + if skip_output_reshape: + return out.permute(0, 2, 1, 3) else: - out = out.reshape(b, -1, heads * dim_head) - - return out + return out.reshape(b, -1, heads * dim_head) -def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): +def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): if skip_reshape: b, _, _, dim_head = q.shape else: @@ -271,12 +295,55 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha (q, k, v), ) - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) - out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) + if mask is not None: + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(1) + + if SDP_BATCH_LIMIT >= b: + out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) + if skip_output_reshape: + return out + else: + return out.transpose(1, 2).reshape(b, -1, heads * dim_head) + + out = torch.empty((b, q.shape[2], heads * dim_head), dtype=q.dtype, layout=q.layout, device=q.device) + + for i in range(0, b, SDP_BATCH_LIMIT): + m = mask + if mask is not None: + if mask.shape[0] > 1: + m = mask[i : i + SDP_BATCH_LIMIT] + + out[i : i + SDP_BATCH_LIMIT] = ( + torch.nn.functional.scaled_dot_product_attention( + q[i : i + SDP_BATCH_LIMIT], + k[i : i + SDP_BATCH_LIMIT], + v[i : i + SDP_BATCH_LIMIT], + attn_mask=m, + dropout_p=0.0, + is_causal=False, + ) + .transpose(1, 2) + .reshape(-1, q.shape[2], heads * dim_head) + ) + return out -def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): +if IS_SAGE_2 and args.sage2_function is not SageAttentionFuncs.auto: + from functools import partial + import sageattention + + _function = getattr(sageattention, f"sageattn_qk_int8_pv_{args.sage2_function.value}") + if args.sage2_function is SageAttentionFuncs.fp16_triton: + sageattn = partial(_function, quantization_backend=args.sage_quantization_backend.value) + else: + sageattn = partial(_function, qk_quant_gran=args.sage_quant_gran.value, pv_accum_dtype=args.sage_accum_dtype.value) + + +def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): if skip_reshape: b, _, _, dim_head = q.shape tensor_layout = "HND" @@ -287,9 +354,9 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= if (IS_SAGE_2 and dim_head > 128) or ((not IS_SAGE_2) and (dim_head not in (64, 96, 128))): if memory_management.xformers_enabled(): - return attention_xformers(q, k, v, heads, mask, attn_precision, skip_reshape) + return attention_xformers(q, k, v, heads, mask, attn_precision, skip_reshape, skip_output_reshape, **kwargs) else: - return attention_pytorch(q, k, v, heads, mask, attn_precision, skip_reshape) + return attention_pytorch(q, k, v, heads, mask, attn_precision, skip_reshape, skip_output_reshape, **kwargs) if not skip_reshape: q, k, v = map( @@ -315,19 +382,24 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= (q, k, v), ) if memory_management.xformers_enabled(): - return attention_xformers(q, k, v, heads, mask=mask, skip_reshape=True) + return attention_xformers(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape, **kwargs) else: - return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True) + return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape, **kwargs) if tensor_layout == "HND": - out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) + if skip_output_reshape: + return out + else: + out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) + else: - out = out.reshape(b, -1, heads * dim_head) - - return out + if skip_output_reshape: + return out.transpose(1, 2) + else: + return out.reshape(b, -1, heads * dim_head) -def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False): +def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): if skip_reshape: b, _, _, dim_head = q.shape else: @@ -359,13 +431,24 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape display_once(e, "attention_flash") out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) - out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) - return out + if skip_output_reshape: + return out + else: + out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) if memory_management.sage_enabled(): - print(f"Using SageAttention {'2' if IS_SAGE_2 else ''}") attention_function = attention_sage + match args.sage2_function: + case SageAttentionFuncs.auto: + print(f"Using SageAttention {'2' if IS_SAGE_2 else ''}") + case SageAttentionFuncs.fp16_triton: + print("Using SageAttention (fp16 Triton)") + case SageAttentionFuncs.fp16_cuda: + print("Using SageAttention (fp16 CUDA)") + case SageAttentionFuncs.fp8_cuda: + print("Using SageAttention (fp8 CUDA)") + elif memory_management.flash_enabled(): print("Using FlashAttention") attention_function = attention_flash @@ -425,22 +508,26 @@ def slice_attention_single_head_spatial(q, k, v): def normal_attention_single_head_spatial(q, k, v): # compute attention - b, c, h, w = q.shape + orig_shape = q.shape + b = orig_shape[0] + c = orig_shape[1] - q = q.reshape(b, c, h * w) + q = q.reshape(b, c, -1) q = q.permute(0, 2, 1) # b,hw,c - k = k.reshape(b, c, h * w) # b,c,hw - v = v.reshape(b, c, h * w) + k = k.reshape(b, c, -1) # b,c,hw + v = v.reshape(b, c, -1) r1 = slice_attention_single_head_spatial(q, k, v) - h_ = r1.reshape(b, c, h, w) + h_ = r1.reshape(orig_shape) del r1 return h_ def xformers_attention_single_head_spatial(q, k, v): # compute attention - B, C, H, W = q.shape + orig_shape = q.shape + B = orig_shape[0] + C = orig_shape[1] q, k, v = map( lambda t: t.view(B, C, -1).transpose(1, 2).contiguous(), (q, k, v), @@ -448,15 +535,17 @@ def xformers_attention_single_head_spatial(q, k, v): try: out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=get_xformers_flash_attention_op(q, k, v)) - out = out.transpose(1, 2).reshape(B, C, H, W) + out = out.transpose(1, 2).reshape(orig_shape) except NotImplementedError: - out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W) + out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(orig_shape) return out def pytorch_attention_single_head_spatial(q, k, v): # compute attention - B, C, H, W = q.shape + orig_shape = q.shape + B = orig_shape[0] + C = orig_shape[1] q, k, v = map( lambda t: t.view(B, 1, C, -1).transpose(2, 3).contiguous(), (q, k, v), @@ -464,10 +553,10 @@ def pytorch_attention_single_head_spatial(q, k, v): try: out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) - out = out.transpose(2, 3).reshape(B, C, H, W) + out = out.transpose(2, 3).reshape(orig_shape) except memory_management.OOM_EXCEPTION as e: display_once(e, "pytorch_attention_single_head_spatial") - out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W) + out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(orig_shape) return out diff --git a/backend/memory_management.py b/backend/memory_management.py index c587ab58..0b7bfc24 100644 --- a/backend/memory_management.py +++ b/backend/memory_management.py @@ -1057,14 +1057,12 @@ def pytorch_attention_flash_attention(): def force_upcast_attention_dtype(): upcast = args.force_upcast_attention try: - if platform.mac_ver()[0] in ["14.5"]: # black image bug on OSX Sonoma 14.5 + if tuple(int(n) for n in platform.mac_ver()[0].split(".")) >= (14, 5): upcast = True except Exception: pass - if upcast: - return torch.float32 - else: - return None + + return {torch.float16: torch.float32} if upcast else None def get_free_memory(dev=None, torch_free_too=False): diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 0ec48f50..49d23cc2 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -22,9 +22,6 @@ parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to g parser.add_argument("--ngrok-region", type=str, help="does not do anything.", default="") parser.add_argument("--ngrok-options", type=json.loads, help='The options to pass to ngrok in JSON format, e.g.: \'{"authtoken_from_env":true, "basic_auth":"user:password", "oauth_provider":"google", "oauth_allow_emails":"user@asdf.com"}\'', default=dict()) parser.add_argument("--enable-insecure-extension-access", action="store_true", help="enable extensions tab regardless of other options") -parser.add_argument("--xformers", action="store_true", help="install xformers for cross attention") -parser.add_argument("--sage", action="store_true", help="install sageattention") -parser.add_argument("--flash", action="store_true", help="install flash_attn") parser.add_argument("--listen", action="store_true", help="launch gradio with 0.0.0.0 as server name, allowing to respond to network requests") parser.add_argument("--port", type=int, help="launch gradio with given server port, you need root/admin rights for ports < 1024, defaults to 7860 if available", default=None) parser.add_argument("--ui-config-file", type=str, help="filename to use for ui configuration", default=os.path.join(data_path, "ui-config.json"))