[diffusion] feat: rebuild minimax-h3 adaln outputs on demand (#34650)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -1,5 +1,6 @@
 import copy
 import logging
+from collections.abc import Callable
 from contextlib import nullcontext
 from typing import Any
 
@@ -50,6 +51,10 @@ def _resolve_checkpoint_load_device(
     return runtime_device
 
 
+def _minimax_h3_adaln_cache_key_filter(name: str) -> bool:
+    return ".adaln_proj.linear." not in name
+
+
 def _default_quantized_attention_backend(
     quant_spec: TransformerQuantLoadSpec, server_args: ServerArgs
 ) -> AttentionBackendEnum | None:
@@ -203,6 +208,40 @@ class TransformerLoader(ComponentLoader):
             "hf_config": config,
             "quant_config": quant_spec.runtime_quant_config,
         }
+        checkpoint_key_filter: Callable[[str], bool] | None = None
+        adaln_cache_path = component_server_args.minimax_h3_adaln_cache_path
+        if adaln_cache_path is not None:
+            if cls_name != "MiniMaxH3DiTModel":
+                raise ValueError(
+                    "--minimax-h3-adaln-cache-path is only supported by MiniMax H3"
+                )
+            if component_server_args.model_variant not in ("fl2va", "ref2va"):
+                raise ValueError(
+                    "MiniMax H3 AdaLN cache requires --model-variant fl2va or ref2va"
+                )
+            init_params["adaln_cache_path"] = adaln_cache_path
+            init_params["adaln_cache_model_variant"] = (
+                component_server_args.model_variant
+            )
+            checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter
+        if component_server_args.minimax_h3_adaln_online:
+            if cls_name != "MiniMaxH3DiTModel":
+                raise ValueError(
+                    "--minimax-h3-adaln-online is only supported by MiniMax H3"
+                )
+            if adaln_cache_path is not None:
+                raise ValueError(
+                    "--minimax-h3-adaln-online and --minimax-h3-adaln-cache-path "
+                    "are mutually exclusive"
+                )
+            # Keep the weights off-device; the model rebuilds the AdaLN
+            # outputs from the checkpoint for each request's timestep plan.
+            init_params["adaln_weight_files"] = safetensors_list
+            init_params["adaln_plan_width"] = (
+                component_server_args.minimax_h3_adaln_plan_width
+            )
+            checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter
+
         if (
             init_params["quant_config"] is None
             and component_server_args.transformer_weights_path is not None
@@ -273,6 +312,7 @@ class TransformerLoader(ComponentLoader):
                 output_dtype=None,
                 strict=False,
                 weight_load_plan=weight_load_plan,
+                checkpoint_key_filter=checkpoint_key_filter,
             )
 
         # post-hooks (e.g., patch scales (nunchaku))

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -237,6 +237,7 @@ def maybe_load_fsdp_model(
     pin_cpu_memory: bool = True,
     strict: bool = True,
     weight_load_plan: WeightLoadPlan | None = None,
+    checkpoint_key_filter: Callable[[str], bool] | None = None,
 ) -> torch.nn.Module:
     """Load a model with optional FSDP (Fully Sharded Data Parallel) support.
 
@@ -347,6 +348,7 @@ def maybe_load_fsdp_model(
         and use_fsdp
         and weight_dir_list
         and preprocess_loaded_state_dict is None
+        and checkpoint_key_filter is None
         and not is_bnb_quantized
     ):
         preconverted_state_dict = (
@@ -361,6 +363,7 @@ def maybe_load_fsdp_model(
         and not use_fsdp
         and weight_dir_list
         and preprocess_loaded_state_dict is None
+        and checkpoint_key_filter is None
         and not is_bnb_quantized
     ):
         preconverted_state_dict = (
@@ -375,10 +378,14 @@ def maybe_load_fsdp_model(
         if weight_load_plan.load_full_state_dict_on_device:
             weight_iterator = safetensors_weights_iterator(
                 weight_dir_list,
+                key_filter=checkpoint_key_filter,
                 weight_load_plan=weight_load_plan,
             )
         else:
-            weight_iterator = safetensors_weights_iterator(weight_dir_list)
+            weight_iterator = safetensors_weights_iterator(
+                weight_dir_list,
+                key_filter=checkpoint_key_filter,
+            )
         if preprocess_loaded_state_dict is not None:
             weight_iterator = preprocess_loaded_state_dict(weight_iterator)
         if is_bnb_quantized:

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -228,6 +228,10 @@ def safetensors_weights_iterator(
         use_runai_model_streamer = (
             HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER
         )
+    if key_filter is not None:
+        # streamer filters after materializing all tensors, so it cannot skip
+        # a checkpoint partition at load time
+        use_runai_model_streamer = False
 
     # Validate files before loading
     corrupted_files, duplicate_files_by_key = _scan_safetensors_files(hf_weights_files)

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -8,10 +8,14 @@ contract accepts packed inference keyword arguments and returns packed logits.
 from __future__ import annotations
 
 import math
-from typing import Any
+import os
+import struct
+from contextlib import ExitStack
+from typing import Any, Callable
 
 import torch
 import torch.nn as nn
+from safetensors.torch import safe_open
 
 from sglang.kernels.ops.activation.activation import (
     silu_and_mul_with_activation_rounding_,
@@ -40,6 +44,7 @@ from sglang.multimodal_gen.runtime.distributed import (
 )
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_ring_ctx,
+    get_tp_rank,
     get_ulysses_ctx,
 )
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
@@ -64,10 +69,13 @@ from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
     eager_on_graph,
 )
 
+logger = init_logger(__name__)
+
 _ARCH_DEFAULTS = MiniMaxH3DiTArchConfig()
 _BF16_DTYPE = torch.bfloat16
 _FP32_DTYPE = torch.float32
@@ -802,6 +810,289 @@ class MiniMaxH3AdalnProj(nn.Module):
         return self.split_output(x)
 
 
+# A ref2va request carrying both a visual and an audio reference reaches four
+# distinct timesteps in one step: video, audio, the imgvid condition and the
+# audio reference. That is the widest case, so it is the default; a deployment
+# serving only narrower tasks (t2va reaches 2, fl2va 3) can shrink the slab
+# proportionally via --minimax-h3-adaln-plan-width.
+MINIMAX_H3_ADALN_MAX_PLAN_WIDTH = 4
+
+
+def _plan_key(timesteps: torch.Tensor) -> tuple[int, ...]:
+    """One denoise step's unique timesteps as their exact fp32 bit patterns."""
+    return tuple(
+        struct.unpack("<I", struct.pack("<f", float(value)))[0]
+        for value in timesteps.tolist()
+    )
+
+
+class MiniMaxH3AdalnCache(nn.Module):
+    """Precomputed AdaLN outputs for fixed FP32 timestep plans."""
+
+    _FORMAT_VERSION = "2"
+    plan_timesteps: torch.Tensor
+    plan_lengths: torch.Tensor
+    block_params: torch.Tensor
+    final_params: torch.Tensor
+
+    def __init__(
+        self,
+        arch: MiniMaxH3DiTArchConfig,
+        *,
+        path: str | None = None,
+        model_variant: str | None = None,
+        weight_files: list[str] | None = None,
+        max_plans: int = 64,
+        max_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH,
+    ) -> None:
+        super().__init__()
+        if (path is None) == (weight_files is None):
+            raise ValueError(
+                "MiniMax H3 AdaLN cache takes exactly one of path (prebuilt "
+                "sidecar) or weight_files (rebuild from the checkpoint)"
+            )
+        self.path = path
+        self.model_variant = model_variant
+        self.weight_files = weight_files
+        self.max_plans = max_plans
+        self.max_plan_width = max_plan_width
+        self.num_layers = arch.num_layers
+        self.hidden_size = arch.hidden_size
+        self.block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * arch.hidden_size
+        self.final_width = 2 * arch.hidden_size
+        # Rebuild path only: plan bit pattern -> slot, tracked on the host.
+        self._slots: dict[tuple[int, ...], int] = {}
+        self.rebuilds = 0
+
+    def load(self, device: torch.device) -> None:
+        if self.path is None:
+            self._allocate(device)
+            return
+        if not os.path.isfile(self.path):
+            raise ValueError(f"MiniMax H3 AdaLN cache does not exist: {self.path}")
+
+        with safe_open(self.path, framework="pt", device="cpu") as cache_file:
+            metadata = cache_file.metadata() or {}
+            if metadata.get("format_version") != self._FORMAT_VERSION:
+                raise ValueError(
+                    "MiniMax H3 AdaLN cache has an unsupported or missing format_version"
+                )
+            cache_variant = metadata.get("model_variant")
+            if self.model_variant is not None and cache_variant != self.model_variant:
+                raise ValueError(
+                    "MiniMax H3 AdaLN cache model_variant does not match the loaded "
+                    f"variant ({cache_variant!r} != {self.model_variant!r})"
+                )
+            plan_timesteps = cache_file.get_tensor("plan_timesteps")
+            plan_lengths = cache_file.get_tensor("plan_lengths")
+            block_params = cache_file.get_tensor("block_params")
+            final_params = cache_file.get_tensor("final_params")
+
+        expected_block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * self.hidden_size
+        expected_final_width = 2 * self.hidden_size
+        if (
+            plan_timesteps.dtype != _FP32_DTYPE
+            or plan_timesteps.ndim != 2
+            or plan_lengths.dtype != torch.int64
+            or plan_lengths.shape != (plan_timesteps.shape[0],)
+            or (plan_lengths < 1).any()
+            or (plan_lengths > plan_timesteps.shape[1]).any()
+        ):
+            raise ValueError("MiniMax H3 AdaLN cache has invalid timestep plans")
+        if block_params.dtype != _BF16_DTYPE or block_params.shape != (
+            plan_timesteps.shape[0],
+            plan_timesteps.shape[1],
+            self.num_layers,
+            expected_block_width,
+        ):
+            raise ValueError("MiniMax H3 AdaLN cache has invalid block_params")
+        if final_params.dtype != _BF16_DTYPE or final_params.shape != (
+            plan_timesteps.shape[0],
+            plan_timesteps.shape[1],
+            expected_final_width,
+        ):
+            raise ValueError("MiniMax H3 AdaLN cache has invalid final_params")
+
+        self.register_buffer("plan_timesteps", plan_timesteps.to(device))
+        self.register_buffer("plan_lengths", plan_lengths.to(device))
+        self.register_buffer("block_params", block_params.to(device))
+        self.register_buffer("final_params", final_params.to(device))
+
+    def _allocate(self, device: torch.device) -> None:
+        """Empty slab for the rebuild path; its pointers must never move.
+
+        ``plan_lengths`` starts at zero and that is what keeps unused slots out
+        of ``lookup``: a real plan always has at least one timestep, so a zero
+        length can never match. Breakable CUDA graph keys its replay signature
+        on tensor pointers, so this is allocated once and only written in place.
+        """
+        width = self.max_plan_width
+        self.register_buffer(
+            "plan_timesteps",
+            torch.zeros((self.max_plans, width), dtype=_FP32_DTYPE, device=device),
+        )
+        self.register_buffer(
+            "plan_lengths",
+            torch.zeros((self.max_plans,), dtype=torch.int64, device=device),
+        )
+        self.register_buffer(
+            "block_params",
+            torch.zeros(
+                (self.max_plans, width, self.num_layers, self.block_width),
+                dtype=_BF16_DTYPE,
+                device=device,
+            ),
+        )
+        self.register_buffer(
+            "final_params",
+            torch.zeros(
+                (self.max_plans, width, self.final_width),
+                dtype=_BF16_DTYPE,
+                device=device,
+            ),
+        )
+        logger.info(
+            "MiniMax H3 AdaLN rebuild slab: %d plans x %d timesteps = %.2f GiB",
+            self.max_plans,
+            width,
+            self.block_params.numel() * 2 / 2**30,
+        )
+
+    def build(
+        self,
+        step_timesteps: list[torch.Tensor],
+        *,
+        embed: Callable[[torch.Tensor], torch.Tensor],
+    ) -> None:
+        """Fill every plan this request will look up, in one streaming pass.
+
+        Each plan keeps its own timestep count as the GEMM batch size, because
+        cuBLAS selects kernels by shape and the selection is not monotonic in M:
+        against the runtime's M == 2, results at M == 4/8/16/64/96 are
+        bit-identical while M == 32 differs in 11760 of 96768 elements and
+        M == 1 (the GEMV path the first denoise step takes) differs in 69.
+        Rebuilding a plan at any other batch size silently perturbs the output.
+
+        The pass reads all 50 adaln_proj layers regardless of how many plans are
+        missing, so a request builds everything it needs before denoising rather
+        than filling in step by step.
+        """
+        wanted: dict[tuple[int, ...], torch.Tensor] = {}
+        for timesteps in step_timesteps:
+            wanted.setdefault(_plan_key(timesteps), timesteps)
+        missing = {k: v for k, v in wanted.items() if k not in self._slots}
+        if not missing:
+            return
+        if len(self._slots) + len(missing) > self.max_plans:
+            # Every plan a request looks up has to stay resident for the whole
+            # denoise loop, so an overflow means the capacity is too small --
+            # evicting part of it would only move the failure into lookup().
+            self._slots.clear()
+            self.plan_lengths.zero_()
+        if len(missing) > self.max_plans:
+            raise ValueError(
+                f"MiniMax H3 AdaLN rebuild needs {len(missing)} plans but "
+                f"max_plans is {self.max_plans}"
+            )
+        widest = max(timesteps.numel() for timesteps in missing.values())
+        if widest > self.max_plan_width:
+            raise ValueError(
+                f"MiniMax H3 AdaLN rebuild hit a {widest}-timestep plan but the "
+                f"slab was allocated for {self.max_plan_width}; raise "
+                "--minimax-h3-adaln-plan-width (t2va needs 2, fl2va 3, ref2va 4)"
+            )
+
+        device = self.block_params.device
+        slots = []
+        for key, timesteps in missing.items():
+            slot = len(self._slots)
+            self._slots[key] = slot
+            slots.append((slot, timesteps.numel(), embed(timesteps.to(device))))
+            self.plan_timesteps[slot, : timesteps.numel()] = timesteps.to(device)
+
+        # adaln_proj is a ColumnParallelLinear: each rank owns a slice of the
+        # output features and all-gathers afterwards. The rebuild has to do the
+        # same rather than read the full width in one go -- a sharded GEMM has a
+        # different N, so cuBLAS picks a different kernel and the outputs stop
+        # matching. It also cuts per-rank checkpoint reads to 1/tp.
+        tp_size = get_tp_world_size()
+        tp_rank = get_tp_rank() if tp_size > 1 else 0
+
+        with ExitStack() as stack:
+            handles = [
+                stack.enter_context(safe_open(f, framework="pt", device=str(device)))
+                for f in self.weight_files
+            ]
+            index = {name: h for h in handles for name in h.keys()}
+
+            def read_shard(name: str, out_features: int) -> torch.Tensor:
+                if tp_size == 1:
+                    return index[name].get_tensor(name)
+                shard = out_features // tp_size
+                start = tp_rank * shard
+                return index[name].get_slice(name)[start : start + shard]
+
+            def project(adaln_input: torch.Tensor, weight, bias) -> torch.Tensor:
+                out = nn.functional.linear(adaln_input, weight, bias)
+                return tensor_model_parallel_all_gather(out) if tp_size > 1 else out
+
+            for layer in range(self.num_layers):
+                prefix = f"blocks.{layer}.adaln_proj.linear"
+                weight = read_shard(f"{prefix}.weight", self.block_width)
+                bias = read_shard(f"{prefix}.bias", self.block_width)
+                for slot, length, adaln_input in slots:
+                    self.block_params[slot, :length, layer] = project(
+                        adaln_input, weight, bias
+                    )
+                del weight, bias
+            prefix = "final_layer.adaln_proj.linear"
+            weight = read_shard(f"{prefix}.weight", self.final_width)
+            bias = read_shard(f"{prefix}.bias", self.final_width)
+            for slot, length, adaln_input in slots:
+                self.final_params[slot, :length] = project(adaln_input, weight, bias)
+            del weight, bias
+
+        for slot, length, _ in slots:
+            self.plan_lengths[slot] = length
+        self.rebuilds += 1
+        logger.info(
+            "MiniMax H3 AdaLN: rebuilt %d plan(s), %d/%d resident, pass #%d",
+            len(missing),
+            len(self._slots),
+            self.max_plans,
+            self.rebuilds,
+        )
+
+    def lookup(self, unique_timesteps: torch.Tensor) -> torch.Tensor:
+        num_timesteps = unique_timesteps.shape[0]
+        matches = self.plan_lengths.eq(num_timesteps) & self.plan_timesteps[
+            :, :num_timesteps
+        ].eq(unique_timesteps).all(dim=-1)
+        if not bool(matches.any()):
+            raise ValueError(
+                "MiniMax H3 AdaLN cache does not cover the request timestep plan"
+            )
+        return matches.to(torch.int64).argmax()
+
+    def block(
+        self,
+        index: int,
+        cache_plan_index: torch.Tensor,
+        num_timesteps: int,
+    ) -> tuple[torch.Tensor, ...]:
+        params = self.block_params[cache_plan_index, :num_timesteps, index]
+        params = params.reshape(-1, 6, self.hidden_size)
+        return tuple(params.unbind(dim=1))
+
+    def final(
+        self,
+        cache_plan_index: torch.Tensor,
+        num_timesteps: int,
+    ) -> tuple[torch.Tensor, ...]:
+        params = self.final_params[cache_plan_index, :num_timesteps]
+        return tuple(params.reshape(-1, 2, self.hidden_size).unbind(dim=1))
+
+
 class MiniMaxH3TokenRefinerBlock(nn.Module):
     """Standard pre-norm transformer block without AdaLN or RoPE."""
 
@@ -890,6 +1181,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
+        use_adaln_cache: bool = False,
     ) -> None:
         super().__init__()
         self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
@@ -900,13 +1192,17 @@ class MiniMaxH3DiTBlock(nn.Module):
             prefix=f"{prefix}.attn",
         )
         self.mlp = MiniMaxH3MLP(arch, quant_config, prefix=f"{prefix}.mlp")
-        self.adaln_proj = MiniMaxH3AdalnProj(
-            arch,
-            arch.adaln_out_features,
-            quant_config,
-            prefix=f"{prefix}.adaln_proj",
-            expand_ratio=6,
-            modality_num=MINIMAX_H3_ADALN_MODALITY_NUM,
+        self.adaln_proj = (
+            None
+            if use_adaln_cache
+            else MiniMaxH3AdalnProj(
+                arch,
+                arch.adaln_out_features,
+                quant_config,
+                prefix=f"{prefix}.adaln_proj",
+                expand_ratio=6,
+                modality_num=MINIMAX_H3_ADALN_MODALITY_NUM,
+            )
         )
         self.preserve_input_for_cache_dit = False
 
@@ -932,6 +1228,8 @@ class MiniMaxH3DiTBlock(nn.Module):
         norm2 -> scale/shift -> MLP -> gated residual.
         """
         if adaln_params is None:
+            if self.adaln_proj is None:
+                raise ValueError("MiniMax H3 AdaLN cache parameters are required")
             adaln_params = self.adaln_proj(adaln_input)
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln_params
         # Cache-DiT retains the inputs to its Fn and Mn block ranges. Only the
@@ -984,6 +1282,7 @@ class MiniMaxH3FinalLayer(nn.Module):
         quant_config: QuantizationConfig | None,
         *,
         prefix: str,
+        use_adaln_cache: bool = False,
     ) -> None:
         super().__init__()
         video_patch_dim = (
@@ -993,13 +1292,17 @@ class MiniMaxH3FinalLayer(nn.Module):
             * arch.patch_size[2]
         )
         self.norm = _norm(arch.hidden_size, eps=arch.final_norm_eps)
-        self.adaln_proj = MiniMaxH3AdalnProj(
-            arch,
-            arch.final_adaln_out_features,
-            quant_config,
-            prefix=f"{prefix}.adaln_proj",
-            expand_ratio=2,
-            modality_num=1,
+        self.adaln_proj = (
+            None
+            if use_adaln_cache
+            else MiniMaxH3AdalnProj(
+                arch,
+                arch.final_adaln_out_features,
+                quant_config,
+                prefix=f"{prefix}.adaln_proj",
+                expand_ratio=2,
+                modality_num=1,
+            )
         )
         self.video_out = ColumnParallelLinear(
             arch.hidden_size,
@@ -1026,6 +1329,7 @@ class MiniMaxH3FinalLayer(nn.Module):
         *,
         adaln_input: torch.Tensor,
         inverse_indices: torch.Tensor,
+        adaln_params: tuple[torch.Tensor, ...] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Project all rows into TP-local video/audio output shards.
 
@@ -1034,7 +1338,11 @@ class MiniMaxH3FinalLayer(nn.Module):
         The model gathers output columns only after selecting live media rows,
         preserving the GEMM shape while reducing collective payload.
         """
-        shift, scale = self.adaln_proj(adaln_input)
+        if adaln_params is None:
+            if self.adaln_proj is None:
+                raise ValueError("MiniMax H3 AdaLN cache parameters are required")
+            adaln_params = self.adaln_proj(adaln_input)
+        shift, scale = adaln_params
         h = self.norm(x)
         h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE)
         # Preserve full precision through both final output projections.
