mlx_fastwan_bench ¶

Prove-out benchmark for the MLX FastWan runtime (Apple Silicon).

Sweeps {dtype/quant} x {decoder}, generates a clip per cell, and records the latency breakdown, peak unified memory, and MS-SSIM (optionally LPIPS) against a reference video. It emits a JSON blob and a markdown table -- the artifact that turns "int8 + TAEHV looks good" into defensible numbers, and (via --assert-min-ssim) a regression gate for the mx.compile work.

Design notes: - Generation reuses the hybrid POC helpers in examples/inference/basic/mlx_wan_prompt_to_video.py (torch-MPS UMT5 encode and Wan-VAE/TAEHV decode) plus the on-device MLX DMD sampler (fastvideo/mlx_runtime/sampling.py); the denoise loop never leaves the device. - Quality reuses the tested MS-SSIM primitive fastvideo/tests/utils.py::compute_video_ssim_torchvision. - Reference: by default each cell is scored against the highest-fidelity cell in the sweep (fp16 + wan-vae), which needs no CUDA box and answers "how much does int8/int4/TAEHV degrade vs the best local config". Pass --reference PATH to score against an external clip instead (e.g. the torch-MPS or CUDA FastVideo output of the same model) for a "vs. the original model" column.

Run on an Apple Silicon Mac (needs mlx + a torch build with MPS):

python fastvideo/benchmarks/mlx_fastwan_bench.py         --modes fp16,bf16,int8,int4 --decoders taehv,wan-vae

Functions:¶

fastvideo.benchmarks.mlx_fastwan_bench.denoise_dmd_on_device ¶

denoise_dmd_on_device(*, mx, dit, latents, encoder_hidden_states, freqs_cis, timesteps: list[int], renoise_by_step: list[ndarray], schedule, dmd_step, mx_dtype) -> tuple[ndarray, list[float]]

Run the FastWan DMD loop entirely on the MLX device.

Mirrors the loop in mlx_wan_prompt_to_video.py (fp32 affine math, MLX RNG re-noise) so the benchmark measures exactly the shipped path.

Returns the final latents plus per-step wall times. The first step carries one-time costs (mx.compile tracing, kernel warm-up), so first-vs-steady step timing is how the benchmark separates cold-start from steady-state denoise throughput.

All host-side tensors (timesteps, re-noise draws) are uploaded before the loop starts, so the per-step body performs no bulk host->device transfers and step timings measure device work rather than staging copies.

Source code in fastvideo/benchmarks/mlx_fastwan_bench.py

def denoise_dmd_on_device(
    *,
    mx,
    dit,
    latents,
    encoder_hidden_states,
    freqs_cis,
    timesteps: list[int],
    renoise_by_step: list[np.ndarray],
    schedule,
    dmd_step,
    mx_dtype,
) -> tuple[np.ndarray, list[float]]:
    """Run the FastWan DMD loop entirely on the MLX device.

    Mirrors the loop in ``mlx_wan_prompt_to_video.py`` (fp32 affine math, MLX RNG
    re-noise) so the benchmark measures exactly the shipped path.

    Returns the final latents plus per-step wall times. The first step carries
    one-time costs (mx.compile tracing, kernel warm-up), so first-vs-steady
    step timing is how the benchmark separates cold-start from steady-state
    denoise throughput.

    All host-side tensors (timesteps, re-noise draws) are uploaded before the
    loop starts, so the per-step body performs no bulk host->device transfers
    and step timings measure device work rather than staging copies.
    """
    timesteps_mx = [mx.array([float(timestep)]).astype(mx.float32) for timestep in timesteps]
    renoise_mx = [mx.array(renoise).astype(mx.float32) for renoise in renoise_by_step]
    if timesteps_mx or renoise_mx:
        mx.eval(*timesteps_mx, *renoise_mx)

    step_times: list[float] = []
    for step_index, timestep in enumerate(timesteps):
        step_start = time.perf_counter()
        noise_input_latent = latents
        noise_pred = dit(latents.astype(mx_dtype), encoder_hidden_states, timesteps_mx[step_index], freqs_cis)

        noise_input_f32 = noise_input_latent.astype(mx.float32)
        pred_noise_f32 = noise_pred.astype(mx.float32)
        if step_index < len(timesteps) - 1:
            next_ts: float | None = float(timesteps[step_index + 1])
            renoise = renoise_mx[step_index]
        else:
            next_ts, renoise = None, None
        latents = dmd_step(
            latents=noise_input_f32,
            noise_input_latent=noise_input_f32,
            pred_noise=pred_noise_f32,
            schedule=schedule,
            timestep=float(timestep),
            next_timestep=next_ts,
            noise=renoise,
        ).astype(mx_dtype)
        mx.eval(latents)
        step_times.append(time.perf_counter() - step_start)
    return np.array(latents.astype(mx.float32)), step_times