[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):