--- title: minimax_h3 ---

MiniMax H3 joint audio-video DiT for the Apple Silicon MLX runtime.

A faithful MLX port of the upstream CUDA reference:

Apple Silicon memory controls:

Checkpoint layout: the released H3 checkpoint uses the diffusers reference module names 1:1 (transformer_blocks.{i}.attn.to_out.0.weight, ff.net.0.proj.weight, time_embedder.linear_1.weight, ...). This module keeps those names as the MLX weight keys — no renaming contract. proj_in / audio_proj_in / time_embedder / proj_out / audio_proj_out are fp32 in the release and are kept fp32 here.

Nothing in this file requires the CUDA stack; it imports fastvideo.logger only. Parity tests against the torch reference live in fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py.

fastvideo.mlx_runtime.minimax_h3.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3PreparedReference

MiniMaxH3PreparedReference(media_type: Literal['image', 'video', 'audio'], has_audio: bool = False, image: Image.Image | None = None, frames: np.ndarray | None = None, waveform: np.ndarray | None = None, block_timestamps: list[float] = list(), num_latent_frames: int = 1, latent_height: int = 0, latent_width: int = 0, num_audio_latents: int = 0)

Decoded Ref2VA media and the latent geometry resolved by condition encoding.

waveform is (channels, samples) float32, not a tensor: the MLX side never needs a device-resident waveform, the audio VAE converts on encode.

fastvideo.mlx_runtime.minimax_h3.MLXQuantizationSpec

MLXQuantizationSpec(mode: str, bits: int | None = None, group_size: int | None = None)

MLX quantized-matmul configuration for DiT linear weights.

fastvideo.mlx_runtime.minimax_h3.MLXQuantizationSpec.from_name

from_name(name: str | None) -> MLXQuantizationSpec | None

fastvideo.mlx_runtime.minimax_h3.QuantizedMatrix

QuantizedMatrix(weight: mx.array, scales: mx.array, biases: mx.array | None, spec: MLXQuantizationSpec, dequantized_dtype: mx.Dtype)

fastvideo.mlx_runtime.minimax_h3.ensure_quantization_supported

ensure_quantization_supported(spec: MLXQuantizationSpec | None) -> None

Raise :class:UnsupportedMLXQuantizationError if spec cannot run here.

fastvideo.mlx_runtime.minimax_h3.quantize_matrix

quantize_matrix(weight, spec: MLXQuantizationSpec | None)

fastvideo.mlx_runtime.minimax_h3.silu

silu(x)

fastvideo.mlx_runtime.minimax_h3.timestep_embedding

timestep_embedding(t, dim: int, max_period: int = 10000)

fastvideo.mlx_runtime.minimax_h3.weight_dtype

weight_dtype(weight)

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAConfig

MiniMaxH3VSAConfig(enabled: bool = False, sparsity: float = 0.9, tile_size: int = 64, prefix_mode: PrefixMode = 'exempt', dense_first_n_steps: int = 0, dense_layers: tuple[int, ...] = (), impl: VSAImpl = 'auto', ref_policy: str | None = None, ref_keep_rate: float | None = None)

Runtime VSA knobs. Defaults preserve dense MLX H3 behavior.

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAConfig.effective_ref_keep_rate

effective_ref_keep_rate: float | None

The per-reference-video keep rate, or None when references stay dense.

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAConfig.layer_sparsity

layer_sparsity(layer_idx: int, step_index: int) -> float

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAGeometry

MiniMaxH3VSAGeometry(prefix_segments: tuple[int, ...], dit_seq_shape: tuple[int, int, int], tile_shape: tuple[int, int, int], tile_elems: int, total_seq_length: int, num_prefix_tiles: int, num_video_tiles: int, variable_block_sizes: np.ndarray, untile_combined_index: np.ndarray, tile_partition_indices: np.ndarray, video_tile_spans: tuple[tuple[int, int], ...] = (), segments: tuple[int | tuple[int, int, int], ...] = (), reference_regions: tuple[tuple[int, int, int], ...] = ())

Packed-sequence tile map shared by routing, reference, and Metal paths.

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAGeometry.dense_rows

dense_rows: np.ndarray

Packed positions of the dense rows, in prefix-tile order.

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAGeometry.dense_rows_lead

dense_rows_lead: bool

True when the dense rows are the first packed rows (no interleaved region).

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAStats

MiniMaxH3VSAStats(configured_sparsity: float = 0.0, layer_sparsity: float = 0.0, tile_size: int = 64, prefix_mode: str = 'exempt', impl: str = 'dense', num_prefix_tiles: int = 0, num_video_tiles: int = 0, video_keep: float = 0.0, achieved_sparsity: float = 0.0, dense_fallback_reason: str | None = None, attention_calls: int = 0, sparse_calls: int = 0, impl_counts: dict[str, int] = dict(), fallback_reasons: list[str] = list())

Filled during a sparse forward so the pipeline can report achieved sparsity.

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3VSAStats.record

record(call: MiniMaxH3VSAStats) -> None

Aggregate equally sized video-query tile maps across blocks and steps.

fastvideo.mlx_runtime.minimax_h3.build_h3_segment_geometry

build_h3_segment_geometry(segments: tuple[int | tuple[int, int, int], ...], tile_size: int = 64) -> MiniMaxH3VSAGeometry

Tile a packed sequence of interleaved dense segments and video regions.

segments lists, in packed order, dense row counts (int) and video token grids ((t, h, w)). Rows are permuted to [dense chunks][region 0 tiles][region 1 tiles]... exactly as the CUDA _h3_segment_tile_geometry does; the last region is the generated video.

fastvideo.mlx_runtime.minimax_h3.dense_only_vsa_error

dense_only_vsa_error(checkpoint_dir: str | Any) -> DenseOnlyVSACheckpointError

fastvideo.mlx_runtime.minimax_h3.dit_seq_shape_from_layout

dit_seq_shape_from_layout(layout: Any, patch_size: tuple[int, int, int]) -> tuple[int, int, int]

fastvideo.mlx_runtime.minimax_h3.h3_vsa_attention

h3_vsa_attention(query, key, value, geometry: MiniMaxH3VSAGeometry, *, sparsity: float, exempt: bool = True, gate_compress = None, impl: VSAImpl = 'auto', stats: MiniMaxH3VSAStats | None = None, ref_keep_rate: float | None = None)

Packed [S, H, D] VSA attention. Falls back to dense SDPA when sparsity is 0.

ref_keep_rate is the p2_multi_region keep rate of each reference-video region of geometry; geometries with reference regions require it.

fastvideo.mlx_runtime.minimax_h3.prefix_segments_from_layout

prefix_segments_from_layout(layout: Any, patch_size: tuple[int, int, int]) -> tuple[int, ...]

Segment sizes preceding the generated-video tail, matching the PyTorch stage.

Prefix tiles are chunked per segment, so a segment boundary must fall where the packed rows change kind. A Ref2VA layout interleaves a video reference's audio rows with the references themselves, so it is split per reference from layout.reference_segments; T2VA/FL2VA keep the single condition segment they have always had.

fastvideo.mlx_runtime.minimax_h3.ref2va_vsa_segments_from_layout

ref2va_vsa_segments_from_layout(layout: Any, patch_size: tuple[int, int, int]) -> tuple[int | tuple[int, int, int], ...]

Packed segments with every reference VIDEO as a sparse region (p2_multi_region).

Mirrors _h3_vsa_ref2va_segments in the CUDA denoising stage: [text | reference spans... | target audio | target video] where each reference video is a token grid and text, reference audio, image references and the target audio are dense row counts.

fastvideo.mlx_runtime.minimax_h3.linear

linear(x, weight, bias = None)

Run an H3 linear with its measured wide-row affine dispatch enabled.

Source
def linear(x, weight, bias=None):
    """Run an H3 linear with its measured wide-row affine dispatch enabled."""
    return _shared_linear(x, weight, bias, use_affine_dq_gemm=True)

fastvideo.mlx_runtime.minimax_h3.align_num_frames

align_num_frames(num_frames: int) -> int
Source
def align_num_frames(num_frames: int) -> int:
    if num_frames < 1:
        raise ValueError(f"`num_frames` must be positive, got {num_frames}.")
    while num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK:
        num_frames += 1
    return num_frames

fastvideo.mlx_runtime.minimax_h3.video_latent_num_frames

video_latent_num_frames(num_frames: int) -> int
Source
def video_latent_num_frames(num_frames: int) -> int:
    if num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK:
        raise ValueError(f"`num_frames` must be of the form 17 * n + 5, got {num_frames}.")
    return (num_frames - MINIMAX_H3_LATENTS_PER_CHUNK) // MINIMAX_H3_FRAMES_PER_CHUNK * MINIMAX_H3_LATENTS_PER_CHUNK + 2

fastvideo.mlx_runtime.minimax_h3.audio_latent_num_frames

audio_latent_num_frames(num_frames: int) -> int
Source
def audio_latent_num_frames(num_frames: int) -> int:
    return int(round(num_frames / MINIMAX_H3_FPS * MINIMAX_H3_AUDIO_LATENTS_PER_SECOND))

fastvideo.mlx_runtime.minimax_h3.spatial_position_grid

spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> np.ndarray
Source
def spatial_position_grid(dim: int, patch: int, sqrt_area: float) -> np.ndarray:
    ratio = dim / sqrt_area
    left = (1.0 - ratio) / 2.0
    # NumPy endpoint=False float64 arithmetic is the reference contract.
    return np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE

fastvideo.mlx_runtime.minimax_h3.temporal_position_grid

temporal_position_grid(num_latent_frames: int, origin: float) -> np.ndarray
Source
def temporal_position_grid(num_latent_frames: int, origin: float) -> np.ndarray:
    spans = np.array(
        [
            MINIMAX_H3_ROPE_FRAME_RESCALE *
            MINIMAX_H3_ROPE_FRAMES_PER_LATENT[index % len(MINIMAX_H3_ROPE_FRAMES_PER_LATENT)]
            for index in range(num_latent_frames)
        ],
        dtype=np.float64,
    )
    return origin + np.concatenate([np.zeros(1, dtype=np.float64), spans[:-1].cumsum()])

fastvideo.mlx_runtime.minimax_h3.patchify_video_latents

patchify_video_latents(latents: np.ndarray, patch_size: tuple[int, int, int]) -> np.ndarray

