Fused all-gather matmul on multi GPUs with PyTorch Symmetric Memory

[kwen2501]

Description

Each rank gathers inputs from (all) peer GPUs, and perform a matrix multiplication with its local weight.

Peer inputs are made visible via PyTorch Symmetric Memory, i.e.

import torch.distributed._symmetric_memor as symm_mem
symm_mem.empty(...)

The fused kernel is equivalent to:

dist.all_gather_into_tensor(ag_out, inp, group)
out = ag_out @ w

The fusion overlaps communication and computation in fine grain.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.