vllm-omni/.claude/skills/add-diffusion-model/references/cache-dit-patterns.md

Cache-DiT Patterns Reference

Overview

Cache-DiT accelerates Diffusion Transformers by caching intermediate computation results across denoising steps. Adjacent steps produce similar features, so redundant computations can be skipped.

Three caching strategies:

• DBCache: Dynamic block-level caching — selectively computes or caches transformer blocks based on residual differences
• TaylorSeer: Calibration-based prediction using Taylor expansion to estimate block outputs
• SCM (Step Computation Masking): Dynamic step skipping based on configurable policies

Typical speedup: 1.5-2.5x depending on model and configuration.

Official docs: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit

Architecture

vLLM-Omni integrates cache-dit through CacheDiTBackend:

Component	Purpose
CacheDiTBackend	Unified backend — auto-selects enabler (standard or custom)
enable_cache_for_dit()	Default enabler for standard single-transformer models
CUSTOM_DIT_ENABLERS dict	Registry of custom enablers keyed by pipeline class name
BlockAdapter	Wraps complex architectures (multi-block-list or multi-transformer)
ForwardPattern	Specifies block forward signature: Pattern_0, Pattern_1, Pattern_2
ParamsModifier	Per-transformer or per-block-list config customization
DBCacheConfig	Configuration for DBCache parameters
cache_dit.refresh_context()	Updates cache context when num_inference_steps changes

Source files:

• vllm_omni/diffusion/cache/cache_dit_backend.py — CacheDiTBackend, enablers, CUSTOM_DIT_ENABLERS
• vllm_omni/diffusion/cache/ — cache backend implementations

Standard Models: Automatic Support

Most DiT models follow this pattern:

• Single transformer with one nn.ModuleList of blocks
• Standard forward signature
• Compatible with cache-dit's automatic detection

Examples: Qwen-Image, Z-Image, FLUX

No code changes needed. CacheDiTBackend automatically uses enable_cache_for_dit():

from vllm_omni import Omni

omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="cache_dit",
    cache_config={
        "Fn_compute_blocks": 1,
        "Bn_compute_blocks": 0,
        "max_warmup_steps": 4,
    }
)

What happens automatically:

def enable_cache_for_dit(pipeline, cache_config):
    db_cache_config = DBCacheConfig(
        num_inference_steps=None,
        Fn_compute_blocks=cache_config.Fn_compute_blocks,
        Bn_compute_blocks=cache_config.Bn_compute_blocks,
        max_warmup_steps=cache_config.max_warmup_steps,
        max_cached_steps=cache_config.max_cached_steps,
        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
        residual_diff_threshold=cache_config.residual_diff_threshold,
    )

    cache_dit.enable_cache(pipeline.transformer, cache_config=db_cache_config)

    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
        cache_dit.refresh_context(
            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
        )
    return refresh_cache_context

Custom Architectures: Writing Custom Enablers

When you need a custom enabler

• Model has multiple block lists in one transformer (e.g., transformer_blocks + single_transformer_blocks)
• Model has two transformers (e.g., high-noise + low-noise like Wan2.2)
• Model uses non-standard block forward signature

Pattern 1: Multi-Block-List (LongCat-Image style)

Single transformer with two block lists:

import cache_dit
from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig

def enable_cache_for_your_model(pipeline, cache_config):
    db_cache_config = DBCacheConfig(
        num_inference_steps=None,
        Fn_compute_blocks=cache_config.Fn_compute_blocks,
        Bn_compute_blocks=cache_config.Bn_compute_blocks,
        max_warmup_steps=cache_config.max_warmup_steps,
        max_cached_steps=cache_config.max_cached_steps,
        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
        residual_diff_threshold=cache_config.residual_diff_threshold,
    )

    cache_dit.enable_cache(
        BlockAdapter(
            transformer=pipeline.transformer,
            blocks=[
                pipeline.transformer.transformer_blocks,
                pipeline.transformer.single_transformer_blocks,
            ],
            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
            params_modifiers=[ParamsModifier(...)],
        ),
        cache_config=db_cache_config,
    )

    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
        cache_dit.refresh_context(
            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
        )
    return refresh_cache_context

