--- title: minimax_h3_vsa ---

MiniMax H3 Video Sparse Attention for the native MLX runtime.

Ports the packed-sequence VSA-H3 contract from fastvideo/attention/backends/video_sparse_attn_h3.py:

Execution backends:

Dense fused SDPA remains the default when VSA is disabled, the geometry is unsupported, or a dense-only checkpoint is loaded.

fastvideo.mlx_runtime.minimax_h3_vsa.init_logger

init_logger(name: str) -> _FastvideoLogger

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

fastvideo.mlx_runtime.minimax_h3_vsa.validate_ref_policy

validate_ref_policy(ref_policy: str | None, ref_keep_rate: float | None) -> None

p2_multi_region needs a keep rate in (0, 1]; no policy takes no keep rate.

Source
def validate_ref_policy(ref_policy: str | None, ref_keep_rate: float | None) -> None:
    """``p2_multi_region`` needs a keep rate in (0, 1]; no policy takes no keep rate."""
    if ref_policy is None:
        if ref_keep_rate is not None:
            raise ValueError("VSA ref_keep_rate needs ref_policy='p2_multi_region'.")
        return
    if ref_policy not in VSA_REF_POLICIES:
        raise ValueError(f"VSA ref_policy must be one of {VSA_REF_POLICIES} or None, got {ref_policy!r}.")
    if (isinstance(ref_keep_rate, bool) or not isinstance(ref_keep_rate, int | float)
            or not math.isfinite(ref_keep_rate) or not 0.0 < float(ref_keep_rate) <= 1.0):
        raise ValueError(f"VSA ref_keep_rate must be in (0, 1] for {ref_policy}, got {ref_keep_rate!r}.")

fastvideo.mlx_runtime.minimax_h3_vsa.DenseOnlyVSACheckpointError

DenseOnlyVSACheckpointError()

VSA was requested for a checkpoint that dropped the gate weights.

Source
class DenseOnlyVSACheckpointError(ValueError):
    """VSA was requested for a checkpoint that dropped the gate weights."""

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAConfig

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

Runtime VSA knobs. Defaults preserve dense MLX H3 behavior.

Source
@dataclass(frozen=True)
class MiniMaxH3VSAConfig:
    """Runtime VSA knobs. Defaults preserve dense MLX H3 behavior."""

    enabled: bool = False
    sparsity: float = 0.9
    tile_size: int = 64
    prefix_mode: PrefixMode = "exempt"
    dense_first_n_steps: int = 0
    dense_layers: tuple[int, ...] = ()
    impl: VSAImpl = "auto"
    # ``p2_multi_region`` makes each reference video its own sparse region that
    # keeps ``ref_keep_rate`` of its tiles; None keeps references dense prefix.
    ref_policy: str | None = None
    ref_keep_rate: float | None = None

    def __post_init__(self) -> None:
        validate_ref_policy(self.ref_policy, self.ref_keep_rate)
        if not 0.0 <= self.sparsity < 1.0:
            raise ValueError(f"VSA sparsity must be in [0, 1), got {self.sparsity}.")
        if self.tile_size not in VSA_H3_TILE_SHAPES:
            raise ValueError(f"VSA tile_size must be one of {sorted(VSA_H3_TILE_SHAPES)}, got {self.tile_size}.")
        if self.prefix_mode not in VSA_PREFIX_MODES:
            raise ValueError(f"VSA prefix_mode must be one of {VSA_PREFIX_MODES}, got {self.prefix_mode!r}.")
        if self.impl not in VSA_IMPLS:
            raise ValueError(f"VSA impl must be one of {VSA_IMPLS}, got {self.impl!r}.")
        if self.dense_first_n_steps < 0:
            raise ValueError(f"vsa_dense_first_n_steps must be >= 0, got {self.dense_first_n_steps}.")
        if isinstance(self.dense_layers, str) or any(
                not isinstance(layer, Integral) or isinstance(layer, bool) or layer < 0 for layer in self.dense_layers):
            raise ValueError("VSA dense_layers must be a sequence of non-negative integer indices.")
        object.__setattr__(self, "dense_layers", tuple(int(layer) for layer in self.dense_layers))

    @property
    def exempt(self) -> bool:
        return self.prefix_mode == "exempt"

    @property
    def effective_ref_keep_rate(self) -> float | None:
        """The per-reference-video keep rate, or None when references stay dense."""
        return self.ref_keep_rate if self.ref_policy == "p2_multi_region" else None

    def layer_sparsity(self, layer_idx: int, step_index: int) -> float:
        if not self.enabled:
            return 0.0
        if step_index < self.dense_first_n_steps:
            return 0.0
        if layer_idx in self.dense_layers:
            return 0.0
        return self.sparsity

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAConfig.effective_ref_keep_rate

