Merge d52cd3e356f1ee9f33b6856807d8e78fc882c098 into d6b977b2e680e98ad18a37ee13783da4f30e15f4

comfy/cli_args.py

@@ -111,6 +111,7 @@ attn_group.add_argument("--use-split-cross-attention", action="store_true", help
 attn_group.add_argument("--use-quad-cross-attention", action="store_true", help="Use the sub-quadratic cross attention optimization . Ignored when xformers is used.")
 attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
 attn_group.add_argument("--use-sage-attention", action="store_true", help="Use sage attention.")
+attn_group.add_argument("--use-sage-attention3", action="store_true", help="Use sage attention 3. Supported only on blackwell GPUs.")
 attn_group.add_argument("--use-flash-attention", action="store_true", help="Use FlashAttention.")
 
 parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")

comfy/ldm/modules/attention.py

@@ -20,6 +20,7 @@ if model_management.xformers_enabled():
 if model_management.sage_attention_enabled():
     try:
         from sageattention import sageattn
+        logging.info("Found SageAttention 1.x/2.x (sageattention package)")
     except ModuleNotFoundError as e:
         if e.name == "sageattention":
             logging.error(f"\n\nTo use the `--use-sage-attention` feature, the `sageattention` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install sageattention")
@@ -27,6 +28,17 @@ if model_management.sage_attention_enabled():
             raise e
         exit(-1)
 
+if model_management.sage_attention3_enabled():
+    try:
+        from sageattn import sageattn_blackwell
+        logging.info("Found SageAttention3 (sageattn package)")
+    except ModuleNotFoundError as e:
+        if e.name == "sageattn":
+            logging.error(f"\n\nTo use the `--use-sage-attention3` feature, the `sageattn` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install sageattn")
+        else:
+            raise e
+        exit(-1)
+
 if model_management.flash_attention_enabled():
     try:
         from flash_attn import flash_attn_func
@@ -470,7 +482,6 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
             ).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, skip_output_reshape=False):
     if skip_reshape:
         b, _, _, dim_head = q.shape
@@ -502,7 +513,6 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
                 (q, k, v),
             )
         return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape)
-
     if tensor_layout == "HND":
         if not skip_output_reshape:
             out = (
@@ -516,6 +526,66 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
     return out
 
 
+def attention_sage3(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+    if skip_reshape:
+        b, _, _, dim_head = q.shape
+        tensor_layout = "HND"
+    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),
+        )
+        tensor_layout = "NHD"
+
+    if mask is not None:
+        # add a batch dimension if there isn't already one
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        # add a heads dimension if there isn't already one
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+
+    try:
+        if dim_head >= 256:
+            # SageAttention3 doesn't support head_dim >= 256, fall back to pytorch
+            logging.warning(f"SageAttention3 doesn't support head_dim >= 256 (got {dim_head}), falling back to pytorch attention")
+            if tensor_layout == "NHD":
+                q, k, v = map(
+                    lambda t: t.transpose(1, 2),
+                    (q, k, v),
+                )
+            return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape)
+
+        # SageAttention3 expects tensor layout as (batch, heads, seq_len, head_dim)
+        if tensor_layout == "NHD":
+            q_sa3, k_sa3, v_sa3 = map(lambda t: t.transpose(1, 2), (q, k, v))
+        else:
+            q_sa3, k_sa3, v_sa3 = q, k, v
+
+        out = sageattn_blackwell(q_sa3, k_sa3, v_sa3, attn_mask=mask, is_causal=False, per_block_mean=False)
+
+        # Convert back to expected layout
+        if tensor_layout == "HND":
+            if not skip_output_reshape:
+                out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+        else:
+            if skip_output_reshape:
+                out = out.transpose(1, 2)
+            else:
+                out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+    except Exception as e:
+        logging.error("Error running sage attention 3: {}, using pytorch attention instead.".format(e))
+        if tensor_layout == "NHD":
+            q, k, v = map(
+                lambda t: t.transpose(1, 2),
+                (q, k, v),
+            )
+        return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape)
+    return out
+
+
 try:
     @torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
     def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
@@ -575,8 +645,11 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
 
 optimized_attention = attention_basic
 
-if model_management.sage_attention_enabled():
-    logging.info("Using sage attention")
+if model_management.sage_attention3_enabled():
+    print("Using sage attention 3")
+    optimized_attention = attention_sage3
+elif model_management.sage_attention_enabled():
+    print("Using sage attention 1.x/2.x")
     optimized_attention = attention_sage
 elif model_management.xformers_enabled():
     logging.info("Using xformers attention")

comfy/model_management.py

@@ -1079,6 +1079,9 @@ def cast_to_device(tensor, device, dtype, copy=False):
 def sage_attention_enabled():
     return args.use_sage_attention
 
+def sage_attention3_enabled():
+    return args.use_sage_attention3
+
 def flash_attention_enabled():
     return args.use_flash_attention