From a9d44d9869978dec96d461842af8d61a056021e3 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Fri, 25 Jul 2025 12:00:11 -0400 Subject: [PATCH 1/9] Preliminar support for sageattention3 --- comfy/ldm/modules/attention.py | 72 +++++++++++++++++++++++++++------- 1 file changed, 57 insertions(+), 15 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 35d2270ee..e49679e07 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -18,9 +18,24 @@ if model_management.xformers_enabled(): import xformers.ops if model_management.sage_attention_enabled(): + sage_attention_available = False + SAGE_ATTENTION_3_AVAILABLE = False + try: - from sageattention import sageattn - except ModuleNotFoundError as e: + from sageattn import sageattn_blackwell + SAGE_ATTENTION_3_AVAILABLE = True + sage_attention_available = True + print("Found SageAttention3 (sageattn package)") + + except ImportError: + try: + from sageattention import sageattn + sage_attention_available = True + print("Found SageAttention2 (sageattention package)") + except ModuleNotFoundError as e: + pass + + if not sage_attention_available: 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") else: @@ -470,7 +485,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 @@ -493,7 +507,46 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= mask = mask.unsqueeze(1) try: - out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout) + if SAGE_ATTENTION_3_AVAILABLE and dim_head < 256: + # 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=True) + + # 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) + elif not SAGE_ATTENTION_3_AVAILABLE: + # Fall back to SageAttention2 if available + out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_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.reshape(b, -1, heads * dim_head) + else: + # SageAttention3 is available but 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) except Exception as e: logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e)) if tensor_layout == "NHD": @@ -502,17 +555,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 = ( - out.transpose(1, 2).reshape(b, -1, heads * dim_head) - ) - else: - if skip_output_reshape: - out = out.transpose(1, 2) - else: - out = out.reshape(b, -1, heads * dim_head) return out From 2c7bb101b41948ed4e4476ae144550ff8c2e8180 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Fri, 25 Jul 2025 13:01:40 -0400 Subject: [PATCH 2/9] Change per_block_mean to False Not sure what are the implications of this, but it seems to make sage3 actually run, --- comfy/ldm/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index e49679e07..200c07678 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -514,7 +514,7 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= 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=True) + 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": From bb7393565ba09b64729bddf9077dcd65fa25b734 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Fri, 25 Jul 2025 21:18:42 -0400 Subject: [PATCH 3/9] Ruff fixes --- comfy/ldm/modules/attention.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 200c07678..6a2290629 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -20,13 +20,13 @@ if model_management.xformers_enabled(): if model_management.sage_attention_enabled(): sage_attention_available = False SAGE_ATTENTION_3_AVAILABLE = False - + try: from sageattn import sageattn_blackwell SAGE_ATTENTION_3_AVAILABLE = True sage_attention_available = True print("Found SageAttention3 (sageattn package)") - + except ImportError: try: from sageattention import sageattn @@ -34,12 +34,9 @@ if model_management.sage_attention_enabled(): print("Found SageAttention2 (sageattention package)") except ModuleNotFoundError as e: pass - + if not sage_attention_available: - 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") - else: - raise e + 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") exit(-1) if model_management.flash_attention_enabled(): From 9a7042397347671ce905a5de5d4ba4068a534699 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Fri, 25 Jul 2025 21:22:34 -0400 Subject: [PATCH 4/9] More ruff fixes --- comfy/ldm/modules/attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 6a2290629..cc6a84f91 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -510,9 +510,9 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= 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: From 27720a1e455e3d2cef371a7fad548d731a7d9169 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Fri, 25 Jul 2025 21:24:47 -0400 Subject: [PATCH 5/9] change print to logging.info --- comfy/ldm/modules/attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index cc6a84f91..50978cd06 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -25,13 +25,13 @@ if model_management.sage_attention_enabled(): from sageattn import sageattn_blackwell SAGE_ATTENTION_3_AVAILABLE = True sage_attention_available = True - print("Found SageAttention3 (sageattn package)") + logging.info("Found SageAttention3 (sageattn package)") except ImportError: try: from sageattention import sageattn sage_attention_available = True - print("Found SageAttention2 (sageattention package)") + logging.info("Found SageAttention2 (sageattention package)") except ModuleNotFoundError as e: pass From 01934bf6397c225120b24665925cb387446cb913 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Fri, 25 Jul 2025 21:28:26 -0400 Subject: [PATCH 6/9] remove unused variable --- comfy/ldm/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 50978cd06..9c7e280e2 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -32,7 +32,7 @@ if model_management.sage_attention_enabled(): from sageattention import sageattn sage_attention_available = True logging.info("Found SageAttention2 (sageattention package)") - except ModuleNotFoundError as e: + except ModuleNotFoundError: pass if not sage_attention_available: From 4a9235007d44cfb1af7a0a29bba86ced583d6bf5 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Sat, 26 Jul 2025 14:21:13 -0400 Subject: [PATCH 7/9] Separate sage 1.x/2.x and sage 3.x on different functions and flags 1. --- comfy/ldm/modules/attention.py | 134 +++++++++++++++++++++------------ 1 file changed, 84 insertions(+), 50 deletions(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 9c7e280e2..75b9fcb5a 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -18,25 +18,25 @@ if model_management.xformers_enabled(): import xformers.ops if model_management.sage_attention_enabled(): - sage_attention_available = False - SAGE_ATTENTION_3_AVAILABLE = False + 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") + else: + raise e + exit(-1) +if model_management.sage_attention3_enabled(): try: from sageattn import sageattn_blackwell - SAGE_ATTENTION_3_AVAILABLE = True - sage_attention_available = True logging.info("Found SageAttention3 (sageattn package)") - - except ImportError: - try: - from sageattention import sageattn - sage_attention_available = True - logging.info("Found SageAttention2 (sageattention package)") - except ModuleNotFoundError: - pass - - if not sage_attention_available: - 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") + 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(): @@ -504,39 +504,52 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= mask = mask.unsqueeze(1) try: - if SAGE_ATTENTION_3_AVAILABLE and dim_head < 256: - # 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) - elif not SAGE_ATTENTION_3_AVAILABLE: - # Fall back to SageAttention2 if available - out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_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.reshape(b, -1, heads * dim_head) + out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout) + except Exception as e: + logging.error("Error running sage attention: {}, 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) + 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: - # SageAttention3 is available but head_dim >= 256, fall back to pytorch + out = out.reshape(b, -1, heads * dim_head) + 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( @@ -544,8 +557,26 @@ 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) + + # 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: {}, using pytorch attention instead.".format(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), @@ -614,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") From 5e1865c1c7cc75601b6c4cfbb2ebd717df306f9a Mon Sep 17 00:00:00 2001 From: Panchovix Date: Sat, 26 Jul 2025 14:21:48 -0400 Subject: [PATCH 8/9] Separate sage 1.x/2.x and sage 3.x on different functions and flags 2. --- comfy/cli_args.py | 1 + 1 file changed, 1 insertion(+) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 0d760d524..e511cee60 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.") From 3c3009c02037685fe73f3f22ed9a46e53da55289 Mon Sep 17 00:00:00 2001 From: Panchovix Date: Sat, 26 Jul 2025 14:22:22 -0400 Subject: [PATCH 9/9] Separate sage 1.x/2.x and sage 3.x on different functions and flags 3. --- comfy/model_management.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/comfy/model_management.py b/comfy/model_management.py index 232d363aa..1f6445c9f 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1044,6 +1044,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