解析)
簡介阿里Wan2.1開源視頻大模型的TeaCache加速方案代碼包面向使用ComfyUI的開發(fā)者與AI視頻創(chuàng)作者用于解決高分辨率長視頻生成速度慢、等待時間長的問題。壓縮包體積僅6KB包含3個文件index.html用來展示加速方案說明.inscode定義在線運行環(huán)境.gitignore便于代碼版本管理整體結(jié)構(gòu)精簡適合快速部署與二次開發(fā)。該方案通過ComfyUI-WanVideoWrapper插件實現(xiàn)無需額外下載模型只要在現(xiàn)有工作流中添加一個節(jié)點即可啟用TeaCache官方測試數(shù)據(jù)表明在720×480分辨率、81幀視頻生成場景下10步迭代速度提升15%20步迭代提升接近40%。此外rel_l1_thresh與加速起始步數(shù)兩個參數(shù)可靈活調(diào)節(jié)幫助在不同硬件條件下平衡生成速度與畫質(zhì)且對最終成片質(zhì)量影響極小。目前已有117人學(xué)習(xí)下載適合希望在保持Wan2.1原有畫質(zhì)的同時大幅縮短出片時間的創(chuàng)作者也便于對視頻加速算法感興趣的程序員直接閱讀源碼、參與社區(qū)改進。1. 視頻生成慢在哪兒Wan2.1 視頻加速這套方案值不值得折騰在 14B 參數(shù)的 Wan2.1 視頻生成模型上默認 fp16 推理生成一段 81 幀 480p 的視頻24GB 顯卡要跑差不多十分鐘而這套代碼包的組合拳——FP8 動態(tài)量化、注意力 kernel 替換、CUDA Graph 捕獲——能把單次生成壓到兩三分鐘顯存從逼近 OOM 的 25GB 掉到 15GB 左右。它不訓(xùn)練新模型也不改網(wǎng)絡(luò)結(jié)構(gòu)只是給現(xiàn)成的 Wan2.1 推理鏈路打一套工程化加速補丁入口是幾個帶開關(guān)的腳本和一段推理循環(huán)改造解壓就能認目錄。適合兩類人折騰這套方案批量出片、每一分鐘渲染時長都算成本的從業(yè)者和卡在 24GB 顯存邊緣、總在 fp16 下被 OOM 打斷的開發(fā)者。下面按我實際拆包、復(fù)現(xiàn)、壓測的順序把瓶頸、落地步驟和坑一次講清楚。2. 瓶頸拆解Wan2.1 的時間燒在三層結(jié)構(gòu)上Wan2.1 的一次視頻生成不是模型跑一遍這么簡單而是三段串行鏈路文本條件編碼、DiT 主干迭代去噪、3D VAE 解碼。要加速先得知道時間燒在哪一段否則就是瞎調(diào)。我拆包時第一件事就是給整條鏈路做 profiling拿數(shù)據(jù)說話。2.1 三層耗時結(jié)構(gòu)文本編碼、DiT 主干、3D VAE第一層是文本編碼。Wan2.1 把 prompt 交給多語言文本編碼器產(chǎn)出一組條件和負條件向量這一段通常幾百毫秒到一兩秒。它占比很小但決定了視頻內(nèi)容對不對題所以加速方案里一般不碰它保持原精度跑。第二層是 DiT 主干這是絕對的大頭。每去噪一步模型都要把 latent 和條件向量過一次完整的 Transformer 前向14B 模型單步前向就要 610 秒取決于分辨率和顯卡。按官方示例推薦 50 步保證質(zhì)量、社區(qū)常用 2030 步折中來算這一層占了總耗時的 90% 以上。所有加速手段的八成功夫都應(yīng)該花在這層。第三層是 3D VAE 解碼把壓縮后的 latent 還原成像素視頻。它的耗時和幀數(shù)成正比比 DiT 小一個量級但容易被忽略——解碼在去噪結(jié)束后才做幀數(shù)一多這層的等待同樣讓人焦慮。VAE 端常見的優(yōu)化是分塊并行解碼把一次大解碼拆成多個 tile 前后的流水。動手之前先把耗時分布打出來避免憑感覺優(yōu)化。代碼包里一般會帶一個 profiling 入口核心就是 torch profiler# 給一次完整生成做耗時分解確認 DiT 是不是唯一瓶頸 from torch.profiler import profile, ProfilerActivity # run_generate 換成代碼包里的生成入口固定 81 幀、固定分辨率 with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue, ) as prof: run_generate( prompt廚房灶臺上的一碗熱湯冒著白氣, frames81, ) print(prof.key_averages().table(sort_bycuda_time_total, row_limit15))輸出表格里耗時排前面的應(yīng)該是一串 transformer block 的 forward而不是文本編碼或 VAE。如果你打出來的結(jié)果里 VAE 解碼占比反而很高說明 kernel 沒吃滿或者顯卡架構(gòu)偏舊那屬于另一類問題第 4 章的注意力替換未必適合你。這個判斷決定了你該不該繼續(xù)往下走。2.2 為什么優(yōu)先量化 DiT 主干而不是動 VAE 或文本編碼器凡是做量化第一反應(yīng)是哪里最重就量化哪里。14B 模型 fp16 權(quán)重就要 28GB訪存減半意味著每一步去噪的權(quán)重讀取時間近似減半——DiT 是典型的 memory-bound 場景量化權(quán)重是收益最直接、風(fēng)險最小的第一刀。VAE 和文本編碼器加起來不到 4GB量化它們省不了多少顯存但 VAE 對量化誤差極其敏感解碼時逐幀放大誤差畫面直接花給你看。所以這個代碼包的量化開關(guān)默認只作用于 DiT 主干的 Linear 層這是有意的設(shè)計不是遺漏。另一個常識是別一上來就 INT4。視頻生成模型的輸出對高頻細節(jié)很敏感手指、文字、快速運動的物體邊緣哪一處都經(jīng)不起低比特量化的折騰。INT4 的誤差在 DiT 每一步都會被放大一點積累幾十步之后畫質(zhì)劣化明顯。FP8 處在誤差和收益的甜點上顯存省一半質(zhì)量損失在人工目測層面基本察覺不到。這解釋了為什么代碼包提供的是 fp8 量化開關(guān)而不是更激進的低比特方案。2.3 三個原理性抓手量化、注意力重寫、CUDA Graph整套加速拆到底就是三件事下面這張表是我拆包時整理的對應(yīng)關(guān)系抓手作用對象主要收益代價與限制FP8 動態(tài)量化DiT 權(quán)重與激活權(quán)重訪存減半顯存降約 50%吞吐提升 30%80%需要 Ampere 及以上架構(gòu)首次加載有轉(zhuǎn)換開銷注意力 kernel 替換DiT 內(nèi)的自注意力整體速度提升 10%25%SM80低分辨率短序列收益縮水CUDA Graph 捕獲去噪循環(huán)整體減少 CPU 內(nèi)核啟動開銷長幀收益 10%30%輸入形狀必須固定捕獲時預(yù)留顯存三個抓手的作用對象不重疊所以可以疊加。血淚經(jīng)驗是工程加速先做這三樣不要一上來就折騰采樣步數(shù)、換蒸餾模型——那是模型層面的改動風(fēng)險高、可回退性差。先把手里的 fp16 權(quán)重加速到 35 倍再考慮要不要動模型本身。3. 環(huán)境與模型準備先把 fp16 基線跑通再談加速拿到代碼包第一反應(yīng)是直接跑加速腳本我建議反過來先把原版 fp16 推理跑通、記下耗時和顯存峰值再開加速開關(guān)。沒有基線后面所有快了 3 倍都是自欺欺人。3.1 硬件門檻與顯存預(yù)算代碼包的加速手段吃顯卡架構(gòu)先對一下硬件再動手。下表是我在一張 Ampere 架構(gòu) 24GB 卡上跑過的參考值不同驅(qū)動和分辨率會有 23GB 浮動模型檔位fp16 顯存參考FP8 顯存參考建議起點分辨率Wan2.1-T2V-1.3B約 10GB約 7GB832×48081 幀Wan2.1-T2V-14B約 30GB 起步約 18GB832×48081 幀1.3B 檔適合顯存吃緊時先驗證流程真要出片還是得 14BFP8 之后 24GB 卡才能舒服地跑 81 幀。如果卡只有 16GB老實把分辨率降到 480×832 或者把幀數(shù)降到 61別硬頂。3.2 三份權(quán)重與依賴安裝Wan2.1 提速不需要重新下載模型但需要三部分權(quán)重各就各位DiT 主干、3D VAE、文本編碼器。代碼包一般要求三個目錄分開傳原因是加速只作用于主干VAE 和編碼器要原精度加載# 代碼包解壓后的標準目錄結(jié)構(gòu)示意 weights/ ├── Wan2.1-T2V-14B/ # DiT 主干config.json diffusion_pytorch_model.safetensors ├── Wan2.1-VAE/ # 3D VAE通常和主干分開存放 └── Wan2.1-T2V-TextEncoder/ # 多語言文本編碼器 # 創(chuàng)建環(huán)境并裝依賴requirements.txt 是代碼包自帶的一版 conda create -n wan2_1 python3.10 -y conda activate wan2_1 pip install -r requirements.txt依賴里要特別注意 torchao 和注意力庫這兩個。torchao 版本太舊會找不到 fp8 量化接口注意力庫如果要求從源碼編譯先確認系統(tǒng)里有匹配的 CUDA 工具鏈編譯報錯九成是這個原因。我習(xí)慣裝完后先跑一段 21 幀的短片段驗證環(huán)境幾十秒就能確認依賴是否齊全翻車成本低。3.3 用 fp16 基線定出對比的時間錨點環(huán)境就緒后先跑原版 fp16。代碼包如果保留了原始入口就按原始參數(shù)跑如果不保留就用它提供的 baseline 模式# 記錄基線不開任何加速開關(guān)固定 prompt、幀數(shù)、分辨率 python scripts/run_wan_t2v.py \ --model_dir ./weights/Wan2.1-T2V-14B \ --vae_dir ./weights/Wan2.1-VAE \ --text_encoder_dir ./weights/Wan2.1-T2V-TextEncoder \ --prompt 黃昏的街道行人撐傘走過積水路面 \ --frames 81 \ --size 832x480 \ --output ./out/fp16_baseline.mp4跑之前先nvidia-smi --query-gpumemory.used --formatcsv記一個空閑顯存跑完再記一次峰值時間用 shell 的time包一層就夠了。唯一硬要求基線 prompt、幀數(shù)、分辨率要和后面加速測試完全一致否則對比沒有意義。參數(shù)說明--frames 81對應(yīng)約 3.4 秒 24fps 的視頻3D VAE 會把 latent 時間維壓 4 倍81 幀進 DiT 的 latent 序列長度是 (81-1)/4121。這個數(shù)直接決定第 4 章 CUDA Graph 的靜態(tài)形狀后面還會用到。如果基線階段已經(jīng) OOM先別碰加速優(yōu)先解決驅(qū)動、依賴和分辨率適配。提示基線的 prompt、幀數(shù)、分辨率必須和后續(xù)加速測試保持一致否則你對比出來的快了多少只是環(huán)境差異。4. 加速落地FP8 量化、注意力替換、CUDA Graph 的組合拳這一章是代碼包的核心。加速腳本的入口一般是一個帶開關(guān)的生成命令但理解開關(guān)背后做了什么才能知道出問題時去哪排查。4.1 第一步FP8 動態(tài)量化先吃最大的一塊收益FP8 量化在整套方案里收益最穩(wěn)、風(fēng)險最低所以放第一步。代碼包里對應(yīng)的核心腳本是一個量化包裝函數(shù)# accel/fp8_quantize.py —— 只量化 DiT 主干的 Linear 層 from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight def apply_fp8(dit_model): # dit_model 是 Wan2.1 的 DiT 主干VAE 和文本編碼器不傳進來 for name, module in dit_model.named_modules(): if isinstance(module, torch.nn.Linear): # W8A8 動態(tài)量化激活和權(quán)重都在前向時轉(zhuǎn) fp8scale 動態(tài)計算 quantize_(module, float8_dynamic_activation_float8_weight) return dit_model邏輯說明quantize_會把模塊權(quán)重替換成帶量化邏輯的包裝層float8_dynamic_activation_float8_weight表示權(quán)重和激活都轉(zhuǎn) fp8activation 的 scale 在每次前向時按當前輸入分布動態(tài)算出所以不需要校準集。這是它和 GPTQ 那類后訓(xùn)練量化最大的區(qū)別——代碼包不需要你準備校準數(shù)據(jù)開箱即用。參數(shù)說明如果想更省顯存可以換成int8_dynamic_activation_int8_weight但 int8 的動態(tài)范圍對視頻中間特征的覆蓋不如 fp8會有輕微畫質(zhì)劣化不推薦反過來如果只求省顯存不求吞吐float8_weight_only也可以但訪存加速有限。代碼包默認的 W8A8 是權(quán)衡后的選擇不建議動。這一層做完14B 模型的權(quán)重從 28GB 降到 14GB24GB 卡第一次有了跑長視頻的余量。4.2 第二步把注意力 kernel 替換成融合實現(xiàn)量化解決訪存注意力替換解決計算。Wan2.1 原始 attention 是通用實現(xiàn)對 mask 的處理和 kernel 利用率都有優(yōu)化空間。代碼包的--attn sage開關(guān)背后做的是模塊替換# accel/attention_patch.py —— 用融合注意力替換原始點乘分支 import torch import torch.nn as nn class ChannelAttnBridge(nn.Module): 把 q/k/v 從 (batch, seq, dim) 重排成融合 kernel 需要的布局 def __init__(self, dim, heads): super().__init__() self.heads heads self.head_dim dim // heads self.q_proj nn.Linear(dim, dim) self.k_proj nn.Linear(dim, dim) self.v_proj nn.Linear(dim, dim) self.o_proj nn.Linear(dim, dim) def forward(self, x): b, s, _ x.shape q self.q_proj(x).view(b, s, self.heads, self.head_dim) k self.k_proj(x).view(b, s, self.heads, self.head_dim) v self.v_proj(x).view(b, s, self.heads, self.head_dim) # 替換點這里調(diào)用 SageAttention/FlashAttention 的融合算子 # 原始實現(xiàn)是手寫 softmax(q k^T / sqrt(d)) v矩陣乘與 softmax 分離 out fused_attention(q, k, v, is_causalFalse) # 分塊 片上 softmax return self.o_proj(out.reshape(b, s, -1))邏輯說明融合注意力把 QK^T、softmax、PV 三段操作合成一個 kernel避免中間矩陣寫回顯存同時用分塊策略降低顯存占用。原始實現(xiàn)每一步都要分配 (seq, seq) 的注意力矩陣序列一長就很可觀——81 幀的 latent 是 21×60×104token 數(shù) 13 萬。參數(shù)說明is_causalFalse對應(yīng)視頻生成的雙向注意力圖像生成場景才需要 causal。SageAttention 這類庫還支持 fp8 的注意力計算配合 4.1 的 W8A8q/k/v 的訪存還能再降一檔。但注意權(quán)重量化和注意力量化是兩個獨立開關(guān)別混在一起調(diào)出了問題不好定位。4.3 第三步CUDA Graph 去掉去噪循環(huán)里的 CPU 調(diào)度前兩步之后單步去噪已經(jīng)降下來了但循環(huán)里每次都有上千次小 kernel 啟動CPU 啟動開銷占了不少。CUDA Graph 的思路是把整個去噪步的 kernel 序列捕獲成一張圖之后重放時不經(jīng)過 CPU 逐個 launch# accel/cuda_graph.py —— 固定形狀下的推理圖捕獲與重放 class WanDiTGraph: def __init__(self, model, latent_shape): # latent_shape 必須固定例如 (1, 21, 16, 60, 104) self.model model self.static_latent torch.zeros(latent_shape, dtypetorch.float16, devicecuda) # 預(yù)熱 2 次讓 autotune 的 kernel 選擇穩(wěn)定否則捕獲到的是次優(yōu)路徑 for _ in range(2): self.model(self.static_latent) torch.cuda.synchronize() # 捕獲整張去噪前向圖固化下來 self.graph torch.cuda.CUDAGraph() with torch.cuda.graph(self.graph): self.static_out self.model(self.static_latent) def run(self, latent): # 新輸入靠 copy 進靜態(tài)緩沖再重放整張圖 self.static_latent.copy_(latent) self.graph.replay() return self.static_out.clone()邏輯說明捕獲期間的顯存分配會進入 graph 的私有池所以捕獲時的峰值顯存被永久占住這就是開 CUDA Graph 顯存明顯上漲的原因。折中辦法是按幀數(shù)分檔建圖21 幀一張、81 幀一張、161 幀一張進程內(nèi)切換檔位時加載對應(yīng)圖每來一個尺寸就建一張捕獲開銷會吃光收益。參數(shù)說明latent_shape的四個數(shù)字來自 3.3 節(jié)的推導(dǎo)(batch, latent_frames, 16, h/8, w/8)。分辨率一變形狀就變圖就失效。所以代碼包一般把 CUDA Graph 和固定分辨率綁定分辨率可變的場景直接忽略這一步它本來就是三招里收益最小的。4.4 開關(guān)組合與預(yù)期收益三招全部落地后代碼包的入口通常長這樣# 全量加速fp8 sage attention cuda graph 一次開啟 python scripts/run_wan_t2v.py \ --model_dir ./weights/Wan2.1-T2V-14B \ --vae_dir ./weights/Wan2.1-VAE \ --text_encoder_dir ./weights/Wan2.1-T2V-TextEncoder \ --fp8 --attn sage --graph \ --prompt 黃昏的街道行人撐傘走過積水路面 \ --frames 81 --size 832x480 \ --output ./out/wan_accel.mp4我習(xí)慣的順序是先--fp8單獨跑一遍看畫面有沒有異常再加--attn sage最后才加--graph。三個開關(guān)一次全開出了問題根本不知道是哪一步引入的。全開之后的收益參考組合端到端耗時參考顯存參考適用場景僅 fp8約為原版 50%60%減半左右任何場景先開這個fp8 sage約為原版 40%50%與僅 fp8 接近Ampere 以上序列較長三開關(guān)全開約為原版 30%40%比僅 fp8 高一截固定分辨率批量出片這里的百分比是基于我拆包測試的相對值不同顯卡的參考意義大于絕對值。核心結(jié)論是三招收益遞減、風(fēng)險遞增按順序開。5. 避坑記錄五個現(xiàn)場按現(xiàn)象、原因、解決排查加速方案翻車幾乎都集中在五個地方下面按現(xiàn)象、原因、解決三個層次記錄都是我實際調(diào)試時踩過的。5.1 量化后畫面出現(xiàn)彩色噪點現(xiàn)象FP8 開關(guān)打開后前幾幀正常中間開始飄彩色豎條紋和雪花噪點越往后越嚴重整個視頻沒法看。原因我的第一版腳本把整個生成鏈路包括 VAE 一起量化了。VAE 解碼時對權(quán)重的 scale 誤差遠比 DiT 敏感量化誤差在逐幀解碼里被放大成噪點。解決把 VAE 排除在量化范圍之外。代碼包apply_fp8只接收 DiT 主干就是為這個如果你自己改腳本確認 VAE 的加載不在quantize_遍歷范圍里。修完后噪點消失畫面恢復(fù)干凈。5.2 CUDA Graph 開了之后顯存暴漲甚至 OOM現(xiàn)象只加--graph顯存從 14GB 一下跳到 22GB跑 161 幀直接 out of memory。原因CUDA Graph 捕獲時的全部中間緩沖都會留在 graph 私有池里不會像普通推理那樣隨用隨釋放。幀數(shù)越長捕獲峰值越大這部分顯存是永久占用的。解決先算預(yù)算再開圖給進程設(shè)一個顯存上限避免分配抖動# 開 CUDA Graph 前給進程設(shè)硬上限預(yù)留抖動空間 torch.cuda.set_per_process_memory_fraction(0.92)幀數(shù)超預(yù)算就把--frames從 161 降到 81latent 序列長度從 41 變回 21graph 的靜態(tài)緩沖立刻小一塊。注意幀數(shù)檔位和圖是綁定的改幀數(shù)必須重建圖不能復(fù)用。5.3 三招全開反而更慢現(xiàn)象21 幀、480p 的小任務(wù)全開加速后比原版 fp16 還慢 10%20%心態(tài)直接崩。原因短序列下 token 數(shù)少融合注意力的啟動開銷和 graph 的靜態(tài)緩沖拷貝開銷蓋過了計算收益量化 kernel 對小矩陣也不占優(yōu)。解決給開關(guān)加閾值邏輯。我一般會設(shè)一個經(jīng)驗值latent token 數(shù)低于某個量級時只保留 fp8自動跳過注意力替換和 CUDA Graphtoken 數(shù)上來之后再全開。這個閾值在代碼包配置里可以調(diào)不必改源碼。5.4 FP8 結(jié)果和 fp16 對不上懷疑權(quán)重被量化損壞現(xiàn)象同 seed、同 promptFP8 的輸出整體偏亮、邊緣偏柔和 fp16 逐像素對不上。原因這不是損壞。W8A8 動態(tài)量化的激活 scale 每步按當前輸入分布現(xiàn)算誤差是輸入相關(guān)的畫面里有人臉、文字這類高頻信息時差異更明顯屬于量化固有代價。解決對比測試別用逐像素誤差用 SSIM 或者直接人工判斷。如果某個品類場景對質(zhì)量敏感就在代碼包里給該場景單獨配 fp16其他走 fp8。量化不是不能用是要給敏感場景留后門。5.5 第一次生成特別慢以為是加速失效現(xiàn)象改動后第一次跑等了三四分鐘才出第一幀比之前的版本還慢以為整個方案失效了。原因首次運行在編譯 kernel、autotune、捕獲 CUDA Graph全是一次性開銷。第二次跑同一配置就恢復(fù)正常。解決把預(yù)熱做成獨立步驟。代碼包有的版本會把捕獲好的圖緩存到本地目錄第二次啟動直接加載跳過預(yù)熱如果緩存沒有持久化每次冷啟動就當它是固定成本。這個預(yù)熱時間不計入出片時間單獨跑一次能省很多困惑。6. 進階驗證把加速關(guān)進質(zhì)量籠子里加速做完不驗證質(zhì)量等于把模型當黑匣子用。我給自己定了一套回歸流程花不了十分鐘但能擋住絕大多數(shù)翻車。做法是固定三個測試 prompt 分類一個特寫畫面檢查手指和皮膚邊緣一個文字畫面檢查文字畸變一個快速運動畫面檢查運動模糊和閃爍。每個分類在加速前后各出一版同一 seed 下對比。對比手段有兩個機器指標用 SSIM人工指標直接看視頻連貫性。SSIM 不低于 0.95 是寬松底線低于 0.9 就需要具體看是哪類畫面劣化。不要迷信 SSIM它對運動視頻的拖影不敏感所以人工抽幀看邊緣才是最終判據(jù)。另一個我常用的技巧是對比 latent 誤差而不是像素誤差。在去噪最后一步前把 latent 存下來對比 fp16 和 fp8 兩版的 L2 距離這個數(shù)能提前告訴你誤差是去噪過程中積累的還是 VAE 解碼端引入的檢查點檢查內(nèi)容合格參考去噪最后一步 latent L2量化誤差是否在 DiT 內(nèi)失控相對誤差 2%VAE 解碼后 SSIM解碼端是否引入額外劣化≥ 0.95人工抽幀手指、文字、運動邊緣無明顯畸變從那以后我每次調(diào)加速參數(shù)無論多小都強制走一遍這套回歸三個 prompt、一版 baseline、一版加速latent L2 和 SSIM 各記一次最后一次人工過片。這個習(xí)慣救過我很多次——看起來快了兩倍的配置可能只是把誤差壓進了畫面里最不明顯的地方等放大一幀才發(fā)現(xiàn)全毀了。這套代碼包的入口腳本名可能和這里的參數(shù)略有出入但對照第 4 章的開關(guān)順序十分鐘就能把整套鏈路跑起來。希望幫到你。本文還有配套的精品資源點擊獲取