[MiniMax-H3] TeaCache support and Cache-DiT validation (#5840)

Signed-off-by: dpeng123 <diyapeng123@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/cache/test_cache_backends.py

@@ -471,6 +471,66 @@ class TestTeaCacheBackend:
         assert backend.enabled is True
         mock_apply_hook.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_enable_uses_generic_default_threshold(self, mock_apply_hook):
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "QwenImagePipeline"
+        pipeline.transformer = Mock()
+        pipeline.transformer.__class__.__name__ = "QwenImageTransformer2DModel"
+
+        TeaCacheBackend(DiffusionCacheConfig()).enable(pipeline)
+        assert mock_apply_hook.call_args.args[1].rel_l1_thresh == 0.2
+
+    @pytest.mark.parametrize("partition", ["fl2va", "combined"])
+    @pytest.mark.parametrize(
+        ("configured_threshold", "expected_threshold"),
+        [(None, 0.17), (0.2, 0.2)],
+    )
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_minimax_h3_only_enables_fl2va_teacache(
+        self,
+        mock_apply_hook,
+        partition,
+        configured_threshold,
+        expected_threshold,
+    ):
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
+        pipeline.partition = partition
+        pipeline.transformer = Mock()
+        pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
+        pipeline.transformers_ref = Mock()
+
+        config = (
+            DiffusionCacheConfig()
+            if configured_threshold is None
+            else DiffusionCacheConfig(rel_l1_thresh=configured_threshold)
+        )
+        backend = TeaCacheBackend(config)
+        backend.enable(pipeline)
+
+        mock_apply_hook.assert_called_once()
+        transformer, teacache_config = mock_apply_hook.call_args.args
+        assert transformer is pipeline.transformer
+        assert transformer is not pipeline.transformers_ref
+        assert teacache_config.rel_l1_thresh == expected_threshold
+
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_minimax_h3_rejects_ref2va_teacache(self, mock_apply_hook):
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
+        pipeline.partition = "ref2va"
+        pipeline.transformer = Mock()
+        pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
+
+        backend = TeaCacheBackend(DiffusionCacheConfig())
+
+        with pytest.raises(ValueError, match="only supports the MiniMax-H3 FL2VA partition"):
+            backend.enable(pipeline)
+
+        assert backend.enabled is False
+        mock_apply_hook.assert_not_called()
+
     @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
     def test_enable_with_coefficients(self, mock_apply_hook):
         """Test enabling TeaCache with custom coefficients."""