RFC: Ref-Target KV Cache for Video Diffusion Models (V3)

Abstract

This RFC proposes a Ref-Target KV Cache mechanism for video diffusion models in vLLM-Omni, inspired by FullDiT2's selective context caching. Unlike the original model-agnostic proposal, empirical validation on Wan2.1 VACE reveals that lossless caching is structurally impossible for VACE's reference hints due to full bidirectional attention (Ref tokens are "contaminated" by Target tokens at every layer). Consequently, this V3 RFC adopts a three-tiered approach (P0/P1/P2) that accurately reflects architectural realities:

Tier	Scope	Lossless?	Default	Speedup
P0	Text Cross-Attention K/V	✅ Bit-exact	❌ OFF	~0% (regression test only)
P1	VACE Hint Reuse (Refresh + Subset)	❌ Lossy (~8% mean diff)	❌ OFF	~15–25%
P2	Generic Lossless (Future)	✅	❌ OFF	TBD

Key empirical findings (single A800, Wan2.1-VACE-1.3B, 480×832, 17 frames, g=1):

• Reference hint L2 drift vs. T0 climbs monotonically to ~66% by the final step.
• P0 (text cross-attn K/V) is bit-exact but yields ~0% speedup.
• P1 (VACE hint cache) delivers ~19% speedup with mean pixel diff of 8.2% (max single-frame: 72%).

First cut = single-GPU Wan2.1-VACE V2V. TP/SP/CFG-parallel as a follow-up.

1. Motivation

1.1 The V2V Use Case

In Video-to-Video (V2V) generation, a source video serves as the reference context. The model must preserve the source video's content, structure, and motion while applying transformations specified by a text prompt.

The source video is encoded once but its features are attended to at every denoising step—creating significant step-wise computational redundancy. FullDiT2 demonstrates that selective caching can yield 2–3× average time cost reduction per diffusion step with minimal quality degradation.

1.2 T0 vs T1...TN: The Key Distinction

Aspect	T0 (First Step)	T1...TN (Subsequent Steps)
Reference K/V projection + LN + RoPE	✅ Compute per layer and cache	❌ Skipped (read cached post-RoPE K/V)
Target K/V projection + LN + RoPE	✅ Compute	✅ Compute
Input sequence composition	✅ [ref_tokens, tgt_tokens]	❌ [tgt_tokens] only
Cross-attention	✅ Q_target → [K_target_rope, K_ref_rope]	✅ Q_target → [K_target_rope, K_ref_rope_cached]
MoE processing	✅ Processes full sequence	✅ Processes only [tgt_tokens] (P2-only for ref dropping)

1.3 VACE Architectural Constraint: Ref Contamination

Critical distinction from FullDiT2: VACE uses full bidirectional attention—Ref and Target tokens interact bidirectionally at every layer. Empirical profiling (vace_drift_diff.py) confirms:

• Layer 1: Ref is immediately "contaminated" by Target features.
• This contaminated Ref propagates to Layer 2, getting further contaminated.
• By deep layers, cached K_ref/V_ref is heavily biased by Target information.
• Reusing "old contaminated" Ref at a later timestep (where Target has evolved) causes feature drift.

Therefore, lossless caching of VACE hints is structurally impossible. This is not an implementation artifact—it is inherent to the architecture. The original model-agnostic assumption (reference context remains perfectly static) does not hold for VACE.

1.4 MoE and the P2-Only Caveat

The statement "MoE only processes tgt from T1…TN" (from the original RFC) is P2-only, not applicable to VACE:

• In VACE, reference is entangled into the latent via vace_blocks (control = proj_in(control) + hidden_states).
• Dropping reference tokens from the MoE/FFN path changes the target's representation.
• Only in a decoupled-attention model (FullDiT2-style, P2) are reference tokens a separable sequence segment that can be dropped after T0.

Section 1.3 of any referenced internal documentation should be marked P2-only.

1.5 Design Goals (Revised)

