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:
- Tiles are
[segment-pure prefix chunks] + [3D video tiles]. - Tile sizes 64
(4, 4, 4), 128(4, 4, 8)and 256(4, 8, 8). p2_multi_regionRef2VA policy: every reference VIDEO is its own sparse region with its own top-k atref_keep_rate(image/audio/text rows stay dense prefix), as in the CUDA backend's_build_block_mask.- Per-head pooled Q/K scoring and top-k routing.
- Prefix queries are always dense; prefix keys are
exempt(always kept) orcompete(FLOP-matched top-k). - Optional dense-first steps and per-layer dense overrides.
- Trained
to_gate_compresspooled-compression branch.
Execution backends:
- reference — grouped gather plus batched
mx.fast.scaled_dot_product_attention. Correctness baseline;autofor tiles other than 128 andimpl="reference". - simd — flash-style SIMD-group 8x8 matrix operations over the reference tile map for tile sizes 64 and 128 and head dimension 128. Unsupported shapes and kernel failures fall back to the reference backend.
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) -> _FastvideoLoggerThe 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) -> Nonep2_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</code></pre></details>
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)</code></pre></details>
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</code></pre></details>
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)</code></pre></details>
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)</code></pre></details>
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]),
)</code></pre></details>
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), )</code></pre></details>
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</code></pre></details>
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"</code></pre></details>
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)</code></pre></details>
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