[perf] Speed up exact MiniMax H3 MLX inference

[aryan5v]

Purpose

Uncached MiniMax H3 generation on Apple Silicon spends 79–82 seconds streaming and expanding the 66 GB BF16 conditioner before denoising. The H3 transformer also repeats row-expanded AdaLN scale and shift operations in every block.

This PR reduces both costs without changing the checkpoint, quantization, sampling schedule, spatial or temporal workload, conditioning values, or decoder path.

Changes

• Read each requested conditioner BF16 tensor with one bounded file operation, reinterpret its storage in MLX, and materialize the FP32 conversion before building the decoder-layer graph. Embedding access remains row-scoped, and only one language-model layer is resident at a time.
• Apply gathered AdaLN scale and shift values through one reusable shapeless mx.compile graph.
• Add storage-bit, malformed-read, source-ownership, dtype, stride, and eager-parity tests for both paths.

The module documentation and comments record the streaming memory contract and explain why the BF16 cast is synchronized before layer evaluation. No user-facing option or cache format changes.

Measured results

Environment: Apple M4 Max with 36 GB unified memory, macOS 27.0 build 26A5416b. Tests used the affine INT6 group-64 H3 checkpoint at 832×480, 124 frames, four denoising steps, and no spatial or temporal fast mode.

Conditioner phase

The same 70-token speech prompt and layer-50 output were compared within each MLX runtime.

MLX	Original	This PR	Reduction	Parity
0.31.2	78.509 s	15.228 s	80.6%	Hidden features and tags exact
0.32.2	82.462 s	15.569 s	81.1%	Hidden features and tags exact

The conditioner MLX active-memory peak increases from about 1.88 GiB to 2.06 GiB. It remains well below the approximately 19.6 GiB denoise peak and does not load a full shard, embedding table, or conditioner model.

AdaLN modulation

Each real-block result is the median of eight alternating baseline/candidate pairs at packed sequence length 14,937.

MLX	Eager	Compiled	Reduction	Parity
0.31.2	1.843257 s	1.796039 s	2.56%	Exact
0.32.2	1.799464 s	1.772093 s	1.52%	Exact

Two order-reversed four-step comparisons also favored the compiled path: 393.597 versus 443.541 seconds and 366.016 versus 392.436 seconds. Those runs paged heavily, so the 6.7–11.3% full-denoise range is directional evidence; the isolated real-block medians are the stable performance claim.

A complete uncached seed-2027 generation took 524.062 seconds: 17.163 conditioning, 396.085 denoise, 109.730 video decode, 0.696 audio decode, and 0.389 mux. It produced a 5.17-second 832×480 H.264 clip with stereo AAC audio. An eager replay produced bit-identical final video and audio latents.

Test Plan

/Users/aryank/fvrelease/.venv/bin/python -m pytest \
  fastvideo/tests/mlx/test_mlx_minimax_h3_conditioner.py \
  fastvideo/tests/mlx/test_mlx_minimax_h3_weight_reads.py \
  fastvideo/tests/mlx/test_mlx_minimax_h3_modulation.py \
  fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py -q

PYTHONPATH=/Users/aryank/os/FastVideoH3YukonPerf/outputs/yukon/mlx0322 \
  /Users/aryank/fvrelease/.venv/bin/python -m pytest \
  fastvideo/tests/mlx/test_mlx_minimax_h3_conditioner.py \
  fastvideo/tests/mlx/test_mlx_minimax_h3_weight_reads.py \
  fastvideo/tests/mlx/test_mlx_minimax_h3_modulation.py \
  fastvideo/tests/mlx/test_mlx_minimax_h3_parity.py -q

/Users/aryank/os/FastVideo/.venv/bin/pre-commit run --all-files

Manual full-workload validation used seeds 2026 and 2027 with 832×480, 124 frames, four steps, and the unchanged audio/video decoders.

Test Results

• MLX 0.31.2: 51 passed, 14 Torch deprecation warnings.
• MLX 0.32.2: 51 passed, 14 Torch deprecation warnings.
• Configured pre-commit: all hooks passed across the repository.
• Seed 2026: all four video velocities, all four audio velocities, and both final latents match the eager path bit-for-bit in both run orders.
• Seed 2027: final video and audio latents match the eager path bit-for-bit.
• Generated media: container, stream dimensions, duration, frame finiteness, audio finiteness, clipping, and representative frames checked.
• Not run: the four-GPU CUDA MiniMax H3 SSIM test. It does not exercise the MLX runtime changed here; the MLX comparisons instead gate exact conditioner, block, per-step, and final-latent values.

Review notes

The conditioner gain applies to uncached prompts. Cached prompts bypass that phase. The compiled AdaLN graph is lazy and process-local, and it preserves the original multiply-then-add expression and output dtype.

The branch is based on 620bc36dc and includes only these two performance commits. PR #1788's wide-matrix affine INT6 dispatch and PR #1789's optional spatial fast mode are already present in main and are not duplicated here.

Checklist

• I ran pre-commit run --all-files and fixed all issues
• I added or updated tests for my changes
• I updated documentation if needed
• I considered GPU memory impact of my changes

For model/pipeline changes:

• I verified SSIM regression tests pass (not run: the existing test is a four-GPU CUDA path and does not exercise MLX)
• I confirmed no support-matrix change is required because this PR adds no model

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]