[Feat] DreamZero fix tp & cross attn cache (#4154)

tests/dreamzero/test_crossattn_cache.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused.
+
+DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a
+per-layer ``crossattn_cache`` dict: they are computed on the first forward of a
+session (``is_init`` False -> True) and reused on every later step, because the
+text/image context is session-invariant. Only the query (which depends on the
+per-step hidden state) is recomputed.
+
+These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache
+on the first call and reuses it (not recompute) on later calls, and recomputes
+every call when no cache is provided.
+
+CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it
+builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub
+the projections/attention to count calls and exercise only the caching logic.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_CLIP_LEN = 257  # image tokens the I2V cross-attn splits off the front of context
+
+
+def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention:
+    """A WanI2VCrossAttention whose submodules are call-counting stubs.
+
+    ``__init__`` is bypassed (it needs a TP group); we only wire what
+    ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so
+    the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so
+    we can assert how many times it was invoked.
+    """
+    inner = n_heads * head_dim
+    module = WanI2VCrossAttention.__new__(WanI2VCrossAttention)
+    module.tp_num_heads = n_heads
+    module.head_dim = head_dim
+
+    def _proj() -> MagicMock:
+        return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner))
+
+    module.q = _proj()
+    module.k = _proj()
+    module.v = _proj()
+    module.k_img = _proj()
+    module.v_img = _proj()
+    module.norm_q = lambda t: t
+    module.norm_k = lambda t: t
+    module.norm_k_img = lambda t: t
+    # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy.
+    module.attn = lambda q, k, v: q
+    module.o = lambda t: t
+    return module
+
+
+def _fresh_cache() -> dict:
+    return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None}
+
+
+def _inputs(inner: int = 8):
+    x = torch.randn(1, 3, inner)  # B=1, query_len=3
+    context = torch.randn(1, _CLIP_LEN + 5, inner)  # 257 image tokens + 5 text tokens
+    return x, context
+
+
+def test_cache_inits_on_first_call_and_reuses_on_second() -> None:
+    module = _make_cross_attn()
+    cache = _fresh_cache()
+    x, context = _inputs()
+
+    # First call: cache is initialised and every k/v projection runs once.
+    module.forward(x, context, crossattn_cache=cache)
+    assert cache["is_init"] is True
+    for key in ("k", "v", "k_img", "v_img"):
+        assert cache[key] is not None, f"{key} should be cached after first call"
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+
+    cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")}
+
+    # Second call: k/v/k_img/v_img are read from cache, NOT recomputed.
+    module.forward(x, context, crossattn_cache=cache)
+    assert module.k.call_count == 1
+    assert module.v.call_count == 1
+    assert module.k_img.call_count == 1
+    assert module.v_img.call_count == 1
+    for key, tensor in cached.items():
+        assert cache[key] is tensor, f"{key} should be the same cached object"
+
+    # The query always depends on the per-step hidden state, so it is recomputed.
+    assert module.q.call_count == 2
+
+
+def test_no_cache_recomputes_every_call() -> None:
+    module = _make_cross_attn()
+    x, context = _inputs()
+
+    module.forward(x, context, crossattn_cache=None)
+    module.forward(x, context, crossattn_cache=None)
+
+    # Without a cache, the image/text k/v are recomputed on every call.
+    assert module.k_img.call_count == 2
+    assert module.v_img.call_count == 2
+    assert module.k.call_count == 2
+    assert module.v.call_count == 2

tests/dreamzero/test_fused_qk_rms_norm.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the TP world size and the collective are mocked, so a single process
+simulates one rank (all-reduce is identity). This exercises the cat/split and
+count-scaling logic that a real multi-rank run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.dreamzero.causal_wan_model import (
+    DistributedRMSNorm,
+    fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model"
+
+
+def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]:
+    """Two norms with distinct random weights so a swapped q/k slice is caught."""
+    torch.manual_seed(0)
+    norm_q = DistributedRMSNorm(dim)
+    norm_k = DistributedRMSNorm(dim)
+    norm_q.weight.data.normal_()
+    norm_k.weight.data.normal_()
+    return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+    """Single-rank all-reduce: the reduction over one rank is the identity."""
+    return tensor
+
+
+def test_fused_matches_separate_tp1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_matches_separate_tp_gt1() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    # tp_size=2 exercises the reduce branch and the `global_count *= tp_size`
+    # path in both the per-norm forward and the fused helper.
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce),
+    ):
+        ref_q, ref_k = norm_q(q), norm_k(k)
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_single_all_reduce_over_packed_pair() -> None:
+    dim = 64
+    norm_q, norm_k = _make_norms(dim)
+    q = torch.randn(2, 5, dim)
+    k = torch.randn(2, 5, dim)
+
+    with (
+        patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+        patch(
+            f"{_MODULE}.tensor_model_parallel_all_reduce",
+            side_effect=_identity_all_reduce,
+        ) as mock_ar,
+    ):
+        fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    # Exactly one collective, over a tensor whose last dim packs the q+k pair.
+    assert mock_ar.call_count == 1
+    packed = mock_ar.call_args[0][0]
+    assert packed.shape[-1] == 2
+
+
+def test_fused_identity_fallback_when_qk_norm_disabled() -> None:
+    # qk_norm=False makes the norms nn.Identity; the helper must fall back to
+    # plain passthrough without touching the collective.
+    norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity()
+    q = torch.randn(2, 5, 64)
+    k = torch.randn(2, 5, 64)
+
+    with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar:
+        fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+    torch.testing.assert_close(fused_q, q, rtol=0, atol=0)
+    torch.testing.assert_close(fused_k, k, rtol=0, atol=0)
+    mock_ar.assert_not_called()

