[Model][Feat] Support Minimax-H3 quality grading requests through dynamic loading/unloading with Cache-DiT (#5853)

Signed-off-by: mglyn <1203789601@qq.com>
2026-08-07 09:07:23 +00:00
1 parent e906f41dc8
commit d8c35663b4
@@ -93,8 +93,8 @@ class TestCacheDiTBackend:

    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
-    def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
-        """Combined pipelines enable and refresh Cache-DiT for every DiT."""
+    def test_lifecycle_applies_to_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
+        """Combined pipelines manage every exact Cache-DiT target."""
        pipeline = Mock()
        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
        pipeline._dit_modules = ["transformer", "transformers_ref"]
@@ -104,18 +104,24 @@ class TestCacheDiTBackend:
            transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
                block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
            )
+        installed_adapters = [Mock(), Mock()]
+        mock_block_adapter.side_effect = installed_adapters

        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
        backend.enable(pipeline)
        backend.refresh(pipeline, num_inference_steps=20)
+        backend.disable(pipeline)

-        assert len(backend._refresh_funcs) == 2
        assert mock_cache_dit.enable_cache.call_count == 2
        assert mock_cache_dit.refresh_context.call_count == 2
        assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == {
            pipeline.transformer,
            pipeline.transformers_ref,
        }
+        assert [call.args[0] for call in mock_cache_dit.disable_cache.call_args_list] == installed_adapters
+        assert backend._refresh_funcs == []
+        assert backend._cache_targets == []
+        assert not backend.is_enabled()

    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
    def test_summary_skips_uncached_nested_dit(self, mock_cache_dit):
@@ -291,6 +297,32 @@ class TestCacheDiTBackend:

        assert mock_cache_dit.refresh_context.call_count == 2

+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_refresh_rearms_force_refresh_hint(self, mock_cache_dit, mock_block_adapter):
+        """A once hint must be restored for every repeated request."""
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "DiTPipeline"
+        pipeline.transformer = Mock()
+        pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={"layers": ForwardPattern.Pattern_0}
+        )
+
+        config = DiffusionCacheConfig(
+            force_refresh_step_hint=7,
+            force_refresh_step_policy="once",
+        )
+        backend = CacheDiTBackend(config)
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=20)
+        backend.refresh(pipeline, num_inference_steps=20)
+
+        assert mock_cache_dit.refresh_context.call_count == 2
+        for call in mock_cache_dit.refresh_context.call_args_list:
+            refresh_config = call.kwargs["cache_config"]
+            assert refresh_config.force_refresh_step_hint == 7
+            assert refresh_config.force_refresh_step_policy == "once"
+
    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
    def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
@@ -81,6 +81,8 @@ def test_cachedit_public_api_is_explicit():
        "CacheDiTAdapterConfig",
        "CacheDiTBackend",
        "CacheDiTConfig",
+        "CacheDiTRequestSpec",
+        "RequestScopedCacheDiTRuntime",
        "SensenovaCachedAdapter",
        "cache_summary",
        "enable_cache_for_dit",
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.cache.cachedit import runtime as runtime_module
+from vllm_omni.diffusion.cache.cachedit.runtime import (
+    CacheDiTRequestSpec,
+    RequestScopedCacheDiTRuntime,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _spec(key: str, steps: int = 50) -> CacheDiTRequestSpec:
+    return CacheDiTRequestSpec(
+        installation_key=key,
+        cache_config=DiffusionCacheConfig(),
+        num_inference_steps=steps,
+    )
+
+
+def _recording_backend(events):
+    class FakeBackend:
+        def __init__(self, config):
+            events.append(("create", config))
+            self.enabled = False
+
+        def enable(self, pipeline):
+            events.append(("enable", pipeline))
+            self.enabled = True
+
+        def refresh(self, pipeline, num_inference_steps):
+            events.append(("refresh", pipeline, num_inference_steps))
+
+        def disable(self, pipeline):
+            events.append(("disable", pipeline))
+            self.enabled = False
+
+        def is_enabled(self):
+            return self.enabled
+
+    return FakeBackend
+
+
+def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch):
+    events = []
+    FakeBackend = _recording_backend(events)
+    monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend)
+    pipeline = object()
+    config = DiffusionCacheConfig()
+    runtime = RequestScopedCacheDiTRuntime(pipeline)
+    startup_backend = FakeBackend(config)
+    startup_backend.enabled = True
+
+    runtime.adopt(startup_backend, installation_key="generic")
+    runtime.prepare(_spec("generic", 50))  # omitted: keep startup cache
+    runtime.prepare(None)  # lossless: uninstall
+    runtime.prepare(None)  # repeated lossless: no-op
+    runtime.prepare(_spec("high", 40))  # high: install H3 profile
+    runtime.prepare(_spec("high", 30))  # repeated high: refresh only
+    runtime.prepare(_spec("generic", 20))  # omitted: restore startup profile
+
+    assert [event[0] for event in events] == [
+        "create",
+        "refresh",
+        "disable",
+        "create",
+        "enable",
+        "refresh",
+        "refresh",
+        "disable",
+        "create",
+        "enable",
+        "refresh",
+    ]
+    assert events[1][-1] == 50
+    assert events[5][-1] == 40
+    assert events[6][-1] == 30
+    assert events[10][-1] == 20
+    assert runtime.is_enabled
+
+
+def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch):
+    events = []
+    monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events))
+    pipeline = object()
+    runtime = RequestScopedCacheDiTRuntime(pipeline)
+
+    runtime.prepare(_spec("high", 40))
+    runtime.prepare(None)
+
+    assert [event[0] for event in events] == [
+        "create",
+        "enable",
+        "refresh",
+        "disable",
+    ]
+    assert events[2][-1] == 40
+    assert not runtime.is_enabled
@@ -13,6 +13,70 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]


