[Perf] [MiniMax H3] Precompute AdaLN modulation schedules

Signed-off-by: Sy03 <1370724210@qq.com>

vllm_omni/diffusion/models/minimax_h3/adaln_schedule_cache.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact schedule identity and policy for MiniMax-H3 AdaLN reuse."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV = "VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE"
+MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV = "VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS"
+_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"})
+
+MiniMaxH3AdalnScheduleKey = tuple[tuple[int, ...], ...]
+
+
+@dataclass(frozen=True)
+class MiniMaxH3AdalnScheduleCache:
+    """One model-local exact modulation schedule."""
+
+    key: MiniMaxH3AdalnScheduleKey
+    t_embs: tuple[torch.Tensor, ...]
+    step_cursor: torch.Tensor
+    steps: int
+    max_unique_timesteps: int
+    table_bytes: int
+    projection_bytes: int = 0
+    offloaded_projection_bytes: int = 0
+    net_memory_saved_bytes: int = 0
+
+
+def minimax_h3_adaln_schedule_cache_enabled() -> bool:
+    """Return whether persistent exact-schedule AdaLN reuse is enabled."""
+    return os.getenv(MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV, "").strip().lower() in _TRUTHY_ENV_VALUES
+
+
+def minimax_h3_adaln_weight_offload_enabled() -> bool:
+    """Return whether exact tables should take ownership of AdaLN weights."""
+    return minimax_h3_adaln_schedule_cache_enabled() and (
+        os.getenv(MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV, "").strip().lower() in _TRUTHY_ENV_VALUES
+    )
+
+
+def minimax_h3_float32_bits(values: Any) -> tuple[int, ...]:
+    """Return the signed int32 bit patterns of a CPU float32 tensor."""
+    if not isinstance(values, torch.Tensor) or values.device.type != "cpu" or values.dtype != torch.float32:
+        raise ValueError("schedule key values must be a CPU float32 tensor")
+    return tuple(int(value) for value in values.contiguous().view(torch.int32).tolist())
+
+
+def minimax_h3_adaln_schedule_key(
+    metadata: Iterable[Any],
+) -> MiniMaxH3AdalnScheduleKey:
+    """Build an immutable schedule key without reading a device tensor."""
+    key: list[tuple[int, ...]] = []
+    for step in metadata:
+        bits = getattr(step, "timestep_bits", None)
+        if bits is None:
+            raise ValueError("timestep metadata is missing host timestep_bits")
+        key.append(tuple(int(value) for value in bits))
+    if not key:
+        raise ValueError("AdaLN schedule must contain at least one denoise step")
+    return tuple(key)
+
+
+def minimax_h3_build_adaln_table(
+    projection: Any,
+    t_embs: tuple[torch.Tensor, ...],
+) -> torch.Tensor:
+    """Project each step at its original shape and pad only unused rows."""
+    rows_by_step = tuple(projection.compute_flat(t_emb) for t_emb in t_embs)
+    max_rows = max(int(rows.shape[0]) for rows in rows_by_step)
+    padded: list[torch.Tensor] = []
+    for rows in rows_by_step:
+        missing = max_rows - int(rows.shape[0])
+        if missing:
+            rows = torch.cat(
+                [rows, rows.new_zeros((missing, rows.shape[1]))],
+                dim=0,
+            )
+        padded.append(rows)
+    return torch.stack(padded, dim=0)
+
+
+__all__ = [
+    "MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV",
+    "MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV",
+    "MiniMaxH3AdalnScheduleCache",
+    "MiniMaxH3AdalnScheduleKey",
+    "minimax_h3_build_adaln_table",
+    "minimax_h3_adaln_schedule_cache_enabled",
+    "minimax_h3_adaln_weight_offload_enabled",
+    "minimax_h3_adaln_schedule_key",
+    "minimax_h3_float32_bits",
+]

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -10,10 +10,17 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass
 from typing import Any
 
 import torch
 