tests/dreamzero/test_qkv_fusion.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence test for DreamZero self-attn QKV fusion.
+
+CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and
+splits the result, instead of three separate ColumnParallelLinear. The DreamZero
+checkpoint still ships separate q/k/v weights, so load_weights routes each into
+the packed `qkv` param with its shard id. This test proves the chain
+"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three
+independent projections, given identical weights.
+
+CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and
+the GEMM dispatcher is forced to the default (CPU-compatible) path, following
+tests/diffusion/layers/test_adalayernorm.py.
+"""
+
+import os
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Minimal world_size=1 TP group required to build QKVParallelLinear."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29507")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for CPU test tensors."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_fused_qkv_load_and_split_matches_separate_projections():
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    torch.manual_seed(0)
+    dim, head_dim, num_heads = 16, 4, 4  # inner_dim == dim, no GQA
+
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    # Distinct random weight + bias per projection, loaded the exact way
+    # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id).
+    refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+    for shard_id in ("q", "k", "v"):
+        w = torch.randn(dim, dim)
+        b = torch.randn(dim)
+        refs[shard_id] = (w, b)
+        qkv.weight.weight_loader(qkv.weight, w, shard_id)
+        qkv.bias.weight_loader(qkv.bias, b, shard_id)
+
+    x = torch.randn(2, 3, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True):
+        w, b = refs[shard_id]
+        ref = torch.nn.functional.linear(x, w, b)
+        torch.testing.assert_close(got, ref)
+
+
+def test_fused_qkv_shard_ids_are_not_transposed():
+    # Guard against q/k/v being routed to the wrong slice: give each projection a
+    # constant, distinct weight so a swapped shard id would surface immediately.
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    dim, head_dim, num_heads = 16, 4, 4
+    qkv = QKVParallelLinear(
+        hidden_size=dim,
+        head_size=head_dim,
+        total_num_heads=num_heads,
+        bias=True,
+    )
+
+    scale = {"q": 1.0, "k": 2.0, "v": 3.0}
+    for shard_id, s in scale.items():
+        qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id)
+        qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id)
+
+    x = torch.ones(1, dim)
+    fused, _ = qkv(x)
+    q, k, v = fused.split([dim, dim, dim], dim=-1)
+
+    # F.linear(ones, full(s)) == s * dim per element.
+    torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim))
+    torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim))
+    torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim))

vllm_omni/diffusion/models/dreamzero/causal_wan_model.py

@@ -21,11 +21,12 @@ import torch.nn as nn
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
-    get_tp_group,
+    tensor_model_parallel_all_reduce,
 )
 from vllm.model_executor.layers.conv import Conv3dLayer
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
+    QKVParallelLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.utils import set_weight_attrs
@@ -174,22 +175,82 @@ class DistributedRMSNorm(nn.Module):
             raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
         param.data.copy_(shard)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        tp_size = get_tensor_model_parallel_world_size()
+    def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
+        """Per-rank float32 activation, local sum-of-squares, and local width."""
         x_float = x.float()
         local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
-        local_count = x.shape[-1]
+        return x_float, local_sum_sq, x.shape[-1]
+
+    def _scale(
+        self,
+        x_float: torch.Tensor,
+        global_sum_sq: torch.Tensor,
+        global_count: int,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply the (already reduced) RMS and the per-rank weight shard."""
+        mean_sq = global_sum_sq / global_count
+        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_float, local_sum_sq, local_count = self._local_sum_sq(x)
 
         if tp_size > 1:
