minimax_h3

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)

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

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

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

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

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()])

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))

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))

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))

@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&#91;tuple&#91;str, int, tuple&#91;int, int, int&#93;&#93;, ...&#93; = ()</code></pre></details>
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&#91;:num_text_tokens, 0&#93; = 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=&quot;ij&quot;)
frame_grid = np.stack(&#91;mesh_h.reshape(-1), mesh_w.reshape(-1)&#93;, axis=-1)

for index, anchor in enumerate(keyframe_anchors):
    if anchor == &quot;first&quot;:
        anchor_time = float(num_text_tokens)
    elif anchor == &quot;last&quot;:
        anchor_time = (float(num_text_tokens) + _temporal_position_span(num_latent_frames) -
                       MINIMAX_H3_ROPE_FRAME_RESCALE)
    else:
        raise ValueError(f&quot;A keyframe anchor must be &#x27;first&#x27; or &#x27;last&#x27;, got &#123;anchor!r&#125;.&quot;)
    rows = slice(condition_start + index * rows_per_frame, condition_start + (index + 1) * rows_per_frame)
    position_ids&#91;rows, 0&#93; = anchor_time
    position_ids&#91;rows, 1:&#93; = frame_grid

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

frame_time = temporal_position_grid(num_latent_frames, 0.0) * video_temporal_scale + float(num_text_tokens)
position_ids&#91;video_start:, 0&#93; = np.repeat(frame_time, rows_per_frame)
position_ids&#91;video_start:, 1:&#93; = np.tile(frame_grid, (num_latent_frames, 1))

video_indices = np.concatenate(&#91;
    np.arange(condition_start, audio_start),
    np.arange(video_start, sequence_length),
&#93;)
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&#91;text_indices&#93; = text_token_tags
token_tags&#91;audio_indices&#93; = MINIMAX_H3_AUDIO_TAG
token_tags&#91;video_indices&#93; = 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,
)</code></pre></details>
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.
&quot;&quot;&quot;
if patch_size != (1, 2, 2):
    raise ValueError(f&quot;MiniMax-H3 Ref2VA requires patch_size=(1, 2, 2), got &#123;patch_size&#125;.&quot;)
if text_token_tags.ndim != 1:
    raise ValueError(f&quot;text_token_tags must be one-dimensional, got &#123;text_token_tags.shape&#125;.&quot;)
if not np.isin(text_token_tags, (MINIMAX_H3_TEXT_TAG, MINIMAX_H3_VIDEO_TAG)).all():
    raise ValueError(&quot;text_token_tags may contain only text and vision tags.&quot;)
if not references:
    raise ValueError(&quot;Ref2VA requires at least one prepared reference.&quot;)

patch_t, patch_h, patch_w = patch_size
target_geometry = (num_latent_frames, latent_height, latent_width)
if any(value &lt;= 0 for value in target_geometry) or num_audio_latents &lt;= 0:
    raise ValueError(&quot;Ref2VA target latent geometry must be positive.&quot;)
if any(value % patch for value, patch in zip(target_geometry, patch_size, strict=True)):
    raise ValueError(f&quot;Target geometry &#123;target_geometry&#125; is not divisible by patch &#123;patch_size&#125;.&quot;)

visual_row_counts = &#91;
    0 if reference.media_type == &quot;audio&quot; else _num_reference_video_rows(reference, patch_size)
    for reference in references
&#93;
for reference in references:
    if reference.media_type == &quot;audio&quot; and not reference.has_audio:
        raise ValueError(&quot;An audio reference must carry decoded waveform latents.&quot;)
    if reference.has_audio and reference.num_audio_latents &lt;= 0:
        raise ValueError(&quot;An audio-bearing reference has no resolved audio latents.&quot;)

num_text_tokens = int(text_token_tags.shape&#91;0&#93;)
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&#91;:num_text_tokens, 0&#93; = 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&#91;np.ndarray&#93; = &#91;&#93;
audio_indices: list&#91;np.ndarray&#93; = &#91;&#93;
reference_segments: list&#91;tuple&#91;str, int, tuple&#91;int, int, int&#93;&#93;&#93; = &#91;&#93;
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 == &quot;image&quot;:
        reference_segments.append((&quot;image&quot;, 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&#91;rows, 0&#93; = rotary_time
        position_ids&#91;rows, 1:&#93; = frame_grid
        rotary_time += 1.0
    elif reference.media_type == &quot;audio&quot;:
        count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
        reference_segments.append((&quot;audio&quot;, 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 == &quot;video&quot;:
        audio_count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS if reference.has_audio else 0
        if audio_count:
            reference_segments.append((&quot;audio&quot;, audio_count, (0, 0, 0)))
        reference_segments.append((&quot;video&quot;, 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&#91;0&#93;
        position_ids&#91;video_rows, 0&#93; = np.repeat(frame_time, rows_per_frame)
        position_ids&#91;video_rows, 1:&#93; = 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&quot;Unsupported prepared reference type: &#123;reference.media_type!r&#125;.&quot;)

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&#91;video_start:, 0&#93; = np.repeat(frame_time, target_frame_grid.shape&#91;0&#93;)
position_ids&#91;video_start:, 1:&#93; = 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&#91;text_indices&#93; = text_token_tags.astype(np.int64)
token_tags&#91;packed_audio_indices&#93; = MINIMAX_H3_AUDIO_TAG
token_tags&#91;packed_video_indices&#93; = 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),
)</code></pre></details>
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)``.
&quot;&quot;&quot;
row_timesteps = np.full(layout.sequence_length, video_timestep, dtype=np.float64)
if layout.num_condition_video_rows:
    row_timesteps&#91;layout.video_indices&#91;:layout.num_condition_video_rows&#93;&#93; = condition_video_timestep
row_timesteps&#91;layout.audio_indices&#91;layout.num_condition_audio_rows:&#93;&#93; = audio_timestep
if layout.num_condition_audio_rows:
    row_timesteps&#91;layout.audio_indices&#91;:layout.num_condition_audio_rows&#93;&#93; = condition_audio_timestep
unique, inverse = np.unique(row_timesteps, return_inverse=True)
return unique.astype(np.float32), inverse.astype(np.int64)</code></pre></details>
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.
&quot;&quot;&quot;
if num_denoise_steps &lt; 1:
    raise ValueError(f&quot;num_denoise_steps must be &gt;= 1, got &#123;num_denoise_steps&#125;.&quot;)
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)</code></pre></details>
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=&quot;pdd&quot;`` and a FastH3 8-Step V2 DMD export (which carries
``dmd_denoising_steps``) is tagged ``kind=&quot;dmd&quot;``. A shifted sidecar that
cannot be parsed fails instead of silently baking the four-step grid.
&quot;&quot;&quot;
path = find_fasth3_inference_contract(model_root)
if path is None:
    return None
try:
    contract = json.loads(path.read_text(encoding=&quot;utf-8&quot;))
except json.JSONDecodeError as exc:
    raise ValueError(f&quot;Malformed FastH3 inference contract &#123;path&#125;: &#123;exc&#125;&quot;) from exc
if not isinstance(contract, dict):
    raise ValueError(f&quot;Malformed FastH3 inference contract &#123;path&#125;.&quot;)
if _is_pdd_contract(contract):
    return _load_pdd_contract(contract, path)
declares_shifts = &quot;video_scheduler_shift&quot; in contract or &quot;audio_scheduler_shift&quot; in contract
if not declares_shifts:
    return None
if contract.get(&quot;schema_version&quot;) != FASTH3_INFERENCE_CONTRACT_SCHEMA:
    raise ValueError(f&quot;Unsupported FastH3 inference contract schema in &#123;path&#125;.&quot;)
if &quot;video_scheduler_shift&quot; not in contract or &quot;audio_scheduler_shift&quot; not in contract:
    raise ValueError(f&quot;&#123;path&#125; must declare both video_scheduler_shift and audio_scheduler_shift.&quot;)
steps = validate_fasth3_dmd_rungs(contract.get(&quot;dmd_denoising_steps&quot;))
forwards = contract.get(&quot;transformer_forwards&quot;)
grid_points = contract.get(&quot;num_inference_steps&quot;)
if type(forwards) is not int or forwards != len(steps):
    raise ValueError(&quot;FastH3 transformer_forwards must equal the DMD rung count.&quot;)
if type(grid_points) is not int or grid_points != len(steps) + 1:
    raise ValueError(&quot;FastH3 num_inference_steps must equal the DMD rung count plus one.&quot;)
contract&#91;&quot;video_scheduler_shift&quot;&#93; = _positive_shift(contract&#91;&quot;video_scheduler_shift&quot;&#93;, name=&quot;video_scheduler_shift&quot;)
contract&#91;&quot;audio_scheduler_shift&quot;&#93; = _positive_shift(contract&#91;&quot;audio_scheduler_shift&quot;&#93;, name=&quot;audio_scheduler_shift&quot;)
contract&#91;&quot;dmd_denoising_steps&quot;&#93; = steps
contract&#91;&quot;kind&quot;&#93; = &quot;dmd&quot;
return contract</code></pre></details>
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</code></pre></details>
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 = &#123;
    &quot;video&quot;: PDDModalitySchedule(shift=contract&#91;&quot;video_scheduler_shift&quot;&#93;),
    &quot;audio&quot;: PDDModalitySchedule(shift=contract&#91;&quot;audio_scheduler_shift&quot;&#93;),
&#125;
return build_pdd_sampling_plan(contract&#91;&quot;pdd_step_indices&quot;&#93;, schedules)</code></pre></details>
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&#x27;s node sigmas.
&quot;&quot;&quot;
if type(num_steps) is not int or num_steps &lt; 1:
    raise ValueError(f&quot;num_steps must be a positive integer, got &#123;num_steps!r&#125;.&quot;)
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&quot;MLX H3 checkpoint has a fixed AdaLN ladder that does not support --steps &quot;
                         f&quot;&#123;num_steps&#125;. Use the step count used during conversion (normally 4), or re-export &quot;
                         &quot;the checkpoint.&quot;)
    return ResolvedH3Schedule(video, audio, union, num_steps, &quot;uniform&quot;)

if contract&#91;&quot;kind&quot;&#93; == &quot;pdd&quot;:
    return _resolve_pdd_schedule(contract, num_steps, cached_timesteps)

rungs = contract&#91;&quot;dmd_denoising_steps&quot;&#93;
forwards = len(rungs)
if num_steps not in (forwards, forwards + 1):
    raise ValueError(f&quot;Requested &#123;num_steps&#125; steps, but the FastH3 contract trains &#123;forwards&#125; transformer &quot;
                     f&quot;forwards (&#123;forwards + 1&#125; sigma-grid points).&quot;)
video = MiniMaxH3SchedulerState.from_dmd_steps(contract&#91;&quot;video_scheduler_shift&quot;&#93;, rungs)
audio = MiniMaxH3SchedulerState.from_dmd_steps(contract&#91;&quot;audio_scheduler_shift&quot;&#93;, 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, &quot;contract-dmd&quot;)</code></pre></details>
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&#91;:-1&#93;, fp32

@classmethod
def create(cls, shift: float, num_denoise_steps: int) -&gt; MiniMaxH3SchedulerState:
    if shift &lt;= 0:
        raise ValueError(f&quot;`shift` must be positive, got &#123;shift&#125;.&quot;)
    sigmas = minimax_h3_sigmas(shift, num_denoise_steps)
    return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas&#91;:-1&#93;).astype(np.float32))

@classmethod
def from_dmd_steps(cls, shift: float, dmd_steps: list&#91;int&#93;) -&gt; MiniMaxH3SchedulerState:
    &quot;&quot;&quot;Shift explicit DMD rungs the way the CUDA denoiser does.&quot;&quot;&quot;
    shift = _positive_shift(shift, name=&quot;shift&quot;)
    rungs = validate_fasth3_dmd_rungs(dmd_steps)
    base = np.array(&#91;step / 1000.0 for step in rungs&#93; + &#91;0.0&#93;, 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 &lt; 2 or not bool(np.all(sigmas64&#91;1:&#93; &lt; sigmas64&#91;:-1&#93;)) or float(sigmas64&#91;-1&#93;) != 0.0:
        raise ValueError(&quot;DMD sigmas must be strictly decreasing and end at 0.&quot;)
    return cls(shift=shift, sigmas=sigmas64.astype(np.float32), timesteps=(1.0 - sigmas64&#91;:-1&#93;).astype(np.float32))

@classmethod
def from_pdd_node_sigmas(cls, shift: float, node_sigmas: np.ndarray) -&gt; MiniMaxH3SchedulerState:
    &quot;&quot;&quot;Carry a PDD plan&#x27;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&#x27;s total integration weight.
    &quot;&quot;&quot;
    shift = _positive_shift(shift, name=&quot;shift&quot;)
    sigmas = np.asarray(node_sigmas, dtype=np.float32)
    if sigmas.size &lt; 2 or float(sigmas&#91;-1&#93;) != 0.0 or not bool(np.all(sigmas&#91;1:&#93; &lt; sigmas&#91;:-1&#93;)):
        raise ValueError(&quot;PDD node sigmas must be strictly decreasing and end at 0.&quot;)
    return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas&#91;:-1&#93;).astype(np.float32))

@property
def num_steps(self) -&gt; int:
    return int(self.timesteps.shape&#91;0&#93;)

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

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

    sigma_from_timestep = float(1.0 - self.timesteps&#91;step_index&#93;)
    denoised = sample + sigma_from_timestep * model_output
    sigma = float(self.sigmas&#91;step_index&#93;)
    sigma_next = float(self.sigmas&#91;step_index + 1&#93;)
    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):
    &quot;&quot;&quot;Conditioning noise-aug: t*sample + (1-t)*noise (t=0.999 for keyframes).&quot;&quot;&quot;
    return timestep * sample + (1.0 - timestep) * noise</code></pre></details>
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&#x27;s total integration weight.
&quot;&quot;&quot;
shift = _positive_shift(shift, name=&quot;shift&quot;)
sigmas = np.asarray(node_sigmas, dtype=np.float32)
if sigmas.size &lt; 2 or float(sigmas&#91;-1&#93;) != 0.0 or not bool(np.all(sigmas&#91;1:&#93; &lt; sigmas&#91;:-1&#93;)):
    raise ValueError(&quot;PDD node sigmas must be strictly decreasing and end at 0.&quot;)
return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas&#91;:-1&#93;).astype(np.float32))</code></pre></details>
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.
&quot;&quot;&quot;
import mlx.core as mx

sigma_from_timestep = float(1.0 - self.timesteps&#91;step_index&#93;)
denoised = sample + sigma_from_timestep * model_output
sigma = float(self.sigmas&#91;step_index&#93;)
sigma_next = float(self.sigmas&#91;step_index + 1&#93;)
ratio = sigma_next / sigma
prev = ratio * sample.astype(mx.float32) + (1.0 - ratio) * denoised.astype(mx.float32)
return prev.astype(sample.dtype)</code></pre></details>
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).
&quot;&quot;&quot;
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&#91;:, :, None&#93; * inv_freq&#91;None, None, :&#93;  # (S, 3, F)
freqs_t, freqs_h, freqs_w = freqs&#91;:, 0&#93;, freqs&#91;:, 1&#93;, freqs&#91;:, 2&#93;
freqs = mx.concatenate(&#91;freqs_t, freqs_h, freqs_w&#93;, axis=-1)
freqs = mx.concatenate(&#91;freqs, freqs&#93;, axis=-1)
return mx.cos(freqs), mx.sin(freqs)</code></pre></details>
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 &lt;= 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.
&quot;&quot;&quot;
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&#91;-1&#93;
out = _compiled_rotate_half(x&#91;..., :rotary_dim&#93;, cos, sin)
return mx.concatenate(&#91;out, x&#91;..., rotary_dim:&#93;&#93;, axis=-1)</code></pre></details>
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&#x27;s video/audio/condition timesteps). Blocks then index
tables directly and the per-block ``adaln_proj`` weights can be freed —
~40% of H3&#x27;s parameters never need to be resident during denoise.
&quot;&quot;&quot;

timesteps: np.ndarray  # (n,) fp32, sorted
block_tables: list&#91;tuple&#91;Any, ...&#93;&#93;  # 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) -&gt; np.ndarray:
    &quot;&quot;&quot;Map a step&#x27;s unique timesteps to rows in the cached union.&quot;&quot;&quot;
    positions = np.searchsorted(self.timesteps, step_timesteps)
    if (positions.size and (np.any(positions &gt;= len(self.timesteps))
                            or not np.allclose(self.timesteps&#91;positions&#93;, step_timesteps, atol=1e-6))):
        raise ValueError(f&quot;Step timesteps &#123;step_timesteps&#125; are not in the cached schedule union.&quot;)
    return positions.astype(np.int64)</code></pre></details>
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&#x27;s key names. ``blocks`` holds
the main transformer blocks, ``refiner`` the text refiner blocks; each is
a dict of arrays / :class:`QuantizedMatrix`.
&quot;&quot;&quot;

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

def _resolve_pdd_grid_size(self) -&gt; int | None:
    &quot;&quot;&quot;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&#x27;s
    ``config&#91;&quot;pdd_steps&quot;&#93;`` when present; otherwise it is inferred from the
    projection&#x27;s leading dimension and cross-checked against the audio head.
    A widened head that would silently mis-project raises here.
    &quot;&quot;&quot;
    from fastvideo.mlx_runtime.minimax_h3_pdd import resolved_grid_size

    declared = self.config.get(&quot;pdd_steps&quot;)
    if declared is not None and (isinstance(declared, bool) or not isinstance(declared, int) or declared &lt; 2):
        raise ValueError(f&quot;MiniMax H3 pdd_steps must be an int &gt;= 2, got &#123;declared!r&#125;.&quot;)
    video_weight = self.weights.get(&quot;proj_out.weight&quot;)
    audio_weight = self.weights.get(&quot;audio_proj_out.weight&quot;)
    video_grid = None if video_weight is None else resolved_grid_size(
        video_weight, self.weights.get(&quot;proj_out.bias&quot;), self.patch_dim)
    audio_grid = None if audio_weight is None else resolved_grid_size(
        audio_weight, self.weights.get(&quot;audio_proj_out.bias&quot;), self.audio_in_channels)
    if video_grid is not None and audio_grid is not None and video_grid != audio_grid:
        raise ValueError(f&quot;MiniMax H3 PDD heads disagree between proj_out (grid &#123;video_grid&#125;) and &quot;
                         f&quot;audio_proj_out (grid &#123;audio_grid&#125;).&quot;)
    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&quot;MiniMax H3 checkpoint config pdd_steps=&#123;declared&#125; disagrees with its widened heads &quot;
                         f&quot;(grid &#123;grid&#125;).&quot;)
    return grid if grid is not None else declared

def _fused_output_params(
    self,
    key: str,
    modality: str,
    pdd_block: tuple&#91;int, int&#93; | None,
    pdd_integration_weights: dict&#91;str, np.ndarray&#93; | None,
):
    &quot;&quot;&quot;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.
    &quot;&quot;&quot;
    weight = self.weights&#91;key&#93;
    bias = self.weights.get(f&quot;&#123;key.rsplit(&#x27;.&#x27;, 1)&#91;0&#93;&#125;.bias&quot;)
    if pdd_block is None:
        if self.pdd_grid_size is not None:
            raise ValueError(f&quot;MiniMax H3 checkpoint has a PDD-widened `&#123;key&#125;` but no PDD block was supplied; &quot;
                             &quot;pass pdd_block= and pdd_integration_weights= to the forward.&quot;)
        return weight, bias
    if self.pdd_grid_size is None:
        raise ValueError(f&quot;A PDD block was supplied but `&#123;key&#125;` is not widened (single-headed checkpoint).&quot;)
    if pdd_integration_weights is None or modality not in pdd_integration_weights:
        raise ValueError(f&quot;A PDD block needs this modality&#x27;s integration weights (&#123;modality!r&#125;).&quot;)
    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&#91;0&#93;,
        end=pdd_block&#91;1&#93;,
        weights=pdd_integration_weights&#91;modality&#93;,
    )

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

def compute_temb(self, timesteps):
    &quot;&quot;&quot;(n,) timesteps -&gt; (n, time_embed_dim) through time_proj + MLP.&quot;&quot;&quot;
    t_freq = timestep_embedding(timesteps,
                                self.freq_dim).astype(weight_dtype(self.weights&#91;&quot;time_embedder.linear_1.weight&quot;&#93;))
    temb = linear(
        t_freq,
        self.weights&#91;&quot;time_embedder.linear_1.weight&quot;&#93;,
        self.weights&#91;&quot;time_embedder.linear_1.bias&quot;&#93;,
    )
    temb = linear(
        silu(temb),
        self.weights&#91;&quot;time_embedder.linear_2.weight&quot;&#93;,
        self.weights&#91;&quot;time_embedder.linear_2.bias&quot;&#93;,
    )
    if self.adaln_rank is not None:
        temb = linear(
            silu(temb).astype(weight_dtype(self.weights&#91;&quot;adaln_basis.weight&quot;&#93;)), self.weights&#91;&quot;adaln_basis.weight&quot;&#93;)
    return temb

def refine_text(self, text_rows):
    hidden = linear(
        text_rows.astype(weight_dtype(self.weights&#91;&quot;context_embedder.weight&quot;&#93;)),
        self.weights&#91;&quot;context_embedder.weight&quot;&#93;,
        self.weights&#91;&quot;context_embedder.bias&quot;&#93;,
    )
    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&#91;&quot;token_refiner.final_norm.weight&quot;&#93;, self.final_norm_eps)

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

def precompute_adaln(self, timesteps: np.ndarray, *, drop_weights: bool = True) -&gt; MiniMaxH3StepCache:
    &quot;&quot;&quot;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.
    &quot;&quot;&quot;
    import mlx.core as mx

    timesteps = np.unique(np.asarray(timesteps, dtype=np.float32))
    temb = self.compute_temb(mx.array(timesteps))
    block_tables = &#91;_adaln_tables(block, temb, apply_silu=self.adaln_rank is None) for block in self.blocks&#93;
    shift_scale = linear(
        (silu(temb) if self.adaln_rank is None else temb).astype(
            weight_dtype(self.weights&#91;&quot;norm_out.linear.weight&quot;&#93;)),
        self.weights&#91;&quot;norm_out.linear.weight&quot;&#93;,
        self.weights&#91;&quot;norm_out.linear.bias&quot;&#93;,
    )
    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&#91;&quot;adaln_proj.linear.weight&quot;&#93; = None
            block&#91;&quot;adaln_proj.linear.bias&quot;&#93; = None
    return self._adaln_cache

@property
def vsa_capable(self) -&gt; bool:
    return self._vsa_capable

def configure_vsa(self, config: MiniMaxH3VSAConfig | None) -&gt; None:
    config = config or MiniMaxH3VSAConfig()
    if config.enabled and not self._vsa_capable:
        raise dense_only_vsa_error(&quot;this MLX H3 DiT&quot;)
    if any(layer &gt;= len(self.blocks) for layer in config.dense_layers):
        raise ValueError(f&quot;VSA dense_layers must be below the model&#x27;s block count (&#123;len(self.blocks)&#125;).&quot;)
    self.vsa_config = config
    self._vsa_geometry = None
    self.reset_vsa_stats()

def reset_vsa_stats(self) -&gt; 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) -&gt; 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 &gt; 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&quot;H3 VSA cannot run with this packed layout (segments=&#123;segments&#125;, &quot;
                         f&quot;tile_size=&#123;config.tile_size&#125;): &#123;error&#125;. &quot;
                         &quot;Disable VSA or pick a supported tile size (64, 128 or 256).&quot;) from error
    self._vsa_geometry = geometry
    return self._vsa_geometry

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

def _vsa_block_kwargs(self, block_index: int, step_index: int) -&gt; dict&#91;str, Any&#93;:
    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 &#123;
        &quot;vsa_geometry&quot;: geometry,
        &quot;vsa_sparsity&quot;: sparsity,
        &quot;vsa_exempt&quot;: self.vsa_config.exempt,
        &quot;vsa_ref_keep_rate&quot;: self.vsa_config.effective_ref_keep_rate,
        &quot;vsa_impl&quot;: self.vsa_config.impl,
        &quot;vsa_stats&quot;: stats,
        &quot;use_gate_compress&quot;: geometry is not None and self._block_gate_active(block_index),
    &#125;

# -- 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&#91;int, int&#93; | None = None,
    pdd_integration_weights: dict&#91;str, np.ndarray&#93; | None = None,
):
    &quot;&quot;&quot;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&#x27;s ``pdd_block`` and ``pdd_integration_weights`` so the widened
    heads are fused.
    &quot;&quot;&quot;
    import mlx.core as mx

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

    video_embeds = linear(
        video_rows.astype(weight_dtype(self.weights&#91;&quot;proj_in.weight&quot;&#93;)),
        self.weights&#91;&quot;proj_in.weight&quot;&#93;,
        self.weights&#91;&quot;proj_in.bias&quot;&#93;,
    )
    audio_embeds = linear(
        audio_rows.astype(weight_dtype(self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;)),
        self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;,
        self.weights&#91;&quot;audio_proj_in.bias&quot;&#93;,
    )
    text_embeds = self.refine_text(text_rows)

    packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
    packed&#91;_as_mx_indices(text_indices)&#93; = text_embeds
    packed&#91;_as_mx_indices(video_indices)&#93; = video_embeds.astype(text_embeds.dtype)
    packed&#91;_as_mx_indices(audio_indices)&#93; = 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&#91;&quot;norm_out.linear.weight&quot;&#93;)),
        self.weights&#91;&quot;norm_out.linear.weight&quot;&#93;,
        self.weights&#91;&quot;norm_out.linear.bias&quot;&#93;,
    )
    shift_rows, scale_rows = mx.split(shift_scale, 2, axis=-1)
    normed = _h3_rms_norm(packed, self.weights&#91;&quot;norm_out.norm.weight&quot;&#93;, self.final_norm_eps)
    normed = normed * (1.0 + scale_rows&#91;timestep_indices&#93;) + shift_rows&#91;timestep_indices&#93;
    video_weight, video_bias = self._fused_output_params(&quot;proj_out.weight&quot;, &quot;video&quot;, pdd_block,
                                                         pdd_integration_weights)
    audio_weight, audio_bias = self._fused_output_params(&quot;audio_proj_out.weight&quot;, &quot;audio&quot;, pdd_block,
                                                         pdd_integration_weights)
    video_output = linear(
        normed&#91;_as_mx_indices(video_indices)&#93;.astype(weight_dtype(video_weight)),
        video_weight,
        video_bias,
    )
    audio_output = linear(
        normed&#91;_as_mx_indices(audio_indices)&#93;.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&#91;int, int&#93; | None = None,
    pdd_integration_weights: dict&#91;str, np.ndarray&#93; | None = None,
    record_prefix_kv: Any | None = None,
):
    &quot;&quot;&quot;Denoise-step forward served entirely from the AdaLN cache.

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

    ``step_timesteps`` are this step&#x27;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&#x27;s
    ``pdd_block`` and ``pdd_integration_weights`` so the widened output
    heads are fused into the block&#x27;s weighted mean.
    &quot;&quot;&quot;
    import mlx.core as mx

    if self._adaln_cache is None:
        raise RuntimeError(&quot;precompute_adaln() must run before forward_with_cache().&quot;)
    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&#91;&quot;proj_in.weight&quot;&#93;)),
        self.weights&#91;&quot;proj_in.weight&quot;&#93;,
        self.weights&#91;&quot;proj_in.bias&quot;&#93;,
    )
    audio_embeds = linear(
        audio_rows.astype(weight_dtype(self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;)),
        self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;,
        self.weights&#91;&quot;audio_proj_in.bias&quot;&#93;,
    )
    text_embeds = self.refine_text(text_rows)

    packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
    packed&#91;mx.array(layout.text_indices)&#93; = text_embeds
    packed&#91;mx.array(layout.video_indices)&#93; = video_embeds.astype(text_embeds.dtype)
    packed&#91;mx.array(layout.audio_indices)&#93; = audio_embeds.astype(text_embeds.dtype)

    token_tags = mx.array(layout.token_tags)
    row_inverse = mx.array(row_timestep_inverse)
    adaln_indices = (positions&#91;row_inverse&#93; * 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&#91;Any&#93; = &#91;&#93;
    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&quot;Prefix K/V cache is sized for &#123;record_prefix_kv.num_layers&#125; layers x &quot;
                             f&quot;&#123;record_prefix_kv.rows&#125; rows, the step has &#123;len(self.blocks)&#125; x &#123;frozen_rows&#125;.&quot;)
        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&#91;:frozen_rows&#93;, v&#91;:frozen_rows&#93;))

            vsa_kwargs&#91;&quot;kv_sink&quot;&#93; = 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&#91;&quot;vsa_stats&quot;&#93;)

    row_positions = positions&#91;row_inverse&#93;
    normed = _h3_rms_norm(packed, self.weights&#91;&quot;norm_out.norm.weight&quot;&#93;, self.final_norm_eps)
    normed = normed * (1.0 + cache.norm_out_scale&#91;row_positions&#93;) + cache.norm_out_shift&#91;row_positions&#93;
    video_weight, video_bias = self._fused_output_params(&quot;proj_out.weight&quot;, &quot;video&quot;, pdd_block,
                                                         pdd_integration_weights)
    audio_weight, audio_bias = self._fused_output_params(&quot;audio_proj_out.weight&quot;, &quot;audio&quot;, pdd_block,
                                                         pdd_integration_weights)
    video_output = linear(
        normed&#91;mx.array(layout.video_indices)&#93;.astype(weight_dtype(video_weight)),
        video_weight,
        video_bias,
    )
    audio_output = linear(
        normed&#91;mx.array(layout.audio_indices)&#93;.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):
    &quot;&quot;&quot;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.
    &quot;&quot;&quot;
    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</code></pre></details>
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.
&quot;&quot;&quot;
import mlx.core as mx

timesteps = np.unique(np.asarray(timesteps, dtype=np.float32))
temb = self.compute_temb(mx.array(timesteps))
block_tables = &#91;_adaln_tables(block, temb, apply_silu=self.adaln_rank is None) for block in self.blocks&#93;
shift_scale = linear(
    (silu(temb) if self.adaln_rank is None else temb).astype(
        weight_dtype(self.weights&#91;&quot;norm_out.linear.weight&quot;&#93;)),
    self.weights&#91;&quot;norm_out.linear.weight&quot;&#93;,
    self.weights&#91;&quot;norm_out.linear.bias&quot;&#93;,
)
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&#91;&quot;adaln_proj.linear.weight&quot;&#93; = None
        block&#91;&quot;adaln_proj.linear.bias&quot;&#93; = None
return self._adaln_cache</code></pre></details>
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&#x27;s ``pdd_block`` and ``pdd_integration_weights`` so the widened
heads are fused.
&quot;&quot;&quot;
import mlx.core as mx

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

video_embeds = linear(
    video_rows.astype(weight_dtype(self.weights&#91;&quot;proj_in.weight&quot;&#93;)),
    self.weights&#91;&quot;proj_in.weight&quot;&#93;,
    self.weights&#91;&quot;proj_in.bias&quot;&#93;,
)
audio_embeds = linear(
    audio_rows.astype(weight_dtype(self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;)),
    self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;,
    self.weights&#91;&quot;audio_proj_in.bias&quot;&#93;,
)
text_embeds = self.refine_text(text_rows)

packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
packed&#91;_as_mx_indices(text_indices)&#93; = text_embeds
packed&#91;_as_mx_indices(video_indices)&#93; = video_embeds.astype(text_embeds.dtype)
packed&#91;_as_mx_indices(audio_indices)&#93; = 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&#91;&quot;norm_out.linear.weight&quot;&#93;)),
    self.weights&#91;&quot;norm_out.linear.weight&quot;&#93;,
    self.weights&#91;&quot;norm_out.linear.bias&quot;&#93;,
)
shift_rows, scale_rows = mx.split(shift_scale, 2, axis=-1)
normed = _h3_rms_norm(packed, self.weights&#91;&quot;norm_out.norm.weight&quot;&#93;, self.final_norm_eps)
normed = normed * (1.0 + scale_rows&#91;timestep_indices&#93;) + shift_rows&#91;timestep_indices&#93;
video_weight, video_bias = self._fused_output_params(&quot;proj_out.weight&quot;, &quot;video&quot;, pdd_block,
                                                     pdd_integration_weights)
audio_weight, audio_bias = self._fused_output_params(&quot;audio_proj_out.weight&quot;, &quot;audio&quot;, pdd_block,
                                                     pdd_integration_weights)
video_output = linear(
    normed&#91;_as_mx_indices(video_indices)&#93;.astype(weight_dtype(video_weight)),
    video_weight,
    video_bias,
)
audio_output = linear(
    normed&#91;_as_mx_indices(audio_indices)&#93;.astype(weight_dtype(audio_weight)),
    audio_weight,
    audio_bias,
)
return video_output, audio_output</code></pre></details>
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&#x27;s K/V of the frozen text/reference rows for later
``forward_target_rows`` steps; the outputs are unchanged.

``step_timesteps`` are this step&#x27;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&#x27;s
``pdd_block`` and ``pdd_integration_weights`` so the widened output
heads are fused into the block&#x27;s weighted mean.
&quot;&quot;&quot;
import mlx.core as mx

if self._adaln_cache is None:
    raise RuntimeError(&quot;precompute_adaln() must run before forward_with_cache().&quot;)
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&#91;&quot;proj_in.weight&quot;&#93;)),
    self.weights&#91;&quot;proj_in.weight&quot;&#93;,
    self.weights&#91;&quot;proj_in.bias&quot;&#93;,
)
audio_embeds = linear(
    audio_rows.astype(weight_dtype(self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;)),
    self.weights&#91;&quot;audio_proj_in.weight&quot;&#93;,
    self.weights&#91;&quot;audio_proj_in.bias&quot;&#93;,
)
text_embeds = self.refine_text(text_rows)

packed = mx.zeros((sequence_length, self.hidden_size), dtype=text_embeds.dtype)
packed&#91;mx.array(layout.text_indices)&#93; = text_embeds
packed&#91;mx.array(layout.video_indices)&#93; = video_embeds.astype(text_embeds.dtype)
packed&#91;mx.array(layout.audio_indices)&#93; = audio_embeds.astype(text_embeds.dtype)

token_tags = mx.array(layout.token_tags)
row_inverse = mx.array(row_timestep_inverse)
adaln_indices = (positions&#91;row_inverse&#93; * 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&#91;Any&#93; = &#91;&#93;
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&quot;Prefix K/V cache is sized for &#123;record_prefix_kv.num_layers&#125; layers x &quot;
                         f&quot;&#123;record_prefix_kv.rows&#125; rows, the step has &#123;len(self.blocks)&#125; x &#123;frozen_rows&#125;.&quot;)
    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&#91;:frozen_rows&#93;, v&#91;:frozen_rows&#93;))

        vsa_kwargs&#91;&quot;kv_sink&quot;&#93; = 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&#91;&quot;vsa_stats&quot;&#93;)

row_positions = positions&#91;row_inverse&#93;
normed = _h3_rms_norm(packed, self.weights&#91;&quot;norm_out.norm.weight&quot;&#93;, self.final_norm_eps)
normed = normed * (1.0 + cache.norm_out_scale&#91;row_positions&#93;) + cache.norm_out_shift&#91;row_positions&#93;
video_weight, video_bias = self._fused_output_params(&quot;proj_out.weight&quot;, &quot;video&quot;, pdd_block,
                                                     pdd_integration_weights)
audio_weight, audio_bias = self._fused_output_params(&quot;audio_proj_out.weight&quot;, &quot;audio&quot;, pdd_block,
                                                     pdd_integration_weights)
video_output = linear(
    normed&#91;mx.array(layout.video_indices)&#93;.astype(weight_dtype(video_weight)),
    video_weight,
    video_bias,
)
audio_output = linear(
    normed&#91;mx.array(layout.audio_indices)&#93;.astype(weight_dtype(audio_weight)),
    audio_weight,
    audio_bias,
)
return video_output, audio_output</code></pre></details>
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.
&quot;&quot;&quot;
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)</code></pre></details>
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.
&quot;&quot;&quot;
import mlx.core as mx
transformer_path = Path(transformer_path)
if config is None:
    config_path = transformer_path / &quot;config.json&quot; if transformer_path.is_dir() else transformer_path.with_name(
        &quot;config.json&quot;)
    config = json.loads(Path(config_path).read_text())
total_blocks = int(config&#91;&quot;num_layers&quot;&#93;)
if num_blocks is None:
    num_blocks = total_blocks
num_refiner = int(config&#91;&quot;num_refiner_layers&quot;&#93;)
cast_dtype = &#123;&quot;fp16&quot;: mx.float16, &quot;bf16&quot;: mx.bfloat16, &quot;fp32&quot;: mx.float32&#125;&#91;dtype&#93;
spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None
                                                       or isinstance(quantization, str)) else quantization
ensure_quantization_supported(spec)

weights: dict&#91;str, Any&#93; = &#123;&#125;
blocks: list&#91;dict&#91;str, Any&#93; | None&#93; = &#91;None&#93; * num_blocks
refiner: list&#91;dict&#91;str, Any&#93; | None&#93; = &#91;None&#93; * num_refiner
cache_timesteps = None
temb = None
cached_block_tables: list&#91;tuple&#91;Any, ...&#93; | None&#93; | None = None
pending_adaln: dict&#91;int, dict&#91;str, Any&#93;&#93; = &#123;&#125;

def assign(key: str, value) -&gt; None:
    if key.startswith(&quot;transformer_blocks.&quot;):
        _, index_str, sub = key.split(&quot;.&quot;, 2)
        index = int(index_str)
        if index &lt; num_blocks:
            block = blocks&#91;index&#93;
            if block is None:
                block = &#123;&#125;
                blocks&#91;index&#93; = block
            block&#91;sub&#93; = value
    elif key.startswith(&quot;token_refiner.refiner_blocks.&quot;):
        _, _, index_str, sub = key.split(&quot;.&quot;, 3)
        index = int(index_str)
        refiner_block = refiner&#91;index&#93;
        if refiner_block is None:
            refiner_block = &#123;&#125;
            refiner&#91;index&#93; = refiner_block
        refiner_block&#91;sub&#93; = value
    else:
        weights&#91;key&#93; = 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(&quot;time_embedder.&quot;) or key == &quot;adaln_basis.weight&quot;):
                continue
            keep_fp32 = key.split(&quot;.&quot;, 1)&#91;0&#93; in FP32_MODULE_PREFIXES
            target_dtype = mx.float32 if keep_fp32 else (mx.float16 if key == &quot;adaln_basis.weight&quot; else cast_dtype)
            assign(key, _load_array(source, target_dtype))
        del shard_arrays
    required_time_keys = &#123;
        &quot;time_embedder.linear_1.weight&quot;,
        &quot;time_embedder.linear_1.bias&quot;,
        &quot;time_embedder.linear_2.weight&quot;,
        &quot;time_embedder.linear_2.bias&quot;,
    &#125;
    missing_time_keys = sorted(required_time_keys - weights.keys())
    if missing_time_keys:
        raise KeyError(f&quot;Missing time-embedder weights needed for the AdaLN cache: &#123;missing_time_keys&#125;&quot;)
    timestep_rows = mx.array(cache_timesteps)
    t_freq = timestep_embedding(timestep_rows, int(config&#91;&quot;freq_dim&quot;&#93;)).astype(
        weight_dtype(weights&#91;&quot;time_embedder.linear_1.weight&quot;&#93;))
    temb = linear(t_freq, weights&#91;&quot;time_embedder.linear_1.weight&quot;&#93;, weights&#91;&quot;time_embedder.linear_1.bias&quot;&#93;)
    temb = linear(silu(temb), weights&#91;&quot;time_embedder.linear_2.weight&quot;&#93;, weights&#91;&quot;time_embedder.linear_2.bias&quot;&#93;)
    if config.get(&quot;adaln_rank&quot;) is not None:
        if &quot;adaln_basis.weight&quot; not in weights:
            raise KeyError(&quot;Rank-reduced AdaLN checkpoint is missing adaln_basis.weight&quot;)
        temb = linear(silu(temb).astype(weight_dtype(weights&#91;&quot;adaln_basis.weight&quot;&#93;)), weights&#91;&quot;adaln_basis.weight&quot;&#93;)
    mx.eval(temb)
    cached_block_tables = &#91;None&#93; * 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(&quot;.weight&quot;):
            raise ValueError(f&quot;Packed transformer weight &#123;key&#125; is not a floating-point source. &quot;
                             &quot;Convert the released BF16 transformer to MLX INT6 instead.&quot;)
        if temb is not None and (key.startswith(&quot;time_embedder.&quot;) or key == &quot;adaln_basis.weight&quot;):
            continue
        if key.startswith(&quot;transformer_blocks.&quot;):
            index = int(key.split(&quot;.&quot;)&#91;1&#93;)
            if index &gt;= num_blocks:
                continue
        if key.startswith(&quot;rope.&quot;):
            continue  # non-persistent analytic buffer, rebuilt on the fly
        keep_fp32 = key.split(&quot;.&quot;, 1)&#91;0&#93; in FP32_MODULE_PREFIXES
        factorized_adaln = config.get(&quot;adaln_rank&quot;) is not None and (&quot;.adaln_proj.&quot; in key or key.startswith(
            (&quot;norm_out.linear.&quot;, &quot;adaln_basis.&quot;)))
        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 &quot;.adaln_proj.linear.&quot; in key:
            _, index_str, sub = key.split(&quot;.&quot;, 2)
            index = int(index_str)
            block_pending = pending_adaln.setdefault(index, &#123;&#125;)
            block_pending&#91;sub&#93; = array
            if &#123;&quot;adaln_proj.linear.weight&quot;, &quot;adaln_proj.linear.bias&quot;&#125; &lt;= block_pending.keys():
                tables = _adaln_tables(block_pending, temb, apply_silu=config.get(&quot;adaln_rank&quot;) is None)
                mx.eval(tables)
                assert cached_block_tables is not None
                cached_block_tables&#91;index&#93; = tables
                block = blocks&#91;index&#93;
                if block is None:
                    block = &#123;&#125;
                    blocks&#91;index&#93; = block
                block&#91;&quot;adaln_proj.linear.weight&quot;&#93; = None
                block&#91;&quot;adaln_proj.linear.bias&quot;&#93; = None
                del pending_adaln&#91;index&#93;
            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&quot;Incomplete AdaLN projection pairs for blocks &#123;sorted(pending_adaln)&#125;&quot;)
    assert cache_timesteps is not None and cached_block_tables is not None
    missing_cache_blocks = &#91;index for index, tables in enumerate(cached_block_tables) if tables is None&#93;
    if missing_cache_blocks:
        raise KeyError(f&quot;Missing AdaLN cache tables for blocks &#123;missing_cache_blocks&#125;&quot;)
    shift_scale = linear(
        (silu(temb) if config.get(&quot;adaln_rank&quot;) is None else temb).astype(
            weight_dtype(weights&#91;&quot;norm_out.linear.weight&quot;&#93;)),
        weights&#91;&quot;norm_out.linear.weight&quot;&#93;,
        weights&#91;&quot;norm_out.linear.bias&quot;&#93;,
    )
    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=&#91;tables for tables in cached_block_tables if tables is not None&#93;,
        norm_out_shift=norm_out_shift,
        norm_out_scale=norm_out_scale,
    )

if any(block is None for block in blocks):
    missing = &#91;i for i, block in enumerate(blocks) if block is None&#93;
    raise KeyError(f&quot;Missing transformer block weights for indices &#123;missing&#125;.&quot;)
if any(block is None for block in refiner):
    raise KeyError(&quot;Missing token refiner weights.&quot;)
loaded_blocks = &#91;block for block in blocks if block is not None&#93;
if include_vsa:
    missing_gates = &#91;
        index for index, block in enumerate(loaded_blocks) if &quot;attn.to_gate_compress.weight&quot; not in block
    &#93;
    if missing_gates:
        raise KeyError(&quot;VSA conversion requested but gate projections are missing for transformer blocks &quot;
                       f&quot;&#123;missing_gates&#125;. The source checkpoint must contain &quot;
                       f&quot;`transformer_blocks.*.&#123;VSA_GATE_KEY_SUFFIX&#125;`.&quot;)
    logger.info(&quot;Retained %d VSA gate projections (%s).&quot;, len(loaded_blocks), VSA_GATE_KEY_SUFFIX)
dit = MLXMiniMaxH3DiT(
    weights,
    loaded_blocks,
    &#91;block for block in refiner if block is not None&#93;,
    config,
)
dit._adaln_cache = cache
return dit</code></pre></details>
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&#91;str, Any&#93; = &#123;&#125;
quantized: dict&#91;str, dict&#91;str, Any&#93;&#93; = &#123;&#125;
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&quot;Mixed quantization specs in one checkpoint (&#123;spec&#125; vs &#123;value.spec&#125; at &#x27;&#123;key&#125;&#x27;).&quot;)
        spec = value.spec
        arrays&#91;key&#93; = value.weight
        arrays&#91;f&quot;&#123;key&#125;.scales&quot;&#93; = value.scales
        if value.biases is not None:
            arrays&#91;f&quot;&#123;key&#125;.biases&quot;&#93; = value.biases
        quantized&#91;key&#93; = &#123;
            &quot;dequantized_dtype&quot;: _dtype_name(value.dequantized_dtype),
            &quot;has_biases&quot;: value.biases is not None,
        &#125;
    else:
        arrays&#91;key&#93; = value

cache_manifest = None
if dit._adaln_cache is not None:
    cache = dit._adaln_cache
    cache_manifest = &#123;
        &quot;timesteps&quot;: cache.timesteps.tolist(),
        &quot;num_blocks&quot;: len(cache.block_tables),
        &quot;tables_per_block&quot;: 6,
    &#125;
    for block_index, tables in enumerate(cache.block_tables):
        for table_index, table in enumerate(tables):
            arrays&#91;f&quot;__adaln_cache.block.&#123;block_index&#125;.table.&#123;table_index&#125;&quot;&#93; = table
    arrays&#91;&quot;__adaln_cache.norm_out_shift&quot;&#93; = cache.norm_out_shift
    arrays&#91;&quot;__adaln_cache.norm_out_scale&quot;&#93; = cache.norm_out_scale

manifest = &#123;
    &quot;format_version&quot;: H3_FORMAT_VERSION,
    &quot;config&quot;: dit.config,
    &quot;num_blocks&quot;: len(dit.blocks),
    &quot;num_refiner_blocks&quot;: len(dit.refiner),
    &quot;quantization&quot;: None if spec is None else &#123;
        &quot;mode&quot;: spec.mode,
        &quot;bits&quot;: spec.bits,
        &quot;group_size&quot;: spec.group_size,
    &#125;,
    &quot;quantized_keys&quot;: quantized,
    &quot;adaln_cache&quot;: cache_manifest,
    &quot;vsa&quot;: &#123;
        &quot;capable&quot;:
        bool(dit.vsa_capable),
        &quot;num_gate_matrices&quot;:
        sum(1 for block in dit.blocks if &quot;attn.to_gate_compress.weight&quot; in block),
        &quot;gate_key_suffix&quot;:
        VSA_GATE_KEY_SUFFIX,
        &quot;attention_activations&quot;:
        &quot;bf16&quot;,
        &quot;note&quot;: (&quot;Gate projections are quantized with the checkpoint&#x27;s affine INT8/INT6/INT4 &quot;
                 &quot;weight-only grid. VSA attention Q/K/V remain unquantized activations.&quot;),
    &#125;,
&#125;
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(&quot;Saved MLX H3 DiT checkpoint (%d arrays, quantization=%s) to %s&quot;, len(arrays),
            spec.label if spec else &quot;none&quot;, checkpoint_dir)
return checkpoint_dir</code></pre></details>
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.
&quot;&quot;&quot;
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&quot;Not an MLX H3 checkpoint directory: &#123;checkpoint_dir&#125; &quot;
                            f&quot;(expected &#123;H3_MANIFEST_FILENAME&#125; and &#123;H3_WEIGHTS_FILENAME&#125;).&quot;)

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

spec = None
if manifest&#91;&quot;quantization&quot;&#93; is not None:
    spec = MLXQuantizationSpec(**manifest&#91;&quot;quantization&quot;&#93;)
    ensure_quantization_supported(spec)

arrays = mx.load(str(weights_path))
quantized_keys: dict&#91;str, dict&#91;str, Any&#93;&#93; = manifest&#91;&quot;quantized_keys&quot;&#93;

def rebuild(key: str):
    if key not in quantized_keys:
        return arrays&#91;key&#93;
    info = quantized_keys&#91;key&#93;
    assert spec is not None, f&quot;Quantized key &#x27;&#123;key&#125;&#x27; in a checkpoint without a quantization spec&quot;
    return QuantizedMatrix(
        weight=arrays&#91;key&#93;,
        scales=arrays&#91;f&quot;&#123;key&#125;.scales&quot;&#93;,
        biases=arrays&#91;f&quot;&#123;key&#125;.biases&quot;&#93; if info&#91;&quot;has_biases&quot;&#93; else None,
        spec=spec,
        dequantized_dtype=_name_to_dtype(info&#91;&quot;dequantized_dtype&quot;&#93;),
    )

weights: dict&#91;str, Any&#93; = &#123;&#125;
blocks: list&#91;dict&#91;str, Any&#93; | None&#93; = &#91;None&#93; * int(manifest&#91;&quot;num_blocks&quot;&#93;)
refiner: list&#91;dict&#91;str, Any&#93; | None&#93; = &#91;None&#93; * int(manifest&#91;&quot;num_refiner_blocks&quot;&#93;)
for key in arrays:
    if key.startswith(&quot;__adaln_cache.&quot;):
        continue
    if key.endswith(&quot;.scales&quot;) or key.endswith(&quot;.biases&quot;):
        continue
    value = rebuild(key)
    if key.startswith(&quot;blocks.&quot;):
        _, index_str, sub = key.split(&quot;.&quot;, 2)
        index = int(index_str)
        block = blocks&#91;index&#93;
        if block is None:
            block = &#123;&#125;
            blocks&#91;index&#93; = block
        block&#91;sub&#93; = value
    elif key.startswith(&quot;refiner.&quot;):
        _, index_str, sub = key.split(&quot;.&quot;, 2)
        index = int(index_str)
        refiner_block = refiner&#91;index&#93;
        if refiner_block is None:
            refiner_block = &#123;&#125;
            refiner&#91;index&#93; = refiner_block
        refiner_block&#91;sub&#93; = value
    else:
        weights&#91;key&#93; = value

if any(block is None for block in blocks) or any(block is None for block in refiner):
    raise ValueError(f&quot;MLX H3 checkpoint &#123;checkpoint_dir&#125; is missing block weights.&quot;)
dit = MLXMiniMaxH3DiT(
    weights,
    &#91;block for block in blocks if block is not None&#93;,
    &#91;block for block in refiner if block is not None&#93;,
    manifest&#91;&quot;config&quot;&#93;,
)
cache_info = manifest.get(&quot;adaln_cache&quot;)
if cache_info is not None:
    block_tables = &#91;&#93;
    for block_index in range(int(cache_info&#91;&quot;num_blocks&quot;&#93;)):
        block_tables.append(
            tuple(arrays&#91;f&quot;__adaln_cache.block.&#123;block_index&#125;.table.&#123;table_index&#125;&quot;&#93;
                  for table_index in range(int(cache_info&#91;&quot;tables_per_block&quot;&#93;))))
    dit._adaln_cache = MiniMaxH3StepCache(
        timesteps=np.asarray(cache_info&#91;&quot;timesteps&quot;&#93;, dtype=np.float32),
        block_tables=block_tables,
        norm_out_shift=arrays&#91;&quot;__adaln_cache.norm_out_shift&quot;&#93;,
        norm_out_scale=arrays&#91;&quot;__adaln_cache.norm_out_scale&quot;&#93;,
    )
return dit</code></pre></details>