Add comm split support for XCCL

.github/scripts/setup_env.sh

@@ -0,0 +1,92 @@
+#!/bin/bash
+# Common environment setup for CI jobs
+# Usage: source setup_env.sh [--with-cmake] [--cuda-version <version>] [--torch-version <version>] <torch-channel>
+#   --with-cmake: Install cmake and ninja-build
+#   --cuda-version: CUDA version (e.g., "12.8") - required for nightly builds
+#   --torch-version: Exact torch version to install (e.g., "2.6.0.dev20250101")
+#   torch-channel: "stable" or "nightly"
+
+set -ex
+
+INSTALL_CMAKE=false
+TORCH_CHANNEL=""
+CUDA_VERSION=""
+TORCH_VERSION=""
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --with-cmake)
+      INSTALL_CMAKE=true
+      shift
+      ;;
+    --cuda-version)
+      CUDA_VERSION="$2"
+      shift 2
+      ;;
+    --torch-version)
+      TORCH_VERSION="$2"
+      shift 2
+      ;;
+    *)
+      TORCH_CHANNEL="$1"
+      shift
+      ;;
+  esac
+done
+
+if [ -z "$TORCH_CHANNEL" ]; then
+  echo "Error: torch-channel argument required (stable or nightly)"
+  exit 1
+fi
+
+# Install system packages
+dnf config-manager --set-enabled powertools
+dnf install -y almalinux-release-devel
+
+if [ "$INSTALL_CMAKE" = true ]; then
+  dnf install -y ninja-build cmake
+  # Remove old cmake/ninja from conda/local
+  rm -f "/opt/conda/bin/ninja" || true
+  rm -f "/opt/conda/bin/cmake" || true
+  rm -f "/usr/local/bin/cmake" || true
+fi
+
+# Set up conda environment
+conda config --set solver libmamba
+conda create -n venv python=3.12 -y
+conda activate venv
+python -m pip install --upgrade pip
+
+# Nuke conda libstd++ to avoid conflicts with system toolset
+rm -f "$CONDA_PREFIX/lib/libstdc"* || true
+
+# Install torch (nightly or stable)
+if [ -n "$CUDA_VERSION" ]; then
+  # Convert CUDA version (e.g., "12.8") to PyTorch format (e.g., "cu128")
+  CUDA_TAG="cu$(echo "$CUDA_VERSION" | tr -d '.')"
+
+  if [ "$TORCH_CHANNEL" = "nightly" ]; then
+    INDEX_URL="https://download.pytorch.org/whl/nightly/${CUDA_TAG}"
+    if [ -n "$TORCH_VERSION" ]; then
+      pip install --pre torch=="${TORCH_VERSION}" --index-url "$INDEX_URL"
+    else
+      pip install --pre torch --index-url "$INDEX_URL"
+    fi
+  else
+    # Stable with CUDA - use PyTorch wheel index
+    INDEX_URL="https://download.pytorch.org/whl/${CUDA_TAG}"
+    if [ -n "$TORCH_VERSION" ]; then
+      pip install torch=="${TORCH_VERSION}" --index-url "$INDEX_URL"
+    else
+      pip install torch --index-url "$INDEX_URL"
+    fi
+  fi
+else
+  # No CUDA version - install from PyPI (stable only)
+  if [ -n "$TORCH_VERSION" ]; then
+    pip install torch=="${TORCH_VERSION}"
+  fi
+fi
+
+pip install -r requirements.txt

.github/scripts/xpu_test.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+set -ex
+# Install oneAPI DLE
+ONEAPI_URL="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/9065c156-58ab-41b0-bbee-9b0e229ffca5/intel-deep-learning-essentials-2025.3.1.15_offline.sh"
+wget -qO /tmp/intel-deep-learning-essentials.sh ${ONEAPI_URL}
+chmod +x /tmp/intel-deep-learning-essentials.sh
+/tmp/intel-deep-learning-essentials.sh -a --silent --eula accept
+rm -f /tmp/intel-deep-learning-essentials.sh
+
+export INTEL_ONEAPI=/opt/intel/oneapi
+source $INTEL_ONEAPI/compiler/latest/env/vars.sh
+source $INTEL_ONEAPI/ccl/latest/env/vars.sh
+
+conda create -yn xpu_torchcomms_ci python=3.10
+source activate xpu_torchcomms_ci
+conda install conda-forge::glog=0.4.0 conda-forge::gflags conda-forge::fmt -y
+
+export USE_XCCL=ON
+export USE_NCCL=OFF
+export USE_NCCLX=OFF
+export USE_TRANSPORT=OFF
+export USE_SYSTEM_LIBS=1
+
+python3 -m pip install --pre torch pytorch-triton-xpu --index-url https://download.pytorch.org/whl/nightly/xpu --force-reinstall --no-cache-dir
+cd torchcomms && pip install . --no-build-isolation && cd ..
+
+python3 -c "import torch; import torchcomms; print(f'Torch version: {torch.__version__}')"

