From 535bd35b4252e3c031d4b64c60a6250cb78a579c Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:30:42 +0800 Subject: [PATCH] [Perf] Add reusable exact projection caching with MiniMax-H3 integration (#7987) Signed-off-by: lishunyang12 Signed-off-by: Armaan Amatya Co-authored-by: lishunyang12 Co-authored-by: Armaan Amatya <38102001+armaanamatya@users.noreply.github.com> --- .../diffusion/diffusion_model_integration.md | 55 +++ recipes/MiniMaxAI/MiniMax-H3.md | 78 ++++ .../cache/test_exact_projection_cache.py | 270 +++++++++++ .../models/minimax_h3/test_adaln_cache.py | 439 ++++++++++++++++++ .../minimax_h3/test_adaln_cache_device.py | 127 +++++ tools/minimax_h3/build_adaln_cache.py | 132 ++++++ vllm_omni/diffusion/cache/__init__.py | 4 +- .../diffusion/cache/exact_projection_cache.py | 180 +++++++ .../models/minimax_h3/adaln_cache.py | 438 +++++++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 42 +- .../models/minimax_h3/pipeline_minimax_h3.py | 124 +++++ 11 files changed, 1886 insertions(+), 3 deletions(-) create mode 100644 tests/diffusion/cache/test_exact_projection_cache.py create mode 100644 tests/diffusion/models/minimax_h3/test_adaln_cache.py create mode 100644 tests/diffusion/models/minimax_h3/test_adaln_cache_device.py create mode 100644 tools/minimax_h3/build_adaln_cache.py create mode 100644 vllm_omni/diffusion/cache/exact_projection_cache.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/adaln_cache.py diff --git a/docs/design/module/diffusion/diffusion_model_integration.md b/docs/design/module/diffusion/diffusion_model_integration.md index d6eb44561..1143246aa 100644 --- a/docs/design/module/diffusion/diffusion_model_integration.md +++ b/docs/design/module/diffusion/diffusion_model_integration.md @@ -57,3 +57,58 @@ cross-stage routing. Test registry selection, checkpoint loading, minimal inference, input and output contracts, and every declared optional capability. + +## Exact conditioning projection reuse + +Models with repeated conditioning projections can use `ExactProjectionCache` +from `vllm_omni.diffusion.cache`. It stores actual projection outputs for identical +conditioning tensors and unchanged weights, with a byte budget and TP-wide hit +coordination. It has no model names, sampling schedules, projection layouts, or +artifact formats built in. MiniMax-H3 is its first model integration. + +Keep one cache per model instance and prepare each conditioning tensor once +before its projections. For a leaf linear with fixed SiLU preprocessing: + +```python +from torch.nn.functional import silu + +from vllm_omni.diffusion.cache import ExactProjectionCache + +cache = ExactProjectionCache(max_bytes=256 * 1024**2) + +# Inside an inference forward, with gradients disabled: +cache.prepare(conditioning) +modulation = cache.project( + "block.0.modulation", + projection, + conditioning, + lambda: projection(silu(conditioning)), +) +``` + +The original computation must depend only on the supplied conditioning tensor, +the leaf projection's weights/buffers, and fixed preprocessing. Include prompt, +guidance, modality, or other request-dependent conditioning in the actual input; +a timestep-only key is unsafe when those affect the result. Do not mutate the +input between `prepare` and `project`, or modify outputs returned on cache hits. +Use separate cache instances for concurrent forwards with different inputs. + +Call `clear()` when loading or switching adapters, changing projection behavior, +or moving the model. Weight-version checks invalidate individual entries. +Compilation, gradient-enabled execution, projection hooks, and unversioned +inference weights use the original computation. Validate cold/warm parity, +invalidation, and rank-local misses when integrating another model. +Hook checks include global hooks and hooks on wrapped submodules. Numerical +settings are captured at each projection call, so entering an autocast context +after `prepare()` cannot reuse an output from a different precision. + +Validated offline results can be supplied through `_lookup_precomputed()`; +this extension is consulted only at TP1 on runtime-cache misses. Artifact +validation, weight binding, and model-specific layouts remain in the adapter. +H3's adapter retains its main/Ref2VA sidecars and their eager-only admission. + +The cache retains projection results and leaves all original weights available. +It does not reduce weight residency or cancel layerwise prefetch. A future +integration that skips weight transfers must coordinate with `OffloadPlan` +before prefetch and provide a safe miss path; moving weight storage currently +causes conservative cache invalidation. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 9eac85ece..dbd5de289 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -1017,6 +1017,84 @@ balanced switch order. > workload. The values above apply to this deployment and are not universal > guarantees. `lossless` remains the exact reference path. +## Exact AdaLN reuse + +H3 enables exact AdaLN projection reuse by default, independently of the task, +hardware, sampling schedule, quantization setting, and request `quality`. +This applies to both original H3 (including its default 50-point schedule) and +FastH3; no distilled adapter or explicit few-step ladder is required. +The first occurrence of an input runs the existing projection; later identical +time embeddings reuse its output only while the projection weights and numerical +settings remain unchanged. This does not approximate neighboring timesteps or +change attention, precision, sampling, or the generated audio/video contract. + +Each DiT keeps at most 256 MiB of runtime projection outputs. If a long schedule +exceeds that budget, it retains a reusable subset and computes other entries +normally, avoiding cyclic eviction on repeated original-H3 requests. All original weights +remain loaded so new schedules and adapters can compute normally. Adapter changes, +weight reloads, and model device moves invalidate the cache; parameter versions +also guard individual entries. TP ranks coordinate hits before skipping a +projection collective. Gradient-enabled execution, compilation, custom linear +hooks, and tensors without weight version counters use the original computation. +Offload paths that replace weight storage can therefore reduce the hit rate. + +The runtime cache uses the shared `ExactProjectionCache` implementation; H3 keeps +only its optional sidecar adaptation. Other models can integrate the same +[projection cache interface](../../docs/design/module/diffusion/diffusion_model_integration.md#exact-conditioning-projection-reuse). +This cache retains projection outputs, not offloaded weights, and does not skip +block weight prefetch. + +For an A/B comparison, disable only this reuse at server startup: + +```bash +--cache-config '{"minimax_h3_adaln_cache": false}' +``` + +### Optional offline sidecar + +An offline sidecar can seed the first projection results; it is not required to +enable the default cache. Sidecars require `--enforce-eager`, native BF16 TP1 +math, and the same numerical environment as the builder. With the default +compiled execution, sidecars are rejected before reading their payloads: compiled +H3 blocks bypass cached projections, so retaining those payloads would waste GPU +memory. This applies to both the main and Ref2VA sidecars. Other serving +configurations retain the default runtime-cache behavior described above. +From the repository root, for a fixed FastH3 adapter and its own four-step schedule: + +```bash +PYTHONPATH=. python tools/minimax_h3/build_adaln_cache.py \ + --transformer-path "${MODEL_ROOT}/FL2VA/transformer" \ + --model-variant fl2va --mode t2va \ + --fasth3-adapter "${FASTH3_ADAPTER}" \ + --num-inference-steps 4 --flow-shift 12 --audio-flow-shift 3 \ + --device cuda --output /path/to/h3-adaln.safetensors +``` + +Omit `--fasth3-adapter` for base weights and use the serving step count and shifts. +The builder accepts a native transformer directory with `config.json` and indexed +or single-file safetensors. It streams the required inputs and refuses to overwrite +an existing output. It does not instantiate the full DiT. + +Pass the resulting local artifact at eager server startup: + +```bash +--enforce-eager \ +--cache-config '{"minimax_h3_adaln_cache_path": "/path/to/h3-adaln.safetensors"}' +``` + +Combined servers can also use `minimax_h3_ref_adaln_cache_path` for their Ref2VA +transformer. Each artifact binds the model architecture, task, schedule, shifts, +fixed adapter file, effective time-embedding/AdaLN weights, payload checksums, and +numerical environment. A missing, corrupt, incompatible, or outdated sidecar is +rejected with a warning; the model still loads every weight and computes through +the runtime path. A request with different settings similarly falls back. +Build sidecars from trusted local inputs: checksums verify identity and integrity, +not the correctness of an untrusted generator. + +This implementation saves repeated projection work. It does not remove AdaLN +weights or claim a GPU memory reduction. End-to-end gains depend on the workload, +offload behavior, embedding fingerprint cost, and TP coordination overhead. + ## LoRA ### Turbo LoRA diff --git a/tests/diffusion/cache/test_exact_projection_cache.py b/tests/diffusion/cache/test_exact_projection_cache.py new file mode 100644 index 000000000..cae9e97f4 --- /dev/null +++ b/tests/diffusion/cache/test_exact_projection_cache.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact projection reuse without model-specific schedules or artifacts.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.cache.exact_projection_cache import ExactProjectionCache + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def tp1(monkeypatch): + import vllm.distributed + + monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 1) + + +@pytest.fixture(autouse=True) +def deterministic_weights(): + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + yield + + +@pytest.mark.parametrize("batch,in_features,out_features", [(2, 4, 8), (1, 8, 48), (4, 6, 18)]) +def test_runtime_reuses_exact_outputs_and_tracks_changed_weights(mocker, batch, in_features, out_features): + cache = ExactProjectionCache() + linear = torch.nn.Linear(in_features, out_features) + x = torch.randn(batch, in_features) + compute = mocker.Mock(side_effect=lambda: linear(x)) + with torch.no_grad(): + cache.prepare(x) + first = cache.project("block", linear, x, compute) + first_copy = first.clone() + first.zero_() # The caller/provider cannot corrupt the owned cache. + assert torch.equal(cache.project("block", linear, x, compute), first_copy) + assert compute.call_count == 1 + linear.weight.add_(0.25) + assert torch.equal(cache.project("block", linear, x, compute), linear(x)) + assert compute.call_count == 2 + other = x.clone() + other[0, 0] += 1 + cache.prepare(other) + cache.project("block", linear, other, lambda: linear(other)) + cache.prepare(x) + cache.project("block", linear, x, compute) + assert compute.call_count == 2 + assert cache.hits == 2 + + +def test_runtime_budget_retains_reusable_entries_without_schedule_thrashing(mocker): + cache = ExactProjectionCache(max_bytes=32) + linear = torch.nn.Linear(2, 4) + x = torch.ones(2, 2) + compute = mocker.Mock(side_effect=lambda: linear(x)) + with torch.no_grad(): + cache.prepare(x) + for name in ("first", "second", "first"): + cache.project(name, linear, x, compute) + assert cache._bytes <= 32 + assert compute.call_count == 2 + assert cache.hits == 1 + cache.clear() + assert cache._bytes == 0 and not cache._entries + + +def test_dynamic_lora_buffers_and_suspend_mask_invalidate_entries(mocker): + cache = ExactProjectionCache() + linear = torch.nn.Linear(2, 2) + linear.lora_a_stacked = (torch.ones(1, 2),) + linear.lora_b_stacked = (torch.ones(2, 1),) + linear._diffusion_lora_active_slices = (True,) + x = torch.ones(1, 2) + + def compute(): + result = linear(x) + if linear._diffusion_lora_active_slices[0]: + result = result + x @ linear.lora_a_stacked[0].T @ linear.lora_b_stacked[0].T + return result + + counted = mocker.Mock(side_effect=compute) + with torch.no_grad(): + cache.prepare(x) + cache.project("a", linear, x, counted) + cache.project("a", linear, x, counted) + assert counted.call_count == 1 + linear.lora_b_stacked[0].add_(1) + assert torch.equal(cache.project("a", linear, x, counted), compute()) + linear._diffusion_lora_active_slices = (False,) + assert torch.equal(cache.project("a", linear, x, counted), compute()) + assert counted.call_count == 3 + + +@pytest.mark.parametrize("scope", ["local", "global", "child"]) +@pytest.mark.parametrize("kind", ["pre", "post"]) +def test_linear_hooks_run_even_after_cache_warmup(scope, kind): + cache = ExactProjectionCache() + linear = torch.nn.Linear(2, 2) + projection = torch.nn.Sequential(linear) if scope == "child" else linear + x = torch.ones(1, 2) + with torch.no_grad(): + cache.prepare(x) + expected = projection(x) + cache.project("a", projection, x, lambda: projection(x)) + calls = [] + + def pre_hook(module, args): + if module is linear: + calls.append(module) + return (args[0] + 1,) + + def post_hook(module, args, output): + if module is linear: + calls.append(module) + return output + 1 + + if kind == "pre": + expected = linear(x + 1) + register = ( + torch.nn.modules.module.register_module_forward_pre_hook + if scope == "global" + else linear.register_forward_pre_hook + ) + hook = register(pre_hook) + else: + expected = expected + 1 + register = ( + torch.nn.modules.module.register_module_forward_hook + if scope == "global" + else linear.register_forward_hook + ) + hook = register(post_hook) + try: + for _ in range(2): + assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), expected) + assert len(calls) == 2 + finally: + hook.remove() + assert torch.equal(cache.project("a", projection, x, lambda: projection(x)), projection(x)) + + +@pytest.mark.parametrize("autocast_at_prepare", [False, True]) +def test_numerical_settings_follow_projection_context(mocker, autocast_at_prepare): + cache = ExactProjectionCache() + linear = torch.nn.Linear(4, 4) + x = torch.randn(2, 4) + compute = mocker.Mock(side_effect=lambda: linear(x)) + with torch.no_grad(): + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast_at_prepare): + cache.prepare(x) + for enabled in (False, True, False, True): + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=enabled): + actual = cache.project("a", linear, x, compute) + expected = linear(x) + assert actual.dtype == expected.dtype + assert torch.equal(actual, expected) + assert compute.call_count == 2 + + +@pytest.mark.parametrize( + "setting", + ["allow_tf32", "allow_bf16_reduced_precision_reduction", "allow_fp16_reduced_precision_reduction"], +) +def test_matmul_settings_invalidate_prepared_projection(monkeypatch, mocker, setting): + cache = ExactProjectionCache() + linear = torch.nn.Linear(2, 2) + x = torch.ones(1, 2) + compute = mocker.Mock(side_effect=lambda: linear(x)) + with torch.no_grad(): + cache.prepare(x) + cache.project("a", linear, x, compute) + monkeypatch.setattr(torch.backends.cuda.matmul, setting, not getattr(torch.backends.cuda.matmul, setting)) + assert torch.equal(cache.project("a", linear, x, compute), linear(x)) + assert compute.call_count == 2 + + +def test_disabled_and_grad_paths_keep_original_computation(monkeypatch, mocker): + import vllm_omni.diffusion.cache.exact_projection_cache as module + + monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash"))) + linear = torch.nn.Linear(2, 2) + x = torch.ones(1, 2, requires_grad=True) + for cache in (ExactProjectionCache(max_bytes=0), ExactProjectionCache()): + cache.prepare(x) + cache.project("a", linear, x, lambda: linear(x)).sum().backward() + assert x.grad is not None + with torch.no_grad(): + cache = ExactProjectionCache(max_bytes=0) + cache.prepare(x) + assert torch.equal(cache.project("a", linear, x, lambda: linear(x)), linear(x)) + + +def test_tp_rank_miss_forces_every_rank_to_recompute(monkeypatch, mocker): + import vllm.distributed + + cache = ExactProjectionCache() + linear = torch.nn.Linear(2, 2) + x = torch.ones(1, 2) + compute = mocker.Mock(side_effect=lambda: linear(x)) + with torch.no_grad(): + cache.prepare(x) + cache.project("a", linear, x, compute) + votes = mocker.Mock(side_effect=[torch.tensor([1]), torch.tensor([2])]) + monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes)) + cache.project("a", linear, x, compute) + assert compute.call_count == 2 + cache.project("a", linear, x, compute) + assert compute.call_count == 2 + assert votes.call_count == 2 + + +def test_compiled_execution_bypasses_runtime_and_precomputed_results(monkeypatch, mocker): + import vllm_omni.diffusion.cache.exact_projection_cache as module + + cache = ExactProjectionCache() + linear = torch.nn.Linear(2, 2) + x = torch.ones(1, 2) + compute = mocker.Mock(side_effect=lambda: linear(x)) + with torch.no_grad(): + cache.prepare(x) + cache.project("modulation", linear, x, compute) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + monkeypatch.setattr(module, "tensor_digest", mocker.Mock(side_effect=AssertionError("must not hash"))) + seeded = mocker.patch.object(cache, "_lookup_precomputed", side_effect=AssertionError("must not seed")) + # An input prepared before compilation must also bypass the warm entry. + assert torch.equal(cache.project("modulation", linear, x, compute), linear(x)) + cache.prepare(x) + assert torch.equal(cache.project("modulation", linear, x, compute), linear(x)) + assert compute.call_count == 3 + seeded.assert_not_called() + + +@pytest.mark.parametrize("world_size", [1, 2]) +def test_precomputed_extension_is_used_only_at_tp1(monkeypatch, mocker, world_size): + import vllm.distributed + + cache = ExactProjectionCache() + linear = torch.nn.Linear(2, 2) + x = torch.ones(1, 2) + monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: world_size) + votes = mocker.Mock(side_effect=[torch.tensor([0]), torch.tensor([world_size])]) + monkeypatch.setattr(vllm.distributed, "get_tp_group", lambda: SimpleNamespace(all_reduce=votes)) + with torch.no_grad(): + expected = linear(x) + seeded = mocker.patch.object(cache, "_lookup_precomputed", return_value=expected) + compute = mocker.Mock(side_effect=lambda: linear(x)) + cache.prepare(x) + for _ in range(2): + assert torch.equal(cache.project("conditioning.scale_shift", linear, x, compute), expected) + assert seeded.call_count == (1 if world_size == 1 else 0) + assert compute.call_count == (0 if world_size == 1 else 1) + assert cache.hits == 1 + + +def test_inference_weights_without_versions_use_safe_runtime_path(mocker): + with torch.inference_mode(): + linear = torch.nn.Linear(2, 2) + x = torch.ones(1, 2) + cache = ExactProjectionCache() + compute = mocker.Mock(side_effect=lambda: linear(x)) + cache.prepare(x) + cache.project("a", linear, x, compute) + linear.weight.add_(1) + assert torch.equal(cache.project("a", linear, x, compute), linear(x)) + assert compute.call_count == 2 and not cache._entries diff --git a/tests/diffusion/models/minimax_h3/test_adaln_cache.py b/tests/diffusion/models/minimax_h3/test_adaln_cache.py new file mode 100644 index 000000000..03d4fcd37 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_adaln_cache.py @@ -0,0 +1,439 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import copy +from dataclasses import asdict +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 +from vllm_omni.diffusion.models.minimax_h3.adaln_cache import ( + FORMAT_VERSION, + MODES, + MiniMaxH3AdalnCache, + MiniMaxH3RuntimeAdalnCache, + build_cache, + canonical_json, + input_names, + schedule_contract, + tensor_digest, + timestep_plans, +) +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def tp1(monkeypatch): + import vllm.distributed + + monkeypatch.setattr(vllm.distributed, "get_tensor_model_parallel_world_size", lambda: 1) + + +def _contract(mode="t2va", **updates): + values = dict( + mode=mode, num_steps=4, base_schedule=(0.999, 0.749, 0.5, 0.25, 0.0), flow_shift=12.0, audio_flow_shift=3.0 + ) + return schedule_contract(**(values | updates)) + + +def _fixture(tmp_path, mode="t2va", adapter=None, *, base_model=False): + arch = h3.MiniMaxH3DiTArchConfig( + num_layers=2, hidden_size=8, timestep_input_dim=4, time_embed_hidden_size=8, time_embed_dim=4 + ) + generator = torch.Generator().manual_seed(4) + weights = {} + for name, shape in { + "time_embedder.proj_in.weight": (8, 4), + "time_embedder.proj_in.bias": (8,), + "time_embedder.proj_out.weight": (4, 8), + "time_embedder.proj_out.bias": (4,), + }.items(): + weights[name] = torch.randn(shape, generator=generator) + for name in sorted(input_names(arch.num_layers)): + if name.startswith("time_embedder."): + continue + width = (2 if name.startswith("final_layer.") else 18) * arch.hidden_size + shape = (width, arch.time_embed_dim) if name.endswith("weight") else (width,) + weights[name] = torch.randn(shape, generator=generator).to(torch.bfloat16) + variant = "ref2va" if mode.startswith("ref2va") else "fl2va" + payload, manifest = build_cache( + arch, + weights.items(), + contract=_contract(mode, **({"num_steps": 50, "base_schedule": None} if base_model else {})), + model_variant=variant, + adapter_sha256=adapter, + device=torch.device("cpu"), + ) + path = tmp_path / "adaln.safetensors" + _save(path, payload, manifest) + return arch, weights, payload, manifest, path + + +def _save(path, payload, manifest): + save_file(payload, str(path), metadata={"format_version": FORMAT_VERSION, "manifest": canonical_json(manifest)}) + + +def _ready(arch, weights, path, variant="fl2va", adapter=None): + cache = MiniMaxH3AdalnCache(arch, path=str(path), model_variant=variant) + cache.bind_adapter(adapter) + for name, weight in weights.items(): + cache.verify_weight(name, weight) + cache.finish_loading(torch.device("cpu")) + return cache + + +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("base_model", [False, True], ids=["explicit-ladder", "base-h3-50"]) +def test_builder_projection_parity_all_tasks(tmp_path, mode, base_model): + arch, weights, payload, manifest, path = _fixture(tmp_path, mode, base_model=base_model) + cache = _ready(arch, weights, path, manifest["model_variant"]) + for i, plan in enumerate(timestep_plans(manifest["schedule"])): + blocks, final = cache.lookup(plan) + embedding = payload["time_embeddings"][i, : len(plan)] + activated = torch.nn.functional.silu(embedding).to(torch.bfloat16) + for block_index, pieces in enumerate(blocks): + prefix = f"blocks.{block_index}.adaln_proj.linear" + expected = torch.nn.functional.linear(activated, weights[prefix + ".weight"], weights[prefix + ".bias"]) + actual = torch.stack(pieces, dim=1).reshape(len(plan), -1) + assert torch.equal(actual, expected) + prefix = "final_layer.adaln_proj.linear" + expected = torch.nn.functional.linear(activated, weights[prefix + ".weight"], weights[prefix + ".bias"]) + assert torch.equal(torch.stack(final, dim=1).reshape(len(plan), -1), expected) + assert not cache.state_dict() + + +@pytest.mark.parametrize("change", ["variant", "architecture", "weights", "payload", "plans", "nan", "version"]) +def test_reject_invalid_sidecar(tmp_path, change): + arch, weights, payload, manifest, path = _fixture(tmp_path) + payload = {key: value.clone() for key, value in payload.items()} + manifest = copy.deepcopy(manifest) + if change == "variant": + manifest["model_variant"] = "ref2va" + elif change == "architecture": + manifest["architecture"]["hidden_size"] += 1 + elif change == "weights": + manifest["weights"].pop(next(iter(weights))) + elif change == "payload": + payload["block_params"][0, 0, 0, 0] += 1 + elif change == "plans": + payload["plan_timesteps"][0, 0] += 0.01 + manifest["payload"]["plan_timesteps"] = tensor_digest(payload["plan_timesteps"]) + elif change == "nan": + payload["final_params"][0, 0, 0] = float("nan") + manifest["payload"]["final_params"] = tensor_digest(payload["final_params"]) + _save(path, payload, manifest) + if change == "version": + save_file(payload, str(path), metadata={"format_version": "3"}) + with pytest.raises(ValueError): + MiniMaxH3AdalnCache(arch, path=str(path), model_variant="fl2va") + + +def test_weight_adapter_schedule_and_uncovered_timestep_checks(tmp_path): + arch, weights, _, manifest, path = _fixture(tmp_path, adapter="a" * 64) + cache = MiniMaxH3AdalnCache(arch, path=str(path), model_variant="fl2va") + with pytest.raises(ValueError, match="adapter mismatch"): + cache.bind_adapter(None) + cache.bind_adapter("a" * 64) + with pytest.raises(ValueError, match="Missing"): + cache.finish_loading(torch.device("cpu")) + name = next(iter(weights)) + with pytest.raises(ValueError, match="effective weight mismatch"): + cache.verify_weight(name, weights[name] + 1) + for name, weight in weights.items(): + cache.verify_weight(name, weight) + cache.finish_loading(torch.device("cpu")) + for update in ( + {"flow_shift": 10.0}, + {"audio_flow_shift": 4.0}, + {"mode": "fl2va"}, + {"base_schedule": [0.99, 0.74, 0.5, 0.25, 0]}, + ): + with pytest.raises(ValueError): + cache.check_request(**(manifest["schedule"] | update)) + with pytest.raises(ValueError, match="does not cover"): + cache.lookup(torch.tensor([0.123])) + + +def test_adapter_switch_invalidates_runtime_cache(mocker): + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + cache = mocker.Mock() + pipeline.transformer = SimpleNamespace(adaln_cache=cache) + from vllm.lora.request import LoRARequest + + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sampling = OmniDiffusionSamplingParams(lora_request=LoRARequest("one", 1, "one"), lora_scale=1.0) + pipeline._prepare_adaln_adapter(sampling) + pipeline._prepare_adaln_adapter(sampling) + assert cache.clear.call_count == 1 + sampling.lora_scale = 0.5 + pipeline._prepare_adaln_adapter(sampling) + sampling.lora_request = None + pipeline._prepare_adaln_adapter(sampling) + assert cache.clear.call_count == 3 + + +def test_bad_sidecar_never_drops_runtime_weights(tmp_path): + arch, weights, _, _, path = _fixture(tmp_path) + candidate = MiniMaxH3AdalnCache(arch, path=str(path), model_variant="fl2va") + component = SimpleNamespace(_adaln_sidecar_candidate=candidate) + changed = list(weights.items()) + changed[0] = changed[0][0], changed[0][1] + 1 + result = list(MiniMaxH3Pipeline._verify_adaln_weights(component, iter(changed))) + assert len(result) == len(changed) + assert all(tensor is source for (_, tensor), (_, source) in zip(result, changed)) + assert component._adaln_sidecar_candidate is None + + +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_native_constructor_enables_cache_by_default(monkeypatch, tp_size): + from tests.diffusion.models.minimax_h3.test_minimax_h3_quantization import ( + _FakeAttention, + _FakeLinear, + _small_od_config, + ) + + for name in ("ColumnParallelLinear", "RowParallelLinear", "MergedColumnParallelLinear", "QKVParallelLinear"): + monkeypatch.setattr(h3, name, _FakeLinear) + monkeypatch.setattr(h3, "Attention", _FakeAttention) + monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: tp_size) + config = _small_od_config() + model = h3.MiniMaxH3DiTModel(config, diffusers_weights=False) + assert model.adaln_cache.max_bytes > 0 + assert model.blocks[0].adaln_proj._adaln_cache is model.adaln_cache + assert model.final_layer.adaln_proj._adaln_cache is model.adaln_cache + assert "blocks.0.adaln_proj.linear.weight" in dict(model.named_parameters()) + config.cache_config = {"minimax_h3_adaln_cache": False} + disabled = h3.MiniMaxH3DiTModel(config, diffusers_weights=False) + assert disabled.adaln_cache.max_bytes == 0 + + +def test_optional_sidecar_seeds_runtime_projection(tmp_path, mocker): + arch, weights, payload, _, path = _fixture(tmp_path) + sidecar = _ready(arch, weights, path) + runtime = MiniMaxH3RuntimeAdalnCache() + name = "blocks.0.adaln_proj.linear" + linear = torch.nn.Linear(arch.time_embed_dim, 18 * arch.hidden_size, dtype=torch.bfloat16) + with torch.no_grad(): + linear.weight.copy_(weights[name + ".weight"]) + linear.bias.copy_(weights[name + ".bias"]) + runtime.seed(sidecar, {name: linear}) + x = payload["time_embeddings"][1, :2].clone() + runtime.prepare(x) + compute = mocker.Mock(side_effect=lambda: linear(torch.nn.functional.silu(x).to(torch.bfloat16))) + result = runtime.project(name, linear, x, compute) + assert compute.call_count == 0 + assert torch.equal(result, compute()) + linear.weight.add_(1) + runtime.project(name, linear, x, compute) + assert compute.call_count == 2 + + +def test_failed_seed_does_not_attach_sidecar(tmp_path): + arch, weights, _, _, path = _fixture(tmp_path) + sidecar = _ready(arch, weights, path) + runtime = MiniMaxH3RuntimeAdalnCache() + linear = torch.nn.Linear(arch.time_embed_dim, 18 * arch.hidden_size, dtype=torch.bfloat16) + with torch.inference_mode(): + unversioned = torch.nn.Linear(arch.time_embed_dim, 2 * arch.hidden_size, dtype=torch.bfloat16) + + with pytest.raises(RuntimeError, match="Inference tensors do not track version counter"): + runtime.seed(sidecar, {"blocks.0.adaln_proj.linear": linear, "final_layer.adaln_proj.linear": unversioned}) + + assert runtime.sidecar is None + assert runtime._sidecar_signatures == {} + + +@pytest.mark.parametrize("mode", MODES) +def test_sidecar_prepare_uses_one_fingerprint_without_tensor_scan(tmp_path, mocker, mode): + import vllm_omni.diffusion.cache.exact_projection_cache as cache_module + + arch, weights, payload, manifest, path = _fixture(tmp_path, mode=mode, base_model=True) + sidecar = _ready(arch, weights, path, manifest["model_variant"]) + runtime = MiniMaxH3RuntimeAdalnCache() + name = "final_layer.adaln_proj.linear" + linear = torch.nn.Linear(arch.time_embed_dim, 2 * arch.hidden_size, dtype=torch.bfloat16) + with torch.no_grad(): + linear.weight.copy_(weights[name + ".weight"]) + linear.bias.copy_(weights[name + ".bias"]) + runtime.seed(sidecar, {name: linear}) + count = int(payload["plan_lengths"][-1]) + embedding = payload["time_embeddings"][-1, :count].clone() + expected = linear(torch.nn.functional.silu(embedding).to(torch.bfloat16)) + comparisons = mocker.spy(torch, "equal") + fingerprint = mocker.spy(cache_module, "tensor_digest") + compute = mocker.Mock(side_effect=AssertionError("matching sidecar must seed the projection")) + + runtime.prepare(embedding) + actual = runtime.project(name, linear, embedding, compute) + + comparisons.assert_not_called() + fingerprint.assert_called_once() + assert torch.equal(actual, expected) + compute.assert_not_called() + + +def test_sidecar_rechecks_numerical_environment_at_projection(tmp_path, mocker): + arch, weights, payload, _, path = _fixture(tmp_path) + sidecar = _ready(arch, weights, path) + runtime = MiniMaxH3RuntimeAdalnCache() + name = "final_layer.adaln_proj.linear" + linear = torch.nn.Linear(arch.time_embed_dim, 2 * arch.hidden_size, dtype=torch.bfloat16) + with torch.no_grad(): + linear.weight.copy_(weights[name + ".weight"]) + linear.bias.copy_(weights[name + ".bias"]) + runtime.seed(sidecar, {name: linear}) + count = int(payload["plan_lengths"][-1]) + embedding = payload["time_embeddings"][-1, :count].clone() + runtime.prepare(embedding) + compute = mocker.Mock(side_effect=lambda: linear(torch.nn.functional.silu(embedding).to(torch.bfloat16))) + precision = torch.get_float32_matmul_precision() + try: + torch.set_float32_matmul_precision("high" if precision == "highest" else "highest") + actual = runtime.project(name, linear, embedding, compute) + compute.assert_called_once() + assert torch.equal(actual, compute()) + finally: + torch.set_float32_matmul_precision(precision) + + +def test_builder_rejects_missing_or_duplicate_weights(tmp_path): + arch, weights, _, manifest, _ = _fixture(tmp_path) + for source in (list(weights.items())[:-1], [*weights.items(), next(iter(weights.items()))]): + with pytest.raises(ValueError): + build_cache( + arch, + source, + contract=manifest["schedule"], + model_variant="fl2va", + adapter_sha256=None, + device=torch.device("cpu"), + ) + + +def test_native_checkpoint_builder_cli(tmp_path, monkeypatch): + from tools.minimax_h3.build_adaln_cache import main + + arch, weights, _, _, _ = _fixture(tmp_path) + root = tmp_path / "transformer" + root.mkdir() + (root / "config.json").write_text(canonical_json(asdict(arch))) + save_file(weights, str(root / "model.safetensors")) + output = tmp_path / "cli.safetensors" + argv = [ + "build_adaln_cache", + "--transformer-path", + str(root), + "--output", + str(output), + "--model-variant", + "fl2va", + "--num-inference-steps", + "4", + "--device", + "cpu", + ] + monkeypatch.setattr("sys.argv", argv) + main() + _ready(arch, weights, output) + with pytest.raises(ValueError, match="overwrite"): + main() + + +@pytest.mark.parametrize("base_model", [False, True], ids=["fasth3-four", "base-h3-50"]) +def test_runtime_projection_matches_sidecar_using_actual_h3_forwards(tmp_path, monkeypatch, mocker, base_model): + class Linear(torch.nn.Linear): + def __init__(self, in_features, out_features, *, params_dtype, bias=True, **kwargs): + super().__init__(in_features, out_features, bias=bias, dtype=params_dtype) + + def forward(self, x): + return super().forward(x), None + + monkeypatch.setattr(h3, "ColumnParallelLinear", Linear) + monkeypatch.setattr(h3, "RowParallelLinear", Linear) + arch, weights, payload, manifest, _ = _fixture(tmp_path, base_model=base_model) + embedder = h3.MiniMaxH3TimeEmbedder(arch, prefix="time_embedder") + runtime = MiniMaxH3RuntimeAdalnCache() + proj = h3.MiniMaxH3AdalnProj( + arch, + 18 * arch.hidden_size, + None, + expand_ratio=6, + modality_num=3, + prefix="blocks.0.adaln_proj", + adaln_cache=runtime, + ) + embedder.load_state_dict( + { + name.removeprefix("time_embedder."): value + for name, value in weights.items() + if name.startswith("time_embedder.") + } + ) + proj.load_state_dict( + { + name.removeprefix("blocks.0.adaln_proj."): value + for name, value in weights.items() + if name.startswith("blocks.0.adaln_proj.") + } + ) + calls = mocker.spy(proj.linear, "forward") + with torch.no_grad(): + for _ in range(2): + for i, plan in enumerate(timestep_plans(manifest["schedule"])): + embedding = embedder(plan) + assert torch.equal(embedding, payload["time_embeddings"][i, : len(plan)]) + runtime.prepare(embedding) + actual = torch.stack(proj(embedding), dim=1).reshape(len(plan), -1) + assert torch.equal(actual, payload["block_params"][i, : len(plan), 0]) + plan_count = len(timestep_plans(manifest["schedule"])) + assert calls.call_count == plan_count and runtime.hits == plan_count + + +@pytest.mark.parametrize("enforce_eager", [False, True], ids=["compiled", "eager"]) +@pytest.mark.parametrize( + "key,mode", + [("minimax_h3_adaln_cache_path", "t2va"), ("minimax_h3_ref_adaln_cache_path", "ref2va-mixed")], +) +def test_optional_sidecar_requires_eager_before_reading_payload(tmp_path, mocker, enforce_eager, key, mode): + import vllm_omni.diffusion.models.minimax_h3.adaln_cache as cache_module + + arch, _, _, manifest, path = _fixture(tmp_path, mode=mode) + open_sidecar = mocker.spy(cache_module, "safe_open") + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace(cache_config={key: str(path)}, enforce_eager=enforce_eager) + transformer = SimpleNamespace(arch=arch, adaln_cache=MiniMaxH3RuntimeAdalnCache()) + + pipeline._configure_adaln_sidecar(transformer, key, manifest["model_variant"], None, eligible=True) + + if enforce_eager: + open_sidecar.assert_called_once() + assert isinstance(transformer._adaln_sidecar_candidate, MiniMaxH3AdalnCache) + assert transformer._adaln_sidecar_candidate.path == str(path) + else: + open_sidecar.assert_not_called() + assert not hasattr(transformer, "_adaln_sidecar_candidate") + # The normal load completion must not install or move a skipped payload. + pipeline._finish_adaln_sidecar(transformer) + assert transformer.adaln_cache.sidecar is None + assert transformer.adaln_cache.max_bytes > 0 + + +def test_corrupt_optional_sidecar_falls_back_without_rejecting_model(tmp_path): + path = tmp_path / "truncated.safetensors" + path.write_bytes(b"truncated") + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace(cache_config={"cache": str(path)}, enforce_eager=True) + transformer = SimpleNamespace(arch=h3.MiniMaxH3DiTArchConfig(), adaln_cache=MiniMaxH3RuntimeAdalnCache()) + pipeline._configure_adaln_sidecar(transformer, "cache", "fl2va", None, eligible=True) + assert not hasattr(transformer, "_adaln_sidecar_candidate") diff --git a/tests/diffusion/models/minimax_h3/test_adaln_cache_device.py b/tests/diffusion/models/minimax_h3/test_adaln_cache_device.py new file mode 100644 index 000000000..772104380 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_adaln_cache_device.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real vLLM projection/collective parity without loading an H3 checkpoint.""" + +from datetime import timedelta + +import pytest +import torch +import torch.multiprocessing as mp + +from tests.helpers.mark import hardware_marks +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] + + +def _projection_worker(rank, world_size, rendezvous): + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + + from vllm_omni.diffusion.models.minimax_h3.adaln_cache import MiniMaxH3RuntimeAdalnCache + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3AdalnProj, + MiniMaxH3DiTArchConfig, + ) + + torch.set_num_threads(1) + device = torch.device("cuda", rank) + current_omni_platform.set_device(device) + config = VllmConfig(device_config=DeviceConfig(device="cuda")) + with set_current_vllm_config(config): + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=rank, + distributed_init_method=f"file://{rendezvous}", + backend="nccl", + timeout=timedelta(seconds=45), + ) + initialize_model_parallel(tensor_model_parallel_size=world_size) + try: + arch = MiniMaxH3DiTArchConfig(hidden_size=256, time_embed_dim=128) + cache = MiniMaxH3RuntimeAdalnCache() + generator = torch.Generator(device=device).manual_seed(11) + for name, expansion, modalities in (("blocks.0", 6, 3), ("final_layer", 2, 1)): + cache.clear() + projection = MiniMaxH3AdalnProj( + arch, + arch.hidden_size * expansion * modalities, + None, + expand_ratio=expansion, + modality_num=modalities, + prefix=name + ".adaln_proj", + adaln_cache=cache, + ).to(device) + with torch.no_grad(): + for parameter in projection.parameters(): + parameter.normal_(std=0.05, generator=generator) + with torch.inference_mode(): + for width in (1, 2, 3, 4): + embedding = torch.randn(width, arch.time_embed_dim, device=device, generator=generator) + # No prepared input: run the original vLLM projector. + cache.clear() + expected = projection(embedding) + cache.prepare(embedding) + first = projection(embedding) + hits = cache.hits + second = projection(embedding) + assert cache.hits == hits + 1 + for reference, cold, warm in zip(expected, first, second, strict=True): + assert torch.equal(reference, cold) + assert torch.equal(reference, warm) + + # Only one rank loses its entry. All ranks must execute + # the gather again; a local-only hit would deadlock. + if rank == world_size - 1: + cache.clear() + cache.prepare(embedding) + misses = cache.misses + rebuilt = projection(embedding) + assert cache.misses == misses + 1 + for reference, actual in zip(expected, rebuilt, strict=True): + assert torch.equal(reference, actual) + + if rank == world_size - 1: + projection.linear.weight.add_(0.01) + misses = cache.misses + changed = projection(embedding) + assert cache.misses == misses + 1 + cache.clear() + reference = projection(embedding) + for actual, original in zip(changed, reference, strict=True): + assert torch.equal(actual, original) + finally: + cleanup_dist_env_and_memory() + + +@pytest.mark.parametrize( + "world_size", + [ + pytest.param(1, marks=hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=1)), + pytest.param(2, marks=hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2)), + ], +) +def test_real_projection_cache_parity_and_rank_local_invalidation(tmp_path, world_size): + if not current_omni_platform.is_cuda(): + pytest.skip("Requires CUDA") + context = mp.spawn( + _projection_worker, + args=(world_size, str(tmp_path / "rendezvous")), + nprocs=world_size, + join=False, + ) + try: + if not context.join(timeout=90): + # ProcessContext may reap one successful worker per join. + if not context.join(timeout=30): + pytest.fail("AdaLN projection workers timed out") + finally: + for process in context.processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) diff --git a/tools/minimax_h3/build_adaln_cache.py b/tools/minimax_h3/build_adaln_cache.py new file mode 100644 index 000000000..653332355 --- /dev/null +++ b/tools/minimax_h3/build_adaln_cache.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Build an exact AdaLN sidecar from local native weights and a fixed adapter.""" + +from __future__ import annotations + +import argparse +import json +import os +import tempfile +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from vllm_omni.diffusion.models.minimax_h3.adaln_cache import ( + FORMAT_VERSION, + MODES, + MiniMaxH3AdalnCache, + build_cache, + canonical_json, + file_digest, + input_names, + schedule_contract, +) +from vllm_omni.diffusion.models.minimax_h3.fasth3 import FastH3WeightFusion +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3DiTArchConfig + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + 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(MODES), default="t2va") + parser.add_argument("--num-inference-steps", type=int, default=None) + parser.add_argument("--flow-shift", type=float, default=12.0) + parser.add_argument("--audio-flow-shift", type=float, default=3.0) + parser.add_argument("--base-schedule", type=float, nargs="+") + parser.add_argument("--fasth3-adapter", type=Path) + parser.add_argument("--device", default="cuda", help="Use the serving device type; cpu is for CPU-only fixtures") + args = parser.parse_args() + root = args.transformer_path.expanduser().resolve() + if args.output.exists(): + raise ValueError(f"Refusing to overwrite {args.output}") + config = json.loads((root / "config.json").read_text()) + arch = MiniMaxH3DiTArchConfig.from_mapping(config) + variant = "ref2va" if args.mode.startswith("ref2va") else "fl2va" + if args.model_variant != variant: + raise ValueError("Cache mode does not match --model-variant") + device = torch.device(args.device) + if device.type not in {"cpu", "cuda"}: + raise ValueError("The builder supports CPU fixtures and CUDA serving") + fusion = None + if args.fasth3_adapter is not None: + if args.mode != "t2va" or args.model_variant != "fl2va": + raise ValueError("FastH3 Preview adapters support T2VA only") + fusion = FastH3WeightFusion.from_path( + args.fasth3_adapter, + head_dim=arch.attention_head_dim, + num_blocks=arch.num_layers, + num_refiner_blocks=arch.token_refiner_num_layers, + ) + if fusion is None: + raise ValueError("Not a supported fixed FastH3 adapter") + if args.base_schedule is not None and tuple(args.base_schedule) != fusion.base_schedule: + raise ValueError("Explicit schedule does not match the FastH3 adapter") + args.base_schedule = fusion.base_schedule + num_steps = args.num_inference_steps + if num_steps is None: + num_steps = len(args.base_schedule) - 1 if args.base_schedule is not None else 50 + contract = schedule_contract( + mode=args.mode, + num_steps=num_steps, + base_schedule=args.base_schedule, + flow_shift=args.flow_shift, + audio_flow_shift=args.audio_flow_shift, + ) + adapter_sha = file_digest(fusion.source) if fusion is not None else None + index = root / "model.safetensors.index.json" + if index.is_file(): + weight_map = json.loads(index.read_text())["weight_map"] + else: + single = root / "model.safetensors" + with safe_open(str(single), framework="pt", device="cpu") as source: + weight_map = {name: single.name for name in source.keys()} + required = input_names(arch.num_layers) + if missing := required - set(weight_map): + raise ValueError(f"Native checkpoint is missing AdaLN inputs: {sorted(missing)}") + + def weights(): + ordered = sorted(required, key=lambda name: (not name.startswith("time_embedder."), name)) + for name in ordered: + filename = (root / weight_map[name]).resolve() + if not filename.is_relative_to(root): + # HF snapshot shards may be symlinks to blobs; use the lexical + # index path for containment, while allowing those symlinks. + relative = Path(weight_map[name]) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError("Checkpoint index points outside the transformer directory") + with safe_open(str(filename), framework="pt", device="cpu") as source: + weight = source.get_tensor(name) + if fusion is not None: + weight = fusion.fuse(name, weight.to(device)) + yield name, weight + + payload, manifest = build_cache( + arch, + weights(), + contract=contract, + model_variant=args.model_variant, + adapter_sha256=adapter_sha, + device=device, + ) + if fusion is not None and file_digest(fusion.source) != adapter_sha: + raise ValueError("Adapter changed during cache construction") + args.output.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary = tempfile.mkstemp(prefix=".adaln-", suffix=".safetensors", dir=args.output.parent) + os.close(descriptor) + try: + save_file(payload, temporary, metadata={"format_version": FORMAT_VERSION, "manifest": canonical_json(manifest)}) + MiniMaxH3AdalnCache(arch, path=temporary, model_variant=args.model_variant) + # Atomic no-clobber publication on the same filesystem. + os.link(temporary, args.output) + finally: + Path(temporary).unlink(missing_ok=True) + print(f"Wrote weight-bound AdaLN sidecar: {args.output}") + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/diffusion/cache/__init__.py b/vllm_omni/diffusion/cache/__init__.py index 2d87a089a..6ab3ee905 100644 --- a/vllm_omni/diffusion/cache/__init__.py +++ b/vllm_omni/diffusion/cache/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Cache module for diffusion model inference acceleration. @@ -12,6 +12,7 @@ Cache backends are instantiated directly via their constructors and configured v """ from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.exact_projection_cache import ExactProjectionCache from vllm_omni.diffusion.cache.prompt_embed_cache import ( PromptEmbedCache, install_prompt_embed_cache, @@ -27,6 +28,7 @@ from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend __all__ = [ "CacheBackend", + "ExactProjectionCache", "CacheContext", "TeaCacheBackend", "TeaCacheConfig", diff --git a/vllm_omni/diffusion/cache/exact_projection_cache.py b/vllm_omni/diffusion/cache/exact_projection_cache.py new file mode 100644 index 000000000..896dbc4dd --- /dev/null +++ b/vllm_omni/diffusion/cache/exact_projection_cache.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Model-independent exact caching of small conditioning projections.""" + +from __future__ import annotations + +import hashlib +import json +import weakref +from collections.abc import Callable +from typing import Any + +import torch +from torch import nn +from torch.nn.modules import module as module_hooks + + +class ExactProjectionCache: + """Bounded, exact memoization for deterministic conditioning projections. + + Models call :meth:`prepare` once for an immutable conditioning tensor, then + :meth:`project` for each leaf projection using a stable, model-defined name. + The supplied computation must depend only on that tensor, the projection's + parameters/buffers and fixed preprocessing. Include all request conditioning + in the tensor; timestep alone is insufficient for prompt/guidance-dependent + projections. Outputs returned on cache hits must be treated as read-only. + + The original projection computes misses, preserving its quantization, LoRA + and TP behavior. Parameter/buffer versions guard reuse, and TP ranks vote + before skipping collectives. Call :meth:`clear` for adapter/configuration + changes or model moves. Gradient-enabled and compiled execution bypass reuse. + + Each cache belongs to one model instance executing forwards serially. It + retains a bounded subset across requests, so it is independent of the + request-scoped approximate cache backends. It does not offload weights or + suppress block prefetch. Models may supply validated precomputed results by + overriding :meth:`_lookup_precomputed`; artifact loading stays model-owned. + """ + + def __init__(self, *, max_bytes: int = 256 * 1024**2) -> None: + if type(max_bytes) is not int or max_bytes < 0: + raise ValueError("Projection cache budget must be a nonnegative byte count") + self.max_bytes = max_bytes + self._entries: dict[tuple[str, str, tuple[Any, ...]], tuple[Any, torch.Tensor, int]] = {} + self._bytes = 0 + self._input: weakref.ReferenceType[torch.Tensor] | None = None + self._key: str | None = None + self.hits = self.misses = 0 + + def clear(self) -> None: + self._entries.clear() + self._bytes = 0 + self._input = None + self._key = None + + def prepare(self, embedding: torch.Tensor) -> None: + self._input = None + self._key = None + if not self.max_bytes or torch.is_grad_enabled() or torch.compiler.is_compiling(): + return + self._key = tensor_digest(embedding) + self._input = weakref.ref(embedding) + + @staticmethod + def _numerical_settings(device_type: str) -> tuple[Any, ...]: + # A projection may enter its own autocast context after prepare(). + return ( + torch.get_float32_matmul_precision(), + torch.is_autocast_enabled(device_type), + torch.get_autocast_dtype(device_type), + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + ) + + @staticmethod + def _has_forward_hooks(module: nn.Module) -> bool: + # Match nn.Module's global hooks as well as hooks on wrapped linears. + return bool( + module_hooks._global_forward_pre_hooks + or module_hooks._global_forward_hooks + or any(child._forward_pre_hooks or child._forward_hooks for child in module.modules()) + ) + + @staticmethod + def _signature(module: nn.Module) -> tuple[Any, ...]: + # A replaced or restored storage conservatively causes a miss. The + # containing block's offload hooks still run before this method. + tensors = [*module.named_parameters(), *module.named_buffers()] + # vLLM LoRA stores these tensors outside registered buffers. Its + # suspend/resume mask can change without changing the base weights. + for field in ("lora_a_stacked", "lora_b_stacked"): + tensors.extend((f"{field}.{i}", value) for i, value in enumerate(getattr(module, field, ()))) + return ( + id(module), + id(getattr(module, "quant_method", None)), + getattr(module, "_diffusion_lora_active_slices", None), + *( + (name, id(value), value.data_ptr(), value._version, tuple(value.shape), value.dtype, value.device) + for name, value in tensors + ), + ) + + def project( + self, + name: str, + module: nn.Module, + embedding: torch.Tensor, + compute: Callable[[], torch.Tensor], + ) -> torch.Tensor: + if ( + torch.is_grad_enabled() + or torch.compiler.is_compiling() + or self._key is None + or self._input is None + or self._input() is not embedding + ): + return compute() + key = self._key, name, self._numerical_settings(embedding.device.type) + signature: tuple[Any, ...] | None + try: + signature = self._signature(module) + if self._has_forward_hooks(module): + # A linear hook may mutate weights or transform the output; + # executing it only on misses would change its semantics. + signature = None + except RuntimeError: + # Inference tensors without version counters cannot establish that + # weights stayed unchanged. Still participate in the TP vote. + signature = None + cached = self._entries.get(key) + hit = signature is not None and cached is not None and cached[0] == signature + from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group + + world_size = get_tensor_model_parallel_world_size() + if world_size > 1: + vote = torch.tensor([int(hit)], device=embedding.device, dtype=torch.int32) + hit = int(get_tp_group().all_reduce(vote).item()) == world_size + if hit: + assert cached is not None + self.hits += 1 + return cached[1] + self.misses += 1 + value = ( + self._lookup_precomputed(name, embedding, signature) if world_size == 1 and signature is not None else None + ) + if value is None: + value = compute() + if cached is not None: + self._bytes -= self._entries.pop(key)[2] + size = value.numel() * value.element_size() + if signature is not None and self._bytes + size <= self.max_bytes: + # Retain a reusable subset when a long denoising schedule exceeds + # the budget. LRU would evict every entry before the next request + # reaches it, producing zero hits for a repeated cyclic schedule. + # Own the cached storage even if a provider returns workspace views. + self._entries[key] = (signature, value.detach().clone(), size) + self._bytes += size + return value + + def _lookup_precomputed(self, name: str, embedding: torch.Tensor, signature: Any) -> torch.Tensor | None: + """Return a validated TP1 projection, or None to execute the original. + + Called only on runtime-cache misses with a usable weight signature. + Implementations must also validate inputs, numerical settings, artifact + compatibility and any state not represented in the module signature. + """ + return None + + +def canonical_json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def tensor_digest(tensor: torch.Tensor) -> str: + """Hash dtype, shape and actual bytes, including BF16, without widening.""" + value = tensor.detach().cpu().contiguous() + digest = hashlib.sha256(canonical_json([str(value.dtype), list(value.shape)]).encode()) + digest.update(memoryview(value.reshape(-1).view(torch.uint8).numpy())) + return digest.hexdigest() diff --git a/vllm_omni/diffusion/models/minimax_h3/adaln_cache.py b/vllm_omni/diffusion/models/minimax_h3/adaln_cache.py new file mode 100644 index 000000000..f33016e5d --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/adaln_cache.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Default exact AdaLN memoization and optional weight-bound startup sidecars. + +This is not a cross-step approximation. The main blocks' and final layer's +AdaLN outputs are reused only for identical inputs and unchanged weights. All +weights remain available for new schedules and adapters. Legacy experimental +sidecars without weight identities are rejected. +""" + +from __future__ import annotations + +import hashlib +import json +import math +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from safetensors import safe_open +from torch import nn + +from vllm_omni.diffusion.cache.exact_projection_cache import ( + ExactProjectionCache, + tensor_digest, +) +from vllm_omni.diffusion.cache.exact_projection_cache import ( + canonical_json as canonical_json, +) + +from .time_request import minimax_h3_time_shift_sigmas + +if TYPE_CHECKING: + from .minimax_h3_transformer import MiniMaxH3DiTArchConfig + +FORMAT_VERSION = "4" +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"), +} +PAYLOAD_NAMES = frozenset({"plan_timesteps", "plan_lengths", "time_embeddings", "block_params", "final_params"}) + + +class MiniMaxH3RuntimeAdalnCache(ExactProjectionCache): + """Exact projection caching with optional H3 schedule-bound sidecar results.""" + + def __init__(self, *, max_bytes: int = 256 * 1024**2) -> None: + super().__init__(max_bytes=max_bytes) + self.sidecar: MiniMaxH3AdalnCache | None = None + self._sidecar_signatures: dict[str, Any] = {} + self._sidecar_plan: int | None = None + + def clear(self) -> None: + super().clear() + self.sidecar = None + self._sidecar_signatures.clear() + self._sidecar_plan = None + + def seed(self, sidecar: MiniMaxH3AdalnCache, modules: Mapping[str, nn.Module]) -> None: + # A rejected seed must not retain the sidecar's device payload. + signatures = {name: self._signature(module) for name, module in modules.items()} + self.sidecar = sidecar + self._sidecar_signatures = signatures + + def prepare(self, embedding: torch.Tensor) -> None: + self._sidecar_plan = None + super().prepare(embedding) + if self._key is None or self.sidecar is None: + return + if math_identity(embedding.device) != self.sidecar.manifest["math"]: + self.sidecar = None + return + self._sidecar_plan = self.sidecar._embedding_plans.get(self._key) + + def _lookup_precomputed(self, name: str, embedding: torch.Tensor, signature: Any) -> torch.Tensor | None: + if ( + self.sidecar is None + or self._sidecar_plan is None + or self._sidecar_signatures.get(name) != signature + or torch.is_autocast_enabled(embedding.device.type) + or math_identity(embedding.device) != self.sidecar.manifest["math"] + ): + return None + index, count = self._sidecar_plan, embedding.shape[0] + if name.startswith("blocks."): + return self.sidecar.block_params[index, :count, int(name.split(".")[1])].contiguous() + return self.sidecar.final_params[index, :count].contiguous() + + +def file_digest(path: str | Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as source: + while chunk := source.read(16 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def projection_names(num_layers: int) -> frozenset[str]: + prefixes = [f"blocks.{i}.adaln_proj.linear" for i in range(num_layers)] + prefixes.append("final_layer.adaln_proj.linear") + return frozenset(f"{prefix}.{kind}" for prefix in prefixes for kind in ("weight", "bias")) + + +def input_names(num_layers: int) -> frozenset[str]: + return projection_names(num_layers) | { + f"time_embedder.{module}.{kind}" for module in ("proj_in", "proj_out") for kind in ("weight", "bias") + } + + +def architecture(arch: MiniMaxH3DiTArchConfig) -> dict[str, int]: + return { + key: getattr(arch, key) + for key in ("num_layers", "hidden_size", "timestep_input_dim", "time_embed_hidden_size", "time_embed_dim") + } + + +def math_identity(device: torch.device) -> dict[str, Any]: + """Conservatively bind the builder's numerical environment (TP1 only).""" + result = { + "torch": str(torch.__version__), + "device_type": device.type, + "float32_matmul_precision": torch.get_float32_matmul_precision(), + } + if device.type == "cuda": + from vllm_omni.platforms import current_omni_platform + + capability = current_omni_platform.get_device_capability(device.index or 0) + result.update( + cuda=torch.version.cuda, + capability=list(capability) if capability is not None else None, + allow_tf32=torch.backends.cuda.matmul.allow_tf32, + allow_bf16_reduced_precision_reduction=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + ) + return result + + +def schedule_contract( + *, mode: str, num_steps: int, base_schedule: Sequence[float] | None, flow_shift: float, audio_flow_shift: float +) -> dict[str, Any]: + if mode not in MODES: + raise ValueError(f"Unsupported AdaLN cache mode: {mode}") + if type(num_steps) is not int or num_steps < (2 if base_schedule is None else 1): + raise ValueError("AdaLN cache requires a positive complete denoising schedule") + for value in (flow_shift, audio_flow_shift): + if not math.isfinite(value) or value <= 0: + raise ValueError("AdaLN cache shifts must be finite and positive") + if base_schedule is not None: + from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule + + schedule = DMD2SigmaSchedule.from_positions(base_schedule) + if schedule.num_inference_steps != num_steps: + raise ValueError("AdaLN cache step count does not match the explicit schedule") + base_schedule = schedule.base_schedule + return dict( + mode=mode, + num_steps=num_steps, + base_schedule=list(base_schedule) if base_schedule is not None else None, + flow_shift=float(flow_shift), + audio_flow_shift=float(audio_flow_shift), + ) + + +def timestep_plans(contract: Mapping[str, Any]) -> list[torch.Tensor]: + """Use the serving scheduler and its conditioning timesteps, without interpolation.""" + from .denoise_loop import MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, MINIMAX_H3_IMGVID_COND_TIMESTEP + + video = minimax_h3_time_shift_sigmas( + num_steps=contract["num_steps"], shift_scale=contract["flow_shift"], base_schedule=contract["base_schedule"] + ) + audio = minimax_h3_time_shift_sigmas( + num_steps=contract["num_steps"], + shift_scale=contract["audio_flow_shift"], + base_schedule=contract["base_schedule"], + ) + plans: dict[tuple[float, ...], torch.Tensor] = {} + for video_sigma, audio_sigma in zip(video[:-1], audio[:-1], strict=True): + values = { + "video": 1.0 - video_sigma, + "audio": 1.0 - audio_sigma, + "image": max(1.0 - video_sigma, MINIMAX_H3_IMGVID_COND_TIMESTEP), + "audio_ref": max(1.0 - audio_sigma, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + } + plan = torch.tensor([values[key] for key in MODES[contract["mode"]]], dtype=torch.float32).unique(sorted=True) + plans[tuple(plan.tolist())] = plan + return list(plans.values()) + + +class MiniMaxH3AdalnCache(nn.Module): + """Validate an optional sidecar, then verify its weights at load. + + Hashes are checked against the *post-fusion* stream, before TP conversion. + They cover all inputs of the cached computation, not unrelated QKV/MLP + weights. A fixed adapter has an additional complete-file identity. + """ + + block_params: torch.Tensor + final_params: torch.Tensor + plan_timesteps: torch.Tensor + plan_lengths: torch.Tensor + time_embeddings: torch.Tensor + + def __init__(self, arch: MiniMaxH3DiTArchConfig, *, path: str, model_variant: str) -> None: + super().__init__() + self.path = str(Path(path).expanduser().resolve()) + self.arch = arch + self._verified: set[str] = set() + self._adapter_bound = False + self._ready = False + with safe_open(self.path, framework="pt", device="cpu") as source: + metadata = source.metadata() or {} + if metadata.get("format_version") != FORMAT_VERSION: + raise ValueError("Rebuild the AdaLN cache: a weight-bound format-v4 sidecar is required") + try: + manifest = json.loads(metadata["manifest"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Invalid AdaLN cache manifest") from exc + required = {"architecture", "model_variant", "schedule", "adapter_sha256", "weights", "payload", "math"} + if not isinstance(manifest, dict) or set(manifest) != required: + raise ValueError("Invalid AdaLN cache manifest fields") + if any( + not isinstance(manifest[key], dict) + for key in ("architecture", "schedule", "weights", "payload", "math") + ): + raise ValueError("Invalid AdaLN cache manifest objects") + if manifest["architecture"] != architecture(arch) or manifest["model_variant"] != model_variant: + raise ValueError("AdaLN cache model variant or architecture mismatch") + contract = schedule_contract(**manifest["schedule"]) + variant = "ref2va" if contract["mode"].startswith("ref2va") else "fl2va" + if variant != model_variant: + raise ValueError("AdaLN cache task does not match model variant") + if set(manifest["weights"]) != input_names(arch.num_layers): + raise ValueError("AdaLN cache must cover exactly the time embedding and AdaLN parameters") + if set(source.keys()) != PAYLOAD_NAMES or set(manifest["payload"]) != PAYLOAD_NAMES: + raise ValueError("Invalid AdaLN cache payload keys") + for digest in [*manifest["weights"].values(), *manifest["payload"].values()]: + if not isinstance(digest, str) or len(digest) != 64 or any(c not in "0123456789abcdef" for c in digest): + raise ValueError("Invalid AdaLN cache SHA256 digest") + adapter = manifest["adapter_sha256"] + if adapter is not None and ( + not isinstance(adapter, str) or len(adapter) != 64 or any(c not in "0123456789abcdef" for c in adapter) + ): + raise ValueError("Invalid AdaLN cache adapter identity") + payload = {name: source.get_tensor(name) for name in PAYLOAD_NAMES} + plans = timestep_plans(contract) + width = max(len(plan) for plan in plans) + expected_lengths = torch.tensor([len(plan) for plan in plans], dtype=torch.int64) + expected_plans = torch.zeros((len(plans), width), dtype=torch.float32) + for index, plan in enumerate(plans): + expected_plans[index, : len(plan)] = plan + shapes = { + "block_params": ((len(plans), width, arch.num_layers, 18 * arch.hidden_size), torch.bfloat16), + "final_params": ((len(plans), width, 2 * arch.hidden_size), torch.bfloat16), + "time_embeddings": ((len(plans), width, arch.time_embed_dim), torch.float32), + "plan_timesteps": (tuple(expected_plans.shape), torch.float32), + "plan_lengths": (tuple(expected_lengths.shape), torch.int64), + } + for name, value in payload.items(): + shape, dtype = shapes[name] + if tuple(value.shape) != shape or value.dtype != dtype or not bool(torch.isfinite(value).all()): + raise ValueError(f"Invalid AdaLN cache tensor: {name}") + if tensor_digest(value) != manifest["payload"][name]: + raise ValueError(f"AdaLN cache payload checksum mismatch: {name}") + if not torch.equal(payload["plan_timesteps"], expected_plans) or not torch.equal( + payload["plan_lengths"], expected_lengths + ): + raise ValueError("AdaLN cache timestep plans do not match the declared schedule") + self.manifest = manifest + self.projection_names = projection_names(arch.num_layers) + # Index the validated CPU payload once. Runtime prepare() already hashes + # the actual embedding, so lookup needs no per-plan device comparisons. + self._embedding_plans = { + tensor_digest(payload["time_embeddings"][index, :count]): index + for index, count in enumerate(expected_lengths.tolist()) + } + for name, value in payload.items(): + self.register_buffer(name, value, persistent=False) + + def bind_adapter(self, adapter_sha256: str | None) -> None: + if adapter_sha256 != self.manifest["adapter_sha256"]: + raise ValueError("AdaLN cache adapter mismatch; rebuild for the active fixed adapter") + self._adapter_bound = True + + def verify_weight(self, name: str, weight: torch.Tensor) -> None: + if name not in self.manifest["weights"]: + return + if self._ready or name in self._verified: + raise ValueError(f"AdaLN input was loaded more than once: {name}") + if tensor_digest(weight) != self.manifest["weights"][name]: + raise ValueError(f"AdaLN cache effective weight mismatch: {name}") + self._verified.add(name) + + def finish_loading(self, device: torch.device) -> None: + if not self._adapter_bound: + raise ValueError("AdaLN cache requires the active adapter identity before loading") + missing = input_names(self.arch.num_layers) - self._verified + if missing: + raise ValueError(f"Missing AdaLN cache source weights: {sorted(missing)}") + if math_identity(device) != self.manifest["math"]: + raise ValueError("AdaLN cache numerical environment mismatch; rebuild on the serving configuration") + self.to(device=device) + self._ready = True + + def check_request( + self, + *, + mode: str, + num_steps: int, + base_schedule: Sequence[float] | None, + flow_shift: float, + audio_flow_shift: float, + ) -> None: + contract = schedule_contract( + mode=mode, + num_steps=num_steps, + base_schedule=base_schedule, + flow_shift=flow_shift, + audio_flow_shift=audio_flow_shift, + ) + if contract != self.manifest["schedule"]: + raise ValueError("AdaLN cache request schedule/task mismatch; rebuild or disable the cache") + + def lookup(self, timesteps: torch.Tensor) -> tuple[tuple[tuple[torch.Tensor, ...], ...], tuple[torch.Tensor, ...]]: + if not self._ready: + raise RuntimeError("AdaLN cache has not verified its source weights") + if timesteps.dtype != torch.float32: + raise ValueError("AdaLN cache requires exact FP32 timesteps") + values = timesteps.reshape(-1).to(self.plan_timesteps.device) + count = values.numel() + if not 0 < count <= self.plan_timesteps.shape[1]: + raise ValueError("AdaLN cache does not cover this timestep plan") + matches = self.plan_lengths.eq(count) & self.plan_timesteps[:, :count].eq(values).all(dim=1) + if not bool(matches.any()): + raise ValueError("AdaLN cache does not cover this timestep plan") + index = matches.to(torch.int64).argmax() + blocks = self.block_params.permute(2, 0, 1, 3)[:, index, :count] + blocks = blocks.reshape(self.arch.num_layers, -1, 6, self.arch.hidden_size) + final = self.final_params[index, :count].reshape(-1, 2, self.arch.hidden_size) + return tuple(tuple(layer.unbind(1)) for layer in blocks), tuple(final.unbind(1)) + + +@torch.inference_mode() +def build_cache( + arch: MiniMaxH3DiTArchConfig, + weights: Iterable[tuple[str, torch.Tensor]], + *, + contract: dict[str, Any], + model_variant: str, + adapter_sha256: str | None, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Stream one projection at a time; do not instantiate the full DiT. + + The four time-embedder tensors must precede the AdaLN projections. Callers + fuse any fixed adapter before supplying this stream, just like serving. + """ + from .minimax_h3_transformer import MiniMaxH3TimeEmbedder + + contract = schedule_contract(**contract) + plans = timestep_plans(contract) + width = max(len(plan) for plan in plans) + payload = { + "plan_timesteps": torch.zeros((len(plans), width), dtype=torch.float32), + "plan_lengths": torch.tensor([len(plan) for plan in plans], dtype=torch.int64), + "block_params": torch.zeros((len(plans), width, arch.num_layers, 18 * arch.hidden_size), dtype=torch.bfloat16), + "final_params": torch.zeros((len(plans), width, 2 * arch.hidden_size), dtype=torch.bfloat16), + "time_embeddings": torch.zeros((len(plans), width, arch.time_embed_dim), dtype=torch.float32), + } + for i, plan in enumerate(plans): + payload["plan_timesteps"][i, : len(plan)] = plan + + # Use the runtime embedding implementation, including FP32 boundaries. A + # lightweight module supplies ordinary Linear wrappers with TP1 semantics. + class Linear(nn.Linear): + def forward(self, x): + return super().forward(x), None + + embedder = MiniMaxH3TimeEmbedder.__new__(MiniMaxH3TimeEmbedder) + nn.Module.__init__(embedder) + embedder.frequency_embedding_size = arch.timestep_input_dim + embedder.proj_in = Linear(arch.timestep_input_dim, arch.time_embed_hidden_size, device=device, dtype=torch.float32) + embedder.proj_out = Linear(arch.time_embed_hidden_size, arch.time_embed_dim, device=device, dtype=torch.float32) + fingerprints: dict[str, str] = {} + pending: dict[str, torch.Tensor] = {} + inputs: list[torch.Tensor] | None = None + for name, weight in weights: + if name not in input_names(arch.num_layers) or name in fingerprints: + raise ValueError(f"Unexpected or duplicate AdaLN builder input: {name}") + dtype = torch.float32 if name.startswith("time_embedder.") else torch.bfloat16 + if weight.dtype != dtype or not bool(torch.isfinite(weight).all()): + raise ValueError(f"Invalid AdaLN builder input dtype/values: {name}") + fingerprints[name] = tensor_digest(weight) + if name.startswith("time_embedder."): + target = dict(embedder.named_parameters())[name.removeprefix("time_embedder.")] + if weight.shape != target.shape: + raise ValueError(f"Invalid AdaLN time embedding shape: {name}") + target.copy_(weight) + continue + if inputs is None: + if len([key for key in fingerprints if key.startswith("time_embedder.")]) != 4: + raise ValueError("Supply all four time embedding weights before the AdaLN projections") + embeddings = [embedder(plan.to(device)) for plan in plans] + for index, embedding in enumerate(embeddings): + payload["time_embeddings"][index, : len(embedding)].copy_(embedding.cpu()) + inputs = [nn.functional.silu(embedding).to(torch.bfloat16) for embedding in embeddings] + pending[name] = weight.to(device) + prefix = name.rsplit(".", 1)[0] + if f"{prefix}.weight" not in pending or f"{prefix}.bias" not in pending: + continue + w, bias = pending.pop(f"{prefix}.weight"), pending.pop(f"{prefix}.bias") + is_final = prefix.startswith("final_layer.") + out_dim = (2 if is_final else 18) * arch.hidden_size + if w.shape != (out_dim, arch.time_embed_dim) or bias.shape != (out_dim,): + raise ValueError(f"Invalid AdaLN projection shape: {prefix}") + for i, x in enumerate(inputs): + result = nn.functional.linear(x, w, bias).cpu() + if is_final: + payload["final_params"][i, : len(x)].copy_(result) + else: + payload["block_params"][i, : len(x), int(prefix.split(".")[1])].copy_(result) + del w, bias + if set(fingerprints) != input_names(arch.num_layers) or pending: + raise ValueError("Incomplete AdaLN builder input weights") + if any(not bool(torch.isfinite(value).all()) for value in payload.values()): + raise ValueError("Non-finite AdaLN builder output") + manifest = dict( + architecture=architecture(arch), + model_variant=model_variant, + schedule=contract, + adapter_sha256=adapter_sha256, + weights=fingerprints, + payload={name: tensor_digest(value) for name, value in payload.items()}, + math=math_identity(device), + ) + return payload, manifest diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 5662a9641..115264a48 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -55,6 +55,7 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.platforms import current_omni_platform +from .adaln_cache import MiniMaxH3RuntimeAdalnCache from .fasth3 import _resolve_native_target if TYPE_CHECKING: @@ -759,12 +760,15 @@ class MiniMaxH3AdalnProj(nn.Module): expand_ratio: int, modality_num: int, prefix: str, + adaln_cache: MiniMaxH3RuntimeAdalnCache | None = None, ) -> None: super().__init__() if out_features != expand_ratio * arch.hidden_size * modality_num: raise ValueError( f"adaln out_features mismatch: {out_features} != {expand_ratio}*{arch.hidden_size}*{modality_num}" ) + self._adaln_cache = adaln_cache + self._cache_name = prefix + ".linear" self.expand_ratio = expand_ratio self.modality_num = modality_num self.hidden_size = arch.hidden_size @@ -780,8 +784,16 @@ class MiniMaxH3AdalnProj(nn.Module): 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)) + + def project() -> torch.Tensor: + x = nn.functional.silu(t_emb) + return self.linear(x.to(_BF16_DTYPE))[0] + + x = ( + project() + if self._adaln_cache is None + else self._adaln_cache.project(self._cache_name, self.linear, t_emb, project) + ) 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)) @@ -877,6 +889,7 @@ class MiniMaxH3DiTBlock(nn.Module): quant_config: QuantizationConfig | None, *, prefix: str, + adaln_cache: MiniMaxH3RuntimeAdalnCache | None = None, ) -> None: super().__init__() self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) @@ -900,6 +913,7 @@ class MiniMaxH3DiTBlock(nn.Module): expand_ratio=6, modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, prefix=f"{prefix}.adaln_proj", + adaln_cache=adaln_cache, ) def forward( @@ -975,6 +989,7 @@ class MiniMaxH3FinalLayer(nn.Module): quant_config: QuantizationConfig | None, *, prefix: str, + adaln_cache: MiniMaxH3RuntimeAdalnCache | None = None, ) -> None: super().__init__() video_patch_dim = arch.latents_dim * arch.patch_size[0] * arch.patch_size[1] * arch.patch_size[2] @@ -986,6 +1001,7 @@ class MiniMaxH3FinalLayer(nn.Module): expand_ratio=2, modality_num=1, prefix=f"{prefix}.adaln_proj", + adaln_cache=adaln_cache, ) self.video_out = ColumnParallelLinear( arch.hidden_size, @@ -1164,6 +1180,15 @@ class MiniMaxH3DiTModel(nn.Module): self._rope_theta = float(config_mapping.get("rope_theta", 10000.0)) arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping) self.arch = arch + cache_config = getattr(od_config, "cache_config", {}) + enabled = ( + cache_config.get("minimax_h3_adaln_cache", True) + if isinstance(cache_config, Mapping) + else getattr(cache_config, "minimax_h3_adaln_cache", True) + ) + if type(enabled) is not bool: + raise ValueError("minimax_h3_adaln_cache must be a boolean") + self.adaln_cache = MiniMaxH3RuntimeAdalnCache(max_bytes=256 * 1024**2 if enabled else 0) self.od_config = od_config self.parallel_config = od_config.parallel_config self.hidden_size = arch.hidden_size @@ -1227,6 +1252,7 @@ class MiniMaxH3DiTModel(nn.Module): arch, quant_config, prefix=f"blocks.{i}", + adaln_cache=self.adaln_cache, ) for i in range(arch.num_layers) ] @@ -1239,6 +1265,7 @@ class MiniMaxH3DiTModel(nn.Module): arch, quant_config, prefix="final_layer", + adaln_cache=self.adaln_cache, ) self._mark_missing_params_required() @@ -1311,6 +1338,12 @@ class MiniMaxH3DiTModel(nn.Module): if rope_table.dtype != _BF16_DTYPE: raise ValueError(f"rope_table must be {_BF16_DTYPE}, got {rope_table.dtype}.") + def _apply(self, fn, recurse=True): + # Derived outputs must not keep the old device alive after offload/move. + if getattr(self, "adaln_cache", None) is not None: + self.adaln_cache.clear() + return super()._apply(fn, recurse=recurse) + def post_load_weights(self) -> None: for name, param in self.named_parameters(): if name in MINIMAX_H3_FP32_PARAM_NAMES and param.dtype != _FP32_DTYPE: @@ -1328,6 +1361,8 @@ class MiniMaxH3DiTModel(nn.Module): weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: """Load native or Diffusers H3 weights with the existing TP-aware loaders.""" + if getattr(self, "adaln_cache", None) is not None: + self.adaln_cache.clear() params = dict(self.named_parameters()) params.update(dict(self.named_buffers())) loaded: set[str] = set() @@ -1606,6 +1641,9 @@ class MiniMaxH3DiTModel(nn.Module): local_span=local_span, ) + if getattr(self, "adaln_cache", None) is not None: + self.adaln_cache.prepare(t_emb) + combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device) inverse_indices = inverse_indices.to(device) diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 21c6e2681..be9d359d0 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -63,6 +63,7 @@ from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.errors import OmniClientError, client_error_from_metadata +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -642,6 +643,7 @@ class MiniMaxH3Pipeline( ) -> tuple[LoRAModel, PEFTHelper] | None: # A cache eviction may be followed by a different adapter reusing the # same client-supplied ID. Every real load replaces the classification. + self._clear_adaln_caches() self._turbo_lora_specs.pop(lora_request.lora_int_id, None) self._native_lora_adapter_ids.discard(lora_request.lora_int_id) self._lora_sigma_schedules.pop(lora_request.lora_int_id, None) @@ -975,6 +977,22 @@ class MiniMaxH3Pipeline( audio_shift=self.default_audio_shift, ) + self._configure_adaln_sidecar( + self.transformer, + "minimax_h3_adaln_cache_path", + expected_partition, + self._fasth3.source if self._fasth3 is not None else None, + eligible=transformer_quant_config is None and not modular, + ) + if ref2va_model_path is not None: + self._configure_adaln_sidecar( + self.transformers_ref, + "minimax_h3_ref_adaln_cache_path", + "ref2va", + None, + eligible=transformer_quant_config is None and not modular, + ) + if self.load_text_encoder: self.tokenizer = Qwen2TokenizerFast.from_pretrained( str(model_path), @@ -1103,11 +1121,14 @@ class MiniMaxH3Pipeline( # Fuse before the model shards anything, which is also the only # point where the checkpoint's fused QKV/MLP layouts are intact. stream = self._fasth3.apply(stream) + if getattr(component, "_adaln_sidecar_candidate", None) is not None: + stream = self._verify_adaln_weights(component, stream) loaded = component.load_weights(stream) if prefix == "transformer.": transformer_loaded = set(loaded) if prefix != "text_encoder.": component.post_load_weights() + self._finish_adaln_sidecar(component) loaded_with_prefix.update(prefix + name for name in loaded) # Both VAEs load eagerly in ``__init__`` rather than through # ``weights_sources``. The text encoder uses the shared component @@ -1135,6 +1156,89 @@ class MiniMaxH3Pipeline( """True when --lora-path was consumed as a load-time weight fusion.""" return self._fasth3 is not None + def _configure_adaln_sidecar( + self, + transformer: MiniMaxH3DiTModel, + key: str, + variant: str, + adapter_path: str | Path | None, + *, + eligible: bool, + ) -> None: + from safetensors import SafetensorError + from vllm.distributed import get_tensor_model_parallel_world_size + + from .adaln_cache import MiniMaxH3AdalnCache, file_digest + + config = self.od_config.cache_config + path = config.get(key) if isinstance(config, Mapping) else getattr(config, key, None) + if path is None or not transformer.adaln_cache.max_bytes: + return + try: + # Compiled blocks bypass projection reuse. Reject before reading the + # sidecar so load completion cannot move an unused payload to GPU. + if not self.od_config.enforce_eager: + raise ValueError( + "offline sidecars require --enforce-eager; compiled H3 blocks bypass cached projections" + ) + if not eligible or get_tensor_model_parallel_world_size() != 1: + raise ValueError("offline sidecar uses native BF16 TP1 math; use the default runtime cache here") + sidecar = MiniMaxH3AdalnCache(transformer.arch, path=path, model_variant=variant) + sidecar.bind_adapter(file_digest(adapter_path) if adapter_path is not None else None) + # Keep optional CPU-derived data outside the registered model tree. + object.__setattr__(transformer, "_adaln_sidecar_candidate", sidecar) + except (OSError, RuntimeError, TypeError, ValueError, SafetensorError) as exc: + logger.warning("Rejecting optional H3 AdaLN sidecar; runtime caching remains enabled: %s", exc) + + @staticmethod + def _verify_adaln_weights( + component: MiniMaxH3DiTModel, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, tensor in weights: + candidate = getattr(component, "_adaln_sidecar_candidate", None) + if candidate is not None: + try: + candidate.verify_weight(name, tensor) + except ValueError as exc: + object.__setattr__(component, "_adaln_sidecar_candidate", None) + logger.warning("Rejecting optional H3 AdaLN sidecar: %s", exc) + yield name, tensor + + @staticmethod + def _finish_adaln_sidecar(component: nn.Module) -> None: + from vllm.model_executor.layers.utils import default_unquantized_gemm + + candidate = getattr(component, "_adaln_sidecar_candidate", None) + if candidate is None: + return + try: + modules = { + f"blocks.{i}.adaln_proj.linear": block.adaln_proj.linear for i, block in enumerate(component.blocks) + } + modules["final_layer.adaln_proj.linear"] = component.final_layer.adaln_proj.linear + for module in (*modules.values(), component.time_embedder.proj_in, component.time_embedder.proj_out): + if getattr(module.quant_method, "_gemm_impl", None) is not default_unquantized_gemm: + raise ValueError("offline sidecar requires the builder's torch linear backend") + candidate.finish_loading(component.video_patch_proj.weight.device) + component.adaln_cache.seed(candidate, modules) + except (RuntimeError, ValueError) as exc: + logger.warning("Rejecting optional H3 AdaLN sidecar; using runtime projections: %s", exc) + finally: + object.__setattr__(component, "_adaln_sidecar_candidate", None) + + def _clear_adaln_caches(self) -> None: + for name in ("transformer", "transformers_ref"): + cache = getattr(getattr(self, name, None), "adaln_cache", None) + if cache is not None: + cache.clear() + + def _prepare_adaln_adapter(self, sampling: OmniDiffusionSamplingParams) -> None: + request = getattr(sampling, "lora_request", None) + identity = None if request is None else (request.lora_int_id, request.lora_path, float(sampling.lora_scale)) + if identity != getattr(self, "_adaln_adapter_identity", None): + self._clear_adaln_caches() + self._adaln_adapter_identity = identity + def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel: if task == "ref2va" and hasattr(self, "transformers_ref"): return self.transformers_ref @@ -2431,7 +2535,27 @@ class MiniMaxH3Pipeline( except ValueError as exc: raise OmniClientError(str(exc)) from exc + self._prepare_adaln_adapter(sampling) base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) + transformer = getattr( + self, "transformers_ref" if task == "ref2va" and hasattr(self, "transformers_ref") else "transformer", None + ) + cache = getattr(transformer, "adaln_cache", None) + if cache is not None and cache.sidecar is not None: + mode = task + if task == "ref2va": + mode += "-mixed" if visual_shapes and audio_lengths else "-audio" if audio_lengths else "-image" + try: + cache.sidecar.check_request( + mode=mode, + num_steps=num_steps, + base_schedule=base_schedule, + flow_shift=float(extra.get("flow_shift", self.default_video_shift)), + audio_flow_shift=float(extra.get("audio_flow_shift", self.default_audio_shift)), + ) + except ValueError as exc: + logger.warning("Rejecting optional AdaLN sidecar for this schedule; using runtime cache: %s", exc) + cache.clear() quality_plan = self._quality_policy.resolve( quality=sampling.quality, num_inference_steps=num_steps,