[docs] DGX Spark (GB10) performance & tuning guide + reproduction examples (#1631)

Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
This commit is contained in:
Raghav K
2026-08-06 14:36:38 -07:00
committed by GitHub
co-authored by SolitaryThinker
parent b963a24612
commit a2bfc7cdb2
6 changed files with 635 additions and 0 deletions
+1
View File
@@ -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
@@ -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
@@ -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.
+6
View File
@@ -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
@@ -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()
@@ -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()