+def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "fl2va"
+    pipeline.supported_tasks = frozenset({"t2va"})
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config = SimpleNamespace()
+    cache_spec = object()
+    quality_plan = SimpleNamespace(cache_dit=cache_spec)
+    pipeline._quality_policy = Mock()
+    pipeline._quality_policy.resolve.return_value = quality_plan
+    events = []
+    pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec)))
+    pipeline.encode_prompt = Mock(
+        return_value=(
+            torch.ones(1, 2),
+            torch.ones(1, dtype=torch.long),
+        )
+    )
+
+    def diffuse(**kwargs):
+        events.append(("diffuse", kwargs))
+        return torch.zeros(1), torch.zeros(1)
+
+    pipeline.diffuse = diffuse
+    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    sampling = OmniDiffusionSamplingParams(
+        quality="high",
+        width=1344,
+        height=768,
+        fps=24,
+        num_frames=124,
+        num_inference_steps=50,
+        extra_args={"task": "t2va", "aspect_ratio": "16:9"},
+    )
+    batch = DiffusionRequestBatch(
+        [
+            OmniDiffusionRequest(
+                prompt="quality boundary",
+                sampling_params=sampling,
+                request_id="quality-boundary",
+            )
+        ]
+    )
+
+    output = pipeline.forward(batch)
+
+    assert events[0] == ("prepare", cache_spec)
+    assert events[1][0] == "diffuse"
+    pipeline._quality_policy.resolve.assert_called_once_with(
+        quality="high",
+        num_inference_steps=50,
+        extra_args={"task": "t2va", "aspect_ratio": "16:9"},
+    )
+    assert output.output == pipeline.decode.return_value
+
+
 def test_pipeline_import_registry_and_component_discovery():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.registry import (
@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.errors import OmniClientError
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _quality_od_config(*, cache_backend: str = "cache_dit", **overrides):
+    values = {
+        "cache_backend": cache_backend,
+        "cache_config": DiffusionCacheConfig(),
+        "diffusion_attention_config": SimpleNamespace(default=None, per_role={}),
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": False,
+        "enforce_eager": True,
+        "num_gpus": 4,
+        "quantization_config": None,
+        "parallel_config": SimpleNamespace(
+            allgather_degree=1,
+            cfg_parallel_size=1,
+            data_parallel_size=1,
+            pipeline_parallel_size=1,
+            ring_degree=1,
+            sequence_parallel_size=4,
+            tensor_parallel_size=1,
+            ulysses_degree=4,
+            use_hsdp=False,
+        ),
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _resolve_quality(policy, **overrides):
+    values = {
+        "quality": "high",
+        "num_inference_steps": 50,
+        "extra_args": None,
+    }
+    values.update(overrides)
+    return policy.resolve(**values)
+
+
+def test_high_quality_policy_emits_cache_profile_without_deployment_gates():
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
+        MiniMaxH3QualityPolicy,
+    )
+
+    policy = MiniMaxH3QualityPolicy(
+        _quality_od_config(
+            cache_backend="none",
+            enable_layerwise_offload=True,
+            enforce_eager=False,
+            num_gpus=1,
+        )
+    )
+
+    plan = _resolve_quality(policy, num_inference_steps=17)
+    spec = plan.cache_dit
+
+    assert spec is not None
+    assert spec.installation_key == "minimax_h3.high"
+    assert spec.num_inference_steps == 17
+    assert spec.cache_config.Fn_compute_blocks == 1
+    assert spec.cache_config.Bn_compute_blocks == 0
+    assert spec.cache_config.max_warmup_steps == 4
+    assert spec.cache_config.residual_diff_threshold == pytest.approx(0.04)
+    assert spec.cache_config.max_continuous_cached_steps == 1
+    assert spec.cache_config.enable_taylorseer is False
+    assert spec.cache_config.scm_steps_mask_policy is None
+
+
+def test_h3_force_refresh_is_model_owned_and_request_scoped():
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    policy = MiniMaxH3QualityPolicy(_quality_od_config())
+    plan = _resolve_quality(
+        policy,
+        quality="high",
+        num_inference_steps=17,
+        extra_args={"force_refresh_step_hint": 7},
+    )
+
+    assert plan.cache_dit is not None
+    assert plan.cache_dit.installation_key == "minimax_h3.high:force_refresh=7:once"
+    assert plan.cache_dit.cache_config.force_refresh_step_hint == 7
+    assert plan.cache_dit.cache_config.force_refresh_step_policy == "once"
+
+
+def test_h3_force_refresh_policy_supports_repeat_and_does_not_mutate_generic_config():
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    config = _quality_od_config()
+    plan = _resolve_quality(
+        MiniMaxH3QualityPolicy(config),
+        quality=None,
+        num_inference_steps=20,
+        extra_args={"force_refresh_step_hint": 5, "force_refresh_step_policy": "repeat"},
+    )
+
+    assert plan.cache_dit is not None
+    assert plan.cache_dit.installation_key == "minimax_h3.generic:force_refresh=5:repeat"
+    assert plan.cache_dit.cache_config is not config.cache_config
+    assert plan.cache_dit.cache_config.force_refresh_step_hint == 5
+    assert plan.cache_dit.cache_config.force_refresh_step_policy == "repeat"
+    assert config.cache_config.force_refresh_step_hint is None
+    assert config.cache_config.force_refresh_step_policy == "once"
+
+    restored = _resolve_quality(MiniMaxH3QualityPolicy(config), quality=None)
+    assert restored.cache_dit is not None
+    assert restored.cache_dit.installation_key == "minimax_h3.generic"
+    assert restored.cache_dit.cache_config is config.cache_config
+
+
+@pytest.mark.parametrize(
+    ("extra_args", "message"),
+    [
+        ({"force_refresh_step_hint": True}, "positive integer"),
+        ({"force_refresh_step_hint": 0}, "between 1"),
+        ({"force_refresh_step_hint": 51}, "between 1"),
+        ({"force_refresh_step_hint": 5, "force_refresh_step_policy": "always"}, "one of"),
+        ({"force_refresh_step_policy": "repeat"}, "requires force_refresh_step_hint"),
+    ],
+)
+def test_h3_force_refresh_rejects_invalid_values(extra_args, message):
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    with pytest.raises(OmniClientError, match=message):
+        _resolve_quality(
+            MiniMaxH3QualityPolicy(_quality_od_config()),
+            extra_args=extra_args,
+        )
+
+
+@pytest.mark.parametrize(
+    ("cache_backend", "quality"),
+    [
+        ("cache_dit", "lossless"),
+        ("none", None),
+    ],
+)
+def test_h3_force_refresh_requires_active_cache_target(cache_backend, quality):
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
+
+    policy = MiniMaxH3QualityPolicy(_quality_od_config(cache_backend=cache_backend))
+
+    with pytest.raises(OmniClientError, match="require an active Cache-DiT"):
+        _resolve_quality(
+            policy,
+            quality=quality,
+            extra_args={"force_refresh_step_hint": 7},
+        )
+
+
+@pytest.mark.parametrize(
+    ("cache_backend", "quality", "expected_key"),
+    [
+        ("cache_dit", None, "minimax_h3.generic"),
+        ("cache_dit", "lossless", None),
+        ("cache_dit", "high", "minimax_h3.high"),
+        ("none", None, None),
+        ("none", "lossless", None),
+        ("none", "high", "minimax_h3.high"),
+    ],
+)
+def test_model_policy_owns_request_cache_target(cache_backend, quality, expected_key):
+    from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
+        MiniMaxH3QualityPolicy,
+    )
+
+    config = _quality_od_config(cache_backend=cache_backend)
+    policy = MiniMaxH3QualityPolicy(config)
+
+    plan = _resolve_quality(
+        policy,
+        quality=quality,
+    )
+
+    if expected_key is None:
+        assert plan.cache_dit is None
+    else:
+        assert plan.cache_dit is not None
+        assert plan.cache_dit.installation_key == expected_key
+        if quality is None:
+            assert plan.cache_dit.cache_config is config.cache_config
+
+
+def test_h3_adopts_runner_installed_cache_dit_backend():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    backend = object()
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline._cache_dit_runtime = SimpleNamespace(
+        adopt=lambda adopted, *, installation_key: setattr(
+            pipeline,
+            "adopted",
+            (adopted, installation_key),
+        ),
+        is_enabled=True,
+    )
+
+    pipeline.adopt_cache_dit_backend(backend)
+
+    assert pipeline.adopted == (backend, "minimax_h3.generic")
+    assert pipeline.is_cache_dit_enabled()
@@ -162,6 +162,21 @@ def _fake_platform_for_peak_memory():
    )