For single transformer with multiple block lists, refresh_context works the same as standard models — call it once on the transformer.

Pattern 2: Dual-Transformer (Wan2.2 style)

Two transformers with separate configs:

def enable_cache_for_dual_transformer(pipeline, cache_config):
    db_cache_config = DBCacheConfig(...)

    cache_dit.enable_cache(
        BlockAdapter(
            transformer=[pipeline.transformer, pipeline.transformer_2],
            blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
            forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
            params_modifiers=[
                ParamsModifier(...),  # Config for transformer 1
                ParamsModifier(...),  # Config for transformer 2
            ],
        ),
        cache_config=db_cache_config,
    )

    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
        high_steps, low_steps = _split_inference_steps(num_inference_steps)
        cache_dit.refresh_context(
            pipeline.transformer, num_inference_steps=high_steps, verbose=verbose
        )
        cache_dit.refresh_context(
            pipeline.transformer_2, num_inference_steps=low_steps, verbose=verbose
        )
    return refresh_cache_context

Key difference: refresh_context must be called on each transformer separately with its own step count.

Choosing the ForwardPattern

Pattern	Block forward signature	Example models
Pattern_0	block(hidden_states, **kwargs) → residual added inside block	Default
Pattern_1	block(hidden_states, **kwargs) → returns (hidden_states, ...) tuple	FLUX-style single blocks
Pattern_2	block(hidden_states, **kwargs) → (hidden_states, ...) with different residual pattern	Wan2.2 blocks

Inspect your block's forward() return type and residual connection pattern to choose the right one. See Cache-DiT API Reference for details.

Registering Custom Enablers

Add your enabler to CUSTOM_DIT_ENABLERS in vllm_omni/diffusion/cache/cache_dit_backend.py:

CUSTOM_DIT_ENABLERS = {
    "Wan22Pipeline": enable_cache_for_wan22,
    "LongCatImagePipeline": enable_cache_for_longcat_image,
    "YourModelPipeline": enable_cache_for_your_model,
}

The key must match pipeline.__class__.__name__.

Configuration Parameters

Parameter	Default	Description
Fn_compute_blocks	1	Number of blocks to always compute at the front
Bn_compute_blocks	0	Number of blocks to always compute at the back
max_warmup_steps	4	Steps to run without caching at the beginning
max_cached_steps	—	Max total cached steps
max_continuous_cached_steps	—	Max consecutive cached steps
residual_diff_threshold	0.24	Threshold for deciding whether to cache a block

Tuning for quality vs speed

Goal	Adjustments
More speed, acceptable quality loss	Higher residual_diff_threshold (0.24-0.4), lower max_warmup_steps (2-4)
Better quality, less speed	Lower residual_diff_threshold (0.12-0.18), higher max_warmup_steps (6-8), lower max_continuous_cached_steps (2)

Testing

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

omni = Omni(
    model="your-model-name",
    cache_backend="cache_dit",
    cache_config={
        "Fn_compute_blocks": 1,
        "Bn_compute_blocks": 0,
        "max_warmup_steps": 4,
        "residual_diff_threshold": 0.24,
    }
)
images = omni.generate(
    "a beautiful landscape",
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

CLI (online serving):

vllm serve your-model --omni --port 8098 \
  --cache-backend cache_dit \
  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4}'

Verification checklist:

1. Logs show "Cache-dit enabled successfully on xxx"
2. Performance: 1.5-2x speedup vs no cache
3. Quality: compare output with cache_backend=None

Excluded Models

Models listed in _NO_CACHE_ACCELERATION in vllm_omni/diffusion/registry.py do not support cache-dit (e.g., NextStep11Pipeline, StableDiffusionPipeline). Check this set before attempting to enable cache-dit.

Reference Implementations

Model	Path	Notes
Standard DiT	cache_dit_backend.py::enable_cache_for_dit	Default enabler, automatic
Wan2.2	cache_dit_backend.py::enable_cache_for_wan22	Dual-transformer, auto-detects mode
LongCat	cache_dit_backend.py::enable_cache_for_longcat_image	Multi-block-list
BAGEL	cache_dit_backend.py::enable_cache_for_bagel	Complex omni model