@@ -1056,9 +1364,25 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
     reverse_param_names_mapping = _ARCH_DEFAULTS.reverse_param_names_mapping
     lora_param_names_mapping = _ARCH_DEFAULTS.lora_param_names_mapping
 
+    def prepare_adaln_plans(self, step_timesteps: list[torch.Tensor]) -> None:
+        """Fill the AdaLN cache for this request before denoising starts.
+
+        No-op for a prebuilt sidecar; the rebuild path needs the model's own
+        timestep embedding so a filled plan is bit-identical to what resident
+        adaln_proj weights would have produced.
+        """
+        if self.adaln_cache is None or self.adaln_cache.weight_files is None:
+            return
+
+        def embed(timesteps: torch.Tensor) -> torch.Tensor:
+            return nn.functional.silu(self.time_embedder(timesteps)).to(_BF16_DTYPE)
+
+        self.adaln_cache.build(step_timesteps, embed=embed)
+
     def _can_batch_block_adaln(self) -> bool:
         return (
-            get_tp_world_size() > 1
+            self.adaln_cache is None
+            and get_tp_world_size() > 1
             and not torch.compiler.is_compiling()
             and not envs.SGLANG_CACHE_DIT_ENABLED
             and not hasattr(self, "_sglang_cache_dit_adapter")
@@ -1134,8 +1458,21 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         config: MiniMaxH3DiTConfig,
         hf_config: dict[str, Any],
         quant_config: QuantizationConfig | None = None,
+        adaln_cache_path: str | None = None,
+        adaln_cache_model_variant: str | None = None,
+        adaln_weight_files: list[str] | None = None,
+        adaln_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH,
     ) -> None:
         super().__init__(config=config, hf_config=hf_config)