Goal	P0 (Lossless)	P1 (Lossy)	P2 (Future)
Speedup	~0% (not a speedup line)	15–25%	TBD
Quality	Bit-exact to baseline	DINOv2 mean drop ≤ 8%	Lossless
Default	OFF	OFF	OFF
Scope	Text attn2 K/V only	VACE hint blocks	Decoupled-ref models
First Cut	Single-GPU; TP/SP/CFG-parallel follow-up	Same	Same

2. Design

2.1 Core Principle: Layer-Wise Ref KV Caching

K_ref and V_ref are layer-specific. Each transformer layer produces its own K_ref and V_ref from the reference context. Cache organized by layer index:

ref_cache: dict[int, tuple[Tensor, Tensor]] = {
    0: (K_ref_layer0, V_ref_layer0),
    1: (K_ref_layer1, V_ref_layer1),
    ...
    L-1: (K_ref_layerL-1, V_ref_layerL-1)
}

2.2 Three-Tier Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                         Ref-Target KV Cache Pipeline (V3)                  │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  ┌─────────────┐    ┌───────────────────────┐    ┌────────────────────┐   │
│  │   T0 Step   │───▶│  ┌─────────────────┐  │───▶│  P0: Cache Text   │   │
│  │  (Prefill)  │    │  │ Encode Ref &    │  │    │  attn2 K/V        │   │
│  └─────────────┘    │  │ Compute K/V     │  │    │  (lossless)       │   │
│                     │  └─────────────────┘  │    └────────────────────┘   │
│                     │            │           │                             │
│                     │            ▼           │    ┌────────────────────┐   │
│                     │  ┌─────────────────┐  │───▶│  P1: Cache VACE    │   │
│                     │  │ VACEHintCache   │  │    │  hints (lossy)     │   │
│                     │  │ (Refresh: K)    │  │    │  + refresh every   │   │
│                     │  └─────────────────┘  │    │  K steps           │   │
│                     └───────────────────────┘    └────────────────────┘   │
│                                                                             │
│  ┌─────────────┐    ┌───────────────────────┐    ┌────────────────────┐   │
│  │  T1...TN    │───▶│  Skip Ref Encode      │───▶│  Read from cache   │   │
│  │   Steps     │    │  (Read from cache)    │    │  + Attend to Tgt   │   │
│  └─────────────┘    └───────────────────────┘    └────────────────────┘   │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘

2.3 Configuration Flags

--enable-text-cross-attn-cache          # P0: default OFF, bit-exact regression only
--enable-vace-hint-cache-lossy          # P1: default OFF, requires --vace-hint-cache-warning-acknowledged
--enable-ref-target-kv                  # P2: default OFF, model-specific (future)

2.4 Component: RefTargetAdapter (P0 — Lossless, Default-OFF)

Contract: Lossless caching of text cross-attention K/V only. Serves as a regression test baseline, not a performance optimization.

class RefTargetAdapter:
    """Lossless adapter for text cross-attention K/V. P0 only."""
    
    def prefill(self, text_embeds: Tensor, branch_id: str) -> KVcache:
        """T0: Compute and cache text attn2 K/V. Branch-id distinguishes cond/uncond."""
        pass
    
    def get_cached_kv(self, layer_id: int, branch_id: str) -> Tuple[Tensor, Tensor]:
        """Retrieve cached K/V for a given layer and CFG branch."""
        pass

Key requirements:

• CFG Handling: Must be branch-keyed — conditional and unconditional branches use separate caches. Caching one branch and reusing it for the other corrupts output.
• FP8/Quantized KV: P0 bit-exact assertion must run in non-quantized mode. With FP8 KV enabled, P0 becomes a lossy-quantized path and should be documented separately.
• Acceptance: Bit-exact (MAX_ABS_DIFF = 0), latency-neutral (~0% speedup).

2.5 Component: VACEHintCache (P1 — Lossy, Default-OFF)

Contract: Lossy caching of VACE hints. Deliberately not routed through RefTargetAdapter to keep the lossless contract pure. Requires explicit allow_lossy=True opt-in.

