anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Perf] Add reusable exact projection caching with MiniMax-H3 integration (#7987)

Browse Source 
Signed-off-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Signed-off-by: Armaan Amatya <armaanamatya2014@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@users.noreply.github.com>
Co-authored-by: Armaan Amatya <38102001+armaanamatya@users.noreply.github.com>
This commit is contained in:
SYLARlishunyang12Armaan Amatya
 authored and GitHub committed 2026-09-24 23:30:42 +08:00
1 parent 63bbd19a1d
commit 535bd35b42
11 files changed
+1886 -3
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/minimax_h3/test_adaln_cache.py
+439
View File
@@ -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")
tests/diffusion/models/minimax_h3/test_adaln_cache_device.py
+127
View File
@@ -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)