+from .adaln_schedule_cache import (
+    minimax_h3_adaln_schedule_cache_enabled,
+    minimax_h3_adaln_schedule_key,
+    minimax_h3_adaln_weight_offload_enabled,
+    minimax_h3_float32_bits,
+)
 from .scheduling_minimax_h3_euler_ancestral import (
     minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
@@ -28,6 +35,15 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96
 MINIMAX_H3_AUDIO_ROW_WIDTH = 32
 
 
+@dataclass(frozen=True)
+class MiniMaxH3StepTimestepMetadata:
+    """Device-ready timestep routing for one denoise forward."""
+
+    unique_timesteps: torch.Tensor
+    inverse_indices: torch.Tensor
+    timestep_bits: tuple[int, ...]
+
+
 class MiniMaxH3DenoiseBranch:
     """Static per-branch state: packed layout + fixed forward kwargs.
 
@@ -45,6 +61,7 @@ class MiniMaxH3DenoiseBranch:
         device: torch.device,
     ) -> None:
         seq_len = int(packed["seq_len"])
+        self.device = device
         self.seq_len = seq_len
         self.img_pos = packed["img_pos"].view(-1).to(torch.long)
         self.audio_pos = packed["audio_pos"].view(-1).to(torch.long)
@@ -65,6 +82,12 @@ class MiniMaxH3DenoiseBranch:
         self.audio_pos_dev = self.audio_pos.to(device)
         self.update_mask_dev = self.update_mask.to(device)
         self.audio_update_mask_dev = self.audio_update_mask.to(device)
+        timestep_category_ids = torch.zeros(seq_len, dtype=torch.long)
+        timestep_category_ids[self.img_pos[~self.update_mask]] = 1
+        timestep_category_ids[self.audio_pos[self.audio_update_mask]] = 2
+        timestep_category_ids[self.audio_pos[~self.audio_update_mask]] = 3
+        self.timestep_category_ids = timestep_category_ids
+        self.used_timestep_categories = torch.unique(timestep_category_ids, sorted=True)
         self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
         self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
         text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
@@ -87,43 +110,103 @@ class MiniMaxH3DenoiseBranch:
                 "max_seqlen_q": text_len,
             },
         }
+        self._adaln_schedule_cache_active = False
+
+    def prepare_timestep_metadata(
+        self,
+        *,
+        sigmas_video: list[float],
+        sigmas_audio: list[float],
+        imgvid_cond_noise_aug_for_inference: float,
+        audio_cond_noise_aug_for_inference: float,
+    ) -> list[MiniMaxH3StepTimestepMetadata]:
+        """Precompute the full schedule's packed-row timestep routing on CPU."""
+        if len(sigmas_video) != len(sigmas_audio):
+            raise ValueError("video/audio sigma schedules must have equal length")
+        if len(sigmas_video) < 2:
+            raise ValueError("sigma schedules need at least 2 entries")
+
+        result: list[MiniMaxH3StepTimestepMetadata] = []
+        for step in range(len(sigmas_video) - 1):
+            t_video = 1.0 - float(sigmas_video[step])
+            t_audio = 1.0 - float(sigmas_audio[step])
+            category_timesteps = torch.tensor(
+                [
+                    t_video,
+                    max(t_video, float(imgvid_cond_noise_aug_for_inference)),
+                    t_audio,
+                    max(t_audio, float(audio_cond_noise_aug_for_inference)),
+                ],
+                dtype=torch.float32,
+            )
+            active_timesteps = category_timesteps.index_select(0, self.used_timestep_categories)
+            unique_timesteps, active_inverse = torch.unique(
+                active_timesteps,
+                sorted=True,
+                return_inverse=True,
+            )
+            category_inverse = torch.zeros(4, dtype=torch.long)
+            category_inverse.index_copy_(0, self.used_timestep_categories, active_inverse)
+            inverse_indices = category_inverse.index_select(0, self.timestep_category_ids)
+            result.append(
+                MiniMaxH3StepTimestepMetadata(
+                    unique_timesteps=unique_timesteps.to(self.device),
+                    inverse_indices=inverse_indices.to(self.device),
+                    timestep_bits=minimax_h3_float32_bits(unique_timesteps),
+                )
+            )
+        return result
+
+    def prepare_adaln_schedule_cache(
+        self,
+        model: Any,
+        timestep_metadata: list[MiniMaxH3StepTimestepMetadata],
+        *,
+        enabled: bool,
+    ) -> Any | None:
+        """Prepare exact model-local modulation tables for this schedule."""
+        self._adaln_schedule_cache_active = False
+        if not enabled:
+            clear = getattr(model, "clear_adaln_schedule_cache", None)
+            if clear is not None:
+                clear()
+            return None
+        prepare = getattr(model, "prepare_adaln_schedule_cache", None)
+        if prepare is None:
+            return None
+        with torch.inference_mode():
+            cache = prepare(
+                unique_timestep_plan=tuple(step.unique_timesteps for step in timestep_metadata),
+                schedule_key=minimax_h3_adaln_schedule_key(timestep_metadata),
+                offload_weights=minimax_h3_adaln_weight_offload_enabled(),
+            )
+        self._adaln_schedule_cache_active = cache is not None
+        return cache
 
     def forward_kwargs(
         self,
         *,
         video_rows: torch.Tensor,
         audio_rows: torch.Tensor,
-        t_video: float,
-        t_audio: float,
-        imgvid_cond_timestep: float,
-        audio_ref_cond_timestep: float,
+        timestep_metadata: MiniMaxH3StepTimestepMetadata,
+        step_index: int | None = None,
     ) -> dict[str, Any]:
         x = self.x_base.clone()
         x[0].index_copy_(0, self.img_pos_dev, video_rows)
         audio_x = self.audio_x_base.clone()
         audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows)
