[Test] Cover MiniMax H3 AdaLN offload ownership

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

recipes/MiniMaxAI/MiniMax-H3.md

@@ -115,6 +115,36 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
 1. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
+### Optional exact-schedule AdaLN cache
+
+H3 can precompute every AdaLN projection for the request's known timestep
+schedule and reuse the exact rows in each denoise step:
+
+```bash
+export VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE=1
+```
+
+The cache key uses the exact float32 timestep bit patterns. A request with the
+same schedule reuses the installed tables; a changed schedule rebuilds them.
+This is disabled by default because the tables consume persistent device
+memory proportional to the number of denoise steps and layers.
+
+After the tables are installed, the original projection weights can also move
+to CPU and release their device storage:
+
+```bash
+export VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE=1
+export VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS=1
+```
+
+Weight offload requires the schedule cache. On a schedule change, the weights
+are restored, the replacement tables are built, and the weights are offloaded
+again. The optimization falls back to direct projection when the tables would
+occupy at least as much memory as the projection weights. DTensor/HSDP
+parameters remain under their distributed owner. With layerwise offload, H3's
+AdaLN projection parameters stay outside the per-block transfer set while the
+exact-schedule cache owns them.
+
 ### Text encoder tensor parallelism
 
 The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by

tests/diffusion/models/minimax_h3/test_minimax_h3_adaln_cache.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_branch(layout: str):
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+        MiniMaxH3DenoiseBranch,
+    )
+
+    update_mask = {
+        "t2va": [True, True],
+        "fl2va": [False, True],
+        "ref2va": [True, True],
+    }[layout]
+    audio_update_mask = {
+        "t2va": [True, True],
+        "fl2va": [True, True],
+        "ref2va": [False, True],
+    }[layout]
+    packed = {
+        "seq_len": 8,
+        "text_pos": torch.tensor([0, 1]),
+        "img_pos": torch.tensor([2, 3]),
+        "audio_pos": torch.tensor([4, 5]),
+        "update_mask": torch.tensor(update_mask),
+        "audio_update_mask": torch.tensor(audio_update_mask),
+        "img_position_ids": torch.zeros(8, 3, dtype=torch.long),
+        "cu_seqlens": torch.tensor([0, 6, 8], dtype=torch.int32),
+    }
+    return MiniMaxH3DenoiseBranch(
+        packed=packed,
+        text_embeddings=torch.zeros(2, 4),
+        token_tags=torch.zeros(8, dtype=torch.long),
+        device=torch.device("cpu"),
+    )
+
+
+def _reference_full_vector_unique(
+    branch,
+    *,
+    sigma_video: float,
+    sigma_audio: float,
+    imgvid_cond_timestep: float,
+    audio_ref_cond_timestep: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    t_video = 1.0 - sigma_video
+    t_audio = 1.0 - sigma_audio
+    timesteps = torch.full((branch.seq_len,), t_video, dtype=torch.float32)
+    timesteps[branch.img_pos[branch.update_mask]] = t_video
+    timesteps[branch.img_pos[~branch.update_mask]] = max(t_video, imgvid_cond_timestep)
+    timesteps[branch.audio_pos[branch.audio_update_mask]] = t_audio
+    timesteps[branch.audio_pos[~branch.audio_update_mask]] = max(t_audio, audio_ref_cond_timestep)
+    return torch.unique(timesteps, sorted=True, return_inverse=True)
+
+
+@pytest.mark.parametrize("layout", ["t2va", "fl2va", "ref2va"])
+def test_timestep_metadata_matches_full_vector_unique(layout: str):
+    branch = _make_branch(layout)
+    sigmas_video = [1.0, 0.5, 0.0]
+    sigmas_audio = [1.0, 0.25, 0.0]
+
+    cached = branch.prepare_timestep_metadata(
+        sigmas_video=sigmas_video,
+        sigmas_audio=sigmas_audio,
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+
+    assert len(cached) == 2
+    for step, metadata in enumerate(cached):
+        expected_unique, expected_inverse = _reference_full_vector_unique(
+            branch,
+            sigma_video=sigmas_video[step],
+            sigma_audio=sigmas_audio[step],
+            imgvid_cond_timestep=0.999,
+            audio_ref_cond_timestep=1.0,
+        )
+        torch.testing.assert_close(metadata.unique_timesteps.cpu(), expected_unique, rtol=0, atol=0)
+        torch.testing.assert_close(metadata.inverse_indices.cpu(), expected_inverse, rtol=0, atol=0)
+
+
+def test_adaln_schedule_key_uses_exact_float32_bits():
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_schedule_key,
+    )
+
+    branch = _make_branch("t2va")
+    first = branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+    same_after_float32_rounding = branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, 0.5 + 1e-12, 0.0],
+        sigmas_audio=[1.0, 0.25 + 1e-12, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+    changed = branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25000003, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+
+    first_key = minimax_h3_adaln_schedule_key(first)
+    assert first_key == minimax_h3_adaln_schedule_key(same_after_float32_rounding)
+    assert first_key != minimax_h3_adaln_schedule_key(changed)
+
+
+class _FakeAdalnLinear(torch.nn.Module):
+    def __init__(self, out_features: int):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.arange(out_features * 4, dtype=torch.bfloat16).reshape(out_features, 4))
+        self.bias = torch.nn.Parameter(torch.arange(out_features, dtype=torch.bfloat16))
+        self.out_features = out_features
+        self.calls = 0
+
+    def forward(self, x):
+        assert self.weight.numel() > 0
+        assert self.bias.numel() > 0
+        self.calls += 1
+        rows = torch.arange(
+            x.shape[0] * self.out_features,
+            device=x.device,
+            dtype=torch.float32,
+        ).reshape(x.shape[0], self.out_features)
+        rows = rows + x.float().sum(dim=-1, keepdim=True)
+        return rows.to(torch.bfloat16), None
+
+
+def _make_tiny_adaln_projection(*, expand_ratio: int = 2, modality_num: int = 3):
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3AdalnProj,
+    )
+
+    projection = object.__new__(MiniMaxH3AdalnProj)
+    torch.nn.Module.__init__(projection)
+    projection.expand_ratio = expand_ratio
+    projection.modality_num = modality_num
+    projection.hidden_size = 2
+    projection.linear = _FakeAdalnLinear(out_features=expand_ratio * modality_num * projection.hidden_size)
+    projection.register_buffer("_precomputed_table", None, persistent=False)
+    projection._precomputed_step = None
+    projection._offloaded_projection_state = None
+    return projection
+
+
+def test_adaln_projection_reuses_exact_precomputed_rows_without_linear_calls():
+    projection = _make_tiny_adaln_projection()
+    first_t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+    second_t_emb = torch.tensor([[4.0, 5.0, 6.0, 7.0], [8.0, 9.0, 10.0, 11.0]])
+    first_direct = projection.compute_flat(first_t_emb)
+    second_direct = projection.compute_flat(second_t_emb)
+    first_padded = torch.cat(
+        [first_direct, first_direct.new_zeros(second_direct.shape[0] - first_direct.shape[0], first_direct.shape[1])]
+    )
+    table = torch.stack([first_padded, second_direct])
+    cursor = torch.zeros((), dtype=torch.long)
+    direct_calls = projection.linear.calls
+
+    projection.install_precomputed(table, cursor)
+    first_cached = projection(first_t_emb)
+    cursor.fill_(1)
+    second_cached = projection(second_t_emb)
+
+    assert projection.linear.calls == direct_calls
+    for actual, expected in zip(first_cached, first_direct.chunk(2, dim=-1)):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    for actual, expected in zip(second_cached, second_direct.chunk(2, dim=-1)):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_adaln_projection_offload_preserves_cache_and_clear_restores_weights():
+    projection = _make_tiny_adaln_projection()
+    t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+    direct = projection.compute_flat(t_emb)
+    expected_state = {name: tensor.clone() for name, tensor in projection.linear.state_dict().items()}
+    expected_bytes = sum(tensor.numel() * tensor.element_size() for tensor in expected_state.values())
+    projection.install_precomputed(direct.unsqueeze(0), torch.zeros((), dtype=torch.long))
+
+    freed_bytes = projection.offload_projection()
+
+    assert freed_bytes == expected_bytes
+    assert all(parameter.numel() == 0 for parameter in projection.linear.parameters())
+    cached = projection(t_emb)
+    for actual, expected in zip(cached, direct.chunk(projection.expand_ratio, dim=-1)):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    projection.clear_precomputed()
+    assert not projection.projection_weights_offloaded
+    for name, tensor in projection.linear.state_dict().items():
+        torch.testing.assert_close(tensor, expected_state[name], rtol=0, atol=0)
+    torch.testing.assert_close(projection.compute_flat(t_emb), direct, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for allocator accounting")
+def test_adaln_projection_offload_releases_cuda_allocator_storage():
+    device = torch.device("cuda")
+    torch.accelerator.empty_cache()
+    projection = _make_tiny_adaln_projection().to(device)
+    t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]], device=device)
+    direct = projection.compute_flat(t_emb)
+    projection.install_precomputed(
+        direct.unsqueeze(0),
+        torch.zeros((), dtype=torch.long, device=device),
+    )
+    torch.accelerator.synchronize()
+    allocated_before = torch.cuda.memory_allocated(device)
+
+    freed_bytes = projection.offload_projection()
+    torch.accelerator.synchronize()
+    allocated_after = torch.cuda.memory_allocated(device)
+
+    assert freed_bytes > 0
+    assert allocated_before - allocated_after >= freed_bytes
+    assert all(parameter.is_cuda and parameter.numel() == 0 for parameter in projection.linear.parameters())
+
+    projection.clear_precomputed()
+    assert all(parameter.is_cuda and parameter.numel() > 0 for parameter in projection.linear.parameters())
+
+
+class _CountingTimeEmbedder(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.calls = 0
+
+    def forward(self, timesteps):
+        self.calls += 1
+        return timesteps[:, None]
+
+
+class _TinyAdalnBlock(torch.nn.Module):
+    def __init__(self, *, modality_num: int = 3):
+        super().__init__()
+        self.adaln_proj = _make_tiny_adaln_projection(modality_num=modality_num)
+
+
+def _make_adaln_schedule_model():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    torch.nn.Module.__init__(model)
+    model.time_embedder = _CountingTimeEmbedder()
+    model.blocks = torch.nn.ModuleList([_TinyAdalnBlock(), _TinyAdalnBlock()])
+    model.final_layer = _TinyAdalnBlock(modality_num=1)
+    model._adaln_projections = tuple(block.adaln_proj for block in model.blocks) + (model.final_layer.adaln_proj,)
+    model._adaln_schedule_cache = None
+    return model
+
+
+def _schedule(branch, *, middle_video: float, middle_audio: float):
+    return branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, middle_video, 0.0],
+        sigmas_audio=[1.0, middle_audio, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+
+
+def test_model_cache_hit_and_changed_schedule_rebuild_with_weight_offload():
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_schedule_key,
+    )
+
+    branch = _make_branch("t2va")
+    first_plan = _schedule(branch, middle_video=0.5, middle_audio=0.25)
+    changed_plan = _schedule(branch, middle_video=0.4, middle_audio=0.2)
+    model = _make_adaln_schedule_model()
+    expected_projection_bytes = sum(
+        parameter.numel() * parameter.element_size()
+        for projection in model._adaln_projections
+        for parameter in projection.linear.parameters()
+    )
+
+    first = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in first_plan),
+        schedule_key=minimax_h3_adaln_schedule_key(first_plan),
+        offload_weights=True,
+    )
+    calls_after_first = [projection.linear.calls for projection in model._adaln_projections]
+    hit = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in first_plan),
+        schedule_key=minimax_h3_adaln_schedule_key(first_plan),
+        offload_weights=True,
+    )
+    replaced = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in changed_plan),
+        schedule_key=minimax_h3_adaln_schedule_key(changed_plan),
+        offload_weights=True,
+    )
+
+    assert first is hit
+    assert replaced is not None and replaced is not first
+    assert first is not None
+    assert first.offloaded_projection_bytes == expected_projection_bytes
+    assert first.net_memory_saved_bytes == expected_projection_bytes - first.table_bytes
+    assert calls_after_first == [2, 2, 2]
+    assert [projection.linear.calls for projection in model._adaln_projections] == [4, 4, 4]
+    assert all(projection.projection_weights_offloaded for projection in model._adaln_projections)
+
+
+def test_model_does_not_manually_offload_dtensor_owned_parameters(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_schedule_key,
+    )
+
+    class FakeDTensor(torch.nn.Parameter):
+        def to_local(self):
+            return self
+
+    monkeypatch.setattr(minimax_h3_transformer, "DTensor", FakeDTensor)
+    branch = _make_branch("t2va")
+    plan = _schedule(branch, middle_video=0.5, middle_audio=0.25)
+    model = _make_adaln_schedule_model()
+    projection = model._adaln_projections[0]
+    projection.linear.weight = FakeDTensor(projection.linear.weight.detach())
+
+    cache = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in plan),
+        schedule_key=minimax_h3_adaln_schedule_key(plan),
+        offload_weights=True,
+    )
+
+    assert cache is not None
+    assert cache.offloaded_projection_bytes == 0
+    assert all(not item.projection_weights_offloaded for item in model._adaln_projections)
+    assert projection.linear.weight.numel() > 0
+
+
+def test_model_weight_offload_falls_back_when_tables_are_larger():
+    model = _make_adaln_schedule_model()
+    unique_timestep_plan = tuple(torch.tensor([float(step)]) for step in range(12))
+
+    cache = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=unique_timestep_plan,
+        schedule_key=tuple((step,) for step in range(12)),
+        offload_weights=True,
+    )
+
+    assert cache is None
+    assert all(projection._precomputed_table is None for projection in model._adaln_projections)
+    assert all(not projection.projection_weights_offloaded for projection in model._adaln_projections)
+
+
+def test_denoise_loop_prepares_cache_once_and_passes_step_indices(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
+
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1")
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS", "1")
+
+    class FakeModel:
+        def __init__(self):
+            self.prepare_calls = 0
+            self.prepare_for_request_calls = 0
+            self.step_indices = []
+
+        def prepare_for_request_caches(self):
+            assert torch.is_inference_mode_enabled()
+            self.prepare_for_request_calls += 1
+
+        def prepare_adaln_schedule_cache(self, *, unique_timestep_plan, schedule_key, offload_weights):
+            assert torch.is_inference_mode_enabled()
+            assert len(unique_timestep_plan) == len(schedule_key) == 2
+            assert offload_weights
+            self.prepare_calls += 1
+            return object()
+
+        def __call__(self, **kwargs):
+            self.step_indices.append(kwargs["adaln_step_index"])
+            return torch.zeros(2, 96), torch.zeros(2, 32)
+
+    model = FakeModel()
+    denoise_loop.minimax_h3_denoise_loop(
+        model=model,
+        positive=_make_branch("t2va"),
+        initial_video_rows=torch.zeros(2, 96),
+        initial_audio_rows=torch.zeros(2, 32),
+        keyframe_cond_rows=None,
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25, 0.0],
+        device=torch.device("cpu"),
+    )
+
+    assert model.prepare_for_request_calls == 1
+    assert model.prepare_calls == 1
+    assert model.step_indices == [0, 1]
+
+
+def test_weight_offload_env_requires_schedule_cache(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_weight_offload_enabled,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", raising=False)
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS", "1")
+    assert not minimax_h3_adaln_weight_offload_enabled()
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1")
+    assert minimax_h3_adaln_weight_offload_enabled()

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -25,7 +25,7 @@ def test_grouped_qkv_checkpoint_reorder():
     assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
 
 
-def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
+def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(monkeypatch):
     from cache_dit import ForwardPattern
 
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
@@ -38,6 +38,11 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
     assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
     assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
 
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1")
+    assert MiniMaxH3DiTModel.layerwise_offload_excluded_parameter_prefixes() == ("adaln_proj.linear.",)
+    monkeypatch.delenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE")
+    assert MiniMaxH3DiTModel.layerwise_offload_excluded_parameter_prefixes() == ()
+
     model = object.__new__(MiniMaxH3DiTModel)
     nn.Module.__init__(model)
     model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
@@ -61,6 +66,14 @@ def test_packed_attention_is_a_regional_compile_boundary():
     assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
 
 
+def test_adaln_projection_is_a_regional_compile_boundary():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3AdalnProj,
+    )
+
+    assert getattr(MiniMaxH3AdalnProj.forward, "_torchdynamo_disable", False)
+
+
 @pytest.mark.parametrize(
     ("tp_size", "message"),
     [

tests/diffusion/offloader/test_layerwise_backend.py

@@ -135,6 +135,14 @@ class _DummyBlock(nn.Module):
         self.weight = nn.Parameter(torch.randn(10, 10))
 
 
+class _OwnedProjectionBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        self.adaln_proj = nn.Module()
+        self.adaln_proj.linear = nn.Linear(2, 3)
+
+
 class _SingleBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["blocks"]
 
@@ -224,6 +232,49 @@ class TestGetBlocksFromDit:
         assert len(blocks) == 2
 
 
+def test_layerwise_hook_leaves_excluded_parameters_under_model_ownership(patched_offload_runtime):
+    current_block = _OwnedProjectionBlock()
+    next_block = _OwnedProjectionBlock()
+    current_adaln = {name: tensor.clone() for name, tensor in current_block.adaln_proj.linear.state_dict().items()}
+    next_adaln = {name: tensor.clone() for name, tensor in next_block.adaln_proj.linear.state_dict().items()}
+    hook = LayerwiseOffloadHook(
+        next_block=next_block,
+        device=torch.device("cpu"),
+        stream=DummyStream(),
+        pin_memory=False,
+        excluded_parameter_prefixes=("adaln_proj.linear.",),
+    )
+
+    hook.initialize_hook(current_block)
+
+    assert set(hook.block_parameters) == {"weight"}
+    assert set(hook.next_block_parameters) == {"weight"}
+    assert next_block.weight.numel() == 0
+    for name, tensor in next_block.adaln_proj.linear.state_dict().items():
+        torch.testing.assert_close(tensor, next_adaln[name], rtol=0, atol=0)
+
+    hook.offload_layer()
+    assert current_block.weight.numel() == 0
+    for name, tensor in current_block.adaln_proj.linear.state_dict().items():
+        torch.testing.assert_close(tensor, current_adaln[name], rtol=0, atol=0)
+
+
+def test_layerwise_backend_materializes_excluded_parameters_on_target_device():
+    block = _OwnedProjectionBlock()
+    expected_bytes = sum(
+        parameter.numel() * parameter.element_size() for parameter in block.adaln_proj.linear.parameters()
+    )
+
+    moved_bytes = LayerWiseOffloadBackend._move_excluded_parameters_to_device(
+        block,
+        torch.device("cpu"),
+        ("adaln_proj.linear.",),
+    )
+
+    assert moved_bytes == expected_bytes
+    assert all(parameter.device.type == "cpu" for parameter in block.adaln_proj.linear.parameters())
+
+
 class TestGetBlocksAttrNames:
     def test_get_blocks_attr_names_new_format(self):
         model = _MultiBlockModel()

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)

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -41,6 +41,7 @@ class LayerwiseOffloadHook(ModelHook):
         device: torch.device,
         stream: current_omni_platform.Stream | None = None,
         pin_memory: bool = True,
+        excluded_parameter_prefixes: tuple[str, ...] = (),
     ):
         assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"
 
