Modal × B200 × SageAttention × ComfyUI (Wan2.2)

NVIDIA B200 (Blackwell, sm_100) で SageAttention 2.2.0 を実 engage させて ComfyUI 経由で Wan2.2 動画を生成するための、ビルド済み wheel・1 ファイルの Modal スクリプト・ベンチマーク結果のセット。

pip install sageattention をしただけだと B200 では silently fallback して PyTorch SDPA で動きます。ここでは:

SageAttention を B200 専用 PTX で wheel 化、
dispatcher の core.py に sm_100 分岐 を注入、
その wheel を Modal Volume に永続化 して ComfyUI ランタイムから読み込み

までを一気通貫でやります。

TL;DR

	min	mean	median	max
SDPA baseline	98.95s	102.86s	100.64s	108.98s
SageAttention (this wheel)	96.54s	99.26s	98.59s	102.65s
Δ (mean)		+3.6 %
Per-iteration	2.13 s/it
Per-iteration		2.07 s/it

設定: Wan2.2 5B TI2V / 1280×704 / 121 frames @24fps / 30 steps / cfg 5.0 / shift 8.0 / 3 seeds × 2 modes、1 つのウォーム B200 コンテナ内で連続実行。

1280×704 / 30 step では attention の比率はそこまで大きくないので速度差は **~3.6%**。VRAM 余裕がある B200 なら "engage 確認 + 安定動作" の価値の方が大きいです。長尺 / 高解像度 / 14B モデルで差はさらに広がります (未測)。

なぜ B200 で SageAttention を有効化するのが難しいのか

Blackwell (sm_100) が新しすぎて、依存スタックの どれか一つでも未対応だと silently fallback します。実体験した失敗パターン:

レイヤー	既知の壁
CUDA toolkit	`nvcc` が sm_100 を扱うには CUDA 12.8 以上。それ未満だと "unrecognized arch"
PyTorch	2.7.0+cu128 で初めて Blackwell の Tensor Core を fully 使える。古い wheel だと `torch.cuda.get_device_capability()` は `(10, 0)` を返すが kernel 側が未対応
Modal builder	`nvidia/cuda:12.8.1-devel-ubuntu22.04` がベース。`build-essential` には `g++` は入るが `clang` は別パッケージで、Modal 内 Python は clang でビルドされているため setuptools が clang++ を要求して link 失敗
SageAttention setup.py	拡張モジュール名は `_qattn_sm80` 等だが、これはアルゴリズム種別で target arch とは無関係。`TORCH_CUDA_ARCH_LIST=10.0` で各 .so に sm_100 PTX を埋められる
SageAttention dispatcher	`core.py` の `sageattn()` は `sm80/86/89/90/120/121` を hardcode で分岐し、それ以外は `raise ValueError("Unsupported CUDA architecture: sm100")`。これが ComfyUI 側で catch されて `pytorch attention instead` の log が出る
ComfyUI	`--use-sage-attention` 指定でも、SageAttention の例外を catch して silently SDPA に切替わる仕様。何が起きているのか分からないのが最大の罠
Modal volume mount	`git clone` 直後の `/opt/ComfyUI/models` には `put__here` の placeholder があり非空のままだとボリュームをマウントできない*。`rm -rf && mkdir` で空にしておく必要あり

どう工夫したか

1. Wheel ビルド時の patch — sm_100 dispatch を追加

ビルド中、git clone 後・pip wheel 前に sageattention/core.py の sageattn() 関数本体を sed/Python で書き換えます。新しい branch は sm_89/sm_120/sm_121 と同じ FP8 経路 (sageattn_qk_int8_pv_fp8_cuda, pv_accum_dtype="fp32+fp16") に B200 をルーティング:

elif arch == "sm100":  # Blackwell datacenter (B200)
    return sageattn_qk_int8_pv_fp8_cuda(
        q, k, v, tensor_layout=tensor_layout, is_causal=is_causal,
        qk_quant_gran="per_warp", sm_scale=sm_scale,
        return_lse=return_lse, pv_accum_dtype="fp32+fp16")

完全な diff は patches/sm100_dispatch.diff。

2. Build 時の環境変数

TORCH_CUDA_ARCH_LIST=10.0   # Blackwell datacenter のみ生成 (B200 は sm_100, B100 は sm_100、RTX 50 系は sm_120)
FORCE_CUDA=1                # nvcc が GPU を見つけられなくても compile を強行
MAX_JOBS=4                  # Modal builder の vCPU/RAM に合わせて控えめに (デフォの 32 だと OOM)
NVCC_APPEND_FLAGS=--threads 4
EXT_PARALLEL=2
CUDA_HOME=/usr/local/cuda

3. Build イメージに不足を補う

build_image = (
    modal.Image.from_registry("nvidia/cuda:12.8.1-devel-ubuntu22.04", add_python="3.11")
    .apt_install("git", "build-essential", "ninja-build", "ca-certificates", "wget",
                 "clang")   # ← link 段で clang++ を呼ぶので必須
    .pip_install("packaging", "wheel", "setuptools", "ninja", "build")
    .pip_install("torch==2.7.0", "torchvision==0.22.0", "torchaudio==2.7.0",
                 index_url="https://download.pytorch.org/whl/cu128")
)

4. ビルド検証 (これがないと黙ってフォールバックする問題に気付けない)

