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

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/cache/test_cache_backends.py

@@ -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):

tests/diffusion/cache/test_cache_dit.py

@@ -81,6 +81,8 @@ def test_cachedit_public_api_is_explicit():
         "CacheDiTAdapterConfig",
         "CacheDiTBackend",
         "CacheDiTConfig",
+        "CacheDiTRequestSpec",
+        "RequestScopedCacheDiTRuntime",
         "SensenovaCachedAdapter",
         "cache_summary",
         "enable_cache_for_dit",

tests/diffusion/cache/test_cache_dit_request_runtime.py

@@ -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