From 86b85c078bc041e04aee4c4d9167fb10fb1994c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Mon, 31 Aug 2026 12:04:58 +0800 Subject: [PATCH] [Feature][MiniMax-H3] Fuse the FastH3 four-step adapter at load time (#6714) Signed-off-by: princepride Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- recipes/MiniMaxAI/MiniMax-H3.md | 66 ++ .../diffusion/distributed/test_a2a_permute.py | 29 +- .../minimax_h3/test_minimax_h3_contract.py | 7 +- .../minimax_h3/test_minimax_h3_fasth3.py | 498 +++++++++++++ .../minimax_h3/test_minimax_h3_native_lora.py | 6 +- .../sched/test_dmd2_sigma_schedule.py | 14 +- .../openai_api/test_video_api_utils.py | 33 +- .../diffusion/distributed/a2a_permute.py | 2 +- .../diffusion/models/minimax_h3/fasth3.py | 679 ++++++++++++++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 95 ++- vllm_omni/diffusion/sched/sigma_schedule.py | 9 +- .../diffusion/worker/diffusion_worker.py | 19 +- .../entrypoints/openai/video_api_utils.py | 17 +- 13 files changed, 1420 insertions(+), 54 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/fasth3.py diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 9b306bc2f..71de893fd 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -918,6 +918,72 @@ adapter and twice without it, then compare the four output digests. The adapter is bound and deterministic when each pair matches internally and the two pairs differ from each other. +### FastH3 adapter + +[FastH3](https://haoailab.com/blogs/fasth3-preview/) is FastVideo's four-step +DMD2 student of H3-Base. It reuses H3's text encoder, VAEs, tokenizers, and +schedulers unchanged, replacing 49 denoiser evaluations with four. It is fused +into the checkpoint at load time rather than switched per request, because it +carries full-rank deltas that no LoRA layer can express. + +The bundle publishes four variants, so download one and point `--lora-path` at +it; the repository root is refused rather than guessed at: + +```bash +export FASTH3_DIR=/path/to/fasth3 +hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \ + dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}" +export FASTH3_LORA="${FASTH3_DIR}/dense-datafree/adapter_model.safetensors" +``` + +Add `--task-type fl2va --lora-path "${FASTH3_LORA}"` to a non-offloaded server +command. T2VA is served by the FL2VA partition, so `--task-type fl2va` is +correct even though FastH3 preview v1 distills T2VA only. Because the adapter is +fused, `--lora-backend` does not apply and a request carrying a `lora=` field is +rejected rather than served without the adapter it asked for. + +```bash +-F 'num_inference_steps=4' \ +-F 'extra_params={"task":"t2va","duration":4.4}' +``` + +Requests must ask for `num_inference_steps=4` and `task=t2va`: the release's five +sigma points bound four denoiser evaluations, and that count is what the step +scheduler admits a request on. The server denoises on the release's own ladder, +keeping H3's per-modality shifts at the checkpoint values, so a request that +overrides `flow_shift` or `audio_flow_shift` is rejected - it would sample the +student at noise levels it was never distilled at. + +Only a release that identifies itself as FastH3 is fused; any other +`fastvideo-lora-v2` adapter stays on the dynamic LoRA route. A claimed artifact +is then held to its own metadata: one that misdeclares its tensor counts or +leaves a transformer block unedited is refused at startup instead of serving +mostly base H3 weights on a four-step schedule. Offload is refused for the same +reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and +`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast. + +> [!NOTE] +> Only the **Dense / Data-Free** variant is supported today. The three VSA +> variants carry `.set_weight` compression gates for modules that do not exist +> in vLLM-Omni's dense H3 attention, and the release states dense attention is +> not a supported substitute for them, so they are refused at startup rather +> than silently run wrong. Video Sparse Attention for H3's packed +> `[text | cond | audio | video]` sequence is follow-up work. + +Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at +1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded: + +| Adapter | Steps | End-to-end | Diffusion engine | +| --- | ---: | ---: | ---: | +| none (base H3) | 50 | 25.8 / 26.4 s | 16.22 / 16.28 s | +| FastH3 Dense | 4 (5 sigma points) | 11.7 / 11.8 s | 2.37 / 2.36 s | + +The denoising speedup is 6.9x. End-to-end is 2.2x because text encoding, VAE +decoding and muxing are a fixed cost that dominates a clip this short; longer +generations move the end-to-end figure toward the denoising one. Fusing the +adapter does not measurably change startup: weight loading took 77.3 s with it +against 85.8 s without. + ## Key parameters | Parameter | Recommended value | Notes | diff --git a/tests/diffusion/distributed/test_a2a_permute.py b/tests/diffusion/distributed/test_a2a_permute.py index a21e29217..a74f98378 100644 --- a/tests/diffusion/distributed/test_a2a_permute.py +++ b/tests/diffusion/distributed/test_a2a_permute.py @@ -43,21 +43,16 @@ def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatc @dataclass class _FakeAllocation: + """Byte workspace stand-in that slices and reinterprets like ``symm_mem.empty()``.""" + size: int device: torch.device + def __post_init__(self) -> None: + self.buffer = torch.empty(self.size, dtype=torch.uint8) -class _FakeHandle: - rank = 0 - - def __init__(self, allocation: _FakeAllocation) -> None: - self.allocation = allocation - self.views: list[tuple[tuple[int, ...], torch.dtype]] = [] - - def get_buffer(self, rank, shape, dtype): - assert rank == self.rank - self.views.append((shape, dtype)) - return self.allocation, shape, dtype + def __getitem__(self, item: slice) -> torch.Tensor: + return self.buffer[item] def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None: @@ -93,16 +88,20 @@ def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None: return allocation monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty) - monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda allocation, _group: _FakeHandle(allocation)) + # The handle is only retained to keep the rendezvous registration alive. + monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda _allocation, _group: object()) device = torch.device("cuda:0") first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group") smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group") larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group") - assert first[0].size == 12 - assert smaller[0] is first[0] - assert larger[0].size == 24 + assert (first.shape, first.dtype) == ((2, 3), torch.float16) + assert (smaller.shape, smaller.dtype) == ((1, 4), torch.float16) + assert (larger.shape, larger.dtype) == ((4, 3), torch.float16) + # A sub-capacity request reuses the workspace; only growth allocates again. + assert smaller.untyped_storage().data_ptr() == first.untyped_storage().data_ptr() + assert larger.untyped_storage().data_ptr() != first.untyped_storage().data_ptr() assert [allocation.size for allocation in allocations] == [12, 24] assert len(a2a_permute._SYMM_WORKSPACES) == 1 assert len(all_reduce_calls) == 1 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index e4767cabb..382ecf9aa 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -496,13 +496,16 @@ def test_joint_postprocess_is_multiprocessing_picklable(): postprocess = get_minimax_h3_post_process_func(SimpleNamespace()) postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess)) - video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5) + # decode() hands over quantised frames already laid out as (B, T, H, W, C). + video = torch.arange(2 * 2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 2, 4, 5, 3) audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6) result = postprocess((video, audio), output_type="np") assert isinstance(result["video"], list) + assert result["video"][0].dtype == np.uint8 assert result["video"][0].shape == (2, 4, 5, 3) + np.testing.assert_array_equal(result["video"][0], video[0].numpy()) np.testing.assert_array_equal(result["audio"], audio.numpy()) assert result["audio_sample_rate"] == 32000 assert result["fps"] == 24 @@ -595,7 +598,7 @@ def test_base_schedule_overrides_the_uniform_sigma_positions(): @pytest.mark.parametrize( "base_schedule", [ - [0.9, 0.5, 0.0], + [1.5, 0.5, 0.0], [1.0, 0.5, 0.1], [1.0, 0.5, 0.5, 0.0], [1.0], diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py new file mode 100644 index 000000000..8945b9b47 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py @@ -0,0 +1,498 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from vllm_omni.diffusion.models.minimax_h3.fasth3 import ( + FASTH3_BASE_MODEL, + FASTH3_BASE_SCHEDULE, + FASTH3_DENOISE_STEPS, + FASTH3_FORMAT, + FastH3AdapterError, + FastH3WeightFusion, + resolve_fasth3_fusion, +) +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv +from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas +from vllm_omni.errors import OmniClientError + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_RANK = 2 +_HIDDEN = 4 +_HEAD_DIM = 2 +_HEADS = 3 +_INNER = _HEAD_DIM * _HEADS # attention inner size +_FFN = 5 + + +def _factors(out_dim: int, in_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + """A rank-2 pair whose product is reproducible and non-symmetric.""" + a = torch.arange(_RANK * in_dim, dtype=torch.float32).reshape(_RANK, in_dim) / (in_dim * _RANK) + b = torch.arange(out_dim * _RANK, dtype=torch.float32).reshape(out_dim, _RANK) / (out_dim * _RANK) + return a, b + + +def _write_adapter( + path, + *, + tensors=None, + drop: str | None = None, + blocks: int = 1, + metadata: dict[str, str | None] | None = None, +) -> None: + """Write an artifact in the published ``fastvideo-lora-v2`` shape.""" + payload: dict[str, torch.Tensor] = {} + for block in range(blocks): + for suffix, (out_dim, in_dim) in { + "attn.to_q": (_INNER, _HIDDEN), + "attn.to_k": (_INNER, _HIDDEN), + "attn.to_v": (_INNER, _HIDDEN), + "attn.to_out.0": (_HIDDEN, _INNER), + "ff.net.0.proj": (2 * _FFN, _HIDDEN), + "ff.net.2": (_HIDDEN, _FFN), + }.items(): + a, b = _factors(out_dim, in_dim) + payload[f"transformer_blocks.{block}.{suffix}.lora_A.weight"] = a + payload[f"transformer_blocks.{block}.{suffix}.lora_B.weight"] = b + payload.update(tensors or {}) + if drop is not None: + payload.pop(drop, None) + path.parent.mkdir(parents=True, exist_ok=True) + # The published writer records what it emitted; mirror it so the fixtures + # are as self-describing as the real artifact. + declared = { + "format": FASTH3_FORMAT, + "finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1", + "base_model": FASTH3_BASE_MODEL, + "rank": str(_RANK), + "low_rank_tensors": str(sum(1 for key in payload if key.endswith((".lora_A.weight", ".lora_B.weight")))), + "diff_tensors": str(sum(1 for key in payload if key.endswith((".diff", ".diff_b")))), + "set_weight_tensors": str(sum(1 for key in payload if key.endswith(".set_weight"))), + } + # safetensors metadata is string-only, so a ``None`` override drops the key. + declared.update(metadata or {}) + save_file(payload, str(path), metadata={key: value for key, value in declared.items() if value is not None}) + + +def _claim(path, **kwargs) -> FastH3WeightFusion | None: + """Whatever from_path decides, including declining the artifact.""" + kwargs.setdefault("head_dim", _HEAD_DIM) + kwargs.setdefault("num_blocks", 1) + kwargs.setdefault("num_refiner_blocks", 0) + return FastH3WeightFusion.from_path(path, **kwargs) + + +def _load(path, **kwargs) -> FastH3WeightFusion: + """The same, for the tests that only make sense on a claimed adapter.""" + fusion = _claim(path, **kwargs) + assert fusion is not None, f"{path} was not claimed as a FastH3 adapter" + return fusion + + +def test_only_an_artifact_carrying_the_release_identity_is_claimed(tmp_path): + plain = tmp_path / "peft" / "adapter_model.safetensors" + plain.parent.mkdir(parents=True) + save_file({"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(plain)) + # No fastvideo-lora-v2 metadata: this is somebody else's LoRA and has to + # stay on the dynamic route rather than being fused. + assert _claim(plain.parent) is None + assert _claim(tmp_path / "missing") is None + + claimed = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(claimed) + assert _claim(claimed.parent) is not None + + +@pytest.mark.parametrize( + "identity", + [ + # An ordinary H3 adapter out of FastVideo's own extraction tools. + {"finetuned_model": "someone/minimax-h3-style-lora"}, + # Somebody else's adapter that merely names the student it imitates. + {"finetuned_model": "someone/fasth3-style-lora-for-h3"}, + {"finetuned_model": ""}, + {"finetuned_model": None}, + # A FastH3 name over a different base model. + {"base_model": "Wan-AI/Wan2.2-TI2V-5B"}, + ], +) +def test_a_generic_fastvideo_h3_adapter_stays_on_the_dynamic_route(identity, tmp_path): + """The fixture edits every block, so only the identity separates it from FastH3.""" + path = tmp_path / "generic" / "adapter_model.safetensors" + _write_adapter(path, metadata=identity) + + assert _claim(path.parent) is None + + +def test_a_claimed_artifact_must_declare_every_tensor_count(tmp_path): + """Declining is for other people's adapters; a FastH3 file is held to its word.""" + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, metadata={"diff_tensors": None}) + + with pytest.raises(FastH3AdapterError, match="omits diff_tensors"): + _claim(path.parent) + + +def test_the_published_bundle_root_is_refused_rather_than_guessed(tmp_path): + root = tmp_path / "FastVideo-FastH3-4-step-Preview-v1-LoRA" + for slug in ("dense-datafree", "vsa-datafree"): + _write_adapter(root / slug / "adapter_model.safetensors") + (root / "adapter_manifest.json").write_text(json.dumps({"schema_version": "fasth3-lora-bundle-v1"})) + + with pytest.raises(FastH3AdapterError, match="point --lora-path at one variant"): + _load(root) + # One variant inside it is unambiguous. + assert _claim(root / "dense-datafree") is not None + + +def test_a_multi_shard_non_fasth3_adapter_stays_on_the_dynamic_route(tmp_path): + # A plain PEFT LoRA saved as several shards must not hard-fail startup with a + # FastH3-specific message; it has to fall through to the dynamic LoRA route. + directory = tmp_path / "peft-sharded" + directory.mkdir() + for shard in ("adapter_model-00001-of-00002.safetensors", "adapter_model-00002-of-00002.safetensors"): + save_file( + {"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(directory / shard) + ) + + assert _claim(directory) is None + + +def test_low_rank_factors_reach_the_fused_projections(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + fusion = _load(path.parent) + + qkv = torch.zeros((3 * _INNER, _HIDDEN), dtype=torch.float32) + fused = fusion.fuse("blocks.0.attn.qkv_proj.weight", qkv).cpu() + # The checkpoint stores one head group at a time as [q, k, v], so the delta + # has to survive the loader's own unpacking as three separate projections. + q, k, v = torch.split( + _reorder_grouped_qkv_to_qkv(fused, num_query_groups=_HEADS, heads_per_group=1, head_dim=_HEAD_DIM), + [_INNER, _INNER, _INNER], + ) + a, b = _factors(_INNER, _HIDDEN) + expected = b @ a + for got in (q, k, v): + assert torch.allclose(got, expected, atol=1e-5) + + +def test_the_fused_mlp_delta_is_swapped_into_gate_first_order(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + fusion = _load(path.parent) + + fused = fusion.fuse("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))).cpu() + a, b = _factors(2 * _FFN, _HIDDEN) + value, gate = (b @ a).chunk(2, dim=0) + # The diffusers export is value-first; H3's fc1 is gate-first. + assert torch.allclose(fused, torch.cat((gate, value), dim=0), atol=1e-5) + + +def test_diff_and_diff_b_edit_weights_and_biases(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter( + path, + tensors={ + "transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.25), + "transformer_blocks.0.adaln_proj.linear.diff_b": torch.full((_HIDDEN,), -0.5), + "context_embedder.diff": torch.full((_HIDDEN, _HIDDEN), 0.125), + "context_embedder.diff_b": torch.full((_HIDDEN,), 2.0), + }, + ) + fusion = _load(path.parent) + + # A full-rank delta on an RMSNorm vector is exactly what a LoRA layer + # cannot express, and why this release is fused instead of switched. + assert torch.allclose(fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 1.25)) + assert torch.allclose( + fusion.fuse("blocks.0.adaln_proj.linear.bias", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 0.5) + ) + assert torch.allclose( + fusion.fuse("condition_proj.weight", torch.zeros((_HIDDEN, _HIDDEN))).cpu(), + torch.full((_HIDDEN, _HIDDEN), 0.125), + ) + assert torch.allclose(fusion.fuse("condition_proj.bias", torch.zeros(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 2.0)) + + +def test_a_streamed_checkpoint_keeps_the_parameters_the_adapter_never_edits(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 3.0)}) + fusion = _load(path.parent) + + untouched = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3) + streamed = dict( + fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN)), ("blocks.0.attn.q_norm.weight", untouched)]) + ) + assert torch.allclose(streamed["blocks.0.norm1.weight"].cpu(), torch.full((_HIDDEN,), 3.0)) + assert streamed["blocks.0.attn.q_norm.weight"] is untouched + + # validate_fully_applied() releases the deltas, so a second stream would + # fuse nothing at all; it has to fail rather than serve base H3 weights. + with pytest.raises(FastH3AdapterError, match="already been fused"): + next(fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN))])) + + +def test_the_fused_result_keeps_the_parameter_dtype(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.5)}) + fusion = _load(path.parent) + + fused = fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN, dtype=torch.bfloat16)) + assert fused.dtype == torch.bfloat16 + + +def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path): + dense = tmp_path / "dense" / "adapter_model.safetensors" + _write_adapter(dense) + assert _load(dense.parent).requires_vsa is False + + sparse = tmp_path / "vsa" / "adapter_model.safetensors" + _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))}) + assert _load(sparse.parent).requires_vsa is True + + +def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}) + # Dropping it silently would load a model that is not the distilled student. + with pytest.raises(FastH3AdapterError, match="name no known"): + _load(path.parent) + + +def test_an_unpaired_low_rank_factor_is_an_error(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, drop="transformer_blocks.0.attn.to_q.lora_B.weight") + + with pytest.raises(FastH3AdapterError, match="unpaired factor"): + _load(path.parent) + + +def test_a_delta_the_checkpoint_never_offered_is_reported(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + fusion = _load(path.parent) + + fusion.fuse("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN))) + # An unapplied delta is the failure that matters: the model would load and + # generate, just not as the distilled student. + with pytest.raises(FastH3AdapterError, match="never provided"): + fusion.validate_fully_applied() + + for name, shape in ( + ("blocks.0.attn.out_proj.weight", (_HIDDEN, _INNER)), + ("blocks.0.mlp.fc1.weight", (2 * _FFN, _HIDDEN)), + ("blocks.0.mlp.fc2.weight", (_HIDDEN, _FFN)), + ): + fusion.fuse(name, torch.zeros(shape)) + fusion.validate_fully_applied() + + +def _pipeline_stub(fusion=None): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "combined" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline._turbo_lora_adapter_ids = set() + pipeline._fasth3 = fusion + pipeline.od_config = SimpleNamespace() + pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0 + return pipeline + + +def _sampling(steps=FASTH3_DENOISE_STEPS, **kwargs) -> SimpleNamespace: + fields = {"num_inference_steps": steps, "extra_args": {}, "lora_request": None} + return SimpleNamespace(**{**fields, **kwargs}) + + +def _check_request(fusion, sampling) -> None: + fusion.check_request(sampling, video_shift=12.0, audio_shift=3.0) + + +def _fusion(tmp_path, **kwargs) -> FastH3WeightFusion: + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, **kwargs) + return _load(path.parent) + + +@pytest.mark.parametrize("task", ["fl2va", "ref2va"]) +def test_an_active_fasth3_fusion_restricts_requests_to_t2va(task, tmp_path): + pipeline = _pipeline_stub(_fusion(tmp_path)) + + with pytest.raises(OmniClientError, match="distills \\['t2va'\\] only"): + pipeline._resolve_task(task, {}) + assert pipeline._resolve_task("t2va", {}) == "t2va" + assert _pipeline_stub()._resolve_task(task, {}) == task + + +@pytest.mark.parametrize("num_inference_steps", [None, 5, 50]) +def test_fasth3_requires_one_forward_per_sigma_interval(num_inference_steps, tmp_path): + fusion = _fusion(tmp_path) + + # The five sigma points bound four forwards, and a request states forwards: + # the interval contract the pinned schedules use, and the count the step + # scheduler admits a request on before any pipeline hook runs. + with pytest.raises(OmniClientError, match=f"num_inference_steps={FASTH3_DENOISE_STEPS}"): + _check_request(fusion, _sampling(num_inference_steps)) + _check_request(fusion, _sampling()) + # The pinned branch coerces the string form; this one cannot disagree. + _check_request(fusion, _sampling(str(FASTH3_DENOISE_STEPS))) + + +def test_the_request_denoises_on_the_release_ladder(tmp_path): + pipeline = _pipeline_stub(_fusion(tmp_path)) + + positions, num_steps = pipeline._resolve_sigma_positions("t2va", _sampling()) + + # dmd_denoising_steps [999, 749, 500, 250] out of 1000, closed with the + # terminal 0.0: the pre-shift positions the student was distilled at, not + # the uniform ladder num_inference_steps would otherwise derive. + assert positions == (0.999, 0.749, 0.5, 0.25, 0.0) + # Five sigma points bound four transformer forwards, and forwards is the + # unit num_steps carries downstream (Cache-DiT, quality hints). + assert num_steps == FASTH3_DENOISE_STEPS == len(positions) - 1 + # H3's own per-modality shifts still apply on top, and are what turn those + # positions into the noise levels the four forwards run at. + assert minimax_h3_time_shift_sigmas(shift_scale=12.0, base_schedule=positions) == pytest.approx( + [0.999917, 0.972833, 0.923077, 0.8, 0.0], abs=1e-6 + ) + assert minimax_h3_time_shift_sigmas(shift_scale=3.0, base_schedule=positions) == pytest.approx( + [0.999666, 0.899520, 0.75, 0.5, 0.0], abs=1e-6 + ) + # Without a fused adapter the same request keeps the undistilled ladder. + assert _pipeline_stub()._resolve_sigma_positions("t2va", _sampling()) == (None, FASTH3_DENOISE_STEPS) + + +def test_the_contract_leaves_the_checkpoint_shifts_alone(tmp_path): + pipeline = _pipeline_stub(_fusion(tmp_path)) + pipeline.partition = "fl2va" + + pipeline._fasth3.check_serving_contract( + partition=pipeline.partition, + od_config=pipeline.od_config, + video_shift=pipeline.default_video_shift, + audio_shift=pipeline.default_audio_shift, + ) + + # The ladder is stated pre-shift, so H3's 12/3 shifts are what place the + # student on the levels it was distilled at; the contract must not touch them. + assert pipeline.default_video_shift == 12.0 + assert pipeline.default_audio_shift == 3.0 + + +@pytest.mark.parametrize( + "extra", + [{"flow_shift": 1.0}, {"audio_flow_shift": 1.0}, {"flow_shift": "bad"}], +) +def test_a_request_may_not_move_the_student_off_its_rungs(extra, tmp_path): + fusion = _fusion(tmp_path) + + with pytest.raises(OmniClientError, match="FastH3 requires"): + _check_request(fusion, _sampling(extra_args=extra)) + + _check_request(fusion, _sampling(extra_args={"flow_shift": 12.0, "audio_flow_shift": 3.0})) + + +def test_a_request_may_not_carry_a_lora_on_a_fused_server(tmp_path): + fusion = _fusion(tmp_path) + + # The dynamic LoRA manager is skipped for a fused adapter, so an adapter + # asked for per request would be neither applied nor reported. + with pytest.raises(OmniClientError, match="per-request lora is unavailable"): + _check_request(fusion, _sampling(lora_request=SimpleNamespace(lora_int_id=1))) + + +def test_offload_is_refused_because_it_bypasses_the_fusion(tmp_path): + fusion = _fusion(tmp_path) + for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload"): + # A host-weight plan installs the transformer without load_weights(), + # so the fusion and its completeness check would both be skipped. + with pytest.raises(ValueError, match="cannot be combined with"): + fusion.check_serving_contract( + partition="fl2va", + od_config=SimpleNamespace(**{flag: True}), + video_shift=12.0, + audio_shift=3.0, + ) + + +def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path): + sparse = tmp_path / "vsa" / "adapter_model.safetensors" + _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))}) + contract = {"od_config": SimpleNamespace(), "video_shift": 12.0, "audio_shift": 3.0} + with pytest.raises(ValueError, match="Video Sparse Attention variant"): + _load(sparse.parent).check_serving_contract(partition="fl2va", **contract) + + with pytest.raises(ValueError, match="cannot serve a Ref2VA partition"): + _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract) + + +def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path): + # Only low-rank factors are placed into H3's fused QKV and gate/up layouts, + # so a .diff aimed at one would otherwise be added transposed. + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.attn.to_q.diff": torch.ones((_INNER, _HIDDEN))}) + with pytest.raises(FastH3AdapterError, match="fused layout"): + _claim(path.parent) + + +def test_an_adapter_that_leaves_blocks_untouched_is_refused(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, blocks=2) + + # Every tensor in this file is well formed and pairs up, so the fusion's + # own completeness check passes; only the model's depth reveals that a + # third of it would have been served as base H3 weights. + with pytest.raises(FastH3AdapterError, match="missing=\\[2\\]"): + _claim(path.parent, num_blocks=3) + with pytest.raises(FastH3AdapterError, match="unknown=\\[1\\]"): + _claim(path.parent, num_blocks=1) + assert _claim(path.parent, num_blocks=2) is not None + + +def test_an_adapter_shorter_than_it_declares_is_refused(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, metadata={"low_rank_tensors": "724", "diff_tensors": "85"}) + + # The writer records what it emitted, which is the file's one statement + # about its own completeness. + with pytest.raises(FastH3AdapterError, match="declares low_rank_tensors=724 but carries 12"): + _claim(path.parent) + + +def test_only_a_configured_adapter_path_is_read(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + transformer = SimpleNamespace( + arch=SimpleNamespace(attention_head_dim=_HEAD_DIM, num_layers=1, token_refiner_num_layers=0) + ) + unbuilt = object() # a transformer whose arch must not be touched + + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=str(path)), transformer) is not None + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path)]), transformer) is not None + # Nothing configured, or more than one artifact, stays on the dynamic route + # without reading anything off the model. + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=None), unbuilt) is None + assert resolve_fasth3_fusion(SimpleNamespace(), unbuilt) is None + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path), str(path)]), unbuilt) is None + + +def test_a_pipeline_that_fused_its_adapter_needs_no_lora_manager(tmp_path): + assert _pipeline_stub(_fusion(tmp_path)).lora_is_fused is True + assert _pipeline_stub().lora_is_fused is False + + +def test_the_ladder_is_the_one_the_release_publishes(): + # FASTH3_DENOISE_STEPS is derived from it, so the two cannot drift apart. + assert FASTH3_BASE_SCHEDULE.base_schedule == (0.999, 0.749, 0.5, 0.25, 0.0) + assert FASTH3_BASE_SCHEDULE.num_inference_steps == FASTH3_DENOISE_STEPS == 4 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py index c4d581c2c..c4e581381 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py @@ -259,7 +259,11 @@ def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path): ("", "must not be empty"), (",,", "must not be empty"), ("1.0,oops,0.15,0.0", "is malformed"), - ("0.9,0.7,0.4,0.15,0.0", "is malformed"), + # A capped opening rung below 1.0 is a legal DMD2 schedule, so the + # out-of-range cases are a first position above 1.0 and a tail that + # never reaches clean latents. + ("1.5,0.7,0.4,0.15,0.0", "is malformed"), + ("1.0,0.7,0.4,0.15", "is malformed"), ], ) def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message): diff --git a/tests/diffusion/sched/test_dmd2_sigma_schedule.py b/tests/diffusion/sched/test_dmd2_sigma_schedule.py index 108cbe68c..8fbad959b 100644 --- a/tests/diffusion/sched/test_dmd2_sigma_schedule.py +++ b/tests/diffusion/sched/test_dmd2_sigma_schedule.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 import pytest from vllm_omni.diffusion.sched import DMD2SigmaSchedule @@ -18,7 +18,8 @@ def test_step_count_is_intervals_not_boundaries(): [ [], [1.0], - [0.9, 0.5, 0.0], + [1.5, 0.5, 0.0], + [0.0, 0.5, 0.0], [1.0, 0.5, 0.1], [1.0, 0.4, 0.6, 0.0], [1.0, 0.5, 0.5, 0.0], @@ -31,6 +32,15 @@ def test_malformed_positions_are_rejected(positions): DMD2SigmaSchedule.from_positions(positions) +def test_a_capped_first_position_is_a_schedule_not_a_malformed_one(): + # A DMD2 student trained with max_timestep_ratio < 1 never sees pure noise, + # so its opening rung sits just below 1.0 (FastH3: 999/1000). + schedule = DMD2SigmaSchedule.from_positions([0.999, 0.749, 0.5, 0.25, 0.0]) + + assert schedule.base_schedule[0] == 0.999 + assert schedule.num_inference_steps == 4 + + def test_one_base_schedule_drives_several_shift_scales(): schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0]) diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 03ddb8786..9f643d8ca 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -520,6 +520,32 @@ def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch): assert len(coerce_calls) == 1 +def test_uint8_frames_reach_the_muxer_without_a_float_round_trip(): + """uint8 is the muxer's own dtype, so normalisation must leave it alone. + + MiniMax-H3 quantises on the accelerator to keep the response small across + its two process hops. Re-normalising to float here would undo that at the + last step and cost a full-size conversion in each direction. + """ + video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3) + + frames, frame_shape, common_dtype = video_api_utils._prepare_video_frames(video) + + assert common_dtype == np.dtype(np.uint8) + assert frame_shape == (4, 5, 3) + np.testing.assert_array_equal(np.stack(frames), video) + + +def test_non_uint8_integer_frames_are_still_normalised(): + """Only uint8 is the muxer's dtype; other integer payloads still scale.""" + video = np.full((2, 4, 5, 3), 255, dtype=np.int32) + + frames, _, common_dtype = video_api_utils._prepare_video_frames(video) + + assert np.issubdtype(common_dtype, np.floating) + np.testing.assert_allclose(frames[0], 1.0) + + @pytest.mark.parametrize( "video", [ @@ -552,9 +578,10 @@ def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch, monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux) prepared_frames, frame_shape, _ = video_api_utils._prepare_video_frames(video) - reference = np.rint(np.clip(np.stack([frame[..., :3] for frame in prepared_frames]), 0.0, 1.0) * 255.0).astype( - np.uint8 - ) + stacked = np.stack([frame[..., :3] for frame in prepared_frames]) + # uint8 frames are the muxer's own dtype and reach it unchanged; everything + # else is still quantised out of the normalised [0, 1] range. + reference = stacked if stacked.dtype == np.uint8 else np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8) assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video" assert len(mux_inputs) == 1 diff --git a/vllm_omni/diffusion/distributed/a2a_permute.py b/vllm_omni/diffusion/distributed/a2a_permute.py index b4e951d34..4584fbfd2 100644 --- a/vllm_omni/diffusion/distributed/a2a_permute.py +++ b/vllm_omni/diffusion/distributed/a2a_permute.py @@ -142,7 +142,7 @@ def _get_symm_buffer( stream_id=stream_id, ) _SYMM_WORKSPACES[key] = workspace - return workspace.handle.get_buffer(workspace.handle.rank, symm_shape, dtype) + return workspace.allocation[:required_bytes].view(dtype).view(symm_shape) def clear_a2a_permute_workspaces() -> None: diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py new file mode 100644 index 000000000..21f77828a --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""FastVideo FastH3: a four-step DMD2 student of MiniMax-H3. + +FastH3 replaces H3's 49 denoiser evaluations with four. It ships as an adapter +over the base checkpoint rather than as a full release, so it reuses H3's text +encoder, video VAE, audio VAE, tokenizers and schedulers unchanged. + +The artifact is *not* a PEFT LoRA, and it is not request-switchable. Its own +metadata states the reconstruction as:: + + W = W_base + lora_B @ lora_A; then .diff/.diff_b added and .set_weight assigned + +so besides rank-64 factors it carries full-rank ``.diff``/``.diff_b`` deltas for +RMSNorm weights, biases, patch projections and the final layer - none of which a +LoRA layer can express - and the VSA variants add ``.set_weight`` tensors for +compression gates that do not exist in the base transformer at all. The adapter +is therefore fused into the checkpoint stream at load time, before the weights +are sharded, which is also what the release's model card requires. + +The low-rank factors carry no alpha: the reconstruction adds ``lora_B @ lora_A`` +directly, i.e. a scale of exactly 1. + +Two checkpoint spellings meet here. The adapter is written in the diffusers +namespace (``transformer_blocks.0.attn.to_q``) while vLLM-Omni loads H3's native +one (``blocks.0.attn.qkv_proj``), whose attention and MLP projections are fused. +Every mapping and layout convention below was verified tensor by tensor against +the released full checkpoint (``FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree``): +``W_base + delta`` reproduces it to bf16 rounding. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from safetensors import safe_open +from vllm.logger import init_logger + +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from .minimax_h3_transformer import MiniMaxH3DiTModel + +logger = init_logger(__name__) + +# FastVideo's generic adapter container, not a FastH3 marker: their tools emit +# it for ordinary H3 adapters too. +FASTH3_FORMAT = "fastvideo-lora-v2" +FASTH3_MANIFEST = "adapter_manifest.json" +# The release identity: the distilled student the adapter came from, over the +# base it edits. The name is held to FastVideo's own namespace as well as the +# student's name, so an unrelated adapter that merely mentions FastH3 is not +# claimed. +FASTH3_BASE_MODEL = "MiniMaxAI/MiniMax-H3" +_FASTH3_IDENTITY_KEY = "finetuned_model" +_FASTH3_IDENTITY_NAMESPACE = "fastvideo/" +_FASTH3_IDENTITY_MARKER = "fasth3" + +# The rectified-flow positions the student was distilled at, and the ladder the +# server samples on. The release states them as `dmd_denoising_steps` +# [999, 749, 500, 250] (`sampling.base_timesteps` in the bundle manifest): +# timestep indices out of 1000, i.e. pre-shift positions, closed here with the +# terminal 0.0 every rectified-flow schedule ends on. The opening rung is 0.999 +# rather than 1.0 because training capped the noise level there +# (`max_timestep_ratio`). +# +# These are positions, not final sigmas: H3's schedulers apply their own +# per-modality shift on top (12 for video, 3 for audio), which is what +# reproduces the levels the student saw. Nothing here overrides those shifts. +FASTH3_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.749, 0.5, 0.25, 0.0)) +# Its five points bound four transformer forwards, one per sigma interval. That +# count is the one a request states, the one Cache-DiT is refreshed with and the +# one step execution admits a request on - the same interval contract H3's +# pinned checkpoint schedules and native LoRAs already use. +FASTH3_DENOISE_STEPS = FASTH3_BASE_SCHEDULE.num_inference_steps +# Preview v1 distills the text-to-video-and-audio path only. +FASTH3_SUPPORTED_TASKS = frozenset({"t2va"}) + +_LORA_A = ".lora_A.weight" +_LORA_B = ".lora_B.weight" +_DIFF = ".diff" +_DIFF_B = ".diff_b" +_SET_WEIGHT = ".set_weight" + +# Adapter module prefix -> the native parameter it edits, minus the +# ``.weight``/``.bias`` suffix. +_MODEL_LEVEL_TARGETS = { + "proj_in": "video_patch_proj", + "proj_out": "final_layer.video_out", + "audio_proj_in": "audio_patch_proj", + "audio_proj_out": "final_layer.audio_out", + "context_embedder": "condition_proj", + "time_embedder.linear_1": "time_embedder.proj_in", + "time_embedder.linear_2": "time_embedder.proj_out", + "norm_out.linear": "final_layer.adaln_proj.linear", + "norm_out.norm": "final_layer.norm", +} + +# Per-block adapter suffix -> (native suffix, how a delta enters the native +# parameter). H3 stores attention as one grouped QKV matrix and the MLP as one +# fused gate/up matrix, so those deltas need placing rather than adding. +_PLAIN, _QKV_Q, _QKV_K, _QKV_V, _SWAP_HALVES = "plain", "q", "k", "v", "swap_halves" +_QKV_SLOTS = (_QKV_Q, _QKV_K, _QKV_V) +_BLOCK_TARGETS = { + "attn.to_q": ("attn.qkv_proj", _QKV_Q), + "attn.to_k": ("attn.qkv_proj", _QKV_K), + "attn.to_v": ("attn.qkv_proj", _QKV_V), + "attn.to_out.0": ("attn.out_proj", _PLAIN), + "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES), + "ff.net.2": ("mlp.fc2", _PLAIN), + "adaln_proj.linear": ("adaln_proj.linear", _PLAIN), + "norm1": ("norm1", _PLAIN), + "norm2": ("norm2", _PLAIN), +} + +# Adapter block prefix -> native block prefix. +_BLOCK_PREFIXES = ( + ("token_refiner.refiner_blocks.", "token_refiner.blocks."), + ("transformer_blocks.", "blocks."), +) + + +class FastH3AdapterError(ValueError): + """The artifact is a FastH3 adapter, but it cannot be applied as one.""" + + +@dataclass +class _ParamPatch: + """Everything the adapter contributes to one native parameter.""" + + # layout -> (lora_A, lora_B). A grouped QKV parameter collects three. + low_rank: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = field(default_factory=dict) + diff: torch.Tensor | None = None + layout: str = _PLAIN + + +def _swap_halves(tensor: torch.Tensor) -> torch.Tensor: + """Exchange the two halves of a fused gate/up matrix. + + The diffusers export stores the feed-forward projection value-first while + H3's native ``mlp.fc1`` is gate-first, so a delta computed in the diffusers + layout has to be swapped before it can be added to the native parameter. + """ + if tensor.shape[0] % 2: + raise FastH3AdapterError(f"fused gate/up delta must split evenly, got {tuple(tensor.shape)}") + first, second = tensor.chunk(2, dim=0) + return torch.cat((second, first), dim=0) + + +def _place_in_grouped_qkv(deltas: Mapping[str, torch.Tensor], *, head_dim: int) -> torch.Tensor: + """Interleave per-projection deltas into H3's grouped QKV layout. + + The checkpoint stores one head group at a time as ``[q, k, v]``, which is + what :func:`_reorder_grouped_qkv_to_qkv` unpacks on the way in. A delta + built from the separate diffusers projections has to be folded back into + that order. + """ + missing = sorted(set(_QKV_SLOTS) - set(deltas)) + if missing: + raise FastH3AdapterError(f"grouped QKV delta is missing its {missing} projections") + parts = [] + for slot in _QKV_SLOTS: + delta = deltas[slot] + if delta.shape[0] % head_dim: + raise FastH3AdapterError( + f"QKV {slot} delta rows {delta.shape[0]} are not a multiple of head_dim {head_dim}" + ) + parts.append(delta.reshape(delta.shape[0] // head_dim, head_dim, *delta.shape[1:])) + groups = parts[0].shape[0] + if any(part.shape[0] != groups for part in parts): + raise FastH3AdapterError("QKV projections disagree on the number of head groups") + return torch.cat(parts, dim=1).reshape(groups * 3 * head_dim, *parts[0].shape[2:]) + + +def _resolve_native_target(module: str) -> tuple[str, str, tuple[str, int] | None] | None: + """Map an adapter module path to ``(native path, layout, block)``. + + ``block`` is the ``(native block prefix, index)`` the module sits in, or + ``None`` for a model-level one. Coverage is checked per block, so the index + has to survive the mapping rather than being folded into the path. + """ + native = _MODEL_LEVEL_TARGETS.get(module) + if native is not None: + return native, _PLAIN, None + for adapter_prefix, native_prefix in _BLOCK_PREFIXES: + if not module.startswith(adapter_prefix): + continue + remainder = module[len(adapter_prefix) :] + index, _, suffix = remainder.partition(".") + if not index.isdigit(): + return None + target = _BLOCK_TARGETS.get(suffix) + if target is None: + return None + native_suffix, layout = target + return f"{native_prefix}{index}.{native_suffix}", layout, (native_prefix, int(index)) + return None + + +def _split_adapter_key(name: str) -> tuple[str, str] | None: + """Split an adapter tensor name into ``(module path, role)``.""" + for marker, role in ( + (_LORA_A, "lora_a"), + (_LORA_B, "lora_b"), + (_DIFF_B, "diff_b"), + (_DIFF, "diff"), + (_SET_WEIGHT, "set_weight"), + ): + if name.endswith(marker): + return name[: -len(marker)], role + return None + + +def _is_fasth3_release(metadata: Mapping[str, str]) -> bool: + """Whether this is a FastH3 release rather than merely FastVideo's format. + + Claiming on the container alone would fuse an ordinary H3 adapter into the + checkpoint and put every request on a four-step schedule it was never + distilled for, and so would a bare ``fasth3`` substring: an adapter someone + else named after the student is not the student. + """ + if metadata.get("format") != FASTH3_FORMAT: + return False + identity = metadata.get(_FASTH3_IDENTITY_KEY, "").lower() + if not identity.startswith(_FASTH3_IDENTITY_NAMESPACE) or _FASTH3_IDENTITY_MARKER not in identity: + return False + # Only held against a file that states it. + base_model = metadata.get("base_model") + return base_model is None or base_model.casefold() == FASTH3_BASE_MODEL.casefold() + + +def _check_declared_counts( + metadata: Mapping[str, str], + counted: Mapping[str, int], + *, + weights_path: Path, +) -> None: + """Hold the artifact to the tensor counts its own metadata declares. + + The writer of the ``fastvideo-lora-v2`` format records how many tensors of + each kind it emitted, which is the one statement in the file about its own + completeness. A truncated or partially re-exported artifact is otherwise + indistinguishable from a small one, so a claimed file has to carry all of + them rather than opt out by omission. + """ + for key, seen in counted.items(): + declared = metadata.get(key) + if declared is None: + raise FastH3AdapterError( + f"{weights_path} declares the FastH3 identity but omits {key}; " + "a claimed adapter has to state what it emitted" + ) + try: + expected = int(declared) + except (TypeError, ValueError) as exc: + raise FastH3AdapterError(f"{weights_path} declares a non-numeric {key}={declared!r}") from exc + if expected != seen: + raise FastH3AdapterError( + f"{weights_path} declares {key}={expected} but carries {seen}; the adapter is incomplete " + "and would leave most of the transformer on base H3 weights" + ) + + +def _check_block_coverage( + seen: Mapping[str, set[int]], + *, + expected: Mapping[str, int], + weights_path: Path, +) -> None: + """Every block of the model this adapter is loaded against must be edited. + + The release drops tensors training left unchanged, so per-parameter + coverage is legitimately sparse - but a distilled student touches every + block, so a block with no edits at all means the artifact does not match + this model. + """ + for prefix, count in expected.items(): + indices = seen.get(prefix, set()) + wanted = set(range(count)) + if indices == wanted: + continue + missing = sorted(wanted - indices) + extra = sorted(indices - wanted) + raise FastH3AdapterError( + f"{weights_path} edits {len(indices)} of the model's {count} {prefix}* blocks " + f"(missing={missing[:5]}, unknown={extra[:5]}); it is not an adapter for this checkpoint" + ) + + +class FastH3WeightFusion: + """Fuse a FastH3 adapter into the H3 checkpoint stream as it is loaded.""" + + def __init__( + self, + *, + source: Path, + patches: Mapping[str, _ParamPatch], + head_dim: int, + requires_vsa: bool, + ) -> None: + self._source = source + self._patches = dict(patches) + self._head_dim = head_dim + self.requires_vsa = requires_vsa + self._applied: set[str] = set() + self._device: torch.device | None = None + + @property + def source(self) -> Path: + return self._source + + @property + def base_schedule(self) -> tuple[float, ...]: + """The rectified-flow positions this student samples on. + + The fused checkpoint is a four-step student, so the ladder comes from + the release rather than from the many-step teacher's metadata or from + the uniform one ``num_inference_steps`` would otherwise derive. + """ + return FASTH3_BASE_SCHEDULE.base_schedule + + @classmethod + def from_path( + cls, + path: str | Path, + *, + head_dim: int, + num_blocks: int, + num_refiner_blocks: int, + ) -> FastH3WeightFusion | None: + """Build a fusion from an adapter file or directory, else ``None``. + + Returning ``None`` keeps every other ``--lora-path`` artifact on the + dynamic LoRA route; only a file carrying the FastH3 release identity is + claimed here. + + The block counts are the model's, and the artifact has to cover them: + claiming a partial adapter would switch the server onto the four-step + contract while most of the transformer still held base H3 weights. + """ + weights_path = _resolve_adapter_file(path) + if weights_path is None: + return None + + patches: dict[str, _ParamPatch] = {} + gate_tensors: list[str] = [] + unmapped: list[str] = [] + blocks_seen: dict[str, set[int]] = {} + counted = {"low_rank_tensors": 0, "diff_tensors": 0} + with safe_open(weights_path, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if not _is_fasth3_release(metadata): + return None + for name in checkpoint.keys(): + split = _split_adapter_key(name) + if split is None: + unmapped.append(name) + continue + module, role = split + if role == "set_weight": + # A VSA compression gate: a module the base transformer does + # not have, so there is nothing to fuse it into. + gate_tensors.append(name) + continue + target = _resolve_native_target(module) + if target is None: + unmapped.append(name) + continue + native_module, layout, block = target + if block is not None: + blocks_seen.setdefault(block[0], set()).add(block[1]) + counted["diff_tensors" if role in ("diff", "diff_b") else "low_rank_tensors"] += 1 + native_param = f"{native_module}.{'bias' if role == 'diff_b' else 'weight'}" + patch = patches.setdefault(native_param, _ParamPatch(layout=layout)) + tensor = checkpoint.get_tensor(name) + if role in ("diff", "diff_b"): + if patch.diff is not None: + raise FastH3AdapterError(f"duplicate {role} for {native_param}") + patch.diff = tensor + else: + a, b = patch.low_rank.get(layout, (None, None)) + patch.low_rank[layout] = (tensor, b) if role == "lora_a" else (a, tensor) + + if unmapped: + raise FastH3AdapterError( + f"FastH3 adapter at {weights_path} has {len(unmapped)} tensors that name no known " + f"H3 parameter: {sorted(unmapped)[:5]}" + ) + _check_declared_counts( + metadata, + {**counted, "set_weight_tensors": len(gate_tensors)}, + weights_path=weights_path, + ) + _check_block_coverage( + blocks_seen, + expected={"blocks.": num_blocks, "token_refiner.blocks.": num_refiner_blocks}, + weights_path=weights_path, + ) + for native_param, patch in patches.items(): + for slot, (a, b) in patch.low_rank.items(): + if a is None or b is None: + raise FastH3AdapterError(f"FastH3 adapter has an unpaired factor for {native_param} slot {slot!r}") + # Only the low-rank factors are placed into H3's fused QKV and + # gate/up layouts; a full-rank delta is added as it comes, so one + # aimed at a fused parameter would silently land transposed. + if patch.diff is not None and patch.layout != _PLAIN: + raise FastH3AdapterError( + f"FastH3 adapter carries a full-rank delta for {native_param}, which H3 stores in the " + f"{patch.layout!r} fused layout; this loader can only place low-rank factors there" + ) + + fusion = cls( + source=weights_path, + patches=patches, + head_dim=head_dim, + requires_vsa=bool(gate_tensors), + ) + logger.info( + "FastH3 adapter %s: rank=%s, parameters patched=%d, low-rank=%s, diff=%s, set_weight=%d", + weights_path, + metadata.get("rank", "?"), + len(patches), + metadata.get("low_rank_tensors", "?"), + metadata.get("diff_tensors", "?"), + len(gate_tensors), + ) + return fusion + + def _compute_device(self, weight: torch.Tensor) -> torch.device: + """Where to reconstruct a delta. + + H3's per-block modulation projection is 96768x2688, so rebuilding all + 343 patched parameters is a few TFLOP of rank-64 products. On CPU that + adds minutes to a load that already has a startup deadline, so the + accelerator does the arithmetic whenever there is one. + """ + if weight.device.type != "cpu": + return weight.device + if self._device is None: + # Ask the platform rather than PyTorch's global accelerator + # registry, so an out-of-tree backend controls its own placement. + self._device = current_omni_platform.get_torch_device() + return self._device + + @staticmethod + def _widen(tensor: torch.Tensor, device: torch.device) -> torch.Tensor: + """Move to ``device``, then widen to float32. + + Asking ``Tensor.to`` for a device and a dtype at once converts on the + host and ships twice the bytes; splitting it moves bfloat16 and widens + on the accelerator. + """ + return tensor.to(device, non_blocking=True).to(torch.float32) + + def fuse(self, name: str, weight: torch.Tensor) -> torch.Tensor: + """Return ``weight`` with this adapter's contribution added.""" + patch = self._patches.get(name) + if patch is None: + return weight + self._applied.add(name) + + device = self._compute_device(weight) + + delta: torch.Tensor | None = None + if patch.low_rank: + if patch.layout in _QKV_SLOTS: + per_slot = { + slot: self._widen(b, device) @ self._widen(a, device) for slot, (a, b) in patch.low_rank.items() + } + delta = _place_in_grouped_qkv(per_slot, head_dim=self._head_dim) + else: + a, b = patch.low_rank[patch.layout] + if patch.layout == _SWAP_HALVES: + # Permuting the rows of B permutes the rows of the product, + # so swap the rank-64 factor instead of the full delta. + b = _swap_halves(b) + delta = self._widen(b, device) @ self._widen(a, device) + if patch.diff is not None: + diff = self._widen(patch.diff, device) + delta = diff if delta is None else delta + diff + if delta is None: + return weight + if delta.shape != weight.shape: + raise FastH3AdapterError( + f"FastH3 delta for {name} has shape {tuple(delta.shape)}, parameter is {tuple(weight.shape)}" + ) + # Leave the result on the compute device. These weights are bound for + # the accelerator anyway, so returning them to host memory would pay a + # device-to-host copy of the whole checkpoint only for the loader to + # send it straight back: measured at 152s against 15s for 60 GiB of + # patched projections, against 17s for the unavoidable upload alone. + # Fold the base weight into the freshly built delta in place. Promoting + # the weight to float32 on its own would allocate two more buffers the + # size of the parameter, and H3's largest patched projection is 0.5 GiB. + return delta.add_(weight.to(device, non_blocking=True)).to(weight.dtype) + + def apply(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]: + """Fuse every streamed checkpoint tensor on its way into the model.""" + if self._applied: + # ``validate_fully_applied`` released the deltas, so a second stream + # would fuse nothing and then pass its own completeness check: the + # server would serve base H3 weights on the student's ladder. + raise FastH3AdapterError(f"{self._source} has already been fused into this checkpoint") + for name, weight in weights: + yield name, self.fuse(name, weight) + + def validate_fully_applied(self) -> None: + """Close the fusion: every edit must have met its parameter. + + A silently unapplied delta is the failure mode that matters here: the + model would load and generate, just not as the distilled student. The + weights are loaded once, so the mapped payloads are dropped afterwards + rather than held for the life of the process. + """ + missing = sorted(set(self._patches) - self._applied) + if missing: + raise FastH3AdapterError( + f"FastH3 adapter edits {len(missing)} parameters the checkpoint never provided: {missing[:5]}" + ) + for patch in self._patches.values(): + patch.low_rank.clear() + patch.diff = None + + def check_serving_contract( + self, + *, + partition: str, + od_config: Any, + video_shift: float, + audio_shift: float, + ) -> None: + """Hold a starting server to the ladder this student was trained on.""" + if partition == "ref2va": + raise ValueError("FastH3 preview v1 distills T2VA only, so it cannot serve a Ref2VA partition") + offloads = [ + flag + for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload") + if getattr(od_config, flag, False) + ] + if offloads: + # A host-weight plan installs the transformer without going through + # load_weights(), which is where the fusion and its completeness + # check live. Serving base H3 weights under a four-step schedule + # would otherwise degrade output with nothing to signal it. + raise ValueError( + f"FastH3 is fused while the checkpoint streams in, so it cannot be combined with " + f"{sorted(offloads)}. Serve it without offload." + ) + if self.requires_vsa: + raise ValueError( + f"{self.source} is a Video Sparse Attention variant of FastH3, and its " + "compression gates have no counterpart in vLLM-Omni's dense H3 attention. Use the " + "release's Dense variant until VSA lands for H3's packed sequence." + ) + logger.info( + "FastH3 adapter active: sigma points %s for %d transformer forwards, " + "flow_shift=%g, audio_flow_shift=%g, tasks=%s", + list(self.base_schedule), + FASTH3_DENOISE_STEPS, + video_shift, + audio_shift, + sorted(FASTH3_SUPPORTED_TASKS), + ) + + def check_task(self, task: str) -> None: + """Refuse a task this preview never distilled.""" + if task not in FASTH3_SUPPORTED_TASKS: + raise OmniClientError( + f"FastH3 preview v1 distills {sorted(FASTH3_SUPPORTED_TASKS)} only, got task={task!r}" + ) + + def check_request(self, sampling: Any, *, video_shift: float, audio_shift: float) -> None: + """Refuse a request that would sample the student off its rungs.""" + if sampling.lora_request is not None: + # The adapter is already in the weights and the dynamic LoRA manager + # is skipped, so nothing would apply the requested one. Serving the + # request anyway would quietly ignore it. + raise OmniClientError( + f"this server fused {self.source} into the checkpoint at startup, so per-request " + "lora is unavailable; drop the lora field" + ) + if int(sampling.num_inference_steps or 0) != FASTH3_DENOISE_STEPS: + raise OmniClientError( + f"FastH3 is a four-step student and requires num_inference_steps={FASTH3_DENOISE_STEPS} " + "(one transformer forward per sigma interval)" + ) + # The checkpoint's per-modality shifts turn the release's positions into + # the noise levels the student was distilled at, so a request that moves + # them samples where it was never trained. + extra = sampling.extra_args or {} + for key, expected in (("flow_shift", video_shift), ("audio_flow_shift", audio_shift)): + try: + requested = float(extra.get(key, expected)) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"FastH3 requires {key}={expected:g}") from exc + if not math.isclose(requested, expected): + raise OmniClientError(f"FastH3 requires {key}={expected:g}, got {requested:g}") + + +def resolve_fasth3_fusion(od_config: Any, transformer: MiniMaxH3DiTModel) -> FastH3WeightFusion | None: + """Claim ``--lora-path`` when it points at a FastH3 adapter. + + FastH3 rewrites RMSNorm weights and biases, so it cannot be expressed as a + request-switchable LoRA and is fused into the checkpoint instead. Any other + artifact returns None here and stays on the dynamic LoRA route. + """ + lora_path = getattr(od_config, "lora_path", None) + if isinstance(lora_path, (list, tuple)): + if len(lora_path) != 1: + return None + lora_path = lora_path[0] + if not lora_path: + return None + # Placing a delta into the fused QKV parameter needs the head size, and + # completeness is judged against the model's depth, so the architecture is + # only read once an adapter is actually configured. + arch = transformer.arch + return FastH3WeightFusion.from_path( + lora_path, + head_dim=arch.attention_head_dim, + num_blocks=arch.num_layers, + num_refiner_blocks=arch.token_refiner_num_layers, + ) + + +def _safetensors_metadata(path: Path) -> Mapping[str, str]: + """The header metadata of a safetensors file, or ``{}`` if unreadable.""" + try: + with safe_open(path, framework="pt", device="cpu") as checkpoint: + return checkpoint.metadata() or {} + except Exception: # noqa: BLE001 - a path that will not open is simply not ours + return {} + + +def _resolve_adapter_file(path: str | Path) -> Path | None: + """Find the single adapter file at ``path``, or ``None``.""" + candidate = Path(path) + if candidate.is_file(): + return candidate if candidate.suffix == ".safetensors" else None + if not candidate.is_dir(): + return None + named = candidate / "adapter_model.safetensors" + if named.is_file(): + return named + files = sorted(candidate.glob("*.safetensors")) + if len(files) == 1: + return files[0] + # The published repository bundles four variants under one root, marked by an + # adapter_manifest.json. Such a bundle is ambiguous rather than loadable - but + # only hold that against a directory that actually carries FastH3 artifacts; + # an unrelated multi-shard LoRA stays on the dynamic route via ``None``. + if (candidate / FASTH3_MANIFEST).is_file() or any( + _is_fasth3_release(_safetensors_metadata(file)) for file in files + ): + raise FastH3AdapterError( + f"{candidate} holds several FastH3 adapters; point --lora-path at one variant " + "(for example dense-datafree/adapter_model.safetensors)" + ) + return None + + +__all__ = [ + "FASTH3_BASE_MODEL", + "FASTH3_BASE_SCHEDULE", + "FASTH3_DENOISE_STEPS", + "FASTH3_FORMAT", + "FASTH3_SUPPORTED_TASKS", + "FastH3AdapterError", + "FastH3WeightFusion", + "resolve_fasth3_fusion", +] 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 d0f63fe8e..2a13baa36 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -117,6 +117,7 @@ from .denoise_loop import ( minimax_h3_publish_denoise_progress, ) from .encoder import MiniMaxH3Qwen3VLEncoder +from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion from .lora import load_minimax_h3_turbo_lora from .minimax_h3_transformer import ( MiniMaxH3Attention, @@ -354,17 +355,24 @@ def _minimax_h3_post_process(output, output_type: str = "np"): The callable crosses the multiprocessing result queue, so it must remain a module-level function that the standard pickle module can resolve. + + ``_prepare_minimax_h3_video_output`` already quantises the video to uint8 + frames on the accelerator, so there is nothing left to scale or transpose + here. """ if not isinstance(output, tuple) or len(output) != 2: return output video, audio = output + if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4): + # Float or channel-first frames would reach the muxer as a black or + # banded video rather than as an error. + raise ValueError( + f"MiniMax-H3 post-processing expects (B, T, H, W, C) uint8, got {tuple(video.shape)} {video.dtype}" + ) if output_type == "latent": return output if output_type == "np": - if video.dtype == torch.uint8 and video.ndim == 5 and video.shape[-1] == 3: - video = video.detach().cpu().numpy() - else: - video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy() + video = video.detach().cpu().numpy() audio = audio.detach().float().cpu().numpy() video = [sample for sample in video] return { @@ -678,6 +686,8 @@ class MiniMaxH3Pipeline( # Only distilled releases pin a schedule, so the default keeps the legacy # uniform path available to partially constructed pipelines. _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {} + # Set from --lora-path during construction; absent means no FastH3 adapter. + _fasth3: FastH3WeightFusion | None = None def _load_diffusion_lora_adapter( self, @@ -948,6 +958,15 @@ class MiniMaxH3Pipeline( quant_config=transformer_quant_config, ) + self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer) + if self._fasth3 is not None: + self._fasth3.check_serving_contract( + partition=self.partition, + od_config=od_config, + video_shift=self.default_video_shift, + audio_shift=self.default_audio_shift, + ) + if self.load_text_encoder: self.tokenizer = Qwen2TokenizerFast.from_pretrained( str(model_path), @@ -1056,7 +1075,12 @@ class MiniMaxH3Pipeline( raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous") loaded_prefixes.add(prefix) component = getattr(self, prefix.removesuffix(".")) - loaded = component.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights) + stream = ((name[len(prefix) :], tensor) for name, tensor in grouped_weights) + if prefix == "transformer." and self._fasth3 is not None: + # 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) + loaded = component.load_weights(stream) if prefix != "text_encoder.": component.post_load_weights() loaded_with_prefix.update(prefix + name for name in loaded) @@ -1069,13 +1093,49 @@ class MiniMaxH3Pipeline( if component is None: continue loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) + if self._fasth3 is not None: + self._fasth3.validate_fully_applied() return loaded_with_prefix + @property + def lora_is_fused(self) -> bool: + """True when --lora-path was consumed as a load-time weight fusion.""" + return self._fasth3 is not None + def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel: if task == "ref2va" and hasattr(self, "transformers_ref"): return self.transformers_ref return self.transformer + def _resolve_sigma_positions(self, task: str, sampling: Any) -> tuple[tuple[float, ...] | None, int]: + """Pick the rectified-flow positions this request denoises on. + + Returns them explicitly, or ``None`` to leave the uniform ladder to be + derived from the step count, together with the count the rest of the + request speaks in. + """ + if self._fasth3 is not None: + # A fused student carries its own positions; the checkpoint + # underneath it is the many-step teacher, whose schedule does not + # apply. Its five points bound four transformer forwards, and + # forwards is the unit ``check_request``, the pinned-schedule branch + # below and Cache-DiT all speak in. + positions = self._fasth3.base_schedule + return positions, len(positions) - 1 + sigma_schedule = self._sigma_schedule_for_request(sampling, task) + if sigma_schedule is None: + return None, int(sampling.num_inference_steps or 50) + # The schedule lists sigma boundaries; the denoise loop runs one step + # per interval, and that count is what requests and Cache-DiT speak in. + num_steps = sigma_schedule.num_inference_steps + requested_steps = sampling.num_inference_steps + if requested_steps is not None and int(requested_steps) != num_steps: + raise OmniClientError( + "this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps " + f"must be {num_steps} or omitted, got {int(requested_steps)}" + ) + return sigma_schedule.base_schedule, num_steps + def _base_schedule_for_task(self, task: str) -> DMD2SigmaSchedule | None: """Return the distilled schedule of the partition that serves ``task``.""" partition = "ref2va" if task == "ref2va" else "fl2va" @@ -1109,6 +1169,8 @@ class MiniMaxH3Pipeline( raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only") if has_native_lora and task != "t2va": raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only") + if self._fasth3 is not None: + self._fasth3.check_task(task) return task def _resolve_shape( @@ -2121,6 +2183,12 @@ class MiniMaxH3Pipeline( self._validate_turbo_sampling(sampling) if has_native_lora: self._validate_native_sampling(sampling, task=task) + if self._fasth3 is not None: + self._fasth3.check_request( + sampling, + video_shift=self.default_video_shift, + audio_shift=self.default_audio_shift, + ) raw_image = multi_modal_data.get("image") raw_videos = multi_modal_data.get("video") @@ -2337,22 +2405,7 @@ class MiniMaxH3Pipeline( ref_audio_t = audio_lengths[0] seed = int(sampling.seed if sampling.seed is not None else 42) - sigma_schedule = self._sigma_schedule_for_request(sampling, task) - if sigma_schedule is None: - base_schedule = None - num_steps = int(sampling.num_inference_steps or 50) - else: - # The schedule lists sigma boundaries; the denoise loop runs one - # step per interval, and that count is what requests and Cache-DiT - # speak in. - base_schedule = sigma_schedule.base_schedule - num_steps = sigma_schedule.num_inference_steps - requested_steps = sampling.num_inference_steps - if requested_steps is not None and int(requested_steps) != num_steps: - raise OmniClientError( - "this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps " - f"must be {num_steps} or omitted, got {int(requested_steps)}" - ) + base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) video_shift = float(extra.get("flow_shift", self.default_video_shift)) audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) quality_plan = self._quality_policy.resolve( diff --git a/vllm_omni/diffusion/sched/sigma_schedule.py b/vllm_omni/diffusion/sched/sigma_schedule.py index bea6164d3..7eb3ed38c 100644 --- a/vllm_omni/diffusion/sched/sigma_schedule.py +++ b/vllm_omni/diffusion/sched/sigma_schedule.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 from __future__ import annotations @@ -37,8 +37,11 @@ class DMD2SigmaSchedule: # have to be rejected before the monotonicity check rather than after it. if any(not math.isfinite(value) for value in values): raise ValueError("DMD2 base_schedule entries must be finite") - if values[0] != 1.0 or values[-1] != 0.0: - raise ValueError("DMD2 base_schedule must start at 1.0 and end at 0.0") + # A student whose training capped the noise level (``max_timestep_ratio``) + # starts just below 1.0 rather than at it, so the opening position is + # bounded rather than pinned; the terminal one is always clean latents. + if not 0.0 < values[0] <= 1.0 or values[-1] != 0.0: + raise ValueError("DMD2 base_schedule must start in (0.0, 1.0] and end at 0.0") if any(curr <= nxt for curr, nxt in zip(values, values[1:], strict=False)): raise ValueError("DMD2 base_schedule must be strictly decreasing") diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 6499beee5..28a263e1f 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -538,7 +538,16 @@ class DiffusionWorker: def init_lora_manager(self) -> None: """Initialize the LoRA manager for this worker.""" - if self.model_runner.pipeline is None: + pipeline = self.model_runner.pipeline + if pipeline is None: + return + + # A release whose weights cannot be expressed as switchable LoRA layers + # is fused into the checkpoint while the pipeline loads. There is then + # no adapter left to register, and handing the same path to the manager + # would only fail on a format it does not accept. + if getattr(pipeline, "lora_is_fused", False): + logger.info("LoRA was fused into the checkpoint at load time; skipping the dynamic LoRA manager.") return lora_path = self.od_config.lora_path @@ -739,11 +748,15 @@ class DiffusionWorker: logger.warning("LoRA activation skipped: %s", exc) def remove_lora(self, adapter_id: int) -> bool: + if self.lora_manager is None: + return False return self.lora_manager.remove_adapter(adapter_id) def add_lora(self, lora_request: LoRARequest) -> bool: # NOTE (Alex): We have not implemented the API routing # for the frontend server yet. + if self.lora_manager is None: + return False return self.lora_manager.add_adapter(lora_request) def submit_interaction( @@ -756,9 +769,13 @@ class DiffusionWorker: self.model_runner.submit_interaction(request_id, interaction) def list_loras(self) -> list[int]: + if self.lora_manager is None: + return [] return self.lora_manager.list_adapters() def pin_lora(self, adapter_id: int) -> bool: + if self.lora_manager is None: + return False return self.lora_manager.pin_adapter(adapter_id) def sleep(self, level: int = 1) -> int: diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index b5b64712c..34bc23384 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -396,7 +396,7 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray: # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no # numpy dtype, so ``.numpy()`` below raises on it. video_tensor = video_tensor.float().clamp(-1, 1) * 0.5 + 0.5 - else: + elif video_tensor.dtype != torch.uint8: video_tensor = video_tensor.to(torch.float32) / 255.0 video_array = video_tensor.numpy() return _normalize_single_video_array(video_array) @@ -419,7 +419,7 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray: if np.issubdtype(video_array.dtype, np.floating): if video_array.size and (video_array.min() < 0.0 or video_array.max() > 1.0): video_array = np.clip(video_array, -1.0, 1.0) * 0.5 + 0.5 - elif np.issubdtype(video_array.dtype, np.integer): + elif video_array.dtype != np.uint8 and np.issubdtype(video_array.dtype, np.integer): video_array = video_array.astype(np.float32) / 255.0 return video_array @@ -435,7 +435,7 @@ def _normalize_video_array(video_array: np.ndarray) -> list[np.ndarray] | np.nda def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: - """Normalize a list of frames into numpy arrays with values in [0,1].""" + """Normalize a list of frames into numpy arrays, uint8 ones unchanged.""" normalized: list[np.ndarray] = [] for frame in frames: if isinstance(frame, torch.Tensor): @@ -453,7 +453,7 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: if np.issubdtype(frame_array.dtype, np.floating): if frame_array.size and (frame_array.min() < 0.0 or frame_array.max() > 1.0): frame_array = np.clip(frame_array, -1.0, 1.0) * 0.5 + 0.5 - elif np.issubdtype(frame_array.dtype, np.integer): + elif frame_array.dtype != np.uint8 and np.issubdtype(frame_array.dtype, np.integer): frame_array = frame_array.astype(np.float32) / 255.0 normalized.append(frame_array) @@ -461,7 +461,14 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: def _coerce_video_to_frames(video: Any) -> list[np.ndarray]: - """Convert a video payload into a list of normalized float32 frames.""" + """Convert a video payload into a list of normalized frames. + + Frames are float32 in [0, 1], except uint8 ones, which pass through: the + muxer's own dtype is uint8, so normalising here would only pay for a + full-size conversion each way. The direct planar path additionally needs + contiguous channel planes and falls back for interleaved RGB whatever the + dtype; the standard muxer takes the uint8 frames as they are. + """ if isinstance(video, torch.Tensor): video_array = _normalize_video_tensor(video) return list(video_array)