update Attention

https://github.com/comfyanonymous/ComfyUI/blob/v0.3.62/comfy/ldm/modules/attention.py

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

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
 
 

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

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":"[email protected]"}\'', 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"))