Merge remote-tracking branch 'origin/main' into qwen-image-2.1

# Conflicts:
#	vllm_omni/diffusion/executor/multiproc_executor.py
#	vllm_omni/diffusion/offloader/sequential_backend.py

tests/diffusion/distributed/test_cfg_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for CFG (Classifier-Free Guidance) parallel functionality.
 
 CPU tests mock the CFG process group and validate mixin logic without GPUs.
@@ -32,8 +32,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.platforms import current_omni_platform
 
-_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
-_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
+_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=3)
 
 
 def _set_random_seeds(seed: int) -> None:
@@ -1009,7 +1009,7 @@ def _run_multi_branch_gpu_parity(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2)
 @pytest.mark.parametrize("batch_size", [2])
 @pytest.mark.parametrize("cfg_normalize", [False, True])
 @pytest.mark.parametrize("extra_kwargs", [None, {"step_i": 1}])

tests/diffusion/distributed/test_comm.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.
 
 CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
@@ -24,8 +24,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.platforms import current_omni_platform
 
-_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
-_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4)
 
 DeviceKind = Literal["cpu", "cuda"]
 
@@ -461,3 +461,46 @@ def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
         device_kind="cuda",
         master_port=29501,
     )
+
+
+def _run_fused_qkv_matches_three_4d(
+    local_rank: int,
+    world_size: int,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    """LingBot's fused (B, S/N, 3, H, D) exchange must equal three per-tensor 4D exchanges."""
+    from vllm_omni.diffusion.distributed.comm import all_to_all_5D
+
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
+    try:
+        batch, seq_per_rank, heads, head_size = 1, 6, 4 * world_size, 8
+        # Distinguishable values: rank, q/k/v index, position, head, feature.
+        base = torch.arange(seq_per_rank * heads * head_size, dtype=torch.float32, device=device).reshape(
+            batch, seq_per_rank, heads, head_size
+        )
+        tensors = [base + 1000.0 * local_rank + 100000.0 * which for which in range(3)]
+        separate = [SeqAllToAll4D.apply(sp_group, t, 2, 1, False) for t in tensors]
+        fused = all_to_all_5D(torch.stack(tensors, dim=2), scatter_idx=3, gather_idx=1, group=sp_group)
+        assert fused.shape == (batch, seq_per_rank * world_size, 3, heads // world_size, head_size)
+        for which, expected in enumerate(separate):
+            assert torch.equal(fused.unbind(2)[which], expected), f"q/k/v index {which} differs on rank {local_rank}"
+    finally:
+        destroy_distributed_env()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
+    torch.multiprocessing.spawn(
+        _run_fused_qkv_matches_three_4d,
+        args=(world_size, "cpu", 29613),
+        nprocs=world_size,
+    )

tests/diffusion/distributed/test_expert_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 System test for Expert Parallel (EP).

tests/diffusion/distributed/test_flux2_vae_patch_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """FLUX.2-dev VAE patch-parallel correctness tests.
 

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SP subgroup construction in parallel_state.py."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for SP subgroup construction for diffusion parallel state."""
 
 from __future__ import annotations
 
@@ -8,12 +8,55 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+import vllm.distributed.parallel_state as vllm_parallel_state
 
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.distributed import parallel_state as omni_parallel_state
 from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def ensure_parallel_state_initialized():
+    """Ensure torch distributed is initialized and that all world / count vars are set."""
+    assert torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is not None
+    assert vllm_parallel_state._WORLD is not None
+    assert vllm_parallel_state._NODE_COUNT is not None
+
+
+def ensure_parallel_state_not_initialized():
+    """Ensure torch distributed is not initialized and that all world / count vars are unset."""
+    assert not torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is None
+    assert vllm_parallel_state._WORLD is None
+    assert vllm_parallel_state._NODE_COUNT is None
+
+
+def test_omni_manages_vllm_distributed_state(monkeypatch):
+    """Verify Omni initializes and tears down vLLM's distributed state properly.
+
+    This is a regression test for ensuring vLLM Omni also sets up vLLM's vars correctly
+    to avoid potential misalignment in cases where Omni uses vLLM's native coordinator,
+    e.g., MoE + diffusion.
+    """
+    ensure_parallel_state_not_initialized()
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "get_device_count", lambda: 1)
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "set_device", lambda _device: None)
+    omni_parallel_state.init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=get_distributed_init_method(),
+        backend="gloo",
+    )
+    try:
+        ensure_parallel_state_initialized()
+    finally:
+        omni_parallel_state.destroy_distributed_env()
+    ensure_parallel_state_not_initialized()
+
+
 def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
     def _fake_new_group(ranks, *args, **kwargs):
         group = SimpleNamespace(ranks=list(ranks))

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -31,9 +31,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.platforms import current_omni_platform
 
-_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
-_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
-_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+_L4_TWO_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=3)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": ["L4", "B200"]}, num_cards=4)
 
 DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

tests/diffusion/distributed/test_sana_video_sp_numeric.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Real multi-process numeric equivalence gate for SANA-Video sequence parallel.
 
 Runs the native transformer directly (no pipeline, no topology validation) with
@@ -273,7 +273,7 @@ def _baseline_worker(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2)
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 @pytest.mark.parametrize("num_frames", [21, 5])
 def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None:
@@ -290,7 +290,7 @@ def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4)
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 @pytest.mark.parametrize("num_frames", [21, 5])
 def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None:
@@ -307,7 +307,7 @@ def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4)
 def test_sp2_cfg2_matches_sp1_fp32() -> None:
     """Four processes: CFG branches with distinct inputs each run their own SP
     group; catches CFG/SP group crossover and deadlocks."""
