feat: add KV caching support for Wan and VACE models

Author: Perseus14

src/maxdiffusion/configs/base_wan_14b.yml

@@ -355,6 +355,7 @@ use_cfg_cache: False
 # Batch positive and negative prompts in text encoder to save compute.
 use_batched_text_encoder: False
 
+use_kv_cache: False
 use_magcache: False
 magcache_thresh: 0.12
 magcache_K: 2

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -301,6 +301,7 @@ flow_shift: 3.0
 
 # Diffusion CFG cache (FasterCache-style, WAN 2.1 T2V only)
 use_cfg_cache: False
+use_kv_cache: False
 
 # Batch positive and negative prompts in text encoder to save compute.
 use_batched_text_encoder: False

src/maxdiffusion/configs/base_wan_27b.yml

@@ -331,7 +331,7 @@ use_cfg_cache: False
 # Batch positive and negative prompts in text encoder to save compute.
 use_batched_text_encoder: False
 
-
+use_kv_cache: False
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208) — skip forward pass
 # when predicted output change (based on accumulated latent/timestep drift) is small
 use_sen_cache: False

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -318,7 +318,7 @@ use_cfg_cache: False
 # Batch positive and negative prompts in text encoder to save compute.
 use_batched_text_encoder: False
 
-
+use_kv_cache: False
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208)
 use_sen_cache: False
 use_magcache: False

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -330,7 +330,7 @@ use_cfg_cache: False
 # Batch positive and negative prompts in text encoder to save compute.
 use_batched_text_encoder: False
 
-
+use_kv_cache: False
 # SenCache: Sensitivity-Aware Caching (arXiv:2602.24208)
 use_sen_cache: False
 

src/maxdiffusion/generate_wan.py

@@ -104,6 +104,7 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
           magcache_thresh=config.magcache_thresh,
           magcache_K=config.magcache_K,
           retention_ratio=config.retention_ratio,
+          use_kv_cache=config.use_kv_cache,
       )
     elif model_key == WAN2_2:
       return pipeline(
@@ -118,6 +119,7 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
           guidance_scale_high=config.guidance_scale_high,
           use_cfg_cache=config.use_cfg_cache,
           use_sen_cache=config.use_sen_cache,
+          use_kv_cache=config.use_kv_cache,
       )
     else:
       raise ValueError(f"Unsupported model_name for I2V in config: {model_key}")
@@ -136,6 +138,7 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
           magcache_thresh=config.magcache_thresh,
           magcache_K=config.magcache_K,
           retention_ratio=config.retention_ratio,
+          use_kv_cache=config.use_kv_cache,
       )
     elif model_key == WAN2_2:
       return pipeline(
@@ -149,6 +152,7 @@ def call_pipeline(config, pipeline, prompt, negative_prompt):
           guidance_scale_high=config.guidance_scale_high,
           use_cfg_cache=config.use_cfg_cache,
           use_sen_cache=config.use_sen_cache,
+          use_kv_cache=getattr(config, "use_kv_cache", False),
       )
     else:
       raise ValueError(f"Unsupported model_name for T2Vin config: {model_key}")