(B, C, T, H, W) -> (BT'H'W', Cptphpw), channel-major patch features.

Source
def patchify_video_latents(latents: np.ndarray, patch_size: tuple[int, int, int]) -> np.ndarray:
    """(B, C, T, H, W) -> (B*T'*H'*W', C*pt*ph*pw), channel-major patch features."""
    patch_t, patch_h, patch_w = patch_size
    batch_size, channels, num_frames, height, width = latents.shape
    if num_frames % patch_t or height % patch_h or width % patch_w:
        raise ValueError(f"Latents of shape {latents.shape} are not divisible by the patch {patch_size}.")
    latents = latents.reshape(
        batch_size,
        channels,
        num_frames // patch_t,
        patch_t,
        height // patch_h,
        patch_h,
        width // patch_w,
        patch_w,
    )
    latents = latents.transpose(0, 2, 4, 6, 1, 3, 5, 7)
    return np.ascontiguousarray(latents.reshape(-1, channels * patch_t * patch_h * patch_w))

fastvideo.mlx_runtime.minimax_h3.unpatchify_video_tokens

unpatchify_video_tokens(rows: np.ndarray, num_latent_frames: int, latent_height: int, latent_width: int, channels: int, patch_size: tuple[int, int, int]) -> np.ndarray
Source
def unpatchify_video_tokens(
    rows: np.ndarray,
    num_latent_frames: int,
    latent_height: int,
    latent_width: int,
    channels: int,
    patch_size: tuple[int, int, int],
) -> np.ndarray:
    patch_t, patch_h, patch_w = patch_size
    rows = rows.reshape(
        -1,
        num_latent_frames // patch_t,
        latent_height // patch_h,
        latent_width // patch_w,
        channels,
        patch_t,
        patch_h,
        patch_w,
    )
    rows = rows.transpose(0, 4, 1, 5, 2, 6, 3, 7)
    return np.ascontiguousarray(rows.reshape(-1, channels, num_latent_frames, latent_height, latent_width))

fastvideo.mlx_runtime.minimax_h3.unpack_audio_tokens

unpack_audio_tokens(rows: np.ndarray, num_audio_latents: int) -> np.ndarray

(2*num_audio_latents, feat) channel-major rows -> (2, feat, num_audio_latents).

Source
def unpack_audio_tokens(rows: np.ndarray, num_audio_latents: int) -> np.ndarray:
    """(2*num_audio_latents, feat) channel-major rows -> (2, feat, num_audio_latents)."""
    rows = rows.reshape(MINIMAX_H3_AUDIO_CHANNELS, num_audio_latents, rows.shape[-1])
    return np.ascontiguousarray(rows.transpose(0, 2, 1))

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3PackedLayout

MiniMaxH3PackedLayout(sequence_length: int, position_ids: np.ndarray, token_tags: np.ndarray, video_indices: np.ndarray, audio_indices: np.ndarray, text_indices: np.ndarray, num_condition_video_rows: int, num_condition_audio_rows: int, num_video_latent_frames: int, latent_height: int, latent_width: int, num_audio_latents: int, reference_segments: tuple[tuple[str, int, tuple[int, int, int]], ...] = ())

One packed joint sequence and the geometry needed to interpret it.

Arrays are NumPy (position_ids in float64, indices int64) and converted to MLX at the model boundary.

Source
@dataclass(frozen=True)
class MiniMaxH3PackedLayout:
    """One packed joint sequence and the geometry needed to interpret it.

    Arrays are NumPy (position_ids in float64, indices int64) and converted
    to MLX at the model boundary.
    """

    sequence_length: int
    position_ids: np.ndarray
    token_tags: np.ndarray
    video_indices: np.ndarray
    audio_indices: np.ndarray
    text_indices: np.ndarray
    num_condition_video_rows: int
    num_condition_audio_rows: int
    num_video_latent_frames: int
    latent_height: int
    latent_width: int
    num_audio_latents: int
    # Per-reference conditioning spans in packed order (Ref2VA only; empty
    # elsewhere): (kind, row_count, latent_shape) covering exactly the rows
    # between the text tokens and the target audio. kind is "audio", "image",
    # or "video"; latent_shape is the reference's raw latent (frames, height,
    # width) for visual kinds and (0, 0, 0) for audio. A video reference with
    # sound contributes two entries, its audio rows first. VSA-H3 uses them to
    # tile (and sparsify) reference-video regions in place.
    reference_segments: tuple[tuple[str, int, tuple[int, int, int]], ...] = ()

fastvideo.mlx_runtime.minimax_h3.build_packed_layout

build_packed_layout(num_text_tokens: int, num_latent_frames: int, latent_height: int, latent_width: int, num_audio_latents: int, patch_size: tuple[int, int, int] = (1, 2, 2), keyframe_anchors: tuple[str, ...] = (), text_token_tags: np.ndarray | None = None, video_temporal_scale: float = 1.0) -> MiniMaxH3PackedLayout

Build the [text | condition | audio | video] layout (T2VA / FL2VA).

Source
def build_packed_layout(
    num_text_tokens: int,
    num_latent_frames: int,
    latent_height: int,
    latent_width: int,
    num_audio_latents: int,
    patch_size: tuple[int, int, int] = (1, 2, 2),
    keyframe_anchors: tuple[str, ...] = (),
    text_token_tags: np.ndarray | None = None,
    video_temporal_scale: float = 1.0,
) -> MiniMaxH3PackedLayout:
    """Build the ``[text | condition | audio | video]`` layout (T2VA / FL2VA)."""
    _, patch_h, patch_w = patch_size
    if text_token_tags is None:
        text_token_tags = np.full(num_text_tokens, MINIMAX_H3_TEXT_TAG, dtype=np.int64)
    if text_token_tags.shape != (num_text_tokens, ):
        raise ValueError(f"text_token_tags must have shape ({num_text_tokens},), got {text_token_tags.shape}.")
    if not np.isin(text_token_tags, (MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG)).all():
        raise ValueError("text_token_tags may contain only text and vision tags.")
    if not np.isfinite(video_temporal_scale) or video_temporal_scale <= 0:
        raise ValueError(f"video_temporal_scale must be positive and finite, got {video_temporal_scale}.")
    if keyframe_anchors and video_temporal_scale != 1.0:
        raise ValueError("video_temporal_scale is currently supported only for T2VA without keyframe anchors.")

    rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w)
    num_condition_rows = len(keyframe_anchors) * rows_per_frame
    num_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
    num_video_rows = num_latent_frames * rows_per_frame
    sequence_length = num_text_tokens + num_condition_rows + num_audio_rows + num_video_rows

    condition_start = num_text_tokens
    audio_start = condition_start + num_condition_rows
    video_start = audio_start + num_audio_rows

    position_ids = np.zeros((sequence_length, 3), dtype=np.float64)
    position_ids[:num_text_tokens, 0] = np.arange(num_text_tokens, dtype=np.float64)

    sqrt_area = float(np.sqrt(latent_height * latent_width))
    height_grid = spatial_position_grid(latent_height, patch_h, sqrt_area)
    width_grid = spatial_position_grid(latent_width, patch_w, sqrt_area)
    mesh_h, mesh_w = np.meshgrid(height_grid, width_grid, indexing="ij")
    frame_grid = np.stack([mesh_h.reshape(-1), mesh_w.reshape(-1)], axis=-1)

    for index, anchor in enumerate(keyframe_anchors):
        if anchor == "first":
            anchor_time = float(num_text_tokens)
        elif anchor == "last":
            anchor_time = (float(num_text_tokens) + _temporal_position_span(num_latent_frames) -
                           MINIMAX_H3_ROPE_FRAME_RESCALE)
        else:
            raise ValueError(f"A keyframe anchor must be 'first' or 'last', got {anchor!r}.")
        rows = slice(condition_start + index * rows_per_frame, condition_start + (index + 1) * rows_per_frame)
        position_ids[rows, 0] = anchor_time
        position_ids[rows, 1:] = frame_grid

    audio_time = float(num_text_tokens) + np.arange(num_audio_latents, dtype=np.float64)
    position_ids[audio_start:video_start, 0] = np.tile(audio_time, MINIMAX_H3_AUDIO_CHANNELS)
    position_ids[audio_start:video_start, 2] = np.concatenate([
        np.full(num_audio_latents, float(width_grid[0]), dtype=np.float64),
        np.full(num_audio_rows - num_audio_latents, float(width_grid[-1]), dtype=np.float64),
    ])

    frame_time = temporal_position_grid(num_latent_frames, 0.0) * video_temporal_scale + float(num_text_tokens)
    position_ids[video_start:, 0] = np.repeat(frame_time, rows_per_frame)
    position_ids[video_start:, 1:] = np.tile(frame_grid, (num_latent_frames, 1))

    video_indices = np.concatenate([
        np.arange(condition_start, audio_start),
        np.arange(video_start, sequence_length),
    ])
    audio_indices = np.arange(audio_start, video_start)
    text_indices = np.arange(num_text_tokens)
    token_tags = np.empty(sequence_length, dtype=np.int64)
    token_tags[text_indices] = text_token_tags
    token_tags[audio_indices] = MINIMAX_H3_AUDIO_TAG
    token_tags[video_indices] = MINIMAX_H3_VIDEO_TAG

    return MiniMaxH3PackedLayout(
        sequence_length=sequence_length,
        position_ids=position_ids,
        token_tags=token_tags,
        video_indices=video_indices,
        audio_indices=audio_indices,
        text_indices=text_indices,
        num_condition_video_rows=num_condition_rows,
        num_condition_audio_rows=0,
        num_video_latent_frames=num_latent_frames,
        latent_height=latent_height,
        latent_width=latent_width,
        num_audio_latents=num_audio_latents,
    )

fastvideo.mlx_runtime.minimax_h3.build_ref2va_packed_layout

build_ref2va_packed_layout(text_token_tags: np.ndarray, references: list[MiniMaxH3PreparedReference], num_latent_frames: int, latent_height: int, latent_width: int, num_audio_latents: int, patch_size: tuple[int, int, int] = (1, 2, 2)) -> MiniMaxH3PackedLayout

Build the [text | ordered references | target audio | target video] layout.

Mirrors fastvideo.pipelines.basic.minimax_h3.packing.build_ref2va_packed_sequence so a reference-conditioned request packs the same rows, positions and regions on both runtimes. References keep their declared order and each one advances the shared rotary clock by what it occupies, which is what makes the target rows continue after them rather than overlap.

