Fixed Forge for Sage-Attention-2++

2025年7月24日 00:33

2025年9月3日、修正ファイルにpreload.pyを追加しました。

序

現状、公式ではComfyUIとForge-Classicに実装されているSage-Attentionのロード機能ですが、今回Forgeへの実装と一定の高速化に成功しました。

尚、以下はこの記事の適用が前提です。

SA2++のインストールについては、以下で解説しています。以下記事はPortable環境ComfyUIのPython3.12前提ですが、Python3.11venv環境のForgeでも、基本は同じです。単にPython3.11用のwhlを使う也、自力でビルドする也すれば良いだけの話です。

また、繰り返しますがForgeがpython3.10.6でなければ動かない等は絶対にありません。私はPython3.11.13で使用しています。

起動オプションは、以下の様にxformersとの共存が可能です。

--xformers --use-sage-attention

以下の様に、t2i画面にSA適用のON/OFFボタンを設置しました。SD1.5は、SAに対応しません。また、Flux1とSDXL(IL、Pony含む)の2種類のみを想定しての開発なので、SD3系や他のモデルでの動作は不明です。

また、私はForgeではi2iは一切使用しない(A1111一択)為、i2iでの動作は検証していません。

チェックを付けた状態では、以下の様にSAが適用されます。

チェックを外すと、以下の様にFlash-Attention-3(xformers 0.0.31 post1)が適用されます。

Flux1.devの場合、SAのON/OFFで相当速度差があります。FA3と比較すると2倍近く高速になりますね。これは何気に凄い成果です。激重いFlux1で2倍高速化は大きいですよ。

以下、SageAttention（SA）改造で修正した全ファイル・全箇所・前後対比・解説をまとめます。

尚、前回大絶賛したCursorですが、

今回は余りのバカさ加減にブチキレながらの作業になりました。一歩ずつ障害に対してデバッグしつつ進めていた訳ですが、以前の正しいプロセスを平然とド忘れして、せっかく作ったコードをぶち壊して余計な事してロールバックの為に余計な手間かけてくれやがったので、AI相手だから遠慮なくマジ切れしましたよ。

数時間は無駄にしてます。

「てめえマジでざっけんな、金取っておいてふざけた真似しくさってんじゃねえ。以前のログをちゃんと見ろ、バカかてめえは、記憶力がねえのか。てめえと違って人間様は1秒ずつ死に近づいてんだ。てめえが浪費させられた俺の命を弁償できんのか」て感じでキレましたね。

…

修正ファイル一覧

backend/attention.py

modules/ui.py

modules/cmd_args.py

modules/processing.py

extensions-builtin/sd_forge_lora/preload.py

extensions-builtin/sd_forge_controlnet/preload.py

修正内容

1. backend/attention.py

変更前

def attention_dispatch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
    if getattr(cmd_opts, 'xformers', False):
        return attention_xformers(q, k, v, heads, mask, attn_precision, skip_reshape)
    else:
        return attention_pytorch(q, k, v, heads, mask, attn_precision, skip_reshape)

SageAttention分岐もログ出力も一切なし

変更後（SA分岐・1回ごとのログ出力を追加）

# attention_dispatch関数を追加（603行目〜617行目）
def attention_dispatch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
    from modules.shared import cmd_opts
    if getattr(cmd_opts, 'USE_SAGE_ATTENTION', False):
        print('[SageAttention] SageAttention enabled (HND layout, version: unknown)')
        return attention_sage(q, k, v, heads, mask, attn_precision, skip_reshape)
    elif getattr(cmd_opts, 'xformers', False):
        print('[Forge] xformers.memory_efficient_attention called (FA-3 or best available kernel will be used)')
        return attention_xformers(q, k, v, heads, mask, attn_precision, skip_reshape)
    else:
        return attention_pytorch(q, k, v, heads, mask, attn_precision, skip_reshape)

# attention_functionの設定を変更（622行目）
attention_function = attention_dispatch  # ← 分岐処理を使用

# tensor layout修正（72行目〜84行目）
try:
    out = sageattn(q, k, v, tensor_layout="BSHD", is_causal=False)
except Exception as e:
    # フォールバック処理解決した問題: SAボタンがONでも常にxformersが使用される問題

解決した問題:
SAボタンがONでも常にxformersが使用される問題

2. modules/ui.py

## UIの改造まとめ

### 修正したファイル

- `modules/ui.py`

### 主な変更内容

# 正しいcmd_optsのインポートを使用（22行目で既に定義済み）
from modules.shared import cmd_opts

# SAボタンの実装（285行目〜）
if '--use-sage-attention' in sys.argv:
    cmd_opts.USE_SAGE_ATTENTION = True
initial_sage_attention = getattr(cmd_opts, 'USE_SAGE_ATTENTION', False)
sageattention_toggle = gr.Checkbox(label="SageAttention有効化", value=initial_sage_attention)

def set_sageattention(val):
    cmd_opts.USE_SAGE_ATTENTION = val
    print(f"[UI] set_sageattention: cmd_opts.USE_SAGE_ATTENTION={cmd_opts.USE_SAGE_ATTENTION}")
    return val
