diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 0818b875fa..47e06bd394 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -5,6 +5,7 @@ FastVideo supports the following hardware platforms: - [NVIDIA CUDA](installation/gpu.md) - [NVIDIA DGX Spark / GB10 (ARM64 + CUDA 13)](installation/spark.md) + ([performance & tuning](installation/spark_performance.md)) - [Apple silicon](installation/mps.md) ## Quick Installation diff --git a/docs/getting_started/installation/spark.md b/docs/getting_started/installation/spark.md index 499c211b80..326598826a 100644 --- a/docs/getting_started/installation/spark.md +++ b/docs/getting_started/installation/spark.md @@ -137,6 +137,13 @@ If you hit other issues, please open an issue on our our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ) for additional support. +## Next: performance & tuning + +Installed and verified? See [DGX Spark: Performance & Tuning](spark_performance.md) +for which models are practical on the GB10, what makes them faster, and what +won't help on this hardware (and why) — so you don't spend a night tuning knobs +that can't move here. + ## Development Environment Setup If you're planning to contribute to FastVideo please see the diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md new file mode 100644 index 0000000000..2687f073b3 --- /dev/null +++ b/docs/getting_started/installation/spark_performance.md @@ -0,0 +1,195 @@ +# DGX Spark (GB10): Performance & Tuning + +You have FastVideo [installed on a DGX Spark](spark.md) — this page is what to +run next. It covers **which models are practical on the GB10, what actually +makes them faster, and what won't help (and why)**, so you don't burn a night +tuning knobs that can't move on this hardware. + +!!! tip "TL;DR" + - **Use distilled few-step models** (e.g. `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`). + They run in ~40 s/video. Full-step models are 12–47 min on the GB10. + - On few-step models, **VAE decode is the bottleneck**, not attention — it's + bandwidth-bound on the Spark's unified memory. + - **bf16 VAE decode** is the real, lossless lever (FastVideo already turns it + on for Wan). **FlashAttention, linear quantization, and `torch.compile` of + the VAE give little or nothing here** — see the table below. + - Heavy runs can make the box unreachable — run generations with VAE tiling on + and `nice -n 19`. See [Running safely](#running-safely-dont-lock-the-box). + +## The hardware reality (this explains everything below) + +The GB10 pairs a Blackwell GPU (`sm_121`) with **128 GB of unified LPDDR5X memory +(~270 GB/s) shared between CPU and GPU**. That bandwidth is roughly **10× below a +datacenter GPU's HBM**. Two consequences drive every tuning decision: + +1. **Memory-bandwidth-bound stages hurt disproportionately.** VAE decode moves a + lot of data and becomes the dominant cost on short (few-step) generations. +2. **Compute-bound stages scale with step count.** Full-step diffusion (50+ + steps) is denoise-bound and simply takes a long time here. + +## Use distilled few-step models + +The single biggest lever on the GB10 is **model choice**. A 3-step distilled +model is ~18× faster than the full-step version of the same architecture: + +| Model | Steps | Time / video | Bottleneck | +|---|---|---|---| +| FastWan2.1-T2V-1.3B (distilled) | 3 | **~40 s** | VAE decode | +| Wan2.1-T2V-1.3B (full-step) | 50 | ~12 min | denoise | +| Cosmos-Predict2.5-2B (full-step) | 51 | ~47 min | denoise | +| LTX2.3-distilled (+audio) | 8 | ~6 min | mixed | + +The bottleneck flips from decode to denoise at around **4 steps**. Below that, +you're paying mostly for VAE decode; above it, mostly for the denoising loop. + +!!! note "Few-step timings are noisy — measure in-process" + On a 3-step run, one-time per-process startup (Triton autotune, allocator + warmup) dominates and never amortizes, so single-run totals wobble ~±30%. + Compare levers **back-to-back in one process or as medians**, never as two + separate single runs. The [reproduction script](#reproduce-these-numbers) + does this for you. + +## bf16 VAE decode — the real lever (already on for Wan) + +Because few-step generation is decode-bound, VAE decode precision is where the +time is. Decoding in **bf16 instead of fp32 is essentially lossless** (MS-SSIM +~0.9999 vs fp32 on the identical latent) and ~1.14× faster — worth roughly +5–7% end-to-end on a decode-bound few-step model. + +**FastVideo already defaults Wan's decode to bf16** (`vae_decode_precision="bf16"`, +with encode kept at fp32), so for the recommended Wan/FastWan models there's +nothing to set. If you run a model that still defaults to an fp32 decode, set the +decode-only override yourself: + +```python +from fastvideo.configs.pipelines.base import PipelineConfig + +pipeline_config = PipelineConfig.from_pretrained(model_id) +pipeline_config.vae_decode_precision = "bf16" # decode-only; leaves encode precision alone +``` + +Decode is output-only, so lowering its precision is safe. (Encode seeds the +denoising trajectory for I2V/causal models, so that stays at the pipeline's +default — don't lower `vae_precision` blindly for those.) + +## Memory: one unified 128 GB pool + +The GB10 has **no separate VRAM** — CPU and GPU share one 128 GB LPDDR5X pool +(~118 GB usable). Two practical consequences: + +- **`nvidia-smi` reports memory as `[N/A]`** on the GB10, and the system "used" + figure conflates CPU + GPU + cache, so it's only a soft upper bound — treat the + whole 128 GB as one shared budget. For a per-run figure, use FastVideo's own + `peak_memory_mb` (reported on the generation result and by the performance + benchmark), which is measured inside the worker that runs the model. +- **The 128 GB is a *working-set* ceiling, not storage** — the model cache lives + on the NVMe (3.7 TB, ample). What has to fit in 128 GB is the weights, + activations, and KV cache — and, critically, the **VAE decode buffers**, which + is why tiling matters + (an untiled high-res decode can spike the pool into swap and lock the box). + +The recommended few-step models are comfortable here: their weights are small +(1.3–2 B) and few-step generation keeps activations modest — a Wan2.1-1.3B +few-step generation peaks at **~8.4 GB** (measured), a small fraction of the pool. +The pressure comes from **decode resolution/frames**, not the model — a +1080p×121-frame untiled decode is what pushes the pool toward its ceiling, which +is why VAE tiling stays +on by default. + +## What helps vs. what doesn't on the GB10 + +The honest summary — most "obvious" GPU optimizations don't move the needle on +this hardware, for reasons specific to it: + +| Lever | Effect on the GB10 | Use it? | +|---|---|---| +| Distilled few-step model | ~18× vs full-step | ✅ **the primary lever** | +| bf16 VAE decode | ~1.14×, lossless; ~5–7% e2e on few-step | ✅ default for Wan | +| VSA (video sparse attention) | works out of the box (Triton kernel auto-selects on `sm_121`) | ✅ automatic | +| Building FlashAttention | **no speedup** — Torch SDPA already hits an efficient flash kernel on `sm_121`, and FA2 ties it | ❌ not worth building | +| `torch.compile` of the VAE decode | recompile storm (per-frame varying shapes) → ~1.1× | ❌ dead end | +| Linear (fp8 / nvfp4) quantization on long-sequence models (e.g. Cosmos) | ~nothing — see below | ❌ wrong lever here | +| FP4 attention (`ATTN_QAT_INFER`) | works on `sm_121` (runtime allowlist landed in #1647; kernel build is #1598); helps, but needs a QAT-trained checkpoint | ⚠️ opt-in — see below | +| FP4 linear on short-sequence models (LTX2) | up to −24% denoise at 1080p (#1594) | ⚠️ model/resolution-dependent | + +### Why linear quantization is the wrong lever on long-sequence models + +Quantizing the linear (GEMM) layers is a natural first instinct, but on a +long-sequence video model it buys almost nothing on the GB10. A video-DiT denoise +step is dominated by **O(N²) attention** at these sequence lengths (tens of +thousands of tokens); the linear layers are a small single-digit fraction of the +work. Quantizing them faster leaves the attention-bound total essentially +unchanged — measured at ~1% on Cosmos-2.5, i.e. noise, and full-step CFG models +also lose quality to per-step quantization error. + +The same mechanism **does** help on **short-sequence** models: LTX2's aggressive +VAE compression gives it short attention sequences, so FP4 linear reaches −24% +there (#1594). The rule: **on the GB10, the lever that matters is attention +(sparse or FP4), not the linear layers** — unless the model has short sequences. + +### FP4 on the GB10 (opt-in) + +Block-scaled FP4 works on `sm_121` under CUDA 13: + +- **FP4 attention** (`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER`, #1598) is + numerically correct on the GB10 and ~6% faster end-to-end generation, but it only preserves + quality on a **quantization-aware-distilled checkpoint** (e.g. + `FastVideo/FastWan-QAD-1.3B`) — stock weights aren't trained to tolerate it. +- **FP4 linear** helps only where sequences are short (LTX2, above). + +The [`qad_fp4_ab.py`](#reproduce-these-numbers) harness reproduces the FP4 +attention A/B on the QAD checkpoint. + +## Running safely (don't lock the box) + +The GB10 is easy to make **unreachable** — a heavy build or an untiled high-res +decode starves the ~20 ARM cores and unified memory, `sshd` can't get cycles, and +you're locked out at *"Connection timed out during banner exchange"* until the box +is power-cycled. To avoid it: + +- **Inference:** keep **VAE tiling on** (the default), use sane resolution/frames, + and run under `nice -n 19`: + + ```bash + nice -n 19 nohup python your_script.py > run.log 2>&1 & + ``` + +- **Builds** (flash-attn, kernel): `nice -n 19`, `MAX_JOBS=2`, `nohup`. Never a + bare foreground high-parallelism build. +- Leave `*_cpu_offload` at the example defaults — "CPU" offload is the *same* + unified RAM on the GB10, so the win is tiling + sane resolution, not offloading. + +## Gotchas specific to the GB10 + +A few things that surprise people on this box (beyond the memory notes above): + +- **Don't force `TORCH_SDPA` on a VSA checkpoint** (FastWan, LTX2.3-distilled). + The SDPA path builds a model without the gate weights the checkpoint carries and + fails to load. Run the model natively — VSA auto-routes to its Triton kernel on + `sm_121`. +- **Few-step timings are noisy run-to-run** (~±30%) — one-time startup dominates a + 3-step run. Compare in-process / as medians, never two separate single runs (the + benchmark script does this). +- **`nvidia-smi` shows `[N/A]` for memory** — see [Memory](#memory-one-unified-128-gb-pool). +- **Cosmos-2.5** uses a Qwen2.5-VL text encoder; make sure you're on a FastVideo + build recent enough to include its `transformers`-compatibility handling before + running it. + +## Reproduce these numbers + +Two scripts under `examples/inference/optimizations/` reproduce the claims on +your own GB10: + +```bash +# Headline: few-step generation timing (median) + the bf16-vs-fp32 decode A/B. +# FASTVIDEO_STAGE_LOGGING=1 also prints the denoise / decode / text split. +FASTVIDEO_STAGE_LOGGING=1 nice -n 19 \ + python examples/inference/optimizations/spark_benchmark.py + +# FP4 attention quality/speed A/B on the QAD checkpoint (one arm per run). +QAD_LINEAR=0 FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER nice -n 19 \ + python examples/inference/optimizations/qad_fp4_ab.py +``` + +See also the [Optimizations](../../inference/optimizations.md) reference for the +full list of attention backends and quantization options. diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index de315e6209..1ec8730e76 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -3,6 +3,12 @@ This page describes the various options for speeding up generation times in FastVideo. +!!! note "On a DGX Spark (GB10)?" + Several options on this page behave differently on the GB10's unified-memory + hardware — some give little or nothing there. See + [DGX Spark: Performance & Tuning](../getting_started/installation/spark_performance.md) + for what actually helps on that platform and why. + ## Table of Contents - Optimized Attention Backends diff --git a/examples/inference/optimizations/qad_fp4_ab.py b/examples/inference/optimizations/qad_fp4_ab.py new file mode 100644 index 0000000000..8dd09d9c40 --- /dev/null +++ b/examples/inference/optimizations/qad_fp4_ab.py @@ -0,0 +1,219 @@ +"""QAD FP4 quality A/B/C/D harness — Wan2.1-T2V-1.3B on sm_121 (DGX Spark GB10). + +Tests whether the quantization-aware-distilled checkpoint +``FastVideo/FastWan-QAD-1.3B`` recovers FP4 quality on sm_121, using the +sm_121-enabled ``ATTN_QAT_INFER`` attention kernel. Running FP4 attention on +*stock* Wan weights gives output below bf16 — expected, because stock weights +were never trained to tolerate FP4 attention. This harness runs the checkpoint +that *was* (fake-quant FP4 attention + NVFP4 linear trained into the model). + +Attention (bf16 vs FP4) and linear (bf16 vs FP4) are fully decoupled at +inference, so we can isolate each axis: + + FASTVIDEO_ATTENTION_BACKEND unset -> bf16 attention (SDPA) + FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER -> FP4 attention + QAD_LINEAR=0 -> bf16 linear + QAD_LINEAR=1 -> NVFP4 FP4 linear (flashinfer) + + arm attention linear selects with + A bf16 bf16 QAD_LINEAR=0 (no ATTN env) reference + B FP4 bf16 QAD_LINEAR=0 ATTN_QAT_INFER isolate attn + C bf16 FP4 QAD_LINEAR=1 (no ATTN env) isolate linear + D FP4 FP4 QAD_LINEAR=1 ATTN_QAT_INFER full 4-bit + +All four arms run end-to-end on the GB10 (the full 4-bit path — FP4 linear + +FP4 attention — works on sm_121). This script still runs exactly ONE arm per +invocation and dumps a C stack on any hard crash, so a single misbehaving arm +can never take the others down with it; the runbook loops it four times with +different env. Quality is the eye/ear on the saved mp4 + a matching-frame still; +timing is the mean generation_time (full pipeline: text-encode + +denoise + VAE decode) over the measured runs. + +On the GB10, expect FP4 attention ~6% faster end-to-end generation vs bf16 and quality-neutral +by eye on the QAD checkpoint (both share the 3-step distill's quality ceiling). +FP4 *linear* is roughly break-even at 1.3B/480p (the per-call quantize overhead +eats the small-GEMM saving in eager mode); its win shows at higher resolution +with torch.compile. + +Knobs (env): QAD_MODEL, QAD_DISTILLED, QAD_STEPS (3), QAD_GUIDANCE (1.0), +QAD_HEIGHT (480), QAD_WIDTH (832), QAD_FRAMES (77), QAD_SEED (42), +QAD_WARMUP (1), QAD_RUNS (3), QAD_STILL (20), QAD_OUT (qad_fp4_samples). +""" +from __future__ import annotations + +import faulthandler +import glob +import os +import time + +import torch + +faulthandler.enable() # dump a C stack if any arm hard-crashes. + + +def _env(name: str, default: str) -> str: + return os.environ.get(name, default) + + +def _env_int(name: str, default: int) -> int: + return int(os.environ.get(name, str(default))) + + +def _env_float(name: str, default: float) -> float: + return float(os.environ.get(name, str(default))) + + +# Distilled QAD transformer, loaded on top of the base Wan2.1-1.3B pipeline. +DEFAULT_DISTILLED = "FastVideo/FastWan-QAD-1.3B" +# The repo is a full diffusers pipeline; we overlay only its transformer onto the +# base Wan pipeline (vae/text_encoder are Wan-identical). +DISTILLED_WEIGHTS_FILE = "transformer/diffusion_pytorch_model.safetensors" + +PROMPT = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." +) + + +def resolve_distilled_weights(hf_id: str) -> str: + """Return a local path to the distilled transformer safetensors.""" + if not hf_id: + return "" + if os.path.exists(hf_id): + return hf_id + from huggingface_hub import hf_hub_download + return hf_hub_download(repo_id=hf_id, filename=DISTILLED_WEIGHTS_FILE) + + +def build_generator(fp4_linear: bool): + from fastvideo import VideoGenerator + from fastvideo.configs.pipelines.base import PipelineConfig + + model_id = _env("QAD_MODEL", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + + pipeline_config = PipelineConfig.from_pretrained(model_id) + pipeline_config.dit_precision = "bf16" + pipeline_config.vae_precision = "bf16" + pipeline_config.text_encoder_precisions = ("bf16",) + + if fp4_linear: + # Wan-style config: matches to_q/k/v/out + ffn (the plain NVFP4 config is + # LTX2-specific and would quantize nothing on Wan). + from fastvideo.layers.quantization.nvfp4_qat_config import NVFP4QATConfig + pipeline_config.dit_config.quant_config = NVFP4QATConfig() + + extra_kwargs = {} + distilled = resolve_distilled_weights(_env("QAD_DISTILLED", DEFAULT_DISTILLED)) + if distilled: + print(f"[qad] distilled weights: {distilled}") + extra_kwargs["init_weights_from_safetensors"] = distilled + + # Keep everything resident (the 1.3B QAD model + FP4 fits the GB10's unified + # memory); real Wan VAE decode for a faithful quality read (no TAEHV). + return VideoGenerator.from_pretrained( + model_id, + pipeline_config=pipeline_config, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + pin_cpu_memory=False, + enable_torch_compile=False, # eager: isolate the FP4 effect, no compile noise + **extra_kwargs, + ) + + +def main() -> None: + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required.") + + fp4_linear = _env_int("QAD_LINEAR", 0) == 1 + attn = _env("FASTVIDEO_ATTENTION_BACKEND", "default") + steps = _env_int("QAD_STEPS", 3) + guidance = _env_float("QAD_GUIDANCE", 1.0) + height = _env_int("QAD_HEIGHT", 480) + width = _env_int("QAD_WIDTH", 832) + frames = _env_int("QAD_FRAMES", 77) + seed = _env_int("QAD_SEED", 42) + warmup = _env_int("QAD_WARMUP", 1) + runs = _env_int("QAD_RUNS", 3) + still_idx = _env_int("QAD_STILL", 20) + out_dir = _env("QAD_OUT", "qad_fp4_samples") + prompt = _env("QAD_PROMPT", PROMPT) + + tag = f"lin-{'fp4' if fp4_linear else 'bf16'}_attn-{attn.lower()}" + cap = torch.cuda.get_device_capability() + print(f"[qad] GPU {torch.cuda.get_device_name()} (cc {cap[0]}.{cap[1]})") + print(f"[qad] ARM {tag}: linear={'FP4' if fp4_linear else 'bf16'}, " + f"attention={attn}, {steps} steps, guidance {guidance}, " + f"{height}x{width}x{frames}, seed {seed}") + print(f"[qad] prompt: {prompt[:80]}{'...' if len(prompt) > 80 else ''}") + + arm_dir = os.path.join(out_dir, tag) + os.makedirs(arm_dir, exist_ok=True) + generator = build_generator(fp4_linear) + + def _generate(): + # seed + frame dims live under `sampling` (SamplingConfig); `output` + # only takes output_path/save_video/return_frames (OutputConfig). + return generator.generate(request={ + "prompt": prompt, + "sampling": { + "seed": seed, + "num_inference_steps": steps, + "guidance_scale": guidance, + "height": height, + "width": width, + "num_frames": frames, + }, + "output": {"save_video": True, "output_path": arm_dir, + "return_frames": True}, + }) + + for _ in range(warmup): + _generate() + + denoise_times: list[float] = [] + last = None + for i in range(runs): + torch.cuda.synchronize() + t0 = time.perf_counter() + last = _generate() + torch.cuda.synchronize() + wall = time.perf_counter() - t0 + # generate_video returns a plain dict; attribute access would always + # fall back to wall time. + gen_t = last.get("generation_time") if isinstance(last, dict) else None + denoise_times.append(gen_t if gen_t is not None else wall) + print(f"[qad] {tag} run {i + 1}/{runs}: {wall:.2f}s wall " + f"(gen {denoise_times[-1]:.2f}s)") + + # The pipeline wrote the mp4 (full known-good encode) into arm_dir; report + # it and pull a matching-frame still from the [b,c,t,h,w] samples tensor + # using the same recipe the pipeline's frame builder uses. + mp4s = sorted(glob.glob(os.path.join(arm_dir, "*.mp4")), key=os.path.getmtime) + if mp4s: + print(f"[qad] video: {mp4s[-1]}") + samples = getattr(last, "samples", None) if last is not None else None + if samples is not None and getattr(samples, "ndim", 0) == 5: + import imageio + f = min(still_idx, samples.shape[2] - 1) # samples: [b, c, t, h, w] + still = (samples[0, :, f].permute(1, 2, 0).clamp(0, 1) * 255) + still = still.to(torch.uint8).cpu().numpy() + png = os.path.join(arm_dir, f"raccoon_{tag}_f{f}.png") + imageio.imwrite(png, still) + print(f"[qad] still: {png}") + else: + print("[qad] note: no 5-D samples tensor; grab a frame from the mp4 above") + + mean = sum(denoise_times) / len(denoise_times) + print(f"\n[qad][{tag}] generation mean {mean:.2f}s over {runs} runs " + f"({warmup} warmup, {steps} steps)") + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/optimizations/spark_benchmark.py b/examples/inference/optimizations/spark_benchmark.py new file mode 100644 index 0000000000..0850a6494b --- /dev/null +++ b/examples/inference/optimizations/spark_benchmark.py @@ -0,0 +1,207 @@ +"""DGX Spark (GB10) reproduction benchmark for the performance guide. + +Reproduces the two headline claims in +``docs/getting_started/installation/spark_performance.md`` on your own GB10: + + 1. A distilled few-step model is usable (~30 s/video) and **decode-bound** on + the Spark's unified LPDDR5X memory. + 2. bf16 VAE decode is essentially lossless (MS-SSIM ~1.0 vs fp32) and modestly + faster — which is why FastVideo already defaults Wan's *decode* to bf16. + +Two parts, both in-process so they control for the Spark's run-to-run variance: + + A. **Generation timing** — loads a distilled model once, generates ``--runs`` + videos back-to-back (after ``--warmup``), reports the *median* generation + time. Set ``FASTVIDEO_STAGE_LOGGING=1`` to also see the per-stage + (denoise / VAE decode / text-encode) split that shows the decode bottleneck. + + B. **Decode precision A/B** — decodes ONE fixed latent fp32-vs-bf16 in the same + process and reports MS-SSIM + speedup. This isolates the decode delta with + no denoise non-determinism and no video-codec noise. + +**Why median, not a single run:** on the GB10 a 3-step generation is dominated by +one-time per-process startup (Triton autotune, allocator warmup) that never +amortizes over so few steps, so single-run totals wobble ~±30%. Always compare +few-step levers back-to-back / as medians, never as two separate single runs. + +Run it safely on a shared box (see the best-practices note in the perf guide): + + FASTVIDEO_STAGE_LOGGING=1 nice -n 19 nohup \ + python examples/inference/optimizations/spark_benchmark.py > spark_bench.log 2>&1 & + tail -f spark_bench.log + +Knobs (flags or env): --model, --runs (3), --warmup (1), --steps (3), +--frames (81), --height (448), --width (832), --seed (42), --out, --skip-gen, +--skip-decode. +""" +from __future__ import annotations + +import argparse +import os +import statistics +import time + +import torch + + +def _p(msg: str) -> None: + print(f"[spark-bench] {msg}", flush=True) + + +def bench_generation(args) -> None: + """Part A: median few-step generation time on a distilled model.""" + from fastvideo import VideoGenerator + from fastvideo.api.sampling_param import SamplingParam + + # VSA auto-routes to the Triton kernel on sm_121; do NOT force TORCH_SDPA on a + # VSA checkpoint (the SDPA path builds a model without the gate weights the + # checkpoint carries and fails to load). + os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "VIDEO_SPARSE_ATTN") + + _p(f"loading {args.model} ...") + load_t0 = time.perf_counter() + generator = VideoGenerator.from_pretrained( + args.model, + num_gpus=1, + use_fsdp_inference=False, + # Leave offload at these defaults: "CPU" offload is the same unified RAM + # on the GB10, so the win is tiling + sane resolution, not offloading. + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + dit_cpu_offload=False, + vae_cpu_offload=False, + VSA_sparsity=0.8, + ) + _p(f"model loaded in {time.perf_counter() - load_t0:.1f}s") + + sampling_param = SamplingParam.from_pretrained(args.model) + sampling_param.num_frames = args.frames + sampling_param.height = args.height + sampling_param.width = args.width + sampling_param.num_inference_steps = args.steps + sampling_param.seed = args.seed + + prompt = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, " + "its eyes wide with interest. Soft natural light, warm cheerful tones, " + "mid-shot, cinematic.") + + def _gen(): + torch.cuda.synchronize() + t0 = time.perf_counter() + video = generator.generate_video(prompt, output_path=args.out, + save_video=True, + sampling_param=sampling_param) + torch.cuda.synchronize() + # generate_video returns a plain dict (legacy result), not an object — + # attribute access would silently fall back to wall time / None. + dt = video.get("generation_time") if isinstance(video, dict) else None + if dt is None: + dt = time.perf_counter() - t0 + # Peak memory is measured *inside the worker process* that runs the + # pipeline and surfaced on the result; reading torch's allocator in this + # (main) process would report ~0 because the allocations aren't here. + peak = video.get("peak_memory_mb") if isinstance(video, dict) else None + return dt, peak + + for _ in range(args.warmup): + _gen() + + times, peaks = [], [] + for i in range(args.runs): + dt, peak = _gen() + times.append(dt) + if peak: + peaks.append(peak) + _p(f"gen run {i + 1}/{args.runs}: {dt:.2f}s") + + med = statistics.median(times) + _p(f"median generation time over {args.runs} runs " + f"({args.warmup} warmup, {args.steps} steps): {med:.2f}s") + if peaks: + _p(f"peak GPU memory (worker, reported by pipeline): {max(peaks):.0f} MB " + f"= {max(peaks) / 1024:.1f} GB") + else: + _p("peak GPU memory: not reported by this pipeline build") + _p("set FASTVIDEO_STAGE_LOGGING=1 to see the denoise / decode / text split " + "(few-step generation is VAE-decode-bound on the GB10).") + generator.shutdown() + + +def bench_decode_precision(args) -> None: + """Part B: fp32-vs-bf16 VAE decode of one fixed latent (SSIM + speedup).""" + try: + from diffusers import AutoencoderKLWan + from torchmetrics.functional import ( + multiscale_structural_similarity_index_measure as msssim) + except ImportError as e: # torchmetrics is not a hard FastVideo dep + _p(f"skipping decode A/B (missing dependency: {e}); " + "`uv pip install torchmetrics` to enable it.") + return + + dev = "cuda" + vae = AutoencoderKLWan.from_pretrained( + args.model, subfolder="vae", torch_dtype=torch.float32).to(dev).eval() + + # Wan latent geometry for height x width x frames, patch (4,8,8): + # T_lat = (frames - 1) // 4 + 1, H_lat = height // 8, W_lat = width // 8 + t_lat = (args.frames - 1) // 4 + 1 + z = torch.randn(1, 16, t_lat, args.height // 8, args.width // 8, + device=dev, dtype=torch.float32) + + def _decode(autocast: bool): + torch.cuda.synchronize() + t0 = time.perf_counter() + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16, + enabled=autocast): + out = vae.decode(z, return_dict=False)[0] + torch.cuda.synchronize() + return out.float(), time.perf_counter() - t0 + + def _frames(o): # (1,3,T,H,W) [-1,1] -> (T,3,H,W) [0,1] + return ((o.clamp(-1, 1) + 1) / 2)[0].permute(1, 0, 2, 3).contiguous() + + _decode(False) # warm both paths (excluded from timing) + _decode(True) + + o32, t32 = _decode(autocast=False) # fp32 + o16, t16 = _decode(autocast=True) # bf16 (== vae_decode_precision="bf16") + ssim = msssim(_frames(o16), _frames(o32), data_range=1.0).item() + + _p(f"fp32 decode : {t32 * 1000:8.1f} ms") + _p(f"bf16 decode : {t16 * 1000:8.1f} ms ({t32 / t16:.2f}x faster)") + _p(f"MS-SSIM(bf16, fp32) on identical latent: {ssim:.4f} " + "(>= ~0.99 -> lossless; this is why Wan decode defaults to bf16)") + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers") + ap.add_argument("--runs", type=int, default=3) + ap.add_argument("--warmup", type=int, default=1) + ap.add_argument("--steps", type=int, default=3) + ap.add_argument("--frames", type=int, default=81) + ap.add_argument("--height", type=int, default=448) + ap.add_argument("--width", type=int, default=832) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="spark_bench_samples") + ap.add_argument("--skip-gen", action="store_true", + help="skip Part A (generation timing)") + ap.add_argument("--skip-decode", action="store_true", + help="skip Part B (decode precision A/B)") + args = ap.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required (run this on the GB10).") + cap = torch.cuda.get_device_capability() + _p(f"{torch.cuda.get_device_name()} (cc {cap[0]}.{cap[1]}), " + f"torch {torch.__version__}") + + if not args.skip_gen: + bench_generation(args) + if not args.skip_decode: + bench_decode_precision(args) + + +if __name__ == "__main__": + main()