[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/diffusion/models/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))