Skip to content

Add KV Cache for Autoregressive Inference #12600

Open
Open
Add KV Cache for Autoregressive Inference#12600
Assignees
zucchini-nlp
Labels
advancedconsider-for-modular-diffusersThings to consider adding support for in Modular Diffusers (with the help of community)contributions-welcomemodular-diffusersstaleIssues that haven't received updates

Description

@DN6
DN6
opened on Nov 6, 2025

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.

Activity

  1. DN6
    added
    contributions-welcome
    advanced
    consider-for-modular-diffusersThings to consider adding support for in Modular Diffusers (with the help of community)
    modular-diffusers
    on Nov 6, 2025

  2. tolgacangoz commented on Nov 6, 2025

    @tolgacangoz
    tolgacangoz
    on Nov 6, 2025
    Contributor

    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.

  3. sayakpaul commented on Nov 6, 2025

    @sayakpaul
    sayakpaul
    on Nov 6, 2025
    Member

    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?

  4. 11happy commented on Nov 12, 2025

    @11happy
    11happy
    on Nov 12, 2025

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

  5. 11happy commented on Nov 16, 2025

    @11happy
    11happy
    on Nov 16, 2025

    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 : )

  6. sayakpaul commented on Nov 16, 2025

    @sayakpaul
    sayakpaul
    on Nov 16, 2025
    Member

    @DN6 ^

  7. toilaluan commented on Nov 18, 2025

    @toilaluan
    toilaluan
    on Nov 18, 2025
    Contributor

    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

    Details 
    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')

  8. 11happy commented on Nov 19, 2025

    @11happy
    11happy
    on Nov 19, 2025

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

  9. DN6 commented on Nov 19, 2025

    @DN6
    DN6
    on Nov 19, 2025
    CollaboratorAuthor

    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?

  10. 11happy commented on Nov 19, 2025

    @11happy
    11happy
    on Nov 19, 2025

    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 : )

  11. DN6
    assigned
    zucchini-nlp
    on Nov 21, 2025

  12. github-actions commented on Jan 9, 2026

    @github-actions
    github-actionsBot
    on Jan 9, 2026 – with GitHub Actions
    Contributor

    This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

    Please note that issues that do not follow the contributing guidelines are likely to be ignored.

  13. github-actions
    added
    staleIssues that haven't received updates
    on Jan 9, 2026
  14. yiyixuxu
    removed
    staleIssues that haven't received updates
    on Jan 10, 2026

  15. github-actions commented on Feb 3, 2026

    @github-actions
    github-actionsBot
    on Feb 3, 2026 – with GitHub Actions
    Contributor

    This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

    Please note that issues that do not follow the contributing guidelines are likely to be ignored.

  16. github-actions
    added
    staleIssues that haven't received updates
    on Feb 3, 2026

  17. gueraf commented on May 1, 2026

    @gueraf
    gueraf
    on May 1, 2026

    Hi,

    is there appetite for reviving this line of work?

    The motivation from our side is that autoregressive DiTs (with KV caching) are a clear trend in world models, and we'd like to leverage the HF diffusers ecosystem to ship models to our partners in a familiar framework with minimal custom code.

    I extracted/rebuilt an internal implementation we have at Odyssey as #13569 as an MVP similar to #12773 but a bit tighter in scope.

    Please let me know what you think.

    Cheers,
    Fabian

  18. sayakpaul commented on May 1, 2026

    @sayakpaul
    sayakpaul
    on May 1, 2026
    Member

    I think the models with KVcaching we have added to the core codebase (KleinKV, for example) implement their own KVcache class.

    Once we start seeing some patterns, we could try to uniformize some of the abstractions.

  19. gueraf commented on May 1, 2026

    @gueraf
    gueraf
    on May 1, 2026

    OK, I can definitely make it Wan-specific explicitly. In reality most people work on top of Wan or Wan-like architectures in any way, e.g., Self forcing and the many methods inspired by it, some recent work from Krea, etc.

    Right now the PR does a bit of a ducktyping dance around generic-ish aspects, but in perfect honesty it does replicate Wan rope and assumes Wan-like norm placement.

    I'll send an updated PR early next week if that's ok. Appreciate your help and general openness to this :)

  20. gueraf commented on May 6, 2026

    @gueraf
    gueraf
    on May 6, 2026

    Sorry for the delay. I sent out #13681. Thanks @sayakpaul for your help!

    cc @nicolasgriffiths @jeff-hawke

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

advancedconsider-for-modular-diffusersThings to consider adding support for in Modular Diffusers (with the help of community)contributions-welcomemodular-diffusersstaleIssues that haven't received updates

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions