MiniMax-H3-Image: fused blocks and tabulated AdaLN modulation (0.18.0)

README.md

@@ -30,7 +30,8 @@ The implementation is clean-room and Apache-2.0 licensed.
 - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full
   dequantized weight matrix.
 - Fused transformer checkpoints for Krea 2, FLUX.2, Ideogram 4, Qwen-Image 2.1,
-  MiniMax-H3, Boogu-Image, FLUX.3 Klein, MOSS-Image, Fibo 2.0 and Vimoo 1.0
+  MiniMax-H3, MiniMax-H3-Image, Boogu-Image, FLUX.3 Klein, MOSS-Image, Fibo 2.0 and
+  Vimoo 1.0
   (`orbitquant.fused`): each block runs
   grouped INT8 GEMMs with norm/modulation prologues and SwiGLU/gated-residual
   epilogues plus an INT8 attention kernel, and the checkpoint stores only the
@@ -632,6 +633,7 @@ the same for an `orbitquant-v1` component artifact, and
 | Ideogram 4 | `Ideogram4Transformer2DModel` | sandwich-norm blocks with row-indexed modulation |
 | Qwen-Image 2.1 | `QwenImage21Transformer2DModel` | block-causal prefill and KV-cached decode |
 | MiniMax-H3 | `MiniMaxH3Transformer3DModel` | per-row (timestep, modality) AdaLN tables, partial rotate-half RoPE |