Source
def build_ref2va_packed_layout(
        text_token_tags: np.ndarray,
        references: list[MiniMaxH3PreparedReference],
        num_latent_frames: int,
        latent_height: int,
        latent_width: int,
        num_audio_latents: int,
        patch_size: tuple[int, int, int] = (1, 2, 2),
) -> MiniMaxH3PackedLayout:
    """Build the ``[text | ordered references | target audio | target video]`` layout.

    Mirrors ``fastvideo.pipelines.basic.minimax_h3.packing.build_ref2va_packed_sequence``
    so a reference-conditioned request packs the same rows, positions and
    regions on both runtimes. References keep their declared order and each one
    advances the shared rotary clock by what it occupies, which is what makes
    the target rows continue after them rather than overlap.
    """
    if patch_size != (1, 2, 2):
        raise ValueError(f"MiniMax-H3 Ref2VA requires patch_size=(1, 2, 2), got {patch_size}.")
    if text_token_tags.ndim != 1:
        raise ValueError(f"text_token_tags must be one-dimensional, got {text_token_tags.shape}.")
    if not np.isin(text_token_tags, (MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG)).all():
        raise ValueError("text_token_tags may contain only text and vision tags.")
    if not references:
        raise ValueError("Ref2VA requires at least one prepared reference.")

    patch_t, patch_h, patch_w = patch_size
    target_geometry = (num_latent_frames, latent_height, latent_width)
    if any(value <= 0 for value in target_geometry) or num_audio_latents <= 0:
        raise ValueError("Ref2VA target latent geometry must be positive.")
    if any(value % patch for value, patch in zip(target_geometry, patch_size, strict=True)):
        raise ValueError(f"Target geometry {target_geometry} is not divisible by patch {patch_size}.")

    visual_row_counts = [
        0 if reference.media_type == "audio" else _num_reference_video_rows(reference, patch_size)
        for reference in references
    ]
    for reference in references:
        if reference.media_type == "audio" and not reference.has_audio:
            raise ValueError("An audio reference must carry decoded waveform latents.")
        if reference.has_audio and reference.num_audio_latents <= 0:
            raise ValueError("An audio-bearing reference has no resolved audio latents.")

    num_text_tokens = int(text_token_tags.shape[0])
    num_target_video_rows = (num_latent_frames // patch_t) * (latent_height // patch_h) * (latent_width // patch_w)
    num_target_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
    num_reference_video_rows = sum(visual_row_counts)
    num_reference_audio_rows = sum(reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS for reference in references
                                   if reference.has_audio)
    sequence_length = (num_text_tokens + num_reference_video_rows + num_reference_audio_rows + num_target_audio_rows +
                       num_target_video_rows)

    position_ids = np.zeros((sequence_length, 3), dtype=np.float64)
    position_ids[:num_text_tokens, 0] = np.arange(num_text_tokens, dtype=np.float64)
    target_frame_grid, target_width_grid = _frame_position_grid(latent_height, latent_width, patch_h, patch_w)

    video_indices: list[np.ndarray] = []
    audio_indices: list[np.ndarray] = []
    reference_segments: list[tuple[str, int, tuple[int, int, int]]] = []
    cursor = num_text_tokens
    rotary_time = float(num_text_tokens)
    for reference, visual_row_count in zip(references, visual_row_counts, strict=True):
        visual_latent_shape = (reference.num_latent_frames, reference.latent_height, reference.latent_width)
        if reference.media_type == "image":
            reference_segments.append(("image", visual_row_count, visual_latent_shape))
            rows = slice(cursor, cursor + visual_row_count)
            cursor = rows.stop
            video_indices.append(np.arange(rows.start, rows.stop))
            frame_grid, _ = _frame_position_grid(
                reference.latent_height,
                reference.latent_width,
                patch_h,
                patch_w,
            )
            position_ids[rows, 0] = rotary_time
            position_ids[rows, 1:] = frame_grid
            rotary_time += 1.0
        elif reference.media_type == "audio":
            count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
            reference_segments.append(("audio", count, (0, 0, 0)))
            rows = slice(cursor, cursor + count)
            cursor = rows.stop
            audio_indices.append(np.arange(rows.start, rows.stop))
            _fill_audio_positions(position_ids, rows, reference.num_audio_latents, rotary_time, target_width_grid)
            rotary_time += float(reference.num_audio_latents)
        elif reference.media_type == "video":
            audio_count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS if reference.has_audio else 0
            if audio_count:
                reference_segments.append(("audio", audio_count, (0, 0, 0)))
            reference_segments.append(("video", visual_row_count, visual_latent_shape))
            audio_rows = slice(cursor, cursor + audio_count)
            video_rows = slice(audio_rows.stop, audio_rows.stop + visual_row_count)
            cursor = video_rows.stop
            if audio_count:
                audio_indices.append(np.arange(audio_rows.start, audio_rows.stop))
            video_indices.append(np.arange(video_rows.start, video_rows.stop))

            frame_grid, width_grid = _frame_position_grid(
                reference.latent_height,
                reference.latent_width,
                patch_h,
                patch_w,
            )
            if audio_count:
                _fill_audio_positions(
                    position_ids,
                    audio_rows,
                    reference.num_audio_latents,
                    rotary_time,
                    width_grid,
                )
            frame_time = temporal_position_grid(reference.num_latent_frames, rotary_time)
            rows_per_frame = frame_grid.shape[0]
            position_ids[video_rows, 0] = np.repeat(frame_time, rows_per_frame)
            position_ids[video_rows, 1:] = np.tile(frame_grid, (reference.num_latent_frames // patch_t, 1))
            rotary_time += max(
                float(reference.num_audio_latents if reference.has_audio else 0),
                _reference_temporal_span(reference.num_latent_frames),
            )
        else:
            raise ValueError(f"Unsupported prepared reference type: {reference.media_type!r}.")

    audio_start = cursor
    video_start = audio_start + num_target_audio_rows
    _fill_audio_positions(position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time,
                          target_width_grid)
    frame_time = temporal_position_grid(num_latent_frames, rotary_time)
    position_ids[video_start:, 0] = np.repeat(frame_time, target_frame_grid.shape[0])
    position_ids[video_start:, 1:] = np.tile(target_frame_grid, (num_latent_frames // patch_t, 1))

    video_indices.append(np.arange(video_start, sequence_length))
    audio_indices.append(np.arange(audio_start, video_start))
    packed_video_indices = np.concatenate(video_indices)
    packed_audio_indices = np.concatenate(audio_indices)
    text_indices = np.arange(num_text_tokens)
    token_tags = np.empty(sequence_length, dtype=np.int64)
    token_tags[text_indices] = text_token_tags.astype(np.int64)
    token_tags[packed_audio_indices] = MINIMAX_H3_AUDIO_TAG
    token_tags[packed_video_indices] = MINIMAX_H3_VIDEO_TAG

    return MiniMaxH3PackedLayout(
        sequence_length=sequence_length,
        position_ids=position_ids,
        token_tags=token_tags,
        video_indices=packed_video_indices,
        audio_indices=packed_audio_indices,
        text_indices=text_indices,
        num_condition_video_rows=num_reference_video_rows,
        num_condition_audio_rows=num_reference_audio_rows,
        num_video_latent_frames=num_latent_frames,
        latent_height=latent_height,
        latent_width=latent_width,
        num_audio_latents=num_audio_latents,
        reference_segments=tuple(reference_segments),
    )

fastvideo.mlx_runtime.minimax_h3.build_row_timesteps

build_row_timesteps(layout: MiniMaxH3PackedLayout, video_timestep: float, audio_timestep: float, condition_video_timestep: float = 1.0, condition_audio_timestep: float = 1.0) -> tuple[np.ndarray, np.ndarray]

Per-row timesteps -> (unique timesteps sorted ascending, per-row indices).

Matches torch.unique(..., sorted=True, return_inverse=True).

Source
def build_row_timesteps(
    layout: MiniMaxH3PackedLayout,
    video_timestep: float,
    audio_timestep: float,
    condition_video_timestep: float = 1.0,
    condition_audio_timestep: float = 1.0,
) -> tuple[np.ndarray, np.ndarray]:
    """Per-row timesteps -> (unique timesteps sorted ascending, per-row indices).

    Matches ``torch.unique(..., sorted=True, return_inverse=True)``.
    """
    row_timesteps = np.full(layout.sequence_length, video_timestep, dtype=np.float64)
    if layout.num_condition_video_rows:
        row_timesteps[layout.video_indices[:layout.num_condition_video_rows]] = condition_video_timestep
    row_timesteps[layout.audio_indices[layout.num_condition_audio_rows:]] = audio_timestep
    if layout.num_condition_audio_rows:
        row_timesteps[layout.audio_indices[:layout.num_condition_audio_rows]] = condition_audio_timestep
    unique, inverse = np.unique(row_timesteps, return_inverse=True)
    return unique.astype(np.float32), inverse.astype(np.int64)

fastvideo.mlx_runtime.minimax_h3.minimax_h3_sigmas

minimax_h3_sigmas(shift: float, num_denoise_steps: int) -> np.ndarray

Sigma grid of length num_denoise_steps + 1 (descending, ending at 0).

The reference set_timesteps(num_inference_steps=N) builds N sigmas and N-1 timesteps; this wrapper takes the denoise step count directly, which is what a sampler actually schedules.

Source
def minimax_h3_sigmas(shift: float, num_denoise_steps: int) -> np.ndarray:
    """Sigma grid of length ``num_denoise_steps + 1`` (descending, ending at 0).

    The reference ``set_timesteps(num_inference_steps=N)`` builds N sigmas and
    N-1 timesteps; this wrapper takes the *denoise step count* directly, which
    is what a sampler actually schedules.
    """
    if num_denoise_steps < 1:
        raise ValueError(f"num_denoise_steps must be >= 1, got {num_denoise_steps}.")
    base = np.linspace(1.0, 0.0, num_denoise_steps + 1, dtype=np.float64)
    sigmas = shift * base / (1.0 + (shift - 1.0) * base)
    return _unique_consecutive(sigmas).astype(np.float32)

fastvideo.mlx_runtime.minimax_h3.validate_fasth3_dmd_rungs

validate_fasth3_dmd_rungs(dmd_steps: object) -> list[int]

Strictly decreasing integer rungs on the 1000-step training clock.

Source
def validate_fasth3_dmd_rungs(dmd_steps: object) -> list[int]:
    """Strictly decreasing integer rungs on the 1000-step training clock."""
    if (not isinstance(dmd_steps, list) or not dmd_steps
            or any(type(step) is not int or not 0 < step <= 1000 for step in dmd_steps)
            or any(left <= right for left, right in zip(dmd_steps, dmd_steps[1:], strict=False))):
        raise ValueError("FastH3 DMD rungs must be strictly decreasing integers in (0, 1000].")
    return dmd_steps

fastvideo.mlx_runtime.minimax_h3.find_fasth3_inference_contract

find_fasth3_inference_contract(model_root: str | Path | None) -> Path | None

Return the snapshot contract, also checking the parent of a transformer dir.

Source
def find_fasth3_inference_contract(model_root: str | Path | None) -> Path | None:
    """Return the snapshot contract, also checking the parent of a transformer dir."""
    if model_root is None:
        return None
    root = Path(model_root)
    direct = root / FASTH3_INFERENCE_CONTRACT
    if direct.is_file():
        return direct
    parent = root.parent / FASTH3_INFERENCE_CONTRACT
    if parent.is_file() and parent != direct:
        return parent
    return None

fastvideo.mlx_runtime.minimax_h3.load_fasth3_inference_contract

load_fasth3_inference_contract(model_root: str | Path | None) -> dict[str, Any] | None

Load a shift-declaring FastH3 contract, DMD or Ref2VA PDD.

Four-step exports ship a sidecar without scheduler shifts and keep the uniform ladder. A contract that declares both shifts selects an explicit ladder: a FastH3 Ref2VA PDD export (pdd_steps / pdd_step_indices) is tagged kind="pdd" and a FastH3 8-Step V2 DMD export (which carries dmd_denoising_steps) is tagged kind="dmd". A shifted sidecar that cannot be parsed fails instead of silently baking the four-step grid.

Source
def load_fasth3_inference_contract(model_root: str | Path | None) -> dict[str, Any] | None:
    """Load a shift-declaring FastH3 contract, DMD or Ref2VA PDD.

    Four-step exports ship a sidecar without scheduler shifts and keep the
    uniform ladder. A contract that declares both shifts selects an explicit
    ladder: a FastH3 Ref2VA PDD export (``pdd_steps`` / ``pdd_step_indices``)
    is tagged ``kind="pdd"`` and a FastH3 8-Step V2 DMD export (which carries
    ``dmd_denoising_steps``) is tagged ``kind="dmd"``. A shifted sidecar that
    cannot be parsed fails instead of silently baking the four-step grid.
    """
    path = find_fasth3_inference_contract(model_root)
    if path is None:
        return None
    try:
        contract = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise ValueError(f"Malformed FastH3 inference contract {path}: {exc}") from exc
    if not isinstance(contract, dict):
        raise ValueError(f"Malformed FastH3 inference contract {path}.")
    if _is_pdd_contract(contract):
        return _load_pdd_contract(contract, path)
    declares_shifts = "video_scheduler_shift" in contract or "audio_scheduler_shift" in contract
    if not declares_shifts:
        return None
    if contract.get("schema_version") != FASTH3_INFERENCE_CONTRACT_SCHEMA:
        raise ValueError(f"Unsupported FastH3 inference contract schema in {path}.")
    if "video_scheduler_shift" not in contract or "audio_scheduler_shift" not in contract:
        raise ValueError(f"{path} must declare both video_scheduler_shift and audio_scheduler_shift.")
    steps = validate_fasth3_dmd_rungs(contract.get("dmd_denoising_steps"))
    forwards = contract.get("transformer_forwards")
    grid_points = contract.get("num_inference_steps")
    if type(forwards) is not int or forwards != len(steps):
        raise ValueError("FastH3 transformer_forwards must equal the DMD rung count.")
    if type(grid_points) is not int or grid_points != len(steps) + 1:
        raise ValueError("FastH3 num_inference_steps must equal the DMD rung count plus one.")
    contract["video_scheduler_shift"] = _positive_shift(contract["video_scheduler_shift"], name="video_scheduler_shift")
    contract["audio_scheduler_shift"] = _positive_shift(contract["audio_scheduler_shift"], name="audio_scheduler_shift")
    contract["dmd_denoising_steps"] = steps
    contract["kind"] = "dmd"
    return contract

fastvideo.mlx_runtime.minimax_h3.adaln_timestep_union

adaln_timestep_union(video: MiniMaxH3SchedulerState, audio: MiniMaxH3SchedulerState) -> np.ndarray

Sorted unique AdaLN timesteps for one video/audio pair, plus clean time 1.

Source
def adaln_timestep_union(video: MiniMaxH3SchedulerState, audio: MiniMaxH3SchedulerState) -> np.ndarray:
    """Sorted unique AdaLN timesteps for one video/audio pair, plus clean time 1."""
    return np.unique(
        np.concatenate([
            np.asarray(video.timesteps, dtype=np.float32),
            np.asarray(audio.timesteps, dtype=np.float32),
            np.array([1.0], dtype=np.float32),
        ])).astype(np.float32)

fastvideo.mlx_runtime.minimax_h3.ResolvedH3Schedule

ResolvedH3Schedule(video: MiniMaxH3SchedulerState, audio: MiniMaxH3SchedulerState, adaln_timesteps: np.ndarray, num_steps: int, source: str, pdd_plan: Any = None)

The denoise ladder a request is allowed to run.

Source
@dataclass(frozen=True)
class ResolvedH3Schedule:
    """The denoise ladder a request is allowed to run."""

    video: MiniMaxH3SchedulerState
    audio: MiniMaxH3SchedulerState
    adaln_timesteps: np.ndarray
    num_steps: int
    source: str
    # The fused-block plan of a PDD contract (its per-modality coefficients and
    # blocks); None for every other source.
    pdd_plan: Any = None

fastvideo.mlx_runtime.minimax_h3.pdd_plan_from_contract

pdd_plan_from_contract(contract: dict[str, Any]) -> Any

Build the fused-block plan a FastH3 PDD contract trained.

Source
def pdd_plan_from_contract(contract: dict[str, Any]) -> Any:
    """Build the fused-block plan a FastH3 PDD contract trained."""
    from fastvideo.mlx_runtime.minimax_h3_pdd import PDDModalitySchedule, build_pdd_sampling_plan

    schedules = {
        "video": PDDModalitySchedule(shift=contract["video_scheduler_shift"]),
        "audio": PDDModalitySchedule(shift=contract["audio_scheduler_shift"]),
    }
    return build_pdd_sampling_plan(contract["pdd_step_indices"], schedules)

fastvideo.mlx_runtime.minimax_h3.pdd_adaln_timesteps

pdd_adaln_timesteps(contract: dict[str, Any], video: MiniMaxH3SchedulerState, audio: MiniMaxH3SchedulerState) -> np.ndarray

A PDD contract's AdaLN ladder; Ref2VA adds the 0.999 its reference rows always run at.

Source
def pdd_adaln_timesteps(contract: dict[str, Any], video: MiniMaxH3SchedulerState,
                        audio: MiniMaxH3SchedulerState) -> np.ndarray:
    """A PDD contract's AdaLN ladder; Ref2VA adds the 0.999 its reference rows always run at."""
    union = adaln_timestep_union(video, audio)
    if contract.get("model_type") != "ref2va":
        return union
    return np.unique(np.append(union, np.float32(MINIMAX_H3_KEYFRAME_NOISE_AUG))).astype(np.float32)

fastvideo.mlx_runtime.minimax_h3.resolve_h3_denoise_schedule

resolve_h3_denoise_schedule(model_root: str | Path | None, num_steps: int, *, cached_timesteps: np.ndarray | None = None) -> ResolvedH3Schedule

Keep the uniform ladder unless the snapshot declares its own shifts.

num_steps is the transformer-forward count. A shift-declaring DMD contract also accepts that count plus one, the CUDA sigma-grid convention, and then runs exactly its trained rung count. A FastH3 Ref2VA PDD contract instead runs its trained fused-block count exactly and the schedulers carry the plan's node sigmas.

Source
def resolve_h3_denoise_schedule(
    model_root: str | Path | None,
    num_steps: int,
    *,
    cached_timesteps: np.ndarray | None = None,
) -> ResolvedH3Schedule:
    """Keep the uniform ladder unless the snapshot declares its own shifts.

    ``num_steps`` is the transformer-forward count. A shift-declaring DMD
    contract also accepts that count plus one, the CUDA sigma-grid convention,
    and then runs exactly its trained rung count. A FastH3 Ref2VA PDD contract
    instead runs its trained fused-block count exactly and the schedulers carry
    the plan's node sigmas.
    """
    if type(num_steps) is not int or num_steps < 1:
        raise ValueError(f"num_steps must be a positive integer, got {num_steps!r}.")
    contract = load_fasth3_inference_contract(model_root)
    if contract is None:
        video = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps)
        audio = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps)
        union = adaln_timestep_union(video, audio)
        if cached_timesteps is not None and not np.array_equal(np.asarray(cached_timesteps, dtype=np.float32), union):
            raise ValueError(f"MLX H3 checkpoint has a fixed AdaLN ladder that does not support --steps "
                             f"{num_steps}. Use the step count used during conversion (normally 4), or re-export "
                             "the checkpoint.")
        return ResolvedH3Schedule(video, audio, union, num_steps, "uniform")

    if contract["kind"] == "pdd":
        return _resolve_pdd_schedule(contract, num_steps, cached_timesteps)

    rungs = contract["dmd_denoising_steps"]
    forwards = len(rungs)
    if num_steps not in (forwards, forwards + 1):
        raise ValueError(f"Requested {num_steps} steps, but the FastH3 contract trains {forwards} transformer "
                         f"forwards ({forwards + 1} sigma-grid points).")
    video = MiniMaxH3SchedulerState.from_dmd_steps(contract["video_scheduler_shift"], rungs)
    audio = MiniMaxH3SchedulerState.from_dmd_steps(contract["audio_scheduler_shift"], rungs)
    union = adaln_timestep_union(video, audio)
    if cached_timesteps is not None and not np.array_equal(np.unique(np.asarray(cached_timesteps, dtype=np.float32)),
                                                           union):
        _require_cached_schedule(cached_timesteps, union, num_steps)
    return ResolvedH3Schedule(video, audio, union, forwards, "contract-dmd")

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3SchedulerState

MiniMaxH3SchedulerState(shift: float, sigmas: np.ndarray, timesteps: np.ndarray)

One rectified-flow scheduler (use two: video shift 12, audio shift 3).

Source
@dataclass
class MiniMaxH3SchedulerState:
    """One rectified-flow scheduler (use two: video shift 12, audio shift 3)."""

    shift: float
    sigmas: np.ndarray
    timesteps: np.ndarray  # 1 - sigmas[:-1], fp32

    @classmethod
    def create(cls, shift: float, num_denoise_steps: int) -> MiniMaxH3SchedulerState:
        if shift <= 0:
            raise ValueError(f"`shift` must be positive, got {shift}.")
        sigmas = minimax_h3_sigmas(shift, num_denoise_steps)
        return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32))

    @classmethod
    def from_dmd_steps(cls, shift: float, dmd_steps: list[int]) -> MiniMaxH3SchedulerState:
        """Shift explicit DMD rungs the way the CUDA denoiser does."""
        shift = _positive_shift(shift, name="shift")
        rungs = validate_fasth3_dmd_rungs(dmd_steps)
        base = np.array([step / 1000.0 for step in rungs] + [0.0], dtype=np.float64)
        # Float64 shift, then one cast. Shipped MLX AdaLN tables are keyed by
        # that rounding; casting the sigma before 1 - sigma misses the cache.
        sigmas64 = shift * base / (1.0 + (shift - 1.0) * base)
        if sigmas64.size < 2 or not bool(np.all(sigmas64[1:] < sigmas64[:-1])) or float(sigmas64[-1]) != 0.0:
            raise ValueError("DMD sigmas must be strictly decreasing and end at 0.")
        return cls(shift=shift, sigmas=sigmas64.astype(np.float32), timesteps=(1.0 - sigmas64[:-1]).astype(np.float32))

    @classmethod
    def from_pdd_node_sigmas(cls, shift: float, node_sigmas: np.ndarray) -> MiniMaxH3SchedulerState:
        """Carry a PDD plan's shifted node sigmas as the sampling ladder.

        ``node_sigmas`` are already shifted (one noise level per fused-block
        node, descending to 0), so an ordinary Euler step between consecutive
        nodes applies exactly the block's total integration weight.
        """
        shift = _positive_shift(shift, name="shift")
        sigmas = np.asarray(node_sigmas, dtype=np.float32)
        if sigmas.size < 2 or float(sigmas[-1]) != 0.0 or not bool(np.all(sigmas[1:] < sigmas[:-1])):
            raise ValueError("PDD node sigmas must be strictly decreasing and end at 0.")
        return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32))

    @property
    def num_steps(self) -> int:
        return int(self.timesteps.shape[0])

    def step(self, model_output, step_index: int, sample):
        """Data-ward Euler: x0 = x_t + sigma*v; blend toward x0 by sigma ratio.

        fp32 blend for fp16/bf16 samples, matching the reference.
        """
        import mlx.core as mx

        sigma_from_timestep = float(1.0 - self.timesteps[step_index])
        denoised = sample + sigma_from_timestep * model_output
        sigma = float(self.sigmas[step_index])
        sigma_next = float(self.sigmas[step_index + 1])
        ratio = sigma_next / sigma
        prev = ratio * sample.astype(mx.float32) + (1.0 - ratio) * denoised.astype(mx.float32)
        return prev.astype(sample.dtype)

    def scale_noise(self, sample, timestep: float, noise):
        """Conditioning noise-aug: t*sample + (1-t)*noise (t=0.999 for keyframes)."""
        return timestep * sample + (1.0 - timestep) * noise

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3SchedulerState.create