-        # Packed-sequence timestep semantics: non-media rows (text and
-        # padding) inherit the current video timestep. Later steps must reuse
-        # the previous step's updated rows; re-initializing from zeros is only
-        # valid at step 0.
-        timesteps = torch.full(
-            (self.seq_len,),
-            float(t_video),
-            dtype=torch.float32,
-            device=x.device,
-        )
-        timesteps[self.img_pos_dev[self.update_mask_dev]] = t_video
-        timesteps[self.img_pos_dev[~self.update_mask_dev]] = imgvid_cond_timestep
-        timesteps[self.audio_pos_dev[self.audio_update_mask_dev]] = t_audio
-        timesteps[self.audio_pos_dev[~self.audio_update_mask_dev]] = audio_ref_cond_timestep
-        unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True)
-        return {
+        dynamic_kwargs = {
             **self.static_kwargs,
             "x": x,
             "audio_x": audio_x,
-            "unique_timesteps": unique_timesteps,
-            "inverse_indices": inverse_indices,
+            "unique_timesteps": timestep_metadata.unique_timesteps,
+            "inverse_indices": timestep_metadata.inverse_indices,
         }
+        if self._adaln_schedule_cache_active:
+            if step_index is None:
+                raise ValueError("step_index is required while the AdaLN schedule cache is active")
+            dynamic_kwargs["adaln_step_index"] = step_index
+        return dynamic_kwargs
 
 
 def minimax_h3_denoise_loop(
@@ -153,6 +236,12 @@ def minimax_h3_denoise_loop(
         raise ValueError("video/audio sigma schedules must have equal length")
     if len(sigmas_video) < 2:
         raise ValueError("sigma schedules need at least 2 entries")
+    timestep_metadata = positive.prepare_timestep_metadata(
+        sigmas_video=sigmas_video,
+        sigmas_audio=sigmas_audio,
+        imgvid_cond_noise_aug_for_inference=imgvid_cond_noise_aug_for_inference,
+        audio_cond_noise_aug_for_inference=audio_cond_noise_aug_for_inference,
+    )
     n_cond = int((~positive.update_mask).sum())
     if keyframe_cond_rows is None:
         if n_cond != 0:
@@ -187,6 +276,16 @@ def minimax_h3_denoise_loop(
     if audio_anchor is not None:
         audio_rows[~audio_update] = audio_anchor
 
+    prepare_for_request_caches = getattr(model, "prepare_for_request_caches", None)
+    if prepare_for_request_caches is not None:
+        with torch.inference_mode():
+            prepare_for_request_caches()
+    positive.prepare_adaln_schedule_cache(
+        model,
+        timestep_metadata,
+        enabled=minimax_h3_adaln_schedule_cache_enabled(),
+    )
+
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
@@ -194,16 +293,12 @@ def minimax_h3_denoise_loop(
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
             s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
             t_v, t_a = 1.0 - s_v, 1.0 - s_a
-            imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference))
-            audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference))
 
             fk = positive.forward_kwargs(
                 video_rows=video_rows,
                 audio_rows=audio_rows,
-                t_video=t_v,
-                t_audio=t_a,
-                imgvid_cond_timestep=imgvid_cond_t,
-                audio_ref_cond_timestep=audio_ref_cond_t,
+                timestep_metadata=timestep_metadata[step],
+                step_index=step,
             )
             with torch.inference_mode():
                 v_video, v_audio = model(**fk)
@@ -245,5 +340,6 @@ __all__ = [
     "MINIMAX_H3_IMGVID_COND_TIMESTEP",
     "MINIMAX_H3_VIDEO_ROW_WIDTH",
     "MiniMaxH3DenoiseBranch",
+    "MiniMaxH3StepTimestepMetadata",
     "minimax_h3_denoise_loop",
 ]

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 import torch.nn as nn
 from cache_dit import ForwardPattern
+from torch.distributed.tensor import DTensor
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
@@ -33,6 +34,14 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
+from vllm_omni.platforms import current_omni_platform
+
+from .adaln_schedule_cache import (
+    MiniMaxH3AdalnScheduleCache,
+    MiniMaxH3AdalnScheduleKey,
+    minimax_h3_adaln_schedule_cache_enabled,
+    minimax_h3_build_adaln_table,
+)
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -132,6 +141,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset(
         "img_pos_for_infer_output_info",
         "packed_seq_params",
         "refiner_packed_seq_params",
+        "adaln_step_index",
     }
 )
 
