[Feature]: GEMM + ReduceScatter

[mawad-amd]

Suggestion Description

Implement GEMM+ReduceScatter using Iris. See TritonDistributed example for reference.

Also see: https://github.com/vllm-project/vllm/blob/b78bed1bc5debead116092f429eee51398691fc8/vllm/compilation/collective_fusion.py#L53-L84

Implementation variants:

• Sequentially Fully-Fused (Atomic-based GEMM + ReduceScatter #169)

See comment here: #169 (comment)

Operating System

No response

GPU

No response

ROCm Component

No response

[astroC86]

i would like to try and have a go at this too?

[mawad-amd]

i would like to try and have a go at this too?

Yeah, this should be a fun one. Please give it a try after #99. Especially for this one, if you have any questions, feel free to create draft PRs, discuss here, or reach out via email. We would be happy to provide any guidance and brainstorm ideas if needed.

[mawad-amd]

@astroC86, just FYI, we have a PR for this now #169. If you already started working on this and have a different implementation, we would be happy to accept that as well.