[Question]: 09_gemm_one_shot_all_reduce implementation

[hebais]

Problem Description

Hello, I have some questions and doubts about examples/09_gemm_one_shot_all_reduce/matmul_wrapper.py

1. The matmul wrapper returns c instead of c_global, however from the kernel_gemm impl, the c_global should be the one who holds the final acc output, the code refered below in matmul_wrapper.py:

# starting from line 192
matmul._call(
            a=a,
            b=b,
            c=c,
            c_global=c_global, # c_global holds the acc output
            xxx
        )
return c # but here c is returned

2. Compared to gemm-allscatter varies implementations, the current one-shot gemm-allreduce impl seems a bit more like fused-sequential right ? I've tested it on some other platform and found that the performance is a bit worse than torch naive implementation. I consider several factors:
   1. Torch matmul is faster than triton matmul for some certain shape
   2. What is the side effect of frequent atomic_cas (during signal wait), does it affect heavy L1 intruction STALLs or cache pollution ?
   3. The new added reduce post process might affect the original gemm performance ?
      So I wonder what might be better furture move in terms of performance improvement from your perspective:
3. Add workgroup specializaiton to seperate CUs for compute/reduce
4. Unfused gemm_allscatter (producer) overlap with ring-reduce kernels (consumer), but I think one of the issues of multi-stream kernels is that the atomic-xxx wait loop in the consumer kernel might affect the producer kernel performance (because 2 streams might race for CUs).

I really would like to have a try to implement, however I could only try on ROCm compatible platforms instead. So I guess maybe some future helps about review and validation on AMD platforms

Operating System

Ubuntu 22.04

CPU

AMD Ryzen

GPU

AMD MI300

ROCm Version

6.3

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[mawad-amd]

Hi @hebais,

1. I will admit the examples need a bit of cleanup. It is true the GEMM module returns local C but we use the global one for validation. See:

iris/examples/09_gemm_one_shot_all_reduce/benchmark.py

Line 201 in 0aa9ba8

global_C,

and:

iris/examples/09_gemm_one_shot_all_reduce/benchmark.py

Line 261 in 0aa9ba8

success = validate_gemm(A, B, global_C, shmem, atol=2)

Ideally, the example would be restructured so we have something like:

def dist_gemm_all_reduce(A, B)
	# Do GEMM + AllReduce
	return C # Global C

2. Correct, this is a fused-sequential implementation. This implementation is not a fast one (unlike the all scatter ones).
   i. The GEMM will need the right tiling parameters to achieve good performance (I suggest looking at https://github.com/ROCm/tritonBLAS for that)
   ii. Synchronization here waits for all GPUs to finish one tile which may be expensive.
   iii. I am not sure what you are referring too here.
3. I would generally recommend looking into how to improve the reduction algorithm itself. For example, how to minimize the atomics/synchronization and the remote loads/stores. Exploring other variants of fusion is good too, but I don't think they will have high performance impact.
4. That is an interesting approach. You could potentially use multiple locks to synchronize between different phases.

We’ve been working on improving the performance of AllReduce, and @neoblizz has a great implementation in this pull request that uses ring all-reduce. I think there’s still room to push the performance further. As the team discussed potential solutions for AllReduce, we realized that it can be implemented in many interesting ways — and the more clever minds we have working on it, the better. Please share your thoughts and novel implementations!

I could only try on ROCm-compatible platforms instead.

I’d like to understand this better so we can help you get started. A few questions: do you currently have access to any AMD GPUs? Which GPUs (AMD/NVIDIA) do you have access to? Do you have access to the AMD Developer Cloud?

Feel free to reach out via email if you have any specific questions or details about hardware resources that you’d prefer to discuss privately.

Great questions btw!

[hebais]

Hi @mawad-amd , thanks for the help and suggestions. I will have a look at @neoblizz 's PR.
For point 2.III，I meant the extra acc part at the end of the persistent_gemm kernel:

# Store data to the global result using puts
acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)
for remote_rank in range(world_size):
    # might this load from the remote rank pollute the gemm load/store Cache ?
    # maybe a modifier ".cg" that bypass L1 is better ?
    acc += iris.load(C + sub_offset, cur_rank, remote_rank, heap_bases, mask=sub_mask)

tl.store(c_global + sub_offset, acc, mask=sub_mask, cache_modifier=".wt")

For the hardware platform, I found that the AMD Developer Cloud is currently out-of-resource. What I have is old amdgcn-like platforms. So if there is other env that I can rent that you know, will be very helpful for the validation by myself.

Thanks again