effective_ref_keep_rate: float | None

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

Source
def effective_ref_keep_rate(self) -> float | None:
    """The per-reference-video keep rate, or None when references stay dense."""
    return self.ref_keep_rate if self.ref_policy == "p2_multi_region" else None

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAConfig.layer_sparsity

layer_sparsity(layer_idx: int, step_index: int) -> float
Source
def layer_sparsity(self, layer_idx: int, step_index: int) -> float:
    if not self.enabled:
        return 0.0
    if step_index < self.dense_first_n_steps:
        return 0.0
    if layer_idx in self.dense_layers:
        return 0.0
    return self.sparsity

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAGeometry

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

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

Source
@dataclass(frozen=True)
class MiniMaxH3VSAGeometry:
    """Packed-sequence tile map shared by routing, reference, and Metal paths."""

    prefix_segments: tuple[int, ...]
    dit_seq_shape: tuple[int, int, int]
    tile_shape: tuple[int, int, int]
    tile_elems: int
    total_seq_length: int
    num_prefix_tiles: int
    num_video_tiles: int
    variable_block_sizes: np.ndarray
    untile_combined_index: np.ndarray
    tile_partition_indices: np.ndarray
    # One [start, end) tile range per sparse video region in packed order, the
    # generated video last. Tile numbering includes the dense prefix tiles.
    video_tile_spans: tuple[tuple[int, int], ...] = ()
    segments: tuple[int | tuple[int, int, int], ...] = ()
    reference_regions: tuple[tuple[int, int, int], ...] = ()

    def __post_init__(self) -> None:
        if not self.video_tile_spans:
            object.__setattr__(self, "video_tile_spans",
                               ((self.num_prefix_tiles, self.num_prefix_tiles + self.num_video_tiles), ))

    @property
    def dense_rows(self) -> np.ndarray:
        """Packed positions of the dense rows, in prefix-tile order."""
        return self.tile_partition_indices[:self.prefix_length]

    @cached_property
    def dense_rows_lead(self) -> bool:
        """True when the dense rows are the first packed rows (no interleaved region)."""
        return bool(np.array_equal(self.dense_rows, np.arange(self.prefix_length)))

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

    @property
    def padded_length(self) -> int:
        return self.num_tiles * self.tile_elems

    @property
    def prefix_length(self) -> int:
        return int(sum(self.prefix_segments))

    @cached_property
    def tile_gather_index(self):
        import mlx.core as mx

        index = np.full(self.padded_length, self.total_seq_length, dtype=np.int32)
        index[self.untile_combined_index] = np.arange(self.total_seq_length, dtype=np.int32)
        return mx.array(index)

    @cached_property
    def prefix_gather_index(self):
        import mlx.core as mx

        index = np.full(self.num_prefix_tiles * self.tile_elems, self.prefix_length, dtype=np.int32)
        index[self.untile_combined_index[self.dense_rows]] = np.arange(self.prefix_length, dtype=np.int32)
        return mx.array(index)

    @cached_property
    def dense_row_index(self):
        import mlx.core as mx

        return mx.array(self.dense_rows.astype(np.int32))

    @cached_property
    def untile_index(self):
        import mlx.core as mx

        return mx.array(self.untile_combined_index, dtype=mx.int32)

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAGeometry.dense_rows

dense_rows: np.ndarray

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

Source
def dense_rows(self) -> np.ndarray:
    """Packed positions of the dense rows, in prefix-tile order."""
    return self.tile_partition_indices[:self.prefix_length]

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAGeometry.dense_rows_lead

dense_rows_lead: bool

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

Source
def dense_rows_lead(self) -> bool:
    """True when the dense rows are the first packed rows (no interleaved region)."""
    return bool(np.array_equal(self.dense_rows, np.arange(self.prefix_length)))

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAStats

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

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

Source
@dataclass
class MiniMaxH3VSAStats:
    """Filled during a sparse forward so the pipeline can report achieved sparsity."""

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

    def record(self, call: MiniMaxH3VSAStats) -> None:
        """Aggregate equally sized video-query tile maps across blocks and steps."""
        self.attention_calls += 1
        self.sparse_calls += int(call.layer_sparsity > 0.0)
        self.impl_counts[call.impl] = self.impl_counts.get(call.impl, 0) + 1
        self.impl = next(iter(self.impl_counts)) if len(self.impl_counts) == 1 else "mixed"
        for name in ("layer_sparsity", "video_keep", "achieved_sparsity"):
            old = getattr(self, name)
            setattr(self, name, old + (getattr(call, name) - old) / self.attention_calls)
        self.num_prefix_tiles = call.num_prefix_tiles
        self.num_video_tiles = call.num_video_tiles
        if call.dense_fallback_reason and call.dense_fallback_reason not in self.fallback_reasons:
            self.fallback_reasons.append(call.dense_fallback_reason)
        self.dense_fallback_reason = "; ".join(self.fallback_reasons) or None

