From 40f8e94400fca12ddcefbb46b9043f6a6d6dba96 Mon Sep 17 00:00:00 2001 From: heyuanliu-intel <90166842+heyuanliu-intel@users.noreply.github.com> Date: Tue, 15 Sep 2026 01:02:30 +0800 Subject: [PATCH] [Model][MiniMax-H3] Multi-GPU serving fixes and layer-wise offload component selection (#7047) Signed-off-by: heyuanliu-intel --- docs/design/feature/sequence_parallel.md | 6 +- .../diffusion/regional_compilation.md | 29 ++ recipes/MiniMaxAI/MiniMax-H3.md | 5 +- .../diffusion/attention/test_ring_padding.py | 221 ++++++++ .../minimax_h3/test_minimax_h3_contract.py | 306 ++++++++++- .../minimax_h3/test_minimax_h3_offload.py | 244 ++++++++- .../minimax_h3/test_minimax_h3_packing.py | 89 +++- .../test_minimax_h3_vae_gather_reuse.py | 473 ++++++++++++++++++ .../minimax_h3/test_minimax_h3_vae_ops.py | 2 +- .../openai_api/test_video_server.py | 26 + tests/model_extras/test_model_extras.py | 10 + .../diffusion/attention/backends/abstract.py | 5 +- .../attention/backends/ring/ring_utils.py | 29 +- .../attention/backends/ring_flash_attn.py | 80 ++- .../attention/backends/ring_pytorch_attn.py | 55 +- .../diffusion/attention/parallel/ring.py | 4 + .../minimax_h3/minimax_h3_transformer.py | 25 +- .../models/minimax_h3/packed_sequence.py | 48 +- .../models/minimax_h3/pipeline_minimax_h3.py | 81 ++- vllm_omni/diffusion/models/minimax_h3/vae.py | 160 ++++++ vllm_omni/model_extras/registry.py | 19 + 21 files changed, 1857 insertions(+), 60 deletions(-) create mode 100644 tests/diffusion/attention/test_ring_padding.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py diff --git a/docs/design/feature/sequence_parallel.md b/docs/design/feature/sequence_parallel.md index d0328bcf6..7ae263b7d 100644 --- a/docs/design/feature/sequence_parallel.md +++ b/docs/design/feature/sequence_parallel.md @@ -77,7 +77,9 @@ Use it when plain Ulysses-SP would otherwise fail because: - `auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`. - `advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes. -- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`. +- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume arbitrary + `attention_mask` values. A model with contiguous suffix padding may instead publish `valid_kv_length`; Ring trims + each circulated K/V block against that global valid prefix. - `advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement. --- @@ -416,7 +418,7 @@ _sp_plan = { | Constraint | Description | |------------|-------------| | **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. | -| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. | +| **Ring Attention Limitation** | Ring attention does not support arbitrary `attention_mask` values, so `auto_pad=True` remains incompatible. Contiguous suffix padding is supported when the model publishes the global `valid_kv_length` prefix. | 1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`: ```python diff --git a/docs/user_guide/diffusion/regional_compilation.md b/docs/user_guide/diffusion/regional_compilation.md index 448f19518..364900ff5 100644 --- a/docs/user_guide/diffusion/regional_compilation.md +++ b/docs/user_guide/diffusion/regional_compilation.md @@ -34,5 +34,34 @@ that provide their own `setup_compile()` implementation manage their compilation policy independently. Compilation is lazy, so backend or graph errors can first surface on the initial request. +## Pinning one packed shape (MiniMax-H3) + +Compiled regions are keyed by input shape. MiniMax-H3 packs a request into a +row count that follows the request -- a longer prompt, more frames, a larger +frame -- and rounds it up to the next 64-row boundary, so two requests can land +on two different packed lengths. + +Under the default `diffusion_compile_dynamic: true` a new packed length is +absorbed by the dynamic shape, so it is not by itself a recompilation. It still +hands the backend a shape it has not autotuned, which is the risk this knob +exists to remove: a fixed-shape deployment (`--no-diffusion-compile-dynamic`) +recompiles the transformer blocks per shape, and a first-seen large shape can +push autotuning into its worst-case memory use. + +A request pins the length with `extra_args["pad_seq_len"]`: + +```python +sampling_params = SamplingParams(extra_args={"pad_seq_len": 54080}) +``` + +The value must be a positive multiple of 64 and must cover the rows the request +actually uses; the packed sequence is then padded to it instead of to the next +64-row boundary. Pick a bucket that covers every request dimension the +deployment accepts -- prompt length, frame count, frame size and reference +blocks all change the used row count. The server logs the effective length as +`MiniMax H3 packed sequence: ... pad_seq_len=... used=... seq_len=...` whenever +a request pins it. The padding rows are masked out, so their cost is the +attention and feed-forward work on those rows. + Use `--enforce-eager` to disable the model runner's generic compile setup. Pipelines that compile internally define their own eager-mode behavior. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 812ed6f30..c966b6b2e 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -1240,8 +1240,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ - `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch). - VAE patch parallelism requires size 1 or the full DiT group size and supports the H3 native `tile` mode only. -- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use - pure Ulysses. +- Ulysses x Ring hybrid attention supports H3's single-request contiguous suffix + padding. Arbitrary attention masks and multi-request packed batches remain + unsupported on the Ring path. - Online FP8 with DLO AllGather temporarily materializes the complete FP8 model in host memory on every rank during startup before retaining only each rank's shard. Size startup host memory for that transient peak. diff --git a/tests/diffusion/attention/test_ring_padding.py b/tests/diffusion/attention/test_ring_padding.py new file mode 100644 index 000000000..8994e3fc1 --- /dev/null +++ b/tests/diffusion/attention/test_ring_padding.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from dataclasses import dataclass +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends import ring_flash_attn, ring_pytorch_attn +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType +from vllm_omni.diffusion.attention.backends.ring.ring_utils import ring_kv_block_valid_length +from vllm_omni.diffusion.attention.parallel import ring as ring_parallel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _SingleRankComm: + rank = 0 + world_size = 1 + + +@dataclass(frozen=True) +class _FakeSequenceParallelGroup: + ring_group: object + + +def _fake_attention_result(q, _k, _v, **_kwargs): + return torch.zeros_like(q), torch.zeros(q.shape[0], q.shape[2], q.shape[1]) + + +@pytest.mark.parametrize(("block_rank", "expected"), [(0, 4), (1, 1), (2, 0)]) +def test_global_prefix_is_partitioned_across_ring_blocks(block_rank, expected): + assert ring_kv_block_valid_length(5, 4, block_rank, 3) == expected + + +def test_absent_prefix_length_keeps_the_full_ring_block(): + assert ring_kv_block_valid_length(None, 4, 2, 3) == 4 + + +def test_pytorch_ring_trims_padding_from_each_kv_block(monkeypatch): + captured: dict[str, Any] = {} + + def fake_attention(q, k, v, **kwargs): + captured.update(k=k, v=v, **kwargs) + return _fake_attention_result(q, k, v, **kwargs) + + monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention) + qkv = torch.randn(1, 4, 2, 8) + + output = ring_pytorch_attn.ring_pytorch_attn_func( + qkv, + qkv, + qkv, + group=object(), + valid_kv_length=3, + ) + + assert output.shape == qkv.shape + assert captured["k"].shape[1] == 3 + assert captured["v"].shape[1] == 3 + + +def test_pytorch_ring_trims_circulated_blocks_by_source_rank(monkeypatch): + observed_lengths = [] + # Each ring step sends K and V separately. + blocks = [torch.randn(1, 4, 2, 8) for _ in range(4)] + + class FakeRingComm: + rank = 2 + world_size = 3 + + def __init__(self, _group): + self._next_block = 0 + + def send_recv(self, _tensor): + block = blocks[self._next_block] + self._next_block += 1 + return block + + def commit(self): + pass + + def wait(self): + pass + + def fake_attention(q, k, v, **kwargs): + observed_lengths.append(k.shape[1]) + return _fake_attention_result(q, k, v, **kwargs) + + monkeypatch.setattr(ring_pytorch_attn, "RingComm", FakeRingComm) + monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention) + qkv = torch.randn(1, 4, 2, 8) + + ring_pytorch_attn.ring_pytorch_attn_func( + qkv, + qkv, + qkv, + group=object(), + valid_kv_length=5, + ) + + # Rank 2 starts with global block 2, then receives blocks 1 and 0. + # The first block is empty and therefore does not launch attention. + assert observed_lengths == [1, 4] + + +def test_flash_ring_trims_padding_from_each_kv_block(monkeypatch): + captured: dict[str, Any] = {} + + def fake_attention(q, k, v, **kwargs): + captured.update(k=k, v=v, **kwargs) + return _fake_attention_result(q, k, v, **kwargs) + + monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr(ring_flash_attn, "select_flash_attn_impl", lambda *_args, **_kwargs: fake_attention) + qkv = torch.randn(1, 4, 2, 8) + + output, _ = ring_flash_attn.ring_flash_attn_forward( + object(), + qkv, + qkv, + qkv, + softmax_scale=0.25, + causal=False, + attn_type=AttnType.FA, + valid_kv_length=3, + ) + + assert output.shape == qkv.shape + assert captured["k"].shape[1] == 3 + assert captured["v"].shape[1] == 3 + + +def test_ring_dispatch_forwards_global_valid_kv_length(monkeypatch): + captured: dict[str, Any] = {} + + def fake_ring_attention(*_args, **kwargs): + captured.update(kwargs) + return "output" + + monkeypatch.setattr(ring_pytorch_attn, "ring_pytorch_attn_func", fake_ring_attention) + strategy = ring_parallel.RingParallelAttention( + _FakeSequenceParallelGroup(ring_group=object()), + attn_backend_pref="sdpa", + ) + metadata = AttentionMetadata(extra={"valid_kv_length": 3}) + query = torch.randn(1, 4, 2, 8) + + assert strategy.run_attention(query, query, query, metadata) == "output" + assert captured["valid_kv_length"] == 3 + + +def test_pytorch_ring_rejects_causal_with_a_valid_prefix(monkeypatch): + """Trimming a K/V block under causal=True would silently shift the diagonal. + + The query block keeps the full padded length while the circulated K/V block + is trimmed, and a causal mask over a rectangular tile is bottom-right + aligned, so the combination must be refused rather than computed. + """ + monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", _fake_attention_result) + qkv = torch.randn(1, 4, 2, 8) + + with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"): + ring_pytorch_attn.ring_pytorch_attn_func( + qkv, + qkv, + qkv, + causal=True, + group=object(), + valid_kv_length=3, + ) + + +def test_flash_ring_rejects_causal_with_a_valid_prefix(monkeypatch): + """Same refusal on the Flash-Attention ring path.""" + monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr( + ring_flash_attn, + "select_flash_attn_impl", + lambda *_args, **_kwargs: _fake_attention_result, + ) + qkv = torch.randn(1, 4, 2, 8) + + with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"): + ring_flash_attn.ring_flash_attn_forward( + object(), + qkv, + qkv, + qkv, + softmax_scale=0.25, + causal=True, + attn_type=AttnType.FA, + valid_kv_length=3, + ) + + +def test_ring_paths_still_accept_causal_without_a_valid_prefix(monkeypatch): + """The guard is scoped to the padded case: plain causal Ring is unchanged.""" + monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm()) + monkeypatch.setattr( + ring_flash_attn, + "select_flash_attn_impl", + lambda *_args, **_kwargs: _fake_attention_result, + ) + qkv = torch.randn(1, 4, 2, 8) + + output, _ = ring_flash_attn.ring_flash_attn_forward( + object(), + qkv, + qkv, + qkv, + softmax_scale=0.25, + causal=True, + attn_type=AttnType.FA, + ) + + assert output.shape == qkv.shape 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 7868c255b..67498abc5 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -459,6 +459,7 @@ def test_combined_task_inference_and_transformer_routing(): ("fl2va", "fl2va", 1, "FL2VA", ["FL2VA"], {"t2va", "fl2va"}, False), ], ) +@pytest.mark.parametrize("trust_remote_code", [False, True]) def test_pipeline_loads_task_selected_components_with_encoder_ownership( monkeypatch, tmp_path, @@ -469,6 +470,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( source_partitions, expected_tasks, load_text_encoder, + trust_remote_code, ): from vllm_omni.diffusion.data import ( DiffusionParallelConfig, @@ -500,6 +502,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( (tmp_path / partition_name / component).mkdir() created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} + created_kwargs: dict[str, list[Any]] = {"text_encoder": [], "video_vae": [], "audio_vae": []} class FakeModule(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -513,6 +516,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( def component_factory(name): def create(path, *args, **kwargs): created[name].append(str(path)) + created_kwargs[name].append(kwargs) return FakeModule() return create @@ -589,6 +593,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( cfg_parallel_size=1, text_encoder_tp_size=1, ), + trust_remote_code=trust_remote_code, ) pipeline = pipeline_module.MiniMaxH3Pipeline(od_config=od_config) @@ -598,6 +603,13 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership( component_path = tmp_path / component_partition assert created["video_vae"] == [str(component_path / "video_vae")] assert created["audio_vae"] == [str(component_path / "audio_vae")] + # The checkpoint VAEs execute the modeling code shipped with the weights, so + # the pipeline must hand them the engine's trust_remote_code decision. Both + # VAEs are created for every parametrization, so this stays unconditional. + for component in ("video_vae", "audio_vae"): + assert [kwargs["trust_remote_code"] for kwargs in created_kwargs[component]] == [trust_remote_code] + # The tokenizer/processor "loaded once" assertions are asserted per + # load_text_encoder branch below, matching upstream's own expectations. expected_dit_modules = ["transformer", "transformers_ref"] if expected_dits == 2 else ["transformer"] assert pipeline._dit_modules == expected_dit_modules assert hasattr(pipeline, "transformers_ref") is (expected_dits == 2) @@ -1122,6 +1134,64 @@ def test_packed_attention_keeps_padding_mask_for_other_backends(): ) +def test_ring_packed_attention_publishes_prefix_length_without_global_mask(): + attention = _fake_packed_attention("CUDNN_ATTN") + attention.attention.use_ring = True + q = torch.randn(4, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + metadata = attention.attention.metadata + assert metadata.attn_mask is None + assert metadata.extra["valid_kv_length"] == 5 + + +def test_ring_packed_attention_keeps_mask_when_runtime_sp_is_inactive(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer + + class InactiveContext: + sp_active = False + + attention = _fake_packed_attention("CUDNN_ATTN") + attention.attention.use_ring = True + monkeypatch.setattr( + minimax_h3_transformer, + "is_forward_context_available", + lambda: True, + raising=False, + ) + monkeypatch.setattr( + minimax_h3_transformer, + "get_forward_context", + InactiveContext, + raising=False, + ) + q = torch.randn(8, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + metadata = attention.attention.metadata + assert torch.equal( + metadata.attn_mask, + torch.tensor([[True, True, True, True, True, False, False, False]]), + ) + assert metadata.extra["npu_attn_varlen"] is True + + def _fake_packed_attention( backend_name: str, *, @@ -1143,6 +1213,10 @@ def _fake_packed_attention( def supports_multi_doc_packed_varlen(cls) -> bool: return supports_multi_doc + @classmethod + def supports_packed_mask_free(cls) -> bool: + return False + class FakeAttention(torch.nn.Module): attn_backend = FakeBackend @@ -1712,7 +1786,7 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, *, trust_remote_code: FakeRemote(), ) video_vae = vae_module.MiniMaxH3VideoVAE( @@ -1834,7 +1908,7 @@ def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, *, trust_remote_code: FakeRemote(), ) video_vae = vae_module.MiniMaxH3VideoVAE( @@ -1869,7 +1943,7 @@ def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, *, trust_remote_code: FakeRemote(), ) monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls) @@ -1890,6 +1964,123 @@ def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch): stager.offload.assert_called_once_with() +def test_remote_component_refuses_untrusted_checkpoint_code(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + loaded = [] + + def spy_get_class(class_reference, component_path): + loaded.append((class_reference, component_path)) + raise AssertionError("checkpoint code must not be resolved when untrusted") + + monkeypatch.setattr(vae_module, "get_class_from_dynamic_module", spy_get_class) + + with pytest.raises(ValueError, match="--trust-remote-code"): + vae_module._load_remote_component( + "/checkpoints/video_vae", + {"auto_map": {"AutoModel": "modeling_vae.VideoVAE"}}, + trust_remote_code=False, + ) + + assert loaded == [] + + +def test_remote_component_resolves_checkpoint_code_when_trusted(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeComponent(torch.nn.Module): + @classmethod + def from_pretrained(cls, component_path): + module = cls() + module.loaded_from = component_path + return module + + loaded = [] + + def spy_get_class(class_reference, component_path): + loaded.append((class_reference, component_path)) + return FakeComponent + + monkeypatch.setattr(vae_module, "get_class_from_dynamic_module", spy_get_class) + + module = vae_module._load_remote_component( + "/checkpoints/video_vae", + {"auto_map": {"AutoModel": "modeling_vae.VideoVAE"}}, + trust_remote_code=True, + ) + + # Negative control: the trusted path stays byte-for-byte the pre-gate call. + assert loaded == [("modeling_vae.VideoVAE", "/checkpoints/video_vae")] + assert module.loaded_from == "/checkpoints/video_vae" + + +@pytest.mark.parametrize("vae_cls_name", ["MiniMaxH3VideoVAE", "MiniMaxH3AudioVAE"]) +@pytest.mark.parametrize("trusted", [True, False]) +def test_vae_forwards_trust_remote_code_to_loader(monkeypatch, vae_cls_name, trusted): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1) + + forwarded = [] + + def spy_load(_component_path, _config, *, trust_remote_code): + forwarded.append(trust_remote_code) + return FakeRemote() + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + "sample_rate": 32000, + }, + ) + monkeypatch.setattr(vae_module, "_load_remote_component", spy_load) + + vae_cls = getattr(vae_module, vae_cls_name) + vae_cls( + "unused", + device=torch.device("cpu"), + trust_remote_code=trusted, + ) + + assert forwarded == [trusted] + + +@pytest.mark.parametrize("vae_cls_name", ["MiniMaxH3VideoVAE", "MiniMaxH3AudioVAE"]) +def test_vae_defaults_to_refusing_checkpoint_code(monkeypatch, vae_cls_name): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + forwarded = [] + + def spy_load(_component_path, _config, *, trust_remote_code): + forwarded.append(trust_remote_code) + raise AssertionError("unreachable") + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + "sample_rate": 32000, + }, + ) + monkeypatch.setattr(vae_module, "_load_remote_component", spy_load) + + vae_cls = getattr(vae_module, vae_cls_name) + with pytest.raises(AssertionError, match="unreachable"): + vae_cls("unused", device=torch.device("cpu")) + + assert forwarded == [False] + + def test_video_vae_encode_uses_configured_parallel_tiling(): from vllm_omni.diffusion.models.minimax_h3.vae import ( MiniMaxH3VideoVAE, @@ -2889,3 +3080,112 @@ def test_encoder_load_places_every_source_in_its_own_rows(tmp_path): } for source, rows in slices.items(): assert torch.equal(rows, torch.full_like(rows, _SOURCE_FILL[source])), source + + +def _pad_seq_len_batch(pad_seq_len): + """A t2va request that pins the packed length through ``extra_args``.""" + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + extra_args = {"task": "t2va", "aspect_ratio": "16:9"} + if pad_seq_len is not None: + extra_args["pad_seq_len"] = pad_seq_len + sampling = OmniDiffusionSamplingParams( + quality="lossless", + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=None, + extra_args=extra_args, + ) + return DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="pinned packed length", + sampling_params=sampling, + request_id="pad-seq-len", + ) + ] + ) + + +@pytest.mark.parametrize("pinned, expected", [(None, None), (54080, 54080)]) +def test_pad_seq_len_reaches_diffuse_from_extra_args(pinned, expected): + """The request field must arrive at ``diffuse``, not stop at validation.""" + diffuse_calls: list[dict[str, Any]] = [] + pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": None, "ref2va": None}) + + pipeline.forward(_pad_seq_len_batch(pinned)) + + assert [call["pad_seq_len"] for call in diffuse_calls] == [expected] + + +@pytest.mark.parametrize("task", ["t2va", "ref2va"]) +def test_pad_seq_len_reaches_the_packer_for_every_task_family(task, monkeypatch): + """``diffuse`` must hand the pin to whichever packer the task uses.""" + import vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 as module + + class _CapturedError(Exception): + pass + + recorded = {} + + def _capture(**kwargs): + recorded.update(kwargs) + raise _CapturedError + + monkeypatch.setattr(module, "minimax_h3_packed_sequence", _capture) + monkeypatch.setattr(module, "minimax_h3_packed_sequence_ref2va_blocks", _capture) + + pipeline = _distilled_pipeline([], {"fl2va": None, "ref2va": None}) + pipeline._initial_noise = lambda **kwargs: (torch.zeros(1), torch.zeros(1)) + + with pytest.raises(_CapturedError): + pipeline._build_denoise_inputs( + task=task, + text_embeddings=torch.ones(7, 2), + text_tags=torch.ones(7, dtype=torch.long), + seed=0, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + num_frames=5, + num_steps=2, + video_shift=12.0, + audio_shift=3.0, + base_schedule=None, + visual_condition=None, + visual_condition_shape=(1, 4, 6), + audio_condition=None, + ref_audio_t=1, + pad_seq_len=192, + ) + + assert recorded["seq_len"] == 192 + + +@pytest.mark.parametrize("value", [0, -64, 100, "128", True, 1.5, (1 << 20) + 64]) +def test_pad_seq_len_request_validation_rejects_bad_values(value): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _resolve_pad_seq_len, + ) + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError): + _resolve_pad_seq_len(value) + + +def test_pad_seq_len_request_validation_accepts_aligned_values(): + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + MINIMAX_H3_MAX_PAD_SEQ_LEN, + ) + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _resolve_pad_seq_len, + ) + + assert _resolve_pad_seq_len(None) is None + assert _resolve_pad_seq_len(54080) == 54080 + assert _resolve_pad_seq_len(MINIMAX_H3_MAX_PAD_SEQ_LEN) == MINIMAX_H3_MAX_PAD_SEQ_LEN diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index 865b2248a..31a8916a6 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from contextlib import contextmanager +import weakref +from contextlib import contextmanager, nullcontext from types import SimpleNamespace from typing import Any from unittest.mock import Mock @@ -76,6 +77,247 @@ def test_requested_model_offload_without_backend_uses_resident_components(): pipeline.text_encoder.encode_ids.assert_called_once() +def test_decoded_output_is_unchanged_without_active_model_offload(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [] + output = Mock() + + assert pipeline._offload_model_cpu_stage_output(output) is output + output.cpu.assert_not_called() + + +def test_reply_rank_copies_decoded_output_to_host(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [Mock()] + monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: True) + output = Mock() + host_output = Mock() + output.cpu.return_value = host_output + + assert pipeline._offload_model_cpu_stage_output(output) is host_output + output.cpu.assert_called_once_with() + + +def test_non_reply_rank_drops_decoded_output_storage(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._model_cpu_offload_modules = [Mock()] + monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: False) + output = torch.ones((2, 3), dtype=torch.float16) + + placeholder = pipeline._offload_model_cpu_stage_output(output) + + assert placeholder.device.type == "meta" + assert placeholder.shape == output.shape + assert placeholder.dtype == output.dtype + + +def test_decode_releases_only_audio_before_return(monkeypatch): + """Video is released by the callers of ``decode`` -- after quantization.""" + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + video = torch.ones((1, 1, 4, 4)) + audio = torch.ones((1, 8)) + pipeline = SimpleNamespace( + device=torch.device("cpu"), + video_vae=Mock(decode_latent=Mock(return_value=video)), + audio_vae=Mock(decode_latent=Mock(return_value=audio)), + _component_on_device=lambda component: nullcontext(), + _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value), + ) + monkeypatch.setattr( + module.current_omni_platform, + "create_autocast_context", + lambda **kwargs: nullcontext(), + ) + + actual_video, actual_audio = MiniMaxH3Pipeline.decode( + pipeline, + torch.zeros(1), + torch.zeros(1), + height=2, + width=3, + ) + + assert actual_video.shape[-2:] == (2, 3) + assert actual_audio is audio + calls = pipeline._offload_model_cpu_stage_output.call_args_list + assert len(calls) == 1 + assert calls[0].args[0] is audio + + +def _recording_release(events, releases): + """Stand-in for ``_offload_model_cpu_stage_output`` that records what it got.""" + + def release(value): + events.append("release") + releases.append(value) + return value + + return release + + +def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + # MiniMaxH3VideoVAE.decode_latent ends in frames.float(), so this is what + # the release hook would see if it ran before the quantizer. + decoded_dtype = torch.float32 + events: list[str] = [] + releases: list[torch.Tensor] = [] + + def diffuse(**kwargs): + events.append(f"diffuse:{kwargs['seed']}") + return torch.zeros(1), torch.zeros(1) + + def decode(video_latent, audio_latent, *, height, width): + events.append("decode") + return ( + torch.zeros((1, 3, 2, height, width), dtype=decoded_dtype), + torch.zeros((1, 8)), + ) + + # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the + # tensor-decode branch, which is the one the release hook wraps. + context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False} + pipeline = SimpleNamespace( + od_config=SimpleNamespace(), + _extract_prompt=lambda raw: ("a prompt", {}), + _extract_text_conditioning=lambda raw: None, + _extract_prepared_reference_videos=lambda raw: None, + _prepare_request_inputs=lambda **kwargs: context, + _denoise_kwargs=lambda ctx: {}, + diffuse=diffuse, + decode=decode, + _offload_model_cpu_stage_output=_recording_release(events, releases), + ) + request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace()) + + output = MiniMaxH3Pipeline.forward(pipeline, request) + + # The release must see the quantized frames, not the decoded floats. + assert len(releases) == 2 + assert [tensor.dtype for tensor in releases] == [torch.uint8, torch.uint8] + assert decoded_dtype is not torch.uint8 + # ... and it must still happen before the next seed reloads the DiT. + assert events == ["diffuse:41", "decode", "release", "diffuse:42", "decode", "release"] + assert output.output[0].dtype is torch.uint8 + + +def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypatch): + """Releasing the copy is not enough if a local still pins the original. + + ``_offload_model_cpu_stage_output`` hands back a *new* tensor, so the + decoded frames stay alive for as long as anything still references them. + The whole point of this branch is that they are gone before the next seed + reloads the DiT, and ``diffuse`` is where that reload happens. + """ + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + decoded_refs: list[weakref.ref] = [] + alive_at_diffuse: list[int] = [] + + def diffuse(**kwargs): + alive_at_diffuse.append(sum(1 for ref in decoded_refs if ref() is not None)) + return torch.zeros(1), torch.zeros(1) + + def decode(video_latent, audio_latent, *, height, width): + video = torch.zeros((1, 3, 2, height, width), dtype=torch.float32) + decoded_refs.append(weakref.ref(video)) + return video, torch.zeros((1, 8)) + + # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the + # tensor-decode branch, which is the one the release hook wraps. + context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False} + pipeline = SimpleNamespace( + od_config=SimpleNamespace(), + _extract_prompt=lambda raw: ("a prompt", {}), + _extract_text_conditioning=lambda raw: None, + _extract_prepared_reference_videos=lambda raw: None, + _prepare_request_inputs=lambda **kwargs: context, + _denoise_kwargs=lambda ctx: {}, + diffuse=diffuse, + decode=decode, + _offload_model_cpu_stage_output=lambda value: value.clone(), + ) + request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace()) + + MiniMaxH3Pipeline.forward(pipeline, request) + + assert len(decoded_refs) == 2 + # Seed 1's diffuse sees nothing decoded yet; seed 2's must not see seed 1's + # decoded frames still resident. + assert alive_at_diffuse == [0, 0] + + +def test_post_decode_releases_video_after_quantization(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + events: list[str] = [] + releases: list[torch.Tensor] = [] + + def decode(video_latent, audio_latent, *, height, width): + events.append("decode") + return ( + torch.zeros((1, 3, 2, height, width), dtype=torch.float16), + torch.zeros((1, 8)), + ) + + pipeline = SimpleNamespace( + od_config=SimpleNamespace(), + _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)), + decode=decode, + _offload_model_cpu_stage_output=_recording_release(events, releases), + ) + state = SimpleNamespace( + latents=torch.zeros(1), + extra={ + module._STEP_BRANCH: object(), + module._STEP_AUDIO_ROWS: torch.zeros(1), + module._STEP_SHAPE: { + "latent_t": 1, + "latent_h": 1, + "latent_w": 1, + "audio_t": 1, + "height": 4, + "width": 4, + }, + }, + ) + + output = MiniMaxH3Pipeline.post_decode(pipeline, state) + + assert events == ["decode", "release"] + assert len(releases) == 1 + assert releases[0].dtype is torch.uint8 + assert output.output[0] is releases[0] + + +@pytest.mark.parametrize( + ("available", "initialized", "rank", "expected"), + [(False, False, 7, True), (True, False, 7, True), (True, True, 0, True), (True, True, 1, False)], +) +def test_output_owner_rank_matches_executor_contract(monkeypatch, available, initialized, rank, expected): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + monkeypatch.setattr(module.dist, "is_available", lambda: available) + monkeypatch.setattr(module.dist, "is_initialized", lambda: initialized) + monkeypatch.setattr(module.dist, "get_rank", lambda: rank) + + assert MiniMaxH3Pipeline._is_output_owner_rank() is expected + + def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py index 53348f579..4340e093f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.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 import torch @@ -190,3 +190,90 @@ def test_condition_noise_accepts_a_reference_video_longer_than_the_target(): noise_aug=0.5, ) assert result.shape == rows.shape + + +def test_explicit_seq_len_pins_one_shape_across_prompt_lengths(): + """Prompts of different token counts must land on the same packed length. + + Without the pin they fall into different 64-row buckets, and each bucket is + a packed shape the compiled transformer has not seen before. + """ + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + ) + + common = dict( + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + include_keyframe_cond=False, + ) + # used = text_len + 18 rows: 22 -> bucket 64; 118 -> bucket 128. + short = minimax_h3_packed_sequence(text_len=4, **common) + long = minimax_h3_packed_sequence(text_len=100, **common) + assert int(short["seq_len"]) == 64 + assert int(long["seq_len"]) == 128 + + short_pinned = minimax_h3_packed_sequence(text_len=4, seq_len=192, **common) + long_pinned = minimax_h3_packed_sequence(text_len=100, seq_len=192, **common) + assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192 + # The pin only adds padding rows; the used prefix is untouched. + assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1]) == 22 + assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1]) == 118 + assert short_pinned["img_pos"].tolist() == short["img_pos"].tolist() + + +def test_the_ref2va_packer_pins_one_shape_across_prompt_lengths_too(): + """Both packers share one resolver, so both honour the pin.""" + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence_ref2va_blocks, + ) + + common = dict( + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ref_blocks=[{"kind": "audio", "ref_audio_t": 1}], + ) + short = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, **common) + long = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, **common) + assert int(short["seq_len"]) != int(long["seq_len"]) + + short_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, seq_len=192, **common) + long_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, seq_len=192, **common) + assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192 + assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1]) + assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1]) + + +def test_explicit_seq_len_below_used_rows_is_a_client_error(): + """Both packers reject an under-sized pin as a 4xx, not a server error.""" + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + minimax_h3_packed_sequence_ref2va_blocks, + ) + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match="used rows"): + minimax_h3_packed_sequence( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + include_keyframe_cond=False, + seq_len=8, + ) + + with pytest.raises(OmniClientError, match="used rows"): + minimax_h3_packed_sequence_ref2va_blocks( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ref_blocks=[{"kind": "audio", "ref_audio_t": 1}], + seq_len=8, + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py new file mode 100644 index 000000000..116f44ef9 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_gather_reuse.py @@ -0,0 +1,473 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU tests for the pinned tiled-VAE gather buffer in the MiniMax-H3 video VAE. + +The checkpoint's ``_all_gather_tiled_results`` allocates its gather output on +every call, so with a per-request ``empty_cache()`` the collective lands on a +new device address every request and XCCL's per-address registrations pile up. +On XPU the adapter replaces that method with an equal-shape gather into a +buffer it keeps alive. These tests pin the properties that make it a fix +rather than a rewrite: the address is stable, the tiles come back +bit-identical and in the same order, the caller owns what it gets back, and no +other device is touched. +""" + +import logging +import sys +import types +from typing import Any + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module +from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +TILE = (1, 3, 2, 4, 4) + + +class _CheckpointGatherSentinel: + """Stands in for the checkpoint's own bound method. + + Identity is the point: a test asserts this exact object is still installed + on the devices the adapter must leave alone. + """ + + def __call__(self, tasks, num_tiles): # pragma: no cover - never invoked + raise AssertionError("the checkpoint gather must not run in these tests") + + +class _FakeCheckpointModel: + def __init__(self): + self._all_gather_tiled_results = _CheckpointGatherSentinel() + + +def _local_tiles(all_tiles, num_tiles, rank, sp_size): + """One rank's round-robin share, as ``klvae.py:275`` computes it.""" + return [all_tiles[index] for index in range(rank, num_tiles, sp_size)] + + +def _reference_gather(all_tiles, num_tiles, sp_size): + """The checkpoint's variable-shape list gather, reimplemented verbatim. + + Written from the control flow of ``_all_gather_tiled_results`` in the + checkpoint's ``klvae.py`` (lines 252-273) -- stack each rank's share, gather + the list, walk it back into global tile order -- so that it is an + independent oracle rather than a restatement of the code under test. + """ + gathered = [torch.stack(_local_tiles(all_tiles, num_tiles, rank, sp_size), dim=0) for rank in range(sp_size)] + results = [None] * num_tiles + for rank, rank_tensors in enumerate(gathered): + for k in range(rank_tensors.shape[0]): + global_index = k * sp_size + rank + if global_index >= num_tiles: + break + results[global_index] = rank_tensors[k] + return results + + +def _peer_stacks(all_tiles, num_tiles, sp_size): + """What every rank hands to the collective, padded the way the adapter pads.""" + max_tasks = -(-num_tiles // sp_size) + stacks = [] + for rank in range(sp_size): + share = _local_tiles(all_tiles, num_tiles, rank, sp_size) + stacked = share[0].new_zeros((max_tasks, *share[0].shape)) + stacked[: len(share)] = torch.stack(share, dim=0) + stacks.append(stacked) + return stacks + + +def _tiles(num_tiles, seed=0, dtype=torch.float32): + generator = torch.Generator().manual_seed(seed) + return [torch.rand(TILE, generator=generator).to(dtype) for _ in range(num_tiles)] + + +def _workspace_bytes(num_tiles, sp_size, dtype=torch.float32): + """Bytes one gather needs: ``sp_size`` copies of the padded local stack.""" + max_tasks = -(-num_tiles // sp_size) + element_size = torch.empty(0, dtype=dtype).element_size() + numel = 1 + for dim in TILE: + numel *= dim + return sp_size * max_tasks * numel * element_size + + +@pytest.fixture +def fake_collective(monkeypatch): + """Replace the collective with a single-process stand-in. + + ``all_gather_into_tensor`` is the only distributed call on this path, so + filling the output buffer from a caller-supplied peer list exercises the + real padding, keying, addressing and unpacking code without a process + group. + """ + calls: list[dict[str, Any]] = [] + + def fake_all_gather_into_tensor(output, input_tensor, group=None, async_op=False): + peers = calls[-1]["peers"] + assert output.shape == (len(peers), *input_tensor.shape) + calls[-1]["sent"] = input_tensor.clone() + calls[-1]["buf_ptr"] = output.data_ptr() + for rank, peer in enumerate(peers): + output[rank].copy_(peer) + + monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather_into_tensor) + return calls + + +class _Records(logging.Handler): + """vllm's logger sets ``propagate=False``, so caplog never sees these.""" + + def __init__(self): + super().__init__(level=logging.DEBUG) + self.messages = [] + self.records = [] + + def emit(self, record): + self.messages.append(record.getMessage()) + self.records.append((record.levelno, record.getMessage())) + + +@pytest.fixture +def gather_log(): + handler = _Records() + logger = vae_module.logger + previous = logger.level + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + try: + yield handler + finally: + logger.removeHandler(handler) + logger.setLevel(previous) + + +def _vae(sp_size, sp_rank=0): + """An adapter instance with the override installed, no checkpoint needed.""" + state = { + "sp_size": sp_size, + "sp_rank": sp_rank, + "sp_enabled": sp_size > 1, + "sp_process_group": object(), + } + module = types.ModuleType("fake_gather_ckpt.parallel") + module.get_parallel_state = lambda: state # type: ignore[attr-defined] + sys.modules.setdefault("fake_gather_ckpt", types.ModuleType("fake_gather_ckpt")) + sys.modules["fake_gather_ckpt.parallel"] = module + + vae = object.__new__(MiniMaxH3VideoVAE) + vae.remote = type("Remote", (), {"__module__": "fake_gather_ckpt.klvae"})() + vae.model = _FakeCheckpointModel() + vae._device_target = torch.device("xpu") + vae._tile_gather_workspace = None + vae._tile_gather_stats = {"hits": 0, "allocs": 0, "workspace_bytes": 0} + vae._checkpoint_tile_gather = None + vae._install_persistent_tile_gather() + return vae + + +def _run(vae, calls, all_tiles, num_tiles, sp_size, sp_rank): + calls.append({"peers": _peer_stacks(all_tiles, num_tiles, sp_size)}) + return vae.model._all_gather_tiled_results(_local_tiles(all_tiles, num_tiles, sp_rank, sp_size), num_tiles) + + +# -------------------------------------------------------------------------- +# The address is stable across calls. This is what the accumulation is about: +# drop the workspace and every call registers a new receive address. And the +# workspace is one grow-only block, so varying request geometries cannot keep +# adding resident buffers the way a per-geometry cache would. +# -------------------------------------------------------------------------- + + +def test_the_landing_buffer_address_is_stable_across_calls(fake_collective): + vae = _vae(4) + pointers = [] + for request in range(3): + _run(vae, fake_collective, _tiles(8, seed=request), 8, 4, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + + assert len(set(pointers)) == 1, f"gather buffer moved across calls: {pointers}" + assert vae._tile_gather_stats == { + "hits": 2, + "allocs": 1, + "workspace_bytes": _workspace_bytes(8, 4), + } + assert vae._tile_gather_workspace.numel() == _workspace_bytes(8, 4) + assert vae._tile_gather_workspace.dtype is torch.uint8 + + +def test_distinct_shapes_share_the_one_workspace(fake_collective): + """One decode gathers several temporal clips, and the head/tail clips differ.""" + vae = _vae(4) + _run(vae, fake_collective, _tiles(8), 8, 4, 0) + first = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(12), 12, 4, 0) + second = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(8, seed=5), 8, 4, 0) + third = fake_collective[-1]["buf_ptr"] + + # The larger geometry did not fit, so the block was replaced -- not added to. + assert second == third + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["hits"] == 1 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(12, 4) + # ``first`` is deliberately not compared: the replaced block is freed, so + # the allocator is free to hand its address back for the larger one. + assert isinstance(first, int) + + +def test_workspace_is_bounded_across_geometries(fake_collective): + """Alternating geometries must not each retain a buffer of their own.""" + sp_size = 4 + large, small, largest = 12, 8, 20 + vae = _vae(sp_size) + + pointers = [] + for round_index in range(3): + for num_tiles in (large, small): + _run(vae, fake_collective, _tiles(num_tiles, seed=round_index), num_tiles, sp_size, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + + assert len(pointers) == 6 + assert len(set(pointers)) == 1, f"workspace moved across geometries: {pointers}" + assert vae._tile_gather_stats["allocs"] == 1 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(large, sp_size) + assert vae._tile_gather_workspace.numel() == _workspace_bytes(large, sp_size) + + # A geometry that does not fit grows the single block; it does not add one. + _run(vae, fake_collective, _tiles(largest), largest, sp_size, 0) + grown = fake_collective[-1]["buf_ptr"] + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(largest, sp_size) + assert vae._tile_gather_workspace.numel() == _workspace_bytes(largest, sp_size) + + # ...and the geometries seen before it now fit, so nothing reallocates again. + for num_tiles in (large, small): + _run(vae, fake_collective, _tiles(num_tiles, seed=7), num_tiles, sp_size, 0) + assert fake_collective[-1]["buf_ptr"] == grown + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(largest, sp_size) + + +def test_workspace_grows_only_upward(fake_collective): + """Resident bytes track the largest geometry seen, not the sum of them.""" + sp_size = 4 + small, large = 8, 12 + vae = _vae(sp_size) + + _run(vae, fake_collective, _tiles(small), small, sp_size, 0) + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(small, sp_size) + + _run(vae, fake_collective, _tiles(large), large, sp_size, 0) + + assert vae._tile_gather_stats["allocs"] == 2 + assert vae._tile_gather_stats["hits"] == 0 + assert vae._tile_gather_stats["workspace_bytes"] == _workspace_bytes(large, sp_size) + assert vae._tile_gather_stats["workspace_bytes"] < _workspace_bytes(small, sp_size) + _workspace_bytes( + large, sp_size + ) + assert vae._tile_gather_workspace.numel() == _workspace_bytes(large, sp_size) + + +def test_mixed_dtype_shares_one_workspace(fake_collective): + """The workspace is bytes, so a narrower dtype reuses it instead of adding one.""" + sp_size, num_tiles = 4, 8 + vae = _vae(sp_size) + fp32_bytes = _workspace_bytes(num_tiles, sp_size, dtype=torch.float32) + + pointers = [] + for round_index in range(2): + for dtype in (torch.float32, torch.bfloat16): + tiles = _tiles(num_tiles, seed=round_index, dtype=dtype) + produced = _run(vae, fake_collective, tiles, num_tiles, sp_size, 0) + pointers.append(fake_collective[-1]["buf_ptr"]) + assert all(tile.dtype is dtype for tile in produced) + for index, tile in enumerate(produced): + assert torch.equal(tile, tiles[index]) + + assert len(set(pointers)) == 1, f"workspace moved across dtypes: {pointers}" + assert vae._tile_gather_stats["allocs"] == 1 + assert vae._tile_gather_stats["hits"] == 3 + assert vae._tile_gather_stats["workspace_bytes"] == fp32_bytes + assert vae._tile_gather_workspace.numel() == fp32_bytes + + +# -------------------------------------------------------------------------- +# Equivalence with the checkpoint's own gather, padded branch included. A +# reordered tile grid is silent at runtime, so order is asserted as well. +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_tiles", [5, 8, 9, 28]) +@pytest.mark.parametrize("sp_rank", [0, 3]) +def test_gathered_tiles_match_the_checkpoint_gather(fake_collective, num_tiles, sp_rank): + sp_size = 4 + all_tiles = _tiles(num_tiles, seed=num_tiles) + vae = _vae(sp_size, sp_rank=sp_rank) + + produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, sp_rank) + expected = _reference_gather(all_tiles, num_tiles, sp_size) + + assert len(produced) == num_tiles + assert all(tile is not None for tile in produced) + for index, (got, want) in enumerate(zip(produced, expected)): + assert torch.equal(got, want), f"tile {index} differs from the checkpoint gather" + # ...and it is the tile that belongs at that index, not a permutation that + # happens to agree with an equally permuted oracle. + for index, tile in enumerate(produced): + assert torch.equal(tile, all_tiles[index]) + + +def test_the_padded_branch_is_actually_exercised(fake_collective): + """9 tiles over 4 ranks: ranks 1-3 are one short, so the padding runs.""" + sp_size, num_tiles = 4, 9 + all_tiles = _tiles(num_tiles, seed=9) + assert len(_local_tiles(all_tiles, num_tiles, 1, sp_size)) < -(-num_tiles // sp_size) + + vae = _vae(sp_size, sp_rank=1) + produced = _run(vae, fake_collective, all_tiles, num_tiles, sp_size, 1) + + sent = fake_collective[-1]["sent"] + assert sent.shape[0] == 3 + assert torch.equal(sent[2], torch.zeros_like(sent[2])) + for index, tile in enumerate(produced): + assert torch.equal(tile, all_tiles[index]) + + +def test_an_empty_local_share_still_raises(fake_collective): + """A rank with no tiles fails loudly, exactly as the checkpoint does.""" + vae = _vae(4) + with pytest.raises(ValueError, match="empty tasks on sp rank"): + vae.model._all_gather_tiled_results([], 8) + + +def test_more_tiles_than_round_robin_allows_is_refused(fake_collective): + """The unpacking arithmetic is only valid for the checkpoint's ownership.""" + vae = _vae(4) + with pytest.raises(ValueError, match="round-robin"): + vae.model._all_gather_tiled_results(_tiles(5), 8) + + +# -------------------------------------------------------------------------- +# Ownership: the buffer is overwritten by the next gather, so returned tiles +# must not alias it. +# -------------------------------------------------------------------------- + + +def test_returned_tiles_do_not_alias_the_pinned_buffer(fake_collective): + sp_size, num_tiles = 4, 8 + first_tiles = _tiles(num_tiles, seed=1) + vae = _vae(sp_size) + produced = _run(vae, fake_collective, first_tiles, num_tiles, sp_size, 0) + + buffer_storage = vae._tile_gather_workspace.untyped_storage().data_ptr() + for tile in produced: + assert tile.untyped_storage().data_ptr() != buffer_storage + + # The next gather reuses that memory; the first call's tiles must survive it. + _run(vae, fake_collective, _tiles(num_tiles, seed=2), num_tiles, sp_size, 0) + for index, tile in enumerate(produced): + assert torch.equal(tile, first_tiles[index]) + + +# -------------------------------------------------------------------------- +# Non-XPU behaviour is unchanged: the checkpoint's own method stays installed. +# -------------------------------------------------------------------------- + + +def _construct(monkeypatch, device, model): + """Run the real ``__init__`` against a stub checkpoint component.""" + + class _Remote(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1)) + self.model = model + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda path: {"latent_channels": 1, "latents_mean": [0.0], "latents_std": [1.0]}, + ) + # **kwargs: the loader gained a trust_remote_code argument at some point, + # and this stub does not care either way. + monkeypatch.setattr(vae_module, "_load_remote_component", lambda path, config, **kwargs: _Remote()) + # The adapter only stores this handle; resolving it for real needs an + # accelerator, which a CPU test runner does not have. + monkeypatch.setattr(torch, "get_device_module", lambda *args, **kwargs: types.SimpleNamespace()) + return MiniMaxH3VideoVAE("unused", device=torch.device(device)) + + +@pytest.mark.parametrize("device", ["cpu", "meta"]) +def test_non_xpu_devices_keep_the_checkpoint_gather(monkeypatch, device): + model = _FakeCheckpointModel() + original = model._all_gather_tiled_results + + vae = _construct(monkeypatch, device, model) + + assert vae._tile_gather_reuse_enabled() is False + assert vae.model._all_gather_tiled_results is original + assert vae._checkpoint_tile_gather is None + + +@pytest.mark.parametrize( + ("device", "enabled"), + [("xpu", True), ("cpu", False), ("cuda", False), ("meta", False)], +) +def test_only_xpu_asks_for_the_pinned_buffer(device, enabled): + vae = object.__new__(MiniMaxH3VideoVAE) + vae._device_target = torch.device(device) + assert vae._tile_gather_reuse_enabled() is enabled + + +def test_construction_installs_the_override_when_the_device_asks_for_it(monkeypatch): + """Close the loop: ``__init__`` must consult the predicate, not ignore it.""" + monkeypatch.setattr(MiniMaxH3VideoVAE, "_tile_gather_reuse_enabled", lambda self: True) + model = _FakeCheckpointModel() + original = model._all_gather_tiled_results + + vae = _construct(monkeypatch, "cpu", model) + + assert vae.model._all_gather_tiled_results is not original + assert vae._checkpoint_tile_gather is original + + +# -------------------------------------------------------------------------- +# The decision-point receipt carries quantities, not just presence: a line +# that only says "installed" cannot tell reuse from reallocation. +# -------------------------------------------------------------------------- + + +def test_the_receipt_line_reports_hits_allocs_and_the_address(fake_collective, gather_log): + vae = _vae(4) + install = [ + (level, line) for level, line in gather_log.records if "persistent tile gather installed device=xpu" in line + ] + # The one-shot install line is what an operator needs at info level. + assert install and all(level == logging.INFO for level, _ in install) + gather_log.messages.clear() + gather_log.records.clear() + + _run(vae, fake_collective, _tiles(8), 8, 4, 0) + first_ptr = fake_collective[-1]["buf_ptr"] + _run(vae, fake_collective, _tiles(8, seed=3), 8, 4, 0) + + records = [(level, line) for level, line in gather_log.records if "[H3_VAE_GATHER] reuse=" in line] + # The per-gather line fires once per decoder tile batch, so it is debug. + assert all(level == logging.DEBUG for level, _ in records) + lines = [line for _, line in records] + assert len(lines) == 2 + assert "reuse=alloc" in lines[0] and "hits=0 allocs=1" in lines[0] + assert "reuse=hit" in lines[1] and "hits=1 allocs=1" in lines[1] + # Resident bytes are reported, not just the bytes this one gather needed. + needed = _workspace_bytes(8, 4) / (1024.0 * 1024.0) + for line in lines: + assert f"need_mib={needed:.2f}" in line + assert f"workspace_mib={needed:.2f}" in line + assert f"buf_ptr=0x{first_ptr:x}" in lines[0] + assert f"buf_ptr=0x{first_ptr:x}" in lines[1] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py index 7a1c54f7b..86ccaa27b 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py @@ -198,7 +198,7 @@ def test_video_vae_installs_exact_optimizations(monkeypatch): monkeypatch.setattr( vae_module, "_load_remote_component", - lambda _path, _config: FakeRemote(), + lambda _path, _config, **_kwargs: FakeRemote(), ) install = Mock(return_value=True) monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 01332beec..57b2daaf8 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -635,6 +635,32 @@ def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): } +def test_i2v_minimax_h3_preserves_reference_geometry(): + engine = FakeAsyncOmni() + engine.model_class_name = "MiniMaxH3Pipeline" + handler = OmniOpenAIServingVideo.for_diffusion( + diffusion_engine=engine, + model_name="MiniMaxAI/MiniMax-H3", + ) + image = Image.new("RGB", (48, 32)) + + asyncio.run( + handler._run_and_extract( + VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64), + "minimax-h3-reference-geometry", + reference_image=ReferenceImage(image), + ) + ) + + assert engine.captured_prompt is not None + assert engine.captured_sampling_params_list is not None + input_image = engine.captured_prompt["multi_modal_data"]["image"] + assert isinstance(input_image, Image.Image) + assert input_image.size == (48, 32) + sampling_params = engine.captured_sampling_params_list[0] + assert (sampling_params.width, sampling_params.height) == (96, 64) + + def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture): image_bytes = _make_test_image_bytes((48, 48)) mocker.patch( diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index dcdf13678..a187c570c 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -290,6 +290,16 @@ def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.Monkey assert captured == {"model": "org/model", "revision": "pinned-revision"} +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("model_class_name", ["MiniMaxH3Pipeline", "MiniMaxH3ModularPipeline"]) +def test_minimax_h3_preserves_reference_image_size(model_class_name: str) -> None: + assert should_preserve_reference_image_size( + model_class_name, + model="MiniMaxAI/MiniMax-H3", + ) + + @pytest.mark.core_model @pytest.mark.cpu def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None: diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index e903960d6..592f5ba9b 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -215,8 +215,9 @@ class AttentionMetadata: # variable-length query/key sequences for FlashAttention. # "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with # the packed cu_seqlens tensors. - # "valid_kv_length": int — contiguous valid K/V prefix length for a - # backend that advertises supports_prefix_kv_slicing. + # "valid_kv_length": int — contiguous valid K/V prefix length. Local + # backends may slice it directly; ring backends trim each circulated + # K/V block from the same global prefix. # "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path # (TND npu_fusion_attention driven by cu_seqlens, mask never read). # Requires the [real, pad] two-document packing contract; see diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py index b8743fe13..a898c7729 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention @@ -9,7 +10,33 @@ import torch.nn.functional as F LseLayout = Literal["bhs", "bsh"] -__all__ = ["update_out_and_lse", "flatten_varlen_lse", "unflatten_varlen_lse"] +__all__ = [ + "ring_kv_block_valid_length", + "update_out_and_lse", + "flatten_varlen_lse", + "unflatten_varlen_lse", +] + + +def ring_kv_block_valid_length( + valid_kv_length: int | None, + block_size: int, + block_rank: int, + world_size: int, +) -> int: + """Return the valid prefix length in one globally ordered ring K/V block.""" + if valid_kv_length is None: + return block_size + if isinstance(valid_kv_length, bool) or not isinstance(valid_kv_length, int): + raise ValueError("valid_kv_length must be a Python integer") + total_length = block_size * world_size + if not 0 < valid_kv_length <= total_length: + raise ValueError( + "valid_kv_length must be within the global ring K/V sequence, " + f"got {valid_kv_length} for length {total_length}" + ) + block_start = block_rank * block_size + return max(0, min(block_size, valid_kv_length - block_start)) def _normalize_lse(block_lse: torch.Tensor, block_out: torch.Tensor, lse_layout: LseLayout) -> torch.Tensor: diff --git a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py b/vllm_omni/diffusion/attention/backends/ring_flash_attn.py index 1fce4f1b4..87c3b9e3d 100644 --- a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/ring_flash_attn.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention @@ -6,7 +7,10 @@ import torch from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType, select_flash_attn_impl -from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse +from vllm_omni.diffusion.attention.backends.ring.ring_utils import ( + ring_kv_block_valid_length, + update_out_and_lse, +) from vllm_omni.diffusion.distributed.comm import RingComm @@ -27,6 +31,7 @@ def ring_flash_attn_forward( joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): # Validate causal + joint_strategy combination # When causal=True and joint_strategy="rear", the causal mask would incorrectly @@ -42,6 +47,21 @@ def ring_flash_attn_forward( "from seeing the joint tokens." ) + # Trimming the circulated K/V blocks changes their length while the query + # block keeps the full padded length. Flash-Attention aligns a causal mask + # to the bottom-right of a rectangular (seqlen_q, seqlen_k) tile, so a + # trimmed block would shift the diagonal and let queries attend to the + # wrong key positions -- silently, with no shape error. The two are + # therefore an unsupported combination rather than a slow path: the caller + # must either drop the padding before the ring or run non-causal. + if causal and valid_kv_length is not None: + raise ValueError( + "valid_kv_length is not supported with causal=True in Ring Attention. " + "Trimming a circulated K/V block shortens seqlen_k while seqlen_q keeps " + "the padded length, and the causal diagonal is bottom-right aligned, so " + "the mask would silently shift. Unpad before the ring, or use causal=False." + ) + comm = RingComm(process_group) out = None @@ -66,8 +86,15 @@ def ring_flash_attn_forward( comm.commit() if not causal or step <= comm.rank: - step_k = k - step_v = v + block_rank = (comm.rank - step) % comm.world_size + block_valid_length = ring_kv_block_valid_length( + valid_kv_length, + k.shape[1], + block_rank, + comm.world_size, + ) + step_k = k[:, :block_valid_length] + step_v = v[:, :block_valid_length] if step == 0 and joint_tensor_key is not None: if joint_strategy == "front": step_k = torch.cat([joint_tensor_key, step_k], dim=1) @@ -76,27 +103,28 @@ def ring_flash_attn_forward( step_k = torch.cat([step_k, joint_tensor_key], dim=1) step_v = torch.cat([step_v, joint_tensor_value], dim=1) - fn = select_flash_attn_impl(attn_type, stage="fwd-only", attn_processor=attn_processor) - block_out, block_lse = fn( - q, - step_k, - step_v, - dropout_p=dropout_p, - softmax_scale=softmax_scale, - causal=causal and step == 0, - window_size=window_size, - softcap=softcap, - alibi_slopes=alibi_slopes, - return_softmax=True and dropout_p > 0, - ) + if step_k.shape[1] > 0: + fn = select_flash_attn_impl(attn_type, stage="fwd-only", attn_processor=attn_processor) + block_out, block_lse = fn( + q, + step_k, + step_v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal and step == 0, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + return_softmax=True and dropout_p > 0, + ) - # Ensure block_out is contiguous if needed, though usually it is from FA + # Ensure block_out is contiguous if needed, though usually it is from FA - if attn_type == AttnType.SPARSE_SAGE: - out, lse = block_out, block_lse - else: - # Ring kernel wrappers canonicalize LSE to (B, H, S). - out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") + if attn_type == AttnType.SPARSE_SAGE: + out, lse = block_out, block_lse + else: + # Ring kernel wrappers canonicalize LSE to (B, H, S). + out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") if step + 1 != comm.world_size: comm.wait() @@ -132,6 +160,7 @@ class RingFlashAttnFunc(torch.autograd.Function): joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) @@ -158,6 +187,7 @@ class RingFlashAttnFunc(torch.autograd.Function): joint_tensor_key=joint_tensor_key, joint_tensor_value=joint_tensor_value, joint_strategy=joint_strategy, + valid_kv_length=valid_kv_length, ) return out if not return_softmax else (out, softmax_lse, None) @@ -193,6 +223,7 @@ def ring_flash_attn_qkvpacked_func( None, # joint_tensor_key None, # joint_tensor_value "front", # joint_strategy + None, # valid_kv_length ) @@ -228,6 +259,7 @@ def ring_flash_attn_kvpacked_func( None, # joint_tensor_key None, # joint_tensor_value "front", # joint_strategy + None, # valid_kv_length ) @@ -249,6 +281,7 @@ def ring_flash_attn_func( joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, None]: """Ring Attention forward pass using Flash Attention backend. @@ -290,6 +323,8 @@ def ring_flash_attn_func( Defaults to None. joint_strategy (str): Concatenation strategy ("front" or "back"). Defaults to "front". + valid_kv_length (int | None): Length of the contiguous valid prefix in + the global ring K/V sequence. Defaults to the full sequence. Returns: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, None]]: @@ -314,4 +349,5 @@ def ring_flash_attn_func( joint_tensor_key, joint_tensor_value, joint_strategy, + valid_kv_length, ) diff --git a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py index 7ab60b784..65f5d8ad7 100644 --- a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py +++ b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention @@ -10,7 +11,10 @@ import torch from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.ring.ring_kernels import pytorch_attn_forward -from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse +from vllm_omni.diffusion.attention.backends.ring.ring_utils import ( + ring_kv_block_valid_length, + update_out_and_lse, +) from vllm_omni.diffusion.distributed.comm import RingComm logger = init_logger(__name__) @@ -33,6 +37,7 @@ def ring_pytorch_attn_func( joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): return RingAttentionFunc.apply( group, @@ -45,6 +50,7 @@ def ring_pytorch_attn_func( joint_tensor_key, joint_tensor_value, joint_strategy, + valid_kv_length, ) @@ -64,6 +70,7 @@ class RingAttentionFunc(torch.autograd.Function): joint_tensor_key=None, joint_tensor_value=None, joint_strategy="front", + valid_kv_length: int | None = None, ): # Validate causal + joint_strategy combination # When causal=True and joint_strategy="rear", the causal mask would incorrectly @@ -79,6 +86,22 @@ class RingAttentionFunc(torch.autograd.Function): "from seeing the joint tokens." ) + # Trimming the circulated K/V blocks changes their length while the + # query block keeps the full padded length. A causal SDPA mask over a + # rectangular (seqlen_q, seqlen_k) tile is aligned to the bottom right, + # so a trimmed block would shift the diagonal and let queries attend to + # the wrong key positions -- silently, with no shape error. The two are + # therefore an unsupported combination rather than a slow path: the + # caller must either drop the padding before the ring or run + # non-causal. + if is_causal and valid_kv_length is not None: + raise ValueError( + "valid_kv_length is not supported with causal=True in Ring Attention. " + "Trimming a circulated K/V block shortens seqlen_k while seqlen_q keeps " + "the padded length, and the causal diagonal is bottom-right aligned, so " + "the mask would silently shift. Unpad before the ring, or use causal=False." + ) + comm = RingComm(group) # Ensure tensors are contiguous for P2P communication q = q.contiguous() @@ -98,8 +121,15 @@ class RingAttentionFunc(torch.autograd.Function): comm.commit() if not is_causal or step <= comm.rank: - step_k = k - step_v = v + block_rank = (comm.rank - step) % comm.world_size + block_valid_length = ring_kv_block_valid_length( + valid_kv_length, + k.shape[1], + block_rank, + comm.world_size, + ) + step_k = k[:, :block_valid_length] + step_v = v[:, :block_valid_length] if step == 0 and joint_tensor_key is not None: if joint_strategy == "front": step_k = torch.cat([joint_tensor_key, step_k], dim=1) @@ -108,15 +138,16 @@ class RingAttentionFunc(torch.autograd.Function): step_k = torch.cat([step_k, joint_tensor_key], dim=1) step_v = torch.cat([step_v, joint_tensor_value], dim=1) - block_out, block_lse = pytorch_attn_forward( - q, - step_k, - step_v, - softmax_scale=sm_scale, - causal=is_causal and step == 0, - op_type=op_type, - ) - out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") + if step_k.shape[1] > 0: + block_out, block_lse = pytorch_attn_forward( + q, + step_k, + step_v, + softmax_scale=sm_scale, + causal=is_causal and step == 0, + op_type=op_type, + ) + out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs") if step + 1 != comm.world_size: comm.wait() diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index d71359ae7..1eb811636 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/vllm_omni/diffusion/attention/parallel/ring.py @@ -156,9 +156,11 @@ class RingParallelAttention: # Extract joint tensors joint_key, joint_value = None, None joint_strategy = "front" + valid_kv_length = None if attn_metadata is not None: joint_key = attn_metadata.joint_key joint_value = attn_metadata.joint_value + valid_kv_length = attn_metadata.extra.get("valid_kv_length") if attn_metadata.joint_strategy is not None: joint_strategy = attn_metadata.joint_strategy @@ -176,6 +178,7 @@ class RingParallelAttention: joint_tensor_key=joint_key, joint_tensor_value=joint_value, joint_strategy=joint_strategy, + valid_kv_length=valid_kv_length, ) # Ring only implements local Flash/AITER and SDPA kernels. HuggingFace @@ -261,4 +264,5 @@ class RingParallelAttention: joint_tensor_key=joint_key, joint_tensor_value=joint_value, joint_strategy=joint_strategy, + valid_kv_length=valid_kv_length, ) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 3abc90476..f0839429b 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -44,6 +44,10 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, +) from vllm_omni.diffusion.layers.activation import SiluAndMul from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm @@ -67,6 +71,15 @@ logger = init_logger(__name__) # packing, and ``_run_packed_attention`` re-checks it per forward; a name-only # gate would let FLASH_ATTN's NPU/XPU code paths through even though those # variants would silently attend across request boundaries. +def _ring_sequence_parallel_is_active(attention_layer: Attention) -> bool: + """Match :meth:`Attention._get_active_parallel_strategy` for Ring.""" + if not getattr(attention_layer, "use_ring", False) or getattr(attention_layer, "skip_sequence_parallel", False): + return False + if is_forward_context_available() and not get_forward_context().sp_active: + return False + return True + + def _attention_isolates_packed_requests(attention_layer: Any) -> bool: """True if this attention layer keeps N-document packed boundaries. @@ -78,7 +91,7 @@ def _attention_isolates_packed_requests(attention_layer: Any) -> bool: backend = getattr(attention_layer, "attn_backend", None) if backend is None or not backend.supports_multi_doc_packed_varlen(): return False - return not getattr(attention_layer, "use_ring", False) + return not _ring_sequence_parallel_is_active(attention_layer) @dataclass @@ -508,6 +521,7 @@ class MiniMaxH3Attention(nn.Module): ) attn_mask = None mask_free_packed_padding = False + use_ring = _ring_sequence_parallel_is_active(self.attention) if num_requests > 1: # A step-mode batch packs one document per request, so its valid # rows are block-diagonal rather than a prefix: neither a KV prefix @@ -534,12 +548,15 @@ class MiniMaxH3Attention(nn.Module): # supports_packed_mask_free: backend consumes the packed metadata # without ever reading attn_mask (CUDA packed varlen, NPU # npu_attn_varlen opt-in with its own fallback rebuild). - use_ring = getattr(self.attention, "use_ring", False) mask_free_packed_padding = not use_ring and self.attention.attn_backend.supports_packed_mask_free() no_mask = not use_ring and ( self.attention.attn_backend.supports_prefix_kv_slicing or mask_free_packed_padding ) - if used < packed_total and not no_mask: + # Hybrid Ulysses reshards Q to one ring partition before the ring + # kernel runs, so a global [packed_total] mask cannot pass its + # query-length check. Ring consumes valid_kv_length directly and + # trims the circulated K/V blocks instead. + if used < packed_total and not no_mask and not use_ring: attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, @@ -563,7 +580,7 @@ class MiniMaxH3Attention(nn.Module): # quadratic full_qk mask is never materialized. Ring attention # is excluded: it keeps the aligned padding rows for its # fixed-size P2P buffers and still needs the mask. - "npu_attn_varlen": not getattr(self.attention, "use_ring", False), + "npu_attn_varlen": not use_ring, # fp16-range protection for the ascend_laser_attention kernel # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other # backend/path. diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py index 0dcbd19b8..ecf797a9c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py +++ b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 packed-sequence materialization for FL2VA and Ref2VA tasks. Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P]. @@ -18,6 +19,8 @@ from collections.abc import Mapping, Sequence import numpy as np import torch +from vllm_omni.errors import OmniClientError + MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: tuple[tuple[int, ...], ...] = ( (0,), (-1,), @@ -38,13 +41,42 @@ _INTERP = 32 _T_GROUP = 5 _FRAME_PER_TOKEN = (1, 4, 4, 4, 4) _FRAME_RESCALE = 5.0 / 3.0 -_SEQ_ALIGN = 64 +MINIMAX_H3_SEQ_ALIGN = 64 +"""Row alignment of the packed sequence. + +An explicit ``seq_len`` must be a multiple of this, so that pinning a length +lands on a bucket the default rounding could also produce. +""" + +MINIMAX_H3_MAX_PAD_SEQ_LEN = 1 << 20 +"""Ceiling for a request-pinned packed length. + +Every packed row costs a fixed number of bytes across the structural tensors +built below, so an unbounded request value would size those allocations. The +largest shape this pipeline is documented against -- 672x384 at 209 frames -- +packs about 16k rows, so this ceiling leaves roughly two orders of magnitude of +headroom while keeping the structural tensors in the tens of megabytes. +""" # MiniMax H3 packs video latents with a fixed [1, 2, 2] (t, h, w) patch. _PATCH_T = 1 _PATCH_H = 2 _PATCH_W = 2 +def _resolve_padded_seq_len(used: int, seq_len: int | None) -> int: + """Resolve the padded row count for a packed sequence. + + Without an explicit ``seq_len`` the used rows are rounded up to the next + :data:`MINIMAX_H3_SEQ_ALIGN` boundary. An explicit value is honored as-is + and must cover the used rows. + """ + if seq_len is None: + return ((used + MINIMAX_H3_SEQ_ALIGN - 1) // MINIMAX_H3_SEQ_ALIGN) * MINIMAX_H3_SEQ_ALIGN + if seq_len < used: + raise OmniClientError(f"seq_len {seq_len} < used rows {used}") + return seq_len + + def _keyframe_cond_frame_indices( *, include_keyframe_cond: bool, @@ -124,10 +156,15 @@ def minimax_h3_packed_sequence( include_keyframe_cond: bool, keyframe_frame_indices: list[int] | tuple[int, ...] | None = None, frame_count: int | None = None, + seq_len: int | None = None, ) -> dict[str, object]: """Build the packed-sequence structural fields for one CFG branch. - The used length is padded up to a multiple of 64. + The used length is padded up to a multiple of + :data:`MINIMAX_H3_SEQ_ALIGN`, or to an explicit ``seq_len`` that covers the + used rows -- the same contract + :func:`minimax_h3_packed_sequence_ref2va_blocks` already offers. Pinning + the length keeps requests of different prompt lengths on one packed shape. """ ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W frame_rows = ph * pw @@ -143,7 +180,7 @@ def minimax_h3_packed_sequence( video_rows = latent_t * frame_rows audio_rows = audio_t * audio_channel used = text_len + cond_rows + audio_rows + video_rows - seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN + seq_len = _resolve_padded_seq_len(used, seq_len) text_sl = slice(0, text_len) cond_sl = slice(text_len, text_len + cond_rows) @@ -376,10 +413,7 @@ def minimax_h3_packed_sequence_ref2va_blocks( audio_rows = audio_t * audio_channel ref_rows = ref_visual_rows + ref_audio_rows used = text_len + ref_rows + audio_rows + video_rows - if seq_len is None: - seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN - if seq_len < used: - raise ValueError(f"seq_len {seq_len} < used rows {used}") + seq_len = _resolve_padded_seq_len(used, seq_len) text_sl = slice(0, text_len) cursor = text_len 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 f80513d4e..8f814a27a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -137,6 +137,8 @@ from .npu.lora import ( load_minimax_h3_native_lora, ) from .packed_sequence import ( + MINIMAX_H3_MAX_PAD_SEQ_LEN, + MINIMAX_H3_SEQ_ALIGN, minimax_h3_packed_sequence, minimax_h3_packed_sequence_ref2va_blocks, ) @@ -288,6 +290,7 @@ _MINIMAX_H3_DENOISE_INPUT_KEYS = ( "visual_condition_shapes", "audio_condition_lengths", "keyframe_frame_indices", + "pad_seq_len", ) # ``StepRequestState.extra`` keys owned by the step-execution path. @@ -543,6 +546,28 @@ def _resolve_minimax_h3_num_outputs(value: Any) -> int: return value +def _resolve_pad_seq_len(value: object) -> int | None: + """Validate the optional packed-length pin from ``extra_args``. + + The packer itself only needs a value that covers the used rows. A request + field needs two more guards: an unaligned value would silently open yet + another compiled shape, and an unbounded one would size every structural + tensor the packer allocates. + """ + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, np.integer)): + raise OmniClientError("MiniMax H3 pad_seq_len must be an integer") + pinned = int(value) + if pinned <= 0: + raise OmniClientError(f"MiniMax H3 pad_seq_len must be positive, got {pinned}") + if pinned % MINIMAX_H3_SEQ_ALIGN: + raise OmniClientError(f"MiniMax H3 pad_seq_len must be a multiple of {MINIMAX_H3_SEQ_ALIGN}, got {pinned}") + if pinned > MINIMAX_H3_MAX_PAD_SEQ_LEN: + raise OmniClientError(f"MiniMax H3 pad_seq_len must be at most {MINIMAX_H3_MAX_PAD_SEQ_LEN}, got {pinned}") + return pinned + + def _minimax_h3_output_seeds(seed: int, num_outputs: int) -> list[int]: return [int(seed) + output_index for output_index in range(int(num_outputs))] @@ -1077,11 +1102,13 @@ class MiniMaxH3Pipeline( os.path.join(model_path, "video_vae"), device=self.device, load_device=component_load_device, + trust_remote_code=od_config.trust_remote_code, ) self.audio_vae = MiniMaxH3AudioVAE( os.path.join(model_path, "audio_vae"), device=self.device, load_device=component_load_device, + trust_remote_code=od_config.trust_remote_code, ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae @@ -1674,6 +1701,31 @@ class MiniMaxH3Pipeline( logger.exception("Failed to release retained allocator cache after offload failure") raise + @staticmethod + def _is_output_owner_rank() -> bool: + """Whether this rank's output is returned by the diffusion executor.""" + return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + + def _offload_model_cpu_stage_output(self, tensor: torch.Tensor) -> torch.Tensor: + """Release a decoded output's storage before a later seed reloads the DiT. + + Model-level CPU offload makes the VAE hook evict the DiT before decode. + For a multi-output request, however, the next seed reloads the DiT while + the preceding decoded output would otherwise still occupy the device. + The reply rank performs the D2H copy early; ranks whose outputs are not + consumed retain only a metadata placeholder for the final concatenation. + + Call this on the tensor that is actually returned to the engine. Audio is + released inside ``decode``; video is released by the callers of ``decode`` + only after ``_prepare_minimax_h3_video_output`` has quantized it, so the + early D2H copy carries ``uint8`` frames rather than the decoded floats. + """ + if not getattr(self, "_model_cpu_offload_modules", None): + return tensor + if self._is_output_owner_rank(): + return tensor.cpu() + return torch.empty(tensor.shape, dtype=tensor.dtype, device="meta") + def _encode_visual_conditions( self, images: list[Image.Image], @@ -1932,6 +1984,7 @@ class MiniMaxH3Pipeline( visual_condition_shapes: list[tuple[int, int, int]] | None = None, audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, + pad_seq_len: int | None = None, ) -> dict[str, Any]: """Build the packed layout, initial rows, anchors, and sigma schedules. @@ -1961,6 +2014,7 @@ class MiniMaxH3Pipeline( latent_w=latent_w, audio_t=audio_t, ref_blocks=ref_blocks, + seq_len=pad_seq_len, ) else: packed = minimax_h3_packed_sequence( @@ -1972,7 +2026,19 @@ class MiniMaxH3Pipeline( include_keyframe_cond=task == "fl2va", keyframe_frame_indices=keyframe_frame_indices if task == "fl2va" else None, frame_count=num_frames if task == "fl2va" else None, + seq_len=pad_seq_len, ) + # Report the effective shape at info level exactly when a request pins + # it, so a deployment can confirm the pin landed without raising the + # log level for every request. + log = logger.info if pad_seq_len is not None else logger.debug + log( + "MiniMax H3 packed sequence: task=%s pad_seq_len=%s used=%d seq_len=%d", + task, + pad_seq_len, + int(packed["cu_seqlens"][1]), + int(packed["seq_len"]), + ) tags = packed["token_tags"].clone() tags[packed["text_pos"]] = text_tags.cpu() @@ -2108,6 +2174,7 @@ class MiniMaxH3Pipeline( visual_condition_shapes: list[tuple[int, int, int]] | None = None, audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, + pad_seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = self._build_denoise_inputs( task=task, @@ -2131,6 +2198,7 @@ class MiniMaxH3Pipeline( visual_condition_shapes=visual_condition_shapes, audio_condition_lengths=audio_condition_lengths, keyframe_frame_indices=keyframe_frame_indices, + pad_seq_len=pad_seq_len, ) branch = inputs["branch"] transformer = self._transformer_for_task(task) @@ -2253,6 +2321,7 @@ class MiniMaxH3Pipeline( video = video[..., :height, :width].contiguous() with self._component_on_device(self.audio_vae): audio = self.audio_vae.decode_latent(audio_latent) + audio = self._offload_model_cpu_stage_output(audio) return video, audio @staticmethod @@ -2561,6 +2630,7 @@ class MiniMaxH3Pipeline( 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)) + pad_seq_len = _resolve_pad_seq_len(extra.get("pad_seq_len")) quality_plan = self._quality_policy.resolve( quality=quality, num_inference_steps=num_steps, @@ -2587,6 +2657,7 @@ class MiniMaxH3Pipeline( "visual_condition_shapes": visual_shapes, "audio_condition_lengths": audio_lengths, "keyframe_frame_indices": keyframe_frame_indices, + "pad_seq_len": pad_seq_len, "seed": seed, "num_steps": num_steps, "video_shift": video_shift, @@ -2643,7 +2714,13 @@ class MiniMaxH3Pipeline( height=context["height"], width=context["width"], ) - videos.append(_prepare_minimax_h3_video_output(video)) + # Rebind rather than append the expression: the local would + # otherwise keep the decoded frames' device storage alive for the + # rest of the iteration, and the next seed's diffuse() -- the DiT + # reload this release exists to make room for -- runs before the + # loop rebinds it. post_decode() rebinds for the same reason. + video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) + videos.append(video) audios.append(audio) if videos and isinstance(videos[0], bytes): video = videos[0] if len(videos) == 1 else videos @@ -2966,7 +3043,7 @@ class MiniMaxH3Pipeline( height=shape["height"], width=shape["width"], ) - video = _prepare_minimax_h3_video_output(video) + video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video)) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index ca8359904..e40251dff 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -75,11 +75,23 @@ def _load_component_config(component_path: str) -> dict[str, Any]: def _load_remote_component( component_path: str, config: dict[str, Any], + *, + trust_remote_code: bool, ) -> nn.Module: auto_map = config.get("auto_map") or {} class_reference = auto_map.get("AutoModel") if not isinstance(class_reference, str): raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel") + if not trust_remote_code: + raise ValueError( + f"Loading {component_path} executes the modeling code shipped with " + f"the checkpoint (auto_map.AutoModel = {class_reference}). Pass " + "--trust-remote-code (or trust_remote_code=True) to allow it." + ) + # ``trust_remote_code`` is checked here rather than forwarded to + # ``get_class_from_dynamic_module``: that helper takes no such argument and + # would silently absorb it into ``**kwargs``, so forwarding it would read + # like a gate while executing the remote code unconditionally. component_cls = get_class_from_dynamic_module( class_reference, component_path, @@ -144,6 +156,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): *, device: torch.device, load_device: torch.device | None = None, + trust_remote_code: bool = False, ) -> None: super().__init__() self._device_target = device @@ -151,6 +164,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self.remote = _load_remote_component( component_path, self.config_dict, + trust_remote_code=trust_remote_code, ) # Match the reference loader contract before installing inference-only # decoder fast paths. Keyframe encoding remains FP32; decoder Linear @@ -176,6 +190,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self.use_slicing = False self.parallel_size = 1 self.device_module = torch.get_device_module() + self._tile_gather_workspace: torch.Tensor | None = None + self._tile_gather_stats = {"hits": 0, "allocs": 0, "workspace_bytes": 0} + self._checkpoint_tile_gather = None + if self._tile_gather_reuse_enabled(): + self._install_persistent_tile_gather() def load_to_device(self) -> None: if self._stager is not None: @@ -241,6 +260,145 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): parallel_module = importlib.import_module(f"{package}.parallel") return parallel_module.get_parallel_state() + def _tile_gather_reuse_enabled(self) -> bool: + """Whether this device needs the tiled-VAE gather address kept stable. + + XPU only. The accumulation this guards against is an XCCL-side + registration that is kept for every distinct receive-buffer address and + never reclaimed; no other backend in tree does that, so everywhere else + the checkpoint's own method is left in place and behaviour is unchanged. + The stable address comes from one grow-only byte workspace, so what the + adapter holds resident is one buffer for the largest geometry it has + seen rather than one per geometry. + """ + + return self._device_target.type == "xpu" + + def _install_persistent_tile_gather(self) -> None: + """Route the checkpoint's tiled-VAE gather through a bounded workspace. + + The checkpoint's ``_all_gather_tiled_results`` allocates its gather + output afresh on every call. Model-level offload calls ``empty_cache()`` + once per request, so the next request's output lands on a new device + address; XCCL registers a non-reclaimable resource per new receive + address, and the registrations accumulate until the card is full. The + replacement below keeps a single byte workspace alive on this adapter + and carves every gather out of its front, so the address the collective + writes to is the same one every request. + + The workspace only ever grows: a geometry that fits in the current + block reuses it as is, and a larger one replaces the block, which + drops the previous allocation before allocating its own. Resident + memory is therefore bounded by one buffer for the largest geometry the + adapter has seen, not by one buffer per distinct geometry. + + The override is bound on the checkpoint *instance*, not its class: the + class is remote code shared by every component loaded from the same + checkpoint, and only this adapter knows the device it runs on. + """ + + self._checkpoint_tile_gather = self.model._all_gather_tiled_results + self.model._all_gather_tiled_results = self._persistent_tile_gather + logger.info( + "[H3_VAE_GATHER] persistent tile gather installed device=%s", + self._device_target.type, + ) + + def _persistent_tile_gather( + self, + tasks: list[torch.Tensor], + num_tiles: int, + ) -> list[torch.Tensor]: + """Equal-shape replacement for the checkpoint's tiled-result gather. + + Contract kept identical to the checkpoint's method: return a list of + ``num_tiles`` tensors in global tile order, and raise on an empty local + share so a rank that owns no tile cannot silently skip the collective. + + Tile ownership is round-robin (``range(sp_rank, num_tiles, sp_size)``), + so every rank can compute every other rank's task count from + ``num_tiles`` and ``sp_size`` alone. That makes the per-rank payloads + equal once the leading task dimension is padded to ``max_tasks``, which + is what lets a single ``all_gather_into_tensor`` into a stable buffer + replace the variable-shape gather. + + The landing buffer is a view onto the front of one byte workspace that + this adapter keeps alive and only ever grows, so its address stays + stable across requests while resident memory stays bounded by the + largest geometry seen rather than growing per geometry. A byte + workspace is dtype-agnostic, so mixed-precision requests share it too. + + Returned tiles are cloned out of the buffer: the buffer is overwritten + by the next call and callers hold the tiles past that point. + """ + + state = self._native_parallel_state() + group = state["sp_process_group"] + sp_size = int(state["sp_size"]) + sp_rank = int(state["sp_rank"]) + + if not tasks: + raise ValueError(f"Found empty tasks on sp rank {sp_rank}") + + max_tasks = -(-num_tiles // sp_size) + if len(tasks) > max_tasks: + raise ValueError( + f"sp rank {sp_rank} holds {len(tasks)} tiles but round-robin " + f"ownership of {num_tiles} tiles across {sp_size} ranks allows " + f"at most {max_tasks}" + ) + if len(tasks) == max_tasks: + stacked = torch.stack(tasks, dim=0) + else: + # Pad the leading (task) dimension only. The padded slots belong to + # ranks whose share is short by construction, and the unpacking loop + # below never reads them back. + stacked = tasks[0].new_empty((max_tasks, *tasks[0].shape)) + torch.stack(tasks, dim=0, out=stacked[: len(tasks)]) + stacked[len(tasks) :].zero_() + + need_bytes = sp_size * stacked.numel() * stacked.element_size() + workspace = self._tile_gather_workspace + if workspace is None or workspace.device != stacked.device or workspace.numel() < need_bytes: + # Drop the previous block before allocating its replacement so the + # two are never resident at once; only one workspace is ever held. + workspace = None + self._tile_gather_workspace = None + workspace = torch.empty(need_bytes, dtype=torch.uint8, device=stacked.device) + self._tile_gather_workspace = workspace + self._tile_gather_stats["allocs"] += 1 + self._tile_gather_stats["workspace_bytes"] = need_bytes + reuse = "alloc" + else: + self._tile_gather_stats["hits"] += 1 + reuse = "hit" + buffer = workspace[:need_bytes].view(stacked.dtype).view(sp_size, *stacked.shape) + # Quantities, not just presence: a line that only says "installed" + # cannot distinguish a buffer that is being reused from one that is + # reallocated every request, which is the whole failure being fixed. + # Debug level, because this fires once per decoder tile batch (12 times + # per request on the canonical 1344x768 geometry) and the one-shot + # "installed" line above is what an operator needs at info level. + logger.debug( + "[H3_VAE_GATHER] reuse=%s shape=%s/%s need_mib=%.2f workspace_mib=%.2f buf_ptr=0x%x hits=%d allocs=%d", + reuse, + tuple(stacked.shape), + stacked.dtype, + need_bytes / (1024.0 * 1024.0), + self._tile_gather_stats["workspace_bytes"] / (1024.0 * 1024.0), + workspace.data_ptr(), + self._tile_gather_stats["hits"], + self._tile_gather_stats["allocs"], + ) + dist.all_gather_into_tensor(buffer, stacked, group=group) + + results: list[torch.Tensor] = [None] * num_tiles # type: ignore[list-item] + for rank in range(sp_size): + num_rank_tasks = -(-(num_tiles - rank) // sp_size) + for k in range(num_rank_tasks): + results[k * sp_size + rank] = buffer[rank][k].clone() + return results + def _decoder_tile_count(self, latent: torch.Tensor) -> int: """Number of decoder tiles the checkpoint will split ``latent`` into. @@ -464,6 +622,7 @@ class MiniMaxH3AudioVAE(nn.Module): *, device: torch.device, load_device: torch.device | None = None, + trust_remote_code: bool = False, ) -> None: super().__init__() self._device_target = device @@ -471,6 +630,7 @@ class MiniMaxH3AudioVAE(nn.Module): self.remote = _load_remote_component( component_path, self.config_dict, + trust_remote_code=trust_remote_code, ) # The checkpoint's audio VAE contract is FP32 for both reference # encoding and waveform decoding. diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index f21691eab..9c4a6eed6 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -189,6 +189,16 @@ def default_image_to_image_prompt( return result +def _always_preserve_reference_image_size( + *, + model: str | None, + revision: str | None = None, +) -> bool: + """Declare that a pipeline owns reference-image geometry.""" + del model, revision + return True + + _EXTRA_SPECS: dict[str, dict[str, Any]] = { "BagelPipeline": { "extra_body_params": BAGEL_EXTRA_BODY_PARAMS, @@ -230,6 +240,15 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { # Shared T2I/I2V envelopes select the output modality. LingBot's # pipeline owns model-specific validation and normalization. }, + **{ + model_class_name: { + "reference_image_size_resolver": _always_preserve_reference_image_size, + } + for model_class_name in ( + "MiniMaxH3Pipeline", + "MiniMaxH3ModularPipeline", + ) + }, **{ model_class_name: { "extra_body_params": LTX_EXTRA_BODY_PARAMS,