[KMCompiler][TLERaw] Add multi-node support for AG-GEMM

[lizhangyu258]

Startup Guide

# torch config
export NCCL_SOCKET_IFNAME=bond0
export NCCL_IB_HCA='=mlx5_1:1'

# nvshmem config
export NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=bond0
export NVSHMEM_BOOTSTRAP_UID_SOCK_FAMILY=AF_INET
export NVSHMEM_REMOTE_TRANSPORT=ibrc
export NVSHMEM_ENABLE_NIC_PE_MAPPING=1
export NVSHMEM_HCA_LIST=mlx5_1:1

export NVSHMEM_IB_ENABLE_IBGDA=1
export NVSHMEM_IBGDA_NIC_HANDLER=auto

export NVSHMEM_IB_ADDR_FAMILY=AF_INET
export NVSHMEM_IB_ROCE_VERSION_NUM=2
export NVSHMEM_IB_ADDR_RANGE=0.0.0.0/0

# node-0
torchrun --node_rank=0 --nproc_per_node=4 --nnodes=2 --rdzv_endpoint=10.0.9.3:23456 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/benchmark.py  --M 8192 --iters 50 --warmup_iters 20 --dump_csv

# node-1
torchrun --node_rank=1 --nproc_per_node=4 --nnodes=2 --rdzv_endpoint=10.0.9.3:23456 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/benchmark.py  --M 8192 --iters 50 --warmup_iters 20 --dump_csv

Perf

[i3wanna2]

LGTM