@@ -528,14 +538,118 @@ class MiniMaxH3AdalnProj(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.linear",
         )
+        self.register_buffer("_precomputed_table", None, persistent=False)
+        self._precomputed_step: torch.Tensor | None = None
+        self._offloaded_projection_state: dict[str, tuple[torch.Tensor, torch.device, torch.Size]] | None = None
 
+    @property
+    def projection_weights_offloaded(self) -> bool:
+        return self._offloaded_projection_state is not None
+
+    @property
+    def projection_nbytes(self) -> int:
+        state = self._offloaded_projection_state
+        if state is not None:
+            return sum(tensor.numel() * tensor.element_size() for tensor, _, _ in state.values())
+        total = 0
+        for parameter in self.linear.parameters():
+            local = parameter.to_local() if isinstance(parameter, DTensor) else parameter
+            total += local.numel() * local.element_size()
+        return total
+
+    def can_offload_projection(self) -> bool:
+        """Whether this projection can safely own ordinary parameter storage."""
+        if self.projection_weights_offloaded:
+            return True
+        parameters = tuple(self.linear.parameters())
+        return bool(parameters) and all(
+            not isinstance(parameter, DTensor) and not parameter.is_meta and parameter.numel() > 0
+            for parameter in parameters
+        )
+
+    def offload_projection(self) -> int:
+        """Move exact projection parameters to CPU and release device storage."""
+        if self.projection_weights_offloaded:
+            return 0
+        if self._precomputed_table is None:
+            raise RuntimeError("AdaLN projection weights require an installed precomputed table before offload")
+        if not self.can_offload_projection():
+            raise RuntimeError("AdaLN projection weight offload does not support DTensor, meta, or empty parameters")
+
+        state: dict[str, tuple[torch.Tensor, torch.device, torch.Size]] = {}
+        try:
+            for name, parameter in self.linear.named_parameters():
+                cpu_tensor = parameter.detach().to(device="cpu", copy=True)
+                state[name] = (cpu_tensor, parameter.device, parameter.shape)
+                parameter.data = torch.empty((0,), device=parameter.device, dtype=parameter.dtype)
+        except Exception:
+            for name, (cpu_tensor, device, shape) in state.items():
+                self.linear.get_parameter(name).data = cpu_tensor.to(device=device).view(shape)
+            raise
+        self._offloaded_projection_state = state
+        return sum(tensor.numel() * tensor.element_size() for tensor, _, _ in state.values())
+
+    def restore_projection(self) -> None:
+        """Restore a previously offloaded projection to its original devices."""
+        state = self._offloaded_projection_state
+        if state is None:
+            return
+        restored = {
+            name: cpu_tensor.to(device=device).view(shape) for name, (cpu_tensor, device, shape) in state.items()
+        }
+        for name, tensor in restored.items():
+            self.linear.get_parameter(name).data = tensor
+        self._offloaded_projection_state = None
+
+    def compute_flat(self, t_emb: torch.Tensor) -> torch.Tensor:
+        """Run the original projection and return `[M*modalities, ratio*H]`."""
+        x = nn.functional.silu(t_emb)
+        x, _ = self.linear(x.to(self.linear.weight.dtype))
+        m = x.shape[0]
+        return x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
+
+    def install_precomputed(
+        self,
+        table: torch.Tensor,
+        step_cursor: torch.Tensor,
+    ) -> None:
+        """Atomically select a non-persistent exact-schedule modulation table."""
+        expected_width = self.expand_ratio * self.hidden_size
+        if table.dim() != 3 or table.shape[-1] != expected_width:
+            raise ValueError(
+                f"precomputed AdaLN table must be [steps, rows, {expected_width}], got {list(table.shape)}"
+            )
+        if table.dtype != self.linear.weight.dtype:
+            raise ValueError(f"precomputed AdaLN table dtype mismatch: {table.dtype} != {self.linear.weight.dtype}")
+        if step_cursor.dtype != torch.long or step_cursor.dim() != 0:
+            raise ValueError("precomputed AdaLN step cursor must be a scalar torch.long tensor")
+        if table.device != step_cursor.device:
+            raise ValueError(f"precomputed AdaLN table/cursor device mismatch: {table.device} != {step_cursor.device}")
+        self._buffers["_precomputed_table"] = table
+        self._precomputed_step = step_cursor
+
+    def clear_precomputed(self) -> None:
+        """Restore direct projection without changing checkpoint parameters."""
+        self.restore_projection()
+        self._buffers["_precomputed_table"] = None
+        self._precomputed_step = None
+
+    @torch.compiler.disable
     def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]:
         """t_emb: [M, t_dim] -> expand_ratio tensors of [M*modality_num, H]."""