fastvideo.mlx_runtime.minimax_h3_vsa.MiniMaxH3VSAStats.record

record(call: MiniMaxH3VSAStats) -> None

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

Source
def record(self, call: MiniMaxH3VSAStats) -> None:
    """Aggregate equally sized video-query tile maps across blocks and steps."""
    self.attention_calls += 1
    self.sparse_calls += int(call.layer_sparsity > 0.0)
    self.impl_counts[call.impl] = self.impl_counts.get(call.impl, 0) + 1
    self.impl = next(iter(self.impl_counts)) if len(self.impl_counts) == 1 else "mixed"
    for name in ("layer_sparsity", "video_keep", "achieved_sparsity"):
        old = getattr(self, name)
        setattr(self, name, old + (getattr(call, name) - old) / self.attention_calls)
    self.num_prefix_tiles = call.num_prefix_tiles
    self.num_video_tiles = call.num_video_tiles
    if call.dense_fallback_reason and call.dense_fallback_reason not in self.fallback_reasons:
        self.fallback_reasons.append(call.dense_fallback_reason)
    self.dense_fallback_reason = "; ".join(self.fallback_reasons) or None

fastvideo.mlx_runtime.minimax_h3_vsa.vsa_gate_key

vsa_gate_key(block_index: int) -> str
Source
def vsa_gate_key(block_index: int) -> str:
    return f"transformer_blocks.{block_index}.{VSA_GATE_KEY_SUFFIX}"

fastvideo.mlx_runtime.minimax_h3_vsa.expected_vsa_gate_keys

expected_vsa_gate_keys(num_layers: int) -> tuple[str, ...]
Source
def expected_vsa_gate_keys(num_layers: int) -> tuple[str, ...]:
    return tuple(vsa_gate_key(index) for index in range(num_layers))

fastvideo.mlx_runtime.minimax_h3_vsa.is_vsa_gate_key

is_vsa_gate_key(key: str) -> bool
Source
def is_vsa_gate_key(key: str) -> bool:
    return key.endswith(VSA_GATE_KEY_SUFFIX)

fastvideo.mlx_runtime.minimax_h3_vsa.dense_only_vsa_error

dense_only_vsa_error(checkpoint_dir: str | Any) -> DenseOnlyVSACheckpointError
Source
def dense_only_vsa_error(checkpoint_dir: str | Any) -> DenseOnlyVSACheckpointError:
    return DenseOnlyVSACheckpointError(
        f"VSA was requested but {checkpoint_dir} is a dense-only MLX H3 checkpoint "
        f"(no `{VSA_GATE_KEY_SUFFIX}` matrices / manifest vsa.capable). "
        "Reconvert with `--include-vsa`, for example:\n"
        "  python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \\\n"
        "    --model-root <FastH3 transformer dir> --out <new dir> --formats int6 --include-vsa")

fastvideo.mlx_runtime.minimax_h3_vsa.compute_topk

compute_topk(sparsity: float, num_blocks: int) -> int

Blocks to keep for a sparsity level, clamped to [1, num_blocks].

Source
def compute_topk(sparsity: float, num_blocks: int) -> int:
    """Blocks to keep for a sparsity level, clamped to [1, num_blocks]."""
    if num_blocks <= 0:
        return 0
    return max(1, min(math.ceil((1.0 - sparsity) * num_blocks), num_blocks))

fastvideo.mlx_runtime.minimax_h3_vsa.parse_dense_layers

parse_dense_layers(value: str | None) -> tuple[int, ...]
Source
def parse_dense_layers(value: str | None) -> tuple[int, ...]:
    if value is None or value.strip() == "":
        return ()
    layers = tuple(int(part.strip()) for part in value.split(",") if part.strip())
    if any(layer < 0 for layer in layers):
        raise ValueError(f"vsa dense layers must be non-negative, got {value!r}.")
    return layers

fastvideo.mlx_runtime.minimax_h3_vsa.prefix_segments_from_layout

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

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

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