pip wheel 後、ビルド関数内でその場で:

pip install でローカルに入れて
inspect.getsource(c.sageattn) で 'sm100' in src を assert (パック後にも patch が残っているか)
torch.randn(device="cuda") + sageattn(...) で B200 上で forward を実行、shape/dtype を確認

しています。これで「engage できるか」が build と同時に分かる。

5. ComfyUI 連携

/opt/ComfyUI/models と /opt/ComfyUI/output を image build 時に rm -rf & mkdir してから Modal Volume をマウント。これで ComfyUI が起動時に必要なサブディレクトリを自動生成、HF からダウンロードしたモデルは /opt/ComfyUI/models/diffusion_models/*.safetensors 等の標準パスに置けば自動認識。

ランタイム関数では:

whl = sorted(glob.glob("/wheels/sageattention-*.whl"))[-1]
subprocess.check_call(["pip", "install", "--force-reinstall", "--no-deps", whl])
subprocess.Popen(["python", "main.py", "--listen", "0.0.0.0", "--port", "8188",
                  "--use-sage-attention", "--disable-auto-launch"], cwd="/opt/ComfyUI")

ComfyUI のログに Unsupported CUDA architecture: sm100 が 1 度も出ないこと が engage 確認の決定打。

中身

modal_B200_sageattetion_comfyUI/
├── README.md                                       ← この文書
├── wheels/
│   └── sageattention-2.2.0-cp311-cp311-linux_x86_64.whl   ← 主成果物 (17 MB)
├── patches/
│   └── sm100_dispatch.diff                         ← core.py への sm_100 分岐パッチ
├── scripts/
│   ├── comfy_b200.py                               ← wheel build + ComfyUI web server
│   ├── wan22_demo.py                               ← Wan2.2 モデル DL + 動画生成 + ベンチ
│   └── push_to_hf.py                               ← この repo を再アップする用
├── benchmarks/
│   ├── wan22_sage_vs_sdpa.json                     ← 全 6 run の生データ
│   ├── sample_sdpa.webp                            ← SDPA 出力サンプル
│   └── sample_sage.webp                            ← SageAttention 出力サンプル
├── videos/
│   ├── wan22_00001_.webp                           ← 1 回目 (SDPA fallback だった)
│   ├── wan22_00001_.json
│   ├── wan22_00002_.webp                           ← 2 回目 (SageAttention engaged)
│   └── wan22_00002_.json
└── logs/
    └── build_success.log                           ← B200 forward sanity 込みのビルドログ

使い方

0. 前提

Modal アカウント (pip install modal && modal token new)
Cloudflare R2 とかは不要、Modal Volume だけ
Hugging Face のモデルダウンロードに anonymous でアクセス可 (anon でも rate limit はあり)

1. ビルド

modal run scripts/comfy_b200.py::build_wheel
# 初回 ~10 min。`./wheels/sageattention-*.whl` と Modal Volume `sage-wheels` に保存される。
# 既にあるときは skip、上書きしたいなら --force。

2. ComfyUI を web server として起動

modal serve scripts/comfy_b200.py
# 印字された https://<...>.modal.run を開けば ComfyUI UI。
# モデルは `/opt/ComfyUI/models/...` 配下に置けば認識。
# Modal Volume `comfy-models` に置くなら:
modal volume put comfy-models <local-file> /diffusion_models/foo.safetensors

3. Wan2.2 5B TI2V 動画生成 (T2V)

# 一度だけモデルをボリュームに DL (合計 ~17 GB、~5-10 min)
modal run scripts/wan22_demo.py::download_models

# 生成
modal run scripts/wan22_demo.py::generate
modal run scripts/wan22_demo.py::generate --prompt "..." --seed 1234 --steps 30
# → ./videos/wan22_*.webp + 同名 .json (メタデータ)

4. SDPA vs SageAttention ベンチマーク

modal run scripts/wan22_demo.py::benchmark
modal run scripts/wan22_demo.py::benchmark --seeds 42,1234,7777 --steps 30
# → ./benchmarks/wan22_sage_vs_sdpa.json

既知の制約 / 注意点

対象 GPU は B200 (sm_100) のみ。B100 (同 sm_100) と RTX 50 系 (sm_120) は patch 不要、上流の対応あり。RTX 50 で使うなら TORCH_CUDA_ARCH_LIST=12.0 で再ビルド推奨。
ベンチ wall-clock には ComfyUI 起動 + Wan2.2 5B モデル load (~25s) が 含まれている 場合があります。pure attention の差はおそらく 5-8% / step。
SageAttention 2.2.0 ベース (HEAD d1a57a5)。upstream で sm_100 が正式サポートされたら本 patch は不要に。
14B モデル (Wan2.2-T2V-A14B 等) は未検証。VRAM 192 GB あれば動くはずだがメモリ管理に注意。

ライセンス / クレジット

SageAttention: thu-ml/SageAttention (Apache-2.0)
ComfyUI: comfyanonymous/ComfyUI (GPL-3.0)
Wan2.2 model: Wan-AI/Wan2.2-TI2V-5B (config / weights は HF 側のライセンスに従う)
Modal: modal.com (デプロイ環境)

本 repo の Modal スクリプト・patch・README は Apache-2.0 で公開。

Downloads last month: -

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support