Memory leak: `WanI2VCrossAttention._fp8_cache` accumulates across prompts when transformer object is reused

[xiaoghoul]

WanI2VCrossAttention._fp8_cache (and the equivalent in WanT2VCrossAttention) is a dict that caches cross-attention KV (text + image context) keyed by id(context). It is intended to be reused across denoising steps within a single prompt (since context is static across steps), but it is never cleared.

In the default workflow this is masked because loadmodel reconstructs the transformer object each run, so the dict starts fresh. But whenever the transformer object is reused across prompts (e.g. when keeping model weights resident in VRAM to avoid reloading ~17GB each run, or any scenario where ComfyUI's output cache reuses the WanVideoModelLoader output), the dict grows unbounded, leading to ~0.8GB VRAM leak per prompt and OOM after ~15 prompts.

Root cause

wanvideo/modules/model.py, class WanI2VCrossAttention (and WanT2VCrossAttention):

# __init__
self._fp8_cache = {}          # initialized once, never cleared

# forward
ctx_id = id(context)          # new context object each prompt -> new key
if ctx_id not in self._fp8_cache:
    k_txt = ...; v_txt = ...
    self._fp8_cache[ctx_id] = (k_txt.to(fp8), v_txt.to(fp8), k_txt.dtype)  # accumulates

There is no .clear() / reassignment anywhere except __init__. With 40 blocks times (text + image) entries per prompt, each prompt adds ~291 tensors (~0.8GB at 640x480 i2v) that are never released.

Reproduction

• Wan2.2 i2v 14B GGUF Q8, 640x480, 17 frames
• Keep the WanVideoModelLoader output cached across prompts (same model loader inputs, vary the sampler seed so the sampler actually re-runs)
• Run ~5 prompts, observe nvidia-smi baseline growing ~0.8GB each prompt with no convergence

How I tracked it down

I added a quick gc-scan in WanVideoSampler.process that walks gc.get_objects() and buckets all live CUDA tensors by element-count. The "millions of elements" bucket grew by ~291 tensors / ~0.8GB per prompt while the weight bucket (483 tensors / 16.5GB) stayed flat, so the leak is in per-block cross-attention caches, not weights or activations.

Workaround I tried

Clearing _fp8_cache at the start of each WanVideoSampler.process run fixes it for me. This keeps the within-prompt step-to-step reuse (which I believe is the cache's intended purpose) while stopping cross-prompt accumulation:

# In WanVideoSampler.process, before sampling begins:
for m in transformer.modules():
    if hasattr(m, "_fp8_cache") and isinstance(getattr(m, "_fp8_cache"), dict):
        m._fp8_cache.clear()

After this, 4 consecutive prompts (17f, varying seed): the "millions of elements" tensor bucket stayed flat at 324 tensors / 822MB (was growing +291/~0.8GB per prompt before), and torch.cuda.memory_allocated stable at 17.24GB. Not sure if this is the right place to put the clear, just sharing what worked for me.

Environment

• WanVideoWrapper (recent main)
• ComfyUI-MultiGPU (T5 on cuda:1, DiT resident on cuda:0)
• 2x Tesla V100 32GB, torch 2.7.1+cu126, GGUF Q8 i2v

Note

This looks like a latent bug in WanVideoWrapper (the code seems to assume the transformer object is not reused across prompts), not a ComfyUI or MultiGPU bug. It only shows up when the model is kept resident, which is probably why it's gone unnoticed: most setups reload each run and the dict gets reset implicitly. Also worth noting, the _fp8_cache name is a bit misleading since the accumulation happens regardless of whether fp8 is actually used (the dict entry is written on every new context even on non-fp8 paths).