One shot all reduce Example

[joydddd]

Stacked PRs:

• [Benchmark] Add all gather matmul benchmark #400
• Add all reduce benchmark #393
• [Pending multi device CI] Add symmetric memory sync test. #375
• ->One shot all reduce Example #245
• [BC breaking] Add MulticastTensor support to hl.signal & hl.wait (as_ptrs) #261
• Add multicast tensor #346

---

One shot all reduce Example

[jansel]

One shot All Reduce Performance on 8xH100.

Benchmark: joydddd/kraken@helion_bench

Performance drop at 1m due to hardcoded block_size.

shape dtype nccl helion_1shot triton_1shot dist_1shot Speedup over nccl Best Backend
(4k) torch.bfloat16 21.408 14.304 13.440 16.032 1.593 triton_1shot
(8k) torch.bfloat16 22.112 14.432 13.280 15.968 1.665 triton_1shot
(16k) torch.bfloat16 23.904 14.912 13.696 16.960 1.745 triton_1shot
(32k) torch.bfloat16 24.032 15.232 14.688 18.080 1.636 triton_1shot
(64k) torch.bfloat16 24.128 17.312 17.088 19.840 1.412 triton_1shot
(128k) torch.bfloat16 24.416 20.704 21.376 25.088 1.179 helion_1shot
(256k) torch.bfloat16 24.736 29.184 29.312 30.400 1.000 nccl
(512k) torch.bfloat16 34.016 47.488 42.016 43.328 1.000 nccl
(1m) torch.bfloat16 55.808 91.168 63.488 65.376 1.000 nccl

What explains the gap between triton and helion here? How are you autotuning?

[joydddd]

One shot All Reduce Performance on 8xH100.

Benchmark: joydddd/kraken@helion_bench
Performance drop at 1m due to hardcoded block_size.
shape dtype nccl helion_1shot triton_1shot dist_1shot Speedup over nccl Best Backend
(4k) torch.bfloat16 21.408 14.304 13.440 16.032 1.593 triton_1shot
(8k) torch.bfloat16 22.112 14.432 13.280 15.968 1.665 triton_1shot
(16k) torch.bfloat16 23.904 14.912 13.696 16.960 1.745 triton_1shot
(32k) torch.bfloat16 24.032 15.232 14.688 18.080 1.636 triton_1shot
(64k) torch.bfloat16 24.128 17.312 17.088 19.840 1.412 triton_1shot
(128k) torch.bfloat16 24.416 20.704 21.376 25.088 1.179 helion_1shot
(256k) torch.bfloat16 24.736 29.184 29.312 30.400 1.000 nccl
(512k) torch.bfloat16 34.016 47.488 42.016 43.328 1.000 nccl
(1m) torch.bfloat16 55.808 91.168 63.488 65.376 1.000 nccl

What explains the gap between triton and helion here? How are you autotuning?

No, we are not autotuning. I'd expect some performance boost once that is ready.

Helion is launching an additional kernel to turn a pointer to dev_ptr array into a tensor.

https://github.com/pytorch-labs/helion/pull/245/files#diff-4e027817dc836f1b7a352b698bd5cfdaf58189b065e007c80ff7fffc80383891R17

There is unfortunately no native way to get signal_pad/buffer pointers as a tensor from symmetric memory handlers now. If we decide this is what we'll need, we can talk to Ke Wen and see if they can add support.

[joydddd]

Workaround before symmetric memory handler natively supports getting dev_ptrs as a tensor, instead of a pointer to the array.

[jansel]

There is unfortunately no native way to get signal_pad/buffer pointers as a tensor from symmetric memory handlers now. If we decide this is what we'll need, we can talk to Ke Wen and see if they can add support.

What are the baselines doing? Passing the pointers as kernel args? If that is needed for perf I think we can make that work.

[joydddd]

There is unfortunately no native way to get signal_pad/buffer pointers as a tensor from symmetric memory handlers now. If we decide this is what we'll need, we can talk to Ke Wen and see if they can add support.

What are the baselines doing? Passing the pointers as kernel args? If that is needed for perf I think we can make that work.

The Triton baseline is passing
signal_pad_ptrs_dev - Device pointer to array of signal pad pointers _distributed_c10d.pyi:776
as a triton kernel arg, and then the number of signal pad.

If we want to keep the nice Helion abstraction of taking Tensors as arguments, we need to find a way to construct a Tensor out of shape, stride, & data_ptr directly without CPU sync or additional kernel launch.

Let me ask Natalia and Ke to see if they plan to add native support for getting the dev_ptrs array as a tensor.

[joydddd]

There is unfortunately no native way to get signal_pad/buffer pointers as a tensor from symmetric memory handlers now. If we decide this is what we'll need, we can talk to Ke Wen and see if they can add support.

What are the baselines doing? Passing the pointers as kernel args? If that is needed for perf I think we can make that work.

See #393 for a more complete benchmark result & analysis/ after the multicastTensor implementation.