@@ -48,6 +49,7 @@ class LayerwiseOffloadHook(ModelHook):
         self.device = device
         self.copy_stream = stream or current_omni_platform.current_stream()
         self.pin_memory = pin_memory
+        self.excluded_parameter_prefixes = tuple(excluded_parameter_prefixes)
 
         # Per-block synchronization primitive: set after H2D copy completes.
         self._prefetch_done: current_omni_platform.Event | None = None
@@ -90,10 +92,10 @@ class LayerwiseOffloadHook(ModelHook):
         # the input module is kept intact
         module = super().initialize_hook(module)
 
-        self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters())
+        self.block_parameters = self._owned_parameters(module)
         self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers())
 
-        self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters())
+        self.next_block_parameters = self._owned_parameters(self.next_block)
         self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers())
 
         # Pre-allocate gpu tensors in a flattened way
@@ -106,6 +108,13 @@ class LayerwiseOffloadHook(ModelHook):
 
         return module
 
+    def _owned_parameters(self, module: nn.Module) -> dict[str, nn.Parameter]:
+        return {
+            name: parameter
+            for name, parameter in module.named_parameters()
+            if not name.startswith(self.excluded_parameter_prefixes)
+        }
+
     @staticmethod
     def _to_cpu(
         params: dict[str, nn.Parameter],
@@ -255,9 +264,10 @@ def apply_block_hook(
     device: torch.device,
     stream: current_omni_platform.Stream | None = None,
     pin_memory: bool = True,
+    excluded_parameter_prefixes: tuple[str, ...] = (),
 ) -> LayerwiseOffloadHook:
     registry = HookRegistry.get_or_create(module)
-    hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory)
+    hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory, excluded_parameter_prefixes)
     registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook)
 
     return hook