Source
def prefix_segments_from_layout(layout: Any, patch_size: tuple[int, int, int]) -> tuple[int, ...]:
    """Segment sizes preceding the generated-video tail, matching the PyTorch stage.

    Prefix tiles are chunked per segment, so a segment boundary must fall where
    the packed rows change kind. A Ref2VA layout interleaves a video
    reference's audio rows with the references themselves, so it is split per
    reference from ``layout.reference_segments``; T2VA/FL2VA keep the single
    condition segment they have always had.
    """
    n_text = int(layout.text_indices.shape[0])
    n_video = ((layout.num_video_latent_frames // patch_size[0]) * (layout.latent_height // patch_size[1]) *
               (layout.latent_width // patch_size[2]))
    reference_segments = getattr(layout, "reference_segments", ())
    if reference_segments:
        # Packed order is [text | ordered reference segments | target audio |
        # target video]; everything before the video tail is prefix.
        n_target_audio = int(layout.audio_indices.shape[0]) - int(layout.num_condition_audio_rows)
        segments = (n_text, *(int(rows) for _, rows, _ in reference_segments), n_target_audio)
        prefix_length = sum(segments)
    else:
        n_cond = int(layout.num_condition_video_rows)
        n_audio = int(layout.audio_indices.shape[0])
        segments = (n_text, n_cond, n_audio)
        prefix_length = n_text + n_cond + n_audio
    if prefix_length + n_video != int(layout.sequence_length):
        raise ValueError("VSA-H3 supports the standard [text|cond|audio|video] packing only; "
                         f"segments {segments} + video {n_video} do not sum to "
                         f"sequence length {layout.sequence_length}.")
    return tuple(segments)

fastvideo.mlx_runtime.minimax_h3_vsa.ref2va_vsa_segments_from_layout

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

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

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

Source
def ref2va_vsa_segments_from_layout(
    layout: Any,
    patch_size: tuple[int, int, int],
) -> tuple[int | tuple[int, int, int], ...]:
    """Packed segments with every reference VIDEO as a sparse region (``p2_multi_region``).

    Mirrors ``_h3_vsa_ref2va_segments`` in the CUDA denoising stage:
    ``[text | reference spans... | target audio | target video]`` where each
    reference video is a token grid and text, reference audio, image
    references and the target audio are dense row counts.
    """
    reference_segments = getattr(layout, "reference_segments", ())
    if not reference_segments:
        raise ValueError("Ref2VA VSA sparsification needs per-reference spans on the layout; "
                         "this layout carries none (T2VA/FL2VA packing).")
    cursor = int(layout.text_indices.shape[0])
    segments: list[int | tuple[int, int, int]] = [cursor]
    for kind, rows, latent_shape in reference_segments:
        rows = int(rows)
        if kind == "video":
            if any(int(axis) <= 0 or int(axis) % patch for axis, patch in zip(latent_shape, patch_size, strict=True)):
                raise ValueError(f"Reference video latent shape {tuple(latent_shape)} is not a positive multiple of "
                                 f"patch {tuple(patch_size)}.")
            grid = tuple(int(axis) // patch for axis, patch in zip(latent_shape, patch_size, strict=True))
            if math.prod(grid) != rows:
                raise ValueError(f"Reference video span of {rows} rows does not match its grid {grid}.")
            segments.append(grid)  # type: ignore[arg-type]
        else:
            segments.append(rows)
        cursor += rows
    n_target_audio = int(layout.audio_indices.shape[0]) - int(layout.num_condition_audio_rows)
    segments.append(n_target_audio)
    cursor += n_target_audio
    target = dit_seq_shape_from_layout(layout, patch_size)
    segments.append(target)
    if cursor + math.prod(target) != int(layout.sequence_length):
        raise ValueError(f"Ref2VA reference spans {reference_segments} + target segments do not tile the packed "
                         f"sequence: reached row {cursor} + video {math.prod(target)} != {layout.sequence_length}.")
    return tuple(segments)

fastvideo.mlx_runtime.minimax_h3_vsa.dit_seq_shape_from_layout

dit_seq_shape_from_layout(layout: Any, patch_size: tuple[int, int, int]) -> tuple[int, int, int]
Source
def dit_seq_shape_from_layout(layout: Any, patch_size: tuple[int, int, int]) -> tuple[int, int, int]:
    return (
        int(layout.num_video_latent_frames // patch_size[0]),
        int(layout.latent_height // patch_size[1]),
        int(layout.latent_width // patch_size[2]),
    )

fastvideo.mlx_runtime.minimax_h3_vsa.validate_h3_tile_geometry

validate_h3_tile_geometry(prefix_segments: tuple[int | tuple[int, int, int], ...], dit_seq_shape: tuple[int, int, int] | None, variable_block_sizes: np.ndarray, untile_combined_index: np.ndarray, tile_elems: int) -> None
Source
def validate_h3_tile_geometry(
    prefix_segments: tuple[int | tuple[int, int, int], ...],
    dit_seq_shape: tuple[int, int, int] | None,
    variable_block_sizes: np.ndarray,
    untile_combined_index: np.ndarray,
    tile_elems: int,
) -> None:
    total = _segments_total(prefix_segments) + (math.prod(dit_seq_shape) if dit_seq_shape is not None else 0)
    n_pad = int(variable_block_sizes.size) * tile_elems
    sizes_min = int(variable_block_sizes.min()) if variable_block_sizes.size else 0
    sizes_max = int(variable_block_sizes.max()) if variable_block_sizes.size else 0
    sizes_sum = int(variable_block_sizes.sum())
    if sizes_min < 1 or sizes_max > tile_elems or sizes_sum != total:
        raise ValueError(f"VSA-H3 tile sizes out of bounds for prefix={prefix_segments}, video={dit_seq_shape}, "
                         f"tile_elems={tile_elems}: min={sizes_min}, max={sizes_max}, sum={sizes_sum}, "
                         f"expected sum={total}.")
    if int(untile_combined_index.size) != total:
        raise ValueError(f"VSA-H3 untile index has {untile_combined_index.size} entries for a packed "
                         f"sequence of {total} rows (prefix={prefix_segments}, video={dit_seq_shape}).")
    idx_min = int(untile_combined_index.min())
    idx_max = int(untile_combined_index.max())
    if idx_min < 0 or idx_max >= n_pad:
        raise ValueError(f"VSA-H3 untile index is not an injective map into non-pad slots: range "
                         f"[{idx_min}, {idx_max}] vs padded length {n_pad} "
                         f"(prefix={prefix_segments}, video={dit_seq_shape}).")
    in_tile_offset = untile_combined_index % tile_elems
    maps_into_pad = bool((in_tile_offset >= variable_block_sizes[untile_combined_index // tile_elems]).any())
    if maps_into_pad or int(np.unique(untile_combined_index).size) != total:
        raise ValueError(f"VSA-H3 untile index is not an injective map into non-pad slots: "
                         f"pad-slot hit={maps_into_pad} "
                         f"(prefix={prefix_segments}, video={dit_seq_shape}).")

fastvideo.mlx_runtime.minimax_h3_vsa.build_h3_tile_geometry

build_h3_tile_geometry(prefix_segments: tuple[int, ...], dit_seq_shape: tuple[int, int, int], tile_size: int = 64) -> MiniMaxH3VSAGeometry

Tile the packed sequence: segment-pure prefix chunks, then video tiles.

Source
def build_h3_tile_geometry(
    prefix_segments: tuple[int, ...],
    dit_seq_shape: tuple[int, int, int],
    tile_size: int = 64,
) -> MiniMaxH3VSAGeometry:
    """Tile the packed sequence: segment-pure prefix chunks, then video tiles."""
    return build_h3_segment_geometry((*prefix_segments, _grid(dit_seq_shape)), tile_size)

fastvideo.mlx_runtime.minimax_h3_vsa.build_h3_segment_geometry

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

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

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

Source
def build_h3_segment_geometry(
    segments: tuple[int | tuple[int, int, int], ...],
    tile_size: int = 64,
) -> MiniMaxH3VSAGeometry:
    """Tile a packed sequence of interleaved dense segments and video regions.

    ``segments`` lists, in packed order, dense row counts (``int``) and video
    token grids (``(t, h, w)``). Rows are permuted to ``[dense chunks][region 0
    tiles][region 1 tiles]...`` exactly as the CUDA ``_h3_segment_tile_geometry``
    does; the last region is the generated video.
    """
    tile_shape = VSA_H3_TILE_SHAPES.get(int(tile_size))
    if tile_shape is None:
        raise ValueError(f"VSA-H3 tile_size must be one of {sorted(VSA_H3_TILE_SHAPES)}, got {tile_size!r}")
    tile_elems = math.prod(tile_shape)
    dense_sizes: list[int] = []
    dense_rows: list[np.ndarray] = []
    regions: list[tuple[int, tuple[int, int, int]]] = []
    kept: list[int | tuple[int, int, int]] = []
    cursor = 0
    for segment in segments:
        if isinstance(segment, tuple):
            grid = tuple(int(axis) for axis in segment)
            if len(grid) != 3 or any(axis <= 0 for axis in grid):
                raise ValueError(f"VSA-H3 video axes must be positive, got {segment}.")
            regions.append((cursor, grid))  # type: ignore[arg-type]
            kept.append(grid)  # type: ignore[arg-type]
            cursor += math.prod(grid)
            continue
        rows = int(segment)
        if rows <= 0:
            continue
        full, rem = divmod(rows, tile_elems)
        dense_sizes.extend([tile_elems] * full)
        if rem:
            dense_sizes.append(rem)
        dense_rows.append(np.arange(cursor, cursor + rows, dtype=np.int64))
        kept.append(rows)
        cursor += rows
    if not regions:
        raise ValueError("VSA-H3 needs at least one video region.")
    num_prefix_tiles = len(dense_sizes)

    index_parts: list[np.ndarray] = []
    size_parts: list[np.ndarray] = []
    spans: list[tuple[int, int]] = []
    tile_cursor = num_prefix_tiles
    for start, grid in regions:
        sizes = _video_tile_sizes(grid, tile_shape)
        index_parts.append(_video_tile_partition_indices(grid, tile_shape) + start)
        size_parts.append(sizes.astype(np.int64))
        spans.append((tile_cursor, tile_cursor + int(sizes.size)))
        tile_cursor += int(sizes.size)
    num_video_tiles = tile_cursor - num_prefix_tiles
    tile_partition_indices = np.concatenate([*dense_rows, *index_parts], axis=0)
    variable_block_sizes = np.concatenate([np.asarray(dense_sizes, dtype=np.int64), *size_parts], axis=0)
    non_pad_index = _non_pad_index(variable_block_sizes, tile_elems)
    untile_combined_index = non_pad_index[np.argsort(tile_partition_indices, kind="stable")]
    validate_h3_tile_geometry(tuple(kept), None, variable_block_sizes, untile_combined_index, tile_elems)
    prefix_segments = tuple(segment for segment in kept if not isinstance(segment, tuple))
    return MiniMaxH3VSAGeometry(
        prefix_segments=prefix_segments,  # type: ignore[arg-type]
        dit_seq_shape=regions[-1][1],
        tile_shape=tile_shape,
        tile_elems=tile_elems,
        total_seq_length=cursor,
        num_prefix_tiles=num_prefix_tiles,
        num_video_tiles=num_video_tiles,
        variable_block_sizes=variable_block_sizes,
        untile_combined_index=untile_combined_index,
        tile_partition_indices=tile_partition_indices,
        video_tile_spans=tuple(spans),
        segments=tuple(kept),
        reference_regions=tuple(grid for _, grid in regions[:-1]),
    )

fastvideo.mlx_runtime.minimax_h3_vsa.token_tile_and_valid

token_tile_and_valid(variable_block_sizes: np.ndarray, tile_elems: int) -> tuple[np.ndarray, np.ndarray]
Source
def token_tile_and_valid(variable_block_sizes: np.ndarray, tile_elems: int) -> tuple[np.ndarray, np.ndarray]:
    token_tile = np.repeat(np.arange(variable_block_sizes.size, dtype=np.int64), tile_elems)
    token_valid = (np.arange(tile_elems, dtype=np.int64)[None, :] < variable_block_sizes[:, None]).reshape(-1)
    return token_tile, token_valid

fastvideo.mlx_runtime.minimax_h3_vsa.region_sparsities

region_sparsities(num_regions: int, sparsity: float, ref_keep_rate: float | None) -> tuple[float, ...]

Per-region sparsity: reference videos keep ref_keep_rate, the target follows sparsity.

Mirrors MiniMaxH3VSAMetadataBuilder.build: a dense step (sparsity <= 0) keeps every region dense, whatever the reference keep rate.

Source
def region_sparsities(num_regions: int, sparsity: float, ref_keep_rate: float | None) -> tuple[float, ...]:
    """Per-region sparsity: reference videos keep ``ref_keep_rate``, the target follows ``sparsity``.

    Mirrors ``MiniMaxH3VSAMetadataBuilder.build``: a dense step (sparsity <= 0)
    keeps every region dense, whatever the reference keep rate.
    """
    if num_regions < 1:
        raise ValueError("VSA-H3 needs at least one video region.")
    if sparsity <= 0.0:
        return (0.0, ) * num_regions
    if num_regions == 1:
        return (float(sparsity), )
    if ref_keep_rate is None:
        raise ValueError(f"A VSA-H3 build with {num_regions - 1} reference-video region(s) needs ref_keep_rate.")
    return (1.0 - float(ref_keep_rate), ) * (num_regions - 1) + (float(sparsity), )

fastvideo.mlx_runtime.minimax_h3_vsa.build_block_mask

build_block_mask(scores: np.ndarray, num_prefix_tiles: int, num_video_tiles: int, sparsity: float, exempt: bool, spans: tuple[tuple[int, int], ...] | None = None, span_sparsities: tuple[float, ...] | None = None) -> np.ndarray

scores: [..., n_tiles, n_tiles] -> bool mask, same shape.

Mirrors _build_block_mask in the PyTorch H3 backend: every video query keeps its own top-k of EACH video region (spans / span_sparsities).

Source
def build_block_mask(
    scores: np.ndarray,
    num_prefix_tiles: int,
    num_video_tiles: int,
    sparsity: float,
    exempt: bool,
    spans: tuple[tuple[int, int], ...] | None = None,
    span_sparsities: tuple[float, ...] | None = None,
) -> np.ndarray:
    """scores: [..., n_tiles, n_tiles] -> bool mask, same shape.

    Mirrors ``_build_block_mask`` in the PyTorch H3 backend: every video query
    keeps its own top-k of EACH video region (``spans`` / ``span_sparsities``).
    """
    n_tiles = scores.shape[-1]
    spans, topks = _region_plan(num_prefix_tiles, num_video_tiles, sparsity, exempt, spans, span_sparsities)
    if all(k == end - start for k, (start, end) in zip(topks, spans, strict=True)):
        return np.ones_like(scores, dtype=bool)
    mask = np.zeros_like(scores, dtype=bool)
    if exempt or num_prefix_tiles == 0:
        for (start, end), k_span in zip(spans, topks, strict=True):
            if k_span == end - start:
                mask[..., start:end] = True
                continue
            idx = np.argsort(-scores[..., start:end], axis=-1, kind="stable")[..., :k_span] + start
            np.put_along_axis(mask, idx, True, axis=-1)
        mask[..., :num_prefix_tiles] = True
    else:
        k_total = min(topks[0] + num_prefix_tiles, n_tiles)
        idx = np.argsort(-scores, axis=-1, kind="stable")[..., :k_total]
        np.put_along_axis(mask, idx, True, axis=-1)
    mask[..., :num_prefix_tiles, :] = True
    return mask

fastvideo.mlx_runtime.minimax_h3_vsa.resolve_impl

resolve_impl(requested: VSAImpl, head_dim: int, tile_elems: int) -> str
Source
def resolve_impl(requested: VSAImpl, head_dim: int, tile_elems: int) -> str:
    if requested == "reference" or (requested == "auto" and tile_elems not in _AUTO_SIMD_TILES):
        return "reference"

    from fastvideo.mlx_runtime.minimax_h3_vsa_simd import (SIMD_HEAD_DIM, SIMD_TILE_SIZES, simd_kernel_available,
                                                           simd_kernel_error)

    if tile_elems not in SIMD_TILE_SIZES or head_dim != SIMD_HEAD_DIM:
        if requested == "simd":
            logger.warning_once(
                f"SIMD VSA requires tile {' or '.join(map(str, SIMD_TILE_SIZES))} and head dim {SIMD_HEAD_DIM} "
                f"(got tile={tile_elems} head_dim={head_dim}); using reference VSA")
        return "reference"
    if not simd_kernel_available():
        reason = simd_kernel_error() or "SIMD-group VSA kernel is unavailable"
        logger.warning_once(f"SIMD VSA kernel unavailable ({reason}); using reference VSA")
        return "reference"
    return "simd"

fastvideo.mlx_runtime.minimax_h3_vsa.h3_vsa_attention

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

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

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

Source
def h3_vsa_attention(
    query,
    key,
    value,
    geometry: MiniMaxH3VSAGeometry,
    *,
    sparsity: float,
    exempt: bool = True,
    gate_compress=None,
    impl: VSAImpl = "auto",
    stats: MiniMaxH3VSAStats | None = None,
    ref_keep_rate: float | None = None,
):
    """Packed ``[S, H, D]`` VSA attention. Falls back to dense SDPA when sparsity is 0.

    ``ref_keep_rate`` is the ``p2_multi_region`` keep rate of each reference-video
    region of ``geometry``; geometries with reference regions require it.
    """
    import mlx.core as mx

    _, heads, dim = query.shape
    scale = dim**-0.5
    if stats is not None:
        stats.configured_sparsity = sparsity
        stats.layer_sparsity = sparsity
        stats.tile_size = geometry.tile_elems
        stats.prefix_mode = "exempt" if exempt else "compete"
        stats.num_prefix_tiles = geometry.num_prefix_tiles
        stats.num_video_tiles = geometry.num_video_tiles

    if sparsity <= 0.0 and gate_compress is None:
        if stats is not None:
            stats.impl = "dense"
            stats.achieved_sparsity = 0.0
            stats.video_keep = geometry.num_video_tiles
        return _dense_sdpa(query, key, value, scale)

    q_tiled = _tile_hidden(query, geometry)
    k_tiled = _tile_hidden(key, geometry)
    v_tiled = _tile_hidden(value, geometry)
    q_pooled = _pool_tiles(q_tiled, geometry.variable_block_sizes, geometry.tile_elems)
    k_pooled = _pool_tiles(k_tiled, geometry.variable_block_sizes, geometry.tile_elems)
    scores = (q_pooled @ k_pooled.transpose(0, 2, 1)) / (dim**0.5)

    span_sparsities = region_sparsities(len(geometry.video_tile_spans), sparsity, ref_keep_rate)
    spans = geometry.video_tile_spans
    k_vid = sum(compute_topk(sp, end - start) for (start, end), sp in zip(spans, span_sparsities, strict=True))
    if stats is not None:
        stats.video_keep = k_vid if sparsity > 0.0 else geometry.num_video_tiles
        stats.achieved_sparsity = (1.0 - (k_vid / geometry.num_video_tiles) if geometry.num_video_tiles else 0.0)

    if sparsity <= 0.0:
        out_tiled = _tile_hidden(_dense_sdpa(query, key, value, scale), geometry)
        chosen = "dense"
    else:
        block_idx, block_num = _block_indices_from_scores(
            scores,
            geometry.num_prefix_tiles,
            geometry.num_video_tiles,
            sparsity,
            exempt,
            spans,
            span_sparsities,
        )
        if stats is not None and not exempt:
            stats.video_keep = float(mx.mean(mx.sum(block_idx >= geometry.num_prefix_tiles, axis=-1)).item())
            stats.achieved_sparsity = 1.0 - stats.video_keep / geometry.num_video_tiles
        chosen = resolve_impl(impl, dim, geometry.tile_elems)
        if stats is not None and impl == "simd" and chosen != "simd":
            from fastvideo.mlx_runtime.minimax_h3_vsa_simd import simd_kernel_error

            stats.dense_fallback_reason = ("SIMD requires tile 64 or 128 and head dim 128" if
                                           geometry.tile_elems not in (64, 128) or dim != 128 else simd_kernel_error())
        if not geometry.prefix_length:
            prefix_out = query[:0]
        else:
            dense_query = (query[:geometry.prefix_length]
                           if geometry.dense_rows_lead else query[geometry.dense_row_index])
            prefix_out = _dense_sdpa(dense_query, key, value, scale)
        n_prefix_pad = geometry.num_prefix_tiles * geometry.tile_elems
        if chosen == "simd":
            from fastvideo.mlx_runtime.minimax_h3_vsa_simd import disable_simd_kernel, simd_block_sparse

            try:
                video_tiled = simd_block_sparse(q_tiled, k_tiled, v_tiled, block_idx, block_num, geometry, scale)
                # MLX compiles and executes custom Metal kernels lazily.
                mx.eval(video_tiled)
                video_tiled = video_tiled[n_prefix_pad:]
            except Exception as error:  # noqa: BLE001 - keep generation alive on kernel failure
                disable_simd_kernel(error)
                logger.warning_once(f"SIMD VSA kernel failed ({error}); falling back to reference gather+SDPA")
                chosen = "reference"
                if stats is not None:
                    stats.dense_fallback_reason = f"simd kernel failed: {error}"
                video_tiled = _reference_gather_sdpa(q_tiled, k_tiled, v_tiled, block_idx, geometry, scale)
        elif _reference_full_mask_fits(geometry, heads):
            scores_np = np.array(scores, dtype=np.float32)
            mask = build_block_mask(
                scores_np,
                geometry.num_prefix_tiles,
                geometry.num_video_tiles,
                sparsity,
                exempt,
                spans,
                span_sparsities,
            )
            video_tiled = _reference_token_sdpa(q_tiled, k_tiled, v_tiled, mask, geometry, scale)[n_prefix_pad:]
        else:
            video_tiled = _reference_gather_sdpa(q_tiled, k_tiled, v_tiled, block_idx, geometry, scale)
        video_tiled = video_tiled.astype(query.dtype)
        prefix_tiled = mx.concatenate([
            prefix_out,
            mx.zeros((1, heads, dim), dtype=query.dtype),
        ])[geometry.prefix_gather_index]
        # Prefix query tiles are dense; keep fused-SDPA prefix rows and sparse video tiles.
        out_tiled = mx.concatenate([prefix_tiled, video_tiled], axis=0)

    if gate_compress is not None:
        gate_tiled = _tile_hidden(gate_compress, geometry)
        out_tiled = out_tiled + _gate_compress_output(scores, v_tiled, gate_tiled, geometry).astype(out_tiled.dtype)

    if stats is not None:
        stats.impl = chosen if sparsity > 0.0 else "dense"
    return _untile_hidden(out_tiled, geometry)

fastvideo.mlx_runtime.minimax_h3_vsa.geometry_is_supported

geometry_is_supported(prefix_segments: tuple[int | tuple[int, int, int], ...], dit_seq_shape: tuple[int, int, int], tile_size: int) -> str | None

Return a fallback reason, or None when VSA can run.

Source
def geometry_is_supported(prefix_segments: tuple[int | tuple[int, int, int], ...], dit_seq_shape: tuple[int, int, int],
                          tile_size: int) -> str | None:
    """Return a fallback reason, or None when VSA can run."""
    try:
        build_h3_segment_geometry((*prefix_segments, _grid(dit_seq_shape)), tile_size)
    except ValueError as error:
        return str(error)
    return None