create(shift: float, num_denoise_steps: int) -> MiniMaxH3SchedulerState
Source
@classmethod
def create(cls, shift: float, num_denoise_steps: int) -> MiniMaxH3SchedulerState:
    if shift <= 0:
        raise ValueError(f"`shift` must be positive, got {shift}.")
    sigmas = minimax_h3_sigmas(shift, num_denoise_steps)
    return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32))

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3SchedulerState.from_dmd_steps

from_dmd_steps(shift: float, dmd_steps: list[int]) -> MiniMaxH3SchedulerState

Shift explicit DMD rungs the way the CUDA denoiser does.

Source
@classmethod
def from_dmd_steps(cls, shift: float, dmd_steps: list[int]) -> MiniMaxH3SchedulerState:
    """Shift explicit DMD rungs the way the CUDA denoiser does."""
    shift = _positive_shift(shift, name="shift")
    rungs = validate_fasth3_dmd_rungs(dmd_steps)
    base = np.array([step / 1000.0 for step in rungs] + [0.0], dtype=np.float64)
    # Float64 shift, then one cast. Shipped MLX AdaLN tables are keyed by
    # that rounding; casting the sigma before 1 - sigma misses the cache.
    sigmas64 = shift * base / (1.0 + (shift - 1.0) * base)
    if sigmas64.size < 2 or not bool(np.all(sigmas64[1:] < sigmas64[:-1])) or float(sigmas64[-1]) != 0.0:
        raise ValueError("DMD sigmas must be strictly decreasing and end at 0.")
    return cls(shift=shift, sigmas=sigmas64.astype(np.float32), timesteps=(1.0 - sigmas64[:-1]).astype(np.float32))

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3SchedulerState.from_pdd_node_sigmas

from_pdd_node_sigmas(shift: float, node_sigmas: np.ndarray) -> MiniMaxH3SchedulerState

Carry a PDD plan's shifted node sigmas as the sampling ladder.

node_sigmas are already shifted (one noise level per fused-block node, descending to 0), so an ordinary Euler step between consecutive nodes applies exactly the block's total integration weight.

Source
@classmethod
def from_pdd_node_sigmas(cls, shift: float, node_sigmas: np.ndarray) -> MiniMaxH3SchedulerState:
    """Carry a PDD plan's shifted node sigmas as the sampling ladder.

    ``node_sigmas`` are already shifted (one noise level per fused-block
    node, descending to 0), so an ordinary Euler step between consecutive
    nodes applies exactly the block's total integration weight.
    """
    shift = _positive_shift(shift, name="shift")
    sigmas = np.asarray(node_sigmas, dtype=np.float32)
    if sigmas.size < 2 or float(sigmas[-1]) != 0.0 or not bool(np.all(sigmas[1:] < sigmas[:-1])):
        raise ValueError("PDD node sigmas must be strictly decreasing and end at 0.")
    return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32))

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3SchedulerState.step

step(model_output, step_index: int, sample)

Data-ward Euler: x0 = x_t + sigma*v; blend toward x0 by sigma ratio.

fp32 blend for fp16/bf16 samples, matching the reference.

Source
def step(self, model_output, step_index: int, sample):
    """Data-ward Euler: x0 = x_t + sigma*v; blend toward x0 by sigma ratio.

    fp32 blend for fp16/bf16 samples, matching the reference.
    """
    import mlx.core as mx

    sigma_from_timestep = float(1.0 - self.timesteps[step_index])
    denoised = sample + sigma_from_timestep * model_output
    sigma = float(self.sigmas[step_index])
    sigma_next = float(self.sigmas[step_index + 1])
    ratio = sigma_next / sigma
    prev = ratio * sample.astype(mx.float32) + (1.0 - ratio) * denoised.astype(mx.float32)
    return prev.astype(sample.dtype)

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3SchedulerState.scale_noise

scale_noise(sample, timestep: float, noise)

Conditioning noise-aug: tsample + (1-t)noise (t=0.999 for keyframes).

Source
def scale_noise(self, sample, timestep: float, noise):
    """Conditioning noise-aug: t*sample + (1-t)*noise (t=0.999 for keyframes)."""
    return timestep * sample + (1.0 - timestep) * noise

fastvideo.mlx_runtime.minimax_h3.rope_cos_sin

rope_cos_sin(position_ids, rope_freq_dim: int, rope_theta: float)

(S, 3) positions -> (cos, sin) each (S, 6*rope_freq_dim), fp32.

Shared 16-freq inv_freq across the three axes; the (t, h, w) blocks are concatenated and doubled, so the rotary width is 6 * rope_freq_dim (96 of the 128 head dims at full size).

Source
def rope_cos_sin(position_ids, rope_freq_dim: int, rope_theta: float):
    """(S, 3) positions -> (cos, sin) each (S, 6*rope_freq_dim), fp32.

    Shared 16-freq inv_freq across the three axes; the (t, h, w) blocks are
    concatenated and doubled, so the rotary width is 6 * rope_freq_dim (96 of
    the 128 head dims at full size).
    """
    import mlx.core as mx

    positions = position_ids.astype(mx.float32)
    inv_freq = 1.0 / (rope_theta**(mx.arange(0, 2 * rope_freq_dim, 2, dtype=mx.float32) / (2 * rope_freq_dim)))
    freqs = positions[:, :, None] * inv_freq[None, None, :]  # (S, 3, F)
    freqs_t, freqs_h, freqs_w = freqs[:, 0], freqs[:, 1], freqs[:, 2]
    freqs = mx.concatenate([freqs_t, freqs_h, freqs_w], axis=-1)
    freqs = mx.concatenate([freqs, freqs], axis=-1)
    return mx.cos(freqs), mx.sin(freqs)

fastvideo.mlx_runtime.minimax_h3.apply_h3_rotary

apply_h3_rotary(x, cos, sin)

Rotate the RoPE prefix of each head, preserving the rest.

x: (S, H, D); cos/sin: (S, R) with R <= D. Half-split (GPT-NeoX style) rotation inside the rotary prefix, computed in fp32. The fp32 math runs in one compiled kernel so the (S, H, R) fp32 intermediates never hit memory.

Source
def apply_h3_rotary(x, cos, sin):
    """Rotate the RoPE prefix of each head, preserving the rest.

    x: (S, H, D); cos/sin: (S, R) with R <= D. Half-split (GPT-NeoX style)
    rotation inside the rotary prefix, computed in fp32. The fp32 math runs in
    one compiled kernel so the (S, H, R) fp32 intermediates never hit memory.
    """
    import mlx.core as mx

    global _compiled_rotate_half
    if _compiled_rotate_half is None:
        _compiled_rotate_half = mx.compile(_rotate_half_fp32)
    rotary_dim = cos.shape[-1]
    out = _compiled_rotate_half(x[..., :rotary_dim], cos, sin)
    return mx.concatenate([out, x[..., rotary_dim:]], axis=-1)

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3StepCache

MiniMaxH3StepCache(timesteps: np.ndarray, block_tables: list[tuple[Any, ...]], norm_out_shift: Any, norm_out_scale: Any)

Precomputed AdaLN tables + norm_out modulation for a fixed schedule.