class VACEHintCache:
    """Lossy cache for VACE hints. P1 only. Default-OFF."""
    
    def __init__(
        self,
        refresh_interval: int = 4,        # K: Refresh every K steps
        cached_layers: List[int] = None,  # Default: deep layers only
        allow_lossy: bool = False         # Explicit opt-in required
    ):
        self.refresh_interval = refresh_interval
        self.cached_layers = cached_layers or []  # Shallow layers change too rapidly.
        self.allow_lossy = allow_lossy
        self.step_counter = 0
    
    def should_refresh(self) -> bool:
        return self.step_counter % self.refresh_interval == 0
    
    def refresh(self, step: int) -> None:
        """Recompute and update cache at step."""
        pass

Why deep layers only?

• Shallow layers: more susceptible to contamination and rapid texture changes; caching them causes significant artifacts.
• Deep layers: carry more stable semantic features; safer to cache, albeit still lossy.

Why refresh every K steps?

• Cuts the long error accumulation chain into short segments.
• Prevents video collapse in later generation stages.
• Provides a loss↔speed dial: smaller K = higher quality but less speedup.

Acceptance criteria (under real serving CFG, not g=1):

• Default-OFF; when ON, reproduce ~15–25% speedup.
• Quality gate: DINOv2 feature similarity drop ≤ 8% (mean), max single-frame drop ≤ 15%.
• Explicit --vace-hint-cache-warning-acknowledged flag required.

Important: The ~8% mean diff from g=1 runs is a lower bound. Quality gate must be measured under real serving CFG (g>1), where two branches drift independently, potentially amplifying errors.

2.6 Component: Wan2_1VACEAdapter (Model Wiring)

Wires both P0 and P1 caches into the VACE attention layers.

class Wan2_1VACEAdapter:
    def forward_attention(
        self, 
        q, k_target, v_target, 
        layer_id, step, 
        use_text_cache=False,    # P0
        use_hint_cache=False,    # P1
        branch_id="cond"
    ):
        # P0: Text cross-attn cache
        if use_text_cache and self.text_cache.has(layer_id, branch_id):
            k_text, v_text = self.text_cache.retrieve(layer_id, branch_id)
            k = torch.cat([k_target, k_text], dim=-2)
            v = torch.cat([v_target, v_text], dim=-2)
        # P1: VACE hint cache (separate path)
        elif use_hint_cache and self.hint_cache.has(layer_id):
            k_hint, v_hint = self.hint_cache.retrieve(layer_id)
            k = torch.cat([k_target, k_hint], dim=-2)
            v = torch.cat([v_target, v_hint], dim=-2)
        else:
            k, v = k_target, v_target
        return scaled_dot_product_attention(q, k, v)

2.7 P2 — Future Lossless (Generic)

P2 targets a future lossless implementation via architectural decoupling (similar to FullDiT2):

• Causal mask prevents Ref from seeing Target.
• Ref tokens are a separable sequence segment.
• MoE routing can drop ref tokens after T0.

This will be a drop-in replacement via --enable-ref-target-kv when such a model variant lands. Candidate models: Hunyuan-Image reference-gen, CausalWan-style world models.

2.8 Cache Lifecycle

• Per-request — cache must be cleared between requests.
• Warmup-safe — vLLM's warmup fires a dummy denoise step; cache must be keyed by request ID to avoid warmup K/V leaking into real runs.
• Multi-worker — config must be set at process launch and live inside the worker.

3. Configuration

3.1 Parameters

@dataclass
class RefTargetKVCacheConfig:
    # P0: Lossless text cross-attn cache
    enable_text_cross_attn_cache: bool = False
    
    # P1: Lossy VACE hint cache
    enable_vace_hint_cache_lossy: bool = False
    vace_hint_cache_warning_acknowledged: bool = False
    vace_hint_refresh_interval: int = 4          # K
    vace_hint_cached_layers: List[int] = None    # Default: deep layers only
    
    # P2: Future lossless (model-specific)
    enable_ref_target_kv: bool = False
    model_type: str = "wan2_1"
    
    # Cache lifecycle
    ref_cache_on_device: bool = True
    enable_transfer: bool = False

3.2 Usage

# P0: Lossless regression test only (default-off)
vllm serve Wan2.1 --omni --enable-text-cross-attn-cache

# P1: Lossy VACE hint cache (requires explicit ack)
vllm serve Wan2.1 --omni \
    --enable-vace-hint-cache-lossy \
    --vace-hint-cache-warning-acknowledged \
    --vace-hint-refresh-interval 4

