Add KV Cache for Autoregressive Inference

[DN6]

Autoregressive Diffusion Techniques such as Self Forcing rely on a rolling KV Cache across video frame chunks to transfer information from past context frames to the current frames being denoised.

This rolling KV Cache design (or variants similar to it) is likely to show up in other types of long video generation/ world models, so it would be good to see if we can support it natively in Diffusers.

Tasks

• Implement rolling KV Cache seen in Self Forcing using Diffusers' cache hooks design.
• Add a Modular Block to Wan Modular Pipelines that uses this rolling KV Cache to perform autoregressive inference.

[tolgacangoz]

I was just about opening this issue; it would be great!
@kuantuna, you might be interested in this since your MAGI-1 PR-continuation is related to KV-Caching too.

[sayakpaul]

Add a Modular Block to Wan Modular Pipelines that uses this rolling KV Cache to perform autoregressive inference.

@DN6 for clarity, to be able to perform inference, the one also needs a trained model right that supports rolling KV caching?

[11happy]

Hii @sayakpaul , I would like to work on this as a part of Diffusers MVP program.
Thank you : )

[11happy]

Hello @sayakpaul , Here's my proposal:
First Before moving to the implementation I wanted to make sure that I understood Rolling KV cache properly , I went through the linked paper for Self-Forcing https://arxiv.org/abs/2506.08009 & also one another which was referenced inside it for inspiration https://arxiv.org/abs/2309.17453 streaming language model with attention sink:

the core algorithm ( sink based ) looks like this:

1. compute queries, keys, value for the current step.
2. identify sink tokens (fixed)
3. if adding new tokens would exceed the kv cache size, we remove the oldest kv entries after the sink region , shift the remainig to left & insert the new tokens

I also went through the self forcing paper linked codebase & found this relavant code section
Link: https://github.com/guandeh17/Self-Forcing/blob/33593df3e81fa3ec10239271dd2c100facac6de1/wan/modules/causal_model.py#L206

frame_seqlen = math.prod(grid_sizes[0][1:]).item()
current_start_frame = current_start // frame_seqlen
roped_query = causal_rope_apply(
    q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
roped_key = causal_rope_apply(
    k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)

current_end = current_start + roped_query.shape[1]
sink_tokens = self.sink_size * frame_seqlen
# If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
kv_cache_size = kv_cache["k"].shape[1]
num_new_tokens = roped_query.shape[1]
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
        num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
    # Calculate the number of new tokens added in this step
    # Shift existing cache content left to discard oldest tokens
    # Clone the source slice to avoid overlapping memory error
    num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
    num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
    kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
        kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
    kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
        kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
    # Insert the new keys/values at the end
    local_end_index = kv_cache["local_end_index"].item() + current_end - \
        kv_cache["global_end_index"].item() - num_evicted_tokens
    local_start_index = local_end_index - num_new_tokens
    kv_cache["k"][:, local_start_index:local_end_index] = roped_key
    kv_cache["v"][:, local_start_index:local_end_index] = v

however rolling kv cache has mutiple possible variants- like naive sliding window, StreamingLLM style with sink tokens etc So do we want to implement a single recommended implementation or a configurable system supporting multiple variants.

Proposed implementation : hook based design
I went through the FirstBlockCache & FasterCache implementations based on hook designs, for Rolling KV Cache we can also follow a similar structure,
at high level It would require implementing RollingKVCacheState to store & manage kv tensors, an AttentionHook performing the core logic.

Also for a cleaner implementation I would like to work on Modular block for wan pipelines as a follow up PR.

Please let me know you thoughts and I will start working on PR accordingly.

Thank you : )

[sayakpaul]

@DN6 ^

[toilaluan]

I converted self-forcing checkpoint to diffusers format so it can be loaded directly by

from diffusers import WanTransformer3DModel

transformer = WanTransformer3DModel.from_pretrained("toilaluan/self-forcing-diffusers")

Code

import torch
import re

