diff --git a/comfy/cli_args.py b/comfy/cli_args.py index cc1f12482..dfdfa9482 100644 --- a/comfy/cli_args.py +++ b/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.") diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 043df28df..ca0d54117 100644 --- a/comfy/ldm/modules/attention.py +++ b/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") diff --git a/comfy/model_management.py b/comfy/model_management.py index bbfc3c7a1..f45c67a64 100644 --- a/comfy/model_management.py +++ b/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