# P2: Future decoupled-reference model
vllm serve Wan2.1 --omni --enable-ref-target-kv --model-type wan2_1

4. Implementation Roadmap

First cut = single-GPU Wan2.1-VACE. TP/SP/CFG-parallel as a follow-up.

Phase	Tasks	Duration	Deliverables	Owner
Phase 1	Implement RefTargetAdapter + RefTargetKVCacheManager	1 Week	P0: Bit-exact regression test harness	@linzhenpl07
Phase 2	Implement VACEHintCache + refresh logic	1 Week	P1: Lossy hint cache with K-step refresh
Phase 3	Integrate with vLLM-Omni (OmniKVTransferManager)	1 Week	Seamless KV transfer layer
Phase 4a	Performance benchmarking + P1 K-value enumeration	0.5 Week	Speedup numbers + raw quality metrics	@linzhenpl07
Phase 4b	P1 Quality Gate (CI + Subjective)	0.5 Week (parallel)	DINOv2 CI gate, Pareto analysis, A/B test report → go/no-go decision	@linzhenpl07

Total Duration: 4 Weeks

Phase 4b — Quality Gate Detailed Workflow (Owned by @linzhenpl07)

Step	Task	Output
1	Integrate DINOv2 evaluation script (#4497 M0 harness) into CI	CI gate: mean drop ≤ 8%, max drop ≤ 15%
2	Run P1 evaluation under real serving CFG (g>1, not g=1)	Per-K quality table (K=1,2,3,4,6,8)
3	Plot Pareto frontier (speedup vs. quality loss)	Elbow point selection recommendation
4	A/B subjective test (N=20-30, 5-10 raters)	≥90% "no noticeable degradation" → gate green

Decision Branch

• Gate green (all five items in §5.3 ✅) → P1 implementation accepted for merge.
• Gate red (e.g., quality collapse at K=4, no acceptable K-value, or subjective test fails) → P1 abandoned; team pivots to P2 (FullDiT2-style decoupled attention) as the recommended path forward. A follow-up RFC will be drafted for P2.

5. Validation Checklist

5.1 Functional Validation

• T0 computes and caches K_ref/V_ref correctly.
• T1...TN skip reference encoding and read from cache.
• Cross-attention correctly combines K_target + K_ref_cached.
• P0: Output quality matches baseline bit-exactly (MAX_ABS_DIFF = 0).
• P0: CFG branches use branch-keyed caches (cond vs. uncond).
• P1: Explicit --vace-hint-cache-warning-acknowledged required to enable.
• P1: Cache cleared per request; warmup does not poison real runs.
• P2: Section 1.3 of internal docs marked P2-only (MoE ref dropping not applicable to VACE).

5.2 Performance Validation

• Achieve 15–25% speedup for V2V short videos on single H100/A800.
• Multi-GPU (TP/SP/CFG-parallel) deferred to follow-up.

5.3 P1 Quality Evaluation (3-Layer Funnel) — Owned by @linzhenpl07

Stage	Metric	Threshold	Status
CI Gate	DINOv2 Feature Similarity (vs. baseline)	Mean drop ≤ 8%, Max single-frame drop ≤ 15%	⬜
CI Gate	Latency Speedup	≥ 10% (otherwise cache not worth the quality loss)	⬜
Offline Tuning	Pareto Frontier (K vs. Quality)	Enumerate K ∈ {1, 2, 3, 4, 6, 8}. Select elbow point. Report raw table.	⬜
Offline Tuning	Temporal Error Heatmap	Ensure no sharp decay in the second half of long videos (>5s).	⬜
Subjective	A/B Preference Test (MOS proxy)	N=20-30, 5-10 raters, ≥90% samples show "no noticeable degradation"	⬜

Final Decision: All five items above must be ✅ for P1 to be accepted. If any fails, P1 is rejected and the team pivots to P2 (FullDiT2-style decoupled attention).

Critical Notes:

• K-Value Monotonicity: Quality loss does not strictly increase monotonically with K. Use discrete enumeration with error bars (multiple seeds) and prioritize the largest K that remains within the quality redline.
• CFG Measurement: All P1 quality gates must be measured under real serving CFG (g>1), not g=1. The g=1 ~8% mean diff is a lower bound; actual serving loss may be higher.

5.4 FP8 / Quantized KV Compatibility

• P0 bit-exact assertion must run in non-quantized mode.
• With FP8 KV enabled, P0 becomes a lossy-quantized path and should be documented separately.
• P1 (already lossy) has no conflict with FP8 KV.

6. Parallelism Support

First cut = single-GPU. Multi-GPU support (TP/SP/CFG-parallel) to be validated in follow-up.

Strategy	Status	Notes
Tensor Parallelism (TP)	✅ (follow-up)	Must return rank-local shard of K/V.
Sequence Parallelism (SP)	✅ (follow-up)	Shards reference tokens by sequence.
Data Parallelism (DP)	✅ (follow-up)	Standard DP.
CFG-Parallel	✅ (follow-up)	Cond/uncond on different ranks; maps to branch-keyed caches.
MoE Models	✅ (follow-up)	P2-only; VACE entanglement prevents dropping ref tokens.
Quantized KV Cache	✅ (conditional)	P0: bit-exact only in non-quantized mode.

7. Risks & Mitigations

Risk	Impact	Mitigation
VACE Ref Contamination (Full Attention)	Intrinsic ~8% quality loss in P1.	P1 uses refresh (K) + deep-layer subset to control drift. If quality drops catastrophically at K=4, pivot to P2 (FullDiT2-style decoupled attention).
Non-monotonic K-value Loss	Pareto analysis becomes complex.	Use discrete enumeration with statistical significance (error bars). Treat K=3 and K=4 as equivalent if their losses overlap within variance, and pick the larger K.
CFG Double-branch OOM	H100/910B memory limits.	Offload unconditional branch K/V to CPU memory as fallback.
Long Video Error Accumulation	Later frames collapse.	Include >5-second videos in test set. Track slope of quality decay in second half. If slope > 0.5/sec, reject config.
Warmup Poisoning	Dummy warmup K/V leaks into real run.	Key cache by request ID; clear per-request.
P0 Memory Overhead	Memory increase for no speedup.	P0 default-OFF; use only for regression testing.
Internal Doc §1.3 Misinterpretation	Team assumes MoE ref dropping works for VACE.	Explicitly mark §1.3 as P2-only in the RFC and review process.

8. Reproducibility

All measurement scripts are available at: https://gist.github.com/linzhenpl07/69f129766cb9fbfeddd4091d98cde131

Script	Purpose
vace_drift_diff.py	Measures Ref L2 drift vs. T0 (~66% by last step).
vace_p0_diff2.py + p0_patch2.py	P0: Bit-exact text cross-attn cache.
vace_p1_diff.py	P1: Lossy VACE hint cache with refresh.

Results Summary (Single A800, Wan2.1-VACE-1.3B, 480×832, 17 frames, g=1):

Tier	Speedup	Quality	Hits/Misses
P0	~0% (15.63s → 15.69s @ 30 steps)	Bit-exact (MAX_ABS_DIFF=0)	6705 hits / 45 miss
P1	~19% (6.36s → 5.12s @ 10 steps)	Mean pixel diff 8.2%, Max 72%	405 hits / 45 miss

9. References

1. FullDiT2: Efficient In-Context Conditioning for Video Diffusion Transformers. Xuanhua He, Quande Liu, et al. arXiv:2506.04213, 2025.

   • Selective Context Caching.
   • Decoupled Masked Attention (Ref cannot see Tgt).
   • 2–3× speedup per diffusion step.

2. Wan2.1 VACE Pipeline Documentation (Internal).

   • Note: Section 1.3 (MoE ref dropping) applies to P2 only, not VACE.

3. vLLM-Omni OmniKVTransferManager (Internal).

4. Repro Scripts: https://gist.github.com/linzhenpl07/69f129766cb9fbfeddd4091d98cde131

10. CC List

@hsliuustc0106 @Gaohan123 @david6666666 @bjf-frz @asukaqaq-s @linzhenpl07 @evanchueng