+| MiniMax-H3-Image | `MiniMaxH3ImageTransformer2DModel` | the MiniMax-H3 blocks; `minimax_h3.tabulate_modulation(model, timesteps)` replaces the per-block AdaLN projections and the timestep MLP with tables for a fixed timestep set (the six-step schedule drops 13.0B of the 33.1B parameters); a timestep outside the set raises `ValueError` |
 | Boogu-Image | `BooguImageTransformer2DModel` | refiner, single- and double-stream blocks; 8-bit activations |
 | FLUX.3 Klein | `flux3.dit.DiT` (Black Forest Labs' `flux3` package) | per-mode and single-stream blocks: Q\|K\|V, SwiGLU, gated residuals; one sample at a time |
 | MOSS-Image | `MOSSImageTransformer` (the model's `pipeline_moss_image.py`) | joint image/text blocks: RMSNorm and shared AdaLN prologues, Q\|K\|V, SiTU gate\|up, gated residuals; one sample at a time |

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "orbitquant"
-version = "0.17.1"
+version = "0.18.0"
 description = "Calibration-free OrbitQuant for transformer linear projections"
 readme = "README.md"
 requires-python = ">=3.11"

src/orbitquant/__init__.py

@@ -24,7 +24,7 @@
 from orbitquant.quantizer import register_hf_quantizers
 from orbitquant.recipes import recipe
 
-__version__ = "0.17.1"
+__version__ = "0.18.0"
 
 register_hf_quantizers()
 

src/orbitquant/fused/families/minimax_h3.py

@@ -1,26 +1,39 @@
-"""MiniMax-H3 (``MiniMaxH3Transformer3DModel``) main blocks.
+"""MiniMax-H3 (``MiniMaxH3Transformer3DModel``) and MiniMax-H3-Image
+(``MiniMaxH3ImageTransformer2DModel``) main blocks.
 
 One packed video/audio/text sequence; the AdaLN modulation of every row is a row of a
 (timestep, modality) table. The prologues and residual epilogues read that table through the
 per-row index directly instead of materializing ``[seq, hidden]`` modulation tensors. Q/K get
 RMSNorm and a rotate-half RoPE on the leading 96 channels of each 128-wide head. The token
 refiner blocks (text only, once per forward) keep their modules.
 
+MiniMax-H3-Image samples a fixed grid of six timesteps (plus one noise level for reference rows),
+and its AdaLN projections read nothing but the timestep embedding. ``tabulate_modulation`` stores
+their outputs for those timesteps (``modulation.json`` lists them) and removes the projections:
+13.0 of the 33.1 billion parameters. ``install`` then runs the model forward with table lookups;
+a timestep outside the table raises an error.
+
 Padding rows form their own attention document (the model's mask pairs live rows with live rows
 and padding rows with padding rows); the two documents run as separate mask-free attentions.
 """
 
 from __future__ import annotations
 
+import types
+
 import torch
+import torch.nn.functional as F
+from torch import nn
 
 from orbitquant.fused.families import int8_attention_allowed, specs
 from orbitquant.fused.groups import RowSource
 from orbitquant.kernels import triton_activation as act
 from orbitquant.kernels import triton_int8_gemm as gemm
 
 FAMILY = "minimax_h3"
-MODEL_CLASSES = ("MiniMaxH3Transformer3DModel",)
+MODEL_CLASSES = ("MiniMaxH3Transformer3DModel", "MiniMaxH3ImageTransformer2DModel")
+# Rows per timestep in the AdaLN tables: video/image, text, audio.
+MODALITIES = 3
 
 
 def blocks(model):
@@ -143,3 +156,141 @@ def _forward(self, hidden_states, temb, adaln_indices, rotary_emb, attention_mas
         out=h,
     )
     return h[None]
+
+
+class ModulationTable(nn.Module):
+    """The outputs of one block's ``adaln_proj`` for a fixed set of timesteps.
+
+    ``table`` holds ``adaln_proj.linear(silu(temb))`` reshaped to ``(timesteps * MODALITIES,
+    6 * hidden)``, the row layout the block indexes with ``timestep_index * MODALITIES + tag``.
+    ``forward`` takes the table row of each distinct timestep of the call."""
+
+    def __init__(self, timesteps: int, width: int, dtype=torch.bfloat16, device=None):
+        super().__init__()
+        table = torch.empty(timesteps * MODALITIES, width, dtype=dtype, device=device)
+        self.register_buffer("table", table)
+
+    def forward(self, rows: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        tags = torch.arange(MODALITIES, device=rows.device)
+        index = (rows[:, None] * MODALITIES + tags).reshape(-1)
+        return self.table.index_select(0, index).chunk(6, dim=-1)
+
+
+class ModulationKeys(nn.Module):
+    """The tabulated timesteps and the output norm's shift/scale for each of them."""
+
+    def __init__(self, timesteps: int, hidden: int, dtype=torch.bfloat16, device=None):
+        super().__init__()
+        keys = torch.empty(timesteps, dtype=torch.float32, device=device)
+        self.register_buffer("timesteps", keys)
+        output = torch.empty(timesteps, 2 * hidden, dtype=dtype, device=device)
+        self.register_buffer("output", output)
+
+    def rows(self, timestep: torch.Tensor) -> torch.Tensor:
+        """Table row of every timestep (one host sync)."""
+        timestep = timestep.to(device=self.timesteps.device, dtype=torch.float32)
+        match = timestep[:, None] == self.timesteps[None]
+        if not bool(match.any(dim=1).all()):
+            raise ValueError(
+                f"timesteps {timestep.tolist()} are not all in the modulation table "
+                f"{self.timesteps.tolist()}; this checkpoint supports the six-step sampling grid "
+                "of MiniMax-H3-Image only"
+            )
+        return match.to(torch.int32).argmax(dim=1)
+
+
+def add_modulation_tables(model, timesteps: int, device=None) -> None:
+    """Replace the AdaLN projections and the timestep MLP by empty tables (before a tabulated
+    checkpoint loads)."""
+    hidden = model.config.hidden_size
+    dtype = model.transformer_blocks[0].adaln_proj.linear.weight.dtype
+    if dtype not in (torch.bfloat16, torch.float16, torch.float32):
+        dtype = torch.bfloat16
+    for block in model.transformer_blocks:
+        block.adaln_proj = ModulationTable(timesteps, 6 * hidden, dtype=dtype, device=device)
+    model.oq_modulation = ModulationKeys(timesteps, hidden, dtype=dtype, device=device)
+    del model.time_embedder
+    del model.time_proj
+    del model.norm_out.linear
+
+
+@torch.no_grad()
+def tabulate_modulation(model, timesteps) -> list[float]:
+    """Compute the AdaLN tables of ``timesteps`` (unscaled, as the pipeline passes them) from the
+    loaded weights and remove the projections. Each timestep is projected on its own, as a
+    text-to-image call projects its one timestep."""
+    from diffusers.models.modeling_utils import get_parameter_dtype
+
+    keys = sorted({float(torch.tensor(t, dtype=torch.float32)) for t in timesteps})
+    device = model.transformer_blocks[0].adaln_proj.linear.weight.device
+    embedder_dtype = get_parameter_dtype(model.time_embedder)
+    tables = [[] for _ in model.transformer_blocks]
+    output = []
+    for t in keys:
+        timestep = torch.tensor([t], dtype=torch.float32, device=device)
+        temb = model.time_embedder(model.time_proj(timestep).to(embedder_dtype))
+        for i, block in enumerate(model.transformer_blocks):
+            shift_scale = block.adaln_proj(temb)
+            tables[i].append(torch.cat(shift_scale, dim=-1))
+        linear = model.norm_out.linear
+        output.append(linear(F.silu(temb).to(linear.weight.dtype)))
+    dtype = tables[0][0].dtype
+    add_modulation_tables(model, len(keys), device=device)
+    for block, rows in zip(model.transformer_blocks, tables, strict=True):
+        block.adaln_proj.table.copy_(torch.cat(rows, dim=0))
+    model.oq_modulation.timesteps.copy_(torch.tensor(keys, dtype=torch.float32))
+    model.oq_modulation.output.copy_(torch.cat(output, dim=0).to(dtype))
+    install(model)
+    return keys
+
+
+def install(model):
+    image_model = type(model).__name__ == "MiniMaxH3ImageTransformer2DModel"
+    if image_model and hasattr(model, "oq_modulation"):
+        model.forward = types.MethodType(_image_forward, model)
+
+
+def _image_forward(
+    self,
+    hidden_states,
+    encoder_hidden_states,
+    timestep,
+    timestep_indices,
+    token_tags,
+    position_ids,
+    image_indices,
+    text_indices,
+    pdd_step_index,
+    return_dict=True,
+):
+    """``MiniMaxH3ImageTransformer2DModel.forward`` with tabulated modulation."""
+    from diffusers.models.modeling_outputs import Transformer2DModelOutput
+    from diffusers.models.modeling_utils import get_parameter_dtype
+
+    rotary_emb = self.rope(position_ids)
+    image_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))
+    text_embeds = self.token_refiner(
+        self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
+    )
+    sequence_length = position_ids.shape[0]
+    batch, width = text_embeds.shape[0], text_embeds.shape[-1]
+    hidden_states = text_embeds.new_zeros((batch, sequence_length, width))
+    hidden_states = hidden_states.index_copy(1, text_indices, text_embeds)
+    hidden_states = hidden_states.index_copy(1, image_indices, image_embeds.to(text_embeds.dtype))
+
+    rows = self.oq_modulation.rows(timestep)
+    adaln_indices = timestep_indices * MODALITIES + token_tags
+    for block in self.transformer_blocks:
+        hidden_states = block(hidden_states, rows, adaln_indices, rotary_emb)
+
+    proj_out = self.pdd_proj_out[pdd_step_index]
+    shift, scale = self.oq_modulation.output.index_select(0, rows).chunk(2, dim=-1)
+    hidden_states = self.norm_out.norm(hidden_states)
+    scale = scale.index_select(0, timestep_indices)
+    shift = shift.index_select(0, timestep_indices)
+    hidden_states = hidden_states * (1.0 + scale) + shift
+    output = proj_out(hidden_states.to(get_parameter_dtype(proj_out)))
+    output = output.index_select(1, image_indices)
+    if not return_dict:
+        return (output,)
+    return Transformer2DModelOutput(sample=output)

