Doubt about repeated joint_tensor_query in all_to_all

[Edenzzzz]

It seems that in this line, the image conditioning for query is concatenated before all_to_all, which means it's repeated by ulysses_size times. I wonder if that might affect correctness?

xDiT/xfuser/core/long_ctx_attention/hybrid/attn_layer.py

Line 108 in 57eb27f

query = torch.cat([joint_tensor_query, query], dim=1)

[feifeibear]

It is absolutely correct. We have two different parallel for mm-dit SP.

1. split both text and image.
2. split image replicate text for small text or dim of text is not divisible by sp_degree.

What you see is 2.

[Edenzzzz]

Thanks for your reply.
In case 2, joint_tensor_query has full sequence length before all_to_all, doesn't that mean it's duplicated in a interleaved fashion with image embedding along the seq dim (text_cond, image_seq_1, test_cond, image_seq_2, ...)?
This also seems to prevent pack qkv (can only pack kv now) due to seqlen diff.
Would it be better to cat after all_to_all inside ring comp?

[feifeibear]

It does waste some communication bandwidth for dimension flexibility. I encourage you to give a try for supporting var len, which is definitely valuable by avoiding text padding.