Merge branch 'main' into codex-vae-spatial-decode-auto

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -64,6 +64,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         local_rank,
         backend,
         group_name,
+        use_all2all=False,
     ):
         del backend
         group = _FakeGroup(
@@ -73,6 +74,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
             device_communicator=object(),
             reduce_scatter=lambda tensor, **kwargs: tensor,
         )
+        group.use_all2all = use_all2all
         created_groups.append(group)
         return group
 
@@ -121,6 +123,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
     assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
     assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
+    # vLLM 0.27's MoE oracle asserts the EP communicator has an all2all
+    # manager, which vLLM only builds for groups created with use_all2all.
+    assert parallel_state.vllm_parallel_state._EP.use_all2all is True
     assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
     assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
     assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
@@ -313,3 +318,35 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
     assert parallel_state.vllm_parallel_state._PCP is None
     assert parallel_state.vllm_parallel_state._EP is None
     assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_init_vllm_group_forwards_use_all2all_only_when_supported(monkeypatch):
+    """use_all2all reaches vLLM only on versions whose API accepts it.
+
+    vLLM 0.27 builds the EP all2all manager only for groups created with
+    use_all2all=True (its MoE oracle then asserts the manager exists —
+    build 2953, HunyuanImage3-DIT accuracy). Older vLLM has no such kwarg,
+    so forwarding it unconditionally would TypeError at init.
+    """
+    received = {}
+
+    def modern(group_ranks, local_rank, backend, group_name=None, use_device_communicator=True, use_all2all=False):
+        received["modern"] = use_all2all
+        return object()
+
+    monkeypatch.setattr(parallel_state.vllm_parallel_state, "init_model_parallel_group", modern)
+    parallel_state.init_vllm_model_parallel_group([[0]], 0, "gloo", "ep", use_all2all=True)
+    assert received["modern"] is True
+
+    parallel_state.init_vllm_model_parallel_group([[0]], 0, "gloo", "tp")
+    assert received["modern"] is False
+
+    def legacy(group_ranks, local_rank, backend, group_name=None, use_device_communicator=True):
+        received["legacy"] = True
+        return object()
+
+    monkeypatch.setattr(parallel_state.vllm_parallel_state, "init_model_parallel_group", legacy)
+    parallel_state.init_vllm_model_parallel_group([[0]], 0, "gloo", "ep", use_all2all=True)
+    assert received["legacy"] is True

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -251,6 +251,128 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon
     assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("split_dim", ["height", "width"])
+@pytest.mark.parametrize(
+    ("full_extent", "world_size"),
+    [
+        pytest.param(5, 4, id="four-ranks-empty-tail"),
+        pytest.param(20, 8, id="production-eight-ranks-empty-tail"),
+        pytest.param(39, 2, id="two-ranks-partial-tail"),
+    ],
+)
+def test_reshard_from_trimmed_extent_clamps_and_reconstructs(
+    monkeypatch: pytest.MonkeyPatch,
+    split_dim: str,
+    full_extent: int,
+    world_size: int,
+):
+    local_extent = (full_extent + world_size - 1) // world_size
+    shape = (1, 1, 1, full_extent, 1) if split_dim == "height" else (1, 1, 1, 1, full_extent)
+    x = torch.arange(full_extent, dtype=torch.float32).reshape(shape)
+    outputs = []
+
+    for rank in range(world_size):
+        monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group, rank=rank: (rank, world_size))
+        token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
+            wan_spatial_shard.SpatialShardContext(
+                input_extent=full_extent,
+                local_input_extent=local_extent,
+                split_dim=split_dim,
+                rank=rank,
+                world_size=world_size,
+            )
+        )
+        try:
+            out = wan_spatial_shard.reshard_from_trimmed_extent(
+                x,
+                local_extent=local_extent,
+                split_dim=split_dim,
+                group=object(),
+            )
+        finally:
+            wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
+
+        start = rank * local_extent
+        valid_extent = min(local_extent, max(0, full_extent - start))
+        expected = torch.zeros(local_extent, dtype=torch.float32)
+        if valid_extent:
+            expected[:valid_extent] = torch.arange(start, start + valid_extent, dtype=torch.float32)
+        assert out.shape[_spatial_dim_for_test(split_dim)] == local_extent
+        assert torch.equal(out.flatten(), expected)
+        outputs.append(out)
+
+    reconstructed = torch.cat(outputs, dim=_spatial_dim_for_test(split_dim))
+    reconstructed = reconstructed.narrow(_spatial_dim_for_test(split_dim), 0, full_extent)
+    assert torch.equal(reconstructed, x)
+
+
+def _spatial_dim_for_test(split_dim: str) -> int:
+    return -2 if split_dim == "height" else -1
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_reshard_from_trimmed_extent_rejects_rank_mismatch(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
+    x = torch.zeros((1, 1, 1, 1, 4), dtype=torch.float32)
+    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
+        wan_spatial_shard.SpatialShardContext(
+            input_extent=4,
+            local_input_extent=2,
+            split_dim="width",
+            rank=0,
+            world_size=2,
+        )
+    )
+    try:
+        with pytest.raises(RuntimeError, match="group_rank=1, context_rank=0"):
+            wan_spatial_shard.reshard_from_trimmed_extent(
+                x,
+                local_extent=2,
+                split_dim="width",
+                group=object(),
+            )
+    finally:
+        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.MonkeyPatch):
+    class ShrinkingWanAttentionBlock(torch.nn.Module):
+        def forward(self, x):
+            return x[..., :-1]
+
+    monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2))
+    monkeypatch.setattr(
+        wan_spatial_shard,
+        "all_gather_along_dim",
+        lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
+    )
+    module = ShrinkingWanAttentionBlock()
+    wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
+    x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
+    token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
+        wan_spatial_shard.SpatialShardContext(
+            input_extent=4,
+            local_input_extent=2,
+            split_dim="width",
+            rank=0,
+            world_size=2,
+        )
+    )
+    try:
+        with pytest.raises(
+            RuntimeError,
+            match=r"local_extent=2, gathered_extent=4, trimmed_extent=4, output_extent=3",
+        ):
+            module(x)
+    finally:
+        wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):