Kandinsky6 (#8537)

Signed-off-by: leffff <levnovitskiy@gmail.com>
Signed-off-by: Artem Borisov <divotionchannel@gmail.com>
Co-authored-by: leffff <levnovitskiy@gmail.com>
Co-authored-by: Artem Borisov <divotionchannel@gmail.com>

vllm_omni/diffusion/cache/magcache/backend.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
 
 """
 MagCache backend implementation.
@@ -73,6 +73,19 @@ class MagCacheBackend(CacheBackend):
         transformer = pipeline.transformer
         transformer_type = transformer.__class__.__name__
 
+        if transformer_type == "Kandinsky6Transformer3DModel":
+            from vllm_omni.diffusion.models.kandinsky6.cache_accel import attach_mag_cache
+
+            attach_mag_cache(
+                transformer,
+                threshold=self.config.mag_threshold,
+                max_skip_steps=self.config.mag_max_skip_steps,
+                retention_ratio=self.config.mag_retention_ratio,
+            )
+            self._registered = True
+            logger.info("MagCache: Kandinsky 6 step-cache attached (visual/audio velocity reuse)")
+            return
+
         num_inference_steps = self.config.num_inference_steps or 28
 
         mag_ratios = self.config.mag_ratios

vllm_omni/diffusion/cache/magcache/hook.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
 
 """
 Hook-based MagCache implementation for vLLM-Omni.
@@ -543,9 +543,12 @@ def apply_mag_cache_hook(
     state_manager = StateManager(MagCacheState, (), {})
     remaining_blocks = []
 
+    allowed = getattr(module, "_magcache_block_attrs", None)
     for name, submodule in module.named_children():
         if not isinstance(submodule, torch.nn.ModuleList):
             continue
+        if allowed is not None and name not in allowed:
+            continue
         for index, block in enumerate(submodule):
             remaining_blocks.append((f"{name}.{index}", block))