multimodalart / minimax-h3-aoti

MiniMax-H3 AoTI packages

AOT-Inductor packages for one MiniMaxH3TransformerBlock , keyed <width>/torch<X.Y>/sm<cc>/<shape> . One package serves all 50 blocks of a transformer: it carries no weights , and spaces.aoti_patch binds each block's own live state_dict() on its first forward. Built and consumed by h3_aoti.py in the MiniMax-H3 Spaces; built off-Space by job_bf16_aoti.py on an rtx-pro-6000 Job.

What is here

bf16/torch2.11/sm120/dynamic is the one that matters: a dynamic sequence dimension , so it serves every canvas, every duration and every prompt length. The <HxWxF> keys hold byte-identical copies of it, for a Space configured to look artifacts up per canvas. _inputs/ holds cached conditioner output ( prompt_embeds + text_token_tags ) so a build job never has to load the 62 GiB Qwen3-VL. _reports/ holds each build's full JSON report.

A dynamic package is not a nicety. build_packed_sequence pads nothing, so S = num_text_tokens + condition_rows + audio_rows + video_rows moves with the prompt : the same canvas is S = 37726 for a 16-token prompt and S = 37728 for an 18-token one. A static package serves exactly one prompt length.

Measured, unquantized bfloat16, 124 frames, everything resident on one RTX PRO 6000 Blackwell

The absolute saving is near-constant at ~0.5 s/step, which is what AoTI can remove: 50 blocks of kernel-launch overhead plus the norm / rotary / AdaLN-gather epilogues. It cannot touch the matmuls, and at S = 37726 one block is ~70 TFLOP of GEMM and attention, so the released 768x1344 canvas is compute bound. The smaller the canvas, the more this pays.

Proven end to end on the pool: the MiniMax-H3 Space loads bf16/torch2.11/sm120/dynamic at startup, patches all 50 blocks and generates.

The trap, recorded because the symptom is a segfault with no message

spaces.zero.torch.aoti.LazyAOTIModel binds a package's constants by name , intersecting the block's state_dict() with get_constant_fqns() , and silently keeps whatever it cannot match. A shallow clone ( _shallow_clone_module , which aoti_patch uses and which the export side was mirroring) exported in torch.export 's default non-strict mode lifts every weight twice — once as a named PARAMETER , once as an anonymous lifted_tensor_<N> CONSTANT_TENSOR sharing the same data_ptr() . The anonymous half binds to nothing and the compiled kernel reads pointers nobody set. Measured on the real block, at full size, in plain bfloat16 with no tensor subclass anywhere:

So it is neither accelerate's offload hooks nor torchao's tensor subclasses, which were both blamed first. Export the live block when it has nothing to flatten; use strict=True when a subclass forces the clone.

temb is padded to a fixed 4 rows

temb is (num_distinct_timesteps, time_embed_dim) — one row at step 0, where video and audio still share a noise level, two from step 1 where their sigma schedules diverge — and the block gathers from a 3 * rows AdaLN table, so the row count is part of the compiled shape. Padding to a fixed 4 rows makes it constant; rows past the live ones are never gathered, so the output is unchanged. Validated rather than assumed: the build replays a real 1-row and a real 2-row call through the compiled block and diffs both against eager.