Holds one row set per distinct timestep in timesteps (the union of every denoise step's video/audio/condition timesteps). Blocks then index tables directly and the per-block adaln_proj weights can be freed — ~40% of H3's parameters never need to be resident during denoise.

Source
@dataclass
class MiniMaxH3StepCache:
    """Precomputed AdaLN tables + norm_out modulation for a fixed schedule.

    Holds one row set per distinct timestep in ``timesteps`` (the union of
    every denoise step's video/audio/condition timesteps). Blocks then index
    tables directly and the per-block ``adaln_proj`` weights can be freed —
    ~40% of H3's parameters never need to be resident during denoise.
    """

    timesteps: np.ndarray  # (n,) fp32, sorted
    block_tables: list[tuple[Any, ...]]  # per block: 6 tables of (n*3, hidden)
    norm_out_shift: Any  # (n, hidden)
    norm_out_scale: Any  # (n, hidden)

    def positions(self, step_timesteps: np.ndarray) -> np.ndarray:
        """Map a step's unique timesteps to rows in the cached union."""
        positions = np.searchsorted(self.timesteps, step_timesteps)
        if (positions.size and (np.any(positions >= len(self.timesteps))
                                or not np.allclose(self.timesteps[positions], step_timesteps, atol=1e-6))):
            raise ValueError(f"Step timesteps {step_timesteps} are not in the cached schedule union.")
        return positions.astype(np.int64)

fastvideo.mlx_runtime.minimax_h3.MiniMaxH3StepCache.positions

positions(step_timesteps: np.ndarray) -> np.ndarray

Map a step's unique timesteps to rows in the cached union.

Source
def positions(self, step_timesteps: np.ndarray) -> np.ndarray:
    """Map a step's unique timesteps to rows in the cached union."""
    positions = np.searchsorted(self.timesteps, step_timesteps)
    if (positions.size and (np.any(positions >= len(self.timesteps))
                            or not np.allclose(self.timesteps[positions], step_timesteps, atol=1e-6))):
        raise ValueError(f"Step timesteps {step_timesteps} are not in the cached schedule union.")
    return positions.astype(np.int64)

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT

MLXMiniMaxH3DiT(weights: dict[str, Any], blocks: list[dict[str, Any]], refiner: list[dict[str, Any]], config: dict[str, Any])

MiniMax H3 joint audio-video DiT in MLX (batch-1 packed forward).

Weight dicts keep the released checkpoint's key names. blocks holds the main transformer blocks, refiner the text refiner blocks; each is a dict of arrays / :class:QuantizedMatrix.

Source
class MLXMiniMaxH3DiT:
    """MiniMax H3 joint audio-video DiT in MLX (batch-1 packed forward).

    Weight dicts keep the released checkpoint's key names. ``blocks`` holds
    the main transformer blocks, ``refiner`` the text refiner blocks; each is
    a dict of arrays / :class:`QuantizedMatrix`.
    """

    def __init__(
        self,
        weights: dict[str, Any],
        blocks: list[dict[str, Any]],
        refiner: list[dict[str, Any]],
        config: dict[str, Any],
    ) -> None:
        self.weights = weights
        self.blocks = blocks
        self.refiner = refiner
        self.config = config
        self.hidden_size = int(config["hidden_size"])
        self.num_heads = int(config["num_attention_heads"])
        self.head_dim = int(config["attention_head_dim"])
        self.ffn_dim = int(config["ffn_dim"])
        self.in_channels = int(config["in_channels"])
        self.audio_in_channels = int(config["audio_in_channels"])
        self.patch_size = tuple(config["patch_size"])
        self.text_dim = int(config["text_dim"])
        self.freq_dim = int(config["freq_dim"])
        self.time_embed_dim = int(config["time_embed_dim"])
        self.rope_freq_dim = int(config["rope_freq_dim"])
        self.rope_theta = float(config["rope_theta"])
        self.norm_eps = float(config["norm_eps"])
        self.qk_norm_eps = float(config["qk_norm_eps"])
        self.final_norm_eps = float(config["final_norm_eps"])
        self.patch_dim = self.in_channels * math.prod(self.patch_size)
        self.adaln_rank = config.get("adaln_rank")
        self._adaln_cache: MiniMaxH3StepCache | None = None
        self.vsa_config = MiniMaxH3VSAConfig()
        self._vsa_geometry: MiniMaxH3VSAGeometry | None = None
        self._vsa_capable = any("attn.to_gate_compress.weight" in block for block in blocks)
        self._gate_active: list[bool | None] = [None] * len(blocks)
        self.last_vsa_stats: MiniMaxH3VSAStats | None = None
        self.pdd_grid_size = self._resolve_pdd_grid_size()

    def _resolve_pdd_grid_size(self) -> int | None:
        """The PDD head count of this checkpoint, or None for a single-headed DiT.

        A PDD student widens ``proj_out`` / ``audio_proj_out`` head-major to
        ``(grid, output_size)``. The grid comes from the packed manifest's
        ``config["pdd_steps"]`` when present; otherwise it is inferred from the
        projection's leading dimension and cross-checked against the audio head.
        A widened head that would silently mis-project raises here.
        """
        from fastvideo.mlx_runtime.minimax_h3_pdd import resolved_grid_size

        declared = self.config.get("pdd_steps")
        if declared is not None and (isinstance(declared, bool) or not isinstance(declared, int) or declared < 2):
            raise ValueError(f"MiniMax H3 pdd_steps must be an int >= 2, got {declared!r}.")
        video_weight = self.weights.get("proj_out.weight")
        audio_weight = self.weights.get("audio_proj_out.weight")
        video_grid = None if video_weight is None else resolved_grid_size(
            video_weight, self.weights.get("proj_out.bias"), self.patch_dim)
        audio_grid = None if audio_weight is None else resolved_grid_size(
            audio_weight, self.weights.get("audio_proj_out.bias"), self.audio_in_channels)
        if video_grid is not None and audio_grid is not None and video_grid != audio_grid:
            raise ValueError(f"MiniMax H3 PDD heads disagree between proj_out (grid {video_grid}) and "
                             f"audio_proj_out (grid {audio_grid}).")
        grid = video_grid if video_grid is not None else audio_grid
        if declared is not None and grid is not None and grid != declared:
            raise ValueError(f"MiniMax H3 checkpoint config pdd_steps={declared} disagrees with its widened heads "
                             f"(grid {grid}).")
        return grid if grid is not None else declared

    def _fused_output_params(
        self,
        key: str,
        modality: str,
        pdd_block: tuple[int, int] | None,
        pdd_integration_weights: dict[str, np.ndarray] | None,
    ):
        """The projection (weight, bias) to use, fusing a PDD block when one is supplied.

        Without a block a single-headed head passes through and a widened head
        (which would otherwise project with the raw ``grid * output_size`` rows)
        is refused. With a block an unwidened head is refused.
        """
        weight = self.weights[key]
        bias = self.weights.get(f"{key.rsplit('.', 1)[0]}.bias")
        if pdd_block is None:
            if self.pdd_grid_size is not None:
                raise ValueError(f"MiniMax H3 checkpoint has a PDD-widened `{key}` but no PDD block was supplied; "
                                 "pass pdd_block= and pdd_integration_weights= to the forward.")
            return weight, bias
        if self.pdd_grid_size is None:
            raise ValueError(f"A PDD block was supplied but `{key}` is not widened (single-headed checkpoint).")
        if pdd_integration_weights is None or modality not in pdd_integration_weights:
            raise ValueError(f"A PDD block needs this modality's integration weights ({modality!r}).")
        from fastvideo.mlx_runtime.minimax_h3_pdd import fuse_block_params

        return fuse_block_params(
            weight,
            bias,
            grid_size=self.pdd_grid_size,
            start=pdd_block[0],
            end=pdd_block[1],
            weights=pdd_integration_weights[modality],
        )

    # -- conditioning ------------------------------------------------------

    def compute_temb(self, timesteps):
        """(n,) timesteps -> (n, time_embed_dim) through time_proj + MLP."""
        t_freq = timestep_embedding(timesteps,
                                    self.freq_dim).astype(weight_dtype(self.weights["time_embedder.linear_1.weight"]))
        temb = linear(
            t_freq,
            self.weights["time_embedder.linear_1.weight"],
            self.weights["time_embedder.linear_1.bias"],
        )
        temb = linear(
            silu(temb),
            self.weights["time_embedder.linear_2.weight"],
            self.weights["time_embedder.linear_2.bias"],
        )
        if self.adaln_rank is not None:
            temb = linear(
                silu(temb).astype(weight_dtype(self.weights["adaln_basis.weight"])), self.weights["adaln_basis.weight"])
        return temb

    def refine_text(self, text_rows):
        hidden = linear(
            text_rows.astype(weight_dtype(self.weights["context_embedder.weight"])),
            self.weights["context_embedder.weight"],
            self.weights["context_embedder.bias"],
        )
        for block in self.refiner:
            hidden = _refiner_block(
                block,
                hidden,
                num_heads=self.num_heads,
                head_dim=self.head_dim,
                norm_eps=self.norm_eps,
                qk_norm_eps=self.qk_norm_eps,
            )
        return _h3_rms_norm(hidden, self.weights["token_refiner.final_norm.weight"], self.final_norm_eps)

    # -- AdaLN cache --------------------------------------------------------

    def precompute_adaln(self, timesteps: np.ndarray, *, drop_weights: bool = True) -> MiniMaxH3StepCache:
        """Evaluate every modulation table for a fixed schedule, once.

        ``timesteps`` is the union (sorted, unique) of all per-step
        video/audio/condition timesteps the sampler will use. With
        ``drop_weights=True`` the per-block ``adaln_proj.linear`` weights are
        released afterwards — the memory win this runtime exists for.
        """
        import mlx.core as mx

        timesteps = np.unique(np.asarray(timesteps, dtype=np.float32))
        temb = self.compute_temb(mx.array(timesteps))
        block_tables = [_adaln_tables(block, temb, apply_silu=self.adaln_rank is None) for block in self.blocks]
        shift_scale = linear(
            (silu(temb) if self.adaln_rank is None else temb).astype(
                weight_dtype(self.weights["norm_out.linear.weight"])),
            self.weights["norm_out.linear.weight"],
            self.weights["norm_out.linear.bias"],
        )
        norm_out_shift, norm_out_scale = mx.split(shift_scale, 2, axis=-1)
        mx.eval(block_tables, norm_out_scale, norm_out_shift)
        self._adaln_cache = MiniMaxH3StepCache(
            timesteps=timesteps,
            block_tables=block_tables,
            norm_out_shift=norm_out_shift,
            norm_out_scale=norm_out_scale,
        )
        if drop_weights:
            for block in self.blocks:
                block["adaln_proj.linear.weight"] = None
                block["adaln_proj.linear.bias"] = None
        return self._adaln_cache

    @property
    def vsa_capable(self) -> bool:
        return self._vsa_capable

    def configure_vsa(self, config: MiniMaxH3VSAConfig | None) -> None:
        config = config or MiniMaxH3VSAConfig()
        if config.enabled and not self._vsa_capable:
            raise dense_only_vsa_error("this MLX H3 DiT")
        if any(layer >= len(self.blocks) for layer in config.dense_layers):
            raise ValueError(f"VSA dense_layers must be below the model's block count ({len(self.blocks)}).")
        self.vsa_config = config
        self._vsa_geometry = None
        self.reset_vsa_stats()

    def reset_vsa_stats(self) -> None:
        self.last_vsa_stats = MiniMaxH3VSAStats(
            configured_sparsity=self.vsa_config.sparsity,
            tile_size=self.vsa_config.tile_size,
            prefix_mode=self.vsa_config.prefix_mode,
        ) if self.vsa_config.enabled else None

    def prepare_vsa_geometry(self, layout: MiniMaxH3PackedLayout) -> MiniMaxH3VSAGeometry | None:
        if not self.vsa_config.enabled:
            self._vsa_geometry = None
            return None
        config = self.vsa_config
        ref_keep_rate = config.effective_ref_keep_rate
        if ref_keep_rate is not None and layout.reference_segments:
            segments = ref2va_vsa_segments_from_layout(layout, self.patch_size)
        else:
            segments = (*prefix_segments_from_layout(layout, self.patch_size),
                        dit_seq_shape_from_layout(layout, self.patch_size))
        segments = tuple(segment for segment in segments if isinstance(segment, tuple) or segment > 0)
        geometry = self._vsa_geometry
        if geometry is not None and (geometry.segments, geometry.tile_elems) == (segments, config.tile_size):
            return geometry
        try:
            geometry = build_h3_segment_geometry(segments, config.tile_size)
        except ValueError as error:
            raise ValueError(f"H3 VSA cannot run with this packed layout (segments={segments}, "
                             f"tile_size={config.tile_size}): {error}. "
                             "Disable VSA or pick a supported tile size (64, 128 or 256).") from error
        self._vsa_geometry = geometry
        return self._vsa_geometry

    def _block_gate_active(self, block_index: int) -> bool:
        cached = self._gate_active[block_index]
        if cached is not None:
            return cached
        weight = self.blocks[block_index].get("attn.to_gate_compress.weight")
        active = _gate_weight_is_active(weight)
        self._gate_active[block_index] = active
        if weight is not None and not active:
            logger.info_once("MiniMax H3 VSA compression gate is all-zero; skipping the pooled branch.")
        return active

    def _vsa_block_kwargs(self, block_index: int, step_index: int) -> dict[str, Any]:
        geometry = self._vsa_geometry if self.vsa_config.enabled else None
        sparsity = self.vsa_config.layer_sparsity(block_index, step_index) if geometry is not None else 0.0
        stats = MiniMaxH3VSAStats() if self.vsa_config.enabled else None
        return {
            "vsa_geometry": geometry,
            "vsa_sparsity": sparsity,
            "vsa_exempt": self.vsa_config.exempt,
            "vsa_ref_keep_rate": self.vsa_config.effective_ref_keep_rate,
            "vsa_impl": self.vsa_config.impl,
            "vsa_stats": stats,
            "use_gate_compress": geometry is not None and self._block_gate_active(block_index),
        }

    # -- forward ------------------------------------------------------------

    def forward(
        self,
        video_rows,
        audio_rows,
        text_rows,
        *,
        position_ids,
        token_tags,
        timestep_indices,
        timesteps,
        video_indices,
        audio_indices,
        text_indices,
        pdd_block: tuple[int, int] | None = None,
        pdd_integration_weights: dict[str, np.ndarray] | None = None,
    ):
        """Faithful port of the torch forward (batch-1, rows already patchified).

        Returns (video_output, audio_output) rows. The torch reference
        projects *all* packed rows through both heads and then selects; this
        selects first and projects only the relevant rows — row-wise
        identical math at a fraction of the cost.

        This entrypoint is dense-only. VSA requires ``forward_with_cache``
        with a packed layout and step index. With a PDD checkpoint pass the
        step's ``pdd_block`` and ``pdd_integration_weights`` so the widened
        heads are fused.
        """
        import mlx.core as mx

        if self.vsa_config.enabled:
            raise ValueError("VSA requires forward_with_cache() with layout and step_index; forward() is dense-only.")
        sequence_length = position_ids.shape[0]
        cos, sin = rope_cos_sin(position_ids, self.rope_freq_dim, self.rope_theta)

        video_embeds = linear(
            video_rows.astype(weight_dtype(self.weights["proj_in.weight"])),
            self.weights["proj_in.weight"],
            self.weights["proj_in.bias"],
        )
        audio_embeds = linear(
            audio_rows.astype(weight_dtype(self.weights["audio_proj_in.weight"])),
            self.weights["audio_proj_in.weight"],
            self.weights["audio_proj_in.bias"],
        )
        text_embeds = self.refine_text(text_rows)

        packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
        packed[_as_mx_indices(text_indices)] = text_embeds
        packed[_as_mx_indices(video_indices)] = video_embeds.astype(text_embeds.dtype)
        packed[_as_mx_indices(audio_indices)] = audio_embeds.astype(text_embeds.dtype)

        temb = self.compute_temb(timesteps)
        adaln_indices = (timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags).astype(mx.int32)

        for block_index, block in enumerate(self.blocks):
            tables = _adaln_tables(block, temb, apply_silu=self.adaln_rank is None)
            packed = _transformer_block(
                block,
                packed,
                tables,
                adaln_indices,
                cos,
                sin,
                num_heads=self.num_heads,
                head_dim=self.head_dim,
                norm_eps=self.norm_eps,
                qk_norm_eps=self.qk_norm_eps,
                **self._vsa_block_kwargs(block_index, 0),
            )
            mx.eval(packed)  # per-block sync: see forward_with_cache note

        shift_scale = linear(
            (silu(temb) if self.adaln_rank is None else temb).astype(
                weight_dtype(self.weights["norm_out.linear.weight"])),
            self.weights["norm_out.linear.weight"],
            self.weights["norm_out.linear.bias"],
        )
        shift_rows, scale_rows = mx.split(shift_scale, 2, axis=-1)
        normed = _h3_rms_norm(packed, self.weights["norm_out.norm.weight"], self.final_norm_eps)
        normed = normed * (1.0 + scale_rows[timestep_indices]) + shift_rows[timestep_indices]
        video_weight, video_bias = self._fused_output_params("proj_out.weight", "video", pdd_block,
                                                             pdd_integration_weights)
        audio_weight, audio_bias = self._fused_output_params("audio_proj_out.weight", "audio", pdd_block,
                                                             pdd_integration_weights)
        video_output = linear(
            normed[_as_mx_indices(video_indices)].astype(weight_dtype(video_weight)),
            video_weight,
            video_bias,
        )
        audio_output = linear(
            normed[_as_mx_indices(audio_indices)].astype(weight_dtype(audio_weight)),
            audio_weight,
            audio_bias,
        )
        return video_output, audio_output

    def forward_with_cache(
        self,
        video_rows,
        audio_rows,
        text_rows,
        *,
        layout: MiniMaxH3PackedLayout,
        step_timesteps: np.ndarray,
        row_timestep_inverse: np.ndarray,
        step_index: int = 0,
        pdd_block: tuple[int, int] | None = None,
        pdd_integration_weights: dict[str, np.ndarray] | None = None,
        record_prefix_kv: Any | None = None,
    ):
        """Denoise-step forward served entirely from the AdaLN cache.

        ``record_prefix_kv`` (a ``MiniMaxH3PrefixKVCache``) additionally keeps
        every block's K/V of the frozen text/reference rows for later
        ``forward_target_rows`` steps; the outputs are unchanged.

        ``step_timesteps`` are this step's unique timesteps (sorted);
        ``row_timestep_inverse`` maps each packed row to one of them (the
        ``build_row_timesteps`` inverse). The cache must cover every value in
        ``step_timesteps``. With a PDD checkpoint pass the step's
        ``pdd_block`` and ``pdd_integration_weights`` so the widened output
        heads are fused into the block's weighted mean.
        """
        import mlx.core as mx

        if self._adaln_cache is None:
            raise RuntimeError("precompute_adaln() must run before forward_with_cache().")
        if self.vsa_config.enabled:
            self.prepare_vsa_geometry(layout)
        cache = self._adaln_cache
        positions = mx.array(cache.positions(step_timesteps))

        sequence_length = layout.sequence_length
        position_ids = mx.array(layout.position_ids)
        cos, sin = rope_cos_sin(position_ids, self.rope_freq_dim, self.rope_theta)

        video_embeds = linear(
            video_rows.astype(weight_dtype(self.weights["proj_in.weight"])),
            self.weights["proj_in.weight"],
            self.weights["proj_in.bias"],
        )
        audio_embeds = linear(
            audio_rows.astype(weight_dtype(self.weights["audio_proj_in.weight"])),
            self.weights["audio_proj_in.weight"],
            self.weights["audio_proj_in.bias"],
        )
        text_embeds = self.refine_text(text_rows)

        packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
        packed[mx.array(layout.text_indices)] = text_embeds
        packed[mx.array(layout.video_indices)] = video_embeds.astype(text_embeds.dtype)
        packed[mx.array(layout.audio_indices)] = audio_embeds.astype(text_embeds.dtype)

        token_tags = mx.array(layout.token_tags)
        row_inverse = mx.array(row_timestep_inverse)
        adaln_indices = (positions[row_inverse] * MINIMAX_H3_MODALITY_NUM + token_tags).astype(mx.int32)

        # Sync per block: without this, MLX enqueues the entire 50-block graph
        # while the GPU lags behind, allocating every intermediate at once
        # (observed as an OOM kill on 36 GiB Macs at 480x832).
        recorded: list[Any] = []
        if record_prefix_kv is not None:
            from fastvideo.mlx_runtime.minimax_h3_prefix_kv import frozen_prefix_rows

            frozen_rows = frozen_prefix_rows(layout)
            if record_prefix_kv.rows != frozen_rows or record_prefix_kv.num_layers != len(self.blocks):
                raise ValueError(f"Prefix K/V cache is sized for {record_prefix_kv.num_layers} layers x "
                                 f"{record_prefix_kv.rows} rows, the step has {len(self.blocks)} x {frozen_rows}.")
            record_prefix_kv.hidden_dtype = packed.dtype
            record_prefix_kv.recorded_step = step_index
        for block_index, (block, tables) in enumerate(zip(self.blocks, cache.block_tables, strict=True)):
            vsa_kwargs = self._vsa_block_kwargs(block_index, step_index)
            if record_prefix_kv is not None:

                def kv_sink(k, v, _layer=block_index):
                    recorded.extend(record_prefix_kv.store(_layer, k[:frozen_rows], v[:frozen_rows]))

                vsa_kwargs["kv_sink"] = kv_sink
            packed = _transformer_block(
                block,
                packed,
                tables,
                adaln_indices,
                cos,
                sin,
                num_heads=self.num_heads,
                head_dim=self.head_dim,
                norm_eps=self.norm_eps,
                qk_norm_eps=self.qk_norm_eps,
                **vsa_kwargs,
            )
            mx.eval(packed, *recorded)
            recorded.clear()
            if self.last_vsa_stats is not None:
                self.last_vsa_stats.record(vsa_kwargs["vsa_stats"])

        row_positions = positions[row_inverse]
        normed = _h3_rms_norm(packed, self.weights["norm_out.norm.weight"], self.final_norm_eps)
        normed = normed * (1.0 + cache.norm_out_scale[row_positions]) + cache.norm_out_shift[row_positions]
        video_weight, video_bias = self._fused_output_params("proj_out.weight", "video", pdd_block,
                                                             pdd_integration_weights)
        audio_weight, audio_bias = self._fused_output_params("audio_proj_out.weight", "audio", pdd_block,
                                                             pdd_integration_weights)
        video_output = linear(
            normed[mx.array(layout.video_indices)].astype(weight_dtype(video_weight)),
            video_weight,
            video_bias,
        )
        audio_output = linear(
            normed[mx.array(layout.audio_indices)].astype(weight_dtype(audio_weight)),
            audio_weight,
            audio_bias,
        )
        return video_output, audio_output

    def forward_target_rows(self, video_target_rows, audio_target_rows, **kwargs):
        """Denoise-step forward over the target rows only, reusing recorded prefix K/V.

        See ``minimax_h3_prefix_kv_forward.forward_target_rows``; returns the
        (video, audio) outputs of the target rows alone.
        """
        from fastvideo.mlx_runtime.minimax_h3_prefix_kv_forward import forward_target_rows

        return forward_target_rows(self, video_target_rows, audio_target_rows, **kwargs)

    __call__ = forward

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.compute_temb

compute_temb(timesteps)

(n,) timesteps -> (n, time_embed_dim) through time_proj + MLP.

Source
def compute_temb(self, timesteps):
    """(n,) timesteps -> (n, time_embed_dim) through time_proj + MLP."""
    t_freq = timestep_embedding(timesteps,
                                self.freq_dim).astype(weight_dtype(self.weights["time_embedder.linear_1.weight"]))
    temb = linear(
        t_freq,
        self.weights["time_embedder.linear_1.weight"],
        self.weights["time_embedder.linear_1.bias"],
    )
    temb = linear(
        silu(temb),
        self.weights["time_embedder.linear_2.weight"],
        self.weights["time_embedder.linear_2.bias"],
    )
    if self.adaln_rank is not None:
        temb = linear(
            silu(temb).astype(weight_dtype(self.weights["adaln_basis.weight"])), self.weights["adaln_basis.weight"])
    return temb

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.refine_text

refine_text(text_rows)
Source
def refine_text(self, text_rows):
    hidden = linear(
        text_rows.astype(weight_dtype(self.weights["context_embedder.weight"])),
        self.weights["context_embedder.weight"],
        self.weights["context_embedder.bias"],
    )
    for block in self.refiner:
        hidden = _refiner_block(
            block,
            hidden,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            norm_eps=self.norm_eps,
            qk_norm_eps=self.qk_norm_eps,
        )
    return _h3_rms_norm(hidden, self.weights["token_refiner.final_norm.weight"], self.final_norm_eps)

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.precompute_adaln

precompute_adaln(timesteps: np.ndarray, *, drop_weights: bool = True) -> MiniMaxH3StepCache

Evaluate every modulation table for a fixed schedule, once.

timesteps is the union (sorted, unique) of all per-step video/audio/condition timesteps the sampler will use. With drop_weights=True the per-block adaln_proj.linear weights are released afterwards — the memory win this runtime exists for.

Source
def precompute_adaln(self, timesteps: np.ndarray, *, drop_weights: bool = True) -> MiniMaxH3StepCache:
    """Evaluate every modulation table for a fixed schedule, once.

    ``timesteps`` is the union (sorted, unique) of all per-step
    video/audio/condition timesteps the sampler will use. With
    ``drop_weights=True`` the per-block ``adaln_proj.linear`` weights are
    released afterwards — the memory win this runtime exists for.
    """
    import mlx.core as mx

    timesteps = np.unique(np.asarray(timesteps, dtype=np.float32))
    temb = self.compute_temb(mx.array(timesteps))
    block_tables = [_adaln_tables(block, temb, apply_silu=self.adaln_rank is None) for block in self.blocks]
    shift_scale = linear(
        (silu(temb) if self.adaln_rank is None else temb).astype(
            weight_dtype(self.weights["norm_out.linear.weight"])),
        self.weights["norm_out.linear.weight"],
        self.weights["norm_out.linear.bias"],
    )
    norm_out_shift, norm_out_scale = mx.split(shift_scale, 2, axis=-1)
    mx.eval(block_tables, norm_out_scale, norm_out_shift)
    self._adaln_cache = MiniMaxH3StepCache(
        timesteps=timesteps,
        block_tables=block_tables,
        norm_out_shift=norm_out_shift,
        norm_out_scale=norm_out_scale,
    )
    if drop_weights:
        for block in self.blocks:
            block["adaln_proj.linear.weight"] = None
            block["adaln_proj.linear.bias"] = None
    return self._adaln_cache

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.configure_vsa

configure_vsa(config: MiniMaxH3VSAConfig | None) -> None
Source
def configure_vsa(self, config: MiniMaxH3VSAConfig | None) -> None:
    config = config or MiniMaxH3VSAConfig()
    if config.enabled and not self._vsa_capable:
        raise dense_only_vsa_error("this MLX H3 DiT")
    if any(layer >= len(self.blocks) for layer in config.dense_layers):
        raise ValueError(f"VSA dense_layers must be below the model's block count ({len(self.blocks)}).")
    self.vsa_config = config
    self._vsa_geometry = None
    self.reset_vsa_stats()

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.reset_vsa_stats

reset_vsa_stats() -> None
Source
def reset_vsa_stats(self) -> None:
    self.last_vsa_stats = MiniMaxH3VSAStats(
        configured_sparsity=self.vsa_config.sparsity,
        tile_size=self.vsa_config.tile_size,
        prefix_mode=self.vsa_config.prefix_mode,
    ) if self.vsa_config.enabled else None

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.prepare_vsa_geometry

prepare_vsa_geometry(layout: MiniMaxH3PackedLayout) -> MiniMaxH3VSAGeometry | None
Source
def prepare_vsa_geometry(self, layout: MiniMaxH3PackedLayout) -> MiniMaxH3VSAGeometry | None:
    if not self.vsa_config.enabled:
        self._vsa_geometry = None
        return None
    config = self.vsa_config
    ref_keep_rate = config.effective_ref_keep_rate
    if ref_keep_rate is not None and layout.reference_segments:
        segments = ref2va_vsa_segments_from_layout(layout, self.patch_size)
    else:
        segments = (*prefix_segments_from_layout(layout, self.patch_size),
                    dit_seq_shape_from_layout(layout, self.patch_size))
    segments = tuple(segment for segment in segments if isinstance(segment, tuple) or segment > 0)
    geometry = self._vsa_geometry
    if geometry is not None and (geometry.segments, geometry.tile_elems) == (segments, config.tile_size):
        return geometry
    try:
        geometry = build_h3_segment_geometry(segments, config.tile_size)
    except ValueError as error:
        raise ValueError(f"H3 VSA cannot run with this packed layout (segments={segments}, "
                         f"tile_size={config.tile_size}): {error}. "
                         "Disable VSA or pick a supported tile size (64, 128 or 256).") from error
    self._vsa_geometry = geometry
    return self._vsa_geometry

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.forward

forward(video_rows, audio_rows, text_rows, *, position_ids, token_tags, timestep_indices, timesteps, video_indices, audio_indices, text_indices, pdd_block: tuple[int, int] | None = None, pdd_integration_weights: dict[str, np.ndarray] | None = None)

Faithful port of the torch forward (batch-1, rows already patchified).

Returns (video_output, audio_output) rows. The torch reference projects all packed rows through both heads and then selects; this selects first and projects only the relevant rows — row-wise identical math at a fraction of the cost.

This entrypoint is dense-only. VSA requires forward_with_cache with a packed layout and step index. With a PDD checkpoint pass the step's pdd_block and pdd_integration_weights so the widened heads are fused.

Source
def forward(
    self,
    video_rows,
    audio_rows,
    text_rows,
    *,
    position_ids,
    token_tags,
    timestep_indices,
    timesteps,
    video_indices,
    audio_indices,
    text_indices,
    pdd_block: tuple[int, int] | None = None,
    pdd_integration_weights: dict[str, np.ndarray] | None = None,
):
    """Faithful port of the torch forward (batch-1, rows already patchified).

    Returns (video_output, audio_output) rows. The torch reference
    projects *all* packed rows through both heads and then selects; this
    selects first and projects only the relevant rows — row-wise
    identical math at a fraction of the cost.

    This entrypoint is dense-only. VSA requires ``forward_with_cache``
    with a packed layout and step index. With a PDD checkpoint pass the
    step's ``pdd_block`` and ``pdd_integration_weights`` so the widened
    heads are fused.
    """
    import mlx.core as mx

    if self.vsa_config.enabled:
        raise ValueError("VSA requires forward_with_cache() with layout and step_index; forward() is dense-only.")
    sequence_length = position_ids.shape[0]
    cos, sin = rope_cos_sin(position_ids, self.rope_freq_dim, self.rope_theta)

    video_embeds = linear(
        video_rows.astype(weight_dtype(self.weights["proj_in.weight"])),
        self.weights["proj_in.weight"],
        self.weights["proj_in.bias"],
    )
    audio_embeds = linear(
        audio_rows.astype(weight_dtype(self.weights["audio_proj_in.weight"])),
        self.weights["audio_proj_in.weight"],
        self.weights["audio_proj_in.bias"],
    )
    text_embeds = self.refine_text(text_rows)

    packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
    packed[_as_mx_indices(text_indices)] = text_embeds
    packed[_as_mx_indices(video_indices)] = video_embeds.astype(text_embeds.dtype)
    packed[_as_mx_indices(audio_indices)] = audio_embeds.astype(text_embeds.dtype)

    temb = self.compute_temb(timesteps)
    adaln_indices = (timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags).astype(mx.int32)

    for block_index, block in enumerate(self.blocks):
        tables = _adaln_tables(block, temb, apply_silu=self.adaln_rank is None)
        packed = _transformer_block(
            block,
            packed,
            tables,
            adaln_indices,
            cos,
            sin,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            norm_eps=self.norm_eps,
            qk_norm_eps=self.qk_norm_eps,
            **self._vsa_block_kwargs(block_index, 0),
        )
        mx.eval(packed)  # per-block sync: see forward_with_cache note

    shift_scale = linear(
        (silu(temb) if self.adaln_rank is None else temb).astype(
            weight_dtype(self.weights["norm_out.linear.weight"])),
        self.weights["norm_out.linear.weight"],
        self.weights["norm_out.linear.bias"],
    )
    shift_rows, scale_rows = mx.split(shift_scale, 2, axis=-1)
    normed = _h3_rms_norm(packed, self.weights["norm_out.norm.weight"], self.final_norm_eps)
    normed = normed * (1.0 + scale_rows[timestep_indices]) + shift_rows[timestep_indices]
    video_weight, video_bias = self._fused_output_params("proj_out.weight", "video", pdd_block,
                                                         pdd_integration_weights)
    audio_weight, audio_bias = self._fused_output_params("audio_proj_out.weight", "audio", pdd_block,
                                                         pdd_integration_weights)
    video_output = linear(
        normed[_as_mx_indices(video_indices)].astype(weight_dtype(video_weight)),
        video_weight,
        video_bias,
    )
    audio_output = linear(
        normed[_as_mx_indices(audio_indices)].astype(weight_dtype(audio_weight)),
        audio_weight,
        audio_bias,
    )
    return video_output, audio_output

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.forward_with_cache

forward_with_cache(video_rows, audio_rows, text_rows, *, layout: MiniMaxH3PackedLayout, step_timesteps: np.ndarray, row_timestep_inverse: np.ndarray, step_index: int = 0, pdd_block: tuple[int, int] | None = None, pdd_integration_weights: dict[str, np.ndarray] | None = None, record_prefix_kv: Any | None = None)

Denoise-step forward served entirely from the AdaLN cache.

record_prefix_kv (a MiniMaxH3PrefixKVCache) additionally keeps every block's K/V of the frozen text/reference rows for later forward_target_rows steps; the outputs are unchanged.

step_timesteps are this step's unique timesteps (sorted); row_timestep_inverse maps each packed row to one of them (the build_row_timesteps inverse). The cache must cover every value in step_timesteps. With a PDD checkpoint pass the step's pdd_block and pdd_integration_weights so the widened output heads are fused into the block's weighted mean.

Source
def forward_with_cache(
    self,
    video_rows,
    audio_rows,
    text_rows,
    *,
    layout: MiniMaxH3PackedLayout,
    step_timesteps: np.ndarray,
    row_timestep_inverse: np.ndarray,
    step_index: int = 0,
    pdd_block: tuple[int, int] | None = None,
    pdd_integration_weights: dict[str, np.ndarray] | None = None,
    record_prefix_kv: Any | None = None,
):
    """Denoise-step forward served entirely from the AdaLN cache.

    ``record_prefix_kv`` (a ``MiniMaxH3PrefixKVCache``) additionally keeps
    every block's K/V of the frozen text/reference rows for later
    ``forward_target_rows`` steps; the outputs are unchanged.

    ``step_timesteps`` are this step's unique timesteps (sorted);
    ``row_timestep_inverse`` maps each packed row to one of them (the
    ``build_row_timesteps`` inverse). The cache must cover every value in
    ``step_timesteps``. With a PDD checkpoint pass the step's
    ``pdd_block`` and ``pdd_integration_weights`` so the widened output
    heads are fused into the block's weighted mean.
    """
    import mlx.core as mx

    if self._adaln_cache is None:
        raise RuntimeError("precompute_adaln() must run before forward_with_cache().")
    if self.vsa_config.enabled:
        self.prepare_vsa_geometry(layout)
    cache = self._adaln_cache
    positions = mx.array(cache.positions(step_timesteps))

    sequence_length = layout.sequence_length
    position_ids = mx.array(layout.position_ids)
    cos, sin = rope_cos_sin(position_ids, self.rope_freq_dim, self.rope_theta)

    video_embeds = linear(
        video_rows.astype(weight_dtype(self.weights["proj_in.weight"])),
        self.weights["proj_in.weight"],
        self.weights["proj_in.bias"],
    )
    audio_embeds = linear(
        audio_rows.astype(weight_dtype(self.weights["audio_proj_in.weight"])),
        self.weights["audio_proj_in.weight"],
        self.weights["audio_proj_in.bias"],
    )
    text_embeds = self.refine_text(text_rows)

    packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
    packed[mx.array(layout.text_indices)] = text_embeds
    packed[mx.array(layout.video_indices)] = video_embeds.astype(text_embeds.dtype)
    packed[mx.array(layout.audio_indices)] = audio_embeds.astype(text_embeds.dtype)

    token_tags = mx.array(layout.token_tags)
    row_inverse = mx.array(row_timestep_inverse)
    adaln_indices = (positions[row_inverse] * MINIMAX_H3_MODALITY_NUM + token_tags).astype(mx.int32)

    # Sync per block: without this, MLX enqueues the entire 50-block graph
    # while the GPU lags behind, allocating every intermediate at once
    # (observed as an OOM kill on 36 GiB Macs at 480x832).
    recorded: list[Any] = []
    if record_prefix_kv is not None:
        from fastvideo.mlx_runtime.minimax_h3_prefix_kv import frozen_prefix_rows

        frozen_rows = frozen_prefix_rows(layout)
        if record_prefix_kv.rows != frozen_rows or record_prefix_kv.num_layers != len(self.blocks):
            raise ValueError(f"Prefix K/V cache is sized for {record_prefix_kv.num_layers} layers x "
                             f"{record_prefix_kv.rows} rows, the step has {len(self.blocks)} x {frozen_rows}.")
        record_prefix_kv.hidden_dtype = packed.dtype
        record_prefix_kv.recorded_step = step_index
    for block_index, (block, tables) in enumerate(zip(self.blocks, cache.block_tables, strict=True)):
        vsa_kwargs = self._vsa_block_kwargs(block_index, step_index)
        if record_prefix_kv is not None:

            def kv_sink(k, v, _layer=block_index):
                recorded.extend(record_prefix_kv.store(_layer, k[:frozen_rows], v[:frozen_rows]))

            vsa_kwargs["kv_sink"] = kv_sink
        packed = _transformer_block(
            block,
            packed,
            tables,
            adaln_indices,
            cos,
            sin,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            norm_eps=self.norm_eps,
            qk_norm_eps=self.qk_norm_eps,
            **vsa_kwargs,
        )
        mx.eval(packed, *recorded)
        recorded.clear()
        if self.last_vsa_stats is not None:
            self.last_vsa_stats.record(vsa_kwargs["vsa_stats"])

    row_positions = positions[row_inverse]
    normed = _h3_rms_norm(packed, self.weights["norm_out.norm.weight"], self.final_norm_eps)
    normed = normed * (1.0 + cache.norm_out_scale[row_positions]) + cache.norm_out_shift[row_positions]
    video_weight, video_bias = self._fused_output_params("proj_out.weight", "video", pdd_block,
                                                         pdd_integration_weights)
    audio_weight, audio_bias = self._fused_output_params("audio_proj_out.weight", "audio", pdd_block,
                                                         pdd_integration_weights)
    video_output = linear(
        normed[mx.array(layout.video_indices)].astype(weight_dtype(video_weight)),
        video_weight,
        video_bias,
    )
    audio_output = linear(
        normed[mx.array(layout.audio_indices)].astype(weight_dtype(audio_weight)),
        audio_weight,
        audio_bias,
    )
    return video_output, audio_output

fastvideo.mlx_runtime.minimax_h3.MLXMiniMaxH3DiT.forward_target_rows

forward_target_rows(video_target_rows, audio_target_rows, **kwargs)

Denoise-step forward over the target rows only, reusing recorded prefix K/V.

See minimax_h3_prefix_kv_forward.forward_target_rows; returns the (video, audio) outputs of the target rows alone.

Source
def forward_target_rows(self, video_target_rows, audio_target_rows, **kwargs):
    """Denoise-step forward over the target rows only, reusing recorded prefix K/V.

    See ``minimax_h3_prefix_kv_forward.forward_target_rows``; returns the
    (video, audio) outputs of the target rows alone.
    """
    from fastvideo.mlx_runtime.minimax_h3_prefix_kv_forward import forward_target_rows

    return forward_target_rows(self, video_target_rows, audio_target_rows, **kwargs)

fastvideo.mlx_runtime.minimax_h3.mlx_h3_dit_from_diffusers_safetensors

mlx_h3_dit_from_diffusers_safetensors(transformer_path: str | Path, config: dict[str, Any] | None = None, *, dtype: str = 'fp16', num_blocks: int | None = None, quantization: str | MLXQuantizationSpec | None = None, adaln_cache_timesteps: np.ndarray | None = None, include_vsa: bool = False) -> MLXMiniMaxH3DiT

Load the released H3 transformer (diffusers layout) into MLX.

transformer_path is the transformer/ directory of the HF repo (or a single safetensors file, e.g. a student checkpoint). config defaults to config.json next to the weights. fp32-release modules stay fp32; attention/FFN matrices are quantized when quantization is set.

When adaln_cache_timesteps is provided, AdaLN tables are computed one block at a time while the checkpoint is read. The 13B projection weights are never retained together and are omitted from the returned model. This is the memory-bounded build path for the released 33B student.

include_vsa keeps attn.to_gate_compress matrices and quantizes them with the same affine grid as the other linear weights. Dense conversion continues to drop them.

Source
def mlx_h3_dit_from_diffusers_safetensors(
    transformer_path: str | Path,
    config: dict[str, Any] | None = None,
    *,
    dtype: str = "fp16",
    num_blocks: int | None = None,
    quantization: str | MLXQuantizationSpec | None = None,
    adaln_cache_timesteps: np.ndarray | None = None,
    include_vsa: bool = False,
) -> MLXMiniMaxH3DiT:
    """Load the released H3 transformer (diffusers layout) into MLX.

    ``transformer_path`` is the ``transformer/`` directory of the HF repo (or
    a single safetensors file, e.g. a student checkpoint). ``config`` defaults
    to ``config.json`` next to the weights. fp32-release modules stay fp32;
    attention/FFN matrices are quantized when ``quantization`` is set.

    When ``adaln_cache_timesteps`` is provided, AdaLN tables are computed one
    block at a time while the checkpoint is read. The 13B projection weights
    are never retained together and are omitted from the returned model. This
    is the memory-bounded build path for the released 33B student.

    ``include_vsa`` keeps ``attn.to_gate_compress`` matrices and quantizes them
    with the same affine grid as the other linear weights. Dense conversion
    continues to drop them.
    """
    import mlx.core as mx
    transformer_path = Path(transformer_path)
    if config is None:
        config_path = transformer_path / "config.json" if transformer_path.is_dir() else transformer_path.with_name(
            "config.json")
        config = json.loads(Path(config_path).read_text())
    total_blocks = int(config["num_layers"])
    if num_blocks is None:
        num_blocks = total_blocks
    num_refiner = int(config["num_refiner_layers"])
    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)

    weights: dict[str, Any] = {}
    blocks: list[dict[str, Any] | None] = [None] * num_blocks
    refiner: list[dict[str, Any] | None] = [None] * num_refiner
    cache_timesteps = None
    temb = None
    cached_block_tables: list[tuple[Any, ...] | None] | None = None
    pending_adaln: dict[int, dict[str, Any]] = {}

    def assign(key: str, value) -> None:
        if key.startswith("transformer_blocks."):
            _, index_str, sub = key.split(".", 2)
            index = int(index_str)
            if index < num_blocks:
                block = blocks[index]
                if block is None:
                    block = {}
                    blocks[index] = block
                block[sub] = value
        elif key.startswith("token_refiner.refiner_blocks."):
            _, _, index_str, sub = key.split(".", 3)
            index = int(index_str)
            refiner_block = refiner[index]
            if refiner_block is None:
                refiner_block = {}
                refiner[index] = refiner_block
            refiner_block[sub] = value
        else:
            weights[key] = value

    if adaln_cache_timesteps is not None:
        cache_timesteps = np.unique(np.asarray(adaln_cache_timesteps, dtype=np.float32))
        for shard in _safetensors_shards(transformer_path):
            shard_arrays = mx.load(str(shard))
            for key, source in shard_arrays.items():
                if not (key.startswith("time_embedder.") or key == "adaln_basis.weight"):
                    continue
                keep_fp32 = key.split(".", 1)[0] in FP32_MODULE_PREFIXES
                target_dtype = mx.float32 if keep_fp32 else (mx.float16 if key == "adaln_basis.weight" else cast_dtype)
                assign(key, _load_array(source, target_dtype))
            del shard_arrays
        required_time_keys = {
            "time_embedder.linear_1.weight",
            "time_embedder.linear_1.bias",
            "time_embedder.linear_2.weight",
            "time_embedder.linear_2.bias",
        }
        missing_time_keys = sorted(required_time_keys - weights.keys())
        if missing_time_keys:
            raise KeyError(f"Missing time-embedder weights needed for the AdaLN cache: {missing_time_keys}")
        timestep_rows = mx.array(cache_timesteps)
        t_freq = timestep_embedding(timestep_rows, int(config["freq_dim"])).astype(
            weight_dtype(weights["time_embedder.linear_1.weight"]))
        temb = linear(t_freq, weights["time_embedder.linear_1.weight"], weights["time_embedder.linear_1.bias"])
        temb = linear(silu(temb), weights["time_embedder.linear_2.weight"], weights["time_embedder.linear_2.bias"])
        if config.get("adaln_rank") is not None:
            if "adaln_basis.weight" not in weights:
                raise KeyError("Rank-reduced AdaLN checkpoint is missing adaln_basis.weight")
            temb = linear(silu(temb).astype(weight_dtype(weights["adaln_basis.weight"])), weights["adaln_basis.weight"])
        mx.eval(temb)
        cached_block_tables = [None] * num_blocks

    for shard in _safetensors_shards(transformer_path):
        shard_arrays = mx.load(str(shard))
        for key, source in shard_arrays.items():
            if _is_ignored_dense_key(key, include_vsa=include_vsa):
                continue
            if source.dtype == mx.uint8 and key.endswith(".weight"):
                raise ValueError(f"Packed transformer weight {key} is not a floating-point source. "
                                 "Convert the released BF16 transformer to MLX INT6 instead.")
            if temb is not None and (key.startswith("time_embedder.") or key == "adaln_basis.weight"):
                continue
            if key.startswith("transformer_blocks."):
                index = int(key.split(".")[1])
                if index >= num_blocks:
                    continue
            if key.startswith("rope."):
                continue  # non-persistent analytic buffer, rebuilt on the fly
            keep_fp32 = key.split(".", 1)[0] in FP32_MODULE_PREFIXES
            factorized_adaln = config.get("adaln_rank") is not None and (".adaln_proj." in key or key.startswith(
                ("norm_out.linear.", "adaln_basis.")))
            target_dtype = mx.float32 if keep_fp32 else (mx.float16 if factorized_adaln else cast_dtype)
            array = _load_array(source, target_dtype)
            if temb is not None and ".adaln_proj.linear." in key:
                _, index_str, sub = key.split(".", 2)
                index = int(index_str)
                block_pending = pending_adaln.setdefault(index, {})
                block_pending[sub] = array
                if {"adaln_proj.linear.weight", "adaln_proj.linear.bias"} <= block_pending.keys():
                    tables = _adaln_tables(block_pending, temb, apply_silu=config.get("adaln_rank") is None)
                    mx.eval(tables)
                    assert cached_block_tables is not None
                    cached_block_tables[index] = tables
                    block = blocks[index]
                    if block is None:
                        block = {}
                        blocks[index] = block
                    block["adaln_proj.linear.weight"] = None
                    block["adaln_proj.linear.bias"] = None
                    del pending_adaln[index]
                continue
            if spec is not None and _is_quantizable(key, include_vsa=include_vsa) and not keep_fp32:
                value = quantize_matrix(array, spec)
                del array
            else:
                value = array
            _eval_value(value)
            assign(key, value)
        del shard_arrays

    cache = None
    if temb is not None:
        if pending_adaln:
            raise KeyError(f"Incomplete AdaLN projection pairs for blocks {sorted(pending_adaln)}")
        assert cache_timesteps is not None and cached_block_tables is not None
        missing_cache_blocks = [index for index, tables in enumerate(cached_block_tables) if tables is None]
        if missing_cache_blocks:
            raise KeyError(f"Missing AdaLN cache tables for blocks {missing_cache_blocks}")
        shift_scale = linear(
            (silu(temb) if config.get("adaln_rank") is None else temb).astype(
                weight_dtype(weights["norm_out.linear.weight"])),
            weights["norm_out.linear.weight"],
            weights["norm_out.linear.bias"],
        )
        norm_out_shift, norm_out_scale = mx.split(shift_scale, 2, axis=-1)
        mx.eval(norm_out_shift, norm_out_scale)
        cache = MiniMaxH3StepCache(
            timesteps=cache_timesteps,
            block_tables=[tables for tables in cached_block_tables if tables is not None],
            norm_out_shift=norm_out_shift,
            norm_out_scale=norm_out_scale,
        )

    if any(block is None for block in blocks):
        missing = [i for i, block in enumerate(blocks) if block is None]
        raise KeyError(f"Missing transformer block weights for indices {missing}.")
    if any(block is None for block in refiner):
        raise KeyError("Missing token refiner weights.")
    loaded_blocks = [block for block in blocks if block is not None]
    if include_vsa:
        missing_gates = [
            index for index, block in enumerate(loaded_blocks) if "attn.to_gate_compress.weight" not in block
        ]
        if missing_gates:
            raise KeyError("VSA conversion requested but gate projections are missing for transformer blocks "
                           f"{missing_gates}. The source checkpoint must contain "
                           f"`transformer_blocks.*.{VSA_GATE_KEY_SUFFIX}`.")
        logger.info("Retained %d VSA gate projections (%s).", len(loaded_blocks), VSA_GATE_KEY_SUFFIX)
    dit = MLXMiniMaxH3DiT(
        weights,
        loaded_blocks,
        [block for block in refiner if block is not None],
        config,
    )
    dit._adaln_cache = cache
    return dit

