minimax_h3_conditioner
Streamed Qwen3-VL prompt conditioner for MiniMax-H3 on Apple Silicon MLX.
Produces exactly the conditioning the H3 DiT consumes: the language-model hidden states after the first 50 language-model layers plus the per-token modality tags for text prompts.
Memory contract for the 36 GiB tier: the released conditioner is ~66 GB of BF16 and never becomes resident. Tensors are read per-key from the safetensors shards and only the pieces a given forward needs are materialized:
- token embedding table row-gathered per batch (full table never copied);
- one decoder layer (~1 GB BF16) resident at a time, computed in FP32, released before the next layer loads; The forward pass uses MLX. Tokenization uses the Transformers tokenizer API.
fastvideo.mlx_runtime.minimax_h3_conditioner.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_conditioner.ConditionerConfig
ConditionerConfig(hidden_size: int = 5120, num_layers: int = 64, num_attention_heads: int = 64, num_key_value_heads: int = 8, head_dim: int = 128, intermediate_size: int = 25600, rms_norm_eps: float = 1e-06, rope_theta: float = 5000000.0, mrope_section: tuple[int, ...] = (24, 20, 20), vocab_size: int = 151936, vision_start_token_id: int = 151652, image_token_id: int = 151655, video_token_id: int = 151656, vision_spatial_merge_size: int = 2)Source
@dataclass
class ConditionerConfig:
hidden_size: int = 5120
num_layers: int = 64
num_attention_heads: int = 64
num_key_value_heads: int = 8
head_dim: int = 128
intermediate_size: int = 25600
rms_norm_eps: float = 1e-6
rope_theta: float = 5_000_000.0
mrope_section: tuple[int, ...] = (24, 20, 20)
vocab_size: int = 151936
vision_start_token_id: int = 151652
image_token_id: int = 151655
video_token_id: int = 151656
# Vision merge size drives the MRoPE image/video layout; the vision tower
# itself re-reads its own ``vision_config`` from config.json.
vision_spatial_merge_size: int = 2
@classmethod
def from_config_json(cls, path: str | Path) -> ConditionerConfig:
raw = json.loads(Path(path).read_text())
text = raw.get("text_config", raw)
vision = raw.get("vision_config", {})
scaling = text.get("rope_scaling", {}) or {}
return cls(
hidden_size=int(text["hidden_size"]),
num_layers=int(text["num_hidden_layers"]),
num_attention_heads=int(text["num_attention_heads"]),
num_key_value_heads=int(text["num_key_value_heads"]),
head_dim=int(text.get("head_dim", text["hidden_size"] // text["num_attention_heads"])),
intermediate_size=int(text["intermediate_size"]),
rms_norm_eps=float(text["rms_norm_eps"]),
rope_theta=float(text["rope_theta"]),
mrope_section=tuple(int(v) for v in scaling.get("mrope_section", (24, 20, 20))),
vocab_size=int(text.get("vocab_size", 151936)),
vision_start_token_id=int(raw.get("vision_start_token_id", text.get("vision_start_token_id", 151652))),
image_token_id=int(raw.get("image_token_id", text.get("image_token_id", 151655))),
video_token_id=int(raw.get("video_token_id", text.get("video_token_id", 151656))),
vision_spatial_merge_size=int(vision.get("spatial_merge_size", 2)),
)</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_conditioner.ConditionerConfig.from_config_json
from_config_json(path: str | Path) -> ConditionerConfig
Source
@classmethod
def from_config_json(cls, path: str | Path) -> ConditionerConfig:
raw = json.loads(Path(path).read_text())
text = raw.get("text_config", raw)
vision = raw.get("vision_config", {})
scaling = text.get("rope_scaling", {}) or {}
return cls(
hidden_size=int(text["hidden_size"]),
num_layers=int(text["num_hidden_layers"]),
num_attention_heads=int(text["num_attention_heads"]),
num_key_value_heads=int(text["num_key_value_heads"]),
head_dim=int(text.get("head_dim", text["hidden_size"] // text["num_attention_heads"])),
intermediate_size=int(text["intermediate_size"]),
rms_norm_eps=float(text["rms_norm_eps"]),
rope_theta=float(text["rope_theta"]),
mrope_section=tuple(int(v) for v in scaling.get("mrope_section", (24, 20, 20))),
vocab_size=int(text.get("vocab_size", 151936)),
vision_start_token_id=int(raw.get("vision_start_token_id", text.get("vision_start_token_id", 151652))),
image_token_id=int(raw.get("image_token_id", text.get("image_token_id", 151655))),
video_token_id=int(raw.get("video_token_id", text.get("video_token_id", 151656))),
vision_spatial_merge_size=int(vision.get("spatial_merge_size", 2)),
)
fastvideo.mlx_runtime.minimax_h3_conditioner.NVFP4Matrix
NVFP4Matrix(weight: mx.array, scales: mx.array, global_scale: float)
MLX row-major E2M1/E4M3 weights and the export's inverse global scale.
Source
@dataclass(frozen=True)
class NVFP4Matrix:
"""MLX row-major E2M1/E4M3 weights and the export's inverse global scale."""
weight: mx.array
scales: mx.array
global_scale: float
def matmul(self, x):
return mx.quantized_matmul(x, self.weight, self.scales, mode="nvfp4") / self.global_scale</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_conditioner.NVFP4Matrix.matmul
matmul(x)
Source
def matmul(self, x):
return mx.quantized_matmul(x, self.weight, self.scales, mode="nvfp4") / self.global_scale
fastvideo.mlx_runtime.minimax_h3_conditioner.StreamedMiniMaxH3TextConditioner
StreamedMiniMaxH3TextConditioner(component_dir: str | Path, tokenizer_dir: str | Path | None = None)
Layer-streaming Qwen3-VL text stack -> H3 conditioning hidden states.
Source
class StreamedMiniMaxH3TextConditioner:
"""Layer-streaming Qwen3-VL text stack -> H3 conditioning hidden states."""
def __init__(self, component_dir: str | Path, tokenizer_dir: str | Path | None = None):
self.component_dir = Path(component_dir)
self.config = ConditionerConfig.from_config_json(self.component_dir / "config.json")
self.index: Any = _ShardIndex(self.component_dir)
self.tokenizer = self._load_tokenizer(tokenizer_dir)
# Vision weights load on first Ref2VA encode; text-only callers never
# pay for the tower.
self._vision: Any = None
def _load_tokenizer(self, tokenizer_dir: str | Path | None):
from transformers import AutoTokenizer
candidates = [Path(tokenizer_dir)] if tokenizer_dir else [
self.component_dir.parent / "tokenizer",
self.component_dir,
]
last_error = None
for candidate in candidates:
try:
return AutoTokenizer.from_pretrained(str(candidate))
except Exception as error: # noqa: BLE001 - fall through to next candidate
last_error = error
raise RuntimeError(f"Could not load an H3 tokenizer from {candidates}: {last_error}")
# -- public API ------------------------------------------------------
def tokenize(self, prompt: str) -> list[int]:
encoded = self.tokenizer(prompt, add_special_tokens=False)
input_ids = encoded["input_ids"]
if input_ids and isinstance(input_ids[0], list):
if len(input_ids) != 1:
raise ValueError("H3 conditioning expects exactly one sequence.")
input_ids = input_ids[0]
return [int(t) for t in input_ids]
def encode_prompt(self, prompt: str) -> tuple[np.ndarray, np.ndarray]:
"""prompt -> (hidden states (S, hidden), token tags (S,)) both fp32."""
token_ids = self.tokenize(prompt)
return self.encode_tokens(token_ids)
def encode_tokens(self, token_ids: list[int]) -> tuple[np.ndarray, np.ndarray]:
import mlx.core as mx
cfg = self.config
seq_len = len(token_ids)
positions = np.stack([
np.arange(seq_len, dtype=np.float64),
np.arange(seq_len, dtype=np.float64),
np.arange(seq_len, dtype=np.float64),
])
cos, sin = _mrope_cos_sin(positions, cfg)
hidden = self._embed_tokens(token_ids)
if cfg.num_layers < TEXT_ENCODER_LAYER:
raise ValueError(f"Conditioner needs at least {TEXT_ENCODER_LAYER} layers, has {cfg.num_layers}.")
for layer in range(TEXT_ENCODER_LAYER):
hidden = self._decoder_layer(layer, hidden, cos, sin)
# Per-layer sync keeps the 50-layer activation graph bounded.
mx.eval(hidden)
gc.collect()
tags = np.full((seq_len, ), 1, dtype=np.int64) # MINIMAX_H3_TEXT_TAG
return np.asarray(hidden).astype(np.float32), tags
def encode_tokens_with_vision(
self,
token_ids: list[int],
token_tags,
*,
pixel_values: np.ndarray | None = None,
image_grid_thw: np.ndarray | None = None,
pixel_values_videos: np.ndarray | None = None,
video_grid_thw: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Ref2VA encode: splice visual tokens into the prompt, then run the stack.
Mirrors ``MiniMaxH3Qwen3VLConditioner.encode_ids``: image/video features
replace the placeholder rows (in placeholder order), the mixed
text/image/video MRoPE positions replace the plain ``arange``, and each
of the first ``len(deepstack)`` layers adds its DeEPSTACK block back onto
the visual rows. Returns ``(hidden (S, hidden) fp32, token_tags (S,) int64)``
with the caller's tags passed through unchanged.
"""
import mlx.core as mx
from fastvideo.mlx_runtime.minimax_h3_vision import qwen3_vl_rope_index
if (pixel_values is None) != (image_grid_thw is None):
raise ValueError("pixel_values and image_grid_thw must be provided together")
if (pixel_values_videos is None) != (video_grid_thw is None):
raise ValueError("pixel_values_videos and video_grid_thw must be provided together")
cfg = self.config
seq_len = len(token_ids)
ids = np.asarray(token_ids, dtype=np.int64)
has_vision = pixel_values is not None or pixel_values_videos is not None
if has_vision:
positions = qwen3_vl_rope_index(
ids,
image_grid_thw,
video_grid_thw,
vision_start_token_id=cfg.vision_start_token_id,
image_token_id=cfg.image_token_id,
video_token_id=cfg.video_token_id,
spatial_merge_size=cfg.vision_spatial_merge_size,
)
else:
positions = np.stack([np.arange(seq_len, dtype=np.float64)] * 3)
cos, sin = _mrope_cos_sin(positions, cfg)
hidden = self._embed_tokens(token_ids)
hidden, deepstack, visual_rows = self._splice_visual_tokens(hidden, ids, pixel_values, image_grid_thw,
pixel_values_videos, video_grid_thw)
if cfg.num_layers < TEXT_ENCODER_LAYER:
raise ValueError(f"Conditioner needs at least {TEXT_ENCODER_LAYER} layers, has {cfg.num_layers}.")
for layer in range(TEXT_ENCODER_LAYER):
hidden = self._decoder_layer(layer, hidden, cos, sin)
if layer < len(deepstack):
merged = np.asarray(hidden, dtype=np.float32)
merged[visual_rows] += deepstack[layer]
hidden = mx.array(merged)
# Per-layer sync keeps the 50-layer activation graph bounded.
mx.eval(hidden)
gc.collect()
return np.asarray(hidden).astype(np.float32), np.asarray(token_tags, dtype=np.int64)
def _vision_model(self):
if self._vision is None:
from fastvideo.mlx_runtime.minimax_h3_vision import MLXQwen3VLVisionModel
weights = getattr(self.index, "weights", None)
if weights is not None and _VISUAL_PROBE_KEY in weights:
# A resident index (and a format-2 cache) carries the whole tower.
self._vision = MLXQwen3VLVisionModel.from_arrays(self.component_dir / "config.json", weights)
else:
self._vision = MLXQwen3VLVisionModel.from_safetensors_dir(self._vision_source_dir())
return self._vision
def _vision_source_dir(self) -> Path:
"""Where the vision tower's weights are read from when the encoder index lacks its biases.
Caches written before format 2 stored only the ``.weight`` tensors the
resident loader kept, so the tower's LayerNorm, patch-embed and
projection biases are read from the checkpoint the cache records. Newer
caches carry the whole tower and never reach this.
"""
manifest_path = self.component_dir / MLX_NVFP4_ENCODER_MANIFEST
if not manifest_path.is_file():
return self.component_dir
source = json.loads(manifest_path.read_text(encoding="utf-8")).get("source_dir")
if not source:
raise ValueError(f"{manifest_path} records no source_dir, so the vision tower cannot be read "
"from the encoder cache.")
if not Path(source).is_dir():
raise FileNotFoundError(f"The checkpoint this older encoder cache was built from, {source}, is gone; "
"the vision tower's biases live there. Re-export the cache to include the "
"vision tower, or point --conditioner-dir at the checkpoint.")
return Path(source)
def _splice_visual_tokens(self, hidden, ids, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw):
"""Return ``(hidden, deepstack blocks, visual rows)`` with features spliced in."""
cfg = self.config
image_features = image_deepstack = video_features = video_deepstack = None
if pixel_values is not None:
image_features, image_deepstack = self._vision_model().encode(pixel_values, image_grid_thw)
if pixel_values_videos is not None:
video_features, video_deepstack = self._vision_model().encode(pixel_values_videos, video_grid_thw)
embedded = np.asarray(hidden, dtype=np.float32)
image_rows = video_rows = None
if image_features is not None:
image_rows = np.where(ids == cfg.image_token_id)[0]
if image_rows.size != image_features.shape[0]:
raise ValueError(f"Qwen3-VL image features and placeholders disagree: "
f"placeholders={image_rows.size}, features={image_features.shape[0]}")
embedded[image_rows] = image_features
if video_features is not None:
video_rows = np.where(ids == cfg.video_token_id)[0]
if video_rows.size != video_features.shape[0]:
raise ValueError(f"Qwen3-VL video features and placeholders disagree: "
f"placeholders={video_rows.size}, features={video_features.shape[0]}")
embedded[video_rows] = video_features
hidden = mx.array(embedded)
if image_rows is None and video_rows is None:
return hidden, [], np.empty((0, ), dtype=np.int64)
if image_rows is not None and video_rows is not None:
assert image_deepstack is not None and video_deepstack is not None
visual_rows = np.sort(np.concatenate([image_rows, video_rows]))
image_joint = np.isin(visual_rows, image_rows)
video_joint = np.isin(visual_rows, video_rows)
blocks = []
for image_block, video_block in zip(image_deepstack, video_deepstack, strict=True):
combined = np.zeros((visual_rows.size, image_block.shape[-1]), dtype=np.float32)
combined[image_joint] = image_block
combined[video_joint] = video_block
blocks.append(combined)
return hidden, blocks, visual_rows
if image_rows is not None:
return hidden, [np.asarray(block, dtype=np.float32) for block in image_deepstack], image_rows
return hidden, [np.asarray(block, dtype=np.float32) for block in video_deepstack], video_rows
def _embed_tokens(self, token_ids: list[int]):
# Embedding rows gathered individually; the full table is never
# materialized by the BF16 streaming path.
rows = []
for token in token_ids:
key = "model.language_model.embed_tokens.weight"
rows.append(self.index.get_row(key, token))
hidden = mx.array(np.stack(rows).astype(np.float32))
del rows
gc.collect()
return hidden
# -- layers ----------------------------------------------------------
def _decoder_layer(self, index: int, hidden, cos, sin):
cfg = self.config
prefix = f"model.language_model.layers.{index}."
def w(name):
return self.index.get_mlx(prefix + name)
# Self-attention block.
residual = hidden
normed = _rms_norm(hidden, w("input_layernorm.weight"), cfg.rms_norm_eps)
query = _linear(normed, w("self_attn.q_proj.weight"))
key = _linear(normed, w("self_attn.k_proj.weight"))
value = _linear(normed, w("self_attn.v_proj.weight"))
heads = cfg.num_attention_heads
kv_heads = cfg.num_key_value_heads
head_dim = cfg.head_dim
seq_len = hidden.shape[0]
def split_heads(t, count):
return t.reshape(seq_len, count, head_dim)
query = split_heads(query, heads)
key = split_heads(key, kv_heads)
value = split_heads(value, kv_heads)
# Per-head QK RMSNorm (Qwen3), weights only.
query = _rms_norm(query, w("self_attn.q_norm.weight"), cfg.rms_norm_eps)
key = _rms_norm(key, w("self_attn.k_norm.weight"), cfg.rms_norm_eps)
# Interleaved MRoPE.
query = _apply_mrope(query, cos, sin)
key = _apply_mrope(key, cos, sin)
# Fused causal attention with native GQA. A reference request carries
# thousands of vision tokens; explicit fp32 scores would need
# heads * seq^2 * 4 bytes (~14 GB at 7.5k tokens) per layer.
attended = mx.fast.scaled_dot_product_attention(query.transpose(1, 0, 2)[None],
key.transpose(1, 0, 2)[None],
value.transpose(1, 0, 2)[None],
scale=head_dim**-0.5,
mask="causal")[0]
# Attention is head-major here: (heads, sequence, head_dim). Restore
# sequence-major ordering before concatenating heads for o_proj.
attended = attended.transpose(1, 0, 2)
attn_out = _linear(attended.reshape(seq_len, -1), w("self_attn.o_proj.weight"))
hidden = residual + attn_out
# MLP block (SwiGLU).
residual = hidden
normed = _rms_norm(hidden, w("post_attention_layernorm.weight"), cfg.rms_norm_eps)
gate = _linear(normed, w("mlp.gate_proj.weight"))
up = _linear(normed, w("mlp.up_proj.weight"))
down = _linear(gate * mx.sigmoid(gate) * up, w("mlp.down_proj.weight"))
return residual + down
def close(self) -> None:
self._vision = None
self.index.close()</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_conditioner.StreamedMiniMaxH3TextConditioner.tokenize
tokenize(prompt: str) -> list[int]
Source
def tokenize(self, prompt: str) -> list[int]:
encoded = self.tokenizer(prompt, add_special_tokens=False)
input_ids = encoded["input_ids"]
if input_ids and isinstance(input_ids[0], list):
if len(input_ids) != 1:
raise ValueError("H3 conditioning expects exactly one sequence.")
input_ids = input_ids[0]
return [int(t) for t in input_ids]
fastvideo.mlx_runtime.minimax_h3_conditioner.StreamedMiniMaxH3TextConditioner.encode_prompt
encode_prompt(prompt: str) -> tuple[np.ndarray, np.ndarray]
prompt -> (hidden states (S, hidden), token tags (S,)) both fp32.
Source
def encode_prompt(self, prompt: str) -> tuple[np.ndarray, np.ndarray]:
"""prompt -> (hidden states (S, hidden), token tags (S,)) both fp32."""
token_ids = self.tokenize(prompt)
return self.encode_tokens(token_ids)
fastvideo.mlx_runtime.minimax_h3_conditioner.StreamedMiniMaxH3TextConditioner.encode_tokens
encode_tokens(token_ids: list[int]) -> tuple[np.ndarray, np.ndarray]
Source
def encode_tokens(self, token_ids: list[int]) -> tuple[np.ndarray, np.ndarray]:
import mlx.core as mx
cfg = self.config
seq_len = len(token_ids)
positions = np.stack([
np.arange(seq_len, dtype=np.float64),
np.arange(seq_len, dtype=np.float64),
np.arange(seq_len, dtype=np.float64),
])
cos, sin = _mrope_cos_sin(positions, cfg)
hidden = self._embed_tokens(token_ids)
if cfg.num_layers < TEXT_ENCODER_LAYER:
raise ValueError(f"Conditioner needs at least {TEXT_ENCODER_LAYER} layers, has {cfg.num_layers}.")
for layer in range(TEXT_ENCODER_LAYER):
hidden = self._decoder_layer(layer, hidden, cos, sin)
# Per-layer sync keeps the 50-layer activation graph bounded.
mx.eval(hidden)
gc.collect()
tags = np.full((seq_len, ), 1, dtype=np.int64) # MINIMAX_H3_TEXT_TAG
return np.asarray(hidden).astype(np.float32), tags</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_conditioner.StreamedMiniMaxH3TextConditioner.encode_tokens_with_vision
encode_tokens_with_vision(token_ids: list[int], token_tags, *, pixel_values: np.ndarray | None = None, image_grid_thw: np.ndarray | None = None, pixel_values_videos: np.ndarray | None = None, video_grid_thw: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]
Ref2VA encode: splice visual tokens into the prompt, then run the stack.
Mirrors MiniMaxH3Qwen3VLConditioner.encode_ids: image/video features
replace the placeholder rows (in placeholder order), the mixed
text/image/video MRoPE positions replace the plain arange, and each
of the first len(deepstack) layers adds its DeEPSTACK block back onto
the visual rows. Returns (hidden (S, hidden) fp32, token_tags (S,) int64)
with the caller's tags passed through unchanged.
Source
def encode_tokens_with_vision(
self,
token_ids: list[int],
token_tags,
*,
pixel_values: np.ndarray | None = None,
image_grid_thw: np.ndarray | None = None,
pixel_values_videos: np.ndarray | None = None,
video_grid_thw: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Ref2VA encode: splice visual tokens into the prompt, then run the stack.
Mirrors ``MiniMaxH3Qwen3VLConditioner.encode_ids``: image/video features
replace the placeholder rows (in placeholder order), the mixed
text/image/video MRoPE positions replace the plain ``arange``, and each
of the first ``len(deepstack)`` layers adds its DeEPSTACK block back onto
the visual rows. Returns ``(hidden (S, hidden) fp32, token_tags (S,) int64)``
with the caller's tags passed through unchanged.
"""
import mlx.core as mx
from fastvideo.mlx_runtime.minimax_h3_vision import qwen3_vl_rope_index
if (pixel_values is None) != (image_grid_thw is None):
raise ValueError("pixel_values and image_grid_thw must be provided together")
if (pixel_values_videos is None) != (video_grid_thw is None):
raise ValueError("pixel_values_videos and video_grid_thw must be provided together")
cfg = self.config
seq_len = len(token_ids)
ids = np.asarray(token_ids, dtype=np.int64)
has_vision = pixel_values is not None or pixel_values_videos is not None
if has_vision:
positions = qwen3_vl_rope_index(
ids,
image_grid_thw,
video_grid_thw,
vision_start_token_id=cfg.vision_start_token_id,
image_token_id=cfg.image_token_id,
video_token_id=cfg.video_token_id,
spatial_merge_size=cfg.vision_spatial_merge_size,
)
else:
positions = np.stack([np.arange(seq_len, dtype=np.float64)] * 3)
cos, sin = _mrope_cos_sin(positions, cfg)
hidden = self._embed_tokens(token_ids)
hidden, deepstack, visual_rows = self._splice_visual_tokens(hidden, ids, pixel_values, image_grid_thw,
pixel_values_videos, video_grid_thw)
if cfg.num_layers < TEXT_ENCODER_LAYER:
raise ValueError(f"Conditioner needs at least {TEXT_ENCODER_LAYER} layers, has {cfg.num_layers}.")
for layer in range(TEXT_ENCODER_LAYER):
hidden = self._decoder_layer(layer, hidden, cos, sin)
if layer < len(deepstack):
merged = np.asarray(hidden, dtype=np.float32)
merged[visual_rows] += deepstack[layer]
hidden = mx.array(merged)
# Per-layer sync keeps the 50-layer activation graph bounded.
mx.eval(hidden)
gc.collect()
return np.asarray(hidden).astype(np.float32), np.asarray(token_tags, dtype=np.int64)</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_conditioner.StreamedMiniMaxH3TextConditioner.close
close() -> None
Source
def close(self) -> None:
self._vision = None
self.index.close()
fastvideo.mlx_runtime.minimax_h3_conditioner.unswizzle_nvfp4_scales
unswizzle_nvfp4_scales(scale: np.ndarray, rows: int, cols: int) -> np.ndarray
FlashInfer 128x4 scale bytes -> MLX row-major group-16 scale bytes.
Source
def unswizzle_nvfp4_scales(scale: np.ndarray, rows: int, cols: int) -> np.ndarray:
"""FlashInfer 128x4 scale bytes -> MLX row-major group-16 scale bytes."""
pad_rows, pad_cols = -(-rows // 128) * 128, -(-cols // 4) * 4
if scale.size < pad_rows * pad_cols:
raise ValueError(f"NVFP4 scales need {pad_rows * pad_cols} bytes, got {scale.size}.")
tiles = scale.reshape(-1)[:pad_rows * pad_cols].reshape(pad_rows // 128, pad_cols // 4, 32, 4, 4)
return np.ascontiguousarray(tiles.transpose(0, 3, 2, 1, 4).reshape(pad_rows, pad_cols)[:rows, :cols])
fastvideo.mlx_runtime.minimax_h3_conditioner.export_mlx_h3_nvfp4_encoder
export_mlx_h3_nvfp4_encoder(component_dir: str | Path, output_dir: str | Path) -> Path
Cache the released packed encoder in MLX layout, without requantization.
Keep original packed nibbles, row-major scale bytes, global scales and
embedding/norm values. Later loads skip CPU scale unswizzling and staging.
Source
def export_mlx_h3_nvfp4_encoder(component_dir: str | Path, output_dir: str | Path) -> Path:
"""Cache the released packed encoder in MLX layout, without requantization.
Keep original packed nibbles, row-major scale bytes, global scales and
embedding/norm values. Later loads skip CPU scale unswizzling and staging.
"""
raw = _read_nvfp4_encoder_config(component_dir)
final_dir = Path(output_dir)
if final_dir.exists() and any(final_dir.iterdir()):
raise FileExistsError(f"Encoder cache output must be empty: {final_dir}")
# Build in a sibling staging directory and publish it with one rename, so an
# out-of-memory or out-of-disk failure never leaves a partial cache behind.
final_dir.parent.mkdir(parents=True, exist_ok=True)
output_dir = Path(tempfile.mkdtemp(prefix=f".{final_dir.name}.partial-", dir=final_dir.parent))
index = _ResidentNVFP4Index(_ShardIndex(Path(component_dir)))
published = False
try:
arrays = {}
matrices = {}
dense_keys = []
for key, value in index.weights.items():
if isinstance(value, NVFP4Matrix):
arrays[key] = value.weight
arrays[key + ".scales"] = value.scales
matrices[key] = {"global_scale": value.global_scale}
else:
arrays[key] = value
dense_keys.append(key)
mx.save_safetensors(str(output_dir / "model.safetensors"), arrays)
(output_dir / "config.json").write_text(json.dumps(raw, indent=2) + "\n")
manifest = {
"format_version": _NVFP4_CACHE_FORMAT_VERSION,
"language_layers": TEXT_ENCODER_LAYER,
"matrices": matrices,
"dense_keys": dense_keys,
"source_dir": str(Path(component_dir).resolve())
}
(output_dir / MLX_NVFP4_ENCODER_MANIFEST).write_text(json.dumps(manifest, indent=2) + "\n")
umask = os.umask(0)
os.umask(umask)
output_dir.chmod(0o777 & ~umask) # mkdtemp is owner-only; shared model dirs need the umask mode
if final_dir.exists():
final_dir.rmdir() # empty, checked above
output_dir.rename(final_dir)
published = True
finally:
index.close()
if not published:
shutil.rmtree(output_dir, ignore_errors=True)
return final_dir</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_conditioner.ResidentNVFP4MiniMaxH3TextConditioner
ResidentNVFP4MiniMaxH3TextConditioner(component_dir: str | Path, tokenizer_dir: str | Path | None = None)
Released NVFP4 encoder weights with floating-point MLX activations.
The packed weights and embedding table stay resident. CUDA quantizes
activations to FP4; this path keeps FP32 activations, so hidden states are
not expected to be bit-exact with the CUDA encoder.
Source
class ResidentNVFP4MiniMaxH3TextConditioner(StreamedMiniMaxH3TextConditioner):
"""Released NVFP4 encoder weights with floating-point MLX activations.
The packed weights and embedding table stay resident. CUDA quantizes
activations to FP4; this path keeps FP32 activations, so hidden states are
not expected to be bit-exact with the CUDA encoder.
"""
def __init__(self, component_dir: str | Path, tokenizer_dir: str | Path | None = None):
_read_nvfp4_encoder_config(component_dir)
# Fail on an older MLX before reading the encoder's large shards.
try:
packed, scales = mx.quantize(mx.ones((1, 64)), mode="nvfp4")
mx.eval(mx.quantized_matmul(mx.ones((1, 64)), packed, scales, mode="nvfp4"))
except (ValueError, RuntimeError) as error:
raise RuntimeError("Native NVFP4 conditioning requires an MLX build with nvfp4 matmul support.") from error
super().__init__(component_dir, tokenizer_dir)
if (Path(component_dir) / MLX_NVFP4_ENCODER_MANIFEST).exists():
self.index.close()
self.index = _ResidentNVFP4Index.from_mlx_checkpoint(component_dir)
else:
self.index = _ResidentNVFP4Index(self.index)
def _embed_tokens(self, token_ids: list[int]):
table = self.index.get_mlx("model.language_model.embed_tokens.weight")
return table[mx.array(token_ids, dtype=mx.int32)].astype(mx.float32)</code></pre></details>