[KMCompiler] [TLERaw]Triton v3.6.x support nvshmem

[lizhangyu258]

Preliminary Triton support for invoking NVSHMEM communication operations via external calls.

run

# 01-simple-shift
cd python/tutorials/tle/raw/nvshmem
## nvcc compiler
export NVCC_FLAGS="-I/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem/include"
export NVLINK_FLAGS="-L/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem/lib -lnvshmem_device"
./run.sh --launcher nvshmrun --np 4 --arch sm_90a 01-simple-shift/simple-shift.py
## clang compiler
export CLANG_FLAGS="-Xclang -mlink-bitcode-file -Xclang /data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem/lib/libnvshmem_device.bc"
./run.sh --launcher nvshmrun --np 4 --arch sm_90a 01-simple-shift/simple-shift.py


# 02-allgather-gemm
cd python/tutorials/tle/raw/nvshmem
export CLANG_FLAGS="-Xclang -mlink-bitcode-file -Xclang /data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem/lib/libnvshmem_device.bc"
## correctness
./run.sh --launcher torchrun --np 4 --arch sm_90a 02-allgather-gemm/ag-gemm.py --m-per-rank 4096 --chunk-m 4096 --n-per-rank 8192 --k 8192 --mode check
## performance
./run.sh --launcher torchrun --np 4 --arch sm_90a 02-allgather-gemm/ag-gemm.py --m-per-rank 4096 --chunk-m 4096 --n-per-rank 8192 --k 8192 --mode perf

allgather-gemm performance

config

• single node; 4 GPUs
• allgather:
  • chunk_m = local_m
  • grid = num_ranks - 1
  • config = {num_warps=32}
• gemm:
  • grid = (cdiv(local_m * num_ranks , BM) * cdiv(n, BN))
  • config = {"BM": 128, "BN": 256, "BK": 64, "stage": 3, "num_warps": 8}

• allgather
  • chunk_m = {128, 256, 512, 1024}
  • grid = (num_ranks - 1) * (local_m / chunk_m)
  • config = {num_warps=32}

[lizhangyu258]

Since the parent class is object, simply calling super().__init__() is sufficient.