fastvideo.mlx_runtime.minimax_h3.mlx_h3_checkpoint_vsa_capable

mlx_h3_checkpoint_vsa_capable(checkpoint_dir: str | Path) -> bool

True when a saved MLX H3 checkpoint retained the VSA gate projections.

Source
def mlx_h3_checkpoint_vsa_capable(checkpoint_dir: str | Path) -> bool:
    """True when a saved MLX H3 checkpoint retained the VSA gate projections."""
    manifest_path = Path(checkpoint_dir) / H3_MANIFEST_FILENAME
    if not manifest_path.is_file():
        return False
    manifest = json.loads(manifest_path.read_text())
    if bool((manifest.get("vsa") or {}).get("capable")):
        return True
    quantized_keys = manifest.get("quantized_keys") or {}
    return any("attn.to_gate_compress.weight" in key for key in quantized_keys)

fastvideo.mlx_runtime.minimax_h3.save_mlx_h3_checkpoint

save_mlx_h3_checkpoint(dit: MLXMiniMaxH3DiT, checkpoint_dir: str | Path) -> Path

Persist a (possibly quantized, possibly AdaLN-dropped) H3 DiT.

Source
def save_mlx_h3_checkpoint(dit: MLXMiniMaxH3DiT, checkpoint_dir: str | Path) -> Path:
    """Persist a (possibly quantized, possibly AdaLN-dropped) H3 DiT."""
    import mlx.core as mx

    checkpoint_dir = Path(checkpoint_dir)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    arrays: dict[str, Any] = {}
    quantized: dict[str, dict[str, Any]] = {}
    spec: MLXQuantizationSpec | None = None
    for key, value in _flatten_h3_weights(dit).items():
        if isinstance(value, QuantizedMatrix):
            if spec is not None and value.spec != spec:
                raise ValueError(f"Mixed quantization specs in one checkpoint ({spec} vs {value.spec} at '{key}').")
            spec = value.spec
            arrays[key] = value.weight
            arrays[f"{key}.scales"] = value.scales
            if value.biases is not None:
                arrays[f"{key}.biases"] = value.biases
            quantized[key] = {
                "dequantized_dtype": _dtype_name(value.dequantized_dtype),
                "has_biases": value.biases is not None,
            }
        else:
            arrays[key] = value

    cache_manifest = None
    if dit._adaln_cache is not None:
        cache = dit._adaln_cache
        cache_manifest = {
            "timesteps": cache.timesteps.tolist(),
            "num_blocks": len(cache.block_tables),
            "tables_per_block": 6,
        }
        for block_index, tables in enumerate(cache.block_tables):
            for table_index, table in enumerate(tables):
                arrays[f"__adaln_cache.block.{block_index}.table.{table_index}"] = table
        arrays["__adaln_cache.norm_out_shift"] = cache.norm_out_shift
        arrays["__adaln_cache.norm_out_scale"] = cache.norm_out_scale

    manifest = {
        "format_version": H3_FORMAT_VERSION,
        "config": dit.config,
        "num_blocks": len(dit.blocks),
        "num_refiner_blocks": len(dit.refiner),
        "quantization": None if spec is None else {
            "mode": spec.mode,
            "bits": spec.bits,
            "group_size": spec.group_size,
        },
        "quantized_keys": quantized,
        "adaln_cache": cache_manifest,
        "vsa": {
            "capable":
            bool(dit.vsa_capable),
            "num_gate_matrices":
            sum(1 for block in dit.blocks if "attn.to_gate_compress.weight" in block),
            "gate_key_suffix":
            VSA_GATE_KEY_SUFFIX,
            "attention_activations":
            "bf16",
            "note": ("Gate projections are quantized with the checkpoint's affine INT8/INT6/INT4 "
                     "weight-only grid. VSA attention Q/K/V remain unquantized activations."),
        },
    }
    weights_path = checkpoint_dir / H3_WEIGHTS_FILENAME
    mx.save_safetensors(str(weights_path), arrays)
    (checkpoint_dir / H3_MANIFEST_FILENAME).write_text(json.dumps(manifest, indent=2))
    logger.info("Saved MLX H3 DiT checkpoint (%d arrays, quantization=%s) to %s", len(arrays),
                spec.label if spec else "none", checkpoint_dir)
    return checkpoint_dir

