[Bugfix] Fix Wan spatial reshard boundary (#6062)

Signed-off-by: Rahul Steiger <rsteiger@nvl72d106-T18.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d106-T18.cm.cluster>

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -231,6 +231,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):

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -203,8 +203,17 @@ def reshard_from_trimmed_extent(
         return x
 
     dim = _spatial_dim(split_dim)
+    ctx = _SPATIAL_SHARD_CONTEXT.get()
+    if ctx is not None and ctx.rank != rank:
+        raise RuntimeError(
+            "Wan VAE spatial-shard rank mismatch while resharding: "
+            f"group_rank={rank}, context_rank={ctx.rank}, split_dim={split_dim!r}."
+        )
+
     valid_extent = _local_valid_extent(local_extent)
-    start = rank * local_extent
+    actual_extent = x.shape[dim]
+    start = min(rank * local_extent, actual_extent)
+    valid_extent = min(valid_extent, actual_extent - start)
     local = _narrow_along_dim(x, dim, start, valid_extent).contiguous()
     if valid_extent < local_extent:
         local = _pad_along_dim(local, local_extent - valid_extent, dim=dim)
@@ -594,10 +603,22 @@ def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_di
         dim = _spatial_dim(split_dim)
         local_extent = x.shape[dim]
         gathered = all_gather_along_dim(x, group=group, dim=dim).contiguous()
+        gathered_extent = gathered.shape[dim]
         full_extent = _current_full_extent(local_extent)
         if full_extent is not None:
-            gathered = _narrow_along_dim(gathered, dim, 0, full_extent).contiguous()
+            trim_extent = min(full_extent, gathered_extent)
+            gathered = _narrow_along_dim(gathered, dim, 0, trim_extent).contiguous()
+        else:
+            trim_extent = gathered_extent
         out = orig_forward(gathered, *args, **kwargs)
+        if out.shape[dim] != trim_extent:
+            raise RuntimeError(
+                "Wan VAE attention changed the spatial extent during global gather: "
+                f"rank={_rank_world(group)[0]}, split_dim={split_dim!r}, "
+                f"local_extent={local_extent}, gathered_extent={gathered_extent}, "
+                f"trimmed_extent={trim_extent}, output_extent={out.shape[dim]}, "
+                f"context={_SPATIAL_SHARD_CONTEXT.get()!r}."
+            )
         return reshard_from_trimmed_extent(out, local_extent=local_extent, split_dim=split_dim, group=group)
 
     module.forward = MethodType(_forward, module)