@@ -324,7 +324,7 @@ def test_sp2_cfg2_matches_sp1_fp32() -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=4)
 def test_tp2_sp2_matches_sp1_fp32() -> None:
     """Four processes: TP shards heads and channels while SP shards frames; the
     combined mesh must still reproduce the serial output."""
@@ -342,7 +342,7 @@ def test_tp2_sp2_matches_sp1_fp32() -> None:
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@hardware_test(res={"cuda": ["L4", "B200"]}, num_cards=2)
 def test_sp2_bf16_drift_informational() -> None:
     """Records the bf16 reduction-reorder drift magnitude; the load-bearing
     numeric gate is the fp32 suite above."""

tests/diffusion/distributed/test_sequence_parallel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.

tests/diffusion/distributed/test_vae_decode_parallelism.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """VAE patch-parallel decode parity tests for diffusion models.
 
@@ -145,7 +145,7 @@ def _run_generate(
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 4})
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards={"cuda": 4})
 @pytest.mark.parametrize("model_case", MODEL_CASES)
 def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
     if current_omni_platform.is_npu():

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import os
 from types import SimpleNamespace
 
@@ -353,6 +356,94 @@ def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.Mon
         wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
 
 
+def _reference_conv_input(conv, x, cache_x):
+    """The conv input the previous implementation built: cat, F.pad, then halo_exchange."""
+    padding = list(conv._padding)
+    if cache_x is not None and padding[4] > 0:
+        x = torch.cat([cache_x, x], dim=2)
+        padding[4] -= cache_x.shape[2]
+    x = torch.nn.functional.pad(x, padding)
+    x_padded, _, _ = wan_spatial_shard.halo_exchange(
+        x, group=object(), halo_size=conv.halo_size, split_dim=conv.split_dim
+    )
+    return x_padded, x.shape[conv.split_tensor_dim]
+
+
+def _dist_conv(split_dim, kernel=(3, 3, 3), padding=(1, 1, 1)):
+    from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d
+
+    torch.manual_seed(0)
+    source = WanCausalConv3d(3, 4, kernel, padding=padding)
+    return source, wan_spatial_shard.WanDistCausalConv3d(source, object(), split_dim=split_dim)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["width", "height"])
+@pytest.mark.parametrize("with_cache", [False, True])
+def test_single_pass_conv_input_matches_cat_pad_halo_on_one_rank(monkeypatch, split_dim, with_cache):
+    """One rank, halo-bearing kernel: same bytes as the previous cat -> F.pad -> halo_exchange path."""
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
+    _, conv = _dist_conv(split_dim)
+    x = torch.randn(1, 3, 2, 6, 8)
+    cache_x = torch.randn(1, 3, 2, 6, 8) if with_cache else None
+
+    expected, expected_extent = _reference_conv_input(conv, x, cache_x)
+    got, got_extent = conv._assemble_input(x, cache_x)
+
+    assert got.shape == expected.shape and got_extent == expected_extent
+    assert torch.equal(got, expected)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_single_pass_conv_rejects_cache_longer_than_causal_padding():
+    _, conv = _dist_conv("width")
+    with pytest.raises(AssertionError, match="temporal cache exceeds"):
+        conv._assemble_input(torch.randn(1, 3, 2, 6, 8), torch.randn(1, 3, 3, 6, 8))
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["width", "height"])
+def test_single_pass_conv_without_a_halo_matches_the_stock_causal_conv(monkeypatch, split_dim):
+    """A kernel with no spatial extent along the split (the temporal conv) needs no halo and equals the stock conv."""
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
+    source, conv = _dist_conv(split_dim, kernel=(3, 1, 1), padding=(1, 0, 0))
+    x = torch.randn(1, 3, 2, 6, 8)
+    cache_x = torch.randn(1, 3, 2, 6, 8)
+
+    assert torch.equal(conv(x, cache_x), source(x, cache_x))
+    assert torch.equal(conv(x, None), source(x, None))
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["width", "height"])
+@pytest.mark.parametrize("rank", [0, 1])
+@pytest.mark.parametrize("kernel,padding", [((3, 3, 3), (1, 1, 1)), ((3, 1, 1), (0, 0, 0))])
+def test_single_pass_conv_input_matches_cat_pad_halo_on_two_ranks(monkeypatch, split_dim, rank, kernel, padding):
+    """Same bytes as cat -> F.pad -> halo_exchange, halos included, on either rank of a pair."""
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
+
+    def fake_p2p(*, rank, world_size, group, top_row_ref, bottom_row_ref, recv_top_buf, recv_bottom_buf):
+        # A neighbour's halo is whatever it sent; stand in with a marker derived from our own edge so
+        # both the reference and the single-pass path see the same exchange. Global edges stay zero.
+        recv_top_buf.copy_(top_row_ref * 0 + 7.0) if rank > 0 else recv_top_buf.zero_()
+        recv_bottom_buf.copy_(bottom_row_ref * 0 + 9.0) if rank < world_size - 1 else recv_bottom_buf.zero_()
+
+    monkeypatch.setattr(wan_spatial_shard, "_halo_exchange_p2p", fake_p2p)
+    _, conv = _dist_conv(split_dim, kernel=kernel, padding=padding)
+    x = torch.randn(1, 3, 1, 6, 8)
+    cache_x = torch.randn(1, 3, 2, 6, 8)
+
+    expected, expected_extent = _reference_conv_input(conv, x, cache_x)
+    got, got_extent = conv._assemble_input(x, cache_x)
+
+    assert got.shape == expected.shape and got_extent == expected_extent
+    assert torch.equal(got, expected)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
@@ -538,7 +629,7 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
 @pytest.mark.parametrize("split_dim", ["height", "width"])
 def test_spatial_shard_decode_matches_reference(split_dim: str):
     manager = mp.get_context("spawn").Manager()