-        x = nn.functional.silu(t_emb)
-        x, _ = self.linear(x.to(_BF16_DTYPE))
-        m = x.shape[0]
-        x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
-        return tuple(x.chunk(self.expand_ratio, dim=-1))
+        table = self._precomputed_table
+        if table is None:
+            rows = self.compute_flat(t_emb)
+            return tuple(rows.chunk(self.expand_ratio, dim=-1))
+
+        step_cursor = self._precomputed_step
+        if step_cursor is None:
+            raise RuntimeError("precomputed AdaLN table is installed without a step cursor")
+        row_count = t_emb.shape[0] * self.modality_num
+        if row_count > table.shape[1]:
+            raise ValueError(f"precomputed AdaLN table row capacity is too small: {table.shape[1]} < {row_count}")
+        rows = table.index_select(0, step_cursor.reshape(1))[0].narrow(0, 0, row_count)
+        return tuple(rows.chunk(self.expand_ratio, dim=-1))
 
 
 class MiniMaxH3TokenRefinerBlock(nn.Module):
@@ -781,6 +895,13 @@ class MiniMaxH3DiTModel(nn.Module):
     _repeated_blocks = ["MiniMaxH3DiTBlock"]
     _layerwise_offload_blocks_attrs = ["blocks"]
 
+    @staticmethod
+    def layerwise_offload_excluded_parameter_prefixes() -> tuple[str, ...]:
+        """Parameters owned by the exact-schedule cache instead of block hooks."""
+        if minimax_h3_adaln_schedule_cache_enabled():
+            return ("adaln_proj.linear.",)
+        return ()
+
     @staticmethod
     def _is_transformer_block(name: str, module: nn.Module) -> bool:
         del module
