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

[lizhangyu258]

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

run

export NVSHMEM_HOME=/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem

# 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
## test script
./test-ag-gemm.sh

ag-gemm performance

config

• 1 node
• 8 GPUs, NVIDIA H800
• dtype: torch.float16
• allgather: copy engine
• gemm kernel
  • triton-distributed: autotune
  • flagtree: use triton-distributed best config
• The latency reported here is the average across all ranks.

FlagTree

(m_per_rank, n_per_rank, k)	GPUs	flagtree(ms)	torch (ms)	speedup
(1024, 1024, 1024)	8	0.418	0.257	0.615 x
(4096, 1024, 8192)	8	3.041	3.803	1.251 x
(4096, 8192, 8192)	8	7.043	9.270	1.316 x
(4096, 28672, 8192)	8	23.440	25.165	1.074 x
(4096, 8192, 28672)	8	23.450	32.963	1.406 x
(8192, 1024, 8192)	8	5.868	7.493	1.277 x
(8192, 8192, 8192)	8	13.783	18.052	1.310 x
(8192, 28672, 8192)	8	45.405	55.212	1.216 x
(8192, 8192, 28672)	8	46.205	66.344	1.436 x

Triton-distributed

# To profile triton-dist，please follow the script below
export NVSHMEM_HOME=/data/zyuli/miniconda3/envs/triton_dist/lib/python3.12/site-packages/nvidia/nvshmem
export NVSHMEM_SYMMETRIC_SIZE=10G

# For single node,  the script/launch.sh script needs to be modified.
Triton-distributed# git diff scripts/launch.sh
diff --git a/scripts/launch.sh b/scripts/launch.sh
index ac6b09d..f92510c 100644
--- a/scripts/launch.sh
+++ b/scripts/launch.sh
@@ -86,6 +86,10 @@ function check_nvshmem_bootstrap_uid_sock() {
   if [ -n "${NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME}" ]; then
     if ! check_if_interface_exists ${NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME}; then
       echo -e "${YELLOW}${BOLD}${WARN_ICON} WARNING: ${RESET}${BOLD}${RESET} NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=${NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME} does not exist..."
+      if [ "${nnodes:-1}" -eq 1 ] && check_if_interface_exists lo; then
+        echo "Single-node launch: use loopback interface lo for NVSHMEM bootstrap"
+        export NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=lo
+      fi
     fi
   fi

# Currently, we are using custom shapes here; additional tests will be added later.
# Need to modify how parameters are read.
## correctness
bash scripts/launch.sh --nproc_per_node=8 /data/zyuli/Triton-distributed/python/triton_dist/test/nvidia/test_ag_gemm.py --m-per-rank 1024 --n-per-rank 1024 --k 1024 --case check
## performance
bash scripts/launch.sh --nproc_per_node=8 /data/zyuli/Triton-distributed/python/triton_dist/test/nvidia/test_ag_gemm.py --m-per-rank 1024 --n-per-rank 1024 --k 1024 --case perf

(m_per_rank, n_per_rank, k)	GPUs	td autotune(ms)	torch(ms)	speedup
(1024, 1024, 1024)	8	0.557	0.326	0.585 x
(4096, 1024, 8192)	8	3.015	3.699	1.227 x
(4096, 8192, 8192)	8	6.368	9.048	1.421 x
(4096, 28672, 8192)	8	22.771	25.241	1.108 x
(4096, 8192, 28672)	8	23.467	32.639	1.391 x
(8192, 1024, 8192)	8	5.949	7.282	1.224 x
(8192, 8192, 8192)	8	13.003	17.979	1.383 x
(8192, 28672, 8192)	8	46.779	52.209	1.116 x
(8192, 8192, 28672)	8	47.092	66.463	1.411 x

[lizhangyu258]

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