+def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
+    events = []
+    plain_pipeline = object()
+    request_scoped_pipeline = SimpleNamespace(
+        adopt_cache_dit_backend=lambda backend: events.append(backend),
+        is_cache_dit_enabled=lambda: True,
+    )
+    backend = object()
+
+    assert not model_runner_module.adopt_request_scoped_cache_dit(plain_pipeline, backend)
+    assert model_runner_module.adopt_request_scoped_cache_dit(request_scoped_pipeline, backend)
+    assert model_runner_module.is_request_scoped_cache_dit_enabled(request_scoped_pipeline)
+    assert events == [backend]
+
+
 def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
    runner = object.__new__(DiffusionModelRunner)
    runner.vllm_config = object()
@@ -419,6 +434,33 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
    assert cache_backend.refresh_calls == [(runner.pipeline, 4, True)]


+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("cache_dit_enabled", [False, True])
+def test_execute_model_cache_summary_follows_pipeline_owned_cache_dit_state(monkeypatch, cache_dit_enabled):
+    runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
+    runner.pipeline.adopt_cache_dit_backend = lambda backend: None
+    runner.pipeline.is_cache_dit_enabled = lambda: cache_dit_enabled
+    req = _make_request()
+
+    cache_summary_calls = []
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    monkeypatch.setattr(
+        model_runner_module,
+        "cache_summary",
+        lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
+    )
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert cache_summary_calls == ([(runner.pipeline, True)] if cache_dit_enabled else [])
+
+
@pytest.mark.core_model
@pytest.mark.cpu
 def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch):
@@ -380,6 +380,63 @@ class TestRequestScheduler:
        assert first.num_running_reqs == 1
        assert first.num_waiting_reqs == 1

+    def test_batches_different_quality_levels_separately(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        high = scheduler.add_request(
+            _make_step_request(
+                "high",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    quality="high",
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "lossless",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    quality="lossless",
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [high]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+    def test_batches_omitted_and_explicit_lossless_separately(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        omitted = scheduler.add_request(
+            _make_step_request(
+                "omitted",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "explicit",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=2,
+                    quality="lossless",
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [omitted]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
    def test_batches_different_request_local_seed_together(self) -> None:
        scheduler = RequestScheduler()
        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -930,6 +987,35 @@ class TestStepScheduler:
        assert second.num_running_reqs == 1
        assert second.num_waiting_reqs == 1

+    def test_step_batch_rejects_different_quality_levels(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        high = scheduler.add_request(
+            _make_step_request(
+                "high",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                    quality="high",
+                ),
+            )
+        )
+        scheduler.add_request(
+            _make_step_request(
+                "lossless",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                    quality="lossless",
+                ),
+            )
+        )
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [high]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
    def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
        """Multiple requests with the same LoRA (id + scale) co-batch."""
        from vllm_omni.lora.request import LoRARequest