.github/workflows/build_test.yaml

@@ -6,8 +6,12 @@ on:
       - main
   pull_request:
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
+
 jobs:
-  build_test:
+  build_pip_wheel:
     strategy:
       fail-fast: false
       matrix:
@@ -22,45 +26,95 @@ jobs:
             torch-version: "nightly"
 
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
     with:
-      timeout: 120
+      timeout: 60
       runner: ${{ matrix.runs-on }}
       gpu-arch-type: ${{ matrix.gpu-arch-type }}
       gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      upload-artifact: build-artifacts-${{ matrix.torch-version }}-${{ matrix.gpu-arch-type }}-${{ matrix.gpu-arch-version }}
       script: |
         set -ex
+        source .github/scripts/setup_env.sh --with-cmake --cuda-version "${{ matrix.gpu-arch-version }}" "${{ matrix.torch-version }}"
+
+        # Build wheel
+        pip install build
+        python -m build --wheel --no-isolation
 
-        # install build tools to system
-        dnf config-manager --set-enabled powertools
-        dnf install -y almalinux-release-devel
-        dnf install -y ninja-build cmake
+        mkdir -p "${RUNNER_ARTIFACT_DIR}"
+        cp dist/*.whl "${RUNNER_ARTIFACT_DIR}/"
 
-        # Remove old cmake/ninja
-        rm -f "/opt/conda/bin/ninja" || true
-        rm -f "/opt/conda/bin/cmake" || true
-        rm -f "/usr/local/bin/cmake" || true
+        python -c "import torch; print(torch.__version__)" > "${RUNNER_ARTIFACT_DIR}/torch_version.txt"
 
-        # use faster libmamba solver
-        conda config --set solver libmamba
+  build_and_run_cpp_tests:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - runs-on: "linux.g5.12xlarge.nvidia.gpu"
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.8"
+            torch-version: "stable"
+          - runs-on: "linux.g5.12xlarge.nvidia.gpu"
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.8"
+            torch-version: "nightly"
 
-        conda create -n venv python=3.12 -y
-        conda activate venv
-        python -m pip install --upgrade pip
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      timeout: 60
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      script: |
+        set -ex
+        source .github/scripts/setup_env.sh --with-cmake --cuda-version "${{ matrix.gpu-arch-version }}" "${{ matrix.torch-version }}"
 
-        # Nuke conda libstd++ to avoid conflicts with system toolset
-        rm -f "$CONDA_PREFIX/lib/libstdc"* || true
+        # Build and run C++ tests
+        cmake -B build_tests -G Ninja -DBUILD_TESTS=ON
+        cmake --build build_tests
+        cd build_tests && ctest --output-on-failure
 
-        # Optionally install torch nightly, pulls latest CUDA from pip otherwise
-        if [ "${{ matrix.torch-version }}" = "nightly" ]; then
-          pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-        fi
+  run_py_tests:
+    needs: build_pip_wheel
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - runs-on: "linux.g5.12xlarge.nvidia.gpu"
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.8"
+            torch-version: "stable"
+          - runs-on: "linux.g5.12xlarge.nvidia.gpu"
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.8"
+            torch-version: "nightly"
+
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      timeout: 120
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      download-artifact: build-artifacts-${{ matrix.torch-version }}-${{ matrix.gpu-arch-type }}-${{ matrix.gpu-arch-version }}
+      script: |
+        set -ex
+        TORCH_VERSION=$(cat "${RUNNER_ARTIFACT_DIR}/torch_version.txt")
+        source .github/scripts/setup_env.sh --cuda-version "${{ matrix.gpu-arch-version }}" --torch-version "$TORCH_VERSION" "${{ matrix.torch-version }}"
 
-        pip install -r requirements.txt
+        # Install from pre-built wheel (skip build step)
+        pip install "${RUNNER_ARTIFACT_DIR}"/*.whl pytest numpy psutil parameterized pydot requests urllib3 tabulate
 
-        # Build pybindings
-        pip install --no-build-isolation -v '.[dev]'
         python -c "import torchcomms; import torchcomms._transport"
 
-        # Run tests
+        # Run Python tests
+        comms/torchcomms/scripts/run_tests_unit_py.sh
         comms/torchcomms/scripts/run_tests_integration_py.sh
-        pytest -v comms/torchcomms/tests/unit/py

.github/workflows/build_test_rccl.yaml

@@ -0,0 +1,85 @@
+name: Build RCCL
+
+on:
+  push:
+    branches:
+      - main
+
+permissions:
+      id-token: write
+      contents: read
+
+jobs:
+  build:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: RCCL
+            runs-on: amd-mi350-runner
+            gpu-arch-type: "rocm"
+            gpu-arch-version: "7.0"
+            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/rocm7.0'
+            docker-image: pytorch/manylinux2_28-builder:rocm7.0
+            cmake-version: "latest"
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    with:
+      timeout: 180
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      script: |
+        set -ex
+        # use faster libmamba solver
+        conda config --set solver libmamba
+
+        # TODO: remove dependency on fbwhoami
+        echo "DEVICE_NAME=$(hostname)" > /etc/fbwhoami
+        cat /etc/fbwhoami
+
+        # For picking C++ 20 compiler from toolset-13
+        export PATH=/opt/rh/gcc-toolset-13/root/usr/bin:$PATH
+        export CXX=/opt/rh/gcc-toolset-13/root/usr/bin/c++
+        export CC=/opt/rh/gcc-toolset-13/root/usr/bin/cc
+
+        conda create -n venv python=3.12 -y
+        conda activate venv
+        python -m pip install --upgrade pip
+        conda install conda-forge::libopenssl-static conda-forge::rsync -y
+        conda install conda-forge::glog=0.4.0 conda-forge::gflags conda-forge::fmt -y
+
+        pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm7.0
+
+        if [ "${{ matrix.cmake-version }}" = "latest" ]; then
+          conda install -y cmake
+        else # default to latest
+          conda install -y cmake==${{ matrix.cmake-version }}
+        fi
+
+        pip install -r requirements.txt
+
+        export ROCM_HOME=/opt/rocm
+        export RCCL_INCLUDE=$ROCM_HOME/include/rccl
+
+        ./build_rccl.sh
+        pip install numpy
+        USE_SYSTEM_LIBS=1 USE_NCCL=0 USE_NCCLX=0 USE_GLOO=0 USE_RCCL=1 USE_TRANSPORT=OFF pip install --no-build-isolation -v -e .
+
+        # Verify installation - TODO re-enable this section after D90936340 stack lands
+        #python -c "import torch; print(f'PyTorch version: {torch.__version__}')"
+        #python -c "import torch; print(f'HIP available: {torch.cuda.is_available()}')"
+        #python -c "import torchcomms; print('TorchComms imported successfully')"
+
+        # Test RCCL backend availability
+        #python -c "
+        #import torchcomms
+        #try:
+        #    comm = torchcomms.new_comm('rccl', torch.device('hip'), 'test_comm')
+        #    print('RCCL backend available')
+        #except Exception as e:
+        #    print(f'RCCL backend test failed: {e}')
+        #"
+
+        # Run integration tests
+        #echo "Running RCCL integration tests..."
+        #comms/torchcomms/scripts/run_tests_integration_rccl_py.sh

.github/workflows/docs.yaml

@@ -22,8 +22,8 @@ jobs:
       - name: Setup conda env
         uses: conda-incubator/setup-miniconda@v2
         with:
-          auto-update-conda: true
           miniconda-version: "latest"
+          conda-version: "25.11.1"
           activate-environment: test
           python-version: '3.10'
           auto-activate: false
@@ -38,14 +38,14 @@ jobs:
         run: python -m pip install --upgrade pip
       - name: Install pytorch
         shell: bash -l {0}
-        run: pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+        run: pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
       - name: Install Dependencies
         shell: bash -l {0}
         run: |
           set -eux
           conda install -y git
           conda install -y -c conda-forge glog==0.4.0 gflags fmt
-          pip install cmake
+          pip install cmake==3.31.6
           pip install -r docs/requirements.txt
           export USE_NCCL=0
           export USE_NCCLX=0