src/orbitquant/fused/linear.py

@@ -47,9 +47,15 @@ def fuse_linears(model: nn.Module, *, activation_bits: int = 8, config=None) ->
 
     runtime = FusedRuntime.from_config(_quant_config(model, config), attention="sdpa")
     count = 0
-    for name, module in list(model.named_modules()):
-        if not isinstance(module, OrbitQuantLinear) or module.packed_weight_indices is None:
-            continue
+    # Names, not modules: a list of the modules would keep every replaced projection's packed
+    # weights alive until the end, so the device would hold the model twice.
+    names = [
+        name
+        for name, module in model.named_modules()
+        if isinstance(module, OrbitQuantLinear) and module.packed_weight_indices is not None
+    ]
+    for name in names:
+        module = model.get_submodule(name)
         holder = nn.Module()
         holder.proj = module
         group = PackedGroup.from_sources(

tests/test_minimax_h3_image.py

@@ -0,0 +1,87 @@
+"""MiniMax-H3-Image: tabulated AdaLN modulation reproduces the original transformer forward."""
+
+import pytest
+import torch
+
+diffusers = pytest.importorskip("diffusers")
+if not hasattr(diffusers, "MiniMaxH3ImageTransformer2DModel"):
+    pytest.skip("diffusers without MiniMaxH3ImageTransformer2DModel", allow_module_level=True)
+
+from orbitquant.fused.families import minimax_h3 as family  # noqa: E402
+
+TIMESTEPS = [0.0, 0.25, 0.5, 0.999]
+
+
+def _model():
+    torch.manual_seed(0)
+    return diffusers.MiniMaxH3ImageTransformer2DModel(
+        num_attention_heads=2,
+        attention_head_dim=32,
+        hidden_size=48,
+        num_layers=2,
+        num_refiner_layers=1,
+        ffn_dim=64,
+        in_channels=4,
+        text_dim=16,
+        freq_dim=32,
+        time_embed_hidden_dim=48,
+        time_embed_dim=24,
+        rope_freq_dim=4,
+        num_pdd_heads=2,
+    ).eval()
+
+
+def _inputs(timestep):
+    image_rows, text_rows = 8, 3
+    seq = image_rows + text_rows
+    grid = torch.meshgrid(torch.arange(2), torch.arange(4), indexing="ij")
+    grid = torch.stack(grid, -1).reshape(-1, 2)
+    positions = torch.zeros(seq, 3, dtype=torch.long)
+    positions[text_rows:, 1:] = grid
+    positions[:text_rows, 0] = torch.arange(text_rows)
+    return dict(
+        hidden_states=torch.randn(1, image_rows, 16),
+        encoder_hidden_states=torch.randn(1, text_rows, 16),
+        timestep=torch.tensor(timestep, dtype=torch.float32),
+        timestep_indices=torch.tensor([0] * text_rows + [len(timestep) - 1] * image_rows),
+        token_tags=torch.tensor([1] * text_rows + [0] * image_rows),
+        position_ids=positions,
+        image_indices=torch.arange(text_rows, seq),
+        text_indices=torch.arange(text_rows),
+        pdd_step_index=1,
+        return_dict=False,
+    )
+
+
+@torch.no_grad()
+def test_tabulated_forward_matches_original():
+    model = _model()
+    calls = [_inputs([0.25]), _inputs([0.5, 0.999])]
+    reference = [model(**kw)[0] for kw in calls]
+    keys = family.tabulate_modulation(model, TIMESTEPS)
+    assert keys == torch.tensor(sorted(TIMESTEPS), dtype=torch.float32).tolist()
+    assert not hasattr(model, "time_embedder")
+    assert all(isinstance(b.adaln_proj, family.ModulationTable) for b in model.transformer_blocks)
+    for kw, ref in zip(calls, reference, strict=True):
+        torch.testing.assert_close(model(**kw)[0], ref, rtol=1e-5, atol=1e-5)
+
+
+@torch.no_grad()
+def test_tables_load_into_empty_skeleton():
+    model = _model()
+    family.tabulate_modulation(model, TIMESTEPS)
+    state = model.state_dict()
+    fresh = _model()
+    family.add_modulation_tables(fresh, len(TIMESTEPS))
+    fresh.load_state_dict(state, strict=True)
+    family.install(fresh)
+    kw = _inputs([0.5])
+    torch.testing.assert_close(fresh(**kw)[0], model(**kw)[0])
+
+
+@torch.no_grad()
+def test_unknown_timestep_raises():
+    model = _model()
+    family.tabulate_modulation(model, TIMESTEPS)
+    with pytest.raises(ValueError, match="modulation table"):
+        model(**_inputs([0.3]))