diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/dreamzero/test_crossattn_cache.py new file mode 100644 index 000000000..94361c090 --- /dev/null +++ b/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 diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/dreamzero/test_fused_qk_rms_norm.py new file mode 100644 index 000000000..dbbb43940 --- /dev/null +++ b/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() diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/dreamzero/test_qkv_fusion.py new file mode 100644 index 000000000..16b78693f --- /dev/null +++ b/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)) diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py index e187054e2..ad3175e1e 100644 --- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py +++ b/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, ) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 30ea2a4ce..1d1429534 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/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) diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py index 0972d1b5a..9c1e7873d 100644 --- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py +++ b/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,