-            global_sum_sq = local_sum_sq.clone()
-            torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group)
+            # Use vLLM's collective (custom all-reduce / symmetric-mem fast path
+            # for small tensors) instead of raw torch.distributed.all_reduce, and
+            # take the return value so the custom-AR path (which may return a new
+            # buffer) is handled correctly. No .clone() needed.
+            global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
             global_count = local_count * tp_size
         else:
             global_sum_sq = local_sum_sq
             global_count = local_count
 
-        mean_sq = global_sum_sq / global_count
-        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
+        return self._scale(x_float, global_sum_sq, global_count, x)
+
+
+def fused_qk_rms_norm(
+    norm_q: nn.Module,
+    norm_k: nn.Module,
+    q: torch.Tensor,
+    k: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply q/k :class:`DistributedRMSNorm` with a SINGLE fused TP all-reduce.
+
+    WHY: self-attention norms q and k every step; on TP each norm issues its own
+    tiny all-reduce of the per-token sum-of-squares. For DreamZero's decode-like
+    per-step forward these latency-bound micro-collectives dominate, so we pack
+    both sum-of-squares into one tensor and reduce once (2 collectives → 1).
+
+    NUMERICALLY IDENTICAL to ``norm_q(q), norm_k(k)``: all-reduce is elementwise,
+    so packing along the last dim reduces each slice independently with the same
+    fp32 accumulation. Requires q and k to share the same shape (true for
+    self-attention — both come from the same hidden states). Falls back to
+    independent application when either norm is not a DistributedRMSNorm
+    (e.g. nn.Identity when qk_norm=False).
+    """
+    if not (isinstance(norm_q, DistributedRMSNorm) and isinstance(norm_k, DistributedRMSNorm)):
+        return norm_q(q), norm_k(k)
+
+    # The fused path reduces one packed sum-of-squares and reuses q's token width
+    # as the RMS count for BOTH q and k, so q and k must share the same shape.
+    # (Self-attention guarantees this: q and k are projected from the same x.)
+    assert q.shape == k.shape, "fused_qk_rms_norm requires q and k to have the same shape."
+
+    tp_size = get_tensor_model_parallel_world_size()
+    q_float, q_sum_sq, count = norm_q._local_sum_sq(q)
+    k_float, k_sum_sq, _ = norm_k._local_sum_sq(k)
+
+    if tp_size > 1:
+        packed = torch.cat([q_sum_sq, k_sum_sq], dim=-1)
+        packed = tensor_model_parallel_all_reduce(packed)
+        q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2]
+        count = count * tp_size
+
+    q_out = norm_q._scale(q_float, q_sum_sq, count, q)
+    k_out = norm_k._scale(k_float, k_sum_sq, count, k)
+    return q_out, k_out
 
 
 # ── Projections ─────────────────────────────────────────────────────
@@ -337,22 +398,35 @@ class WanI2VCrossAttention(nn.Module):
         context = context[:, 257:]
         n, d = self.tp_num_heads, self.head_dim
         q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        # context (text) and context_img (clip features) are constant within a
+        # session, so k/v and k_img/v_img are cached on first call and reused.
+        # context_img == img_emb(clip feature) is session-invariant (state.clip_feas
+        # is set once at current_start_frame==0 and only cleared on reset, which also
+        # rebuilds crossattn_cache). Caching k_img/v_img removes a per-step
+        # norm_k_img all-reduce plus the k_img/v_img projection GEMMs in steady state.
+        # Only q depends on the per-step x and is always recomputed.
         if crossattn_cache is not None:
             if not crossattn_cache["is_init"]:
                 crossattn_cache["is_init"] = True
                 k = self.norm_k(self.k(context)).unflatten(2, (n, d))
                 v = self.v(context).unflatten(2, (n, d))
+                k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
+                v_img = self.v_img(context_img).unflatten(2, (n, d))
                 crossattn_cache["k"] = k
                 crossattn_cache["v"] = v
+                crossattn_cache["k_img"] = k_img
+                crossattn_cache["v_img"] = v_img
             else:
                 k = crossattn_cache["k"]
                 v = crossattn_cache["v"]
+                k_img = crossattn_cache["k_img"]
+                v_img = crossattn_cache["v_img"]
         else:
             k = self.norm_k(self.k(context)).unflatten(2, (n, d))
             v = self.v(context).unflatten(2, (n, d))
+            k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
+            v_img = self.v_img(context_img).unflatten(2, (n, d))
         x = self.attn(q, k, v)
-        k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
-        v_img = self.v_img(context_img).unflatten(2, (n, d))
         img_x = self.attn(q, k_img, v_img)
         x = x.flatten(2)
         img_x = img_x.flatten(2)
@@ -403,9 +477,21 @@ class CausalWanSelfAttention(nn.Module):
         self.num_action_per_block = num_action_per_block
         self.num_state_per_block = num_state_per_block
         self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen
-        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
-        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
-        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
+        # Fused QKV projection: q/k/v all come from x, so a single column-parallel
+        # GEMM replaces three (fewer launches / better GEMM util in the decode-like
+        # per-step forward). No GQA here: total_num_kv_heads defaults to num_heads.
+        self.qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            bias=True,
+        )
+        # The forward splits qkv into three equal q/k/v shards, which is only
+        # correct without GQA. total_num_kv_heads defaults to total_num_heads
+        # here; fail loud if a future checkpoint introduces GQA so the equal
+        # split does not silently misalign k/v.
+        if self.qkv.total_num_kv_heads != self.qkv.total_num_heads:
+            raise ValueError("Self-attn QKV fusion requires no GQA (total_num_kv_heads == total_num_heads).")
         self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
         self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
@@ -431,9 +517,15 @@ class CausalWanSelfAttention(nn.Module):
         """Inference-only forward (KV cache path)."""
         n, d = self.tp_num_heads, self.head_dim
 
-        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
-        k = self.norm_k(self.k(x)).unflatten(2, (n, d))
-        v = self.v(x).unflatten(2, (n, d))
+        # Single fused QKV GEMM, then split into the per-rank q/k/v shards.
+        qkv, _ = self.qkv(x)
+        qk_size = self.tp_num_heads * self.head_dim
+        q, k, v = qkv.split([qk_size, qk_size, qk_size], dim=-1)
+        # Fuse q/k qk-norm into a single TP all-reduce (both come from x, same shape).
+        q, k = fused_qk_rms_norm(self.norm_q, self.norm_k, q, k)
+        q = q.unflatten(2, (n, d))
+        k = k.unflatten(2, (n, d))
+        v = v.unflatten(2, (n, d))
 
         updated_kv_cache: torch.Tensor | None = None
 
@@ -834,6 +926,7 @@ class CausalWanModel(nn.Module):
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
         kv_cache: list[torch.Tensor],
+        crossattn_cache: list[dict] | None,
         current_start_frame: int,
     ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
@@ -884,6 +977,7 @@ class CausalWanModel(nn.Module):
                 context=context,
                 action_register_length=action_register_length,
                 kv_cache=kv_cache[block_index] if kv_cache else None,
+                crossattn_cache=crossattn_cache[block_index] if crossattn_cache else None,
                 current_start_frame=current_start_frame,
             )
             updated_kv_caches.append(updated_kv_cache)
@@ -941,6 +1035,7 @@ class CausalWanModel(nn.Module):
             timestep_action=timestep_action,
             state=state,
             kv_cache=kv_cache,
+            crossattn_cache=crossattn_cache,
             current_start_frame=current_start_frame,
         )
 

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -1104,10 +1104,27 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                     .replace("img_emb.proj.3.", "img_emb.fc2.")
                     .replace("img_emb.proj.4.", "img_emb.norm2.")
                 )
+
+                # Self-attn q/k/v are fused into a single QKVParallelLinear; route
+                # each separate checkpoint weight/bias to the packed `qkv` param
+                # with its shard id. cross_attn keeps separate q/k/v (q from x,
+                # k/v from context — not fusible), so it is left untouched here.
+                qkv_shard_id: str | None = None
+                for shard_id in ("q", "k", "v"):
+                    needle = f".self_attn.{shard_id}."
+                    if needle in new_name:
+                        new_name = new_name.replace(needle, ".self_attn.qkv.")
+                        qkv_shard_id = shard_id
+                        break
+
                 if new_name in params:
                     param = params[new_name]
-                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                    weight_loader(param, tensor)
+                    if qkv_shard_id is not None:
+                        # QKVParallelLinear.weight_loader needs the shard id.
+                        param.weight_loader(param, tensor, qkv_shard_id)
+                    else:
+                        weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                        weight_loader(param, tensor)
                     loaded.add(new_name)
                 elif new_name in buffers:
                     buffers[new_name].data.copy_(tensor)

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -131,8 +131,12 @@ class DreamZeroState:
             torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
         ]
 
-        self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
-        self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+        self.crossattn_cache = [
+            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
+        ]
+        self.crossattn_cache_neg = [
+            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
+        ]
 
     def update_kv_cache(
         self,