sageattention_toggle.change(set_sageattention, inputs=[sageattention_toggle], outputs=[sageattention_toggle])

#### 1. **UI初期化時のコマンドライン引数反映**

1. **初期化時の設定不反映**: `--use-sage-attention`コマンドライン引数が指定されていても、UIの初期化時に`cmd_opts.USE_SAGE_ATTENTION`が`False`になっていた問題

2. **トグル操作が必要**: 最初の生成でSageAttentionが効かず、トグルをOFF→ONにする必要があった問題

### 修正のポイント

- **二重チェック**: UI初期化時とUI初期化完了時の両方でコマンドライン引数を確認

- **確実な反映**: `sys.argv`を直接チェックして`cmd_opts.USE_SAGE_ATTENTION`を設定

- **デバッグ出力**: 設定の状態をログに出力して確認可能

これで、`--use-sage-attention`を指定してWebUIを起動すると、最初の生成からSageAttentionが有効になり、UIのトグルも正しくONの状態で表示されるようになりました！

3. modules/cmd_args.py

変更前（本当に最初の状態）

SA関連の記述は一切なし

変更後（SA対応のための追加）

parser.add_argument("--use-sage-attention", action='store_true', help="enable SageAttention for attention layers")
USE_SAGE_ATTENTION = getattr(args, 'use_sage_attention', False)

コマンドライン引数 `--use-sage-attention` の追加
グローバル変数 `USE_SAGE_ATTENTION` の定義
他ファイルで `cmd_args.USE_SAGE_ATTENTION` を参照し、attention_dispatchの分岐に利用

4. modules/processing.py

変更前（本当に最初の状態）

reset_sage_log/reset_fa3_log などの呼び出しも一切なし

変更後（1生成ごとにログフラグをリセット）

from backend.attention import reset_fa3_log, reset_sage_log
reset_fa3_log()
reset_sage_log()

1生成ごとにログフラグをリセットし、毎回必ず1回だけログが出るようにした

5.extensions-builtin/sd_forge_lora/preload.py

cmd_args.py記述と重複する定義を削除

def preload(parser):
    # Removed --lora-dir argument as it's already defined in modules/cmd_args.py
    # Removed --lyco-dir-backcompat argument as it conflicts with other extensions
    pass

6.extensions-builtin/sd_forge_controlnet/preload.py

目的:
--lyco-dir-backcompatの重複定義を削除

def preload(parser):
    # Removed --lora-dir argument as it's already defined in modules/cmd_args.py
    # Removed --lyco-dir-backcompat argument as it conflicts with other extensions
    pass

環境変数

本来、不要な理屈ですが、現実には以下の環境変数設定が必要になりました。具体的には起動バッチファイルに追記する形です。当然、パスは人それぞれですが。

set INCLUDE=C:\Users\ussoe\AppData\Local\Programs\Python\Python311\include;%INCLUDE%
set LIB=C:\Users\ussoe\AppData\Local\Programs\Python\Python311\libs;%LIB%

尚、上を見て勘の良い人は気づくでしょうが、Embeded Python環境下のComfyUIでtriton-windowsを使用する場合、「見に行く場所が完全に違う」という事です。だから、手動でこの2つのフォルダをpython_embededフォルダに置く必要がある訳です。ComfyUIの場合、環境変数指定しなくても拾いに行ってくれるんですがね…tritonは。

その意味ではForgeが、環境変数で拾わせてやんないとtritonが環境構築してくんねえ…てのは意外でした。まさか、こんな処でコケると思わなかったのですよ。リアタイで環境構築やってる際には。

Forgeそのものは、venv環境で動作させていたから余計に盲点だった…事は事実です。上のパスは一目瞭然グローバル環境に対するパスです。

…

「環境変数（INCLUDE, LIB）」は、PythonやPyTorch、xformersなどのビルドや一部の拡張機能の動作に必須となる場合があるという点について、まとめます。

Windows環境での環境変数（INCLUDE, LIB）設定について

1. 変更前（本当に最初の状態）

INCLUDEやLIBなどの環境変数は特に設定していない
PythonやPyTorch、xformersのビルド・インストール時にエラーが出ることがある
特にxformersや一部のC/C++拡張を使う場合、
「ヘッダファイルが見つからない」「ライブラリが見つからない」等のビルドエラーが発生することがある

2. 変更後（環境変数を明示的に設定）

これにより、xformersやPyTorchの一部機能、Forgeのattention分岐などが
「ビルド失敗」や「DLLロード失敗」などのエラーなしで正常に動作するようになる
特にWindows環境では、これらの環境変数が正しく設定されていないと
SA/FA3の切り替えやログ出力以前に、拡張機能自体が動作しない場合がある

まとめ

Windows環境でForgeやxformers、SageAttention等を使う場合、
「INCLUDE」「LIB」環境変数にPythonのパスを追加することが必須となる場合がある
これを設定しないと、attention分岐やSA/FA3の切り替え以前に、
ビルド・インストール・実行時にエラーが発生することがある