+        if (
+            adaln_cache_path is not None or adaln_weight_files is not None
+        ) and quant_config is not None:
+            raise ValueError(
+                "MiniMax H3 AdaLN cache is only compatible with unquantized weights"
+            )
+        self._adaln_precomputed = (
+            adaln_cache_path is not None or adaln_weight_files is not None
+        )
         arch = self.config
         self.arch = arch
         self.hidden_size = arch.hidden_size
@@ -1197,6 +1534,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
                     arch,
                     quant_config,
                     prefix=f"blocks.{index}",
+                    use_adaln_cache=self._adaln_precomputed,
                 )
                 for index in range(arch.num_layers)
             ]
@@ -1206,6 +1544,18 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
             arch,
             quant_config,
             prefix="final_layer",
+            use_adaln_cache=self._adaln_precomputed,
+        )
+        self.adaln_cache = (
+            MiniMaxH3AdalnCache(
+                arch,
+                path=adaln_cache_path,
+                model_variant=adaln_cache_model_variant,
+                weight_files=adaln_weight_files,
+                max_plan_width=adaln_plan_width,
+            )
+            if self._adaln_precomputed
+            else None
         )
         self._resolved_attention_backend: AttentionBackendEnum | None = None
         self._mark_missing_params_required()
@@ -1256,6 +1606,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
             raise ValueError(
                 f"rope.inv_freq must stay fp32 after load, got {rope_inv_freq.dtype}."
             )
