1052 lines
40 KiB
Python
1052 lines
40 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# mypy: disable-error-code=no-untyped-call
|
||
"""FastWan-oriented helpers for the experimental MLX runtime path."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
import statistics
|
||
import time
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import TYPE_CHECKING, Any
|
||
from collections.abc import Callable
|
||
|
||
import fastvideo.envs as envs
|
||
from fastvideo.logger import init_logger
|
||
|
||
if TYPE_CHECKING:
|
||
import mlx.core as mx
|
||
import torch
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class FastWanShape:
|
||
height: int
|
||
width: int
|
||
num_frames: int
|
||
latent_frames: int
|
||
latent_height: int
|
||
latent_width: int
|
||
patch_frames: int
|
||
patch_height: int
|
||
patch_width: int
|
||
tokens: int
|
||
hidden_size: int
|
||
num_heads: int
|
||
head_dim: int
|
||
|
||
|
||
class UnsupportedMLXQuantizationError(ValueError):
|
||
"""A quantization mode the installed MLX build cannot execute.
|
||
|
||
Raised by :func:`ensure_quantization_supported` before any model weights
|
||
are loaded, so callers (CLI flags, benchmark sweeps) can fail fast with an
|
||
actionable message -- or skip the mode -- instead of crashing deep inside
|
||
``mx.quantize`` mid-load.
|
||
"""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class MLXQuantizationSpec:
|
||
"""MLX quantized-matmul configuration for DiT linear weights."""
|
||
|
||
mode: str
|
||
bits: int | None = None
|
||
group_size: int | None = None
|
||
|
||
@classmethod
|
||
def from_name(cls, name: str | None) -> MLXQuantizationSpec | None:
|
||
if name is None or name in {"", "none", "fp16", "fp32"}:
|
||
return None
|
||
if name == "int8":
|
||
return cls(mode="affine", bits=8, group_size=64)
|
||
if name == "int6":
|
||
return cls(mode="affine", bits=6, group_size=64)
|
||
if name == "int4":
|
||
return cls(mode="affine", bits=4, group_size=64)
|
||
if name == "mxfp8":
|
||
return cls(mode="mxfp8")
|
||
if name == "mxfp4":
|
||
return cls(mode="mxfp4")
|
||
if name == "nvfp4":
|
||
return cls(mode="nvfp4")
|
||
raise ValueError(f"Unsupported MLX quantization mode: {name}")
|
||
|
||
@property
|
||
def label(self) -> str:
|
||
if self.mode == "affine":
|
||
return f"int{self.bits}"
|
||
return self.mode
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class QuantizedMatrix:
|
||
weight: mx.array
|
||
scales: mx.array
|
||
biases: mx.array | None
|
||
spec: MLXQuantizationSpec
|
||
dequantized_dtype: mx.Dtype
|
||
|
||
|
||
def fastwan_shape(
|
||
*,
|
||
height: int,
|
||
width: int,
|
||
num_frames: int,
|
||
vae_temporal_compression: int = 4,
|
||
vae_spatial_compression: int = 8,
|
||
patch_size: tuple[int, int, int] = (1, 2, 2),
|
||
num_heads: int = 12,
|
||
head_dim: int = 128,
|
||
) -> FastWanShape:
|
||
"""Return the approximate DiT token shape for Wan/FastWan T2V inference."""
|
||
latent_frames = (num_frames - 1) // vae_temporal_compression + 1
|
||
latent_height = height // vae_spatial_compression
|
||
latent_width = width // vae_spatial_compression
|
||
patch_frames = latent_frames // patch_size[0]
|
||
patch_height = latent_height // patch_size[1]
|
||
patch_width = latent_width // patch_size[2]
|
||
tokens = patch_frames * patch_height * patch_width
|
||
return FastWanShape(
|
||
height=height,
|
||
width=width,
|
||
num_frames=num_frames,
|
||
latent_frames=latent_frames,
|
||
latent_height=latent_height,
|
||
latent_width=latent_width,
|
||
patch_frames=patch_frames,
|
||
patch_height=patch_height,
|
||
patch_width=patch_width,
|
||
tokens=tokens,
|
||
hidden_size=num_heads * head_dim,
|
||
num_heads=num_heads,
|
||
head_dim=head_dim,
|
||
)
|
||
|
||
|
||
def fastwan_shape_from_config(
|
||
config_path: str | Path,
|
||
*,
|
||
height: int,
|
||
width: int,
|
||
num_frames: int,
|
||
) -> FastWanShape:
|
||
config = json.loads(Path(config_path).read_text())
|
||
return fastwan_shape(
|
||
height=height,
|
||
width=width,
|
||
num_frames=num_frames,
|
||
patch_size=tuple(config["patch_size"]),
|
||
num_heads=int(config["num_attention_heads"]),
|
||
head_dim=int(config["attention_head_dim"]),
|
||
)
|
||
|
||
|
||
def replace_tokens(shape: FastWanShape, tokens: int) -> FastWanShape:
|
||
return FastWanShape(**{**shape.__dict__, "tokens": tokens})
|
||
|
||
|
||
def median_ms(samples: list[float]) -> float:
|
||
return statistics.median(samples) * 1000.0
|
||
|
||
|
||
def benchmark_mlx_attention(shape: FastWanShape, warmup: int, iters: int) -> float:
|
||
import mlx.core as mx
|
||
|
||
q = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
|
||
k = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
|
||
v = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16)
|
||
scale = shape.head_dim**-0.5
|
||
|
||
for _ in range(warmup):
|
||
y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
|
||
mx.eval(y)
|
||
|
||
samples = []
|
||
for _ in range(iters):
|
||
start = time.perf_counter()
|
||
y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
|
||
mx.eval(y)
|
||
samples.append(time.perf_counter() - start)
|
||
return median_ms(samples)
|
||
|
||
|
||
def benchmark_mlx_linear(shape: FastWanShape, warmup: int, iters: int) -> float:
|
||
import mlx.core as mx
|
||
|
||
x = mx.random.normal((shape.tokens, shape.hidden_size), dtype=mx.float16)
|
||
w = mx.random.normal((shape.hidden_size, shape.hidden_size), dtype=mx.float16)
|
||
b = mx.zeros((shape.hidden_size, ), dtype=mx.float16)
|
||
|
||
for _ in range(warmup):
|
||
y = x @ w + b
|
||
mx.eval(y)
|
||
|
||
samples = []
|
||
for _ in range(iters):
|
||
start = time.perf_counter()
|
||
y = x @ w + b
|
||
mx.eval(y)
|
||
samples.append(time.perf_counter() - start)
|
||
return median_ms(samples)
|
||
|
||
|
||
def benchmark_torch_mps_attention(shape: FastWanShape, warmup: int, iters: int) -> float | None:
|
||
try:
|
||
import torch
|
||
import torch.nn.functional as F
|
||
except ImportError:
|
||
return None
|
||
|
||
if not torch.backends.mps.is_available():
|
||
return None
|
||
|
||
device = torch.device("mps")
|
||
q = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
|
||
k = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
|
||
v = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16)
|
||
|
||
for _ in range(warmup):
|
||
y = F.scaled_dot_product_attention(q, k, v)
|
||
torch.mps.synchronize()
|
||
_ = y
|
||
|
||
samples = []
|
||
for _ in range(iters):
|
||
start = time.perf_counter()
|
||
y = F.scaled_dot_product_attention(q, k, v)
|
||
torch.mps.synchronize()
|
||
_ = y
|
||
samples.append(time.perf_counter() - start)
|
||
return median_ms(samples)
|
||
|
||
|
||
def torch_to_mx(tensor) -> mx.array:
|
||
import mlx.core as mx
|
||
|
||
return mx.array(tensor.detach().cpu().float().numpy())
|
||
|
||
|
||
def weight_dtype(weight):
|
||
if isinstance(weight, QuantizedMatrix):
|
||
return weight.dequantized_dtype
|
||
return weight.dtype
|
||
|
||
|
||
_QUANT_SUPPORT_CACHE: dict[tuple[str, int | None, int | None], str | None] = {}
|
||
|
||
|
||
def quantization_support_error(spec: MLXQuantizationSpec) -> str | None:
|
||
"""Probe whether the installed MLX build supports ``spec``.
|
||
|
||
Runs a tiny ``mx.quantize`` + ``mx.quantized_matmul`` with exactly the
|
||
arguments :func:`quantize_matrix` / :func:`linear` use, so the result
|
||
reflects the real runtime path. The affine (int8/int4) modes are stable
|
||
across MLX releases, but the ``mxfp8``/``mxfp4``/``nvfp4`` mode strings
|
||
require newer MLX builds and raise otherwise. Returns ``None`` when the
|
||
mode works, else the underlying error message. Cached per spec.
|
||
"""
|
||
key = (spec.mode, spec.bits, spec.group_size)
|
||
if key not in _QUANT_SUPPORT_CACHE:
|
||
import mlx.core as mx
|
||
|
||
try:
|
||
probe_dim = max(spec.group_size or 0, 64)
|
||
weight = mx.zeros((probe_dim, probe_dim), dtype=mx.float16)
|
||
quantized = quantize_matrix(weight, spec)
|
||
y = linear(mx.zeros((1, probe_dim), dtype=mx.float16), quantized)
|
||
mx.eval(y)
|
||
_QUANT_SUPPORT_CACHE[key] = None
|
||
except Exception as exc: # noqa: BLE001 - MLX raises varied error types per backend/version.
|
||
_QUANT_SUPPORT_CACHE[key] = f"{type(exc).__name__}: {exc}"
|
||
return _QUANT_SUPPORT_CACHE[key]
|
||
|
||
|
||
def ensure_quantization_supported(spec: MLXQuantizationSpec | None) -> None:
|
||
"""Raise :class:`UnsupportedMLXQuantizationError` if ``spec`` cannot run here."""
|
||
if spec is None:
|
||
return
|
||
error = quantization_support_error(spec)
|
||
if error is None:
|
||
return
|
||
import mlx.core as mx
|
||
|
||
mlx_version = getattr(mx, "__version__", "unknown")
|
||
raise UnsupportedMLXQuantizationError(f"MLX quantization mode '{spec.label}' is not supported by the installed mlx "
|
||
f"({mlx_version}): {error}. Upgrade mlx or pick a supported mode "
|
||
f"(int8 is currently the most reliable quality/memory target).")
|
||
|
||
|
||
def quantize_matrix(weight, spec: MLXQuantizationSpec | None):
|
||
if spec is None:
|
||
return weight
|
||
import mlx.core as mx
|
||
|
||
if len(weight.shape) < 2:
|
||
return weight
|
||
q = mx.quantize(weight, group_size=spec.group_size, bits=spec.bits, mode=spec.mode)
|
||
biases = q[2] if len(q) == 3 else None
|
||
eval_args = [q[0], q[1]]
|
||
if biases is not None:
|
||
eval_args.append(biases)
|
||
mx.eval(*eval_args)
|
||
return QuantizedMatrix(
|
||
weight=q[0],
|
||
scales=q[1],
|
||
biases=biases,
|
||
spec=spec,
|
||
dequantized_dtype=weight.dtype,
|
||
)
|
||
|
||
|
||
# Affine quantized_matmul is slower than dequantize + steel GEMM at H3's packed
|
||
# token width. Measured on Apple M4 Max / MLX 0.32.2, INT6 group 64, BF16 acts,
|
||
# Q 5376→7168: M=256 qmm is faster; M=512 dequant+GEMM is +6.4%; M≥1024 ~10%.
|
||
# 832×480×124 packed M is ~14862–14994. H3 opts into this path explicitly;
|
||
# shared FastWan and Wan 2.2 linears stay on quantized_matmul. Do not cache
|
||
# dequantized weights. Override: FASTVIDEO_MLX_DQ_GEMM=0 off, =1 measured
|
||
# floor, =<int> explicit floor.
|
||
_AFFINE_DQ_GEMM_BITS = frozenset({2, 3, 4, 5, 6, 8})
|
||
MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M = 768
|
||
_dq_gemm_engaged = 0
|
||
_dq_gemm_logged = False
|
||
|
||
|
||
def reset_dq_gemm_telemetry() -> None:
|
||
global _dq_gemm_engaged
|
||
_dq_gemm_engaged = 0
|
||
|
||
|
||
def dq_gemm_engaged() -> int:
|
||
return _dq_gemm_engaged
|
||
|
||
|
||
def affine_dq_gemm_min_m() -> int | None:
|
||
raw = envs.FASTVIDEO_MLX_DQ_GEMM.get().strip().lower()
|
||
if raw in {"", "0", "off", "false", "no"}:
|
||
return None
|
||
if raw in {"1", "on", "true", "yes"}:
|
||
return MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M
|
||
try:
|
||
value = int(raw)
|
||
except ValueError:
|
||
return MLX_AFFINE_DQ_GEMM_DEFAULT_MIN_M
|
||
if value <= 0:
|
||
return None
|
||
return value
|
||
|
||
|
||
def _matmul_leading_rows(x) -> int:
|
||
last = int(x.shape[-1]) if x.ndim else 0
|
||
if last <= 0:
|
||
return 0
|
||
return int(x.size) // last
|
||
|
||
|
||
def _quantized_linear(x, weight: QuantizedMatrix, *, use_affine_dq_gemm: bool = False):
|
||
import mlx.core as mx
|
||
|
||
global _dq_gemm_engaged, _dq_gemm_logged
|
||
spec = weight.spec
|
||
min_m = affine_dq_gemm_min_m() if use_affine_dq_gemm else None
|
||
rows = _matmul_leading_rows(x)
|
||
if (min_m is not None and spec.mode == "affine" and spec.bits in _AFFINE_DQ_GEMM_BITS
|
||
and spec.group_size is not None and rows >= min_m):
|
||
dequantized = mx.dequantize(
|
||
weight.weight,
|
||
weight.scales,
|
||
weight.biases,
|
||
group_size=spec.group_size,
|
||
bits=spec.bits,
|
||
mode=spec.mode,
|
||
dtype=x.dtype,
|
||
)
|
||
y = (x @ dequantized.T).astype(x.dtype)
|
||
_dq_gemm_engaged += 1
|
||
if not _dq_gemm_logged:
|
||
_dq_gemm_logged = True
|
||
logger.info("affine dequant+GEMM engaged (rows=%d, floor=%d, bits=%s)", rows, min_m, spec.bits)
|
||
return y
|
||
return mx.quantized_matmul(
|
||
x,
|
||
weight.weight,
|
||
weight.scales,
|
||
weight.biases,
|
||
transpose=True,
|
||
group_size=spec.group_size,
|
||
bits=spec.bits,
|
||
mode=spec.mode,
|
||
).astype(x.dtype)
|
||
|
||
|
||
def linear(x, weight, bias=None, *, use_affine_dq_gemm: bool = False):
|
||
y = (_quantized_linear(x, weight, use_affine_dq_gemm=use_affine_dq_gemm)
|
||
if isinstance(weight, QuantizedMatrix) else x @ weight.T)
|
||
if bias is not None:
|
||
y = y + bias
|
||
return y
|
||
|
||
|
||
def _use_fast_norm() -> bool:
|
||
"""Opt-in to MLX's fused ``mx.fast`` normalization kernels for Wan.
|
||
|
||
Off by default so the numerically-explicit reference path stays the
|
||
baseline. Set ``FASTVIDEO_MLX_FAST_NORM=1`` to route LayerNorm/RMSNorm
|
||
through single fused Metal kernels (fewer intermediates, less memory
|
||
traffic) and benchmark the speedup.
|
||
H3 uses its own fused RMSNorm default and does not consult this toggle.
|
||
"""
|
||
return envs.FASTVIDEO_MLX_FAST_NORM.get()
|
||
|
||
|
||
def layer_norm(x, weight=None, bias=None, eps: float = 1e-6):
|
||
import mlx.core as mx
|
||
|
||
if _use_fast_norm():
|
||
# Compute in fp32 (matching the reference below) so downstream dtype
|
||
# and precision are identical across call sites.
|
||
w = weight.astype(mx.float32) if weight is not None else None
|
||
b = bias.astype(mx.float32) if bias is not None else None
|
||
return mx.fast.layer_norm(x.astype(mx.float32), w, b, eps)
|
||
|
||
x_float = x.astype(mx.float32)
|
||
mean = mx.mean(x_float, axis=-1, keepdims=True)
|
||
var = mx.mean(mx.square(x_float - mean), axis=-1, keepdims=True)
|
||
y = (x_float - mean) * mx.rsqrt(var + eps)
|
||
if weight is not None:
|
||
y = y * weight
|
||
if bias is not None:
|
||
y = y + bias
|
||
return y
|
||
|
||
|
||
def rms_norm(x, weight, eps: float = 1e-6):
|
||
import mlx.core as mx
|
||
|
||
if _use_fast_norm():
|
||
return mx.fast.rms_norm(x, weight, eps)
|
||
|
||
orig_dtype = x.dtype
|
||
x_float = x.astype(mx.float32)
|
||
variance = mx.mean(mx.square(x_float), axis=-1, keepdims=True)
|
||
y = x_float * mx.rsqrt(variance + eps)
|
||
return y.astype(orig_dtype) * weight
|
||
|
||
|
||
def apply_rotary_emb(x, cos, sin, *, is_neox_style: bool = False):
|
||
"""Apply FastVideo's rotary convention to MLX tensors.
|
||
|
||
Args:
|
||
x: [batch, seq, heads, head_dim]
|
||
cos/sin: [seq, head_dim] for Wan's full-dimension rotate-pair style,
|
||
or [seq, head_dim // 2] for traditional RoPE.
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
head_size = x.shape[-1]
|
||
rope_dim = cos.shape[-1]
|
||
cos = cos[None, :, None, :]
|
||
sin = sin[None, :, None, :]
|
||
x_float = x.astype(mx.float32)
|
||
|
||
if rope_dim == head_size:
|
||
x_pairs = x_float.reshape(*x.shape[:-1], -1, 2)
|
||
x_real = x_pairs[..., 0]
|
||
x_imag = x_pairs[..., 1]
|
||
x_rotated = mx.stack([-x_imag, x_real], axis=-1).reshape(*x.shape)
|
||
return (x_float * cos + x_rotated * sin).astype(x.dtype)
|
||
|
||
if is_neox_style:
|
||
x1, x2 = mx.split(x_float, 2, axis=-1)
|
||
o1 = x1 * cos - x2 * sin
|
||
o2 = x2 * cos + x1 * sin
|
||
return mx.concatenate([o1, o2], axis=-1).astype(x.dtype)
|
||
|
||
x1 = x_float[..., ::2]
|
||
x2 = x_float[..., 1::2]
|
||
o1 = x1 * cos - x2 * sin
|
||
o2 = x2 * cos + x1 * sin
|
||
return mx.stack([o1, o2], axis=-1).reshape(*x.shape).astype(x.dtype)
|
||
|
||
|
||
_WINDOWED_ATTENTION_WARNED = False
|
||
|
||
|
||
def _warn_windowed_attention_once(window: int) -> None:
|
||
"""Warn that FASTVIDEO_MLX_WINDOW degrades output on a dense-trained DiT.
|
||
|
||
Sliding-window self-attention is fast (6.6x at a +-3-frame window on 1.3B)
|
||
but these checkpoints were trained with dense attention, and restricting it
|
||
at inference produces heavy colour-block noise: structural agreement with
|
||
the dense baseline drops to 0.25 at +-3 frames and 0.03 at +-5. Sparsity of
|
||
this kind is a training-time method. Kept as a research knob, but it should
|
||
never be on by accident.
|
||
"""
|
||
global _WINDOWED_ATTENTION_WARNED
|
||
if _WINDOWED_ATTENTION_WARNED:
|
||
return
|
||
_WINDOWED_ATTENTION_WARNED = True
|
||
logger.warning(
|
||
"FASTVIDEO_MLX_WINDOW=%d enables sliding-window self-attention. These "
|
||
"checkpoints are trained dense; expect severely degraded output. This is "
|
||
"a research knob, not a speed setting — use --fast-spatial for real "
|
||
"denoise savings.",
|
||
window,
|
||
)
|
||
|
||
|
||
def gelu_tanh(x):
|
||
"""tanh-approximate GELU, as used by Wan's FFN.
|
||
|
||
``mlx.nn.gelu_approx`` is the same tanh approximation behind a fused
|
||
kernel. On the 1.3B FFN shape (32760x8960) it is bit-identical to the
|
||
expanded expression below and 3.3x faster — 28.9ms -> 8.7ms per layer,
|
||
which is 0.6s per denoise step across 30 layers.
|
||
"""
|
||
import mlx.nn as nn
|
||
|
||
return nn.gelu_approx(x)
|
||
|
||
|
||
def silu(x):
|
||
import mlx.core as mx
|
||
|
||
return x * mx.sigmoid(x)
|
||
|
||
|
||
def timestep_embedding(t, dim: int, max_period: int = 10000):
|
||
import mlx.core as mx
|
||
|
||
half = dim // 2
|
||
freqs = mx.exp(-math.log(max_period) * mx.arange(0, half, dtype=mx.float32) / half)
|
||
args = t[:, None].astype(mx.float32) * freqs[None]
|
||
embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
|
||
if dim % 2:
|
||
embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1)
|
||
return embedding
|
||
|
||
|
||
def scale_residual(residual, x, gate):
|
||
return residual + x * gate
|
||
|
||
|
||
def scale_residual_layer_norm_scale_shift(residual, x, gate, shift, scale, weight=None, bias=None, eps: float = 1e-6):
|
||
if isinstance(gate, int):
|
||
assert gate == 1
|
||
residual_output = residual + x
|
||
else:
|
||
residual_output = residual + x * gate
|
||
normalized = layer_norm(residual_output, weight=weight, bias=bias, eps=eps)
|
||
modulated = normalized * (1.0 + scale) + shift
|
||
return modulated, residual_output
|
||
|
||
|
||
class MLXWanT2VCrossAttention:
|
||
|
||
def __init__(self, weights: dict[str, mx.array], *, dim: int, num_heads: int, eps: float = 1e-6) -> None:
|
||
self.weights = weights
|
||
self.dim = dim
|
||
self.num_heads = num_heads
|
||
self.head_dim = dim // num_heads
|
||
self.eps = eps
|
||
|
||
def __call__(self, x, context):
|
||
import mlx.core as mx
|
||
|
||
batch = x.shape[0]
|
||
q = linear(x, self.weights["attn2.to_q.weight"], self.weights.get("attn2.to_q.bias"))
|
||
q = rms_norm(q, self.weights["attn2.norm_q.weight"], eps=self.eps).reshape(batch, -1, self.num_heads,
|
||
self.head_dim)
|
||
|
||
if context.shape[1] == 0:
|
||
attended = mx.zeros_like(q)
|
||
else:
|
||
k = linear(context, self.weights["attn2.to_k.weight"], self.weights.get("attn2.to_k.bias"))
|
||
k = rms_norm(k, self.weights["attn2.norm_k.weight"],
|
||
eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim)
|
||
v = linear(context, self.weights["attn2.to_v.weight"],
|
||
self.weights.get("attn2.to_v.bias")).reshape(batch, -1, self.num_heads, self.head_dim)
|
||
attended = mx.fast.scaled_dot_product_attention(
|
||
q.transpose(0, 2, 1, 3),
|
||
k.transpose(0, 2, 1, 3),
|
||
v.transpose(0, 2, 1, 3),
|
||
scale=self.head_dim**-0.5,
|
||
).transpose(0, 2, 1, 3)
|
||
|
||
attended = attended.reshape(batch, -1, self.dim)
|
||
return linear(attended, self.weights["attn2.to_out.weight"], self.weights.get("attn2.to_out.bias"))
|
||
|
||
|
||
class MLXWanTransformerBlock:
|
||
"""Dense T2V Wan transformer block for the experimental MLX runtime.
|
||
|
||
This mirrors the non-VSA PyTorch block for single-process dense attention.
|
||
Rotary embeddings and sequence-parallel paths are intentionally left out of
|
||
this first parity target.
|
||
"""
|
||
|
||
def __init__(self, weights: dict[str, mx.array], *, dim: int, ffn_dim: int, num_heads: int, eps: float = 1e-6):
|
||
self.weights = weights
|
||
self.dim = dim
|
||
self.ffn_dim = ffn_dim
|
||
self.num_heads = num_heads
|
||
self.head_dim = dim // num_heads
|
||
self.eps = eps
|
||
self.attn2 = MLXWanT2VCrossAttention(weights, dim=dim, num_heads=num_heads, eps=eps)
|
||
|
||
def __call__(self, hidden_states, encoder_hidden_states, temb, freqs_cis=None):
|
||
import mlx.core as mx
|
||
|
||
orig_dtype = hidden_states.dtype
|
||
e = self.weights["scale_shift_table"] + temb.astype(mx.float32)
|
||
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = mx.split(e, 6, axis=1)
|
||
|
||
norm_hidden_states = layer_norm(hidden_states.astype(mx.float32), eps=self.eps)
|
||
norm_hidden_states = (norm_hidden_states * (1.0 + scale_msa) + shift_msa).astype(orig_dtype)
|
||
|
||
query = linear(norm_hidden_states, self.weights["to_q.weight"], self.weights.get("to_q.bias"))
|
||
key = linear(norm_hidden_states, self.weights["to_k.weight"], self.weights.get("to_k.bias"))
|
||
value = linear(norm_hidden_states, self.weights["to_v.weight"], self.weights.get("to_v.bias"))
|
||
|
||
query = rms_norm(query, self.weights["norm_q.weight"],
|
||
eps=self.eps).reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim)
|
||
key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(hidden_states.shape[0], -1,
|
||
self.num_heads, self.head_dim)
|
||
value = value.reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim)
|
||
|
||
if freqs_cis is not None:
|
||
cos, sin = freqs_cis
|
||
query = apply_rotary_emb(query, cos, sin, is_neox_style=False)
|
||
key = apply_rotary_emb(key, cos, sin, is_neox_style=False)
|
||
|
||
# Self-attention only. FASTVIDEO_MLX_WINDOW=0/unset → full SDPA (byte-identical
|
||
# to the historical path). When >0, use chunked symmetric sliding-window
|
||
# attention (see windowed_attention.py). Cross-attn (attn2) stays full.
|
||
# Optional FASTVIDEO_MLX_WINDOW_SINK (default 0) adds global sink tokens.
|
||
q_bh = query.transpose(0, 2, 1, 3) # (B, H, S, D)
|
||
k_bh = key.transpose(0, 2, 1, 3)
|
||
v_bh = value.transpose(0, 2, 1, 3)
|
||
scale = self.head_dim**-0.5
|
||
window = envs.FASTVIDEO_MLX_WINDOW.get()
|
||
if window > 0:
|
||
from fastvideo.mlx_runtime.windowed_attention import windowed_attention
|
||
|
||
_warn_windowed_attention_once(window)
|
||
sink = envs.FASTVIDEO_MLX_WINDOW_SINK.get()
|
||
attn_output = windowed_attention(q_bh, k_bh, v_bh, window=window, sink=sink, scale=scale)
|
||
else:
|
||
attn_output = mx.fast.scaled_dot_product_attention(q_bh, k_bh, v_bh, scale=scale)
|
||
attn_output = attn_output.transpose(0, 2, 1, 3)
|
||
attn_output = attn_output.reshape(hidden_states.shape[0], -1, self.dim)
|
||
attn_output = linear(attn_output, self.weights["to_out.weight"], self.weights.get("to_out.bias"))
|
||
|
||
norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift(
|
||
hidden_states,
|
||
attn_output,
|
||
gate_msa,
|
||
0.0,
|
||
0.0,
|
||
weight=self.weights["self_attn_residual_norm.norm.weight"],
|
||
bias=self.weights["self_attn_residual_norm.norm.bias"],
|
||
eps=self.eps,
|
||
)
|
||
norm_hidden_states = norm_hidden_states.astype(orig_dtype)
|
||
hidden_states = hidden_states.astype(orig_dtype)
|
||
|
||
attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
|
||
norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift(
|
||
hidden_states,
|
||
attn_output,
|
||
1,
|
||
c_shift_msa,
|
||
c_scale_msa,
|
||
eps=self.eps,
|
||
)
|
||
norm_hidden_states = norm_hidden_states.astype(orig_dtype)
|
||
hidden_states = hidden_states.astype(orig_dtype)
|
||
|
||
ff_output = linear(norm_hidden_states, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias"))
|
||
ff_output = gelu_tanh(ff_output)
|
||
ff_output = linear(ff_output, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias"))
|
||
hidden_states = scale_residual(hidden_states, ff_output, c_gate_msa)
|
||
return hidden_states.astype(orig_dtype)
|
||
|
||
|
||
def mlx_block_weights_from_torch(torch_block) -> dict[str, mx.array]:
|
||
return {name: torch_to_mx(value) for name, value in torch_block.state_dict().items()}
|
||
|
||
|
||
class MLXWanDiT:
|
||
"""Experimental FP16 Wan/FastWan DiT forward path in MLX."""
|
||
|
||
def __init__(
|
||
self,
|
||
weights: dict[str, mx.array],
|
||
blocks: list[MLXWanTransformerBlock],
|
||
config: dict,
|
||
*,
|
||
compile: bool = False,
|
||
) -> None:
|
||
|
||
self.weights = weights
|
||
self.blocks = blocks
|
||
self.config = config
|
||
self.num_heads = int(config["num_attention_heads"])
|
||
self.head_dim = int(config["attention_head_dim"])
|
||
self.hidden_size = self.num_heads * self.head_dim
|
||
self.ffn_dim = int(config["ffn_dim"])
|
||
self.in_channels = int(config["in_channels"])
|
||
self.out_channels = int(config["out_channels"])
|
||
self.patch_size = tuple(config["patch_size"])
|
||
self.freq_dim = int(config["freq_dim"])
|
||
# Opt-in graph fusion. With fixed weights and static shapes, the whole
|
||
# denoise-step forward is a pure function of (latents, timestep) -- a
|
||
# good mx.compile target. Off by default so the eager path stays the
|
||
# baseline; enable via constructor or FASTVIDEO_MLX_COMPILE=1 and verify
|
||
# with the benchmark's SSIM ~= 1.0 check.
|
||
self._enable_compile = compile or envs.FASTVIDEO_MLX_COMPILE.get()
|
||
self._compiled_forward: Callable[..., Any] | None = None
|
||
self._compiled_signature: tuple | None = None
|
||
|
||
def patch_embed(self, hidden_states):
|
||
batch, channels, frames, height, width = hidden_states.shape
|
||
pt, ph, pw = self.patch_size
|
||
patch_dim = channels * pt * ph * pw
|
||
x = hidden_states.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw)
|
||
x = x.transpose(0, 2, 4, 6, 1, 3, 5, 7).reshape(batch, -1, patch_dim)
|
||
return linear(x, self.weights["patch_embedding.weight"], self.weights.get("patch_embedding.bias"))
|
||
|
||
def condition(self, timestep, encoder_hidden_states):
|
||
t_freq = timestep_embedding(timestep, self.freq_dim).astype(
|
||
weight_dtype(self.weights["condition_embedder.time_embedder.linear_1.weight"]))
|
||
temb = linear(
|
||
t_freq,
|
||
self.weights["condition_embedder.time_embedder.linear_1.weight"],
|
||
self.weights["condition_embedder.time_embedder.linear_1.bias"],
|
||
)
|
||
temb = silu(temb)
|
||
temb = linear(
|
||
temb,
|
||
self.weights["condition_embedder.time_embedder.linear_2.weight"],
|
||
self.weights["condition_embedder.time_embedder.linear_2.bias"],
|
||
)
|
||
timestep_proj = silu(temb)
|
||
timestep_proj = linear(
|
||
timestep_proj,
|
||
self.weights["condition_embedder.time_proj.weight"],
|
||
self.weights["condition_embedder.time_proj.bias"],
|
||
).reshape(timestep.shape[0], 6, self.hidden_size)
|
||
|
||
encoder_hidden_states = linear(
|
||
encoder_hidden_states,
|
||
self.weights["condition_embedder.text_embedder.linear_1.weight"],
|
||
self.weights["condition_embedder.text_embedder.linear_1.bias"],
|
||
)
|
||
encoder_hidden_states = gelu_tanh(encoder_hidden_states)
|
||
encoder_hidden_states = linear(
|
||
encoder_hidden_states,
|
||
self.weights["condition_embedder.text_embedder.linear_2.weight"],
|
||
self.weights["condition_embedder.text_embedder.linear_2.bias"],
|
||
)
|
||
return temb, timestep_proj, encoder_hidden_states
|
||
|
||
def output(self, hidden_states, temb, *, batch: int, frames: int, height: int, width: int):
|
||
pt, ph, pw = self.patch_size
|
||
post_patch_frames = frames // pt
|
||
post_patch_height = height // ph
|
||
post_patch_width = width // pw
|
||
shift, scale = mx_split_two(self.weights["scale_shift_table"] + temb[:, None, :], axis=1)
|
||
hidden_states = layer_norm(hidden_states, eps=float(self.config["eps"])) * (1.0 + scale) + shift
|
||
hidden_states = hidden_states.astype(weight_dtype(self.weights["proj_out.weight"]))
|
||
hidden_states = linear(hidden_states, self.weights["proj_out.weight"], self.weights["proj_out.bias"])
|
||
hidden_states = hidden_states.reshape(
|
||
batch,
|
||
post_patch_frames,
|
||
post_patch_height,
|
||
post_patch_width,
|
||
pt,
|
||
ph,
|
||
pw,
|
||
self.out_channels,
|
||
)
|
||
hidden_states = hidden_states.transpose(0, 7, 1, 4, 2, 5, 3, 6)
|
||
return hidden_states.reshape(batch, self.out_channels, frames, height, width)
|
||
|
||
def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin):
|
||
"""Pure forward used both eagerly and as the mx.compile target.
|
||
|
||
``cos``/``sin`` are passed as separate array args (rather than a tuple)
|
||
so the function traces cleanly under mx.compile.
|
||
"""
|
||
batch, _, frames, height, width = hidden_states.shape
|
||
freqs_cis = (cos, sin) if cos is not None else None
|
||
hidden_states = self.patch_embed(hidden_states)
|
||
temb, timestep_proj, encoder_hidden_states = self.condition(timestep, encoder_hidden_states)
|
||
for block in self.blocks:
|
||
hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis=freqs_cis)
|
||
return self.output(hidden_states, temb, batch=batch, frames=frames, height=height, width=width)
|
||
|
||
def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis):
|
||
cos, sin = freqs_cis if freqs_cis is not None else (None, None)
|
||
if self._enable_compile and cos is not None:
|
||
import mlx.core as mx
|
||
|
||
# mx.compile keeps one traced graph per input signature, and each
|
||
# graph pins its own materialization of the quantized weights. The
|
||
# two-pass modes (--refine) call the DiT at a second resolution, so
|
||
# keeping both graphs alive doubles resident DiT memory: 14B refine
|
||
# peaked at 34.7 GiB instead of 20.8 GiB. Retire the previous graph
|
||
# when the signature changes; the retrace costs far less than a
|
||
# second copy of the weights.
|
||
signature = (hidden_states.shape, encoder_hidden_states.shape, timestep.shape)
|
||
if self._compiled_forward is not None and signature != self._compiled_signature:
|
||
self._compiled_forward = None
|
||
self._compiled_signature = None
|
||
mx.clear_cache()
|
||
if self._compiled_forward is None:
|
||
self._compiled_forward = mx.compile(self._forward)
|
||
self._compiled_signature = signature
|
||
compiled_forward = self._compiled_forward
|
||
try:
|
||
return compiled_forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
|
||
except Exception as exc: # noqa: BLE001 - some quant graphs may not trace; fall back to eager.
|
||
logger.warning("mx.compile forward failed (%s); falling back to eager execution.", exc)
|
||
self._enable_compile = False
|
||
self._compiled_forward = None
|
||
self._compiled_signature = None
|
||
return self._forward(hidden_states, encoder_hidden_states, timestep, cos, sin)
|
||
|
||
|
||
def mx_split_two(x, *, axis: int):
|
||
import mlx.core as mx
|
||
|
||
left, right = mx.split(x, 2, axis=axis)
|
||
return left, right
|
||
|
||
|
||
def _load_safetensor_value(handle, name: str):
|
||
return handle.get_tensor(name)
|
||
|
||
|
||
def _load_mx_array_from_safetensor(handle, name: str, dtype):
|
||
"""Load a safetensors value and cast before creating the MLX array.
|
||
|
||
The FastWan Diffusers checkpoint is fp32. Creating an MLX array first and
|
||
then casting it to fp16 briefly materializes a large fp32 MLX allocation.
|
||
Casting the CPU tensor before crossing into MLX keeps the transient GPU-side
|
||
footprint lower.
|
||
"""
|
||
import mlx.core as mx
|
||
import torch
|
||
|
||
tensor = handle.get_tensor(name)
|
||
if dtype == mx.float16:
|
||
tensor = tensor.to(torch.float16)
|
||
elif dtype == mx.float32:
|
||
tensor = tensor.to(torch.float32)
|
||
elif dtype == mx.bfloat16:
|
||
# NumPy has no bfloat16, so bridge through fp32 and cast on-device below.
|
||
tensor = tensor.to(torch.float32)
|
||
array = mx.array(tensor.numpy())
|
||
del tensor
|
||
if dtype is not None and array.dtype != dtype:
|
||
array = array.astype(dtype)
|
||
mx.eval(array)
|
||
return array
|
||
|
||
|
||
def _eval_loaded_weight(value) -> None:
|
||
import mlx.core as mx
|
||
|
||
if isinstance(value, QuantizedMatrix):
|
||
eval_args = [value.weight, value.scales]
|
||
if value.biases is not None:
|
||
eval_args.append(value.biases)
|
||
mx.eval(*eval_args)
|
||
else:
|
||
mx.eval(value)
|
||
|
||
|
||
# Diffusers-to-FastVideo key mapping for WanTransformerBlock weights.
|
||
# Shared by both MLX and torch block loaders to keep mappings synchronized.
|
||
_WAN_BLOCK_KEY_MAP = {
|
||
"scale_shift_table": "scale_shift_table",
|
||
"attn1.to_q.weight": "to_q.weight",
|
||
"attn1.to_q.bias": "to_q.bias",
|
||
"attn1.to_k.weight": "to_k.weight",
|
||
"attn1.to_k.bias": "to_k.bias",
|
||
"attn1.to_v.weight": "to_v.weight",
|
||
"attn1.to_v.bias": "to_v.bias",
|
||
"attn1.to_out.0.weight": "to_out.weight",
|
||
"attn1.to_out.0.bias": "to_out.bias",
|
||
"attn1.norm_q.weight": "norm_q.weight",
|
||
"attn1.norm_k.weight": "norm_k.weight",
|
||
"attn2.to_q.weight": "attn2.to_q.weight",
|
||
"attn2.to_q.bias": "attn2.to_q.bias",
|
||
"attn2.to_k.weight": "attn2.to_k.weight",
|
||
"attn2.to_k.bias": "attn2.to_k.bias",
|
||
"attn2.to_v.weight": "attn2.to_v.weight",
|
||
"attn2.to_v.bias": "attn2.to_v.bias",
|
||
"attn2.to_out.0.weight": "attn2.to_out.weight",
|
||
"attn2.to_out.0.bias": "attn2.to_out.bias",
|
||
"attn2.norm_q.weight": "attn2.norm_q.weight",
|
||
"attn2.norm_k.weight": "attn2.norm_k.weight",
|
||
"ffn.net.0.proj.weight": "ffn.fc_in.weight",
|
||
"ffn.net.0.proj.bias": "ffn.fc_in.bias",
|
||
"ffn.net.2.weight": "ffn.fc_out.weight",
|
||
"ffn.net.2.bias": "ffn.fc_out.bias",
|
||
"norm2.weight": "self_attn_residual_norm.norm.weight",
|
||
"norm2.bias": "self_attn_residual_norm.norm.bias",
|
||
}
|
||
|
||
|
||
def mlx_block_weights_from_diffusers_safetensors(
|
||
checkpoint_path: str | Path,
|
||
*,
|
||
block_index: int = 0,
|
||
quantization: str | MLXQuantizationSpec | None = None,
|
||
dtype=None,
|
||
) -> dict[str, mx.array]:
|
||
"""Load one Diffusers-format Wan block into the MLX dense-block key layout."""
|
||
from safetensors import safe_open
|
||
|
||
prefix = f"blocks.{block_index}."
|
||
key_map = _WAN_BLOCK_KEY_MAP
|
||
|
||
spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
|
||
or isinstance(quantization, str)) else quantization
|
||
ensure_quantization_supported(spec)
|
||
matrix_targets = {target for target in key_map.values() if target.endswith(".weight") and "norm" not in target}
|
||
weights = {}
|
||
with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
|
||
available = set(handle.keys())
|
||
for source_name, target_name in key_map.items():
|
||
full = prefix + source_name
|
||
if full not in available:
|
||
# Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
|
||
# The block forward already fetches biases via ``.get(...)``.
|
||
if source_name.endswith(".bias"):
|
||
continue
|
||
raise KeyError(f"missing required block weight: {full}")
|
||
array = _load_mx_array_from_safetensor(handle, full, dtype)
|
||
loaded = quantize_matrix(array, spec) if target_name in matrix_targets else array
|
||
_eval_loaded_weight(loaded)
|
||
weights[target_name] = loaded
|
||
del array
|
||
return weights
|
||
|
||
|
||
def mlx_dit_from_diffusers_safetensors(
|
||
checkpoint_path: str | Path,
|
||
config_path: str | Path,
|
||
*,
|
||
dtype: str = "fp16",
|
||
num_blocks: int | None = None,
|
||
quantization: str | MLXQuantizationSpec | None = None,
|
||
compile: bool = False,
|
||
) -> MLXWanDiT:
|
||
import mlx.core as mx
|
||
from safetensors import safe_open
|
||
|
||
from fastvideo.mlx_runtime.checkpoint_compat import raise_if_unsupported_mlx_checkpoint
|
||
|
||
raise_if_unsupported_mlx_checkpoint(checkpoint_path, config_path)
|
||
|
||
config = json.loads(Path(config_path).read_text())
|
||
total_blocks = int(config["num_layers"])
|
||
if num_blocks is None:
|
||
num_blocks = total_blocks
|
||
cast_dtype = {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[dtype]
|
||
spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
|
||
or isinstance(quantization, str)) else quantization
|
||
ensure_quantization_supported(spec)
|
||
|
||
top_level_names = [
|
||
"patch_embedding.weight",
|
||
"patch_embedding.bias",
|
||
"condition_embedder.time_embedder.linear_1.weight",
|
||
"condition_embedder.time_embedder.linear_1.bias",
|
||
"condition_embedder.time_embedder.linear_2.weight",
|
||
"condition_embedder.time_embedder.linear_2.bias",
|
||
"condition_embedder.time_proj.weight",
|
||
"condition_embedder.time_proj.bias",
|
||
"condition_embedder.text_embedder.linear_1.weight",
|
||
"condition_embedder.text_embedder.linear_1.bias",
|
||
"condition_embedder.text_embedder.linear_2.weight",
|
||
"condition_embedder.text_embedder.linear_2.bias",
|
||
"scale_shift_table",
|
||
"proj_out.weight",
|
||
"proj_out.bias",
|
||
]
|
||
weights = {}
|
||
with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
|
||
available = set(handle.keys())
|
||
for name in top_level_names:
|
||
if name not in available:
|
||
if name.endswith(".bias"):
|
||
continue
|
||
raise KeyError(f"missing required weight: {name}")
|
||
array = _load_mx_array_from_safetensor(handle, name, cast_dtype)
|
||
if name == "patch_embedding.weight":
|
||
array = array.reshape(int(config["num_attention_heads"]) * int(config["attention_head_dim"]), -1)
|
||
if name.endswith(".weight") and name not in {"scale_shift_table"}:
|
||
loaded = quantize_matrix(array, spec)
|
||
else:
|
||
loaded = array
|
||
_eval_loaded_weight(loaded)
|
||
weights[name] = loaded
|
||
del array
|
||
|
||
blocks = []
|
||
for block_index in range(num_blocks):
|
||
block_weights = mlx_block_weights_from_diffusers_safetensors(
|
||
checkpoint_path,
|
||
block_index=block_index,
|
||
quantization=spec,
|
||
dtype=cast_dtype,
|
||
)
|
||
block_weights = {
|
||
name: (value if isinstance(value, QuantizedMatrix) else value.astype(cast_dtype))
|
||
for name, value in block_weights.items()
|
||
}
|
||
for value in block_weights.values():
|
||
_eval_loaded_weight(value)
|
||
blocks.append(
|
||
MLXWanTransformerBlock(
|
||
block_weights,
|
||
dim=int(config["num_attention_heads"]) * int(config["attention_head_dim"]),
|
||
ffn_dim=int(config["ffn_dim"]),
|
||
num_heads=int(config["num_attention_heads"]),
|
||
eps=float(config["eps"]),
|
||
))
|
||
return MLXWanDiT(weights, blocks, config, compile=compile)
|
||
|
||
|
||
def torch_block_state_from_diffusers_safetensors(
|
||
checkpoint_path: str | Path,
|
||
*,
|
||
block_index: int = 0,
|
||
) -> dict[str, torch.Tensor]:
|
||
"""Load one Diffusers-format Wan block into FastVideo's dense block keys."""
|
||
from safetensors import safe_open
|
||
|
||
prefix = f"blocks.{block_index}."
|
||
key_map = _WAN_BLOCK_KEY_MAP
|
||
|
||
state = {}
|
||
with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle:
|
||
available = set(handle.keys())
|
||
for source_name, target_name in key_map.items():
|
||
full = prefix + source_name
|
||
if full not in available:
|
||
# Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN.
|
||
if source_name.endswith(".bias"):
|
||
continue
|
||
raise KeyError(f"missing required block weight: {full}")
|
||
state[target_name] = handle.get_tensor(full).float()
|
||
return state
|