[Model][MiniMax-H3] Multi-GPU serving fixes and layer-wise offload component selection (#7047)

Signed-off-by: heyuanliu-intel <heyuan.liu@intel.com>

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

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

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,
+        )

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]

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)