[Feature] Transfer AR-to-diffusion payloads over NIXL (#6264)

Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -146,3 +146,34 @@ def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
             world_size=4,
             sp_group_ranks=[[0, 2]],
         )
+
+
+@pytest.mark.parametrize("rank", [1, 3])
+def test_tensor_dict_broadcast_preserves_subgroup_source(rank, monkeypatch):
+    from vllm_omni.diffusion.distributed.group_coordinator import GroupCoordinator, TensorMetadata
+
+    coordinator = object.__new__(GroupCoordinator)
+    coordinator.rank = rank
+    coordinator.ranks = [1, 3]
+    coordinator.rank_in_group = [1, 3].index(rank)
+    coordinator.world_size = 2
+    coordinator.device_group = object()
+    coordinator.cpu_group = object()
+    metadata = [("payload", TensorMetadata("cpu", torch.float32, (2,)))]
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+
+    def broadcast_object(value, src):
+        assert src == 0
+        return metadata
+
+    def broadcast(tensor, src, group, async_op):
+        assert src == 1
+        assert group is coordinator.cpu_group
+        assert async_op
+        tensor.fill_(7)
+        return SimpleNamespace(wait=lambda: None)
+
+    monkeypatch.setattr(coordinator, "broadcast_object", broadcast_object)
+    monkeypatch.setattr(torch.distributed, "broadcast", broadcast)
+    result = coordinator.broadcast_tensor_dict({"payload": torch.full((2,), 7.0)} if rank == 1 else None)
+    assert torch.equal(result["payload"], torch.full((2,), 7.0))