def convert_wan_a_to_b(state_dict_a):
    """
    Converts a checkpoint from Modeling A (WanModel) to Modeling B (WanTransformer3DModel).
    Assumes the checkpoint has 'model.' prefix and handles i2v/flf2v specifics if present.
    """
    state_dict_b = {}
    
    for key_a in state_dict_a:
        value = state_dict_a[key_a]
        
        # Remove 'model.' prefix if present
        if key_a.startswith('model.'):
            key = key_a[6:]
        else:
            key = key_a
        
        # Patch embedding
        if key == 'patch_embedding.weight':
            state_dict_b['patch_embedding.weight'] = value
        elif key == 'patch_embedding.bias':
            state_dict_b['patch_embedding.bias'] = value
        
        # Text embedding
        elif key == 'text_embedding.0.weight':
            state_dict_b['condition_embedder.text_embedder.linear_1.weight'] = value
        elif key == 'text_embedding.0.bias':
            state_dict_b['condition_embedder.text_embedder.linear_1.bias'] = value
        elif key == 'text_embedding.2.weight':
            state_dict_b['condition_embedder.text_embedder.linear_2.weight'] = value
        elif key == 'text_embedding.2.bias':
            state_dict_b['condition_embedder.text_embedder.linear_2.bias'] = value
        
        # Time embedding
        elif key == 'time_embedding.0.weight':
            state_dict_b['condition_embedder.time_embedder.linear_1.weight'] = value
        elif key == 'time_embedding.0.bias':
            state_dict_b['condition_embedder.time_embedder.linear_1.bias'] = value
        elif key == 'time_embedding.2.weight':
            state_dict_b['condition_embedder.time_embedder.linear_2.weight'] = value
        elif key == 'time_embedding.2.bias':
            state_dict_b['condition_embedder.time_embedder.linear_2.bias'] = value
        elif key == 'time_projection.1.weight':
            state_dict_b['condition_embedder.time_proj.weight'] = value
        elif key == 'time_projection.1.bias':
            state_dict_b['condition_embedder.time_proj.bias'] = value
        
        # Image embedding (if present, for i2v/flf2v)
        elif key == 'img_emb.proj.0.weight':
            state_dict_b['condition_embedder.image_embedder.norm1.weight'] = value
        elif key == 'img_emb.proj.0.bias':
            state_dict_b['condition_embedder.image_embedder.norm1.bias'] = value
        elif key == 'img_emb.proj.1.weight':
            state_dict_b['condition_embedder.image_embedder.ff.net.0.weight'] = value
        elif key == 'img_emb.proj.1.bias':
            state_dict_b['condition_embedder.image_embedder.ff.net.0.bias'] = value
        elif key == 'img_emb.proj.3.weight':
            state_dict_b['condition_embedder.image_embedder.ff.net.2.weight'] = value
        elif key == 'img_emb.proj.3.bias':
            state_dict_b['condition_embedder.image_embedder.ff.net.2.bias'] = value
        elif key == 'img_emb.proj.4.weight':
            state_dict_b['condition_embedder.image_embedder.norm2.weight'] = value
        elif key == 'img_emb.proj.4.bias':
            state_dict_b['condition_embedder.image_embedder.norm2.bias'] = value
        elif key == 'img_emb.emb_pos':
            state_dict_b['condition_embedder.image_embedder.pos_embed'] = value
        
        # Head / Output
        elif key == 'head.norm.weight':
            state_dict_b['norm_out.weight'] = value
        elif key == 'head.norm.bias':
            state_dict_b['norm_out.bias'] = value
        elif key == 'head.head.weight':
            state_dict_b['proj_out.weight'] = value
        elif key == 'head.head.bias':
            state_dict_b['proj_out.bias'] = value
        elif key == 'head.modulation':
            state_dict_b['scale_shift_table'] = value
        
        # Blocks
        else:
            # Parse block index
            match = re.match(r'blocks\.(\d+)\.(.*)', key)
            if match:
                block_idx = match.group(1)
                subkey = match.group(2)
                prefix_b = f'blocks.{block_idx}.'
                
                # Modulation / scale_shift_table
                if subkey == 'modulation':
                    state_dict_b[prefix_b + 'scale_shift_table'] = value
                
                # Self attention
                elif subkey == 'self_attn.q.weight':
                    state_dict_b[prefix_b + 'attn1.to_q.weight'] = value
                elif subkey == 'self_attn.q.bias':
                    state_dict_b[prefix_b + 'attn1.to_q.bias'] = value
                elif subkey == 'self_attn.k.weight':
                    state_dict_b[prefix_b + 'attn1.to_k.weight'] = value
                elif subkey == 'self_attn.k.bias':
                    state_dict_b[prefix_b + 'attn1.to_k.bias'] = value
                elif subkey == 'self_attn.v.weight':
                    state_dict_b[prefix_b + 'attn1.to_v.weight'] = value
                elif subkey == 'self_attn.v.bias':
                    state_dict_b[prefix_b + 'attn1.to_v.bias'] = value
                elif subkey == 'self_attn.o.weight':
                    state_dict_b[prefix_b + 'attn1.to_out.0.weight'] = value
                elif subkey == 'self_attn.o.bias':
                    state_dict_b[prefix_b + 'attn1.to_out.0.bias'] = value
                elif subkey == 'self_attn.norm_q.weight':
                    state_dict_b[prefix_b + 'attn1.norm_q.weight'] = value
                elif subkey == 'self_attn.norm_k.weight':
                    state_dict_b[prefix_b + 'attn1.norm_k.weight'] = value
                
                # Cross attention
                elif subkey == 'cross_attn.q.weight':
                    state_dict_b[prefix_b + 'attn2.to_q.weight'] = value
                elif subkey == 'cross_attn.q.bias':
                    state_dict_b[prefix_b + 'attn2.to_q.bias'] = value
                elif subkey == 'cross_attn.k.weight':
                    state_dict_b[prefix_b + 'attn2.to_k.weight'] = value
                elif subkey == 'cross_attn.k.bias':
                    state_dict_b[prefix_b + 'attn2.to_k.bias'] = value
                elif subkey == 'cross_attn.v.weight':
                    state_dict_b[prefix_b + 'attn2.to_v.weight'] = value
                elif subkey == 'cross_attn.v.bias':
                    state_dict_b[prefix_b + 'attn2.to_v.bias'] = value
                elif subkey == 'cross_attn.o.weight':
                    state_dict_b[prefix_b + 'attn2.to_out.0.weight'] = value
                elif subkey == 'cross_attn.o.bias':
                    state_dict_b[prefix_b + 'attn2.to_out.0.bias'] = value
                elif subkey == 'cross_attn.norm_q.weight':
                    state_dict_b[prefix_b + 'attn2.norm_q.weight'] = value
                elif subkey == 'cross_attn.norm_k.weight':
                    state_dict_b[prefix_b + 'attn2.norm_k.weight'] = value
                elif subkey == 'cross_attn.k_img.weight':
                    state_dict_b[prefix_b + 'attn2.add_k_proj.weight'] = value
                elif subkey == 'cross_attn.k_img.bias':
                    state_dict_b[prefix_b + 'attn2.add_k_proj.bias'] = value
                elif subkey == 'cross_attn.v_img.weight':
                    state_dict_b[prefix_b + 'attn2.add_v_proj.weight'] = value
                elif subkey == 'cross_attn.v_img.bias':
                    state_dict_b[prefix_b + 'attn2.add_v_proj.bias'] = value
                elif subkey == 'cross_attn.norm_k_img.weight':
                    state_dict_b[prefix_b + 'attn2.norm_added_k.weight'] = value
                
                # Norms
                elif subkey == 'norm1.weight':
                    state_dict_b[prefix_b + 'norm1.weight'] = value
                elif subkey == 'norm1.bias':
                    state_dict_b[prefix_b + 'norm1.bias'] = value
                elif subkey == 'norm3.weight':
                    state_dict_b[prefix_b + 'norm2.weight'] = value
                elif subkey == 'norm3.bias':
                    state_dict_b[prefix_b + 'norm2.bias'] = value
                elif subkey == 'norm2.weight':
                    state_dict_b[prefix_b + 'norm3.weight'] = value
                elif subkey == 'norm2.bias':
                    state_dict_b[prefix_b + 'norm3.bias'] = value
                
                # FFN
                elif subkey == 'ffn.0.weight':
                    state_dict_b[prefix_b + 'ffn.net.0.proj.weight'] = value
                elif subkey == 'ffn.0.bias':
                    state_dict_b[prefix_b + 'ffn.net.0.proj.bias'] = value
                elif subkey == 'ffn.2.weight':
                    state_dict_b[prefix_b + 'ffn.net.2.weight'] = value
                elif subkey == 'ffn.2.bias':
                    state_dict_b[prefix_b + 'ffn.net.2.bias'] = value
        
    return state_dict_b

# Example usage:
# state_dict_a = torch.load('model_a.ckpt', map_location='cpu')
# state_dict_b = convert_wan_a_to_b(state_dict_a)
# torch.save(state_dict_b, 'model_b.ckpt')

[11happy]

Hii @DN6, would like to hear your thoughts on the proposal & questions.
Thank you : )

[DN6]

Hi @11happy I think we have decided to handle this one internally since it might require some changes to the design of the cache system. Apologies for inconvenience. Would you be okay to consider any of the other issues listed in the program?

[11happy]

Hi @11happy I think we have decided to handle this one internally since it might require some changes to the design of the cache system. Apologies for inconvenience. Would you be okay to consider any of the other issues listed in the program?

Sure, no worries, I took a look at the listed issues & issue #11966 seems interesting to me, I will take a deeper look at the issue & get back with the proposal. I will just do a comment there as well to make it easier to track, Thank you : )