@@ -284,6 +294,20 @@ class LayerWiseOffloadBackend(OffloadBackend):
         self.copy_stream = current_omni_platform.Stream()
         self._blocks: list[list[nn.Module]] = []
 
+    @staticmethod
+    def _move_excluded_parameters_to_device(
+        module: nn.Module,
+        device: torch.device,
+        excluded_parameter_prefixes: tuple[str, ...],
+    ) -> int:
+        moved_bytes = 0
+        for name, parameter in module.named_parameters():
+            if not name.startswith(excluded_parameter_prefixes):
+                continue
+            moved_bytes += parameter.numel() * parameter.element_size()
+            parameter.data = parameter.data.to(device, non_blocking=True)
+        return moved_bytes
+
     def enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("LayerWiseOffloadBackend already enabled")
@@ -341,6 +365,30 @@ class LayerWiseOffloadBackend(OffloadBackend):
                 dit_module.to(self.device)
                 continue
 
+            excluded_parameter_prefixes: tuple[str, ...] = ()
+            get_excluded_prefixes = getattr(dit_module, "layerwise_offload_excluded_parameter_prefixes", None)
+            if callable(get_excluded_prefixes):
+                excluded_parameter_prefixes = tuple(get_excluded_prefixes())
+                if excluded_parameter_prefixes:
+                    logger.info(
+                        "Leaving layer-wise parameters matching %s under %s ownership",
+                        excluded_parameter_prefixes,
+                        dit_module.__class__.__name__,
+                    )
+                    moved_bytes = sum(
+                        self._move_excluded_parameters_to_device(
+                            block,
+                            self.device,
+                            excluded_parameter_prefixes,
+                        )
+                        for block in blocks
+                    )
+                    logger.info(
+                        "Materialized %.3f GiB of model-owned layer-wise parameters on %s",
+                        moved_bytes / 1024**3,
+                        self.device,
+                    )
+
             # Move non-block modules to GPU (they stay resident)
             for name, m in dit_module.named_children():
                 if name not in blocks_attr_names:
@@ -368,6 +416,7 @@ class LayerWiseOffloadBackend(OffloadBackend):
                 self.device,
                 self.copy_stream,
                 self.config.pin_cpu_memory,
+                excluded_parameter_prefixes,
             )
             last_hook.prefetch_layer(non_blocking=False)
 
@@ -381,6 +430,7 @@ class LayerWiseOffloadBackend(OffloadBackend):
                     self.device,
                     self.copy_stream,
                     self.config.pin_cpu_memory,
+                    excluded_parameter_prefixes,
                 )
                 block_hooks.append(hook)