@@ -921,6 +1042,8 @@ class MiniMaxH3DiTModel(nn.Module):
             quant_config,
             prefix="final_layer",
         )
+        self._adaln_projections = tuple(block.adaln_proj for block in self.blocks) + (self.final_layer.adaln_proj,)
+        self._adaln_schedule_cache: MiniMaxH3AdalnScheduleCache | None = None
         self._mark_missing_params_required()
 
     def _mark_missing_params_required(self) -> None:
@@ -935,6 +1058,126 @@ class MiniMaxH3DiTModel(nn.Module):
             if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE:
                 raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.")
 
+    @torch.inference_mode()
+    def prepare_for_request_caches(self) -> None:
+        """Materialize a sequentially offloaded DiT before auxiliary cache work."""
+        registry = getattr(self, "_hook_registry", None)
+        if registry is None:
+            return
+        hook = registry.get_hook("sequential_offload")
+        if hook is not None:
+            hook.pre_forward(self)
+
+    @torch.inference_mode()
+    def prepare_adaln_schedule_cache(
+        self,
+        *,
+        unique_timestep_plan: tuple[torch.Tensor, ...],
+        schedule_key: MiniMaxH3AdalnScheduleKey,
+        offload_weights: bool = False,
+    ) -> MiniMaxH3AdalnScheduleCache | None:
+        """Build or reuse the single exact AdaLN schedule held by this model."""
+        if not unique_timestep_plan:
+            raise ValueError("AdaLN schedule must contain at least one denoise step")
+        normalized_key = tuple(tuple(int(value) for value in step) for step in schedule_key)
+        if len(normalized_key) != len(unique_timestep_plan):
+            raise ValueError(
+                f"AdaLN schedule key/plan length mismatch: {len(normalized_key)} != {len(unique_timestep_plan)}"
+            )
+        existing = self._adaln_schedule_cache
+        if existing is not None and existing.key == normalized_key:
+            return existing
+
+        for projection in self._adaln_projections:
+            projection.restore_projection()
+
+        t_embs = tuple(self.time_embedder(timestep.view(-1)) for timestep in unique_timestep_plan)
+        device = t_embs[0].device
+        if any(t_emb.device != device for t_emb in t_embs):
+            raise ValueError("AdaLN timestep embeddings must share one device")
+        step_cursor = torch.zeros((), dtype=torch.long, device=device)
+        tables = tuple(minimax_h3_build_adaln_table(projection, t_embs) for projection in self._adaln_projections)
+        table_bytes = sum(table.numel() * table.element_size() for table in tables)
+        projection_bytes = sum(projection.projection_nbytes for projection in self._adaln_projections)
+        offload_supported = offload_weights and all(
+            projection.can_offload_projection() for projection in self._adaln_projections
+        )
+        if offload_supported and table_bytes >= projection_bytes:
+            for projection in self._adaln_projections:
+                projection.clear_precomputed()
+            self._adaln_schedule_cache = None
+            logger.warning(
+                "MiniMax H3 AdaLN weight offload skipped because table bytes (%d) "
+                "would not save memory versus projection bytes (%d)",
+                table_bytes,
+                projection_bytes,
+            )
+            return None
+
+        # Build every table before changing the live projection path.
+        for projection, table in zip(self._adaln_projections, tables):
+            projection.install_precomputed(table, step_cursor)
+
+        offloaded_projection_bytes = 0
+        if offload_supported:
+            try:
+                offloaded_projection_bytes = sum(
+                    projection.offload_projection() for projection in self._adaln_projections
+                )
+            except Exception:
+                for projection in self._adaln_projections:
+                    projection.restore_projection()
+                logger.exception("MiniMax H3 AdaLN projection offload failed; retaining exact tables and weights")
+        elif offload_weights:
+            logger.warning("MiniMax H3 AdaLN projection offload is unsupported for the active parameter layout")
+
+        if offloaded_projection_bytes:
+            current_omni_platform.empty_cache()
+
+        cache = MiniMaxH3AdalnScheduleCache(
+            key=normalized_key,
+            t_embs=t_embs,
+            step_cursor=step_cursor,
+            steps=len(t_embs),
+            max_unique_timesteps=max(int(t_emb.shape[0]) for t_emb in t_embs),
+            table_bytes=table_bytes,
+            projection_bytes=projection_bytes,
+            offloaded_projection_bytes=offloaded_projection_bytes,
+            net_memory_saved_bytes=offloaded_projection_bytes - table_bytes,
+        )
+        self._adaln_schedule_cache = cache
+        if offloaded_projection_bytes:
+            gib = 1024**3
+            logger.info(
+                "MiniMax H3 AdaLN projection weights offloaded: %.3f GiB weights, "
+                "%.3f GiB tables, %.3f GiB net device-memory reduction",
+                offloaded_projection_bytes / gib,
+                table_bytes / gib,
+                cache.net_memory_saved_bytes / gib,
+            )
+        return cache
+
+    def clear_adaln_schedule_cache(self) -> None:
+        """Drop the model-local table and restore direct AdaLN projections."""
+        if self._adaln_schedule_cache is None:
+            return
+        for projection in self._adaln_projections:
+            projection.clear_precomputed()
+        self._adaln_schedule_cache = None
+
+    def _select_adaln_schedule_step(self, step_index: Any) -> torch.Tensor | None:
+        cache = self._adaln_schedule_cache
+        if cache is None:
+            if step_index is not None:
+                raise ValueError("adaln_step_index was supplied without a prepared AdaLN schedule")
+            return None
+        if isinstance(step_index, bool) or not isinstance(step_index, int):
+            raise ValueError("adaln_step_index must be an integer while the AdaLN schedule is active")
+        if step_index < 0 or step_index >= cache.steps:
+            raise ValueError(f"adaln_step_index {step_index} is outside [0, {cache.steps})")
+        cache.step_cursor.fill_(step_index)
+        return cache.t_embs[step_index]
+
     def load_weights(
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
@@ -1007,6 +1250,7 @@ class MiniMaxH3DiTModel(nn.Module):
         refiner_max_seqlen: int,
         seq_len: int,
         device: torch.device,
+        precomputed_t_emb: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Build packed multimodal embeddings for the TP=1/SP=1 inference path.
 
@@ -1032,7 +1276,7 @@ class MiniMaxH3DiTModel(nn.Module):
         embeddings.index_add_(0, img_pos, video_embed.to(_BF16_DTYPE)[: img_pos.shape[0]])
         embeddings.index_add_(0, audio_pos, audio_embed.to(_BF16_DTYPE)[: audio_pos.shape[0]])
 
-        t_emb = self.time_embedder(unique_timesteps)
+        t_emb = self.time_embedder(unique_timesteps) if precomputed_t_emb is None else precomputed_t_emb
         return embeddings, t_emb
 
     def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1052,6 +1296,8 @@ class MiniMaxH3DiTModel(nn.Module):
                 f"{sorted(_FORWARD_SUPPORTED_KWARGS)}"
             )
 
+        precomputed_t_emb = self._select_adaln_schedule_step(kwargs.get("adaln_step_index"))
+
         x = _required_kwarg(kwargs, "x")
         audio_x = _required_kwarg(kwargs, "audio_x")
         img_position_ids = _required_kwarg(kwargs, "img_position_ids")
@@ -1104,6 +1350,7 @@ class MiniMaxH3DiTModel(nn.Module):
             refiner_max_seqlen=refiner_max,
             seq_len=seq_len,
             device=device,
+            precomputed_t_emb=precomputed_t_emb,
         )
 
         combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)