fastvideo.mlx_runtime.minimax_h3.load_mlx_h3_checkpoint

load_mlx_h3_checkpoint(checkpoint_dir: str | Path) -> MLXMiniMaxH3DiT

Rebuild an H3 DiT saved by :func:save_mlx_h3_checkpoint.

Refuses a quantization grid the installed MLX cannot execute, loudly — never silently dequantizes onto a different grid.

Source
def load_mlx_h3_checkpoint(checkpoint_dir: str | Path) -> MLXMiniMaxH3DiT:
    """Rebuild an H3 DiT saved by :func:`save_mlx_h3_checkpoint`.

    Refuses a quantization grid the installed MLX cannot execute, loudly —
    never silently dequantizes onto a different grid.
    """
    import mlx.core as mx

    checkpoint_dir = Path(checkpoint_dir)
    manifest_path = checkpoint_dir / H3_MANIFEST_FILENAME
    weights_path = checkpoint_dir / H3_WEIGHTS_FILENAME
    if not manifest_path.exists() or not weights_path.exists():
        raise FileNotFoundError(f"Not an MLX H3 checkpoint directory: {checkpoint_dir} "
                                f"(expected {H3_MANIFEST_FILENAME} and {H3_WEIGHTS_FILENAME}).")

    manifest = json.loads(manifest_path.read_text())
    version = manifest.get("format_version")
    if version != H3_FORMAT_VERSION:
        raise ValueError(f"MLX H3 checkpoint {checkpoint_dir} has format_version={version}; "
                         f"this build reads version {H3_FORMAT_VERSION}. Re-export the checkpoint.")

    spec = None
    if manifest["quantization"] is not None:
        spec = MLXQuantizationSpec(**manifest["quantization"])
        ensure_quantization_supported(spec)

    arrays = mx.load(str(weights_path))
    quantized_keys: dict[str, dict[str, Any]] = manifest["quantized_keys"]

    def rebuild(key: str):
        if key not in quantized_keys:
            return arrays[key]
        info = quantized_keys[key]
        assert spec is not None, f"Quantized key '{key}' in a checkpoint without a quantization spec"
        return QuantizedMatrix(
            weight=arrays[key],
            scales=arrays[f"{key}.scales"],
            biases=arrays[f"{key}.biases"] if info["has_biases"] else None,
            spec=spec,
            dequantized_dtype=_name_to_dtype(info["dequantized_dtype"]),
        )

    weights: dict[str, Any] = {}
    blocks: list[dict[str, Any] | None] = [None] * int(manifest["num_blocks"])
    refiner: list[dict[str, Any] | None] = [None] * int(manifest["num_refiner_blocks"])
    for key in arrays:
        if key.startswith("__adaln_cache."):
            continue
        if key.endswith(".scales") or key.endswith(".biases"):
            continue
        value = rebuild(key)
        if key.startswith("blocks."):
            _, index_str, sub = key.split(".", 2)
            index = int(index_str)
            block = blocks[index]
            if block is None:
                block = {}
                blocks[index] = block
            block[sub] = value
        elif key.startswith("refiner."):
            _, index_str, sub = key.split(".", 2)
            index = int(index_str)
            refiner_block = refiner[index]
            if refiner_block is None:
                refiner_block = {}
                refiner[index] = refiner_block
            refiner_block[sub] = value
        else:
            weights[key] = value

    if any(block is None for block in blocks) or any(block is None for block in refiner):
        raise ValueError(f"MLX H3 checkpoint {checkpoint_dir} is missing block weights.")
    dit = MLXMiniMaxH3DiT(
        weights,
        [block for block in blocks if block is not None],
        [block for block in refiner if block is not None],
        manifest["config"],
    )
    cache_info = manifest.get("adaln_cache")
    if cache_info is not None:
        block_tables = []
        for block_index in range(int(cache_info["num_blocks"])):
            block_tables.append(
                tuple(arrays[f"__adaln_cache.block.{block_index}.table.{table_index}"]
                      for table_index in range(int(cache_info["tables_per_block"]))))
        dit._adaln_cache = MiniMaxH3StepCache(
            timesteps=np.asarray(cache_info["timesteps"], dtype=np.float32),
            block_tables=block_tables,
            norm_out_shift=arrays["__adaln_cache.norm_out_shift"],
            norm_out_scale=arrays["__adaln_cache.norm_out_scale"],
        )
    return dit