| |
| """部署加速引擎（ref/gc25d-live-demo 的 FastStream：静态 KV 环形缓冲 + torch.compile + CUDA graph）在 480p / 704p 下的 |
| 每格耗时。耗时与权重数值无关，所以直接用包里的 v8 DMD 权重（ids 版 GcDiT）：只换 lat_hw / 滑窗长度 / 去噪步数。 |
| 稳态 = 缓存已满（n = cap）：每格 = 历史 K 旋转 1 次 + N 趟去噪 + 1 趟提交，另测 TAEHV 解 1 格（4 帧）。 |
| 不含：语义条件嵌入（每格 1 次，很小）、文本 K/V（按 prompt 缓存，换按键组合那格多 6–7 ms）、代码世界、网络。 |
| 只 import 包里的代码（PYTHONPATH 只放包，见 usage），不碰 maliang 的 gc25d。 |
| usage: |
| R=ref/gc25d-live-demo; CUDA_VISIBLE_DEVICES=7 PYTHONPATH=$R/code/gc25d:$R/code/arope:$R/code/diffsynth_src AROPE_ROOT=$R/code/arope \ |
| TAEHV_DIR=$R/code/taehv DIFFSYNTH_SKIP_DOWNLOAD=True TMPDIR=.tmp TORCHINDUCTOR_CACHE_DIR=.tmp/inductor TRITON_CACHE_DIR=.tmp/triton \ |
| .venv/bin/python scripts/bench_fast704.py --out outputs/bench_fast/run.json""" |
| import argparse |
| import json |
| import os |
| import statistics |
| import subprocess |
| import time |
|
|
| import torch |
|
|
| from codeworld.fast_live import FastStream |
| from codeworld.live import TinyStreamVAE, empty_wan_dit |
| from gcdistill.gc_dit import GcDiT, load_gc_ckpt |
|
|
| M = "/opt/dlami/nvme/zhiyangdeng/maliang" |
| R = f"{M}/ref/gc25d-live-demo" |
| RES = {"480p": (30, 52), "704p": (44, 80)} |
| T_GRID = (1000.0, 750.0, 500.0, 250.0) |
|
|
|
|
| def gpu_apps(): |
| """本卡上的计算进程（判断测的时候有没有别的进程在抢卡）。""" |
| try: |
| idx = os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0] |
| q = subprocess.run(["nvidia-smi", "-i", idx, "--query-compute-apps=pid,used_memory", "--format=csv,noheader"], |
| capture_output=True, text=True, timeout=30).stdout.strip() |
| return [ln for ln in q.splitlines() if ln and str(os.getpid()) not in ln] |
| except Exception as e: |
| return [f"nvidia-smi failed: {e}"] |
|
|
|
|
| def cell_ms(fs, x, steps, n_cells): |
| out = [] |
| for _ in range(n_cells): |
| torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| fs._rotate() |
| for t in T_GRID[:steps]: |
| fs.forward(x, t, False) |
| fs.forward(x, 0.0, True) |
| torch.cuda.synchronize() |
| out.append((time.perf_counter() - t0) * 1e3) |
| return out |
|
|
|
|
| def summ(v): |
| v = sorted(v) |
| return {"median": round(statistics.median(v), 1), "p90": round(v[int(0.9 * (len(v) - 1))], 1), "min": round(v[0], 1), |
| "n": len(v)} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default=f"{R}/weights/gc25d_v8_dmd_ema500.safetensors") |
| ap.add_argument("--tae", default=f"{R}/weights/taew2_2_super.pth") |
| ap.add_argument("--configs", default="480p:17:4,704p:17:4,704p:17:3,704p:11:4,704p:11:3", |
| help="res:window:steps，逗号分隔；sink 固定 3") |
| ap.add_argument("--cells", type=int, default=30) |
| ap.add_argument("--out", required=True) |
| a = ap.parse_args() |
| dev = "cuda" |
| rep = {"gpu": torch.cuda.get_device_name(0), "others_on_gpu_before": gpu_apps(), "configs": {}, "taehv": {}} |
| print(rep, flush=True) |
| model = GcDiT(dit=empty_wan_dit(), causal=True).to(dev).eval().requires_grad_(False) |
| model.move_rope_table(dev) |
| load_gc_ckpt(model, a.ckpt) |
| model.to(torch.bfloat16) |
| g = torch.Generator(device=dev).manual_seed(0) |
| fs_cache = {} |
| for spec in a.configs.split(","): |
| res, win, steps = spec.split(":") |
| win, steps = int(win), int(steps) |
| key = (res, win) |
| if key not in fs_cache: |
| fs_cache.clear() |
| torch.cuda.empty_cache() |
| t0 = time.time() |
| fs = FastStream(model, kv_window=win, sink_size=3, rope_cap=3 + win, lat_hw=RES[res], graphs=True, compile=True) |
| fs.ktxt.normal_(generator=g).mul_(0.5) |
| fs.vtxt.normal_(generator=g).mul_(0.5) |
| fs.order, fs.n = list(range(fs.cap)), fs.cap |
| fs.pos.copy_(torch.arange(fs.cap, device=dev)) |
| fs.kraw.normal_(generator=g).mul_(0.5) |
| fs.vbuf.normal_(generator=g).mul_(0.5) |
| x = torch.randn(1, 48, 1, *RES[res], generator=g, device=dev, dtype=torch.bfloat16) |
| cell_ms(fs, x, 4, 3) |
| fs_cache[key] = (fs, x) |
| print(f"[{res} w{win}] build+compile+capture {time.time() - t0:.1f}s bufs {fs.nbytes() / 2**30:.2f} GiB", flush=True) |
| fs, x = fs_cache[key] |
| cell_ms(fs, x, steps, 3) |
| v = cell_ms(fs, x, steps, a.cells) |
| rep["configs"][spec] = {**summ(v), "tokens_per_cell": fs.hw, "kv_cells": fs.cap, |
| "mem_reserved_GiB": round(torch.cuda.memory_reserved() / 2**30, 1)} |
| print(spec, rep["configs"][spec], flush=True) |
| fs_cache.clear() |
| torch.cuda.empty_cache() |
| vae = TinyStreamVAE(a.tae, dev) |
| for res, hw in RES.items(): |
| vae.reset() |
| z = torch.randn(1, 48, 1, *hw, generator=g, device=dev, dtype=torch.bfloat16) |
| for _ in range(4): |
| vae.decode_one(z) |
| v = [] |
| for _ in range(a.cells): |
| torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| fr = vae.decode_one(z) |
| torch.cuda.synchronize() |
| v.append((time.perf_counter() - t0) * 1e3) |
| rep["taehv"][res] = {**summ(v), "frames": list(fr.shape)} |
| print("taehv", res, rep["taehv"][res], flush=True) |
| rep["others_on_gpu_after"] = gpu_apps() |
| os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) |
| json.dump(rep, open(a.out, "w"), indent=1) |
| print(json.dumps(rep, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