+        if self.adaln_cache is not None:
+            self.adaln_cache.load(self.video_patch_proj.weight.device)
 
     @staticmethod
     def _pos_ids(pos_info: Any, key: str) -> torch.Tensor:
@@ -1685,7 +2037,20 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         hidden = decoder_input
         cu_seqlens = cu_seqlens.to(device)
         block_adaln_params = None
-        if self._can_batch_block_adaln():
+        adaln_cache_plan_index = None
+        if self.adaln_cache is not None:
+            adaln_cache_plan_index = self.adaln_cache.lookup(
+                unique_timesteps.view(-1).to(device)
+            )
+            block_adaln_params = tuple(
+                self.adaln_cache.block(
+                    index,
+                    adaln_cache_plan_index,
+                    adaln_input.shape[0],
+                )
+                for index in range(len(self.blocks))
+            )
+        elif self._can_batch_block_adaln():
             local_adaln = torch.stack(
                 [block.adaln_proj.project_local(adaln_input) for block in self.blocks]
             )
@@ -1718,6 +2083,14 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
             hidden,
             adaln_input=adaln_input,
             inverse_indices=block_inverse,
+            adaln_params=(
+                None
+                if adaln_cache_plan_index is None
+                else self.adaln_cache.final(
+                    adaln_cache_plan_index,
+                    adaln_input.shape[0],
+                )
+            ),
         )
         if sp_ws > 1:
             from sglang.multimodal_gen.runtime.distributed.parallel_state import (

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py

@@ -449,6 +449,11 @@ def minimax_h3_denoise_loop(
         imgvid_cond_noise_aug=float(imgvid_cond_noise_aug_for_inference),
         audio_ref_cond_noise_aug=float(audio_cond_noise_aug_for_inference),
     )
+    # Every step's timesteps are settled by now. Rebuilding AdaLN reads all
+    # 24.2 GiB of adaln_proj whatever is missing, so fill the whole request in
+    # one pass here instead of topping up step by step inside the loop.
+    model.prepare_adaln_plans([entry[0] for entry in timestep_plan])
+
     # match the scheduler's device-fp32 math once, then reuse one denoised
     # scratch per modality instead of allocating intermediates every step
     video_sigmas = torch.tensor(sigmas_video, dtype=torch.float32, device=device)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py

@@ -61,7 +61,7 @@ def _cached_decode_mean_std(
     return mean, std
 
 
-def _reverse_normalize_latents_(
+def _reverse_normalize_latents(
     latents: torch.Tensor,
     *,
     mean_values,
@@ -89,7 +89,13 @@ def _reverse_normalize_latents_(
         )
     view_shape = [1] * latents.ndim
     view_shape[1] = int(mean.shape[0])
-    return latents.mul_(std.view(*view_shape)).add_(mean.view(*view_shape))
+    # Out of place on purpose. batch.latents / batch.audio_latents are
+    # inference tensors allocated inside the denoising stage's InferenceMode,
+    # while --vae-cpu-offload runs this stage under torch.inference_mode(False)
+    # (PipelineExecutor._stage_needs_version_counters), where writing to one in
+    # place raises. Keep the mul-then-add rounding order rather than addcmul so
+    # the result does not shift with FMA contraction.
+    return latents * std.view(*view_shape) + mean.view(*view_shape)
 
 
 def _crop_to_target_canvas(batch: Req, frames: torch.Tensor) -> torch.Tensor:
@@ -276,7 +282,7 @@ class MiniMaxH3DecodingStage(DecodingStage):
             if audio_vae.training:
                 audio_vae.eval()
             audio_arch_config = server_args.pipeline_config.audio_vae_config.arch_config
-            audio_decode_latent = _reverse_normalize_latents_(
+            audio_decode_latent = _reverse_normalize_latents(
                 audio_latent,
                 mean_values=audio_arch_config.latents_mean,
                 std_values=audio_arch_config.latents_std,
@@ -332,7 +338,7 @@ class MiniMaxH3DecodingStage(DecodingStage):
             if selected_video_vae.training:
                 selected_video_vae.eval()
             visual_arch_config = server_args.pipeline_config.vae_config.arch_config
-            visual_decode_latent = _reverse_normalize_latents_(
+            visual_decode_latent = _reverse_normalize_latents(
                 visual_latent,
                 mean_values=visual_arch_config.latents_mean,
                 std_values=visual_arch_config.latents_std,

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -293,6 +293,13 @@ class ServerArgs(DisaggServerArgsMixin):
 
     # path to pre-quantized transformer weights (single .safetensors or directory).
     transformer_weights_path: str | None = None
+    # path to precomputed MiniMax H3 AdaLN outputs for inference-only serving.
+    minimax_h3_adaln_cache_path: str | None = None
+    # Rebuild AdaLN outputs per request from the checkpoint, no sidecar needed.
+    minimax_h3_adaln_online: bool = False
+    # Widest timestep plan the rebuild slab is sized for; see
+    # MINIMAX_H3_ADALN_MAX_PLAN_WIDTH.
+    minimax_h3_adaln_plan_width: int = 4
     # Per-component transformer weight overrides (key = model_index.json component name).
     # Pipelines use this when a checkpoint ships separate quantized weights for
     # secondary DiT components; the generic loader consumes it without model-specific
@@ -1495,6 +1502,38 @@ class ServerArgs(DisaggServerArgsMixin):
                 "without exposing repository subfolder layout."
             ),
         )
+        parser.add_argument(
+            "--minimax-h3-adaln-online",
+            action=StoreBoolean,
+            default=ServerArgs.minimax_h3_adaln_online,
+            help=(
+                "Rebuild MiniMax H3 AdaLN outputs from the checkpoint per "
+                "request instead of keeping the 24.2 GiB of adaln_proj weights "
+                "resident. Works with any step count or schedule and needs no "
+                "prebuilt artifact. Requires unquantized weights."
+            ),
+        )
+        parser.add_argument(
+            "--minimax-h3-adaln-plan-width",
+            type=int,
+            default=ServerArgs.minimax_h3_adaln_plan_width,
+            help=(
+                "Widest timestep plan --minimax-h3-adaln-online sizes its slab "
+                "for. The default 4 covers every task; a deployment serving "
+                "only t2va (2) or fl2va (3) can shrink the slab proportionally. "
+                "A request exceeding it is rejected rather than truncated."
+            ),
+        )
+        parser.add_argument(
+            "--minimax-h3-adaln-cache-path",
+            type=str,
+            default=ServerArgs.minimax_h3_adaln_cache_path,
+            help=(
+                "Path to a precomputed MiniMax H3 AdaLN cache. This only "
+                "supports the matching unquantized H3 checkpoint and rejects "
+                "requests whose timestep embeddings are not present in the cache."
+            ),
+        )
         parser.add_argument(
             "--model-id",
             type=str,

python/sglang/multimodal_gen/test/unit/test_fsdp_load.py

@@ -188,6 +188,7 @@ class TestOrdinaryWeightLoading(unittest.TestCase):
         rank_local_load.assert_not_called()
         weight_iterator.assert_called_once_with(
             ["model.safetensors"],
+            key_filter=None,
             weight_load_plan=load_plan,
         )
 

python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from safetensors.torch import save_file
+
+from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
+    MiniMaxH3DiTArchConfig,
+)
+from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
+    MiniMaxH3AdalnCache,
+)
+
+
+def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path):
+    arch = MiniMaxH3DiTArchConfig(
+        num_layers=2,
+        hidden_size=4,
+        time_embed_dim=3,
+    )
+    cache_path = tmp_path / "adaln.safetensors"
+    plan_timesteps = torch.tensor([[0.0, 0.0], [1.0, 2.0]])
+    plan_lengths = torch.tensor([1, 2], dtype=torch.int64)
+    block_params = (
+        torch.arange(2 * 2 * 2 * 72, dtype=torch.float32)
+        .reshape(2, 2, 2, 72)
+        .bfloat16()
+    )
+    final_params = torch.arange(32, dtype=torch.float32).reshape(2, 2, 8).bfloat16()
+    save_file(
+        {
+            "plan_timesteps": plan_timesteps,
+            "plan_lengths": plan_lengths,
+            "block_params": block_params,
+            "final_params": final_params,
+        },
+        cache_path,
+        metadata={"format_version": "2", "model_variant": "fl2va"},
+    )
+
+    cache = MiniMaxH3AdalnCache(
+        arch,
+        path=str(cache_path),
+        model_variant="fl2va",
+    )
+    cache.load(torch.device("cpu"))
+
+    cache_plan_index = cache.lookup(plan_timesteps[1])
+    block = cache.block(1, cache_plan_index, 2)
+    final = cache.final(cache_plan_index, 2)
+
+    # block() hands the forward pass six [num_timesteps * modality, hidden]
+    # chunks, while the checkpoint stores a plan as one flat
+    # [num_timesteps, 6 * modality * hidden] row -- same elements, and the
+    # modality axis folds into the leading one rather than staying separate.
+    assert torch.equal(
+        torch.cat(block, dim=-1).reshape(block_params[1, :, 1].shape),
+        block_params[1, :, 1],
+    )
+    assert torch.equal(torch.cat(final, dim=-1), final_params[1])

python/sglang/multimodal_gen/tools/build_minimax_h3_adaln_cache.py

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Build an inference-only AdaLN cache from a local MiniMax H3 transformer."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from contextlib import ExitStack
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import safe_open, save_file
+
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import (
+    MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
+    MINIMAX_H3_IMGVID_COND_TIMESTEP,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+    minimax_h3_time_shift_sigmas,
+)
+
+_HIDDEN_SIZE = 5376
+_TIMESTEP_INPUT_DIM = 256
+_NUM_BLOCKS = 50
+_BLOCK_PARAM_WIDTH = 18 * _HIDDEN_SIZE
+_FINAL_PARAM_WIDTH = 2 * _HIDDEN_SIZE
+_CACHE_MODES = {
+    "t2va": ("video", "audio"),
+    "fl2va": ("video", "audio", "image"),
+    "ref2va-image": ("video", "audio", "image"),
+    "ref2va-audio": ("video", "audio", "audio_ref"),
+    "ref2va-mixed": ("video", "audio", "image", "audio_ref"),
+}
+_MODE_VARIANTS = {
+    "t2va": "fl2va",
+    "fl2va": "fl2va",
+    "ref2va-image": "ref2va",
+    "ref2va-audio": "ref2va",
+    "ref2va-mixed": "ref2va",
+}
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Build a MiniMax H3 inference-only AdaLN cache from local weights."
+    )
+    parser.add_argument("--transformer-path", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--model-variant", choices=("fl2va", "ref2va"), required=True)
+    parser.add_argument("--mode", choices=tuple(_CACHE_MODES), default="t2va")
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--flow-shift", type=float, default=12.0)
+    parser.add_argument("--audio-flow-shift", type=float, default=3.0)
+    parser.add_argument(
+        "--imgvid-cond-noise-aug",
+        type=float,
+        default=MINIMAX_H3_IMGVID_COND_TIMESTEP,
+    )
+    parser.add_argument(
+        "--audio-cond-noise-aug",
+        type=float,
+        default=MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
+    )
+    parser.add_argument(
+        "--timesteps",
+        type=float,
+        nargs="+",
+        help="Override the scheduler-derived timestep plan.",
+    )
+    parser.add_argument("--device", default="cuda")
+    return parser.parse_args()
+
+
+def _cache_timestep_plans(args: argparse.Namespace) -> list[torch.Tensor]:
+    if args.timesteps is not None:
+        return [torch.tensor(args.timesteps, dtype=torch.float32).unique(sorted=True)]
+
+    video_sigmas = minimax_h3_time_shift_sigmas(
+        num_steps=args.num_inference_steps,
+        shift_scale=args.flow_shift,
+    )
+    audio_sigmas = minimax_h3_time_shift_sigmas(
+        num_steps=args.num_inference_steps,
+        shift_scale=args.audio_flow_shift,
+    )
+    fields = _CACHE_MODES[args.mode]
+    plans = []
+    for video_sigma, audio_sigma in zip(video_sigmas[:-1], audio_sigmas[:-1]):
+        video_timestep = 1.0 - video_sigma
+        audio_timestep = 1.0 - audio_sigma
+        candidates = {
+            "video": video_timestep,
+            "audio": audio_timestep,
+            "image": max(video_timestep, args.imgvid_cond_noise_aug),
+            "audio_ref": max(audio_timestep, args.audio_cond_noise_aug),
+        }
+        plans.append(
+            torch.tensor(
+                [candidates[field] for field in fields], dtype=torch.float32
+            ).unique(sorted=True)
+        )
+
+    deduplicated = []
+    seen = set()
+    for plan in plans:
+        key = tuple(plan.tolist())
+        if key not in seen:
+            seen.add(key)
+            deduplicated.append(plan)
+    return deduplicated
+
+
+def _time_embed(
+    timesteps: torch.Tensor,
+    *,
+    proj_in_weight: torch.Tensor,
+    proj_in_bias: torch.Tensor,
+    proj_out_weight: torch.Tensor,
+    proj_out_bias: torch.Tensor,
+) -> torch.Tensor:
+    half = _TIMESTEP_INPUT_DIM // 2
+    freqs = torch.exp(
+        -math.log(10000.0)
+        * torch.arange(half, dtype=torch.float32, device=timesteps.device)
+        / half
+    )
+    args = timesteps[:, None] * freqs[None]
+    t_freq = torch.cat((torch.cos(args), torch.sin(args)), dim=-1)
+    return F.linear(
+        F.silu(F.linear(t_freq, proj_in_weight, proj_in_bias)),
+        proj_out_weight,
+        proj_out_bias,
+    )
+
+
+def _load_tensor(
+    name: str,
+    *,
+    weight_map: dict[str, str],
+    files: dict[str, Any],
+    device: torch.device,
+) -> torch.Tensor:
+    tensor_file = files[weight_map[name]]
+    return tensor_file.get_tensor(name).to(device)
+
+
+def main() -> None:
+    args = _parse_args()
+    if args.num_inference_steps < 2 and args.timesteps is None:
+        raise ValueError("--num-inference-steps must be at least 2")
+    mode_variant = _MODE_VARIANTS[args.mode]
+    if args.model_variant != mode_variant:
+        raise ValueError(f"--mode {args.mode} requires {mode_variant}")
+    device = torch.device(args.device)
+    if device.type != "cuda" or not torch.cuda.is_available():
+        raise ValueError("MiniMax H3 AdaLN cache must be built on CUDA")
+
+    index_path = args.transformer_path / "model.safetensors.index.json"
+    with index_path.open() as f:
+        weight_map = json.load(f)["weight_map"]
+
+    plans = _cache_timestep_plans(args)
+    if not plans or any(plan.numel() == 0 for plan in plans):
+        raise ValueError("AdaLN cache must cover at least one timestep plan")
+    max_plan_length = max(plan.numel() for plan in plans)
+    plan_timesteps = torch.zeros((len(plans), max_plan_length), dtype=torch.float32)
+    plan_lengths = torch.tensor([plan.numel() for plan in plans], dtype=torch.int64)
+    block_params = torch.empty(
+        (len(plans), max_plan_length, _NUM_BLOCKS, _BLOCK_PARAM_WIDTH),
+        dtype=torch.bfloat16,
+    )
+    final_params = torch.empty(
+        (len(plans), max_plan_length, _FINAL_PARAM_WIDTH), dtype=torch.bfloat16
+    )
+
+    with ExitStack() as stack:
+        files = {
+            filename: stack.enter_context(
+                safe_open(
+                    str(args.transformer_path / filename),
+                    framework="pt",
+                    device="cpu",
+                )
+            )
+            for filename in set(weight_map.values())
+        }
+        time_kwargs = {
+            f"{module}_{name}": _load_tensor(
+                f"time_embedder.{module}.{name}",
+                weight_map=weight_map,
+                files=files,
+                device=device,
+            )
+            for module, name in (
+                ("proj_in", "weight"),
+                ("proj_in", "bias"),
+                ("proj_out", "weight"),
+                ("proj_out", "bias"),
+            )
+        }
+        adaln_inputs = []
+        for plan_index, plan in enumerate(plans):
+            plan_length = plan.numel()
+            plan_timesteps[plan_index, :plan_length].copy_(plan)
+            adaln_inputs.append(
+                F.silu(_time_embed(plan.to(device), **time_kwargs)).to(torch.bfloat16)
+            )
+
+        for index in range(_NUM_BLOCKS):
+            prefix = f"blocks.{index}.adaln_proj.linear"
+            weight = _load_tensor(
+                f"{prefix}.weight",
+                weight_map=weight_map,
+                files=files,
+                device=device,
+            )
+            bias = _load_tensor(
+                f"{prefix}.bias",
+                weight_map=weight_map,
+                files=files,
+                device=device,
+            )
+            for plan_index, adaln_input in enumerate(adaln_inputs):
+                plan_length = adaln_input.shape[0]
+                block_params[plan_index, :plan_length, index].copy_(
+                    F.linear(adaln_input, weight, bias).cpu()
+                )
+
+        prefix = "final_layer.adaln_proj.linear"
+        weight = _load_tensor(
+            f"{prefix}.weight",
+            weight_map=weight_map,
+            files=files,
+            device=device,
+        )
+        bias = _load_tensor(
+            f"{prefix}.bias",
+            weight_map=weight_map,
+            files=files,
+            device=device,
+        )
+        for plan_index, adaln_input in enumerate(adaln_inputs):
+            plan_length = adaln_input.shape[0]
+            final_params[plan_index, :plan_length].copy_(
+                F.linear(adaln_input, weight, bias).cpu()
+            )
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    save_file(
+        {
+            "plan_timesteps": plan_timesteps,
+            "plan_lengths": plan_lengths,
+            "block_params": block_params,
+            "final_params": final_params,
+        },
+        str(args.output),
+        metadata={
+            "format_version": "2",
+            "model_variant": args.model_variant,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()