diff --git a/.azure-pipelines/codecov.yml b/.azure-pipelines/codecov.yml index c4abeaa7..a806445e 100644 --- a/.azure-pipelines/codecov.yml +++ b/.azure-pipelines/codecov.yml @@ -38,6 +38,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -59,6 +61,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -80,6 +84,8 @@ jobs: matrix: rocm6_2: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2 + rocm7_2: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2 container: image: $(containerImage) diff --git a/.azure-pipelines/integration-test.yml b/.azure-pipelines/integration-test.yml index d5d5f9bd..859e2a68 100644 --- a/.azure-pipelines/integration-test.yml +++ b/.azure-pipelines/integration-test.yml @@ -19,21 +19,21 @@ pr: drafts: false paths: exclude: - - .devcontainer/** - - .github/** - - docker/** - - docs/** - - '**/*.md' + - .devcontainer/** + - .github/** + - docker/** + - docs/** + - '**/*.md' jobs: - job: IntegrationTestA100 displayName: Integration test A100 strategy: matrix: - cuda11: - containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda11.8 cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 pool: name: msccl-ci @@ -43,9 +43,9 @@ jobs: steps: - template: templates/integration-test.yml parameters: - subscription: mscclpp-ci - vmssName: mscclpp-ci - gpuArch: '80' + subscription: mscclpp-ci + vmssName: mscclpp-ci + gpuArch: '80' - job: IntegrationTestH100 displayName: Integration test H100 @@ -53,6 +53,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 pool: name: msccl-ci-h100 @@ -62,7 +64,7 @@ jobs: steps: - template: templates/integration-test.yml parameters: - subscription: mscclpp-ci-h100 - vmssName: mscclpp-h100-ci + subscription: mscclpp-ci-h100 + vmssName: mscclpp-h100-ci perfBaselineFile: test/deploy/perf_ndmv5.jsonl - gpuArch: '90' + gpuArch: '90' diff --git a/.azure-pipelines/multi-nodes-test.yml b/.azure-pipelines/multi-nodes-test.yml index 3b3ebe1f..c4e27be4 100644 --- a/.azure-pipelines/multi-nodes-test.yml +++ b/.azure-pipelines/multi-nodes-test.yml @@ -14,7 +14,6 @@ trigger: # Do not run multi-nodes-test for PR, we can trigger it manually pr: none - parameters: - name: vmssName type: string @@ -32,6 +31,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 pool: name: mscclpp-multi-node container: @@ -79,10 +80,10 @@ jobs: - template: templates/deploy.yml parameters: - subscription: mscclpp-ci-h100 - vmssName: ${{ parameters.vmssName }} + subscription: mscclpp-ci-h100 + vmssName: ${{ parameters.vmssName }} resourceGroup: mscclpp - gpuArch: '90' + gpuArch: '90' - template: templates/run-remote-task.yml parameters: @@ -119,6 +120,6 @@ jobs: - template: templates/stop.yml parameters: - subscription: mscclpp-ci-h100 - vmssName: ${{ parameters.vmssName }} + subscription: mscclpp-ci-h100 + vmssName: ${{ parameters.vmssName }} resourceGroup: mscclpp diff --git a/.azure-pipelines/nccl-api-test.yml b/.azure-pipelines/nccl-api-test.yml index 85b466ef..a9fe6b25 100644 --- a/.azure-pipelines/nccl-api-test.yml +++ b/.azure-pipelines/nccl-api-test.yml @@ -35,6 +35,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -56,6 +58,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) diff --git a/.azure-pipelines/rccl-api-test.yml b/.azure-pipelines/rccl-api-test.yml index 43841079..fc793c88 100644 --- a/.azure-pipelines/rccl-api-test.yml +++ b/.azure-pipelines/rccl-api-test.yml @@ -35,6 +35,8 @@ jobs: matrix: rocm6_2: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2 + rocm7_2: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2 container: image: $(containerImage) diff --git a/.azure-pipelines/sglang-multi-node-test.yml b/.azure-pipelines/sglang-multi-node-test.yml new file mode 100644 index 00000000..01fcbb4f --- /dev/null +++ b/.azure-pipelines/sglang-multi-node-test.yml @@ -0,0 +1,143 @@ +# ============================================================================= +# Multi-node SGLang integration test pipeline. +# +# This pipeline runs MSCCL++ SGLang tests across two H100 VMSS GPU nodes. +# High-level flow: +# 1. The pipeline agent runs inside a container on the `mscclpp-multi-node` +# pool. The agent itself has no GPUs. +# 2. SSH/host configuration is generated so the agent can reach the two +# pre-provisioned VMSS GPU nodes. +# 3. `templates/deploy.yml` builds and ships MSCCL++ to the GPU nodes. +# 4. `templates/sglang-multi-test.yml` runs the SGLang multi-node tests. +# 5. `templates/stop.yml` tears down / stops the VMSS nodes. +# +# Docs / non-code changes are excluded from triggering this pipeline. +# ============================================================================= + +trigger: + branches: + include: + - main + - release/* + paths: + exclude: + - .devcontainer/** + - .github/** + - docker/** + - docs/** + - '**/*.md' + +pr: + branches: + include: + - main + - release/* + drafts: false + paths: + exclude: + - .devcontainer/** + - .github/** + - docker/** + - docs/** + - '**/*.md' + +parameters: +# Name of the pre-provisioned Azure VMSS that hosts the GPU test nodes. +# Node hostnames are derived as "${vmssName}000000" and "${vmssName}000001". +- name: vmssName + type: string + default: mscclpp-h100-multinode-ci +# Static /etc/hosts entries mapping VMSS node hostnames to their private IPs. +# These IPs are tied to the specific VMSS above; update both together if the +# VMSS is reprovisioned or renamed. +- name: hostEntries + type: string + default: | + 10.0.0.5 mscclpp-h100-multinode-ci000000 + 10.0.0.4 mscclpp-h100-multinode-ci000001 +# Docker image used for the SGLang test container on the GPU nodes. +- name: sglangImage + type: string + default: lmsysorg/sglang:latest-cu129 + +jobs: +- job: SGLangTestMultiNode + displayName: SGLang Test Multi Node + strategy: + matrix: + cuda12: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 + pool: + name: mscclpp-multi-node + container: + image: $(containerImage) + + steps: + # Ensure the VMSS node hostnames resolve from the pipeline agent container. + # Idempotent: only appends lines that are not already present in /etc/hosts. + - task: Bash@3 + displayName: Add HostEntry + inputs: + targetType: 'inline' + script: | + while IFS= read -r line; do + [ -z "$line" ] && continue + if ! grep -qxF "$line" /etc/hosts; then + echo "Adding to /etc/hosts: $line" + echo "$line" | sudo tee -a /etc/hosts + else + echo "Entry already exists: $line" + fi + done <<< "${{ parameters.hostEntries }}" + + # Generate the SSH config and hostfile consumed by the deploy / test + # templates below: + # - config : SSH client config (custom port + key) for each node + # - hostfile : user@host list used by deploy / test scripts (parallel-ssh) + - task: Bash@3 + displayName: Generate deploy files + inputs: + targetType: 'inline' + script: | + set -e + VMSS="${{ parameters.vmssName }}" + DEPLOY_DIR="$(System.DefaultWorkingDirectory)/test/deploy" + NODE0="${VMSS}000000" + NODE1="${VMSS}000001" + + echo "Host ${NODE0} + Port 22345 + IdentityFile /root/mscclpp/sshkey + StrictHostKeyChecking no + Host ${NODE1} + Port 22345 + IdentityFile /root/mscclpp/sshkey + StrictHostKeyChecking no" > "${DEPLOY_DIR}/config" + + printf '%s\n%s\n' "azureuser@${NODE0}" "azureuser@${NODE1}" > "${DEPLOY_DIR}/hostfile" + + # Build MSCCL++ and deploy it onto the VMSS GPU nodes. + - template: templates/deploy.yml + parameters: + subscription: mscclpp-ci-h100 + vmssName: ${{ parameters.vmssName }} + resourceGroup: mscclpp + gpuArch: '90' + deployArgs: 'multi-node-test true cuda' + containerName: 'sglang-mscclpp-test' + sglangImage: ${{ parameters.sglangImage }} + + # Run the SGLang multi-node tests across the two GPU nodes. + - template: templates/sglang-multi-test.yml + parameters: + subscription: mscclpp-ci-h100 + vmssName: ${{ parameters.vmssName }} + + # Stop/deallocate the VMSS GPU nodes to release resources. + - template: templates/stop.yml + parameters: + subscription: mscclpp-ci-h100 + vmssName: ${{ parameters.vmssName }} + resourceGroup: mscclpp diff --git a/.azure-pipelines/sglang-test.yml b/.azure-pipelines/sglang-test.yml new file mode 100644 index 00000000..0b49a0ce --- /dev/null +++ b/.azure-pipelines/sglang-test.yml @@ -0,0 +1,65 @@ +# ============================================================================= +# Single-node SGLang integration test pipeline. +# +# Runs MSCCL++ SGLang tests on a single H100 GPU node from the `msccl-ci-h100` +# pool. All deploy / run / teardown logic is delegated to +# `templates/sglang-test.yml`. +# +# Docs / non-code changes are excluded from triggering this pipeline. +# ============================================================================= + +trigger: + branches: + include: + - main + - release/* + paths: + exclude: + - .devcontainer/** + - .github/** + - docker/** + - docs/** + - '**/*.md' + +pr: + branches: + include: + - main + - release/* + drafts: false + paths: + exclude: + - .devcontainer/** + - .github/** + - docker/** + - docs/** + - '**/*.md' + +parameters: +# Docker image used for the SGLang test container on the GPU node. +- name: sglangImage + type: string + default: lmsysorg/sglang:latest-cu129 + +jobs: +- job: SGLangTest + displayName: SGLang Test + strategy: + matrix: + cuda12: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 + pool: + name: msccl-ci-h100 + container: + image: $(containerImage) + + steps: + # Deploy MSCCL++ to the GPU node and run the SGLang single-node tests. + - template: templates/sglang-test.yml + parameters: + subscription: mscclpp-ci-h100 + vmssName: mscclpp-h100-ci + gpuArch: '90' + sglangImage: ${{ parameters.sglangImage }} diff --git a/.azure-pipelines/templates/deploy.yml b/.azure-pipelines/templates/deploy.yml index 2f642f1d..9eb46d8b 100644 --- a/.azure-pipelines/templates/deploy.yml +++ b/.azure-pipelines/templates/deploy.yml @@ -32,6 +32,12 @@ parameters: - name: deployArgs type: string default: '' +- name: containerName + type: string + default: 'mscclpp-test' +- name: sglangImage + type: string + default: '' steps: # 0. Ensure Azure CLI exists before running AzureCLI@2 tasks. @@ -147,5 +153,5 @@ steps: inputs: targetType: filePath filePath: test/deploy/deploy.sh - arguments: ${{ parameters.deployArgs }} + arguments: ${{ parameters.deployArgs }} ${{ parameters.containerName }} ${{ parameters.sglangImage }} workingDirectory: '$(System.DefaultWorkingDirectory)' diff --git a/.azure-pipelines/templates/integration-test.yml b/.azure-pipelines/templates/integration-test.yml index b686e4f2..ad95cbc2 100644 --- a/.azure-pipelines/templates/integration-test.yml +++ b/.azure-pipelines/templates/integration-test.yml @@ -15,7 +15,7 @@ steps: subscription: ${{ parameters.subscription }} vmssName: ${{ parameters.vmssName }} gpuArch: ${{ parameters.gpuArch }} - deployArgs: 'single-node-test' + deployArgs: 'single-node-test true cuda' - template: run-remote-task.yml parameters: diff --git a/.azure-pipelines/templates/nccl-test.yml b/.azure-pipelines/templates/nccl-test.yml index fa3900f1..550f5690 100644 --- a/.azure-pipelines/templates/nccl-test.yml +++ b/.azure-pipelines/templates/nccl-test.yml @@ -23,7 +23,7 @@ steps: subscription: ${{ parameters.subscription }} vmssName: ${{ parameters.vmssName }} gpuArch: ${{ parameters.gpuArch }} - deployArgs: 'nccltest-single-node' + deployArgs: 'nccltest-single-node true cuda' - template: run-remote-task.yml parameters: @@ -74,6 +74,15 @@ steps: mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN -x MSCCLPP_ENABLE_NCCL_FALLBACK=TRUE -x MSCCLPP_NCCL_LIB_PATH=/root/nccl/build/lib/libnccl.so -x MSCCLPP_FORCE_NCCL_FALLBACK_OPERATION="broadcast" /root/nccl-tests/build/broadcast_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN -x MSCCLPP_ENABLE_NCCL_FALLBACK=TRUE -x MSCCLPP_NCCL_LIB_PATH=/root/nccl/build/lib/libnccl.so -x MSCCLPP_FORCE_NCCL_FALLBACK_OPERATION="allreduce" /root/nccl-tests/build/broadcast_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 +- template: run-remote-task.yml + parameters: + name: PyBench + displayName: Run Collective Benchmarks + remoteScript: | + mpirun --allow-run-as-root -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3b15 --accum-type float32 --autotune --symmetric-memory + mpirun --allow-run-as-root -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3fn --accum-type float16 --autotune --symmetric-memory + mpirun --allow-run-as-root -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float16 --symmetric-memory --autotune + - template: stop.yml parameters: subscription: ${{ parameters.subscription }} diff --git a/.azure-pipelines/templates/rccl-test.yml b/.azure-pipelines/templates/rccl-test.yml index 8e247161..65f1c984 100644 --- a/.azure-pipelines/templates/rccl-test.yml +++ b/.azure-pipelines/templates/rccl-test.yml @@ -46,16 +46,25 @@ steps: name: RunRcclAllGatherTest displayName: Run RCCL AllGather Test with or without MSCCLPP Lib remoteScript: | - mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 - mpirun -np 8 --bind-to numa --allow-run-as-root /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 + mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 + mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 /root/rocm-systems/projects/rccl-tests/build/all_gather_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 - template: run-remote-task.yml parameters: name: RunRcclAllReduceTest displayName: Run RCCL AllReduce Test with or without MSCCLPP Lib remoteScript: | - mpirun -np 8 --bind-to numa --allow-run-as-root -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 - mpirun -np 8 --bind-to numa --allow-run-as-root /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 + mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x LD_PRELOAD=/root/mscclpp/build/lib/libmscclpp_nccl.so -x MSCCLPP_NCCL_SYMMETRIC_MEMORY=1 -x NCCL_DEBUG=WARN /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 + mpirun -np 8 --bind-to numa --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 /root/rocm-systems/projects/rccl-tests/build/all_reduce_perf -b 1K -e 1G -f 2 -d half -G 20 -w 10 -n 20 + +- template: run-remote-task.yml + parameters: + name: PyBench + displayName: Run Collective Benchmarks + remoteScript: | + mpirun --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3b15 --accum-type float32 --autotune + mpirun --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allreduce --dtype float8_e4m3fnuz --accum-type float32 --autotune + mpirun --allow-run-as-root -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m mscclpp_benchmark.bench_collective --collective allgather --dtype float8_e4m3b15 --autotune --buffer-mode out-of-place - template: stop.yml parameters: diff --git a/.azure-pipelines/templates/sglang-multi-test.yml b/.azure-pipelines/templates/sglang-multi-test.yml new file mode 100644 index 00000000..80e72926 --- /dev/null +++ b/.azure-pipelines/templates/sglang-multi-test.yml @@ -0,0 +1,95 @@ +# ============================================================================= +# SGLang multi-node test template. +# +# Runs on the pipeline agent and dispatches remote steps to the two VMSS GPU +# nodes (via run-remote-task.yml + the SSH config / hostfile produced by the +# caller pipeline). Steps: +# 1. Build and install MSCCL++ on each node. +# 2. Install a (currently forked) SGLang on each node, replacing any +# pre-baked copy from the base image. +# 3. Run a 2-node sglang.bench_one_batch smoke test with MSCCL++ enabled. +# 4. Run the MSCCL++ all-reduce micro-benchmark via torchrun across both +# nodes. +# ============================================================================= + +parameters: +- name: subscription + type: string +- name: vmssName + type: string +- name: containerName + type: string + default: 'sglang-mscclpp-test' + +steps: +# TODO: Switch to the official upstream sglang repo once Caio's PR is merged. +# Tracking: the fork below (`caiomcbr/sglang` @ caiorocha/mscclpp) is a personal +# branch and should not remain a long-term CI dependency. +- template: run-remote-task.yml + parameters: + name: InstallSGLang + displayName: Install SGLang + runRemoteArgs: '--container ${{ parameters.containerName }} --hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --user azureuser' + remoteScript: | + git clone -b main https://github.com/caiomcbr/sglang.git + cd sglang/python + pip install -e . + +# Smoke test: 2-node tensor-parallel benchmark of Qwen3-8B with MSCCL++. +# Port 20003 is the SGLang distributed-init rendezvous port (arbitrary, must +# match across ranks and be free on node 0). +- template: run-remote-task.yml + parameters: + name: RunSGLangMultiBenchOneBatch + displayName: Run SGLang Multi-Node Bench One Batch + runRemoteArgs: '--container ${{ parameters.containerName }} --hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --user azureuser' + remoteScript: | + export FLASHINFER_DISABLE_VERSION_CHECK=1 + VMSS="${{ parameters.vmssName }}" + HOSTNAME=$(hostname) + # Explicit 2-node mapping: hostname suffix -> SGLang node rank. + if [ "$HOSTNAME" = "${VMSS}000000" ]; then + NODE_RANK=0 + elif [ "$HOSTNAME" = "${VMSS}000001" ]; then + NODE_RANK=1 + else + echo "Unknown hostname: $HOSTNAME" + exit 1 + fi + python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 1 2 4 8 16 32 64 128 256 512 --input-len 256 --output-len 256 --tp-size 16 --dist-init-addr ${VMSS}000000:20003 --nnodes 2 --node-rank $NODE_RANK --enable-mscclpp + +# Depends on the `sglang/` source tree cloned by the InstallSGLang step above +# (steps on the same remote share a working directory). +- template: run-remote-task.yml + parameters: + name: RunSGLangMultiTestAllReduce + displayName: Run SGLang Multi-Node Test All Reduce + runRemoteArgs: '--container ${{ parameters.containerName }} --hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --user azureuser' + remoteScript: | + export FLASHINFER_DISABLE_VERSION_CHECK=1 + VMSS="${{ parameters.vmssName }}" + HOSTNAME=$(hostname) + # Explicit 2-node mapping: hostname suffix -> torchrun node rank. + if [ "$HOSTNAME" = "${VMSS}000000" ]; then + NODE_RANK=0 + elif [ "$HOSTNAME" = "${VMSS}000001" ]; then + NODE_RANK=1 + else + echo "Unknown hostname: $HOSTNAME" + exit 1 + fi + + export NODE_SIZE=2 + export WORLD_SIZE=8 + + cd sglang + + # Port 20004 is the torchrun rendezvous port (arbitrary, must match + # across ranks and be free on node 0). Distinct from 20003 used by + # sglang.bench_one_batch above. + torchrun --nproc_per_node $WORLD_SIZE \ + --nnodes $NODE_SIZE \ + --node_rank $NODE_RANK \ + --master_addr ${VMSS}000000 \ + --master_port 20004 \ + benchmark/kernels/all_reduce/benchmark_mscclpp.py diff --git a/.azure-pipelines/templates/sglang-test.yml b/.azure-pipelines/templates/sglang-test.yml new file mode 100644 index 00000000..0d663b71 --- /dev/null +++ b/.azure-pipelines/templates/sglang-test.yml @@ -0,0 +1,87 @@ +# ============================================================================= +# SGLang single-node test template. +# +# Runs on the pipeline agent and dispatches remote steps to a single VMSS GPU +# node (via run-remote-task.yml). Steps: +# 1. Deploy: build the test container and bring the VMSS node online. +# 2. Build and install MSCCL++ on the node. +# 3. Install a (currently forked) SGLang. +# 4. Run sglang.bench_one_batch at several batch sizes. +# 5. Run a longer end-to-end validation: bring up an sglang server and +# drive it with sglang.bench_serving. +# 6. Run the MSCCL++ all-reduce micro-benchmark via torchrun. +# 7. Stop / deallocate the VMSS node. +# ============================================================================= + +parameters: +- name: subscription + type: string +- name: vmssName + type: string +- name: gpuArch + type: string +- name: containerName + type: string + default: 'sglang-mscclpp-test' +- name: sglangImage + type: string + default: 'lmsysorg/sglang:latest' + +steps: +# deployArgs positional fields: +- template: deploy.yml + parameters: + subscription: ${{ parameters.subscription }} + vmssName: ${{ parameters.vmssName }} + gpuArch: ${{ parameters.gpuArch }} + deployArgs: 'single-node-test true cuda' + containerName: ${{ parameters.containerName }} + sglangImage: ${{ parameters.sglangImage }} + +# TODO: Switch to the official upstream sglang repo once Caio's PR is merged. +# Tracking: the fork below (`caiomcbr/sglang` @ caiorocha/mscclpp) is a personal branch and +# should not remain a long-term CI dependency. Also consider pinning to a +# release branch or commit SHA for reproducibility. +- template: run-remote-task.yml + parameters: + name: InstallSGLang + displayName: Install SGLang + runRemoteArgs: '--container ${{ parameters.containerName }}' + remoteScript: | + git clone -b main https://github.com/caiomcbr/sglang.git + cd sglang/python + pip install -e . + +- template: run-remote-task.yml + parameters: + name: RunSGLangBenchOneBatch + displayName: Run SGLang Bench One Batch + runRemoteArgs: '--container ${{ parameters.containerName }}' + remoteScript: | + export FLASHINFER_DISABLE_VERSION_CHECK=1 + python -m sglang.bench_one_batch --model-path Qwen/Qwen3-8B --batch 1 2 4 8 16 32 64 128 256 512 --input-len 256 --output-len 256 --tp-size 8 --enable-mscclpp + +# Depends on the `sglang/` source tree cloned by the InstallSGLang step above +# (steps on the same remote share a working directory). +- template: run-remote-task.yml + parameters: + name: RunSGLangTestAllReduce + displayName: Run SGLang Test All Reduce + runRemoteArgs: '--container ${{ parameters.containerName }}' + remoteScript: | + export FLASHINFER_DISABLE_VERSION_CHECK=1 + export NODE_SIZE=1 + export WORLD_SIZE=8 + export RANK=0 + + cd sglang + + torchrun --nproc_per_node $WORLD_SIZE \ + --nnodes $NODE_SIZE \ + --node_rank $RANK \ + benchmark/kernels/all_reduce/benchmark_mscclpp.py + +- template: stop.yml + parameters: + subscription: ${{ parameters.subscription }} + vmssName: ${{ parameters.vmssName }} diff --git a/.azure-pipelines/templates/ut-no-ib-env.yml b/.azure-pipelines/templates/ut-no-ib-env.yml index a62f1a77..cc7d2018 100644 --- a/.azure-pipelines/templates/ut-no-ib-env.yml +++ b/.azure-pipelines/templates/ut-no-ib-env.yml @@ -13,7 +13,7 @@ steps: vmssName: ${{ parameters.vmssName }} gpuArch: ${{ parameters.gpuArch }} cmakeArgs: '-DMSCCLPP_USE_IB=OFF' - deployArgs: 'single-node-test false' + deployArgs: 'single-node-test false cuda' - template: run-remote-task.yml parameters: diff --git a/.azure-pipelines/templates/ut-npkit.yml b/.azure-pipelines/templates/ut-npkit.yml index 1bd89caf..18934e6b 100644 --- a/.azure-pipelines/templates/ut-npkit.yml +++ b/.azure-pipelines/templates/ut-npkit.yml @@ -14,7 +14,7 @@ steps: vmssName: ${{ parameters.vmssName }} gpuArch: ${{ parameters.gpuArch }} cmakeArgs: '-DMSCCLPP_NPKIT_FLAGS="-DENABLE_NPKIT -DENABLE_NPKIT_EVENT_TIME_SYNC_CPU -DENABLE_NPKIT_EVENT_TIME_SYNC_GPU -DENABLE_NPKIT_EVENT_EXECUTOR_INIT_ENTRY -DENABLE_NPKIT_EVENT_EXECUTOR_INIT_EXIT -DENABLE_NPKIT_EVENT_EXECUTOR_OP_BASE_ENTRY -DENABLE_NPKIT_EVENT_EXECUTOR_OP_BASE_EXIT"' - deployArgs: 'single-node-test' + deployArgs: 'single-node-test true cuda' - template: run-remote-task.yml parameters: diff --git a/.azure-pipelines/templates/ut.yml b/.azure-pipelines/templates/ut.yml index 743c66e6..81a74cf8 100644 --- a/.azure-pipelines/templates/ut.yml +++ b/.azure-pipelines/templates/ut.yml @@ -40,8 +40,8 @@ steps: name: PyTests displayName: Run pytests remoteScript: | - mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_mscclpp.py -x - mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_fp8_accum.py -x + mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_mscclpp.py -x + mpirun --allow-run-as-root -tag-output -x MSCCLPP_HOME=/root/mscclpp -x HSA_NO_SCRATCH_RECLAIM=1 -x GPU_MAX_HW_QUEUES=8 -np 8 python3 -m pytest ./python/test/test_fp8_accum.py -x - template: stop.yml parameters: diff --git a/.azure-pipelines/ut.yml b/.azure-pipelines/ut.yml index 6b8c9eda..2a65f361 100644 --- a/.azure-pipelines/ut.yml +++ b/.azure-pipelines/ut.yml @@ -34,10 +34,10 @@ jobs: name: msccl-ci strategy: matrix: - cuda11: - containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda11.8 cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -55,10 +55,10 @@ jobs: name: msccl-ci strategy: matrix: - cuda11: - containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda11.8 cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -78,6 +78,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -97,6 +99,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -118,6 +122,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) @@ -137,6 +143,8 @@ jobs: matrix: rocm6_2: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2 + rocm7_2: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2 container: image: $(containerImage) @@ -159,6 +167,8 @@ jobs: matrix: cuda12: containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 + cuda13: + containerImage: ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda13.0 container: image: $(containerImage) diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index a5d6bf23..5e41a01b 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -32,7 +32,7 @@ "/usr/include" ], "C_Cpp.default.cStandard": "c17", - "C_Cpp.default.cppStandard": "c++17" + "C_Cpp.default.cppStandard": "c++20" } } }, diff --git a/.devcontainer/devcontainer_amd.json b/.devcontainer/devcontainer_amd.json index 80d47956..7f1e2c2b 100644 --- a/.devcontainer/devcontainer_amd.json +++ b/.devcontainer/devcontainer_amd.json @@ -32,7 +32,7 @@ "/usr/include" ], "C_Cpp.default.cStandard": "c17", - "C_Cpp.default.cppStandard": "c++17" + "C_Cpp.default.cppStandard": "c++20" } } }, diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml index fb065141..b9d3a5c1 100644 --- a/.github/workflows/codeql-analysis.yml +++ b/.github/workflows/codeql-analysis.yml @@ -40,7 +40,7 @@ jobs: fail-fast: false matrix: language: [ 'cpp', 'python' ] - version: [ 'cuda11.8', 'cuda12.9' ] + version: [ 'cuda12.9', 'cuda13.0' ] steps: - name: Checkout repository @@ -85,7 +85,7 @@ jobs: fail-fast: false matrix: language: [ 'cpp', 'python' ] - version: [ 'rocm6.2' ] + version: [ 'rocm6.2', 'rocm7.2' ] steps: - name: Checkout repository diff --git a/.github/workflows/mscclpp-lang.yml b/.github/workflows/mscclpp-lang.yml index 7368cfdf..95bee5c9 100644 --- a/.github/workflows/mscclpp-lang.yml +++ b/.github/workflows/mscclpp-lang.yml @@ -15,7 +15,7 @@ jobs: strategy: fail-fast: false matrix: - version: [ 'cuda11.8', 'cuda12.9' ] + version: [ 'cuda12.9', 'cuda13.0' ] steps: - uses: actions/checkout@v4 diff --git a/CMakeLists.txt b/CMakeLists.txt index 02180fc1..87f10efb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -60,6 +60,7 @@ option(MSCCLPP_BYPASS_GPU_CHECK "Bypass GPU check." OFF) option(MSCCLPP_NPKIT_FLAGS "Set NPKIT flags" OFF) option(MSCCLPP_ENABLE_COVERAGE "Enable code coverage" OFF) option(MSCCLPP_DISABLE_NB_LEAK_WARNINGS "Disable Nanobind leak warnings" ON) +option(MSCCLPP_ROCM_USE_FNUZ_FP8 "Use ROCm FNUZ native FP8 types." ON) set(MSCCLPP_GPU_ARCHS "" CACHE STRING "Specify GPU architectures with delimiters (comma, space, or semicolon).") if(MSCCLPP_BYPASS_GPU_CHECK) @@ -112,11 +113,20 @@ if(MSCCLPP_ENABLE_COVERAGE) if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") message(STATUS "Code coverage enabled") - # Add coverage flags to C++ targets only (not CUDA) - add_compile_options($<$:--coverage>) + # Add coverage flags to C++ targets only (not CUDA). For ROCm, keep + # coverage on the host compilation so HIP device linking does not look + # for gcov runtime symbols. + if(MSCCLPP_USE_ROCM) + add_compile_options($<$:-Xarch_host>) + add_compile_options($<$:--coverage>) + add_link_options($<$:-Xarch_host>) + add_link_options($<$:--coverage>) + else() + add_compile_options($<$:--coverage>) + add_link_options($<$:--coverage>) + endif() add_compile_options($<$:-O0>) add_compile_options($<$:-g>) - add_link_options($<$:--coverage>) # Find lcov find_program(LCOV_PATH lcov) @@ -169,8 +179,8 @@ elseif(MSCCLPP_USE_CUDA) if(NVIDIA_FOUND) set(MSCCLPP_GPU_ARCHS "native") else() - if(CUDAToolkit_VERSION VERSION_LESS "11.8") - message(FATAL_ERROR "CUDA 11.8 or higher required, found ${CUDAToolkit_VERSION}") + if(CUDAToolkit_VERSION VERSION_LESS "12.0") + message(FATAL_ERROR "CUDA 12.0 or higher required (C++20 build), found ${CUDAToolkit_VERSION}") endif() set(MSCCLPP_GPU_ARCHS 80) if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.0") @@ -184,16 +194,26 @@ elseif(MSCCLPP_USE_CUDA) endif() endif() elseif(MSCCLPP_USE_ROCM) - set(MSCCLPP_GPU_ARCHS gfx90a gfx941 gfx942) + set(MSCCLPP_GPU_ARCHS gfx90a gfx942) endif() message(STATUS "GPU architectures: ${MSCCLPP_GPU_ARCHS}") +if(MSCCLPP_USE_ROCM) + if(MSCCLPP_ROCM_USE_FNUZ_FP8) + add_compile_definitions(MSCCLPP_ROCM_FP8_FNUZ) + message(STATUS "ROCm native FP8 aliases: FNUZ (MSCCLPP_ROCM_USE_FNUZ_FP8=ON)") + else() + message(STATUS "ROCm native FP8 aliases: OCP (MSCCLPP_ROCM_USE_FNUZ_FP8=OFF)") + endif() +endif() + # Declare project -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra") if(MSCCLPP_USE_CUDA) - set(CMAKE_CUDA_STANDARD 17) + set(CMAKE_CUDA_STANDARD 20) + set(CMAKE_CUDA_EXTENSIONS OFF) set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wall,-Wextra") enable_language(CUDA) @@ -207,8 +227,12 @@ if(MSCCLPP_USE_CUDA) else() set(GPU_LIBRARIES CUDA::cudart CUDA::cuda_driver) endif() + if(NOT TARGET CUDA::nvml) + message(FATAL_ERROR "CUDA NVML target CUDA::nvml is required for MSCCLPP CUDA builds.") + endif() + list(APPEND GPU_LIBRARIES CUDA::nvml) else() - set(CMAKE_HIP_STANDARD 17) + set(CMAKE_HIP_STANDARD 20) set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Wall -Wextra") set(CMAKE_HIP_ARCHITECTURES ${MSCCLPP_GPU_ARCHS}) @@ -217,6 +241,14 @@ else() set(GPU_INCLUDE_DIRS ${hip_INCLUDE_DIRS}) endif() +message(STATUS "C++ standard: C++${CMAKE_CXX_STANDARD}") +if(MSCCLPP_USE_CUDA) + message(STATUS "CUDA toolkit version: ${CUDAToolkit_VERSION}") + message(STATUS "CUDA language standard: C++${CMAKE_CUDA_STANDARD}") +else() + message(STATUS "HIP language standard: C++${CMAKE_HIP_STANDARD}") +endif() + if(CMAKE_BUILD_TYPE STREQUAL "Debug") add_compile_definitions(DEBUG_BUILD) endif() diff --git a/docker/base-dev-x.dockerfile b/docker/base-dev-x.dockerfile index 47436202..89ea1ec7 100644 --- a/docker/base-dev-x.dockerfile +++ b/docker/base-dev-x.dockerfile @@ -70,7 +70,7 @@ RUN if echo "$TARGET" | grep -q "^cuda"; then \ # Install ROCm-specific packages if building for ROCm RUN if echo "$TARGET" | grep -q "^rocm"; then \ apt-get update -y && \ - apt-get install -y hipblas hipsparse rocsparse rocrand hiprand rocthrust rocsolver rocfft hipfft hipcub rocprim rccl roctracer-dev && \ + apt-get install -y hipblas hipsparse rocsparse rocrand hiprand rocthrust-dev rocsolver rocfft hipfft hipcub-dev rocprim-dev rccl roctracer-dev && \ apt-get autoremove -y && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* /tmp/*; \ @@ -84,7 +84,7 @@ ENV PATH="/root/venv/bin:${PATH}" # Install Python dependencies ADD . /tmp/mscclpp WORKDIR /tmp/mscclpp -RUN target_type=$(echo $TARGET | sed 's/\.[0-9]*$//') && \ +RUN target_type=$(echo "$TARGET" | sed -E 's/^([[:alpha:]]+[0-9]+).*/\1/') && \ if echo "$TARGET" | grep -q "^rocm"; then \ export CUPY_INSTALL_USE_HIP=1 && export ROCM_HOME=/opt/rocm; \ fi && \ diff --git a/docker/build.sh b/docker/build.sh index 651a6122..a65d0867 100755 --- a/docker/build.sh +++ b/docker/build.sh @@ -4,34 +4,35 @@ set -e declare -A baseImageTable baseImageTable=( - ["cuda11.8"]="nvidia/cuda:11.8.0-devel-ubuntu22.04" ["cuda12.4"]="nvidia/cuda:12.4.1-devel-ubuntu22.04" ["cuda12.8"]="nvidia/cuda:12.8.1-devel-ubuntu22.04" ["cuda12.9"]="nvidia/cuda:12.9.1-devel-ubuntu24.04" ["cuda13.0"]="nvidia/cuda:13.0.2-devel-ubuntu24.04" ["rocm6.2"]="rocm/dev-ubuntu-22.04:6.2.2" + ["rocm7.2"]="rocm/dev-ubuntu-24.04:7.2.4" ) declare -A extraLdPathTable extraLdPathTable=( ["rocm6.2"]="/opt/rocm/lib" + ["rocm7.2"]="/opt/rocm/lib" ) declare -A ofedVersionTable ofedVersionTable=( - ["cuda11.8"]="23.07-0.5.1.2" ["cuda12.4"]="23.07-0.5.1.2" ["cuda12.8"]="24.10-1.1.4.0" ["cuda12.9"]="24.10-1.1.4.0" ["cuda13.0"]="24.10-3.2.5.0" ["rocm6.2"]="24.10-1.1.4.0" + ["rocm7.2"]="24.10-3.2.5.0" ) TARGET=${1} OS_ARCH=$(uname -m) print_usage() { - echo "Usage: $0 [cuda11.8|cuda12.4|cuda12.8|cuda12.9|cuda13.0|rocm6.2]" + echo "Usage: $0 [cuda12.4|cuda12.8|cuda12.9|cuda13.0|rocm6.2|rocm7.2]" } if [[ ! -v "baseImageTable[${TARGET}]" ]]; then @@ -75,6 +76,7 @@ docker build -t ${TAG_BASE_DEV} \ --build-arg BASE_IMAGE=${TAG_BASE} \ --build-arg TARGET=${TARGET} . + GHCR="ghcr.io/microsoft/mscclpp/mscclpp" GHCR_TAG_BASE_DEV=${GHCR}:base-dev-${TARGET} GHCR_TAG_BASE_DEV_ARCH=${GHCR}:base-dev-${TARGET}-${OS_ARCH} @@ -107,4 +109,4 @@ echo "" echo " docker buildx imagetools create \\" echo " --tag ${GHCR_TAG_BASE_DEV} \\" echo " ${GHCR_TAG_BASE_DEV_ARCH}" -echo "" +echo "" \ No newline at end of file diff --git a/docs/cpp_api.rst b/docs/cpp_api.rst index 6b5f24c0..a7ebaaf9 100644 --- a/docs/cpp_api.rst +++ b/docs/cpp_api.rst @@ -128,6 +128,8 @@ Utilities .. doxygenclass:: mscclpp::GpuBuffer :members: +.. doxygenenum:: mscclpp::GpuBufferGranularity + .. doxygenclass:: mscclpp::GpuStream :members: diff --git a/docs/quickstart.md b/docs/quickstart.md index f9b3d0a7..9c7f6060 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -14,10 +14,15 @@ * [NDm_A100_v4](https://learn.microsoft.com/en-us/azure/virtual-machines/ndm-a100-v4-series) * [ND_H100_v5](https://learn.microsoft.com/en-us/azure/virtual-machines/nd-h100-v5-series) * Non-Azure Systems - * NVIDIA A100 GPUs + CUDA >= 11.8 + * NVIDIA A100 GPUs + CUDA >= 12.0 * NVIDIA H100 GPUs + CUDA >= 12.0 * AMD MI250X GPUs + ROCm >= 5.7 * AMD MI300X GPUs + ROCm >= 6.0 +* Toolchain + * MSCCL++ is built as **C++20** (both host and device code), so a C++20-capable toolchain is required. + * [CMake](https://cmake.org/) >= 3.25 + * A C++20-capable host compiler, e.g., GCC >= 11 or Clang >= 14 + * On NVIDIA platforms, **CUDA Toolkit >= 12.0** is required. `nvcc` first added `-std=c++20` support in CUDA 12.0, so earlier toolkits (11.x and below) cannot build the project. * OS * Tested on Ubuntu 20.04 and later * Libraries @@ -47,7 +52,7 @@ We provide docker images which package all prerequisites for MSCCL++. You can se # For NVIDIA platforms $ docker run -it --privileged --net=host --ipc=host --gpus all --name mscclpp-dev ghcr.io/microsoft/mscclpp/mscclpp:base-dev-cuda12.9 bash # For AMD platforms -$ docker run -it --privileged --net=host --ipc=host --security-opt=seccomp=unconfined --group-add=video --name mscclpp-dev ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm6.2 bash +$ docker run -it --privileged --net=host --ipc=host --security-opt=seccomp=unconfined --group-add=video --name mscclpp-dev ghcr.io/microsoft/mscclpp/mscclpp:base-dev-rocm7.2 bash ``` See all available images [here](https://github.com/microsoft/mscclpp/pkgs/container/mscclpp%2Fmscclpp). @@ -108,19 +113,16 @@ $ python -m pip install ".[cuda12]" # For NVIDIA platforms with the Expert Parallel (MoE dispatch/combine) extension $ python -m pip install ".[cuda12,ep]" # For AMD platforms -$ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm6]" +$ CXX=/opt/rocm/bin/hipcc python -m pip install ".[rocm7]" ``` -> **Note:** A platform extra (`cuda11`, `cuda12`, `cuda13`, or `rocm6`) is required to install CuPy. -> The CUDA extras install pre-built CuPy wheels. The `rocm6` extra installs CuPy from source, -> which requires ROCm and may take longer. Running `pip install .` without an extra will not install CuPy. +> **Note:** A platform extra (`cuda12`, `cuda13`, `rocm6`, or `rocm7`) is required to install CuPy. +> The CUDA extras install pre-built CuPy wheels and CUDA Python bindings. The ROCm extras install CuPy from source +> and HIP Python for the matching ROCm major version, which require ROCm and may take longer. Running `pip install .` without an extra will not install CuPy. Optional extras can be installed by specifying them in brackets. Available extras: -- **`cuda11`**, **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package for your CUDA version. -- **`rocm6`**: Install CuPy from source for AMD ROCm platforms. -- **`ep`**: Build and install the Expert Parallel extension (`mscclpp.ep`). The extension itself does - not add a PyTorch dependency, but the high-level Python API expects user-provided `torch.Tensor` inputs. - CUDA architectures 90 or newer are required for EP kernels. +- **`cuda12`**, **`cuda13`**: Install a pre-built CuPy package and CUDA Python bindings for your CUDA version. +- **`rocm6`**, **`rocm7`**: Install CuPy from source and HIP Python for AMD ROCm platforms. - **`benchmark`**: Install benchmark dependencies (mpi4py, prettytable, netifaces, matplotlib). - **`test`**: Install test dependencies (pytest, mpi4py, netifaces). @@ -216,15 +218,37 @@ $ mpirun -np 16 -npernode 8 -hostfile hostfile ./bin/mp_unit_tests -ip_port 10.0 ## Performance Benchmark -### Python Benchmark +### Python Benchmark and Tuning -[Install the MSCCL++ Python package](#install-from-source-python-module) and run our Python AllReduce benchmark as follows. It requires MPI on the system. +[Install the MSCCL++ Python package](#install-from-source-python-module) and run the Python collective benchmark as follows. It requires MPI on the system. ```bash # Install with benchmark dependencies and the appropriate CUDA/ROCm extras. -# Replace `cuda12` with your platform: cuda11, cuda12, cuda13, or rocm6. +# Replace `cuda12` with your platform: cuda12, cuda13, rocm6, or rocm7. $ python3 -m pip install ".[cuda12,benchmark,test]" -$ mpirun -tag-output -np 8 python3 ./python/mscclpp_benchmark/allreduce_bench.py + +``` + +To autotune launch parameters and save a tuned config: + +```bash +$ PYTHONPATH=$PWD/python mpirun -np 8 --allow-run-as-root \ + python3 -m mscclpp_benchmark.bench_collective \ + --collective allreduce \ + --dtype float16 \ + --batch-sizes 1,2,4,8 \ + --autotune \ + --write-config /tmp/mscclpp_tuned_configs.json +``` + +Use the tuned config in a benchmark: + +```bash +$ PYTHONPATH=$PWD/python mpirun -np 8 --allow-run-as-root \ + python3 -m mscclpp_benchmark.bench_collective \ + --collective allreduce \ + --dtype float16 \ + --config-path /tmp/mscclpp_tuned_configs.json ``` (nccl-benchmark)= diff --git a/examples/torch-integration/customized_comm_with_tuning.py b/examples/torch-integration/customized_comm_with_tuning.py index 060a0097..b96087c2 100644 --- a/examples/torch-integration/customized_comm_with_tuning.py +++ b/examples/torch-integration/customized_comm_with_tuning.py @@ -70,12 +70,12 @@ class CustomizedComm: _TUNE_N_WARMUP = 5 _TUNE_N_GRAPH_LAUNCHES = 10 _TUNE_N_OPS_PER_GRAPH = 100 - _CANDIDATE_NBLOCKS = [4, 8, 16, 24, 32, 48, 64, 128] + _CANDIDATE_NBLOCKS = [4, 8, 16, 24, 32, 48, 56, 64, 128] _CANDIDATE_NTHREADS = [512, 768, 1024] _NBLOCKS_LIMIT = { "default_allreduce_nvls_packet": 16, "default_allreduce_packet": 56, - "default_allreduce_allpair_packet": 56, + "default_allreduce_allpair_packet": 64, "default_allreduce_fullmesh": 64, "default_allgather_fullmesh2": 32, } diff --git a/include/mscclpp/core.hpp b/include/mscclpp/core.hpp index c8da79df..71ab303e 100644 --- a/include/mscclpp/core.hpp +++ b/include/mscclpp/core.hpp @@ -46,6 +46,10 @@ class Bootstrap { /// @return The total number of ranks per node. virtual int getNranksPerNode() const = 0; + /// Return the number of ranks in this rank's GPU IPC domain. + /// @return The number of ranks in the GPU IPC domain. + virtual int getNranksPerIpcDomain() const; + /// Send arbitrary data to another process. /// /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size, @@ -144,6 +148,9 @@ class TcpBootstrap : public Bootstrap { /// Return the total number of ranks per node. int getNranksPerNode() const override; + /// Return the number of ranks in this rank's GPU IPC domain. + int getNranksPerIpcDomain() const override; + /// Send arbitrary data to another process. /// /// Data sent via `send(senderBuff, size, receiverRank, tag)` can be received via `recv(receiverBuff, size, diff --git a/include/mscclpp/gpu_data_types.hpp b/include/mscclpp/gpu_data_types.hpp index b9d7191a..cd8b0526 100644 --- a/include/mscclpp/gpu_data_types.hpp +++ b/include/mscclpp/gpu_data_types.hpp @@ -16,8 +16,7 @@ using __bfloat16 = __hip_bfloat16; using __bfloat162 = __hip_bfloat162; #define __CUDA_BF16_TYPES_EXIST__ -// AMD FP8 support - Use fnuz types for HIP 6.0 or when HIP_FP8_TYPE_FNUZ is enabled and HIP_FP8_TYPE_OCP is not -// enabled. Otherwise, use the standard FP8 types. +// AMD FP8 support. The CMake MSCCLPP_ROCM_USE_FNUZ_FP8 option controls whether native FP8 aliases use FNUZ. #if defined(HIP_VERSION_MAJOR) && (HIP_VERSION_MAJOR >= 6) #include @@ -25,7 +24,7 @@ using __bfloat162 = __hip_bfloat162; // Define __FP8_E4M3_IS_FNUZ__ / __FP8_E5M2_IS_FNUZ__ when the platform-native FP8 is the // "fnuz" variant (no infinities, NaN-only at 0x80, bias differs from OCP). Dispatch layers // use these macros to throw on unsupported variants requested via DataType. -#if (HIP_VERSION_MAJOR == 6) || (HIP_VERSION_MAJOR > 6 && HIP_FP8_TYPE_FNUZ && !HIP_FP8_TYPE_OCP) +#if defined(MSCCLPP_ROCM_FP8_FNUZ) using __fp8_e4m3 = __hip_fp8_e4m3_fnuz; using __fp8_e5m2 = __hip_fp8_e5m2_fnuz; using __fp8x2_e4m3 = __hip_fp8x2_e4m3_fnuz; @@ -77,7 +76,7 @@ using __bfloat162 = __nv_bfloat162; /// Software float8 with 4 exponent bits, 3 mantissa bits, exponent bias = 15. /// Format (MSB first): [sign:1][exponent:4][mantissa:3] -/// No infinities, no NaN. Encode saturates to ±1.75 (0x7e/0xfe). +/// No infinities, no NaN. Encode saturates to ±1.875 (0x7f/0xff). /// Adapted from the Triton compiler's fp8e4b15 format. struct alignas(1) __fp8_e4m3b15 { uint8_t __x; @@ -109,7 +108,7 @@ struct alignas(1) __fp8_e4m3b15 { /// then convert fp16 → float32. static MSCCLPP_HOST_DEVICE_INLINE float toFloat(uint8_t bits) { // Branch-free decode: fp8 → fp16 → fp32, no special-case handling. - // Encode saturates to ±1.75, so 0x7f/0xff are never produced. + // Every byte maps to a finite value; encode saturates at ±1.875, so 0x7f/0xff decode to ±1.875. // Refer: // https://github.com/triton-lang/triton/blob/cf34004b8a67d290a962da166f5aa2fc66751326/python/triton/language/extra/cuda/utils.py#L34 uint16_t h = (uint16_t)bits << 8; // place fp8 in upper byte of fp16 @@ -138,10 +137,9 @@ struct alignas(1) __fp8_e4m3b15 { } cvt = {h_val}; uint16_t fp16_bits = cvt.u; - // Clamp abs to max encodable value: 1.75 → fp16 = 0x3F00. - // Matches Triton: encode saturates, 0x7f/0xff are never produced. + // Clamp abs to max encodable value: 1.875 → fp16 = 0x3F80 (largest byte 0x7f/0xff). uint16_t abs_fp16 = fp16_bits & 0x7FFFu; - if (abs_fp16 > 0x3F00u) abs_fp16 = 0x3F00u; + if (abs_fp16 > 0x3F80u) abs_fp16 = 0x3F80u; // Reconstruct with sign. uint16_t sign16 = fp16_bits & 0x8000u; @@ -858,27 +856,17 @@ MSCCLPP_DEVICE_INLINE f32x4 to(const f8_e5m2x4& v) { /// f32x2 -> f8_e4m3x2. /// HIP gfx942: float -> fp8 (via __builtin_amdgcn_cvt_pk_fp8_f32). -/// NVIDIA SM90+: float -> half -> fp8 (via __nv_cvt_halfraw2_to_fp8x2). -/// NVIDIA pre-SM90: float -> half -> fp8 (via __nv_cvt_halfraw_to_fp8, element-wise). +/// NVIDIA: float -> fp8 directly (via __nv_cvt_float2_to_fp8x2). On SM89+ this maps to a +/// single hardware round-to-nearest-even instruction; on older arch it falls back to a +/// software direct conversion. template <> MSCCLPP_DEVICE_INLINE f8_e4m3x2 to(const f32x2& v) { #if defined(MSCCLPP_DEVICE_HIP) && defined(__gfx942__) uint32_t packed = __builtin_amdgcn_cvt_pk_fp8_f32(v.data[0], v.data[1], 0, false); return bit_cast(static_cast<__hip_fp8x2_storage_t>(packed)); -#elif defined(MSCCLPP_DEVICE_CUDA) && __CUDA_ARCH__ >= 900 - __half2_raw h2; - h2.x = bit_cast(__float2half_rn(v.data[0])); - h2.y = bit_cast(__float2half_rn(v.data[1])); - __nv_fp8x2_storage_t fp8x2 = __nv_cvt_halfraw2_to_fp8x2(h2, __NV_SATFINITE, __NV_E4M3); - return bit_cast(fp8x2); #elif defined(MSCCLPP_DEVICE_CUDA) - __half_raw h0, h1; - h0.x = bit_cast(__float2half_rn(v.data[0])); - h1.x = bit_cast(__float2half_rn(v.data[1])); - f8_e4m3x2 result; - result.data[0] = bit_cast<__fp8_e4m3>(__nv_cvt_halfraw_to_fp8(h0, __NV_SATFINITE, __NV_E4M3)); - result.data[1] = bit_cast<__fp8_e4m3>(__nv_cvt_halfraw_to_fp8(h1, __NV_SATFINITE, __NV_E4M3)); - return result; + __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(make_float2(v.data[0], v.data[1]), __NV_SATFINITE, __NV_E4M3); + return bit_cast(fp8x2); #else f8_e4m3x2 result; result.data[0] = static_cast<__fp8_e4m3>(v.data[0]); @@ -915,27 +903,17 @@ MSCCLPP_DEVICE_INLINE f8_e4m3x4 to(const f32x4& v) { /// f32x2 -> f8_e5m2x2. /// HIP gfx942: float -> bf8 (via __builtin_amdgcn_cvt_pk_bf8_f32). -/// NVIDIA SM90+: float -> half -> fp8 (via __nv_cvt_halfraw2_to_fp8x2 with __NV_E5M2). -/// NVIDIA pre-SM90: float -> half -> fp8 (via __nv_cvt_halfraw_to_fp8, element-wise). +/// NVIDIA: float -> fp8 directly (via __nv_cvt_float2_to_fp8x2 with __NV_E5M2). On SM89+ this +/// maps to a single hardware round-to-nearest-even instruction; on older arch it falls back to a +/// software direct conversion. template <> MSCCLPP_DEVICE_INLINE f8_e5m2x2 to(const f32x2& v) { #if defined(MSCCLPP_DEVICE_HIP) && defined(__gfx942__) uint32_t packed = __builtin_amdgcn_cvt_pk_bf8_f32(v.data[0], v.data[1], 0, false); return bit_cast(static_cast<__hip_fp8x2_storage_t>(packed)); -#elif defined(MSCCLPP_DEVICE_CUDA) && __CUDA_ARCH__ >= 900 - __half2_raw h2; - h2.x = bit_cast(__float2half_rn(v.data[0])); - h2.y = bit_cast(__float2half_rn(v.data[1])); - __nv_fp8x2_storage_t fp8x2 = __nv_cvt_halfraw2_to_fp8x2(h2, __NV_SATFINITE, __NV_E5M2); - return bit_cast(fp8x2); #elif defined(MSCCLPP_DEVICE_CUDA) - __half_raw h0, h1; - h0.x = bit_cast(__float2half_rn(v.data[0])); - h1.x = bit_cast(__float2half_rn(v.data[1])); - f8_e5m2x2 result; - result.data[0] = bit_cast<__fp8_e5m2>(__nv_cvt_halfraw_to_fp8(h0, __NV_SATFINITE, __NV_E5M2)); - result.data[1] = bit_cast<__fp8_e5m2>(__nv_cvt_halfraw_to_fp8(h1, __NV_SATFINITE, __NV_E5M2)); - return result; + __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(make_float2(v.data[0], v.data[1]), __NV_SATFINITE, __NV_E5M2); + return bit_cast(fp8x2); #else f8_e5m2x2 result; result.data[0] = static_cast<__fp8_e5m2>(v.data[0]); @@ -1109,11 +1087,11 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x2 to(const f16x2& v) { #if defined(MSCCLPP_DEVICE_CUDA) uint32_t in0; asm("mov.b32 %0, %1;" : "=r"(in0) : "r"(*reinterpret_cast(&v))); - // Clamp abs to max encodable e4m3b15 (0x3F00 = 1.75 in fp16). + // Clamp abs to max encodable e4m3b15 (0x3F80 = 1.875 in fp16). uint32_t lo = in0 & 0xFFFFu, hi = in0 >> 16; uint32_t alo = lo & 0x7FFFu, ahi = hi & 0x7FFFu; - alo = alo < 0x3F00u ? alo : 0x3F00u; - ahi = ahi < 0x3F00u ? ahi : 0x3F00u; + alo = alo < 0x3F80u ? alo : 0x3F80u; + ahi = ahi < 0x3F80u ? ahi : 0x3F80u; uint32_t a0 = alo | (ahi << 16); a0 = a0 * 2u + 0x00800080u; uint32_t b0 = a0 | (in0 & 0x80008000u); @@ -1124,7 +1102,7 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x2 to(const f16x2& v) { uint32_t in0 = v.words[0]; uint32_t abs0 = in0 & 0x7fff7fffu; uint32_t a0; - asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F003F00u)); + asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F803F80u)); a0 = a0 * 2u + 0x00800080u; uint32_t b0 = a0 | (in0 & 0x80008000u); uint16_t packed = (uint16_t)(((b0 >> 8) & 0xFFu) | ((b0 >> 16) & 0xFF00u)); @@ -1147,8 +1125,8 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x4 to(const f16x4& v) { asm("mov.b32 %0, %1;" : "=r"(in1) : "r"(v.words[1])); uint32_t abs0 = in0 & 0x7fff7fffu; uint32_t abs1 = in1 & 0x7fff7fffu; - uint32_t a0 = __vminu2(abs0, 0x3F003F00u); - uint32_t a1 = __vminu2(abs1, 0x3F003F00u); + uint32_t a0 = __vminu2(abs0, 0x3F803F80u); + uint32_t a1 = __vminu2(abs1, 0x3F803F80u); a0 = a0 * 2u + 0x00800080u; a1 = a1 * 2u + 0x00800080u; uint32_t b0, b1; @@ -1161,8 +1139,8 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x4 to(const f16x4& v) { uint32_t in0 = v.words[0], in1 = v.words[1]; uint32_t abs0 = in0 & 0x7fff7fffu, abs1 = in1 & 0x7fff7fffu; uint32_t a0, a1; - asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F003F00u)); - asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a1) : "v"(abs1), "v"(0x3F003F00u)); + asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a0) : "v"(abs0), "v"(0x3F803F80u)); + asm volatile("v_pk_min_u16 %0, %1, %2" : "=v"(a1) : "v"(abs1), "v"(0x3F803F80u)); a0 = a0 * 2u + 0x00800080u; a1 = a1 * 2u + 0x00800080u; uint32_t b0 = a0 | (in0 & 0x80008000u); @@ -1274,8 +1252,8 @@ MSCCLPP_DEVICE_INLINE f8_e4m3b15x4 to(const f32x4& v) { return to(h); #elif defined(MSCCLPP_DEVICE_HIP) && defined(__gfx942__) f16x4 h; - h.words[0] = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(v.data[0], v.data[1])); - h.words[1] = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(v.data[2], v.data[3])); + h.words[0] = __builtin_bit_cast(uint32_t, __floats2half2_rn(v.data[0], v.data[1])); + h.words[1] = __builtin_bit_cast(uint32_t, __floats2half2_rn(v.data[2], v.data[3])); return to(h); #else f8_e4m3b15x4 result; diff --git a/include/mscclpp/gpu_utils.hpp b/include/mscclpp/gpu_utils.hpp index b079e0fd..82fa3ec0 100644 --- a/include/mscclpp/gpu_utils.hpp +++ b/include/mscclpp/gpu_utils.hpp @@ -317,6 +317,16 @@ bool isNvlsSupported(); /// @return True if the pointer is allocated by cuMemMap, false otherwise. bool isCuMemMapAllocated(void* ptr); +/// Granularity used to size a `GpuBuffer` allocation so that it is compatible with the multicast (NVLS) API. +enum class GpuBufferGranularity { + /// Minimum multicast granularity. Rounds the allocation up to the minimum granularity required for multicast + /// compatibility, minimizing memory footprint. This is the default. + MultiCastMinimum, + /// Recommended multicast granularity. Rounds the allocation up to the granularity recommended by the driver, + /// which may be larger than the minimum but can yield better performance. + MultiCastRecommended, +}; + /// Allocates a GPU memory space specialized for communication. The memory is zeroed out. Get the device pointer by /// `GpuBuffer::data()`. /// @@ -334,7 +344,11 @@ class GpuBuffer { public: /// Constructs a GpuBuffer with the specified number of elements. /// @param nelems Number of elements to allocate. If it is zero, `data()` will return a null pointer. - GpuBuffer(size_t nelems) : nelems_(nelems) { + /// @param granularity Granularity used to size the allocation for multicast (NVLS) compatibility. Defaults to + /// `GpuBufferGranularity::MultiCastMinimum`, which minimizes memory usage. This is ignored when the buffer is not + /// allocated through the multicast-compatible path. + GpuBuffer(size_t nelems, [[maybe_unused]] GpuBufferGranularity granularity = GpuBufferGranularity::MultiCastMinimum) + : nelems_(nelems) { if (nelems == 0) { bytes_ = 0; return; @@ -342,7 +356,10 @@ class GpuBuffer { MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId_)); #if (CUDA_NVLS_API_AVAILABLE) if (isNvlsSupported()) { - size_t gran = detail::getMulticastGranularity(nelems * sizeof(T), CU_MULTICAST_GRANULARITY_RECOMMENDED); + CUmulticastGranularity_flags granFlag = (granularity == GpuBufferGranularity::MultiCastRecommended) + ? CU_MULTICAST_GRANULARITY_RECOMMENDED + : CU_MULTICAST_GRANULARITY_MINIMUM; + size_t gran = detail::getMulticastGranularity(nelems * sizeof(T), granFlag); bytes_ = (nelems * sizeof(T) + gran - 1) / gran * gran / sizeof(T) * sizeof(T); memory_ = detail::gpuCallocPhysicalShared(nelems, gran); return; diff --git a/include/mscclpp/npkit/npkit_event.hpp b/include/mscclpp/npkit/npkit_event.hpp index 60429193..5084638e 100644 --- a/include/mscclpp/npkit/npkit_event.hpp +++ b/include/mscclpp/npkit/npkit_event.hpp @@ -41,6 +41,6 @@ #define NPKIT_EVENT_KERNEL_ALLREDUCE_EXIT 0x1C #define NPKIT_EVENT_EXECUTOR_OP_BASE_ENTRY 0x1D -#define NPKIT_EVENT_EXECUTOR_OP_BASE_EXIT 0x39 +#define NPKIT_EVENT_EXECUTOR_OP_BASE_EXIT 0x3B #endif diff --git a/include/mscclpp/switch_channel_device.hpp b/include/mscclpp/switch_channel_device.hpp index b52b6572..7b749f7a 100644 --- a/include/mscclpp/switch_channel_device.hpp +++ b/include/mscclpp/switch_channel_device.hpp @@ -37,7 +37,11 @@ struct SwitchChannelDeviceHandle { SwitchChannelDeviceHandle::multimemStore(val, reinterpret_cast(mcPtr) + index); } - template + /// Vectorized multimem load+reduce. The optional `AccumT` template parameter selects the + /// accumulator: when `AccumT == __half` and `VectorType` is an FP8 vector type, the + /// `.acc::f16` variant of the instruction is used (faster but lower precision than the + /// default FP32 accumulator). For all other types `AccumT` is ignored. + template MSCCLPP_DEVICE_INLINE static VectorType multimemLoadReduce(VectorType* ptr) { VectorType val; if constexpr (std::is_same_v) { @@ -81,29 +85,71 @@ struct SwitchChannelDeviceHandle { : "l"(ptr) : "memory"); } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.e4m3x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e4m3x4 %0, [%1];" + : "=r"(val.words[0]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.e4m3x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e4m3x4 {%0,%1}, [%2];" - : "=r"(val.words[0]), "=r"(val.words[1]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v2.e4m3x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e4m3x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e4m3x4 {%0,%1,%2,%3}, [%4];" - : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v4.e4m3x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e4m3x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.e5m2x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e5m2x4 %0, [%1];" + : "=r"(val.words[0]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.e5m2x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e5m2x4 {%0,%1}, [%2];" - : "=r"(val.words[0]), "=r"(val.words[1]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v2.e5m2x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e5m2x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e5m2x4 {%0,%1,%2,%3}, [%4];" - : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v4.e5m2x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e5m2x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } } else { static_assert(dependentFalse, "Not supported type"); } diff --git a/pyproject.toml b/pyproject.toml index 6834ab7f..f4f2e593 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,11 +21,22 @@ dependencies = [ ] [project.optional-dependencies] -cuda11 = ["cupy-cuda11x"] -cuda12 = ["cupy-cuda12x"] -cuda13 = ["cupy-cuda13x"] -rocm6 = ["cupy"] -ep = [] +cuda12 = [ + "cupy-cuda12x", + "cuda-bindings>=12,<13", +] +cuda13 = [ + "cupy-cuda13x", + "cuda-bindings>=13,<14", +] +rocm6 = [ + "cupy", + "hip-python>=6,<7", +] +rocm7 = [ + "cupy", + "hip-python>=7,<8", +] benchmark = [ "mpi4py", "prettytable", diff --git a/python/csrc/core_py.cpp b/python/csrc/core_py.cpp index b6f4c1c9..47eb4c6b 100644 --- a/python/csrc/core_py.cpp +++ b/python/csrc/core_py.cpp @@ -56,6 +56,7 @@ void register_core(nb::module_& m) { .def("get_rank", &Bootstrap::getRank) .def("get_n_ranks", &Bootstrap::getNranks) .def("get_n_ranks_per_node", &Bootstrap::getNranksPerNode) + .def("get_n_ranks_per_ipc_domain", &Bootstrap::getNranksPerIpcDomain) .def( "send", [](Bootstrap* self, uintptr_t ptr, size_t size, int peer, int tag) { diff --git a/python/csrc/gpu_utils_py.cpp b/python/csrc/gpu_utils_py.cpp index 60880456..d6527502 100644 --- a/python/csrc/gpu_utils_py.cpp +++ b/python/csrc/gpu_utils_py.cpp @@ -114,8 +114,13 @@ static nb::capsule toDlpack(GpuBuffer buffer, std::string dataType, std::v void register_gpu_utils(nb::module_& m) { m.def("is_nvls_supported", &isNvlsSupported); + nb::enum_(m, "CppGpuBufferGranularity") + .value("MultiCastMinimum", GpuBufferGranularity::MultiCastMinimum) + .value("MultiCastRecommended", GpuBufferGranularity::MultiCastRecommended); + nb::class_>(m, "CppRawGpuBuffer") - .def(nb::init(), nb::arg("nelems")) + .def(nb::init(), nb::arg("nelems"), + nb::arg("granularity") = GpuBufferGranularity::MultiCastMinimum) .def("nelems", &GpuBuffer::nelems) .def("bytes", &GpuBuffer::bytes) .def("data", [](GpuBuffer& self) { return reinterpret_cast(self.data()); }) diff --git a/python/mscclpp/__init__.py b/python/mscclpp/__init__.py index 5f3a2302..09408171 100644 --- a/python/mscclpp/__init__.py +++ b/python/mscclpp/__init__.py @@ -100,6 +100,7 @@ __all__ = [ "AlgorithmCollection", "CommGroup", "GpuBuffer", + "GpuBufferGranularity", ] diff --git a/python/mscclpp/_core/buffer.py b/python/mscclpp/_core/buffer.py index 0575ca68..e07424f5 100644 --- a/python/mscclpp/_core/buffer.py +++ b/python/mscclpp/_core/buffer.py @@ -6,14 +6,21 @@ from typing import Union, Tuple import cupy as cp import numpy as np -from mscclpp._mscclpp import CppRawGpuBuffer +from mscclpp._mscclpp import CppRawGpuBuffer, CppGpuBufferGranularity -__all__ = ["GpuBuffer"] +__all__ = ["GpuBuffer", "GpuBufferGranularity"] + +GpuBufferGranularity = CppGpuBufferGranularity class GpuBuffer(cp.ndarray): def __new__( - cls, shape: Union[int, Tuple[int]], dtype: cp.dtype = float, strides: Tuple[int] = None, order: str = "C" + cls, + shape: Union[int, Tuple[int]], + dtype: cp.dtype = float, + strides: Tuple[int] = None, + order: str = "C", + granularity: CppGpuBufferGranularity = CppGpuBufferGranularity.MultiCastMinimum, ): # Check if `shape` is valid if isinstance(shape, int): @@ -25,6 +32,6 @@ class GpuBuffer(cp.ndarray): if any(s <= 0 for s in shape): raise ValueError("Shape must be positive.") # Create the buffer - buffer = CppRawGpuBuffer(np.prod(shape) * np.dtype(dtype).itemsize) + buffer = CppRawGpuBuffer(np.prod(shape) * np.dtype(dtype).itemsize, granularity) memptr = cp.cuda.MemoryPointer(cp.cuda.UnownedMemory(buffer.data(), buffer.bytes(), buffer), 0) return cp.ndarray(shape, dtype=dtype, strides=strides, order=order, memptr=memptr) diff --git a/python/mscclpp/_core/comm.py b/python/mscclpp/_core/comm.py index 8e05ca16..aae1a742 100644 --- a/python/mscclpp/_core/comm.py +++ b/python/mscclpp/_core/comm.py @@ -80,6 +80,7 @@ class CommGroup: self.my_rank = self.bootstrap.get_rank() self.nranks = self.bootstrap.get_n_ranks() self.nranks_per_node = self.bootstrap.get_n_ranks_per_node() + self.nranks_per_ipc_domain = self.bootstrap.get_n_ranks_per_ipc_domain() def barrier(self): self.bootstrap.barrier() diff --git a/python/mscclpp/_core/compiler.py b/python/mscclpp/_core/compiler.py index 3b77ce8e..341f3c84 100644 --- a/python/mscclpp/_core/compiler.py +++ b/python/mscclpp/_core/compiler.py @@ -198,7 +198,7 @@ class NativeCodeCompiler: self._is_hip = cp.cuda.runtime.is_hip self._device_arch = get_device_arch() self._compiler = self._get_compiler() - self._default_options = ["-std=c++17", "-O3", "--shared"] + self._default_options = ["-std=c++20", "-O3", "--shared"] python_include = sysconfig.get_path("include") pybind11_include = pybind11.get_include() self._default_options += [f"-I{python_include}", f"-I{pybind11_include}"] diff --git a/python/mscclpp/language/channel.py b/python/mscclpp/language/channel.py index 23d76eda..190e4b25 100644 --- a/python/mscclpp/language/channel.py +++ b/python/mscclpp/language/channel.py @@ -78,6 +78,7 @@ class MemoryChannel: tb_channel_ids = get_program().setup_channel(tb, self) op = SignalOperation(tb_channel_ids, self.channel_type, data_sync, relaxed) get_program().add_operation(self.src_rank, tb, op) + get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type) def wait(self, tb: int, data_sync: SyncType = SyncType.both, relaxed: bool = False): """Wait for a signal through the memory channel. @@ -99,6 +100,7 @@ class MemoryChannel: tb_channel_ids = get_program().setup_channel(tb, self) op = WaitOperation(tb_channel_ids, self.channel_type, data_sync, relaxed) get_program().add_operation(self.src_rank, tb, op) + get_program().register_wait(self.src_rank, self.dst_rank, self.channel_type) def get(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int = None, tb_group: ThreadBlockGroup = None): """Retrieve data from remote memory to local memory. @@ -508,6 +510,7 @@ class PortChannel: tb_channel_ids = get_program().setup_channel(tb, self) op = SignalOperation(tb_channel_ids, self.channel_type, data_sync) get_program().add_operation(self.src_rank, tb, op) + get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type) def wait(self, tb: int, data_sync: SyncType = SyncType.both): """Wait for a signal through the port channel. @@ -527,6 +530,7 @@ class PortChannel: tb_channel_ids = get_program().setup_channel(tb, self) op = WaitOperation(tb_channel_ids, self.channel_type, data_sync) get_program().add_operation(self.src_rank, tb, op) + get_program().register_wait(self.src_rank, self.dst_rank, self.channel_type) def flush(self, tb: int, data_sync: SyncType = SyncType.both): """Flush pending operations through the port channel. @@ -636,6 +640,7 @@ class PortChannel: ) get_program().add_operation(self.src_rank, tb, op) + get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type) def put_with_signal_and_flush(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int): """Send data from local memory to remote memory with signal and flush. @@ -681,6 +686,7 @@ class PortChannel: ) get_program().add_operation(self.src_rank, tb, op) + get_program().register_signal(self.src_rank, self.dst_rank, self.channel_type) def put_packets(self, dst_chunk: Chunk, src_chunk: Chunk, tb: int): """Transfer data from local buffer to remote scratch buffer in packet format. @@ -953,6 +959,54 @@ class SwitchChannel: op = GroupStore(src_chunk, self.buffer_type, buffer_offset, size, tb_channel_ids, self.channel_type) get_program().add_operation(self.src_rank, tb, op) + def broadcast_packets(self, rank, src_chunk: Chunk, buffer_offset, size, tb): + """Broadcast data in packet format from the source chunk to all ranks' scratch buffers in the switch channel. + + Performs a specialized broadcast operation that reads data in packet format + from the source rank's scratch buffer and broadcasts it to each destination rank's + scratch buffer. Both source and destination chunks must be scratch buffers. + + Args: + rank (int): The rank that will execute this broadcast operation. + src_chunk (Chunk): The source scratch chunk containing packet data to broadcast. + buffer_offset (int): The offset in the destination scratch buffer where data will be stored. + size (int): The size of data to broadcast. + tb (int): The thread block ID that will execute this operation. + + Raises: + RuntimeError: If src_chunk rank is not in the rank group, if the source + chunk or destination buffer is not a scratch buffer, if chunk size + doesn't match the required size, or if buffer size is insufficient. + + Example: + >>> channel.broadcast_packets(rank=0, src_chunk=chunk, buffer_offset=0, size=1, tb=0) + """ + self.src_rank = rank + if src_chunk.rank not in self.rank_group.ranks: + raise RuntimeError( + f"Source chunk rank {src_chunk.rank} is not part of the rank group {self.rank_group.ranks}." + ) + if src_chunk.buffer != BufferType.scratch: + raise RuntimeError(f"Source chunk must be of type scratch for the packet broadcast.") + if self.buffer_type != BufferType.scratch: + raise RuntimeError(f"Destination buffer must be of type scratch for the packet broadcast.") + if src_chunk.size != size: + raise RuntimeError(f"Source chunk size {src_chunk.size} does not match the required size {size}.") + + for rank in self.rank_group.ranks: + buffer_size = get_program().gpus[rank].scratch_chunks + + if buffer_size < buffer_offset + size: + raise RuntimeError( + f"Buffer size {buffer_size} is smaller than required size {buffer_offset + size} for rank {rank}." + ) + + tb_channel_ids = get_program().setup_channel(tb, self) + op = GroupStore( + src_chunk, self.buffer_type, buffer_offset, size, tb_channel_ids, self.channel_type, use_packet=True + ) + get_program().add_operation(self.src_rank, tb, op) + class SwitchChannelRankView: """A rank-specific view of a SwitchChannel for performing operations. @@ -1016,3 +1070,23 @@ class SwitchChannel: >>> rank_view.broadcast(src_chunk=chunk, buffer_offset=0, size=1, tb=0) """ return self._channel.broadcast(self._rank, src_chunk, buffer_offset, size, tb) + + def broadcast_packets(self, src_chunk: Chunk, buffer_offset, size, tb): + """Perform a packet broadcast operation from this rank's perspective. + + Convenience method that calls the underlying channel's broadcast_packets + method with this view's rank automatically provided. + + Args: + src_chunk (Chunk): The source chunk containing packet data to broadcast. + buffer_offset (int): The offset in the destination buffer where data will be stored. + size (int): The size of data to broadcast. + tb (int): The thread block ID that will execute this operation. + + Returns: + The result of the underlying channel's broadcast_packets operation. + + Example: + >>> rank_view.broadcast_packets(src_chunk=chunk, buffer_offset=0, size=1, tb=0) + """ + return self._channel.broadcast_packets(self._rank, src_chunk, buffer_offset, size, tb) diff --git a/python/mscclpp/language/collectives.py b/python/mscclpp/language/collectives.py index 55c0e6b6..15d41ad1 100644 --- a/python/mscclpp/language/collectives.py +++ b/python/mscclpp/language/collectives.py @@ -236,3 +236,46 @@ class AllToAll(Collective): } rank_buffers.append(buffers) return rank_buffers + + +class SendRecv(Collective): + """A SendRecv collective communication pattern. + + SendRecv performs a point-to-point send/receive operation. + Each rank sends its input buffer to the next rank and receives data from the + previous rank into its output buffer. + + This operation creates input and output buffers both sized by chunk_factor, + as each rank sends and receives the same amount of data. + """ + + def __init__(self, num_ranks, chunk_factor, inplace): + """Initialize a new SendRecv collective. + + Args: + num_ranks (int): The number of ranks participating in the SendRecv. + chunk_factor (int): The size factor for data chunks. + inplace (bool): Whether the operation should be performed in-place. + + Example: + >>> sendrecv = SendRecv(num_ranks=4, chunk_factor=1, inplace=False) + """ + Collective.__init__(self, num_ranks, chunk_factor, inplace) + self.name = "sendrecv" + + def init_buffers(self): + """Initialize buffers for the SendRecv operation. + + Creates input and output buffers both sized by chunk_factor. + + Returns: + list: A list of buffer dictionaries, one for each rank. + """ + rank_buffers = [] + for rank in range(self.num_ranks): + buffers = { + BufferType.input: BaseBuffer(rank, BufferType.input, 0, self.chunk_factor), + BufferType.output: BaseBuffer(rank, BufferType.output, 0, self.chunk_factor), + } + rank_buffers.append(buffers) + return rank_buffers diff --git a/python/mscclpp/language/internal/operations.py b/python/mscclpp/language/internal/operations.py index 5fb392e3..be6b1f56 100644 --- a/python/mscclpp/language/internal/operations.py +++ b/python/mscclpp/language/internal/operations.py @@ -871,6 +871,7 @@ class GroupLoadReduce(BaseOperation): fused_operation = None if ( isinstance(other, GroupStore) + and other.name == Instruction.group_store and self.buffer_type == other.buffer_type and self.size == other.size and self.dst_chunk == other.src_chunk @@ -911,8 +912,12 @@ class GroupStore(BaseOperation): size: int, channel_ids: List[int], channel_type: ChannelType = ChannelType.switch, + use_packet: bool = False, ): - super().__init__(Instruction.group_store) + if use_packet: + super().__init__(Instruction.group_store_packet) + else: + super().__init__(Instruction.group_store) self.src_chunk = src_chunk self.buffer_type = buffer_type self.buffer_offset = buffer_offset @@ -926,11 +931,10 @@ class GroupStore(BaseOperation): def to_dict(self): result = {"name": self.name.value} - result["src_chunk"] = self.src_chunk.to_dict() - result["buffer_type"] = self.buffer_type.value - result["buffer_offset"] = self.buffer_offset - result["size"] = self.size - result["channel_ids"] = self.channel_ids + result["src_buff"] = [self.src_chunk.to_dict()] + result["dst_buff"] = [ + {"switch_channel_id": self.channel_ids[0], "index": self.buffer_offset, "size": self.size} + ] result["channel_type"] = self.channel_type.value return result diff --git a/python/mscclpp/language/internal/types.py b/python/mscclpp/language/internal/types.py index 9bfe1c76..392282f4 100644 --- a/python/mscclpp/language/internal/types.py +++ b/python/mscclpp/language/internal/types.py @@ -90,6 +90,7 @@ class Instruction(Enum): read_reduce = "rre" read_reduce_send = "rres" group_store = "gstore" + group_store_packet = "gstorepkt" group_load_reduce = "glre" group_load_reduce_store = "glres" pipeline = "pipeline" diff --git a/python/mscclpp/language/program.py b/python/mscclpp/language/program.py index c29e9ab7..825a9d40 100644 --- a/python/mscclpp/language/program.py +++ b/python/mscclpp/language/program.py @@ -10,6 +10,7 @@ from mscclpp.language.rank import Semaphore from mscclpp.language.collectives import * from mscclpp.language.utils import AlgoSpec, ReplicationPolicy from typing import List +from collections import defaultdict import json @@ -112,6 +113,9 @@ class CollectiveProgram: self.loop_context = None + self._signal_counts = defaultdict(int) + self._wait_counts = defaultdict(int) + @classmethod def from_spec(cls, spec: AlgoSpec): """Initialize a new CollectiveProgram from an algorithm specification. @@ -206,7 +210,35 @@ class CollectiveProgram: else: self.gpus[rank].add_operation(tb, operation) + def register_signal(self, src_rank, dst_rank, channel_type): + """Record that `src_rank` issued a signal targeting `dst_rank` over `channel_type`.""" + self._signal_counts[(src_rank, dst_rank, channel_type)] += 1 + + def register_wait(self, src_rank, dst_rank, channel_type): + """Record that `src_rank` performed a wait for `dst_rank` over `channel_type`.""" + self._wait_counts[(src_rank, dst_rank, channel_type)] += 1 + + def validate_signal_wait_pairing(self): + """Validate that every signal issued by a rank is matched by a wait on the peer rank. + + For each (src_rank, dst_rank, channel_type) triple, the number of signals sent by + `src_rank` to `dst_rank` must equal the number of waits performed by `dst_rank` + for `src_rank` on a channel of the same type. Raises RuntimeError on mismatch. + """ + keys = set(self._signal_counts.keys()) | {(dst, src, t) for (src, dst, t) in self._wait_counts.keys()} + for src_rank, dst_rank, channel_type in keys: + signals = self._signal_counts.get((src_rank, dst_rank, channel_type), 0) + waits = self._wait_counts.get((dst_rank, src_rank, channel_type), 0) + if signals != waits: + raise RuntimeError( + f"Signal/Wait mismatch on {channel_type}: rank {src_rank} issues {signals} " + f"signal(s) to rank {dst_rank}, but rank {dst_rank} performs {waits} wait(s) " + f"for rank {src_rank}. Every signal must be matched by a corresponding wait " + f"on the peer rank over a channel of the same type." + ) + def post_process_operations(self): + self.validate_signal_wait_pairing() for gpu in self.gpus: if self.instr_fusion: gpu.optimize_operations() diff --git a/python/mscclpp/language/rank.py b/python/mscclpp/language/rank.py index e5b7aab8..3fd93dc7 100644 --- a/python/mscclpp/language/rank.py +++ b/python/mscclpp/language/rank.py @@ -304,11 +304,24 @@ class BaseBuffer: self.size = offset + size def __getitem__(self, key): - if self.offset + key.stop > self.size: - raise RuntimeError( - f"Index range from {self.offset + key.start} - {self.offset + key.stop} is out of bounds for buffer {self.buffer_type}. Buffer size: {self.size}" + if not isinstance(key, slice): + raise TypeError(f"Buffer indices must be slices, not {type(key).__name__}") + if key.step is not None and key.step != 1: + raise ValueError(f"Buffer slicing does not support step != 1 (got step={key.step})") + buffer_size = self.size - self.offset + start = key.start if key.start is not None else 0 + stop = key.stop if key.stop is not None else buffer_size + if start < 0 or stop < 0: + raise ValueError( + f"Buffer slicing does not support negative indices (got start={key.start}, stop={key.stop})" ) - return Chunk(self.rank, self.buffer_type, self.offset + key.start, key.stop - key.start) + if start > stop: + raise ValueError(f"Buffer slice start ({start}) must be <= stop ({stop})") + if self.offset + stop > self.size: + raise RuntimeError( + f"Index range from {self.offset + start} - {self.offset + stop} is out of bounds for buffer {self.buffer_type}. Buffer size: {self.size}" + ) + return Chunk(self.rank, self.buffer_type, self.offset + start, stop - start) class Buffer(BaseBuffer): diff --git a/python/mscclpp/language/tests/multi_node/send_recv.py b/python/mscclpp/language/tests/multi_node/send_recv.py new file mode 100644 index 00000000..0e898f95 --- /dev/null +++ b/python/mscclpp/language/tests/multi_node/send_recv.py @@ -0,0 +1,95 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import argparse +from mscclpp.language.channel import * +from mscclpp.language.rank import * +from mscclpp.language.general import * +from mscclpp.language.program import * +from mscclpp.language.collectives import * + + +def send_recv(name, nnodes, gpus_per_node, split_mask, instances): + gpu_size = nnodes * gpus_per_node + group_size = split_mask + 1 + if split_mask < 0 or (split_mask & (split_mask + 1)) != 0 or gpu_size % group_size != 0: + raise ValueError( + f"split_mask must be of the form 2^k - 1 and gpu_size ({gpu_size}) must be divisible by " + f"group_size ({group_size}), got split_mask={hex(split_mask)}" + ) + collective = SendRecv(gpu_size, 1, False) + with CollectiveProgram( + name, + collective, + gpu_size, + protocol="Simple", + num_threads_per_block=1024, + use_double_scratch_buffer=False, + min_message_size=0, + max_message_size=2**64 - 1, + instances=instances, + ): + # Creating separate port channels for next and prev directions. + # When prev and next are the same peer (e.g., 2-node ring), both channels go to the same peer + # and get distinct tags. To ensure cross-rank tag matching (rank A's prev_channel signal + # arrives at rank B's next_channel wait), we create channels in opposite order for the + # "higher" rank so that tags cross-match: + # Lower rank: [next(tag0), prev(tag1)] + # Higher rank: [prev(tag0), next(tag1)] + # Then lower.prev(tag1) == higher.next(tag1) and higher.prev(tag0) == lower.next(tag0) + # When prev != next (3+ nodes), each channel targets a different peer so each gets tag 0 + # and this ordering doesn't matter. + group_size = group_size + num_groups = gpu_size // group_size + next_channels = {} # channel for sending to next rank + prev_channels = {} # channel for receiving from prev rank + prev_next_ids = {} + for node in range(nnodes): + for gpu in range(gpus_per_node): + global_rank_id = gpu + gpus_per_node * node + position_in_group = global_rank_id & split_mask + group_id = global_rank_id // group_size + next_group_id = (group_id + 1) % num_groups + next_global_rank_id = next_group_id * group_size + position_in_group + prev_group_id = (group_id - 1 + num_groups) % num_groups + prev_global_rank_id = prev_group_id * group_size + position_in_group + if prev_global_rank_id == next_global_rank_id and global_rank_id > prev_global_rank_id: + # Higher rank: create prev first, then next (swapped order) + prev_channels[global_rank_id] = PortChannel(prev_global_rank_id, global_rank_id) + next_channels[global_rank_id] = PortChannel(next_global_rank_id, global_rank_id) + else: + # Lower rank or different peers: create next first, then prev + next_channels[global_rank_id] = PortChannel(next_global_rank_id, global_rank_id) + prev_channels[global_rank_id] = PortChannel(prev_global_rank_id, global_rank_id) + prev_next_ids[global_rank_id] = (prev_global_rank_id, next_global_rank_id) + + # sync with the next rank and the previous rank in the group + for node in range(nnodes): + for gpu in range(gpus_per_node): + global_rank_id = gpu + gpus_per_node * node + prev_global_rank_id, next_global_rank_id = prev_next_ids[global_rank_id] + prev_channels[global_rank_id].signal(tb=0, data_sync=SyncType.none) + next_channels[global_rank_id].wait(tb=0, data_sync=SyncType.after) + + src_rank = Rank(global_rank_id) + src_buffer = src_rank.get_input_buffer() + dst_rank = Rank(next_global_rank_id) + dst_buffer = dst_rank.get_output_buffer() + + next_channels[global_rank_id].put_with_signal(dst_buffer[:], src_buffer[:], tb=0) + prev_channels[global_rank_id].wait(tb=0, data_sync=SyncType.none) + + print(JSON()) + + +parser = argparse.ArgumentParser() + +parser.add_argument("--name", type=str, help="name of the program") +parser.add_argument("--nnodes", type=int, default=1, help="number of nodes") +parser.add_argument("--gpus_per_node", type=int, help="number of gpus per node") +parser.add_argument("--split_mask", type=lambda x: int(x, 0), default=0x0, help="split mask (e.g. 0x3)") +parser.add_argument("--instances", type=int, default=4, help="number of instances") + +args = parser.parse_args() + +send_recv(args.name, args.nnodes, args.gpus_per_node, args.split_mask, args.instances) diff --git a/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py new file mode 100644 index 00000000..5c138f2a --- /dev/null +++ b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import argparse +from mscclpp.language.channel import * +from mscclpp.language.rank import * +from mscclpp.language.general import * +from mscclpp.language.program import * +from mscclpp.language.collectives import * + + +def allgather_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size, instances): + # Defaults instances=8, num_threads_per_block=256 are tuned for 64-GPU (4x GB200) MNNVL NVLS: + # they give the best busbw across 1MB-1GB (instances saturate at 8; tpb=256 beats 512/1024). + chunksperloop = 1 + collective = AllGather(gpu_size, chunksperloop, True) + with CollectiveProgram( + name, + collective, + gpu_size, + instances=instances, + protocol="Simple", + num_threads_per_block=num_threads_per_block, + use_double_scratch_buffer=False, + min_message_size=min_message_size, + max_message_size=max_message_size, + ): + # NVLS multicast channel over the output buffer. For Allgather each + # rank stores its own chunk to all ranks' output buffers via the switch. + nvls_chan = SwitchChannel(rank_list=[gpu for gpu in range(gpu_size)], buffer_type=BufferType.output) + channels = {} + for gpu in range(gpu_size): + for peer in range(gpu_size): + if peer != gpu: + channels[(peer, gpu)] = MemoryChannel(peer, gpu) + + # Synchronization to ensure all the GPUs are ready + for gpu in range(gpu_size): + src_rank = gpu + for peer in range(gpu_size): + if peer != src_rank: + dst_rank = peer + channels[(dst_rank, src_rank)].signal(tb=0, relaxed=True) + for peer in range(gpu_size): + if peer != src_rank: + dst_rank = peer + channels[(dst_rank, src_rank)].wait(tb=0, relaxed=True, data_sync=SyncType.after) + + # Broadcasting each rank's chunk to every rank via NVLS multimem store. + # Rank `gpu` owns output chunk `gpu` (its input under in-place AllGather) and + # stores it to offset `gpu` across all ranks in the switch group. + for gpu in range(gpu_size): + rank = Rank(gpu) + output_buffer = rank.get_output_buffer() + nvls_chan.at_rank(gpu).broadcast(src_chunk=output_buffer[gpu : gpu + 1], buffer_offset=gpu, size=1, tb=0) + + # Synchronization to ensure the GPUs finished + for gpu in range(gpu_size): + src_rank = gpu + for peer in range(gpu_size): + if peer != src_rank: + dst_rank = peer + channels[(dst_rank, src_rank)].signal(tb=0, relaxed=True, data_sync=SyncType.before) + for peer in range(gpu_size): + if peer != src_rank: + dst_rank = peer + channels[(dst_rank, src_rank)].wait(tb=0, relaxed=True) + + print(JSON()) + + +parser = argparse.ArgumentParser() + +parser.add_argument("--name", type=str, help="name of the program") +parser.add_argument("--num_gpus", type=int, help="number of gpus") +parser.add_argument("--num_threads_per_block", type=int, default=256, help="number of threads per block") +parser.add_argument("--min_message_size", type=int, default=0, help="minimum message size") +parser.add_argument("--max_message_size", type=int, default=2**64 - 1, help="maximum message size") +parser.add_argument("--instances", type=int, default=8, help="number of instances (parallel threadblocks)") + +args = parser.parse_args() + +allgather_example( + args.name, + args.num_gpus, + args.num_threads_per_block, + args.min_message_size, + args.max_message_size, + args.instances, +) diff --git a/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py new file mode 100644 index 00000000..7427e11f --- /dev/null +++ b/python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_pkt.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import argparse +from mscclpp.language.channel import * +from mscclpp.language.rank import * +from mscclpp.language.general import * +from mscclpp.language.program import * +from mscclpp.language.collectives import * + + +def allgather_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size, instances): + # Packet (LL protocol) NVLS AllGather, tuned for small-message latency. + # + # Tuned launch defaults (64-GPU GB200 MNNVL, 1K-1M): instances=1, num_threads_per_block=256 + # + # Unlike allgather_nvls_zero_copy.py (Simple protocol + full-mesh barriers around + # an NVLS multimem store), this variant carries an LL flag inside every packet, so + # the broadcast is self-synchronizing and NO signal/wait barriers are needed. Each + # rank packs its own chunk into scratch, multicasts those packets to every rank's + # scratch via the switch (gstorepkt / MULTI_STORE_PKT), and unpacks locally. + chunksperloop = 1 + collective = AllGather(gpu_size, chunksperloop, True) + with CollectiveProgram( + name, + collective, + gpu_size, + instances=instances, + protocol="LL", + auto_sync=False, + num_threads_per_block=num_threads_per_block, + use_double_scratch_buffer=True, + min_message_size=min_message_size, + max_message_size=max_message_size, + ): + # Scratch holds packet-formatted chunks: gpu_size slots per rank, one per source rank. + scratch_buffer = [] + for gpu in range(gpu_size): + scratch_buffer.append(Buffer(gpu, gpu_size)) + + # NVLS multicast channel bound to the scratch buffer (the packet staging area). + nvls_chan = SwitchChannel(rank_list=[gpu for gpu in range(gpu_size)], buffer_type=BufferType.scratch) + + # Pack each rank's own chunk into its scratch slot `gpu`, then multicast those + # packets to slot `gpu` of every rank's scratch via the switch. + for gpu in range(gpu_size): + rank = Rank(gpu) + output_buffer = rank.get_output_buffer() + rank.copy_packets(scratch_buffer[gpu][gpu : gpu + 1], output_buffer[gpu : gpu + 1], tb=0) + nvls_chan.at_rank(gpu).broadcast_packets( + src_chunk=scratch_buffer[gpu][gpu : gpu + 1], buffer_offset=gpu, size=1, tb=0 + ) + + # Unpack every slot from local scratch into the output buffer. Each unpack waits + # on the packet flag delivered by the owning rank's multicast (no barrier needed). + # Slot j is unpacked on tb=j to parallelize across thread blocks. + for gpu in range(gpu_size): + rank = Rank(gpu) + output_buffer = rank.get_output_buffer() + for j in range(gpu_size): + rank.unpack_packets(output_buffer[j : j + 1], scratch_buffer[gpu][j : j + 1], tb=j) + + print(JSON()) + + +parser = argparse.ArgumentParser() + +parser.add_argument("--name", type=str, help="name of the program") +parser.add_argument("--num_gpus", type=int, help="number of gpus") +parser.add_argument("--num_threads_per_block", type=int, default=256, help="number of threads per block") +parser.add_argument("--min_message_size", type=int, default=1024, help="minimum message size") +parser.add_argument("--max_message_size", type=int, default=1024 * 1024, help="maximum message size") +parser.add_argument("--instances", type=int, default=1, help="number of instances (parallel threadblocks)") + +args = parser.parse_args() + +min_message_size = 1024 * args.instances + +if min_message_size > args.min_message_size: + raise RuntimeError( + f"Minimum message size {args.min_message_size} is too small for the number of instances {args.instances}. The minimum message size must be at least {min_message_size}." + ) + +allgather_example( + args.name, + args.num_gpus, + args.num_threads_per_block, + args.min_message_size, + args.max_message_size, + args.instances, +) diff --git a/python/mscclpp/utils.py b/python/mscclpp/utils.py index 0f0a28d4..99547402 100644 --- a/python/mscclpp/utils.py +++ b/python/mscclpp/utils.py @@ -99,7 +99,7 @@ class KernelBuilder: self._kernel = Kernel(cubin, kernel_name) self.kernel_map[kernel_key] = self._kernel - def _compile_cuda(self, source_file, output_file, std_version="c++17"): + def _compile_cuda(self, source_file, output_file, std_version="c++20"): mscclpp_home = os.environ.get("MSCCLPP_HOME", "/usr/local/mscclpp") include_dir = os.path.join(mscclpp_home, "include") if not cp.cuda.runtime.is_hip: diff --git a/python/mscclpp_benchmark/__init__.py b/python/mscclpp_benchmark/__init__.py index 1ee3f3bf..11e08c9b 100644 --- a/python/mscclpp_benchmark/__init__.py +++ b/python/mscclpp_benchmark/__init__.py @@ -1,4 +1,18 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -from .mscclpp_op import MscclppAllReduce1, MscclppAllReduce2, MscclppAllReduce3, MscclppAllReduce4, MscclppAllReduce5 +__all__ = [ + "MscclppAllReduce1", + "MscclppAllReduce2", + "MscclppAllReduce3", + "MscclppAllReduce4", + "MscclppAllReduce5", +] + + +def __getattr__(name): + if name in __all__: + from . import mscclpp_op + + return getattr(mscclpp_op, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/mscclpp_benchmark/bench_collective.py b/python/mscclpp_benchmark/bench_collective.py new file mode 100644 index 00000000..0cb11739 --- /dev/null +++ b/python/mscclpp_benchmark/bench_collective.py @@ -0,0 +1,648 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from typing import Any + +import cupy as cp +from mpi4py import MPI + +_mscclpp_module = None + +from mscclpp_benchmark.comm import Comm +from mscclpp_benchmark.correctness import ( + CorrectnessStats, + check_correctness as _check_correctness, + fill_case_for_benchmark as _fill_case_for_benchmark, +) +from mscclpp_benchmark.gpu import capture_graph, init_runtime, runtime_name, version +from mscclpp_benchmark.tuner import OfflineTuner +from mscclpp_benchmark.tuning_config import HardwareProfile, TunedConfig, TunedConfigStore, normalize_sku + +_ALLREDUCE = "allreduce" +_ALLGATHER = "allgather" +_DEFAULT_BATCH_SIZES = ( + 1, + 2, + 3, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 96, + 128, + 256, + 512, + 1024, + 1280, + 1536, + 1792, + 2048, + 2560, + 3072, + 3584, + 4096, +) +_DEFAULT_CANDIDATE_NBLOCKS = (1, 4, 8, 16, 24, 32, 48, 56, 64) +_DEFAULT_CANDIDATE_NTHREADS = (256, 512, 768, 1024) + + +def _mscclpp(): + global _mscclpp_module + if _mscclpp_module is None: + import mscclpp + import mscclpp.ext + + _mscclpp_module = mscclpp + return _mscclpp_module + + +@dataclass(frozen=True) +class DTypeSpec: + name: str + cupy_dtype: Any + mscclpp_dtype: Any + accum_dtype: Any | None = None + fp8_format: str | None = None + + +@dataclass(frozen=True) +class CandidateSpec: + algorithm: str + min_message_size: int | None = None + max_message_size: int | None = None + max_nblocks: int | None = None + supported_skus: tuple[str, ...] | None = None + requires_nvls: bool = False + requires_symmetric_memory: bool = False + + +@dataclass +class BenchmarkCase: + collective: str + message_size: int + total_size: int + input: cp.ndarray + output: cp.ndarray + dtype_spec: DTypeSpec + symmetric_memory: bool = False + + +def _device_name() -> str: + props = cp.cuda.runtime.getDeviceProperties(cp.cuda.Device().id) + name = props.get("name", "UNKNOWN") + if isinstance(name, bytes): + return name.decode("utf-8") + return str(name) + + +def _detect_hardware_profile(scale: int) -> HardwareProfile: + return HardwareProfile(sku=normalize_sku(_device_name()), scale=scale) + + +def _parse_dtype(dtype_name: str) -> DTypeSpec: + mscclpp = _mscclpp() + normalized = dtype_name.strip().lower().replace("-", "_") + if normalized in {"float16", "fp16", "half"}: + return DTypeSpec("float16", cp.float16, mscclpp.DataType.float16) + if normalized in {"float32", "fp32", "float"}: + return DTypeSpec("float32", cp.float32, mscclpp.DataType.float32) + if normalized in {"int32", "i32"}: + return DTypeSpec("int32", cp.int32, mscclpp.DataType.int32) + if normalized in {"uint8", "u8"}: + return DTypeSpec("uint8", cp.uint8, mscclpp.DataType.uint8) + if normalized in {"float8_e4m3fn", "fp8_e4m3fn"}: + return DTypeSpec( + "float8_e4m3fn", + cp.uint8, + mscclpp.DataType.float8_e4m3fn, + accum_dtype=mscclpp.DataType.float16, + fp8_format="e4m3fn", + ) + if normalized in {"float8_e4m3fnuz", "fp8_e4m3fnuz"}: + return DTypeSpec( + "float8_e4m3fnuz", + cp.uint8, + mscclpp.DataType.float8_e4m3fnuz, + accum_dtype=mscclpp.DataType.float16, + fp8_format="e4m3fnuz", + ) + if normalized in {"float8_e4m3b15", "fp8_e4m3b15"}: + return DTypeSpec( + "float8_e4m3b15", + cp.uint8, + mscclpp.DataType.float8_e4m3b15, + accum_dtype=mscclpp.DataType.float32, + fp8_format="e4m3b15", + ) + raise ValueError( + f"Unsupported dtype {dtype_name!r}; use float16, float32, int32, uint8, " + "float8_e4m3fn, float8_e4m3fnuz, or float8_e4m3b15" + ) + + +def _with_accum_type(dtype_spec: DTypeSpec, accum_type: str | None) -> DTypeSpec: + if accum_type is None: + return dtype_spec + + mscclpp = _mscclpp() + normalized = accum_type.strip().lower().replace("-", "_") + if normalized in {"native", "same", "auto"}: + accum_dtype = dtype_spec.mscclpp_dtype + elif normalized in {"float16", "fp16", "half"}: + accum_dtype = mscclpp.DataType.float16 + elif normalized in {"float32", "fp32", "float"}: + accum_dtype = mscclpp.DataType.float32 + else: + raise ValueError(f"Unsupported accum type {accum_type!r}; use native, float16, or float32") + + return DTypeSpec( + name=dtype_spec.name, + cupy_dtype=dtype_spec.cupy_dtype, + mscclpp_dtype=dtype_spec.mscclpp_dtype, + accum_dtype=accum_dtype, + fp8_format=dtype_spec.fp8_format, + ) + + +def _human_size(size: int) -> str: + value = float(size) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if value < 1024.0 or unit == "TiB": + return f"{value:.1f} {unit}" + value /= 1024.0 + raise AssertionError("unreachable") + + +def _parse_int_list(raw: str | None, default: tuple[int, ...]) -> tuple[int, ...]: + if raw is None: + return default + values = tuple(sorted({int(item.strip()) for item in raw.split(",") if item.strip()})) + if not values or values[0] <= 0: + raise ValueError(f"Expected a comma-separated list of positive integers, got {raw!r}") + return values + + +def _candidate_specs(collective: str, *, symmetric_memory: bool = False) -> tuple[CandidateSpec, ...]: + if collective == _ALLGATHER: + return (CandidateSpec("default_allgather_fullmesh2", max_nblocks=64, supported_skus=("MI300X",)),) + if collective != _ALLREDUCE: + raise ValueError(f"Unsupported collective: {collective}") + candidates = ( + CandidateSpec( + "default_allreduce_nvls_packet", + max_message_size=512 * 1024, + max_nblocks=16, + supported_skus=("H100", "GB300"), + requires_nvls=True, + ), + CandidateSpec( + "default_allreduce_packet", + max_message_size=4 * 1024 * 1024, + max_nblocks=56, + ), + CandidateSpec( + "default_allreduce_allpair_packet", + max_message_size=4 * 1024 * 1024, + max_nblocks=56, + ), + CandidateSpec( + "default_allreduce_rsag_zero_copy", + min_message_size=512 * 1024 + 1, + ), + CandidateSpec( + "default_allreduce_fullmesh", + min_message_size=512 * 1024 + 1, + max_nblocks=64, + supported_skus=("MI300X",), + ), + ) + if symmetric_memory: + return ( + CandidateSpec( + "default_allreduce_nvls_zero_copy", + max_nblocks=32, + supported_skus=("H100", "GB300"), + requires_nvls=True, + requires_symmetric_memory=True, + ), + *candidates, + ) + return candidates + + +def _candidate_algorithms(comm: Comm, case: BenchmarkCase) -> list[tuple[Any, CandidateSpec]]: + available = comm.algorithms.get(case.collective, {}) + candidates: list[tuple[Any, CandidateSpec]] = [] + seen: set[str] = set() + symmetric_memory = case.symmetric_memory + profile = getattr(comm, "hardware_profile", None) + filtered_out = False + for candidate in _candidate_specs(case.collective, symmetric_memory=symmetric_memory): + if not _candidate_supports_profile(candidate, profile): + filtered_out = True + continue + if not _candidate_supports_message_size(candidate, case.message_size): + filtered_out = True + continue + if candidate.requires_nvls and not _mscclpp().is_nvls_supported(): + filtered_out = True + continue + if candidate.requires_symmetric_memory and not symmetric_memory: + filtered_out = True + continue + algorithm = available.get(candidate.algorithm) + if algorithm is None or algorithm.name in seen: + continue + seen.add(algorithm.name) + candidates.append((algorithm, candidate)) + if candidates: + return candidates + if filtered_out: + return [] + return [(algorithm, CandidateSpec(algorithm.name)) for algorithm in available.values()] + + +def _candidate_supports_profile(candidate: CandidateSpec, profile: HardwareProfile | None) -> bool: + if candidate.supported_skus is None: + return True + sku = None if profile is None else profile.sku + if not sku or sku == "UNKNOWN": + return True + return sku in candidate.supported_skus + + +def _candidate_supports_message_size(candidate: CandidateSpec, message_size: int) -> bool: + if candidate.min_message_size is not None and message_size < candidate.min_message_size: + return False + if candidate.max_message_size is not None and message_size > candidate.max_message_size: + return False + return True + + +def _make_case( + *, + collective: str, + nelems: int, + dtype_spec: DTypeSpec, + comm_group: Any, + buffer_mode: str, + symmetric_memory: bool = False, +) -> BenchmarkCase: + if buffer_mode not in ("in-place", "out-of-place"): + raise ValueError(f"Unsupported buffer mode: {buffer_mode}") + + if collective == _ALLREDUCE: + if buffer_mode == "in-place": + memory = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype) + input_buffer = memory + output = memory + else: + input_buffer = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype) + output = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype) + return BenchmarkCase( + collective=collective, + message_size=input_buffer.nbytes, + total_size=output.nbytes, + input=input_buffer, + output=output, + dtype_spec=dtype_spec, + symmetric_memory=symmetric_memory, + ) + + if collective != _ALLGATHER: + raise ValueError(f"Unsupported collective: {collective}") + + if buffer_mode == "in-place": + output = _mscclpp().GpuBuffer(nelems * comm_group.nranks, dtype=dtype_spec.cupy_dtype) + start = comm_group.my_rank * nelems + input_buffer = output[start : start + nelems] + else: + input_buffer = _mscclpp().GpuBuffer(nelems, dtype=dtype_spec.cupy_dtype) + output = _mscclpp().GpuBuffer(nelems * comm_group.nranks, dtype=dtype_spec.cupy_dtype) + + return BenchmarkCase( + collective=collective, + message_size=input_buffer.nbytes, + total_size=output.nbytes, + input=input_buffer, + output=output, + dtype_spec=dtype_spec, + symmetric_memory=symmetric_memory, + ) + + +def _try_measure_case( + comm: Comm, + case: BenchmarkCase, + config: TunedConfig, + *, + n_warmup: int, + n_graph_launches: int, + n_ops_per_graph: int, +) -> float | None: + try: + return _measure_case( + comm, + case, + config, + n_warmup=n_warmup, + n_graph_launches=n_graph_launches, + n_ops_per_graph=n_ops_per_graph, + ) + except Exception as exc: + if comm.rank == 0: + print( + f"[skip] {config.algorithm} nb={config.nblocks} nt={config.nthreads} " + f"size={case.message_size}: {type(exc).__name__}: {exc}", + flush=True, + ) + return None + + +def _measure_case( + comm: Comm, + case: BenchmarkCase, + config: TunedConfig, + *, + n_warmup: int, + n_graph_launches: int, + n_ops_per_graph: int, +) -> float: + _fill_case_for_benchmark(case, comm.rank) + if comm.run(case, config) != 0: + raise RuntimeError("algorithm returned non-zero status") + cp.cuda.runtime.deviceSynchronize() + comm.comm_group.barrier() + + stream = cp.cuda.Stream(non_blocking=True) + graph = None + + def capture_ops() -> None: + for _ in range(n_ops_per_graph): + ret = comm.run(case, config, stream) + if ret != 0: + raise RuntimeError("algorithm returned non-zero status during graph capture") + + try: + with stream: + graph = capture_graph(stream, capture_ops) + + for _ in range(n_warmup): + graph.launch(stream) + stream.synchronize() + comm.comm_group.barrier() + + start = cp.cuda.Event() + end = cp.cuda.Event() + start.record(stream) + for _ in range(n_graph_launches): + graph.launch(stream) + end.record(stream) + end.synchronize() + + elapsed_us = cp.cuda.get_elapsed_time(start, end) * 1000.0 / (n_graph_launches * n_ops_per_graph) + return float(MPI.COMM_WORLD.allreduce(elapsed_us, op=MPI.MAX)) + finally: + if graph is not None: + graph.close() + + +def _bandwidth_gbps(num_bytes: int, time_us: float) -> float: + return num_bytes / time_us / 1e3 + + +def _busbw_factor(collective: str, nranks: int) -> float: + if nranks <= 1: + return 1.0 + if collective == _ALLREDUCE: + return 2 * (nranks - 1) / nranks + if collective == _ALLGATHER: + return (nranks - 1) / nranks + raise ValueError(f"Unsupported collective: {collective}") + + +def _format_table(headers: list[str], rows: list[list[str]]) -> str: + widths = [len(header) for header in headers] + for row in rows: + widths = [max(width, len(cell)) for width, cell in zip(widths, row)] + header_line = " | ".join(header.ljust(width) for header, width in zip(headers, widths)) + sep_line = "-+-".join("-" * width for width in widths) + row_lines = [" | ".join(cell.ljust(width) for cell, width in zip(row, widths)) for row in rows] + return "\n".join([header_line, sep_line, *row_lines]) + + +def _format_stat(value: float | None) -> str: + if value is None: + return "-" + return f"{value:.6g}" + + +def _format_mismatches(stats: CorrectnessStats | None) -> str: + if stats is None or stats.total == 0: + return "-" + return f"{stats.mismatches}/{stats.total}" + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Benchmark MSCCL++ collectives without PyTorch dependencies") + parser.add_argument("--collective", choices=(_ALLREDUCE, _ALLGATHER), default=_ALLREDUCE) + parser.add_argument("--d-model", type=int, default=5120) + parser.add_argument("--dtype", default="float16") + parser.add_argument("--accum-type", help="Accumulation type for reductions: native, float16, or float32") + parser.add_argument("--batch-sizes", help="Comma-separated batch sizes; default uses the benchmark sweep") + parser.add_argument( + "--buffer-mode", + choices=("in-place", "out-of-place"), + default="in-place", + help="Buffer layout for the collective: in-place (input aliases output) or out-of-place (separate buffers)", + ) + parser.add_argument("--config-path", help="Optional MSCCL++ tuned config JSON") + parser.add_argument("--write-config", help="Write autotuned configs to this JSON path") + parser.add_argument("--autotune", action="store_true", help="Tune each benchmark size before timing it") + parser.add_argument("--skip-correctness", action="store_true") + parser.add_argument("--correctness-iters", type=int, default=1) + parser.add_argument("--scratch-buffer-size", type=int, default=1 << 27) + parser.add_argument("--warmup", type=int, default=5, help="Warmup graph replays before benchmark timing") + parser.add_argument("--graph-launches", type=int, default=10, help="Timed graph replays") + parser.add_argument("--iterations", type=int, default=100, help="Collective operations captured per CUDA graph") + parser.add_argument("--tune-warmup", type=int, default=2) + parser.add_argument("--tune-graph-launches", type=int, default=3) + parser.add_argument("--tune-iterations", type=int, default=20) + parser.add_argument("--candidate-nblocks", help="Comma-separated nblocks tuning candidates") + parser.add_argument("--candidate-nthreads", help="Comma-separated nthreads tuning candidates") + parser.add_argument("--symmetric-memory", action="store_true") + return parser + + +def _validate_args(args: argparse.Namespace) -> None: + for name in ( + "d_model", + "scratch_buffer_size", + "graph_launches", + "iterations", + "tune_graph_launches", + "tune_iterations", + "correctness_iters", + ): + if getattr(args, name) <= 0: + raise ValueError(f"--{name.replace('_', '-')} must be positive") + if args.warmup < 0 or args.tune_warmup < 0: + raise ValueError("warmup counts must be non-negative") + + +def main(argv: list[str] | None = None) -> None: + args = _build_parser().parse_args(argv) + _validate_args(args) + init_runtime() + + local_comm = MPI.COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED, 0, MPI.INFO_NULL) + try: + visible_devices = cp.cuda.runtime.getDeviceCount() + if visible_devices <= 0: + raise RuntimeError("MSCCL++ benchmark requires at least one visible GPU") + cp.cuda.Device(local_comm.Get_rank() % visible_devices).use() + finally: + local_comm.Free() + + dtype_spec = _with_accum_type(_parse_dtype(args.dtype), args.accum_type) + batch_sizes = _parse_int_list(args.batch_sizes, _DEFAULT_BATCH_SIZES) + candidate_nblocks = _parse_int_list(args.candidate_nblocks, _DEFAULT_CANDIDATE_NBLOCKS) + candidate_nthreads = _parse_int_list(args.candidate_nthreads, _DEFAULT_CANDIDATE_NTHREADS) + + comm_group = _mscclpp().CommGroup(MPI.COMM_WORLD) + setattr(comm_group, "_mpi_comm", MPI.COMM_WORLD) + hardware_profile = _detect_hardware_profile(comm_group.nranks) + config_store = TunedConfigStore.load_path(args.config_path) if args.config_path else TunedConfigStore.empty() + comm = Comm( + comm_group, + config_store=config_store, + hardware_profile=hardware_profile, + scratch_buffer_size=args.scratch_buffer_size, + ) + tuner = OfflineTuner( + comm, + candidate_nblocks=candidate_nblocks, + candidate_nthreads=candidate_nthreads, + n_warmup=args.tune_warmup, + n_graph_launches=args.tune_graph_launches, + n_ops_per_graph=args.tune_iterations, + candidate_algorithms=_candidate_algorithms, + check_correctness=_check_correctness, + measure=_try_measure_case, + ) + + rows: list[list[str]] = [] + try: + if comm.rank == 0: + print( + f"MSCCL++ {args.collective} benchmark: profile={hardware_profile} dtype={dtype_spec.name} " + f"graph_launches={args.graph_launches} iterations={args.iterations}", + flush=True, + ) + + for batch_size in batch_sizes: + nelems = batch_size * args.d_model + case = _make_case( + collective=args.collective, + nelems=nelems, + dtype_spec=dtype_spec, + comm_group=comm_group, + buffer_mode=args.buffer_mode, + symmetric_memory=args.symmetric_memory, + ) + config = tuner.tune(case) if args.autotune else comm.resolve_config(case) + if config is None: + continue + if args.autotune: + config_store.upsert(hardware_profile, args.collective, case.message_size, config) + + correctness = "SKIP" + correctness_stats: CorrectnessStats | None = None + if not args.skip_correctness: + correctness_stats = _check_correctness(comm, case, config, niter=args.correctness_iters) + correctness = "PASS" if correctness_stats else "FAIL" + comm.reset(config) + if correctness != "PASS": + raise RuntimeError( + f"Correctness failed for batch_size={batch_size}, message_size={case.message_size}, " + f"config={config}" + ) + + time_us = _measure_case( + comm, + case, + config, + n_warmup=args.warmup, + n_graph_launches=args.graph_launches, + n_ops_per_graph=args.iterations, + ) + comm.reset(config) + + algbw = _bandwidth_gbps(case.total_size, time_us) + busbw = algbw * _busbw_factor(args.collective, comm_group.nranks) + rows.append( + [ + str(batch_size), + _human_size(case.message_size), + _human_size(case.total_size), + config.algorithm, + str(config.nblocks or "auto"), + str(config.nthreads or "auto"), + f"{time_us:.2f}", + f"{algbw:.2f}", + f"{busbw:.2f}", + correctness, + _format_stat(None if correctness_stats is None else correctness_stats.max_abs_diff), + _format_stat(None if correctness_stats is None else correctness_stats.mean_abs_diff), + _format_mismatches(correctness_stats), + ] + ) + if comm.rank == 0: + print(".", end="", flush=True) + if runtime_name() == "hip" and version()[:2] == (7, 2): + # TODO: remove this after ROCm 7.2 HIP IPC export issue is fixed. + del case + comm.comm_group.barrier() + + if args.write_config and comm.rank == 0: + config_store.write_path(args.write_config) + print(f"\nWrote tuned config to {args.write_config}", flush=True) + + if comm.rank == 0: + print( + "\n" + + _format_table( + [ + "batch", + "msg", + "total", + "algorithm", + "nblocks", + "nthreads", + "time_us", + "algBW_GB/s", + "busBW_GB/s", + "check", + "max_diff", + "mean_diff", + "mismatch", + ], + rows, + ), + flush=True, + ) + finally: + comm_group.barrier() + cp.cuda.runtime.deviceSynchronize() + comm.close() + + +if __name__ == "__main__": + main() diff --git a/python/mscclpp_benchmark/comm.py b/python/mscclpp_benchmark/comm.py new file mode 100644 index 00000000..23770ac2 --- /dev/null +++ b/python/mscclpp_benchmark/comm.py @@ -0,0 +1,409 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) +_ALLREDUCE_COLLECTIVE = "allreduce" +_ALLGATHER_COLLECTIVE = "allgather" +_mscclpp_module = None + +from mscclpp_benchmark.gpu import current_device, device_name, set_device +from mscclpp_benchmark.tuning_config import HardwareProfile, TunedConfig, TunedConfigStore, normalize_sku + + +def _mscclpp(): + global _mscclpp_module + if _mscclpp_module is None: + import mscclpp + import mscclpp.ext + + _mscclpp_module = mscclpp + return _mscclpp_module + + +class Buffer: + def __init__( + self, + nbytes: int | None = None, + *, + dtype: str | Any = "float16", + shape: tuple[int, ...] | None = None, + buffer: Any | None = None, + ) -> None: + self.dtype = dtype + self.element_size = _dtype_size(dtype) + if buffer is None: + if nbytes is None: + if shape is None: + raise ValueError("Either nbytes or shape is required") + nbytes = _numel(shape) * self.element_size + _ensure_device() + buffer = _mscclpp().RawGpuBuffer(int(nbytes)) + self.buffer = buffer + self.nbytes = int(buffer.bytes()) + self.shape = shape if shape is not None else (self.nbytes // self.element_size,) + + @property + def ndim(self) -> int: + return len(self.shape) + + @property + def size(self) -> int: + return _numel(self.shape) + + def data_ptr(self) -> int: + return int(self.buffer.data()) + + +class _AllReduceOp: + def __init__(self, comm: "Comm", x: Any, *, symmetric_memory: bool = False) -> None: + self._comm = comm + self._x = x + self._symmetric_memory = symmetric_memory + + def __call__(self, **_: Any) -> Any: + self._comm.run(self._x, symmetric_memory=self._symmetric_memory) + return self._x + + +class _AllGatherOp: + def __init__(self, comm: "Comm", x: Any, *, dim: int, y: Any | None = None, symmetric_memory: bool = False) -> None: + shape = _shape(x) + if len(shape) == 0: + raise ValueError("MSCCL++ allgather requires a non-scalar buffer") + if dim % len(shape) != 0: + raise NotImplementedError("Raw-buffer allgather currently supports only dim=0") + if y is None: + y_shape = (comm._scale() * shape[0], *shape[1:]) + y = Buffer(dtype=_dtype(x), shape=y_shape) + self._comm = comm + self._x = x + self.y = y + self._symmetric_memory = symmetric_memory + + def __call__(self, **_: Any) -> Any: + self._comm.run( + self._x, + collective=_ALLGATHER_COLLECTIVE, + output_tensor=self.y, + symmetric_memory=self._symmetric_memory, + ) + return self.y + + +class Comm: + """Runtime MSCCL++ wrapper that owns algorithm handles and execution without Torch/CuPy tensors.""" + + def __init__( + self, + comm_group: Any, + scratch_buffer_size: int = 1 << 27, + *, + config_store: "TunedConfigStore | None" = None, + hardware_profile: HardwareProfile | None = None, + ) -> None: + self._comm_group = comm_group + self._mpi_comm = getattr(comm_group, "_mpi_comm", None) + self._rank = comm_group.my_rank + self._closed = False + _ensure_device() + self._mscclpp = _mscclpp() + self._scratch_buffer = self._mscclpp.RawGpuBuffer(scratch_buffer_size) + self._config_store = TunedConfigStore.empty() if config_store is None else config_store + self._hardware_profile = ( + _detect_hardware_profile(scale=self._scale()) if hardware_profile is None else hardware_profile + ) + self._default_config_warning_keys: set[tuple[str, str, str, int]] = set() + + algorithms = self._mscclpp.ext.AlgorithmCollectionBuilder().build_default_algorithms( + scratch_buffer=self._scratch_buffer.data(), + scratch_buffer_size=self._scratch_buffer.bytes(), + rank=self._rank, + ) + self._algorithms_by_collective: dict[str, dict[str, Any]] = {} + for algorithm in algorithms: + self._algorithms_by_collective.setdefault(algorithm.collective, {})[algorithm.name] = algorithm + + @property + def comm_group(self) -> Any: + return self._comm_group + + @property + def rank(self) -> int: + return self._rank + + @property + def nranks(self) -> int: + return self._comm_group.nranks + + @property + def algorithms(self) -> dict[str, dict[str, Any]]: + return self._algorithms_by_collective + + @property + def hardware_profile(self) -> HardwareProfile: + return self._hardware_profile + + def make_allreduce(self, x: Any, *, symmetric_memory: bool = False) -> _AllReduceOp: + return _AllReduceOp(self, x, symmetric_memory=symmetric_memory) + + def make_allgather(self, x: Any, dim: int, y: Any | None = None, *, symmetric_memory: bool = False) -> _AllGatherOp: + return _AllGatherOp(self, x, dim=dim, y=y, symmetric_memory=symmetric_memory) + + def _scale(self) -> int: + if self._mpi_comm is not None: + return int(self._mpi_comm.Get_size()) + return 1 + + def resolve_config(self, case: Any, *, symmetric_memory: bool = False) -> TunedConfig: + dtype_override = getattr(getattr(case, "dtype_spec", None), "mscclpp_dtype", None) + accum_dtype = getattr(getattr(case, "dtype_spec", None), "accum_dtype", None) or dtype_override + symmetric_memory = symmetric_memory or bool(getattr(case, "symmetric_memory", False)) + return self._resolve_config( + case.collective, + case.input, + dtype_override=dtype_override, + accum_dtype=accum_dtype, + symmetric_memory=symmetric_memory, + ) + + def _resolve_config( + self, + collective: str, + buffer: Any, + *, + dtype_override: Any | None = None, + accum_dtype: Any | None = None, + symmetric_memory: bool = False, + ) -> TunedConfig: + tuned_config = self._config_store.select(self._hardware_profile, collective, _nbytes(buffer)) + if tuned_config is not None and tuned_config.algorithm in self._algorithms_by_collective.get(collective, {}): + return tuned_config + + if self._rank == 0: + dim = int(_shape(buffer)[1]) if len(_shape(buffer)) > 1 else 1 + warning_key = ( + collective, + str(dtype_override if dtype_override is not None else _dtype(buffer)), + str( + accum_dtype + if accum_dtype is not None + else dtype_override if dtype_override is not None else _dtype(buffer) + ), + dim, + ) + if warning_key not in self._default_config_warning_keys: + self._default_config_warning_keys.add(warning_key) + logger.warning( + "MSCCL++ default config: no tuning for collective=%s profile=%s dtype=%s accum=%s dim=%s; perf may be poor", + collective, + self._hardware_profile, + warning_key[1], + warning_key[2], + dim, + ) + return _default_tuned_config( + collective, + _nbytes(buffer), + self._algorithms_by_collective, + symmetric_memory=symmetric_memory, + ) + + def run( + self, + buffer: Any, + config: TunedConfig | None = None, + stream: Any | None = None, + *, + collective: str = _ALLREDUCE_COLLECTIVE, + output_tensor: Any | None = None, + dtype_override: Any | None = None, + accum_dtype: Any | None = None, + symmetric_memory: bool = False, + ) -> int: + if self._closed: + raise RuntimeError("Cannot use a closed MSCCL++ comm") + + raise_on_error = True + if hasattr(buffer, "input") and hasattr(buffer, "output") and hasattr(buffer, "dtype_spec"): + case = buffer + buffer = case.input + output_tensor = case.output + collective = case.collective + dtype_override = case.dtype_spec.mscclpp_dtype + accum_dtype = case.dtype_spec.accum_dtype or dtype_override + symmetric_memory = symmetric_memory or bool(getattr(case, "symmetric_memory", False)) + raise_on_error = False + + if collective not in self._algorithms_by_collective: + raise RuntimeError(f"No supported MSCCL++ {collective} algorithm is available") + + if config is None: + config = self._resolve_config( + collective, + buffer, + dtype_override=dtype_override, + accum_dtype=accum_dtype, + symmetric_memory=symmetric_memory, + ) + symmetric_memory = symmetric_memory or config.symmetric_memory + algorithm = self._algorithms_by_collective[collective][config.algorithm] + output = buffer if output_tensor is None else output_tensor + dtype = dtype_override if dtype_override is not None else _dtype_to_mscclpp(_dtype(buffer)) + accum = accum_dtype if accum_dtype is not None else dtype + ret = algorithm.execute( + comm=self._comm_group.communicator, + input_buffer=_data_ptr(buffer), + output_buffer=_data_ptr(output), + input_size=_nbytes(buffer), + output_size=_nbytes(output), + dtype=dtype, + op=self._mscclpp.ReduceOp.SUM if collective == _ALLREDUCE_COLLECTIVE else self._mscclpp.ReduceOp.NOP, + stream=_stream_ptr(stream), + nblocks=config.nblocks or 0, + nthreads_per_block=config.nthreads or 0, + symmetric_memory=symmetric_memory, + accum_dtype=accum, + ) + if ret != 0 and raise_on_error: + raise RuntimeError(f"MSCCL++ {collective} failed on rank {self._rank} with error code {ret}") + return ret + + def reset(self, config: TunedConfig | None = None) -> None: + if config is not None: + for algorithms_by_name in self._algorithms_by_collective.values(): + algorithm = algorithms_by_name.get(config.algorithm) + if algorithm is not None: + algorithm.reset() + return + for algorithms_by_name in self._algorithms_by_collective.values(): + for algorithm in algorithms_by_name.values(): + algorithm.reset() + + def close(self) -> None: + self.reset() + self._algorithms_by_collective = {} + self._scratch_buffer = None + self._closed = True + self._mscclpp.ext.AlgorithmCollectionBuilder.reset() + + +def _numel(shape: tuple[int, ...]) -> int: + out = 1 + for dim in shape: + out *= int(dim) + return out + + +def _dtype_size(dtype: Any) -> int: + dtype_name = _dtype_name(dtype) + if dtype_name in {"float16", "bfloat16"}: + return 2 + if dtype_name in {"float32", "int32", "uint32"}: + return 4 + if dtype_name in {"uint8", "float8_e4m3b15", "float8_e4m3fn", "float8_e4m3fnuz"}: + return 1 + raise ValueError(f"Unknown data type size for {dtype}") + + +def _dtype_name(dtype: Any) -> str: + if isinstance(dtype, str): + return dtype.strip().lower().replace("-", "_") + name = str(dtype).rsplit(".", 1)[-1] + return name.strip().lower().replace("-", "_") + + +def _dtype_to_mscclpp(dtype: Any) -> Any: + dtype_name = _dtype_name(dtype) + mapping = { + "float16": _mscclpp().DataType.float16, + "float32": _mscclpp().DataType.float32, + "int32": _mscclpp().DataType.int32, + "uint8": _mscclpp().DataType.uint8, + "float8_e4m3b15": _mscclpp().DataType.float8_e4m3b15, + "float8_e4m3fn": _mscclpp().DataType.float8_e4m3fn, + "float8_e4m3fnuz": _mscclpp().DataType.float8_e4m3fnuz, + } + try: + return mapping[dtype_name] + except KeyError as exc: + raise ValueError(f"Unknown data type: {dtype}") from exc + + +def _data_ptr(buffer: Any) -> int: + if hasattr(buffer, "data_ptr"): + data_ptr = buffer.data_ptr + return int(data_ptr() if callable(data_ptr) else data_ptr) + if hasattr(buffer, "data"): + data = buffer.data + if callable(data): + return int(data()) + if hasattr(data, "ptr"): + return int(data.ptr) + raise TypeError(f"Cannot get device pointer from {type(buffer)!r}") + + +def _stream_ptr(stream: Any | None) -> int: + if stream is None: + return 0 + return int(getattr(stream, "ptr", stream)) + + +def _nbytes(buffer: Any) -> int: + if hasattr(buffer, "nbytes"): + return int(buffer.nbytes) + if hasattr(buffer, "bytes"): + value = buffer.bytes + return int(value() if callable(value) else value) + raise TypeError(f"Cannot get byte size from {type(buffer)!r}") + + +def _shape(buffer: Any) -> tuple[int, ...]: + shape = getattr(buffer, "shape", None) + if shape is None: + return (_nbytes(buffer) // _dtype_size(_dtype(buffer)),) + return tuple(int(dim) for dim in shape) + + +def _dtype(buffer: Any) -> Any: + dtype = getattr(buffer, "dtype", None) + if dtype is None: + return "uint8" + return dtype + + +def _detect_hardware_profile(*, scale: int) -> HardwareProfile: + try: + sku = device_name() + except Exception: + sku = "UNKNOWN" + return HardwareProfile(sku=normalize_sku(sku), scale=scale) + + +def _ensure_device() -> None: + set_device(current_device()) + + +def _default_tuned_config( + collective: str, + message_size: int, + algorithms_by_collective: dict[str, dict[str, Any]], + *, + symmetric_memory: bool = False, +) -> TunedConfig: + if collective == _ALLGATHER_COLLECTIVE: + return TunedConfig("default_allgather_fullmesh2", symmetric_memory=symmetric_memory) + available = algorithms_by_collective.get(collective, {}) + if symmetric_memory and _mscclpp().is_nvls_supported() and "default_allreduce_nvls_zero_copy" in available: + return TunedConfig("default_allreduce_nvls_zero_copy", symmetric_memory=True) + if message_size <= 512 * 1024 and "default_allreduce_packet" in available: + return TunedConfig("default_allreduce_packet", symmetric_memory=symmetric_memory) + if "default_allreduce_rsag_zero_copy" in available: + return TunedConfig("default_allreduce_rsag_zero_copy", symmetric_memory=symmetric_memory) + if available: + return TunedConfig(next(iter(available)), symmetric_memory=symmetric_memory) + raise RuntimeError(f"No MSCCL++ algorithm is available for {collective}") diff --git a/python/mscclpp_benchmark/correctness.py b/python/mscclpp_benchmark/correctness.py new file mode 100644 index 00000000..15de2f5a --- /dev/null +++ b/python/mscclpp_benchmark/correctness.py @@ -0,0 +1,401 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import cupy as cp +from mpi4py import MPI + +_mscclpp_module = None + + +def _mscclpp(): + global _mscclpp_module + if _mscclpp_module is None: + import mscclpp + + _mscclpp_module = mscclpp + return _mscclpp_module + + +@dataclass(frozen=True) +class CorrectnessStats: + ok: bool + max_abs_diff: float = 0.0 + mean_abs_diff: float = 0.0 + mismatches: int = 0 + total: int = 0 + + def __bool__(self) -> bool: + return self.ok + + +def config_accum_dtype(case: Any) -> Any: + return case.dtype_spec.accum_dtype or case.dtype_spec.mscclpp_dtype + + +def fill_case_for_benchmark(case: Any, rank: int) -> None: + values = _benchmark_input_values(case, rank) + encoded = _encode_correctness_input(case, values) + if case.collective == "allreduce": + case.input[...] = encoded + return + case.output.fill(0) + case.input[...] = encoded + + +def check_correctness( + comm: Any, + case: Any, + config: Any, + *, + niter: int = 1, +) -> CorrectnessStats: + all_ok = True + local_max_abs_diff = 0.0 + local_sum_abs_diff = 0.0 + local_mismatches = 0 + local_total = 0 + for iteration in range(niter): + _fill_case_for_correctness(case, comm.rank, iteration) + ret = comm.run(case, config) + cp.cuda.runtime.deviceSynchronize() + comm.comm_group.barrier() + if ret != 0: + all_ok = False + continue + + expected, stats_expected = _expected_outputs(case, comm.nranks, iteration) + iter_stats = _local_diff_stats(case, case.output, expected, comm.nranks, stats_expected=stats_expected) + local_ok = _compare_output(case, case.output, expected, comm.nranks) + all_ok = all_ok and local_ok + local_max_abs_diff = max(local_max_abs_diff, iter_stats.max_abs_diff) + local_sum_abs_diff += iter_stats.mean_abs_diff * iter_stats.total + local_mismatches += iter_stats.mismatches + local_total += iter_stats.total + + if not local_ok: + mismatch = _mismatch_mask(case, case.output, expected, comm.nranks) + print( + "not close: " + f"iter={iteration}, rank={comm.rank}, output={case.output[mismatch][0]}, " + f"expected={expected[mismatch][0]}, max_abs_diff={iter_stats.max_abs_diff:.6g}, " + f"mean_abs_diff={iter_stats.mean_abs_diff:.6g}, mismatches={iter_stats.mismatches}/{iter_stats.total}", + flush=True, + ) + + global_ok = bool(MPI.COMM_WORLD.allreduce(all_ok, op=MPI.LAND)) + global_max_abs_diff = float(MPI.COMM_WORLD.allreduce(local_max_abs_diff, op=MPI.MAX)) + global_sum_abs_diff = float(MPI.COMM_WORLD.allreduce(local_sum_abs_diff, op=MPI.SUM)) + global_mismatches = int(MPI.COMM_WORLD.allreduce(local_mismatches, op=MPI.SUM)) + global_total = int(MPI.COMM_WORLD.allreduce(local_total, op=MPI.SUM)) + global_mean_abs_diff = global_sum_abs_diff / global_total if global_total else 0.0 + return CorrectnessStats( + ok=global_ok, + max_abs_diff=global_max_abs_diff, + mean_abs_diff=global_mean_abs_diff, + mismatches=global_mismatches, + total=global_total, + ) + + +def _fill_case_for_correctness(case: Any, rank: int, iteration: int) -> None: + values = _correctness_input_values(case, rank, iteration) + encoded = _encode_correctness_input(case, values) + if case.collective == "allreduce": + case.input[...] = encoded + return + case.output.fill(0) + case.input[...] = encoded + + +def _correctness_input_values(case: Any, rank: int, iteration: int): + shape = case.input.shape + rng = cp.random.RandomState(_correctness_seed(rank, iteration)) + return _random_input_values(case, rng, shape) + + +def _benchmark_input_values(case: Any, rank: int): + rng = cp.random.RandomState(17_000_003 + rank) + return _random_input_values(case, rng, case.input.shape) + + +def _random_input_values(case: Any, rng, shape): + if case.dtype_spec.fp8_format is not None: + value_range = _fp8_correctness_input_range(case) + return rng.uniform(-value_range, value_range, size=shape).astype(cp.float32) + if case.dtype_spec.cupy_dtype == cp.int32: + return rng.randint(-1, 2, size=shape).astype(cp.int32) + if case.dtype_spec.cupy_dtype == cp.uint8: + return rng.randint(0, 2, size=shape).astype(cp.uint8) + return rng.uniform(-1.0, 1.0, size=shape).astype(cp.float32) + + +def _correctness_seed(rank: int, iteration: int) -> int: + return (iteration + 1) * 1_000_003 + rank + + +def _fp8_correctness_input_range(case: Any) -> float: + if case.collective != "allreduce": + return 1.0 + fp8_format = case.dtype_spec.fp8_format + if fp8_format is None: + return 1.0 + return min(1.0, _fp8_max_abs_value(fp8_format) / max(1, MPI.COMM_WORLD.size)) + + +def _encode_correctness_input(case: Any, values): + if case.dtype_spec.fp8_format is not None: + # FP8 buffers are stored as uint8 raw bytes, so a normal astype(uint8) cast would not produce FP8 bits. + return _encode_fp8_values(case.dtype_spec.fp8_format, values) + return values.astype(case.dtype_spec.cupy_dtype) + + +def _local_diff_stats(case: Any, output, expected, nranks: int, *, stats_expected=None) -> CorrectnessStats: + mismatch = _mismatch_mask(case, output, expected, nranks) + mismatches = int(cp.count_nonzero(mismatch).item()) + total = int(output.size) + if total == 0: + return CorrectnessStats(ok=mismatches == 0) + + output_values = _stats_values(case, output) + expected_values = _stats_values(case, expected) if stats_expected is None else stats_expected.astype(cp.float64) + abs_diff = cp.abs(output_values - expected_values) + return CorrectnessStats( + ok=mismatches == 0, + max_abs_diff=float(cp.max(abs_diff).item()), + mean_abs_diff=float(cp.mean(abs_diff).item()), + mismatches=mismatches, + total=total, + ) + + +def _stats_values(case: Any, values): + # Convert storage buffers into numeric values before computing max/mean diff. + if case.dtype_spec.fp8_format is not None: + return _decode_fp8_array(case.dtype_spec.fp8_format, values) + if cp.issubdtype(values.dtype, cp.floating): + return values.astype(cp.float64) + return values.astype(cp.int64) + + +def _expected_outputs(case: Any, nranks: int, iteration: int): + if case.collective == "allreduce": + encoded_inputs = _encoded_rank_inputs(case, nranks, iteration) + if case.dtype_spec.fp8_format is not None: + stats_expected = _expected_fp8_accum_values(case, encoded_inputs) + return _encode_reduced_output(case, stats_expected), stats_expected + return _encode_reduced_output(case, sum(values.astype(cp.float32) for values in encoded_inputs)), None + + expected = cp.empty_like(case.output) + chunk = case.input.size + for rank, values in enumerate(_encoded_rank_inputs(case, nranks, iteration)): + expected[rank * chunk : (rank + 1) * chunk] = values.reshape(-1) + return expected, None + + +def _encoded_rank_inputs(case: Any, nranks: int, iteration: int) -> list[Any]: + return [_encode_correctness_input(case, _correctness_input_values(case, rank, iteration)) for rank in range(nranks)] + + +def _expected_fp8_accum_values(case: Any, encoded_inputs: list[Any]): + fp8_format = case.dtype_spec.fp8_format + if fp8_format is None: + raise ValueError("FP8 format is required") + + accum_dtype = config_accum_dtype(case) + if accum_dtype == _mscclpp().DataType.float16: + acc = cp.zeros_like(_decode_fp8_array(fp8_format, encoded_inputs[0]), dtype=cp.float16) + for values in encoded_inputs: + acc = (acc + _decode_fp8_array(fp8_format, values).astype(cp.float16)).astype(cp.float16) + return acc.astype(cp.float32) + + if accum_dtype == _mscclpp().DataType.float32: + acc = cp.zeros_like(_decode_fp8_array(fp8_format, encoded_inputs[0]), dtype=cp.float32) + for values in encoded_inputs: + acc += _decode_fp8_array(fp8_format, values).astype(cp.float32) + return acc + + acc = encoded_inputs[0] + for values in encoded_inputs[1:]: + acc = _encode_fp8_values(fp8_format, _decode_fp8_array(fp8_format, acc) + _decode_fp8_array(fp8_format, values)) + return _decode_fp8_array(fp8_format, acc).astype(cp.float32) + + +def _encode_reduced_output(case: Any, values): + if case.dtype_spec.fp8_format is not None: + return _encode_fp8_values(case.dtype_spec.fp8_format, values) + return values.astype(case.output.dtype) + + +def _compare_output(case: Any, output, expected, nranks: int) -> bool: + return bool(cp.all(~_mismatch_mask(case, output, expected, nranks)).item()) + + +def _mismatch_mask(case: Any, output, expected, nranks: int): + tolerance = _comparison_tolerance(case, nranks) + if tolerance is None: + return output != expected + rtol, atol = tolerance + return ~cp.isclose(_stats_values(case, output), _stats_values(case, expected), rtol=rtol, atol=atol) + + +def _comparison_tolerance(case: Any, nranks: int) -> tuple[float, float] | None: + scale = max(1, nranks) if case.collective == "allreduce" else 1 + if case.dtype_spec.fp8_format is not None: + accum_dtype = config_accum_dtype(case) + if accum_dtype == _mscclpp().DataType.float32: + return None + atol = _max_fp8_spacing(case.dtype_spec.fp8_format, float(scale)) + if accum_dtype == _mscclpp().DataType.float16: + return (0.0, atol) + return (0.0, atol * 2) + if case.dtype_spec.cupy_dtype == cp.float16: + return (1.0e-2, 5.0e-4 * scale) + if case.dtype_spec.cupy_dtype == cp.float32: + return (1.0e-5 * scale, 1.0e-6 * scale) + return None + + +_FP8_TABLES: dict[str, list[tuple[int, float]]] = {} +_FP8_LOOKUP_CACHE: dict[str, tuple[Any, Any]] = {} +_FP8_SPACING_CACHE: dict[tuple[str, float], float] = {} + + +def _encode_fp8_values(fp8_format: str, values): + values = values.astype(cp.float32) + if fp8_format == "e4m3b15": + return _encode_e4m3b15_values(values) + + # Round each value to the nearest representable FP8 value (ties to even). + table_values, table_bytes = _fp8_lookup_arrays(fp8_format) + flat_values = values.ravel() + + # For each value find its two surrounding table entries: lower <= value <= upper. + upper = cp.clip(cp.searchsorted(table_values, flat_values), 1, table_values.size - 1) + lower = upper - 1 + + # Pick the closer neighbor; on an exact tie pick the one with an even byte. + dist_to_upper = table_values[upper] - flat_values + dist_to_lower = flat_values - table_values[lower] + upper_is_even = (table_bytes[upper] & cp.uint8(1)) == 0 + pick_upper = (dist_to_upper < dist_to_lower) | ((dist_to_upper == dist_to_lower) & upper_is_even) + + return cp.where(pick_upper, table_bytes[upper], table_bytes[lower]).reshape(values.shape) + + +def _fp8_lookup_arrays(fp8_format: str): + # Cache a sorted (value -> byte) table per format for fast nearest-value lookup. + if fp8_format in _FP8_LOOKUP_CACHE: + return _FP8_LOOKUP_CACHE[fp8_format] + + # Different bytes can decode to the same value (e.g. +0 and -0); keep one byte per value. + byte_for_value: dict[float, int] = {} + for byte, value in _FP8_TABLES.setdefault(fp8_format, _build_fp8_table(fp8_format)): + if value not in byte_for_value or byte < byte_for_value[value]: + byte_for_value[value] = byte + + table = sorted(byte_for_value.items()) + table_values = cp.asarray([value for value, _ in table], dtype=cp.float32) + table_bytes = cp.asarray([byte for _, byte in table], dtype=cp.uint8) + _FP8_LOOKUP_CACHE[fp8_format] = (table_values, table_bytes) + return _FP8_LOOKUP_CACHE[fp8_format] + + +def _max_fp8_spacing(fp8_format: str, max_abs_value: float) -> float: + cache_key = (fp8_format, max_abs_value) + if cache_key in _FP8_SPACING_CACHE: + return _FP8_SPACING_CACHE[cache_key] + + values = sorted( + { + value + for _, value in _FP8_TABLES.setdefault(fp8_format, _build_fp8_table(fp8_format)) + if abs(value) <= max_abs_value + } + ) + if len(values) < 2: + spacing = 0.0 + else: + spacing = max(right - left for left, right in zip(values, values[1:])) + _FP8_SPACING_CACHE[cache_key] = spacing + return spacing + + +def _fp8_max_abs_value(fp8_format: str) -> float: + return max(abs(value) for _, value in _FP8_TABLES.setdefault(fp8_format, _build_fp8_table(fp8_format))) + + +def _encode_e4m3b15_values(values): + # Mirrors the device e4m3b15 encode (gpu_data_types.hpp): clamp the fp16 intermediate + # to 0x3F80 (+/-1.875) so the max encodable byte is 0x7F/0xFF. + fp16_bits = values.astype(cp.float16).view(cp.uint16) + abs_fp16 = fp16_bits & cp.uint16(0x7FFF) + abs_fp16 = cp.minimum(abs_fp16, cp.uint16(0x3F80)).astype(cp.uint32) + sign16 = (fp16_bits & cp.uint16(0x8000)).astype(cp.uint32) + adjusted = abs_fp16 * cp.uint32(2) + cp.uint32(0x0080) + return (((sign16 | adjusted) >> cp.uint32(8)) & cp.uint32(0xFF)).astype(cp.uint8) + + +def _build_fp8_table(fp8_format: str) -> list[tuple[int, float]]: + table = [] + for byte in range(256): + value = _decode_fp8_scalar(fp8_format, byte) + if not math.isnan(value): + table.append((byte, value)) + return table + + +def _decode_fp8_scalar(fp8_format: str, byte: int) -> float: + if fp8_format == "e4m3fnuz" and byte == 0x80: + return float("nan") + sign = -1.0 if byte & 0x80 else 1.0 + return sign * _decode_fp8_positive(fp8_format, byte & 0x7F) + + +def _decode_fp8_positive(fp8_format: str, byte: int) -> float: + exp = (byte >> 3) & 0xF + mant = byte & 0x7 + if fp8_format == "e4m3fn" and exp == 0xF and mant == 0x7: + return float("nan") + if exp == 0 and mant == 0: + return 0.0 + if fp8_format == "e4m3fn": + return math.ldexp(mant / 8.0, -6) if exp == 0 else math.ldexp(1.0 + mant / 8.0, exp - 7) + if fp8_format == "e4m3fnuz": + return math.ldexp(mant / 8.0, -7) if exp == 0 else math.ldexp(1.0 + mant / 8.0, exp - 8) + if fp8_format == "e4m3b15": + return math.ldexp(mant / 8.0, -14) if exp == 0 else math.ldexp(1.0 + mant / 8.0, exp - 15) + raise ValueError(f"Unknown FP8 format: {fp8_format}") + + +def _decode_fp8_array(fp8_format: str, values): + bits = values.astype(cp.int32) + sign = (bits >> 7) & 1 + exp = (bits >> 3) & 0xF + mant = bits & 0x7 + + if fp8_format == "e4m3fn": + subnormal = cp.ldexp(mant.astype(cp.float32) / cp.float32(8.0), cp.int32(-6)) + normal = cp.ldexp(cp.float32(1.0) + mant.astype(cp.float32) / cp.float32(8.0), exp.astype(cp.int32) - 7) + decoded = cp.where(exp == 0, subnormal, normal) + decoded = cp.where((exp == 0xF) & (mant == 0x7), cp.nan, decoded) + elif fp8_format == "e4m3fnuz": + subnormal = cp.ldexp(mant.astype(cp.float32) / cp.float32(8.0), cp.int32(-7)) + normal = cp.ldexp(cp.float32(1.0) + mant.astype(cp.float32) / cp.float32(8.0), exp.astype(cp.int32) - 8) + decoded = cp.where(exp == 0, subnormal, normal) + elif fp8_format == "e4m3b15": + subnormal = cp.ldexp(mant.astype(cp.float32) / cp.float32(8.0), cp.int32(-14)) + normal = cp.ldexp(cp.float32(1.0) + mant.astype(cp.float32) / cp.float32(8.0), exp.astype(cp.int32) - 15) + decoded = cp.where(exp == 0, subnormal, normal) + else: + raise ValueError(f"Unknown FP8 format: {fp8_format}") + + result = cp.where(sign == 1, -decoded, decoded) + if fp8_format == "e4m3fnuz": + result = cp.where(bits == 0x80, cp.float32(float("nan")), result) + return result diff --git a/python/mscclpp_benchmark/gpu.py b/python/mscclpp_benchmark/gpu.py new file mode 100644 index 00000000..248065ef --- /dev/null +++ b/python/mscclpp_benchmark/gpu.py @@ -0,0 +1,203 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +_API_NAMES = { + "get_device_count": ("hipGetDeviceCount", "cudaGetDeviceCount"), + "get_device": ("hipGetDevice", "cudaGetDevice"), + "get_device_properties": ("hipGetDeviceProperties", "cudaGetDeviceProperties"), + "set_device": ("hipSetDevice", "cudaSetDevice"), + "stream_begin_capture": ("hipStreamBeginCapture", "cudaStreamBeginCapture"), + "stream_end_capture": ("hipStreamEndCapture", "cudaStreamEndCapture"), + "graph_instantiate": ("hipGraphInstantiate", "cudaGraphInstantiate"), + "graph_launch": ("hipGraphLaunch", "cudaGraphLaunch"), + "graph_destroy": ("hipGraphDestroy", "cudaGraphDestroy"), + "graph_exec_destroy": ("hipGraphExecDestroy", "cudaGraphExecDestroy"), + "get_error_string": ("hipGetErrorString", "cudaGetErrorString"), + "runtime_get_version": ("hipRuntimeGetVersion", "cudaRuntimeGetVersion"), +} + + +@dataclass(frozen=True) +class _Runtime: + name: str + success: Any + capture_mode_relaxed: Any + funcs: dict[str, Callable[..., Any] | None] + + @classmethod + def create(cls, name: str, module: Any, success: Any, capture_mode_relaxed: Any) -> "_Runtime": + index = 0 if name == "hip" else 1 + funcs = { + attr: (None if names[index] is None else getattr(module, names[index])) + for attr, names in _API_NAMES.items() + } + return cls(name=name, success=success, capture_mode_relaxed=capture_mode_relaxed, funcs=funcs) + + def call(self, name: str, *args: Any) -> tuple[Any, ...]: + fn = self.funcs[name] + if fn is None: + raise RuntimeError(f"{name} is not available for {self.name}") + result = fn(*args) + if not isinstance(result, tuple): + result = (result,) + self.check(result[0], name) + return result[1:] + + def check(self, error: Any, api: str) -> None: + if error == self.success: + return + result = self.funcs["get_error_string"](error) + if not isinstance(result, tuple): + result = (result,) + err, message = result + if err != self.success: + raise RuntimeError(f"{api} failed with error {int(error)}") + decoded = message.decode("utf-8") if isinstance(message, bytes) else str(message) + raise RuntimeError(f"{api} failed: {decoded} ({int(error)})") + + +def _load_runtime() -> _Runtime: + errors: list[str] = [] + + try: + from hip import hip + + runtime = _Runtime.create( + name="hip", + module=hip, + success=hip.hipError_t.hipSuccess, + capture_mode_relaxed=hip.hipStreamCaptureMode.hipStreamCaptureModeRelaxed, + ) + count = runtime.call("get_device_count")[0] + if count and count > 0: + return runtime + errors.append(f"hipGetDeviceCount returned count={count}") + except ImportError as exc: + errors.append(f"hip-python unavailable: {exc}") + + try: + from cuda.bindings import runtime as cuda_runtime + + runtime = _Runtime.create( + name="cuda", + module=cuda_runtime, + success=cuda_runtime.cudaError_t.cudaSuccess, + capture_mode_relaxed=cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeRelaxed, + ) + count = runtime.call("get_device_count")[0] + if count and count > 0: + return runtime + errors.append(f"cudaGetDeviceCount returned count={count}") + except ImportError as exc: + errors.append(f"cuda-bindings unavailable: {exc}") + + raise RuntimeError("No usable CUDA/HIP Python runtime found: " + "; ".join(errors)) + + +_RUNTIME = _load_runtime() + + +class Graph: + def __init__(self, graph_exec: Any) -> None: + self._graph_exec = graph_exec + + def launch(self, stream: Any) -> None: + _api("graph_launch")(self._graph_exec, _stream_ptr(stream)) + + def close(self) -> None: + if self._graph_exec is not None: + _api("graph_exec_destroy")(self._graph_exec) + self._graph_exec = None + + +def init_runtime() -> None: + return None + + +def runtime_name() -> str: + return _RUNTIME.name + + +def _runtime_version_raw() -> int: + return int(_api("runtime_get_version")()[0]) + + +def version() -> tuple[int, int, int]: + version_value = _runtime_version_raw() + if _RUNTIME.name == "hip": + return version_value // 10_000_000, (version_value // 100_000) % 100, version_value % 100_000 + return version_value // 1000, (version_value % 1000) // 10, version_value % 10 + + +def capture_graph(stream: Any, capture_fn: Callable[[], None]) -> Graph: + _api("set_device")(current_device()) + stream_ptr = _stream_ptr(stream) + _api("stream_begin_capture")(stream_ptr, _RUNTIME.capture_mode_relaxed) + + graph = None + try: + capture_fn() + graph = _api("stream_end_capture")(stream_ptr)[0] + except Exception: + try: + _api("stream_end_capture")(stream_ptr) + except Exception: + pass + raise + + try: + graph_exec = _instantiate_graph(graph) + return Graph(graph_exec) + finally: + if graph is not None: + _api("graph_destroy")(graph) + + +def current_device() -> int: + return int(_api("get_device")()[0]) + + +def device_name(device_id: int | None = None) -> str: + if device_id is None: + device_id = current_device() + prop = _api("get_device_properties")(int(device_id))[0] + name = getattr(prop, "name", "UNKNOWN") + return name.decode("utf-8") if isinstance(name, bytes) else str(name) + + +def _stream_ptr(stream: Any) -> int: + return int(getattr(stream, "ptr", stream)) + + +def _instantiate_graph(graph: Any) -> Any: + if _RUNTIME.name == "hip": + return _api("graph_instantiate")(graph, None, 0)[0] + return _api("graph_instantiate")(graph, 0)[0] + + +def _api(name: str) -> Callable[..., tuple[Any, ...]]: + api = globals().get(name) + if api is None: + api = __getattr__(name) + return api + + +def _make_api(name: str) -> Callable[..., tuple[Any, ...]]: + def api(*args: Any) -> tuple[Any, ...]: + return _RUNTIME.call(name, *args) + + api.__name__ = name + return api + + +def __getattr__(name: str) -> Callable[..., tuple[Any, ...]]: + if name in _API_NAMES: + api = _make_api(name) + globals()[name] = api + return api + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/mscclpp_benchmark/tuner.py b/python/mscclpp_benchmark/tuner.py new file mode 100644 index 00000000..42c15ab5 --- /dev/null +++ b/python/mscclpp_benchmark/tuner.py @@ -0,0 +1,83 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from typing import Any, Callable, Iterable + +from mscclpp_benchmark.tuning_config import TunedConfig + + +class OfflineTuner: + def __init__( + self, + comm: Any, + *, + candidate_nblocks: Iterable[int], + candidate_nthreads: Iterable[int], + n_warmup: int, + n_graph_launches: int, + n_ops_per_graph: int, + candidate_algorithms: Callable[[Any, Any], list[tuple[Any, Any]]], + check_correctness: Callable[..., bool], + measure: Callable[..., float | None], + ) -> None: + self.comm = comm + self.candidate_nblocks = tuple(candidate_nblocks) + self.candidate_nthreads = tuple(candidate_nthreads) + self.n_warmup = n_warmup + self.n_graph_launches = n_graph_launches + self.n_ops_per_graph = n_ops_per_graph + self._candidate_algorithms = candidate_algorithms + self._check_correctness = check_correctness + self._measure = measure + + def tune(self, case: Any) -> TunedConfig | None: + best_config: TunedConfig | None = None + best_time_us = float("inf") + symmetric_memory = bool(getattr(case, "symmetric_memory", False)) + candidates = self._candidate_algorithms(self.comm, case) + if not candidates: + if self.comm.rank == 0: + print( + f"[skip] no supported tuning candidates for collective={case.collective} " + f"size={case.message_size}", + flush=True, + ) + return None + for algorithm, candidate_spec in candidates: + for nblocks in self.candidate_nblocks: + if candidate_spec.max_nblocks is not None and nblocks > candidate_spec.max_nblocks: + continue + for nthreads in self.candidate_nthreads: + config = TunedConfig( + algorithm=algorithm.name, + nblocks=nblocks, + nthreads=nthreads, + symmetric_memory=symmetric_memory, + ) + if not self._check_correctness(self.comm, case, config): + self.comm.reset(config) + continue + time_us = self._measure( + self.comm, + case, + config, + n_warmup=self.n_warmup, + n_graph_launches=self.n_graph_launches, + n_ops_per_graph=self.n_ops_per_graph, + ) + self.comm.reset(config) + if time_us is None or time_us >= best_time_us: + continue + best_time_us = time_us + best_config = TunedConfig( + algorithm=algorithm.name, + nblocks=nblocks, + nthreads=nthreads, + symmetric_memory=symmetric_memory, + time_us=time_us, + ) + if best_config is None: + return self.comm.resolve_config(case) + return best_config diff --git a/python/mscclpp_benchmark/tuning_config.py b/python/mscclpp_benchmark/tuning_config.py new file mode 100644 index 00000000..2a914ec9 --- /dev/null +++ b/python/mscclpp_benchmark/tuning_config.py @@ -0,0 +1,242 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +import re +from bisect import bisect_left +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +_KNOWN_GPU_SKUS = ("GB300", "MI300X", "H100", "A100") + + +@dataclass(frozen=True) +class HardwareProfile: + sku: str | None = None + scale: int | None = None + + +@dataclass(frozen=True) +class TunedConfig: + algorithm: str + nblocks: int | None = None + nthreads: int | None = None + symmetric_memory: bool = False + time_us: float | None = None + + +@dataclass(order=True, frozen=True) +class TunedConfigBySize: + message_size: int + config: TunedConfig + + +class TunedConfigStore: + def __init__(self, profiles: dict[HardwareProfile, dict[str, list[TunedConfigBySize]]]) -> None: + self._profiles = profiles + + @classmethod + def empty(cls) -> "TunedConfigStore": + return cls({}) + + @classmethod + def load_path(cls, path: str | Path) -> "TunedConfigStore": + with Path(path).open("r", encoding="utf-8") as handle: + return cls.from_payload(json.load(handle)) + + @classmethod + def from_payload(cls, payload: Any) -> "TunedConfigStore": + if not isinstance(payload, dict): + raise ValueError("MSCCL++ tuned config must be a JSON object") + raw_profiles = payload.get("profiles") + if not isinstance(raw_profiles, list): + raise ValueError("MSCCL++ tuned config must contain a 'profiles' list") + profiles: dict[HardwareProfile, dict[str, list[TunedConfigBySize]]] = {} + for raw_profile in raw_profiles: + profile = _profile_from_payload(raw_profile) + profiles[profile] = _configs_by_collective_from_payload(raw_profile.get("collectives", {})) + return cls(profiles) + + def select(self, profile: HardwareProfile, collective: str, message_size: int) -> TunedConfig | None: + for _, configs_by_collective in _matching_profiles(self._profiles, profile): + config = _select_config(configs_by_collective, collective, message_size) + if config is not None: + return config + return None + + def upsert(self, profile: HardwareProfile, collective: str, message_size: int, config: TunedConfig) -> None: + configs = self._profiles.setdefault(profile, {}).setdefault(collective, []) + for index, existing in enumerate(configs): + if existing.message_size == message_size: + configs[index] = TunedConfigBySize(message_size, config) + break + else: + configs.append(TunedConfigBySize(message_size, config)) + configs.sort(key=lambda item: item.message_size) + + def write_path(self, path: str | Path) -> None: + profiles_payload: list[dict[str, Any]] = [] + for profile, configs_by_collective in sorted( + self._profiles.items(), + key=lambda item: (item[0].sku is None, item[0].sku or "", item[0].scale is None, item[0].scale or 0), + ): + collectives: dict[str, list[dict[str, Any]]] = {} + for collective, configs in sorted(configs_by_collective.items()): + collectives[collective] = [_config_entry_payload(item) for item in sorted(configs)] + profile_payload: dict[str, Any] = {} + if profile.sku is not None: + profile_payload["sku"] = profile.sku + if profile.scale is not None: + profile_payload["scale"] = profile.scale + profile_payload["collectives"] = collectives + profiles_payload.append(profile_payload) + + with Path(path).open("w", encoding="utf-8") as handle: + handle.write(_format_tuned_config_json({"version": 1, "profiles": profiles_payload})) + + +def normalize_sku(raw_sku: str) -> str: + upper_sku = raw_sku.upper() + for known_sku in _KNOWN_GPU_SKUS: + if known_sku in upper_sku: + return known_sku + normalized = re.sub(r"[^A-Z0-9]+", "_", upper_sku).strip("_") + return normalized or "UNKNOWN" + + +def _profile_from_payload(raw_profile: Any) -> HardwareProfile: + if not isinstance(raw_profile, dict): + raise ValueError(f"Invalid tuned config profile: {raw_profile!r}") + raw_sku = raw_profile.get("sku") + return HardwareProfile( + sku=None if raw_sku is None else normalize_sku(str(raw_sku)), + scale=_optional_positive_int(raw_profile.get("scale"), "scale"), + ) + + +def _matching_profiles( + profiles: dict[HardwareProfile, dict[str, list[TunedConfigBySize]]], + runtime_profile: HardwareProfile, +) -> list[tuple[int, dict[str, list[TunedConfigBySize]]]]: + matches: list[tuple[int, dict[str, list[TunedConfigBySize]]]] = [] + for profile, configs_by_collective in profiles.items(): + specificity = _profile_match_specificity(profile, runtime_profile) + if specificity is not None: + matches.append((specificity, configs_by_collective)) + return sorted(matches, key=lambda item: item[0], reverse=True) + + +def _profile_match_specificity(profile: HardwareProfile, runtime_profile: HardwareProfile) -> int | None: + specificity = 0 + if profile.sku is not None: + if profile.sku != runtime_profile.sku: + return None + specificity += 1 + if profile.scale is not None: + if profile.scale != runtime_profile.scale: + return None + specificity += 1 + return specificity + + +def _select_config( + configs_by_collective: dict[str, list[TunedConfigBySize]], collective: str, message_size: int +) -> TunedConfig | None: + configs = configs_by_collective.get(collective, []) + if not configs: + return None + sizes = [item.message_size for item in configs] + index = bisect_left(sizes, message_size) + if index == len(sizes): + return configs[-1].config + if sizes[index] == message_size or index == 0: + return configs[index].config + return configs[index - 1].config + + +def _configs_by_collective_from_payload(payload: Any) -> dict[str, list[TunedConfigBySize]]: + if not isinstance(payload, dict): + raise ValueError("MSCCL++ tuned config collectives must be an object") + + result: dict[str, list[TunedConfigBySize]] = {} + for collective, raw_entries in payload.items(): + if isinstance(raw_entries, dict): + raw_entries = raw_entries.get("configs", []) + if not isinstance(raw_entries, list): + continue + configs = [] + for raw_entry in raw_entries: + if not isinstance(raw_entry, dict): + raise ValueError(f"Invalid tuned config entry for {collective}: {raw_entry!r}") + configs.append( + TunedConfigBySize( + message_size=_parse_positive_int(raw_entry.get("message_size"), "message_size"), + config=TunedConfig( + algorithm=str(raw_entry["algorithm"]), + nblocks=_optional_int(raw_entry.get("nblocks")), + nthreads=_optional_int(raw_entry.get("nthreads")), + symmetric_memory=_optional_bool(raw_entry.get("symmetric_memory", False)), + time_us=_optional_float(raw_entry.get("time_us")), + ), + ) + ) + result[str(collective)] = sorted(configs) + return result + + +def _config_entry_payload(item: TunedConfigBySize) -> dict[str, Any]: + payload: dict[str, Any] = {"message_size": item.message_size, "algorithm": item.config.algorithm} + if item.config.nblocks is not None: + payload["nblocks"] = item.config.nblocks + if item.config.nthreads is not None: + payload["nthreads"] = item.config.nthreads + if item.config.symmetric_memory: + payload["symmetric_memory"] = item.config.symmetric_memory + if item.config.time_us is not None: + payload["time_us"] = item.config.time_us + return payload + + +def _format_tuned_config_json(payload: dict[str, Any]) -> str: + text = json.dumps(payload, indent=2) + pattern = re.compile( + r"(?m)^(?P +)\{\n" + r'(?P(?P=indent) "message_size": [^\n]+,?\n(?:(?P=indent) "[^"]+": [^\n]+,?\n)*)' + r"(?P=indent)\}(?P,?)$" + ) + + def compact(match: re.Match[str]) -> str: + body = " ".join(line.strip() for line in match.group("body").splitlines()) + return f"{match.group('indent')}{{{body}}}{match.group('comma')}" + + return pattern.sub(compact, text) + "\n" + + +def _optional_int(value: Any | None) -> int | None: + return None if value is None else int(value) + + +def _optional_float(value: Any | None) -> float | None: + return None if value is None else float(value) + + +def _optional_positive_int(value: Any | None, name: str) -> int | None: + return None if value is None else _parse_positive_int(value, name) + + +def _optional_bool(value: Any | None) -> bool | None: + if value is None: + return None + if isinstance(value, bool): + return value + raise ValueError(f"Expected boolean value, got {value!r}") + + +def _parse_positive_int(value: Any, name: str) -> int: + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {parsed}") + return parsed diff --git a/python/requirements_cuda12.txt b/python/requirements_cuda12.txt index 71572714..fcc59660 100644 --- a/python/requirements_cuda12.txt +++ b/python/requirements_cuda12.txt @@ -1,5 +1,6 @@ mpi4py cupy-cuda12x +cuda-bindings>=12,<13 prettytable netifaces pytest diff --git a/python/requirements_cuda13.txt b/python/requirements_cuda13.txt index 95e99533..19ad93d7 100644 --- a/python/requirements_cuda13.txt +++ b/python/requirements_cuda13.txt @@ -1,5 +1,6 @@ mpi4py cupy-cuda13x +cuda-bindings>=13,<14 prettytable netifaces pytest diff --git a/python/requirements_rocm6.txt b/python/requirements_rocm6.txt index 757d4e26..bcc22dfb 100644 --- a/python/requirements_rocm6.txt +++ b/python/requirements_rocm6.txt @@ -7,4 +7,5 @@ numpy matplotlib sortedcontainers blake3 -pybind11 \ No newline at end of file +pybind11 +hip-python>=6,<7 \ No newline at end of file diff --git a/python/requirements_cuda11.txt b/python/requirements_rocm7.txt similarity index 71% rename from python/requirements_cuda11.txt rename to python/requirements_rocm7.txt index a9786071..982f9ae6 100644 --- a/python/requirements_cuda11.txt +++ b/python/requirements_rocm7.txt @@ -1,5 +1,5 @@ mpi4py -cupy-cuda11x +cupy prettytable netifaces pytest @@ -7,4 +7,5 @@ numpy matplotlib sortedcontainers blake3 -pybind11 \ No newline at end of file +pybind11 +hip-python>=7,<8 diff --git a/python/test/executor_test.py b/python/test/executor_test.py index 8a309de5..0159b8fa 100644 --- a/python/test/executor_test.py +++ b/python/test/executor_test.py @@ -14,6 +14,7 @@ from mscclpp import CommGroup, GpuBuffer from mscclpp.utils import KernelBuilder, pack import os import struct +from typing import Callable import cupy as cp from mpi4py import MPI @@ -34,13 +35,13 @@ def parse_dtype(dtype_str): raise ValueError(f"Unknown data type: {dtype_str}") -def bench_time(n_iters: int, n_graph_iters: int, func): - # capture cuda graph for n_iters of the kernel launch +def bench_time(n_iters: int, n_graph_iters: int, funcs: list[Callable]): + """Benchmark execution time. `funcs` is a list of callables; iteration i runs funcs[i % len(funcs)].""" stream = cp.cuda.Stream(non_blocking=True) with stream: stream.begin_capture() for i in range(n_iters): - func(stream) + funcs[i % len(funcs)](stream) graph = stream.end_capture() # now run a warm up round @@ -61,15 +62,17 @@ def bench_time(n_iters: int, n_graph_iters: int, func): def bench_correctness( collective: str, - input_buf: cp.ndarray, - result_buf: cp.ndarray, - test_buf: cp.ndarray, + input_bufs: list[cp.ndarray], + result_bufs: list[cp.ndarray], + test_bufs: list[cp.ndarray], dtype_str: str, rank: int, num_ranks: int, n_iters: int, - func, + funcs: list[Callable], + split_mask: int = 0, ): + """Validate correctness. Buffers and funcs are parallel lists; iteration i uses index i % len(funcs).""" type_size = cp.dtype(parse_dtype(dtype_str)).itemsize fill_data_kernel_name = "fill_data_%s" % dtype_str @@ -79,8 +82,12 @@ def bench_correctness( coll = "reduce_scatter" elif "allreduce" in collective: coll = "all_reduce" - else: + elif "alltoall" in collective: coll = "all_to_all" + elif "sendrecv" in collective: + coll = "send_recv" + else: + raise ValueError(f"Unknown collective: {collective}") test_data_kernel_name = "test_data_%s_%s" % (coll, dtype_str) file_dir = os.path.dirname(os.path.abspath(__file__)) @@ -97,11 +104,20 @@ def bench_correctness( with stream: stream.begin_capture() for i in range(n_iters): - fill_data_params = pack(input_buf) + struct.pack("Q", input_buf.nbytes // type_size) + pack(rank, i) + idx = i % len(funcs) + cur_input = input_bufs[idx] + cur_result = result_bufs[idx] + cur_test = test_bufs[idx] + + fill_data_params = ( + pack(cur_input) + struct.pack("Q", cur_input.nbytes // type_size) + pack(rank, i, split_mask) + ) fill_data_kernel.launch_kernel(fill_data_params, nblocks, nthreads, 0, stream) - func(stream) + funcs[idx](stream) test_data_params = ( - pack(result_buf, test_buf) + struct.pack("Q", input_buf.nbytes // type_size) + pack(num_ranks, rank, i) + pack(cur_result, cur_test) + + struct.pack("Q", cur_input.nbytes // type_size) + + pack(num_ranks, rank, i, split_mask) ) test_data_kernel.launch_kernel(test_data_params, nblocks, nthreads, 0, stream) graph = stream.end_capture() @@ -143,10 +159,20 @@ def build_bufs( rank: int, num_ranks: int, ): + """Allocate input/result/test buffers. Returns parallel lists (length 2 for sendrecv double-buffering, + length 1 otherwise) so callers can iterate uniformly.""" type_size = cp.dtype(dtype).itemsize assert (size % type_size) == 0, "size %d not multiple of type size %d" % (size, type_size) nelems = size // type_size + # Sendrecv uses double buffering: build two parallel buffer slots. + if "sendrecv" in collective: + n_slots = 2 + input_bufs = [GpuBuffer(nelems, dtype=dtype) for _ in range(n_slots)] + result_bufs = [GpuBuffer(nelems, dtype=dtype) for _ in range(n_slots)] + test_bufs = [cp.zeros(nelems, dtype=dtype) for _ in range(n_slots)] + return input_bufs, result_bufs, test_bufs, nelems + if "allgather" in collective: assert (nelems % num_ranks) == 0, "nelems %d not multiple of num_ranks %d" % (nelems, num_ranks) nelems_input = nelems if in_place else nelems // num_ranks @@ -173,7 +199,7 @@ def build_bufs( test_buf = cp.zeros(nelems, dtype=dtype) - return input_buf, result_buf, test_buf + return [input_buf], [result_buf], [test_buf], nelems def main( @@ -184,8 +210,14 @@ def main( packet_type: PacketType = PacketType.LL16, n_iters: int = 10, n_graph_iters: int = 10, + split_mask: int = 0, ): mscclpp_group = CommGroup(MPI.COMM_WORLD) + if split_mask < 0 or (split_mask & (split_mask + 1)) != 0 or mscclpp_group.nranks % (split_mask + 1) != 0: + raise ValueError( + f"split_mask must be of the form 2^k - 1 and nranks ({mscclpp_group.nranks}) must be divisible " + f"by group_size ({split_mask + 1}), got split_mask={hex(split_mask)}" + ) cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use() executor = Executor(mscclpp_group.communicator) npkit_dump_dir = env().npkit_dump_dir @@ -195,7 +227,7 @@ def main( collective = execution_plan.collective dtype = parse_dtype(dtype_str) - input_buf, result_buf, test_buf = build_bufs( + input_bufs, result_bufs, test_bufs, nelem = build_bufs( collective, size, in_place, @@ -204,39 +236,48 @@ def main( mscclpp_group.nranks, ) - executor_func = lambda stream: executor.execute( - mscclpp_group.my_rank, - input_buf.data.ptr, - result_buf.data.ptr, - input_buf.nbytes, - result_buf.nbytes, - dtype_to_mscclpp_dtype(dtype_str), - execution_plan, - stream.ptr, - packet_type, - ) + executor_funcs = [ + ( + lambda stream, inp=inp, res=res: executor.execute( + mscclpp_group.my_rank, + inp.data.ptr, + res.data.ptr, + inp.nbytes, + res.nbytes, + dtype_to_mscclpp_dtype(dtype_str), + execution_plan, + stream.ptr, + packet_type, + ) + ) + for inp, res in zip(input_bufs, result_bufs) + ] mscclpp_group.barrier() bench_correctness( collective, - input_buf, - result_buf, - test_buf, + input_bufs, + result_bufs, + test_bufs, dtype_str, mscclpp_group.my_rank, mscclpp_group.nranks, n_iters, - executor_func, + executor_funcs, + split_mask=split_mask, ) mscclpp_group.barrier() - execution_time = bench_time(n_iters, n_graph_iters, executor_func) + execution_time = bench_time(n_iters, n_graph_iters, executor_funcs) if npkit_dump_dir is not None: npkit.dump(npkit_dump_dir) npkit.shutdown() + + result_nbytes = result_bufs[0].nbytes print( f"Rank: {mscclpp_group.my_rank} Execution time: {execution_time} us, " - f"data size: {result_buf.nbytes} bytes data type: {dtype_str} " + f"data size: {result_nbytes} bytes data type: {dtype_str} " + f"bandwidth: {result_nbytes / (execution_time * 1e-6) / (1024**3):.2f} GB/s, " f"packet type: {packet_type}" ) executor = None @@ -252,6 +293,9 @@ if __name__ == "__main__": parser.add_argument("--packet_type", type=str, default="LL16", help="Choose from LL8, LL16") parser.add_argument("--n_iters", type=int, default=10) parser.add_argument("--n_graph_iters", type=int, default=10) + parser.add_argument( + "--split_mask", type=lambda x: int(x, 0), default=0x0, help="split mask for sendrecv (e.g. 0x3)" + ) args = parser.parse_args() packet_type = PacketType.LL16 @@ -267,4 +311,5 @@ if __name__ == "__main__": packet_type, args.n_iters, args.n_graph_iters, + args.split_mask, ) diff --git a/python/test/executor_test_verifier.cu b/python/test/executor_test_verifier.cu index e7749197..96ab25c4 100644 --- a/python/test/executor_test_verifier.cu +++ b/python/test/executor_test_verifier.cu @@ -22,14 +22,19 @@ static __device__ unsigned int ranqd1(unsigned int seed) { // fill/test kernel pairs must have the same thread block size to // match their random number series. -#define FILL_DATA(FuncNameType, DataType) \ - extern "C" __global__ void __launch_bounds__(1024, 1) \ - fill_data_##FuncNameType(DataType* input_buf, size_t num_elems, int rank, int seq) { \ - unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ - seed = ranqd1(seed); \ - input_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ - } \ +// `split_mask` groups ranks together: group_size = split_mask + 1, group_id = rank / group_size. +// Data is seeded by group_id so that all ranks within a group produce the same fill, and ranks +// in different groups produce different fills. With split_mask == 0 this reduces to per-rank +// seeding (group_id == rank). +#define FILL_DATA(FuncNameType, DataType) \ + extern "C" __global__ void __launch_bounds__(1024, 1) \ + fill_data_##FuncNameType(DataType* input_buf, size_t num_elems, int rank, int seq, int split_mask) { \ + int seed_rank = rank / (split_mask + 1); \ + unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ + seed = ranqd1(seed); \ + input_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ + } \ } FILL_DATA(bfloat16, __nv_bfloat16) @@ -37,18 +42,20 @@ FILL_DATA(float16, __half) FILL_DATA(float32, float) FILL_DATA(int32, int) -#define TEST_DATA_ALL_GATHER(FuncNameType, DataType) \ - extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_gather_##FuncNameType( \ - DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \ - for (int rank = 0; rank < num_ranks; rank++) { \ - size_t rank_offset = rank * num_elems; \ - unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ - seed = ranqd1(seed); \ - test_buf[rank_offset + i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ - assert(result_buf[rank_offset + i] == test_buf[rank_offset + i]); \ - } \ - } \ +#define TEST_DATA_ALL_GATHER(FuncNameType, DataType) \ + extern "C" __global__ void __launch_bounds__(1024, 1) \ + test_data_all_gather_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \ + int my_rank, int seq, int split_mask) { \ + for (int rank = 0; rank < num_ranks; rank++) { \ + size_t rank_offset = rank * num_elems; \ + int seed_rank = rank / (split_mask + 1); \ + unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ + seed = ranqd1(seed); \ + test_buf[rank_offset + i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ + assert(result_buf[rank_offset + i] == test_buf[rank_offset + i]); \ + } \ + } \ } TEST_DATA_ALL_GATHER(bfloat16, __nv_bfloat16) @@ -56,25 +63,27 @@ TEST_DATA_ALL_GATHER(float16, __half) TEST_DATA_ALL_GATHER(float32, float) TEST_DATA_ALL_GATHER(int32, int) -#define TEST_DATA_ALL_REDUCE(FuncNameType, DataType, Eps) \ - extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_reduce_##FuncNameType( \ - DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \ - for (int rank = 0; rank < num_ranks; rank++) { \ - unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ - if (rank == 0) { \ - test_buf[i] = 0; \ - } \ - seed = ranqd1(seed); \ - test_buf[i] += DataType(seed % blockDim.x) / DataType(blockDim.x); \ - } \ - } \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ - float expected = float(test_buf[i]); \ - float result = float(result_buf[i]); \ - float tol = Eps * num_ranks * (1.0f + abs(expected)); \ - assert(abs(result - expected) <= tol); \ - } \ +#define TEST_DATA_ALL_REDUCE(FuncNameType, DataType, Eps) \ + extern "C" __global__ void __launch_bounds__(1024, 1) \ + test_data_all_reduce_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \ + int my_rank, int seq, int split_mask) { \ + for (int rank = 0; rank < num_ranks; rank++) { \ + int seed_rank = rank / (split_mask + 1); \ + unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ + if (rank == 0) { \ + test_buf[i] = 0; \ + } \ + seed = ranqd1(seed); \ + test_buf[i] += DataType(seed % blockDim.x) / DataType(blockDim.x); \ + } \ + } \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ + float expected = float(test_buf[i]); \ + float result = float(result_buf[i]); \ + float tol = Eps * num_ranks * (1.0f + abs(expected)); \ + assert(abs(result - expected) <= tol); \ + } \ } TEST_DATA_ALL_REDUCE(bfloat16, __nv_bfloat16, 7.8125e-3f) @@ -83,12 +92,14 @@ TEST_DATA_ALL_REDUCE(float32, float, 1.1920929e-7f) TEST_DATA_ALL_REDUCE(int32, int, 0.0f) #define TEST_DATA_REDUCE_SCATTER(FuncNameType, DataType, Eps) \ - extern "C" __global__ void __launch_bounds__(1024, 1) test_data_reduce_scatter_##FuncNameType( \ - DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \ + extern "C" __global__ void __launch_bounds__(1024, 1) \ + test_data_reduce_scatter_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, \ + int num_ranks, int my_rank, int seq, int split_mask) { \ int nem_elems_per_rank = num_elems / num_ranks; \ int offset = nem_elems_per_rank * my_rank; \ for (int rank = 0; rank < num_ranks; rank++) { \ - unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \ + int seed_rank = rank / (split_mask + 1); \ + unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \ for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ if (rank == 0) { \ test_buf[i] = 0; \ @@ -112,25 +123,51 @@ TEST_DATA_REDUCE_SCATTER(float16, __half, 9.765625e-4f) TEST_DATA_REDUCE_SCATTER(float32, float, 1.1920929e-7f) TEST_DATA_REDUCE_SCATTER(int32, int, 0.0f) -#define TEST_DATA_ALL_TO_ALL(FuncNameType, DataType) \ - extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_to_all_##FuncNameType( \ - DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \ - int nem_elems_per_rank = num_elems / num_ranks; \ - int offset = nem_elems_per_rank * my_rank; \ - for (int rank = 0; rank < num_ranks; rank++) { \ - size_t rank_offset = rank * nem_elems_per_rank; \ - unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \ - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ - seed = ranqd1(seed); \ - if (i >= my_rank * nem_elems_per_rank && i < (my_rank + 1) * nem_elems_per_rank) { \ - test_buf[rank_offset + i - offset] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ - assert(result_buf[rank_offset + i - offset] == test_buf[rank_offset + i - offset]); \ - } \ - } \ - } \ +#define TEST_DATA_ALL_TO_ALL(FuncNameType, DataType) \ + extern "C" __global__ void __launch_bounds__(1024, 1) \ + test_data_all_to_all_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \ + int my_rank, int seq, int split_mask) { \ + int nem_elems_per_rank = num_elems / num_ranks; \ + int offset = nem_elems_per_rank * my_rank; \ + for (int rank = 0; rank < num_ranks; rank++) { \ + size_t rank_offset = rank * nem_elems_per_rank; \ + int seed_rank = rank / (split_mask + 1); \ + unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ + seed = ranqd1(seed); \ + if (i >= my_rank * nem_elems_per_rank && i < (my_rank + 1) * nem_elems_per_rank) { \ + test_buf[rank_offset + i - offset] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ + assert(result_buf[rank_offset + i - offset] == test_buf[rank_offset + i - offset]); \ + } \ + } \ + } \ } TEST_DATA_ALL_TO_ALL(bfloat16, __nv_bfloat16) TEST_DATA_ALL_TO_ALL(float16, __half) TEST_DATA_ALL_TO_ALL(float32, float) -TEST_DATA_ALL_TO_ALL(int32, int) \ No newline at end of file +TEST_DATA_ALL_TO_ALL(int32, int) + +// Sendrecv verification: receive from the prev group in the ring. +// fill_data seeds by group_id (rank / (split_mask + 1)); the receiver in group g expects the +// data produced by group (g - 1 + num_groups) % num_groups, so we recompute that seed here. +#define TEST_DATA_SEND_RECV(FuncNameType, DataType) \ + extern "C" __global__ void __launch_bounds__(1024, 1) \ + test_data_send_recv_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \ + int my_rank, int seq, int split_mask) { \ + int group_size = split_mask + 1; \ + int num_groups = num_ranks / group_size; \ + int my_group_id = my_rank / group_size; \ + int prev_group_id = (my_group_id - 1 + num_groups) % num_groups; \ + unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + prev_group_id + seq); \ + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \ + seed = ranqd1(seed); \ + test_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \ + assert(result_buf[i] == test_buf[i]); \ + } \ + } + +TEST_DATA_SEND_RECV(bfloat16, __nv_bfloat16) +TEST_DATA_SEND_RECV(float16, __half) +TEST_DATA_SEND_RECV(float32, float) +TEST_DATA_SEND_RECV(int32, int) diff --git a/python/test/test_fp8_accum.py b/python/test/test_fp8_accum.py index ba33c085..554e131a 100644 --- a/python/test/test_fp8_accum.py +++ b/python/test/test_fp8_accum.py @@ -167,7 +167,7 @@ else: # --------------------------------------------------------------------------- -# FP8 E4M3B15 helpers (bias=15, encode saturates to ±1.75, no NaN) +# FP8 E4M3B15 helpers (bias=15, float source saturates to ±1.875, no NaN) # Matches Triton's fp8e4b15: all 256 bit patterns are finite. # --------------------------------------------------------------------------- @@ -193,7 +193,7 @@ def float_to_e4m3b15(f32_array, chunk_size=65536): """Encode a cupy float32 array to uint8 E4M3B15 bit patterns. Same lookup-table approach as float_to_e4m3fn. - Saturates to ±1.75 (0x7e/0xfe), matching Triton's fp8e4b15. + Saturates to ±1.875 (0x7f/0xff), matching the device float32 → e4m3b15 path. """ # Build lookup table of all 128 positive E4M3B15 values (0x00..0x7F) all_bytes = cp.arange(128, dtype=cp.uint8) @@ -203,7 +203,7 @@ def float_to_e4m3b15(f32_array, chunk_size=65536): values = f32_array.astype(cp.float32) signs = cp.signbit(values).astype(cp.uint8) absval = cp.abs(values) - absval = cp.clip(absval, cp.float32(0.0), cp.float32(1.75)) + absval = cp.clip(absval, cp.float32(0.0), cp.float32(1.875)) result = cp.zeros(absval.shape, dtype=cp.uint8) n = absval.size @@ -442,8 +442,8 @@ def test_fp8_e4m3b15_accum(mpi_group: MpiGroup, algo_name: str, size: int): bits_r = cp.asarray(rng_r.randint(0, 256, (size,)).astype(np.uint8)) ref_f32 += e4m3b15_to_float(bits_r) - # Clamp reference to e4m3b15 representable range - ref_f32 = cp.clip(ref_f32, -1.75, 1.75) + # Clamp reference to e4m3b15 representable range (float source saturates at ±1.875) + ref_f32 = cp.clip(ref_f32, -1.875, 1.875) # Compute errors abs_err = cp.abs(result_f32 - ref_f32) diff --git a/src/core/bootstrap/bootstrap.cc b/src/core/bootstrap/bootstrap.cc index b3032e50..2f18d70c 100644 --- a/src/core/bootstrap/bootstrap.cc +++ b/src/core/bootstrap/bootstrap.cc @@ -50,6 +50,8 @@ MSCCLPP_API_CPP void Bootstrap::groupBarrier(const std::vector& ranks) { } } +MSCCLPP_API_CPP int Bootstrap::getNranksPerIpcDomain() const { return getNranksPerNode(); } + MSCCLPP_API_CPP void Bootstrap::send(const std::vector& data, int peer, int tag) { size_t size = data.size(); send((void*)&size, sizeof(size_t), peer, tag); @@ -83,6 +85,7 @@ class TcpBootstrap::Impl { int getRank(); int getNranks(); int getNranksPerNode(); + int getNranksPerIpcDomain(); void allGather(void* allData, int size); void broadcast(void* data, int size, int root); void send(void* data, int size, int peer, int tag); @@ -95,6 +98,7 @@ class TcpBootstrap::Impl { int rank_; int nRanks_; int nRanksPerNode_; + int nRanksPerIpcDomain_; bool netInitialized; std::unique_ptr listenSockRoot_; std::unique_ptr listenSock_; @@ -148,6 +152,7 @@ TcpBootstrap::Impl::Impl(int rank, int nRanks) : rank_(rank), nRanks_(nRanks), nRanksPerNode_(0), + nRanksPerIpcDomain_(0), netInitialized(false), peerCommAddresses_(nRanks, SocketAddress()), barrierArr_(nRanks, 0), @@ -451,6 +456,42 @@ int TcpBootstrap::Impl::getNranksPerNode() { return nRanksPerNode_; } +int TcpBootstrap::Impl::getNranksPerIpcDomain() { + if (nRanksPerIpcDomain_ > 0) return nRanksPerIpcDomain_; + std::vector ipcDomainHashes(nRanks_); + ipcDomainHashes[rank_] = getIpcDomainHash(); + allGather(ipcDomainHashes.data(), sizeof(uint64_t)); + + std::unordered_map ipcDomainCounts; + for (uint64_t ipcDomainHash : ipcDomainHashes) { + ipcDomainCounts[ipcDomainHash]++; + } + + const int nRanksPerIpcDomain = ipcDomainCounts[ipcDomainHashes[rank_]]; + const std::string invalidIpcDomainLayout = + "IPC domain ranks must have the same size and be arranged in consecutive rank order"; + if (nRanks_ % nRanksPerIpcDomain != 0) { + throw Error(invalidIpcDomainLayout + ": rank count is not divisible by IPC domain size", ErrorCode::InvalidUsage); + } + + for (const auto& entry : ipcDomainCounts) { + if (entry.second != nRanksPerIpcDomain) { + throw Error(invalidIpcDomainLayout + ": IPC domains have different sizes", ErrorCode::InvalidUsage); + } + } + for (int i = 0; i < nRanks_; ++i) { + const int ipcDomainFirstRank = i - i % nRanksPerIpcDomain; + if (ipcDomainHashes[i] != ipcDomainHashes[ipcDomainFirstRank]) { + throw Error(invalidIpcDomainLayout + ": ranks are not grouped by IPC domain", ErrorCode::InvalidUsage); + } + } + + INFO(MSCCLPP_INIT, "rank %d IPC domain fabric hash 0x%016llx nRanksPerIpcDomain %d", rank_, + static_cast(ipcDomainHashes[rank_]), nRanksPerIpcDomain); + nRanksPerIpcDomain_ = nRanksPerIpcDomain; + return nRanksPerIpcDomain_; +} + void TcpBootstrap::Impl::allGather(void* allData, int size) { char* data = static_cast(allData); int rank = rank_; @@ -592,6 +633,8 @@ MSCCLPP_API_CPP int TcpBootstrap::getNranks() const { return pimpl_->getNranks() MSCCLPP_API_CPP int TcpBootstrap::getNranksPerNode() const { return pimpl_->getNranksPerNode(); } +MSCCLPP_API_CPP int TcpBootstrap::getNranksPerIpcDomain() const { return pimpl_->getNranksPerIpcDomain(); } + MSCCLPP_API_CPP void TcpBootstrap::send(void* data, int size, int peer, int tag) { pimpl_->send(data, size, peer, tag); } diff --git a/src/core/core.cc b/src/core/core.cc index 2d67b988..ee630529 100644 --- a/src/core/core.cc +++ b/src/core/core.cc @@ -75,7 +75,7 @@ MSCCLPP_API_CPP bool TransportFlags::operator==(TransportFlags other) const { } MSCCLPP_API_CPP bool TransportFlags::operator!=(TransportFlags other) const { - return detail::TransportFlagsBase::operator!=(other); + return !detail::TransportFlagsBase::operator==(other); } MSCCLPP_API_CPP detail::TransportFlagsBase TransportFlags::toBitset() const { return *this; } diff --git a/src/core/executor/execution_plan.cc b/src/core/executor/execution_plan.cc index 98ec3ab6..25ae3156 100644 --- a/src/core/executor/execution_plan.cc +++ b/src/core/executor/execution_plan.cc @@ -69,6 +69,10 @@ auto getOpType = [](const std::string& str) { return mscclpp::OperationType::REDUCE_COPY_SEND_PACKETS; } else if (str == "glres") { return mscclpp::OperationType::MULTI_LOAD_REDUCE_STORE; + } else if (str == "gstore") { + return mscclpp::OperationType::MULTI_STORE; + } else if (str == "gstorepkt") { + return mscclpp::OperationType::MULTI_STORE_PKT; } else if (str == "rlxsignal") { return mscclpp::OperationType::RELAXED_SIGNAL; } else if (str == "rlxwait") { @@ -232,7 +236,7 @@ void ExecutionPlan::Impl::loadExecutionPlan(size_t inputSize, size_t outputSize, size_t constDstOffset) { std::ifstream file(this->planPath); json obj = json::parse(file); - if (this->name != obj["name"]) { + if (this->name != obj["name"].get()) { throw Error("Plan name does not match", ErrorCode::ExecutorError); } this->collective = obj["collective"]; @@ -268,7 +272,7 @@ void ExecutionPlan::Impl::lightLoadExecutionPlan(size_t inputSize, size_t output size_t constDstOffset) { std::ifstream file(this->planPath); json obj = json::parse(file); - if (this->name != obj["name"]) { + if (this->name != obj["name"].get()) { throw Error("Plan name does not match", ErrorCode::ExecutorError); } std::string protocol = obj["protocol"]; diff --git a/src/core/executor/executor.cc b/src/core/executor/executor.cc index fcecc4dd..91db3808 100644 --- a/src/core/executor/executor.cc +++ b/src/core/executor/executor.cc @@ -94,9 +94,9 @@ struct hash { namespace { auto hasIBDevices = []() { return mscclpp::getIBDeviceCount() > 0; }; -auto useIB = [](int rank1, int rank2, int nranksPerNode) { - bool inSameNode = rank1 / nranksPerNode == rank2 / nranksPerNode; - return hasIBDevices() && !inSameNode; +auto useIB = [](int rank1, int rank2, int nranksPerIpcDomain) { + bool inSameIpcDomain = rank1 / nranksPerIpcDomain == rank2 / nranksPerIpcDomain; + return hasIBDevices() && !inSameIpcDomain; }; static const mscclpp::Transport IBs[] = {mscclpp::Transport::IB0, mscclpp::Transport::IB1, mscclpp::Transport::IB2, @@ -108,7 +108,7 @@ namespace mscclpp { struct ExecutionContext { std::shared_ptr proxyService; - std::unordered_map connections; + std::vector connections; std::vector> nvlsConnections; MemoryId localMemoryIdBegin = MemoryId(0); @@ -120,8 +120,6 @@ struct ExecutionContext { // local registered memories to keep resources alive std::vector localRegisteredMemories; - std::vector> memorySemaphores; - std::vector proxySemaphores; std::vector memoryChannels; std::vector portChannels; std::vector nvlsChannels; @@ -139,6 +137,7 @@ struct ExecutionContext { struct Executor::Impl { int nranksPerNode; + int nranksPerIpcDomain; int nranks; std::shared_ptr comm; const size_t defaultScratchBufferSize = (1 << 27); @@ -149,6 +148,7 @@ struct Executor::Impl { Impl(std::shared_ptr comm, std::shared_ptr defaultScratchBuffer = nullptr) : comm(comm), defaultScratchBuffer(defaultScratchBuffer) { this->nranksPerNode = comm->bootstrap()->getNranksPerNode(); + this->nranksPerIpcDomain = comm->bootstrap()->getNranksPerIpcDomain(); this->nranks = comm->bootstrap()->getNranks(); this->proxyService = std::make_shared(); this->proxyService->startProxy(true); @@ -218,7 +218,7 @@ struct Executor::Impl { if (type == ChannelType::MEMORY) { flags |= Transport::CudaIpc; } else if (type == ChannelType::PORT) { - if (useIB(rank, info.accessRank, this->nranksPerNode)) { + if (useIB(rank, info.accessRank, this->nranksPerIpcDomain)) { flags |= IBs[rank % this->nranksPerNode]; } else flags |= Transport::CudaIpc; @@ -265,15 +265,36 @@ struct Executor::Impl { } }; - std::vector connectedPeers = plan.impl_->getConnectedPeers(); - std::vector> connectionFutures; - for (int peer : connectedPeers) { - Transport transport = - !useIB(rank, peer, this->nranksPerNode) ? Transport::CudaIpc : IBs[rank % this->nranksPerNode]; - connectionFutures.push_back(this->comm->connect(transport, peer)); + // Create one connection (unique QP) per channel entry. Each channel gets its own + // QP — no shared connections. + // Use per-peer tag counters so that matched connections between pairs of ranks use + // the same tag, regardless of the order peers appear in each rank's connected_to list. + std::unordered_map peerTagCounters; + Transport ibTransport = IBs[rank % this->nranksPerNode]; + std::vector> connFutures; + for (ChannelType channelType : {ChannelType::MEMORY, ChannelType::PORT}) { + std::vector channelInfos = plan.impl_->getChannelInfos(channelType); + for (const auto& info : channelInfos) { + for (int peer : info.connectedPeers) { + Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerIpcDomain) + ? ibTransport + : Transport::CudaIpc; + connFutures.push_back(this->comm->connect(transport, peer, peerTagCounters[peer]++)); + } + } + channelInfos = plan.impl_->getUnpairedChannelInfos(nranks, channelType); + for (const auto& info : channelInfos) { + for (int peer : info.connectedPeers) { + Transport transport = channelType == ChannelType::PORT && useIB(rank, peer, this->nranksPerIpcDomain) + ? ibTransport + : Transport::CudaIpc; + connFutures.push_back(this->comm->connect(transport, peer, peerTagCounters[peer]++)); + } + } } - for (size_t i = 0; i < connectionFutures.size(); i++) { - context.connections[connectedPeers[i]] = connectionFutures[i].get(); + + for (auto& future : connFutures) { + context.connections.push_back(future.get()); } std::vector nvlsInfos = plan.impl_->nvlsInfos.at(rank); @@ -327,10 +348,11 @@ struct Executor::Impl { std::vector> futureProxySemaphores; std::vector> memorySemaphores; std::vector proxySemaphores; + int connIdx = 0; auto processChannelInfos = [&](std::vector& channelInfos) { for (ChannelInfo& info : channelInfos) { - for (int peer : info.connectedPeers) { - auto connection = context.connections.at(peer); + for (size_t i = 0; i < info.connectedPeers.size(); i++) { + auto& connection = context.connections[connIdx++]; if (info.channelType == ChannelType::MEMORY) { futureMemorySemaphores.push_back(this->comm->buildSemaphore( connection, this->comm->remoteRankOf(connection), this->comm->tagOf(connection))); @@ -359,18 +381,15 @@ struct Executor::Impl { proxySemaphores.push_back(context.proxyService->addSemaphore(sem.get())); } - context.memorySemaphores = std::move(memorySemaphores); - context.proxySemaphores = std::move(proxySemaphores); - for (ChannelType channelType : channelTypes) { std::vector channelInfos = plan.impl_->getChannelInfos(channelType); int index = 0; for (ChannelInfo& info : channelInfos) { for (size_t i = 0; i < info.connectedPeers.size(); i++) { if (channelType == ChannelType::MEMORY) { - context.memoryChannels.emplace_back(context.memorySemaphores[index++]); + context.memoryChannels.emplace_back(memorySemaphores[index++]); } else if (channelType == ChannelType::PORT) { - context.portChannels.emplace_back(context.proxyService->basePortChannel(context.proxySemaphores[index++])); + context.portChannels.emplace_back(context.proxyService->basePortChannel(proxySemaphores[index++])); } } } diff --git a/src/core/gpu_ipc_mem.cc b/src/core/gpu_ipc_mem.cc index 0f58ed20..b885a1ac 100644 --- a/src/core/gpu_ipc_mem.cc +++ b/src/core/gpu_ipc_mem.cc @@ -81,8 +81,6 @@ bool isFabricMemHandleAvailable() { #endif // !(CUDA_NVLS_API_AVAILABLE) } -#if defined(MSCCLPP_DEVICE_HIP) - // Custom hash and equality for cudaIpcMemHandle_t struct CudaIpcMemHandleHash { size_t operator()(const cudaIpcMemHandle_t& handle) const { @@ -97,43 +95,59 @@ struct CudaIpcMemHandleEqual { } }; -static std::unordered_map +static std::unordered_map, CudaIpcMemHandleHash, CudaIpcMemHandleEqual> openCudaIpcMemHandleMap; static std::mutex openCudaIpcMemHandleMapMutex; -// Cache open ipc handles to avoid opening multiple times (ROCm may exceed system limit on vm.max_map_count). -static inline cudaError_t cudaIpcOpenMemHandleWrapper(void** addr, cudaIpcMemHandle_t ipcHandle) { +static inline cudaError_t cudaIpcOpenMemHandleWrapper(std::shared_ptr& basePtr, cudaIpcMemHandle_t ipcHandle) { +#if defined(MSCCLPP_DEVICE_HIP) + // Cache open ipc handles to avoid opening multiple times (ROCm may exceed system limit on vm.max_map_count). std::lock_guard lock(openCudaIpcMemHandleMapMutex); auto it = openCudaIpcMemHandleMap.find(ipcHandle); if (it != openCudaIpcMemHandleMap.end()) { - *addr = it->second; - return cudaSuccess; + basePtr = it->second.lock(); + if (basePtr) { + return cudaSuccess; + } + openCudaIpcMemHandleMap.erase(it); } - cudaError_t err = cudaIpcOpenMemHandle(addr, ipcHandle, cudaIpcMemLazyEnablePeerAccess); - if (err == cudaSuccess) { - openCudaIpcMemHandleMap[ipcHandle] = *addr; + + void* rawBasePtr = nullptr; + cudaError_t err = cudaIpcOpenMemHandle(&rawBasePtr, ipcHandle, cudaIpcMemLazyEnablePeerAccess); + if (err != cudaSuccess) { + return err; } - return err; -} -static inline cudaError_t cudaIpcCloseMemHandleWrapper(void* addr, cudaIpcMemHandle_t ipcHandle) { - std::lock_guard lock(openCudaIpcMemHandleMapMutex); - openCudaIpcMemHandleMap.erase(ipcHandle); - return cudaIpcCloseMemHandle(addr); -} - -#else // !defined(MSCCLPP_DEVICE_HIP) - -static inline cudaError_t cudaIpcOpenMemHandleWrapper(void** addr, cudaIpcMemHandle_t ipcHandle) { - return cudaIpcOpenMemHandle(addr, ipcHandle, cudaIpcMemLazyEnablePeerAccess); -} - -static inline cudaError_t cudaIpcCloseMemHandleWrapper(void* addr, [[maybe_unused]] cudaIpcMemHandle_t ipcHandle) { - return cudaIpcCloseMemHandle(addr); -} + basePtr = std::shared_ptr(rawBasePtr, [ipcHandle](void* ptr) { + { + std::lock_guard lock(openCudaIpcMemHandleMapMutex); + openCudaIpcMemHandleMap.erase(ipcHandle); + } + cudaError_t err = cudaIpcCloseMemHandle(ptr); + if (err != cudaSuccess) { + WARN(GPU, "Failed to close CUDA IPC handle at pointer ", ptr, ": ", cudaGetErrorString(err)); + } + }); + openCudaIpcMemHandleMap[ipcHandle] = basePtr; + return cudaSuccess; +#else // !defined(MSCCLPP_DEVICE_HIP) + void* rawBasePtr = nullptr; + cudaError_t err = cudaIpcOpenMemHandle(&rawBasePtr, ipcHandle, cudaIpcMemLazyEnablePeerAccess); + if (err != cudaSuccess) { + return err; + } + basePtr = std::shared_ptr(rawBasePtr, [](void* ptr) { + cudaError_t err = cudaIpcCloseMemHandle(ptr); + if (err != cudaSuccess) { + WARN(GPU, "Failed to close CUDA IPC handle at pointer ", ptr, ": ", cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + }); + return cudaSuccess; #endif // !defined(MSCCLPP_DEVICE_HIP) +} void GpuIpcMemHandle::deleter(GpuIpcMemHandle* handle) { if (handle) { @@ -172,7 +186,10 @@ UniqueGpuIpcMemHandle GpuIpcMemHandle::create(const CUdeviceptr ptr) { if (err == cudaSuccess) { handle->typeFlags |= GpuIpcMemHandle::Type::RuntimeIpc; } else { + // The VMM fallback below handles this failure, so do not leak it as CUDA's last error. (void)cudaGetLastError(); + WARN(GPU, "Failed to create runtime CUDA IPC handle for pointer ", (void*)basePtr, + ": error=", static_cast(err), " (", std::string(cudaGetErrorString(err)), ")"); } #if !defined(MSCCLPP_DEVICE_HIP) // Remove when HIP fully supports virtual memory management APIs @@ -345,17 +362,10 @@ std::shared_ptr GpuIpcMem::map() { if (type_ == GpuIpcMemHandle::Type::RuntimeIpc) { // RuntimeIpc: Open handle and return shared_ptr with cleanup in deleter - void* basePtr = nullptr; - MSCCLPP_CUDATHROW(cudaIpcOpenMemHandleWrapper(&basePtr, handle_.runtimeIpc.handle)); - void* dataPtr = static_cast(static_cast(basePtr) + handle_.offsetFromBase); - cudaIpcMemHandle_t ipcHandle = handle_.runtimeIpc.handle; - return std::shared_ptr(dataPtr, [self = shared_from_this(), basePtr, ipcHandle](void*) { - cudaError_t err = cudaIpcCloseMemHandleWrapper(basePtr, ipcHandle); - if (err != cudaSuccess) { - WARN(GPU, "Failed to close CUDA IPC handle at pointer ", basePtr, ": ", cudaGetErrorString(err)); - (void)cudaGetLastError(); - } - }); + std::shared_ptr basePtr; + MSCCLPP_CUDATHROW(cudaIpcOpenMemHandleWrapper(basePtr, handle_.runtimeIpc.handle)); + void* rawDataPtr = static_cast(static_cast(basePtr.get()) + handle_.offsetFromBase); + return std::shared_ptr(basePtr, rawDataPtr); } size_t pageSize = getpagesize(); diff --git a/src/core/include/execution_common.hpp b/src/core/include/execution_common.hpp index d071ce7d..6ce47d8f 100644 --- a/src/core/include/execution_common.hpp +++ b/src/core/include/execution_common.hpp @@ -66,6 +66,8 @@ enum class OperationType : uint8_t { PIPELINE, SEM_RELEASE, SEM_ACQUIRE, + MULTI_STORE, + MULTI_STORE_PKT, }; struct Channels { diff --git a/src/core/include/execution_kernel.hpp b/src/core/include/execution_kernel.hpp index cb808bc8..43ba8449 100644 --- a/src/core/include/execution_kernel.hpp +++ b/src/core/include/execution_kernel.hpp @@ -174,11 +174,11 @@ MSCCLPP_DEVICE_INLINE void handlePut(const Operation& op, void* input, void* out uint32_t dstOffset = dstOffsets[tid] + getOffset(portChannelBufferTypes_[op.outputBufferRefs[tid].id], offset); uint32_t srcOffset = srcOffsets[tid] + getOffset(op.inputBufferRefs[tid].type, offset); - if constexpr (PutWithSignal) { - portChannels_[channelIndexes[tid]].putWithSignal(dstMemoryId, dstOffset, srcMemoryId, srcOffset, size); - } else if constexpr (PutWithSignalAndFlush) { + if constexpr (PutWithSignalAndFlush) { portChannels_[channelIndexes[tid]].putWithSignalAndFlush(dstMemoryId, (uint64_t)dstOffset, srcMemoryId, - (uint64_t)srcOffsets, size); + (uint64_t)srcOffset, size); + } else if constexpr (PutWithSignal) { + portChannels_[channelIndexes[tid]].putWithSignal(dstMemoryId, dstOffset, srcMemoryId, srcOffset, size); } else { portChannels_[channelIndexes[tid]].put(dstMemoryId, dstOffset, srcMemoryId, srcOffset, size); } @@ -575,6 +575,65 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(const Operation& op, uint3 } } } + +template +MSCCLPP_DEVICE_INLINE void handleMultiStore(const Operation& op, void* input, void* output, void* scratch, + uint32_t offset, uint32_t unitSize) { + const uint32_t size = min(op.inputBufferSizes[0] - offset, unitSize); + if (size == 0) { + return; + } + + const uint32_t srcOffset = op.inputOffsets[0] + getOffset(op.inputBufferRefs[0].type, offset); + const uint32_t dstOffset = op.outputOffsets[0] + getOffset(op.nvlsOutputBufferType, offset); + char* srcBase = static_cast(getBuffer(input, output, scratch, op.inputBufferRefs[0].type)) + srcOffset; + char* dstBase = reinterpret_cast(nvlsChannels_[op.nvlsOutputIndex].mcPtr) + dstOffset; + + size_t processed = 0; + const bool alignedf32x4 = (reinterpret_cast(srcBase) % sizeof(f32x4) == 0) && + (reinterpret_cast(dstBase) % sizeof(f32x4) == 0); + if (alignedf32x4) { + const size_t nf32x4 = size / sizeof(f32x4); + f32x4* src16 = reinterpret_cast(srcBase); + f32x4* dst16 = reinterpret_cast(dstBase); + for (size_t idx = threadIdx.x; idx < nf32x4; idx += blockDim.x) { + SwitchChannelDeviceHandle::multimemStore(src16[idx], dst16 + idx); + } + processed = nf32x4 * sizeof(f32x4); + } + + // Handle remaining data + const size_t startIdx = processed / sizeof(u32x1); + const size_t endIdx = size / sizeof(u32x1); + u32x1* src4 = reinterpret_cast(srcBase); + u32x1* dst4 = reinterpret_cast(dstBase); + for (size_t idx = threadIdx.x + startIdx; idx < endIdx; idx += blockDim.x) { + SwitchChannelDeviceHandle::multimemStore(src4[idx], dst4 + idx); + } +} +#endif + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 +template +MSCCLPP_DEVICE_INLINE void handleMultiStorePkt(const Operation& op, void* input, void* output, void* scratch) { + const uint32_t srcOffset = op.inputOffsets[0]; + const uint32_t dstOffset = op.outputOffsets[0]; + const uint32_t size = op.inputBufferSizes[0]; + uint32_t nPackets = size / sizeof(PacketPayload); + + PacketType* srcPackets = + (PacketType*)((char*)getBuffer(input, output, scratch, op.inputBufferRefs[0].type) + (srcOffset << 1)); + PacketType* multiPkt = + (PacketType*)((char*)nvlsChannels_[op.nvlsOutputIndex].mcPtr + scratchOffset_ + (dstOffset << 1)); + + static_assert(sizeof(PacketType) == 16 || sizeof(PacketType) == 8, "Unsupported packet size for MULTI_STORE_PKT"); + using StoreVec = std::conditional_t; + for (size_t idx = threadIdx.x; idx < nPackets; idx += blockDim.x) { + PacketPayload data = srcPackets[idx].read(flag_); + PacketType pkt(data, flag_); + mscclpp::SwitchChannelDeviceHandle::multimemStore(*(StoreVec*)(&pkt), multiPkt + idx); + } +} #endif template @@ -708,6 +767,10 @@ MSCCLPP_DEVICE_INLINE void executeDeviceFunction(const Operation& op, T* input, #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 else if (opType == OperationType::MULTI_LOAD_REDUCE_STORE) { handleMultiLoadReduceStore(op, offset, unitSize); + } else if (opType == OperationType::MULTI_STORE) { + handleMultiStore(op, input, output, scratch, offset, unitSize); + } else if (opType == OperationType::MULTI_STORE_PKT) { + handleMultiStorePkt(op, input, output, scratch); } #endif else if (opType == OperationType::PIPELINE) { diff --git a/src/core/include/utils_internal.hpp b/src/core/include/utils_internal.hpp index c5c67e26..c6934194 100644 --- a/src/core/include/utils_internal.hpp +++ b/src/core/include/utils_internal.hpp @@ -37,6 +37,7 @@ int64_t busIdToInt64(const std::string busId); uint64_t getHash(const char* string, int n); uint64_t getHostHash(); uint64_t getPidHash(); +uint64_t getIpcDomainHash(); void getRandomData(void* buffer, size_t bytes); struct netIf { diff --git a/src/core/registered_memory.cc b/src/core/registered_memory.cc index f464de2a..c82233ce 100644 --- a/src/core/registered_memory.cc +++ b/src/core/registered_memory.cc @@ -97,11 +97,14 @@ MSCCLPP_API_CPP std::vector RegisteredMemory::serialize() const { for (auto& entry : pimpl_->transportInfos) { detail::serialize(result, entry.transport); if (entry.transport == Transport::CudaIpc) { + if (entry.gpuIpcMemHandle.typeFlags == GpuIpcMemHandle::Type::None) { + THROW(GPU, Error, ErrorCode::InternalError, "GpuIpcMemHandle type is None"); + } detail::serialize(result, entry.gpuIpcMemHandle); } else if (AllIBTransports.has(entry.transport)) { detail::serialize(result, entry.ibMrInfo); } else { - throw Error("Unknown transport", ErrorCode::InternalError); + THROW(GPU, Error, ErrorCode::InternalError, "Unknown transport"); } } return result; diff --git a/src/core/utils_internal.cc b/src/core/utils_internal.cc index 8cc55430..d48c4225 100644 --- a/src/core/utils_internal.cc +++ b/src/core/utils_internal.cc @@ -6,6 +6,10 @@ #include #include +#if defined(MSCCLPP_USE_CUDA) +#include +#endif + #include #include #include @@ -175,6 +179,83 @@ uint64_t getPidHash(void) { return *pidHash; } +#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN) +namespace { + +class NvmlState { + public: + NvmlState() : initialized_(nvmlInit_v2() == NVML_SUCCESS) {} + + NvmlState(const NvmlState&) = delete; + NvmlState& operator=(const NvmlState&) = delete; + NvmlState(NvmlState&&) = delete; + NvmlState& operator=(NvmlState&&) = delete; + + ~NvmlState() { + if (initialized_) { + (void)nvmlShutdown(); + } + } + + bool isInitialized() const { return initialized_; } + + private: + bool initialized_ = false; +}; + +template +uint64_t getFabricHash(const FabricInfo& fabricInfo) { + char hashData[NVML_GPU_FABRIC_UUID_LEN + sizeof(fabricInfo.cliqueId)]; + std::memcpy(hashData, fabricInfo.clusterUuid, NVML_GPU_FABRIC_UUID_LEN); + std::memcpy(hashData + NVML_GPU_FABRIC_UUID_LEN, &fabricInfo.cliqueId, sizeof(fabricInfo.cliqueId)); + return getHash(hashData, sizeof(hashData)); +} + +bool tryGetNvmlIpcDomainHash(uint64_t& ipcDomainHash) { + // Use the current CUDA device; callers must set the rank's device before querying. + int deviceId; + char pciBusId[] = "00000000:00:00.0"; + if (cudaGetDevice(&deviceId) != cudaSuccess || + cudaDeviceGetPCIBusId(pciBusId, sizeof(pciBusId), deviceId) != cudaSuccess) { + return false; + } + + static NvmlState nvml; + nvmlDevice_t nvmlDevice; + if (!nvml.isInitialized() || nvmlDeviceGetHandleByPciBusId_v2(pciBusId, &nvmlDevice) != NVML_SUCCESS) { + return false; + } + +#if defined(nvmlGpuFabricInfo_v2) + nvmlGpuFabricInfoV_t fabricInfo = {}; + fabricInfo.version = nvmlGpuFabricInfo_v2; + nvmlReturn_t result = nvmlDeviceGetGpuFabricInfoV(nvmlDevice, &fabricInfo); +#else + nvmlGpuFabricInfo_t fabricInfo = {}; + nvmlReturn_t result = nvmlDeviceGetGpuFabricInfo(nvmlDevice, &fabricInfo); +#endif + if (result != NVML_SUCCESS || fabricInfo.state != NVML_GPU_FABRIC_STATE_COMPLETED || + fabricInfo.status != NVML_SUCCESS) { + return false; + } + + ipcDomainHash = getFabricHash(fabricInfo); + return true; +} + +} // namespace +#endif + +uint64_t getIpcDomainHash(void) { +#if defined(MSCCLPP_USE_CUDA) && defined(NVML_GPU_FABRIC_UUID_LEN) + uint64_t ipcDomainHash; + if (tryGetNvmlIpcDomainHash(ipcDomainHash)) { + return ipcDomainHash; + } +#endif + return getHostHash(); +} + int parseStringList(const char* string, netIf* ifList, int maxList) { if (!string) return 0; diff --git a/src/ext/collectives/allgather/allgather_fullmesh.cu b/src/ext/collectives/allgather/allgather_fullmesh.cu index fb51a342..d1b4e731 100644 --- a/src/ext/collectives/allgather/allgather_fullmesh.cu +++ b/src/ext/collectives/allgather/allgather_fullmesh.cu @@ -8,6 +8,11 @@ namespace mscclpp { namespace collective { +namespace { +constexpr int kMaxBlocks = 56; +constexpr int kMaxThreadsPerBlock = 1024; +} // namespace + template __global__ void __launch_bounds__(1024, 1) allgatherFullmesh(void* buff, void* scratch, void* resultBuff, DeviceHandle* memoryChannels, @@ -116,12 +121,19 @@ CommResult AllgatherFullmesh::allgatherKernelFunc(const std::shared_ptr ct int rank = ctx->rank; const size_t nElem = inputSize / sizeof(int); std::pair numBlocksAndThreads = {nBlocks, nThreadsPerBlock}; - if (numBlocksAndThreads.first > 56) { - WARN("AllgatherFullmesh: number of blocks exceeds maximum supported blocks, which is 56"); - return mscclpp::CommResult::CommInvalidArgument; - } if (numBlocksAndThreads.first == 0 || numBlocksAndThreads.second == 0) { - numBlocksAndThreads = {56, 1024}; + numBlocksAndThreads = {kMaxBlocks, kMaxThreadsPerBlock}; + } + if (numBlocksAndThreads.first > kMaxBlocks || numBlocksAndThreads.second > kMaxThreadsPerBlock) { + WARN( + "AllgatherFullmesh: number of blocks must be no more than %d and threads per block must be no more than %d; " + "got nBlocks=%d, nThreadsPerBlock=%d", + kMaxBlocks, kMaxThreadsPerBlock, numBlocksAndThreads.first, numBlocksAndThreads.second); + return CommResult::CommInvalidArgument; + } + if (numBlocksAndThreads.second % WARP_SIZE != 0) { + WARN("AllgatherFullmesh: threads per block must be a multiple of warp size %d", WARP_SIZE); + return CommResult::CommInvalidArgument; } if ((char*)input == (char*)output + rank * inputSize) { allgatherFullmesh<<>>( @@ -142,15 +154,13 @@ CommResult AllgatherFullmesh::allgatherKernelFunc(const std::shared_ptr ct std::shared_ptr AllgatherFullmesh::initAllgatherContext(std::shared_ptr comm, const void* input, void*, size_t inputSize, DataType) { - constexpr int nChannelsPerConnection = 56; - auto ctx = std::make_shared(); ctx->rank = comm->bootstrap()->getRank(); ctx->workSize = comm->bootstrap()->getNranks(); ctx->nRanksPerNode = comm->bootstrap()->getNranksPerNode(); // setup semaphores - ctx->memorySemaphores = setupMemorySemaphores(comm, this->conns_, nChannelsPerConnection); + ctx->memorySemaphores = setupMemorySemaphores(comm, this->conns_, kMaxBlocks); // register the memory for the broadcast operation RegisteredMemory localMemory = comm->registerMemory((void*)input, inputSize, Transport::CudaIpc); @@ -159,7 +169,7 @@ std::shared_ptr AllgatherFullmesh::initAllgatherContext(std::shared_ptrmemoryChannels = - setupMemoryChannels(this->conns_, ctx->memorySemaphores, remoteMemories, localMemory, nChannelsPerConnection); + setupMemoryChannels(this->conns_, ctx->memorySemaphores, remoteMemories, localMemory, kMaxBlocks); ctx->memoryChannelDeviceHandles = setupMemoryChannelDeviceHandles(ctx->memoryChannels); // keep registered memories reference @@ -196,4 +206,4 @@ std::shared_ptr AllgatherFullmesh::build() { }); } } // namespace collective -} // namespace mscclpp \ No newline at end of file +} // namespace mscclpp diff --git a/src/ext/collectives/allgather/allgather_fullmesh_2.cu b/src/ext/collectives/allgather/allgather_fullmesh_2.cu index 9d169d68..3500c0c4 100644 --- a/src/ext/collectives/allgather/allgather_fullmesh_2.cu +++ b/src/ext/collectives/allgather/allgather_fullmesh_2.cu @@ -8,7 +8,6 @@ namespace mscclpp { namespace collective { -__device__ DeviceSyncer deviceSyncer; template __global__ void __launch_bounds__(1024, 1) allgatherFullmesh2(void* sendbuff, mscclpp::DeviceHandle* memoryChannels, @@ -17,10 +16,12 @@ __global__ void __launch_bounds__(1024, 1) const size_t tid = threadIdx.x + blockIdx.x * blockDim.x; const size_t lid = tid % WARP_SIZE; const size_t wid = tid / WARP_SIZE; - - const size_t nThread = blockDim.x * gridDim.x; - const size_t nWarp = nThread / WARP_SIZE; const size_t nPeer = nRanksPerNode - 1; + + // Round down to multiple of peer count. + const size_t nThread = (blockDim.x * gridDim.x) / WARP_SIZE / nPeer * nPeer * WARP_SIZE; + bool isWorker = tid < nThread; + const size_t nWarp = nThread / WARP_SIZE; const size_t chanOffset = nPeer * blockIdx.x; auto memChans = memoryChannels + chanOffset; @@ -30,76 +31,80 @@ __global__ void __launch_bounds__(1024, 1) } __syncthreads(); - const size_t bytesPerGPU = nelemsPerGPU * sizeof(int); - const size_t bytes = bytesPerGPU * nPeer; - size_t unitBytesPerThread; - if (bytes >= nThread * 64) { - unitBytesPerThread = 64; - } else { - unitBytesPerThread = 16; - } - const size_t unitBytesPerWarp = unitBytesPerThread * WARP_SIZE; - const size_t unitBytes = unitBytesPerWarp * nWarp; - const size_t nLoop = bytes / unitBytes; - - if (nLoop > 0) { - // First loop unrolling - const size_t peerIdx = wid % nPeer; - const size_t offset = bytesPerGPU * rank + (wid / nPeer) * unitBytesPerWarp; - if constexpr (IsOutOfPlace) { - char* dst = reinterpret_cast(memChans[peerIdx].dst_); - char* src = reinterpret_cast(memChans[peerIdx].src_); - char* buff = reinterpret_cast(sendbuff); - const size_t offsetWithinRank = (wid / nPeer) * unitBytesPerWarp; - mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, - WARP_SIZE); - mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, - WARP_SIZE); + if (isWorker) { + const size_t bytesPerGPU = nelemsPerGPU * sizeof(int); + const size_t bytes = bytesPerGPU * nPeer; + size_t unitBytesPerThread; + if (bytes >= nThread * 64) { + unitBytesPerThread = 64; } else { - memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE); + unitBytesPerThread = 16; } - } + const size_t unitBytesPerWarp = unitBytesPerThread * WARP_SIZE; + const size_t unitBytes = unitBytesPerWarp * nWarp; + const size_t nLoop = bytes / unitBytes; - for (size_t i = 1; i < nLoop; ++i) { - const size_t gWid = wid + i * nWarp; - const size_t peerIdx = gWid % nPeer; - const size_t offset = bytesPerGPU * rank + (gWid / nPeer) * unitBytesPerWarp; - if constexpr (IsOutOfPlace) { - char* dst = reinterpret_cast(memChans[peerIdx].dst_); - char* src = reinterpret_cast(memChans[peerIdx].src_); - char* buff = reinterpret_cast(sendbuff); - const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp; - mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, - WARP_SIZE); - mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, - WARP_SIZE); - } else { - memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE); - } - } - - if (bytes % unitBytes > 0) { - const size_t gWid = wid + nLoop * nWarp; - const size_t peerIdx = gWid % nPeer; - const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp; - const size_t offset = bytesPerGPU * rank + offsetWithinRank; - const size_t remainBytes = (offsetWithinRank + unitBytesPerWarp > bytesPerGPU) - ? ((bytesPerGPU > offsetWithinRank) ? (bytesPerGPU - offsetWithinRank) : 0) - : unitBytesPerWarp; - if (remainBytes > 0) { + if (nLoop > 0) { + // First loop unrolling + const size_t peerIdx = wid % nPeer; + const size_t offset = bytesPerGPU * rank + (wid / nPeer) * unitBytesPerWarp; if constexpr (IsOutOfPlace) { char* dst = reinterpret_cast(memChans[peerIdx].dst_); char* src = reinterpret_cast(memChans[peerIdx].src_); char* buff = reinterpret_cast(sendbuff); - mscclpp::copy<16, true>(src + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid, WARP_SIZE); - mscclpp::copy<16, true>(dst + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid, WARP_SIZE); + const size_t offsetWithinRank = (wid / nPeer) * unitBytesPerWarp; + mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, + WARP_SIZE); + mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, + WARP_SIZE); } else { - memChans[peerIdx].put<16, true>(offset + channelOutOffset, remainBytes, lid, WARP_SIZE); + memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE); + } + } + + for (size_t i = 1; i < nLoop; ++i) { + const size_t gWid = wid + i * nWarp; + const size_t peerIdx = gWid % nPeer; + const size_t offset = bytesPerGPU * rank + (gWid / nPeer) * unitBytesPerWarp; + if constexpr (IsOutOfPlace) { + char* dst = reinterpret_cast(memChans[peerIdx].dst_); + char* src = reinterpret_cast(memChans[peerIdx].src_); + char* buff = reinterpret_cast(sendbuff); + const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp; + mscclpp::copy<16, false>(src + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, + WARP_SIZE); + mscclpp::copy<16, false>(dst + offset + channelOutOffset, buff + offsetWithinRank, unitBytesPerWarp, lid, + WARP_SIZE); + } else { + memChans[peerIdx].put<16, false>(offset + channelOutOffset, unitBytesPerWarp, lid, WARP_SIZE); + } + } + + if (bytes % unitBytes > 0) { + const size_t gWid = wid + nLoop * nWarp; + const size_t peerIdx = gWid % nPeer; + const size_t offsetWithinRank = (gWid / nPeer) * unitBytesPerWarp; + const size_t offset = bytesPerGPU * rank + offsetWithinRank; + const size_t remainBytes = (offsetWithinRank + unitBytesPerWarp > bytesPerGPU) + ? ((bytesPerGPU > offsetWithinRank) ? (bytesPerGPU - offsetWithinRank) : 0) + : unitBytesPerWarp; + if (remainBytes > 0) { + if constexpr (IsOutOfPlace) { + char* dst = reinterpret_cast(memChans[peerIdx].dst_); + char* src = reinterpret_cast(memChans[peerIdx].src_); + char* buff = reinterpret_cast(sendbuff); + mscclpp::copy<16, true>(src + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid, + WARP_SIZE); + mscclpp::copy<16, true>(dst + offset + channelOutOffset, buff + offsetWithinRank, remainBytes, lid, + WARP_SIZE); + } else { + memChans[peerIdx].put<16, true>(offset + channelOutOffset, remainBytes, lid, WARP_SIZE); + } } } } - deviceSyncer.sync(gridDim.x); + __syncthreads(); if (threadIdx.x < nPeer) { memChans[threadIdx.x].signal(); @@ -135,6 +140,24 @@ CommResult AllgatherFullmesh2::allgatherKernelFunc(const std::shared_ptr c numBlocksAndThreads.first = 35; } } + const int nPeer = ctx->nRanksPerNode - 1; + const int nWarp = numBlocksAndThreads.first * numBlocksAndThreads.second / WARP_SIZE; + if (numBlocksAndThreads.first > nChannelsPerConnection_ || numBlocksAndThreads.first <= 0 || + numBlocksAndThreads.second <= 0) { + WARN( + "AllgatherFullmesh2: number of blocks must be a positive multiple of peer count and no more than %d, threads " + "per block must be positive; got nBlocks=%d, nThreadsPerBlock=%d, nPeers=%d", + nChannelsPerConnection_, numBlocksAndThreads.first, numBlocksAndThreads.second, nPeer); + return CommResult::CommInvalidArgument; + } + if (nWarp < nPeer) { + WARN( + "AllgatherFullmesh2: total number of warps must be no less than peer count; got nBlocks=%d, " + "nThreadsPerBlock=%d, " + "nPeers=%d", + numBlocksAndThreads.first, numBlocksAndThreads.second, nPeer); + return CommResult::CommInvalidArgument; + } size_t channelOutOffset = *static_cast(ctx->extras["channel_out_offset"].get()); if ((char*)input == (char*)output + rank * inputSize) { @@ -226,4 +249,4 @@ std::shared_ptr AllgatherFullmesh2::build() { } } // namespace collective -} // namespace mscclpp \ No newline at end of file +} // namespace mscclpp diff --git a/src/ext/collectives/allreduce/allreduce_allpair_packet.cu b/src/ext/collectives/allreduce/allreduce_allpair_packet.cu index 17bcfc33..49058f59 100644 --- a/src/ext/collectives/allreduce/allreduce_allpair_packet.cu +++ b/src/ext/collectives/allreduce/allreduce_allpair_packet.cu @@ -7,7 +7,7 @@ #include "allreduce/allreduce_allpair_packet.hpp" #include "allreduce/common.hpp" #include "collective_utils.hpp" -#include "debug.h" +#include "logger.hpp" namespace mscclpp { namespace collective { @@ -27,22 +27,30 @@ __global__ void allreduceAllPairs(T* buff, T* scratch, T* resultBuff, DeviceHand size_t scratchBaseOffset = (flag % numScratchBuff) ? (scratchBufferSize / numScratchBuff) : 0; size_t channelScratchOffset = scratchBaseOffset; - const int nBlocksPerPeer = gridDim.x / nPeers; - const int localBlockIdx = blockIdx.x % nBlocksPerPeer; - const int tid = threadIdx.x + localBlockIdx * blockDim.x; - const int peerIdx = blockIdx.x / nBlocksPerPeer; - size_t srcOffset = channelDataOffset; + const int tid = threadIdx.x + blockIdx.x * blockDim.x; size_t scratchOffset = channelScratchOffset + rank * nelems * sizeof(LL8Packet); void* scratchBuff = (void*)((char*)scratch + channelScratchOffset); uint32_t* src = (uint32_t*)((char*)buff); uint32_t* dst = (uint32_t*)((char*)resultBuff); - // step 1: write data to each peer's scratch buffer - memoryChannels[peerIdx].putPackets(scratchOffset, srcOffset, nelems * sizeof(uint32_t), tid, - blockDim.x * nBlocksPerPeer, flag); + const int warpId = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + const int nWarpsPerBlock = blockDim.x / WARP_SIZE; + // Assign one warp in every block to each peer. Each peer warp sends the + // same block-owned stripe, so nBlocks only partitions data and no longer + // needs to be grouped by nPeers. + if (warpId < nPeers) { + memoryChannels[warpId].putPackets(scratchOffset, channelDataOffset, nelems * sizeof(uint32_t), + lane + blockIdx.x * WARP_SIZE, gridDim.x * WARP_SIZE, flag); + } + // Safe for in-place allreduce: all peer warps must finish reading src for + // this block's stripe before any warp writes reduced data back to dst/src. + __syncthreads(); - // step 2: Reduce Data - for (size_t idx = threadIdx.x + blockIdx.x * blockDim.x; idx < nelems; idx += blockDim.x * gridDim.x) { + // Split the same sent stream across all warps for reduction. warpId selects + // which strided subset to reduce while lane preserves coalesced packet reads. + for (size_t idx = lane + blockIdx.x * WARP_SIZE + warpId * WARP_SIZE * gridDim.x; idx < nelems; + idx += nWarpsPerBlock * WARP_SIZE * gridDim.x) { uint32_t data = src[idx]; using AccRaw = std::conditional_t, uint32_t, mscclpp::VectorType>; @@ -55,18 +63,17 @@ __global__ void allreduceAllPairs(T* buff, T* scratch, T* resultBuff, DeviceHand } dst[idx] = mscclpp::downcastVector(acc); } - __syncthreads(); if (threadIdx.x == 0) { ((uint32_t*)flags)[blockIdx.x] = flag + 1; } - if (blockIdx.x == 0 && threadIdx.x >= gridDim.x && threadIdx.x < flagSize / sizeof(uint32_t)) { - ((uint32_t*)flags)[threadIdx.x] = flag + 1; + if (tid >= gridDim.x && tid < flagSize / sizeof(uint32_t)) { + ((uint32_t*)flags)[tid] = flag + 1; } } inline std::pair getDefaultBlockNumAndThreadNum(size_t inputSize, int worldSize) { if (inputSize < worldSize * sizeof(int)) { - return {worldSize - 1, 32}; + return {worldSize - 1, (worldSize - 1) * WARP_SIZE}; } return {(worldSize - 1) * 4, 512}; } @@ -80,11 +87,6 @@ struct AllpairAdapter { int nThreadsPerBlock = 0) { using ChannelType = DeviceHandle; const size_t nelems = inputSize / sizeof(T); - // Round nBlocks to multiple of nPeers so every block maps to a valid peer. - const int nPeers = worldSize - 1; - if (nPeers > 0) { - nBlocks = (nBlocks / nPeers) * nPeers; - } allreduceAllPairs<<>>( (T*)buff, (T*)scratch, (T*)resultBuff, (ChannelType*)memoryChannels, channelInOffset, scratchBufferSize, rank, nRanksPerNode, worldSize, nelems, numScratchBuff, flags, flagSize); @@ -110,9 +112,17 @@ CommResult AllreduceAllpairPacket::allreduceKernelFunc(const std::shared_ptrworkSize); } - // nBlocks must be at least nPeers for allpair — each block maps to one peer. + if (blockAndThreadNum.first > maxBlockNum_) { + WARN(ALGO, "Requested block number ", blockAndThreadNum.first, " exceeds the maximum supported block number ", + maxBlockNum_, "."); + return CommResult::CommInvalidArgument; + } const int nPeers = algoCtx->nRanksPerNode - 1; - if (nPeers > 0 && blockAndThreadNum.first < nPeers) { + // The kernel maps peer sends by warpId, so every peer needs a full warp. + if (blockAndThreadNum.second % WARP_SIZE != 0 || blockAndThreadNum.second / WARP_SIZE < nPeers) { + WARN(ALGO, + "Allpair packet requires at least one full warp per peer, but got nThreadsPerBlock=", blockAndThreadNum.second, + " and nPeers=", nPeers, "."); return CommResult::CommInvalidArgument; } size_t sendBytes; @@ -122,7 +132,8 @@ CommResult AllreduceAllpairPacket::allreduceKernelFunc(const std::shared_ptr(op, dtype, accumDtype); if (!allreduce) { - WARN("Unsupported operation or data type for allreduce: op=%d, dtype=%d", op, static_cast(dtype)); + WARN(ALGO, "Unsupported operation or data type for allreduce: op=", static_cast(op), + ", dtype=", static_cast(dtype)); return CommResult::CommInvalidArgument; } cudaError_t error = @@ -131,7 +142,7 @@ CommResult AllreduceAllpairPacket::allreduceKernelFunc(const std::shared_ptrworkSize, inputSize, stream, (void*)flagBuffer_, (uint32_t)flagBufferSize_, this->nSegmentsForScratchBuffer_, blockAndThreadNum.first, blockAndThreadNum.second); if (error != cudaSuccess) { - WARN("AllreducePacket failed with error: %s", cudaGetErrorString(error)); + WARN(ALGO, "AllreducePacket failed with error: ", cudaGetErrorString(error)); return CommResult::CommUnhandledCudaError; } return CommResult::CommSuccess; @@ -189,4 +200,4 @@ std::shared_ptr AllreduceAllpairPacket::build() { }); } } // namespace collective -} // namespace mscclpp \ No newline at end of file +} // namespace mscclpp diff --git a/src/ext/collectives/allreduce/allreduce_fullmesh.cu b/src/ext/collectives/allreduce/allreduce_fullmesh.cu index 24d2a31c..725bdb0d 100644 --- a/src/ext/collectives/allreduce/allreduce_fullmesh.cu +++ b/src/ext/collectives/allreduce/allreduce_fullmesh.cu @@ -10,7 +10,7 @@ namespace mscclpp { namespace collective { template -__global__ void __launch_bounds__(512, 1) +__global__ void __launch_bounds__(1024, 1) allreduceFullmesh(T* buff, T* scratch, T* resultBuff, DeviceHandle* memoryChannels, DeviceHandle* memoryOutChannels, size_t channelOutDataOffset, int rank, int nRanksPerNode, int worldSize, size_t nelems) { @@ -194,17 +194,6 @@ CommResult AllreduceFullmesh::allreduceKernelFunc( MSCCLPP_CUTHROW(cuMemGetAddressRange(&recvBasePtr, &recvBytes, (CUdeviceptr)output)); channelOutOffset = (char*)output - (char*)recvBasePtr; } - std::shared_ptr> inputChannelHandles; - if (this->memoryChannelsMap_.find(input) != this->memoryChannelsMap_.end()) { - inputChannelHandles = this->memoryChannelsMap_[input].second; - } else { - RegisteredMemory localMemory = comm_->registerMemory(const_cast(input), inputSize, Transport::CudaIpc); - std::vector channels = - setupMemoryChannels(this->conns_, this->inputScratchSemaphores_, this->remoteScratchMemories_, localMemory, - nChannelsPerConnection_); - this->memoryChannelsMap_[input] = std::make_pair(channels, setupMemoryChannelDeviceHandles(channels)); - } - inputChannelHandles = this->memoryChannelsMap_[input].second; AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { @@ -220,10 +209,12 @@ CommResult AllreduceFullmesh::allreduceKernelFunc( if (numBlocksAndThreads.first == 0 || numBlocksAndThreads.second == 0) { numBlocksAndThreads = {35, 512}; } + auto inputChannelDeviceHandles = ctx->memoryChannelDeviceHandles.get(); + auto outputChannelDeviceHandles = inputChannelDeviceHandles + ctx->memoryChannels.size() / 2; cudaError_t error = - allreduce(input, this->scratchBuffer_, output, inputChannelHandles.get(), ctx->memoryChannelDeviceHandles.get(), - nullptr, nullptr, 0, channelOutOffset, 0, ctx->rank, ctx->nRanksPerNode, ctx->workSize, inputSize, - stream, nullptr, 0, 0, numBlocksAndThreads.first, numBlocksAndThreads.second); + allreduce(input, this->scratchBuffer_, output, inputChannelDeviceHandles, outputChannelDeviceHandles, nullptr, + nullptr, 0, channelOutOffset, 0, ctx->rank, ctx->nRanksPerNode, ctx->workSize, inputSize, stream, + nullptr, 0, 0, numBlocksAndThreads.first, numBlocksAndThreads.second); if (error != cudaSuccess) { WARN("AllreduceAllconnect failed with error: %s", cudaGetErrorString(error)); return CommResult::CommUnhandledCudaError; @@ -231,20 +222,21 @@ CommResult AllreduceFullmesh::allreduceKernelFunc( return CommResult::CommSuccess; } -AlgorithmCtxKey AllreduceFullmesh::generateAllreduceContextKey(const void*, void* output, size_t, DataType, +AlgorithmCtxKey AllreduceFullmesh::generateAllreduceContextKey(const void* input, void* output, size_t size, DataType, bool symmetricMemory) { - static int tag = 0; - size_t recvBytes; - CUdeviceptr recvBasePtr; - MSCCLPP_CUTHROW(cuMemGetAddressRange(&recvBasePtr, &recvBytes, (CUdeviceptr)output)); symmetricMemory_ = symmetricMemory; if (!symmetricMemory_) { - return AlgorithmCtxKey{nullptr, (void*)recvBasePtr, 0, recvBytes, tag++}; + return AlgorithmCtxKey{const_cast(input), output, size, size, 0}; } - return AlgorithmCtxKey{nullptr, (void*)recvBasePtr, 0, recvBytes, 0}; + + size_t sendBytes, recvBytes; + CUdeviceptr sendBasePtr, recvBasePtr; + MSCCLPP_CUTHROW(cuMemGetAddressRange(&sendBasePtr, &sendBytes, (CUdeviceptr)input)); + MSCCLPP_CUTHROW(cuMemGetAddressRange(&recvBasePtr, &recvBytes, (CUdeviceptr)output)); + return AlgorithmCtxKey{(void*)sendBasePtr, (void*)recvBasePtr, sendBytes, recvBytes, 0}; } -std::shared_ptr AllreduceFullmesh::initAllreduceContext(std::shared_ptr comm, const void*, +std::shared_ptr AllreduceFullmesh::initAllreduceContext(std::shared_ptr comm, const void* input, void* output, size_t size, DataType) { auto ctx = std::make_shared(); ctx->rank = comm->bootstrap()->getRank(); @@ -263,8 +255,19 @@ std::shared_ptr AllreduceFullmesh::initAllreduceContext(std::shared_ptrregisterMemory((void*)recvBasePtr, recvBytes, Transport::CudaIpc); ctx->registeredMemories = setupRemoteMemories(comm, ctx->rank, localMemory); - ctx->memoryChannels = setupMemoryChannels(this->conns_, ctx->memorySemaphores, ctx->registeredMemories, localMemory, - nChannelsPerConnection_); + + RegisteredMemory inputMemory = comm->registerMemory(const_cast(input), size, TransportFlags()); + std::vector inputMemoryChannels = setupMemoryChannels( + this->conns_, this->inputScratchSemaphores_, this->remoteScratchMemories_, inputMemory, nChannelsPerConnection_); + std::vector outputMemoryChannels = setupMemoryChannels( + this->conns_, this->outputSemaphores_, ctx->registeredMemories, localMemory, nChannelsPerConnection_); + ctx->memoryChannels.reserve(inputMemoryChannels.size() + outputMemoryChannels.size()); + for (auto& channel : inputMemoryChannels) { + ctx->memoryChannels.emplace_back(std::move(channel)); + } + for (auto& channel : outputMemoryChannels) { + ctx->memoryChannels.emplace_back(std::move(channel)); + } ctx->memoryChannelDeviceHandles = setupMemoryChannelDeviceHandles(ctx->memoryChannels); return ctx; } diff --git a/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu b/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu index 2d71cd63..95ec384b 100644 --- a/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu +++ b/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu @@ -13,7 +13,7 @@ namespace collective { __device__ DeviceSemaphore deviceSemaphore[NUM_SEMAPHORES]; -template +template __global__ void __launch_bounds__(1024, 1) allreduceNvlsBlockPipeline([[maybe_unused]] const void* src, [[maybe_unused]] void* scratch, [[maybe_unused]] void* dst, @@ -105,7 +105,7 @@ __global__ void __launch_bounds__(1024, 1) deviceSemaphore[oriBid].acquire(); } __syncthreads(); - handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, reduceIterSize, tid, blockDim.x); + handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, reduceIterSize, tid, blockDim.x); __syncthreads(); if (tid == 0) { deviceSemaphore[nBlocksForCopy + bidForReduce * copyReduceRatio + i].release(); @@ -158,24 +158,19 @@ struct NvlsBlockPipelineAdapter { } else if constexpr (std::is_same_v) { // fp8_e4m3b15 is a software-only type with no hardware NVLS support. return cudaErrorNotSupported; - } else -#if defined(__CUDA_ARCH__) // Skip the __CUDA_ARCH__ < 1000 since FP8 has not been supported for NVLS - if constexpr (std::is_same_v || std::is_same_v) { - return cudaErrorNotSupported; - } else -#endif - { - using ChannelType = DeviceHandle; - allreduceNvlsBlockPipeline - <<>>(input, scratch, output, (ChannelType*)memoryChannels, - nvlsChannels, inputSize, scratchBufferSize, rank, nRanksPerNode); - return cudaGetLastError(); - } + } else { + using ChannelType = DeviceHandle; + allreduceNvlsBlockPipeline + <<>>(input, scratch, output, (ChannelType*)memoryChannels, nvlsChannels, + inputSize, scratchBufferSize, rank, nRanksPerNode); + return cudaGetLastError(); + } } }; void AllreduceNvlsBlockPipeline::initialize(std::shared_ptr comm) { nSwitchChannels_ = 8; + fp8NvlsSupported_ = isFp8NvlsSupported(); int nBaseChannels = 64; this->conns_ = setupConnections(comm); // setup semaphores @@ -187,12 +182,15 @@ void AllreduceNvlsBlockPipeline::initialize(std::shared_ptr comm) this->nvlsConnections_ = setupNvlsConnections(comm, nvlsBufferSize_, nSwitchChannels_); } -CommResult AllreduceNvlsBlockPipeline::allreduceKernelFunc(const std::shared_ptr ctx_void, const void* input, - void* output, size_t inputSize, DataType dtype, ReduceOp op, - cudaStream_t stream, int nBlocks, int nThreadsPerBlock, - const std::unordered_map& extras, - DataType accumDtype) { +CommResult AllreduceNvlsBlockPipeline::allreduceKernelFunc( + const std::shared_ptr ctx_void, const void* input, void* output, size_t inputSize, DataType dtype, + ReduceOp op, cudaStream_t stream, int nBlocks, int nThreadsPerBlock, + [[maybe_unused]] const std::unordered_map& extras, DataType accumDtype) { auto ctx = std::static_pointer_cast(ctx_void); + if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) { + WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction."); + return CommResult::CommInvalidArgument; + } AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast(dtype)); diff --git a/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu b/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu index 3bb054da..385aebd5 100644 --- a/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu +++ b/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu @@ -11,7 +11,7 @@ namespace mscclpp { namespace collective { -template +template __global__ void __launch_bounds__(1024, 1) allreduceNvlsWarpPipeline([[maybe_unused]] const void* src, [[maybe_unused]] void* scratch, [[maybe_unused]] void* dst, @@ -85,7 +85,8 @@ __global__ void __launch_bounds__(1024, 1) asm volatile("bar.sync %0, %1;" ::"r"(1), "r"((NCOPY_WARPS + NREDUCE_WARPS) * WARP_SIZE) : "memory"); T* mcBuff = (T*)multicastPtr->mcPtr; size_t offset = blockScratchOffset + (it * copyPerIter) % scratchSizePerBlock; - handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, iterSize, tidInReduce, NREDUCE_WARPS * WARP_SIZE); + handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, iterSize, tidInReduce, + NREDUCE_WARPS * WARP_SIZE); asm volatile("bar.sync %0, %1;" ::"r"(2), "r"((NRECV_COPY_WARPS + NREDUCE_WARPS) * WARP_SIZE) : "memory"); } if (warpId >= startRecvCopyWid && warpId < endRecvCopyWid) { @@ -121,24 +122,19 @@ struct NvlsWarpPipelineAdapter { } else if constexpr (std::is_same_v) { // fp8_e4m3b15 is a software-only type with no hardware NVLS support. return cudaErrorNotSupported; - } else -#if defined(__CUDA_ARCH__) // Skip the __CUDA_ARCH__ < 1000 since FP8 has not been supported for NVLS - if constexpr (std::is_same_v || std::is_same_v) { - return cudaErrorNotSupported; - } else -#endif - { - using ChannelType = DeviceHandle; - allreduceNvlsWarpPipeline - <<>>(input, scratch, output, (ChannelType*)memoryChannels, - nvlsChannels, inputSize, scratchBufferSize, rank, nRanksPerNode); - return cudaGetLastError(); - } + } else { + using ChannelType = DeviceHandle; + allreduceNvlsWarpPipeline + <<>>(input, scratch, output, (ChannelType*)memoryChannels, nvlsChannels, + inputSize, scratchBufferSize, rank, nRanksPerNode); + return cudaGetLastError(); + } } }; void AllreduceNvlsWarpPipeline::initialize(std::shared_ptr comm) { nSwitchChannels_ = 8; + fp8NvlsSupported_ = isFp8NvlsSupported(); int nBaseChannels = 64; this->conns_ = setupConnections(comm); // setup semaphores @@ -155,6 +151,10 @@ CommResult AllreduceNvlsWarpPipeline::allreduceKernelFunc( ReduceOp op, cudaStream_t stream, int nBlocks, int nThreadsPerBlock, [[maybe_unused]] const std::unordered_map& extras, DataType accumDtype) { auto ctx = std::static_pointer_cast(ctx_void); + if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) { + WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction."); + return CommResult::CommInvalidArgument; + } AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast(dtype)); diff --git a/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu b/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu index e7f2028f..48ac67ff 100644 --- a/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu +++ b/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu @@ -13,7 +13,7 @@ namespace collective { constexpr int MAX_NBLOCKS = 32; -template +template __global__ void __launch_bounds__(1024, 1) allreduceNvls([[maybe_unused]] mscclpp::DeviceHandle* memoryChannels, [[maybe_unused]] mscclpp::DeviceHandle* multicast, @@ -56,8 +56,8 @@ __global__ void __launch_bounds__(1024, 1) T* src = (T*)multicastPtr->mcPtr; T* dst = (T*)multicastOutPtr->mcPtr; if (curBlockSize > 0) { - handleMultiLoadReduceStore(src, dst, blockOffset + channelInOffset, blockOffset + channelOutOffset, curBlockSize, - threadIdx.x, blockDim.x); + handleMultiLoadReduceStore(src, dst, blockOffset + channelInOffset, blockOffset + channelOutOffset, + curBlockSize, threadIdx.x, blockDim.x); } __syncthreads(); if (threadIdx.x < nPeers) { @@ -80,17 +80,11 @@ struct NvlsAdapter { } else if constexpr (std::is_same_v) { // fp8_e4m3b15 is a software-only type with no hardware NVLS support. return cudaErrorNotSupported; - } else -#if (!defined(__CUDA_ARCH_SPECIFIC__) && !defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) || (__CUDA_ARCH__ < 1000) - if constexpr (std::is_same_v || std::is_same_v) { - return cudaErrorNotSupported; - } else -#endif - { + } else { using ChannelType = DeviceHandle; - allreduceNvls<<>>((ChannelType*)memoryChannels, nvlsChannels, - nvlsOutChannels, channelInOffset, channelOutOffset, - inputSize, rank, nRanksPerNode); + allreduceNvls + <<>>((ChannelType*)memoryChannels, nvlsChannels, nvlsOutChannels, + channelInOffset, channelOutOffset, inputSize, rank, nRanksPerNode); return cudaGetLastError(); } } @@ -102,6 +96,7 @@ void AllreduceNvls::initialize(std::shared_ptr comm) { cudaDeviceProp deviceProp; MSCCLPP_CUDATHROW(cudaGetDeviceProperties(&deviceProp, device)); computeCapabilityMajor_ = deviceProp.major; + fp8NvlsSupported_ = isFp8NvlsSupported(); nSwitchChannels_ = 32; this->conns_ = setupConnections(comm); // setup semaphores @@ -124,6 +119,10 @@ CommResult AllreduceNvls::allreduceKernelFunc(const std::shared_ptr ctx_vo return CommResult::CommInvalidArgument; } auto ctx = std::static_pointer_cast(ctx_void); + if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) { + WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction."); + return CommResult::CommInvalidArgument; + } AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast(dtype)); diff --git a/src/ext/collectives/allreduce/allreduce_packet.cu b/src/ext/collectives/allreduce/allreduce_packet.cu index 6199f192..414c2b1f 100644 --- a/src/ext/collectives/allreduce/allreduce_packet.cu +++ b/src/ext/collectives/allreduce/allreduce_packet.cu @@ -235,6 +235,18 @@ CommResult AllreducePacket::allreduceKernelFunc(const std::shared_ptr ctx_ if (blockAndThreadNum.first == 0 || blockAndThreadNum.second == 0) { blockAndThreadNum = getDefaultBlockNumAndThreadNum(inputSize, ctx->workSize, ctx->nRanksPerNode, dtype); } + if (blockAndThreadNum.first > maxBlockNum_) { + WARN(ALGO, "Requested block number ", blockAndThreadNum.first, " exceeds the maximum supported block number ", + maxBlockNum_, "."); + return CommResult::CommInvalidArgument; + } + const int nPeers = ctx->nRanksPerNode - 1; + if (blockAndThreadNum.first < nPeers) { + WARN(ALGO, + "AllreducePacket requires block number to be at least peer count, but got nBlocks=", blockAndThreadNum.first, + " and nPeers=", nPeers, "."); + return CommResult::CommInvalidArgument; + } size_t sendBytes; CUdeviceptr sendBasePtr; diff --git a/src/ext/collectives/collective_utils.cc b/src/ext/collectives/collective_utils.cu similarity index 80% rename from src/ext/collectives/collective_utils.cc rename to src/ext/collectives/collective_utils.cu index 016c4a5c..2868c979 100644 --- a/src/ext/collectives/collective_utils.cc +++ b/src/ext/collectives/collective_utils.cu @@ -1,16 +1,93 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -#include "collective_utils.hpp" - #include #include #include +#include #include #include +#include "collective_utils.hpp" + namespace mscclpp { namespace collective { + +namespace { + +#if !defined(MSCCLPP_DEVICE_HIP) +__global__ void fp8NvlsSupportProbeKernel(int* supported) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ + (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) + *supported = 1; +#else + *supported = 0; +#endif +} + +bool detectFp8NvlsSupport() { + AvoidCudaGraphCaptureGuard cgcGuard; + auto supportedDevice = mscclpp::detail::gpuCallocUnique(); + int supportedHost = 0; + auto stream = gpuStreamPool()->getStream(); + + fp8NvlsSupportProbeKernel<<<1, 1, 0, stream>>>(supportedDevice.get()); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + return false; + } + + MSCCLPP_CUDATHROW( + cudaMemcpyAsync(&supportedHost, supportedDevice.get(), sizeof(supportedHost), cudaMemcpyDeviceToHost, stream)); + err = cudaStreamSynchronize(stream); + if (err != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + return supportedHost != 0; +} +#endif + +} // namespace + +bool isFp8DataType(DataType dtype) { + return dtype == DataType::FLOAT8_E4M3FN || dtype == DataType::FLOAT8_E4M3FNUZ || dtype == DataType::FLOAT8_E5M2 || + dtype == DataType::FLOAT8_E5M2FNUZ || dtype == DataType::FLOAT8_E4M3B15; +} + +bool isNativeFp8DataType(DataType dtype) { +#if defined(__FP8_TYPES_EXIST__) +#if defined(__FP8_E4M3_IS_FNUZ__) + if (dtype == DataType::FLOAT8_E4M3FNUZ) { + return true; + } +#else + if (dtype == DataType::FLOAT8_E4M3FN) { + return true; + } +#endif +#if defined(__FP8_E5M2_IS_FNUZ__) + if (dtype == DataType::FLOAT8_E5M2FNUZ) { + return true; + } +#else + if (dtype == DataType::FLOAT8_E5M2) { + return true; + } +#endif +#endif + return false; +} + +bool isFp8NvlsSupported() { +#if defined(MSCCLPP_DEVICE_HIP) + return false; +#else + static const bool supported = detectFp8NvlsSupport(); + return supported; +#endif +} + std::vector setupRemoteMemories(std::shared_ptr comm, int rank, mscclpp::RegisteredMemory localMemory) { std::vector remoteMemories; diff --git a/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp b/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp index 362308b2..64f5ec54 100644 --- a/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_allpair_packet.hpp @@ -29,7 +29,7 @@ class AllreduceAllpairPacket : public AlgorithmBuilder { void* scratchBuffer_; size_t scratchBufferSize_; const int nSegmentsForScratchBuffer_ = 2; - const int maxBlockNum_ = 28; + const int maxBlockNum_ = 64; std::vector conns_; std::vector> memorySemaphores_; std::vector registeredMemories_; diff --git a/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp b/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp index a54352b3..e0c63a3d 100644 --- a/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_fullmesh.hpp @@ -30,8 +30,6 @@ class AllreduceFullmesh : public mscclpp::AlgorithmBuilder { std::vector> inputScratchSemaphores_; std::vector remoteScratchMemories_; RegisteredMemory localScratchMemory_; - std::unordered_map, std::shared_ptr>>> - memoryChannelsMap_; bool symmetricMemory_ = false; }; } // namespace collective diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp index 81b74add..b408c64c 100644 --- a/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp @@ -33,6 +33,7 @@ class AllreduceNvlsBlockPipeline : public AlgorithmBuilder { std::vector baseChannels_; std::vector conns_; std::vector> nvlsConnections_; + bool fp8NvlsSupported_{false}; }; } // namespace collective } // namespace mscclpp diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp index 8f02a873..2ce3a4fb 100644 --- a/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp @@ -33,6 +33,7 @@ class AllreduceNvlsWarpPipeline : public AlgorithmBuilder { std::vector baseChannels_; std::vector conns_; std::vector> nvlsConnections_; + bool fp8NvlsSupported_{false}; }; } // namespace collective } // namespace mscclpp diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp index d53ea180..edc6bceb 100644 --- a/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp @@ -36,6 +36,7 @@ class AllreduceNvls : public AlgorithmBuilder { std::vector> nvlsConnections_; std::vector> nvlsOutConnections_; int computeCapabilityMajor_{0}; + bool fp8NvlsSupported_{false}; }; } // namespace collective diff --git a/src/ext/collectives/include/allreduce/common.hpp b/src/ext/collectives/include/allreduce/common.hpp index 93b18e26..6f9a3d4c 100644 --- a/src/ext/collectives/include/allreduce/common.hpp +++ b/src/ext/collectives/include/allreduce/common.hpp @@ -36,9 +36,21 @@ MSCCLPP_DEVICE_INLINE constexpr std::size_t calcVectorSize() { } } -template +template MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t srcOffset, size_t dstOffset, size_t size, int tid, int nThreads) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ + (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) + constexpr bool fp8NvlsArchSupported = true; +#else + constexpr bool fp8NvlsArchSupported = false; +#endif + constexpr bool nativeFp8Type = std::is_same_v || std::is_same_v; + if constexpr (nativeFp8Type && !fp8NvlsArchSupported) { + MSCCLPP_ASSERT_DEVICE(false, "FP8 NVLS multimem reduction is not supported on this architecture"); + return; + } + // nvls can only handle 4 bytes alignment MSCCLPP_ASSERT_DEVICE(size % 4 == 0, "size must be 4 bytes aligned"); constexpr size_t nElem = calcVectorSize(); @@ -54,7 +66,7 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t src vectorType* src4 = (vectorType*)src; vectorType* dst4 = (vectorType*)dst; for (size_t idx = tid; idx < nVec; idx += nThreads) { - auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce(src4 + srcOffset4 + idx); + auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce(src4 + srcOffset4 + idx); mscclpp::SwitchChannelDeviceHandle::multimemStore(val, dst4 + dstOffset4 + idx); } // handle rest of data @@ -64,7 +76,8 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t src const size_t startIdx = (srcOffset + processed) / sizeof(restVectorType); const size_t endIdx = (srcOffset + size) / sizeof(restVectorType); for (size_t idx = tid + startIdx; idx < endIdx; idx += nThreads) { - auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce((restVectorType*)src + idx); + auto val = + mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce((restVectorType*)src + idx); mscclpp::SwitchChannelDeviceHandle::multimemStore(val, (restVectorType*)dst + idx); } } diff --git a/src/ext/collectives/include/collective_utils.hpp b/src/ext/collectives/include/collective_utils.hpp index f705a9d1..b7f86473 100644 --- a/src/ext/collectives/include/collective_utils.hpp +++ b/src/ext/collectives/include/collective_utils.hpp @@ -32,6 +32,10 @@ constexpr int MAX_NRANKS_PER_NODE = 8; constexpr int SCRATCH_SIZE = 2 * 1024 * 1024 * 70; // double buffer * 35 thread-blocks * 8 ranks * 256KB = 70MB +bool isFp8DataType(DataType dtype); +bool isNativeFp8DataType(DataType dtype); +bool isFp8NvlsSupported(); + std::vector setupRemoteMemories(std::shared_ptr comm, int rank, RegisteredMemory localMemory); diff --git a/src/ext/nccl/CMakeLists.txt b/src/ext/nccl/CMakeLists.txt index 9767e66f..463b7550 100644 --- a/src/ext/nccl/CMakeLists.txt +++ b/src/ext/nccl/CMakeLists.txt @@ -13,6 +13,7 @@ target_include_directories(mscclpp_nccl PRIVATE include ${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/src/core/include + ${PROJECT_SOURCE_DIR}/src/ext/collectives/include ${GPU_INCLUDE_DIRS} ) target_link_libraries(mscclpp_nccl PUBLIC mscclpp mscclpp_collectives) diff --git a/src/ext/nccl/algorithm_selector.cc b/src/ext/nccl/algorithm_selector.cc index c94aab34..1ccac65d 100644 --- a/src/ext/nccl/algorithm_selector.cc +++ b/src/ext/nccl/algorithm_selector.cc @@ -6,6 +6,7 @@ #include #include +#include "collective_utils.hpp" #include "debug.h" namespace mscclpp { @@ -20,24 +21,15 @@ static bool isNvlsSupportedForDataType(const AlgorithmSelectorConfig& config, Da return false; } - const bool isFp8 = dtype == DataType::FLOAT8_E4M3FN || dtype == DataType::FLOAT8_E4M3FNUZ || - dtype == DataType::FLOAT8_E5M2 || dtype == DataType::FLOAT8_E5M2FNUZ; - - if (!isFp8) { + if (!collective::isFp8DataType(dtype)) { return nvlsSupported; } - // FP8 handling #if !defined(__HIP_PLATFORM_AMD__) - // NVLS does not support FP8 on devices with compute capability < 10 - if (config.computeCapability.first < 10) { + if (!collective::isNativeFp8DataType(dtype)) { return false; } -#if (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) - return true; -#else - return false; -#endif + return nvlsSupported && config.fp8NvlsSupported; #else return nvlsSupported; #endif diff --git a/src/ext/nccl/algorithm_selector.hpp b/src/ext/nccl/algorithm_selector.hpp index c8705f8b..2048ea05 100644 --- a/src/ext/nccl/algorithm_selector.hpp +++ b/src/ext/nccl/algorithm_selector.hpp @@ -16,6 +16,7 @@ namespace nccl { struct AlgorithmSelectorConfig { bool symmetricMemory; bool nvlsSupported; + bool fp8NvlsSupported; bool isCuMemMapAllocated; bool inCaptureMode; std::pair computeCapability; diff --git a/src/ext/nccl/nccl.cc b/src/ext/nccl/nccl.cc index 2d6c5f9d..8fcc1bb1 100644 --- a/src/ext/nccl/nccl.cc +++ b/src/ext/nccl/nccl.cc @@ -20,6 +20,7 @@ #include #include "algorithm_selector.hpp" +#include "collective_utils.hpp" #include "datatype_conversion.hpp" static constexpr auto MSCCLPP_NCCL = mscclpp::LogSubsys::NCCL; @@ -239,6 +240,8 @@ static std::shared_ptr algoSelector( static const bool isNvlsSupported = mscclpp::isNvlsSupported(); static const std::pair deviceComputeCapability = getDeviceComputeCapability(); static const bool ncclSymmetricMemory = mscclpp::env()->ncclSymmetricMemory; + const bool fp8NvlsSupported = + mscclpp::collective::isNativeFp8DataType(request.dtype) ? mscclpp::collective::isFp8NvlsSupported() : false; const bool isCuMemMapAllocated = mscclpp::isCuMemMapAllocated(const_cast(request.inputBuffer)) && mscclpp::isCuMemMapAllocated(request.outputBuffer); @@ -249,6 +252,7 @@ static std::shared_ptr algoSelector( mscclpp::nccl::AlgorithmSelectorConfig config{.symmetricMemory = ncclSymmetricMemory, .nvlsSupported = isNvlsSupported, + .fp8NvlsSupported = fp8NvlsSupported, .isCuMemMapAllocated = isCuMemMapAllocated, .inCaptureMode = inCaptureMode, .computeCapability = deviceComputeCapability, diff --git a/test/deploy/deploy.sh b/test/deploy/deploy.sh index 6358787b..02fe4fd2 100644 --- a/test/deploy/deploy.sh +++ b/test/deploy/deploy.sh @@ -1,17 +1,34 @@ +#!/bin/bash +# deploy.sh — Provisions remote hosts, copies sources, and launches Docker containers +# for mscclpp CI/CD test environments. +# +# Usage: deploy.sh [ib_environment] [platform] [container_name] [sglang_image] +# test_name : Test suite to deploy (e.g. single-node-test, nccltest-single-node) +# ib_environment : Enable InfiniBand networking (default: true) +# platform : Target GPU platform — "cuda" or "rocm" (default: cuda) +# container_name : Docker container name (default: mscclpp-test) +# sglang_image : Docker image used for the SGLang test container +# (default: lmsysorg/sglang:latest). Only used when +# container_name is "sglang-mscclpp-test". + set -ex TEST_NAME=$1 IB_ENVIRONMENT="${2:-true}" PLATFORM="${3:-cuda}" +CONTAINER_NAME="${4:-mscclpp-test}" +SGLANG_IMAGE="${5:-lmsysorg/sglang:latest}" KeyFilePath=${SSHKEYFILE_SECUREFILEPATH} ROOT_DIR="${SYSTEM_DEFAULTWORKINGDIRECTORY}/" DST_DIR="/tmp/mscclpp" + if [ "${TEST_NAME}" == "nccltest-single-node" ] || [ "${TEST_NAME}" == "single-node-test" ]; then HOSTFILE="${SYSTEM_DEFAULTWORKINGDIRECTORY}/test/deploy/hostfile_ci" else HOSTFILE="${SYSTEM_DEFAULTWORKINGDIRECTORY}/test/deploy/hostfile" fi + SSH_OPTION="StrictHostKeyChecking=no" chmod 400 ${KeyFilePath} @@ -26,8 +43,8 @@ while true; do echo "Waiting for sshd to start..." sleep 5 done - set -e + parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION "sudo rm -rf ${DST_DIR}" tar czf /tmp/mscclpp.tar.gz -C ${ROOT_DIR} . parallel-scp -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION /tmp/mscclpp.tar.gz /tmp/mscclpp.tar.gz @@ -57,25 +74,38 @@ if [ "${PLATFORM}" == "cuda" ]; then fi" fi -# force to pull the latest image -parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ - "sudo docker pull ${CONTAINERIMAGE}" - -LAUNCH_OPTION="--gpus=all" -if [ "${PLATFORM}" == "rocm" ]; then - LAUNCH_OPTION="--device=/dev/kfd --device=/dev/dri --group-add=video" -fi -if [ "${IB_ENVIRONMENT}" == "true" ]; then +if [ "${CONTAINER_NAME}" == "sglang-mscclpp-test" ]; then + # force to pull the latest image parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ - "sudo docker run --rm -itd --privileged --net=host --ipc=host ${LAUNCH_OPTION} \ - -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=mscclpp-test \ - --entrypoint /bin/bash ${CONTAINERIMAGE}" + "sudo docker pull ${SGLANG_IMAGE}" + + parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ + "sudo docker run --rm -itd --name=${CONTAINER_NAME} --privileged --net=host --ipc=host --gpus=all -w /root -v ${DST_DIR}:/root/mscclpp --entrypoint /bin/bash ${SGLANG_IMAGE}" else + # force to pull the latest image parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ - "sudo docker run --rm -itd --net=host --ipc=host ${LAUNCH_OPTION} --cap-add=SYS_ADMIN --security-opt seccomp=unconfined \ - -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=mscclpp-test \ - --entrypoint /bin/bash ${CONTAINERIMAGE}" -fi -parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ - "sudo docker exec -t --user root mscclpp-test bash '/root/mscclpp/test/deploy/setup.sh' ${PLATFORM}" + "sudo docker pull ${CONTAINERIMAGE}" + # Set GPU passthrough flags based on platform + LAUNCH_OPTION="--gpus=all" + if [ "${PLATFORM}" == "rocm" ]; then + LAUNCH_OPTION="--device=/dev/kfd --device=/dev/dri --group-add=video" + fi + + if [ "${IB_ENVIRONMENT}" == "true" ]; then + # InfiniBand: use --privileged for RDMA device access + parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ + "sudo docker run --rm -itd --privileged --net=host --ipc=host ${LAUNCH_OPTION} \ + -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=${CONTAINER_NAME} \ + --entrypoint /bin/bash ${CONTAINERIMAGE}" + else + # Non-IB: grant SYS_ADMIN and disable seccomp instead of full --privileged + parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ + "sudo docker run --rm -itd --net=host --ipc=host ${LAUNCH_OPTION} --cap-add=SYS_ADMIN --security-opt seccomp=unconfined \ + -w /root -v ${DST_DIR}:/root/mscclpp -v /opt/microsoft:/opt/microsoft --ulimit memlock=-1:-1 --name=${CONTAINER_NAME} \ + --entrypoint /bin/bash ${CONTAINERIMAGE}" + fi +fi + +parallel-ssh -i -t 0 -h ${HOSTFILE} -x "-i ${KeyFilePath}" -O $SSH_OPTION \ + "sudo docker exec -t --user root ${CONTAINER_NAME} bash '/root/mscclpp/test/deploy/setup.sh' ${PLATFORM}" diff --git a/test/deploy/run-remote.sh b/test/deploy/run-remote.sh index 2468243e..9607664f 100755 --- a/test/deploy/run-remote.sh +++ b/test/deploy/run-remote.sh @@ -11,6 +11,7 @@ # --hostfile Override hostfile path (default: test/deploy/hostfile_ci) # --host Run command on a single host (uses parallel-ssh -H) # --user SSH user when using --host or custom hostfile +# --container Docker container name to exec into (default: mscclpp-test) set -e @@ -23,9 +24,10 @@ USE_DOCKER=true USE_LOG=true TARGET_HOST="" REMOTE_USER="" +CONTAINER_NAME="mscclpp-test" usage() { - echo "Usage: $0 [--no-docker] [--no-log] [--hostfile ] [--host ] [--user ] < " >&2 + echo "Usage: $0 [--no-docker] [--no-log] [--hostfile ] [--host ] [--user ] [--container ] < " >&2 } require_value() { @@ -56,6 +58,11 @@ while [[ "$1" == --* ]]; do REMOTE_USER="$2" shift 2 ;; + --container) + require_value "--container" "${2-}" + CONTAINER_NAME="$2" + shift 2 + ;; *) echo "Unknown option: $1" >&2; exit 1 ;; esac done @@ -103,7 +110,7 @@ if $USE_DOCKER; then INNER+=" rm -f \\\"\\\$TMP\\\"" parallel-ssh -i "${PSSH_COMMON[@]}" \ - "sudo docker exec mscclpp-test bash -c \"${INNER}\"" + "sudo docker exec ${CONTAINER_NAME} bash -c \"${INNER}\"" else parallel-ssh -i "${PSSH_COMMON[@]}" \ "set -euxo pipefail; CMD_B64='${CMD_B64}'; TMP=\$(mktemp); printf '%s' \"\$CMD_B64\" | base64 -d > \"\$TMP\"; bash -euxo pipefail \"\$TMP\"; rm -f \"\$TMP\"" diff --git a/test/deploy/setup.sh b/test/deploy/setup.sh index 2a88a310..0cb6f4b8 100644 --- a/test/deploy/setup.sh +++ b/test/deploy/setup.sh @@ -61,14 +61,14 @@ if [ -f "${PIP_CMAKE_ARGS_FILE}" ]; then fi cd /root/mscclpp -if [[ "${CUDA_VERSION}" == *"11."* ]]; then - pip3 install ".[cuda11,benchmark,test]" -elif [[ "${CUDA_VERSION}" == *"12."* ]]; then +if [[ "${CUDA_VERSION}" == *"12."* ]]; then pip3 install ".[cuda12,benchmark,test]" elif [[ "${CUDA_VERSION}" == *"13."* ]]; then pip3 install ".[cuda13,benchmark,test]" elif [ "${PLATFORM}" == "rocm" ]; then - pip3 install ".[rocm6,benchmark,test]" + ROCM_VERSION=$(cat /opt/rocm/.info/version) + ROCM_MAJOR="${ROCM_VERSION%%.*}" + pip3 install ".[rocm${ROCM_MAJOR},benchmark,test]" else pip3 install ".[benchmark,test]" fi diff --git a/test/mp_unit/bootstrap_tests.cc b/test/mp_unit/bootstrap_tests.cc index c28087a4..27ca1f5f 100644 --- a/test/mp_unit/bootstrap_tests.cc +++ b/test/mp_unit/bootstrap_tests.cc @@ -42,10 +42,17 @@ void BootstrapTest::bootstrapTestSendRecv(std::shared_ptr bo } } +void BootstrapTest::bootstrapTestIpcDomain(std::shared_ptr bootstrap) { + int nRanksPerIpcDomain = bootstrap->getNranksPerIpcDomain(); + EXPECT_GT(nRanksPerIpcDomain, 0); + EXPECT_LE(nRanksPerIpcDomain, bootstrap->getNranks()); +} + void BootstrapTest::bootstrapTestAll(std::shared_ptr bootstrap) { bootstrapTestAllGather(bootstrap); bootstrapTestBarrier(bootstrap); bootstrapTestSendRecv(bootstrap); + bootstrapTestIpcDomain(bootstrap); } TEST(BootstrapTest, WithId) { @@ -127,6 +134,7 @@ class MPIBootstrap : public mscclpp::Bootstrap { MPI_Comm_size(shmcomm, &shmrank); return shmrank; } + int getNranksPerIpcDomain() const override { return getNranksPerNode(); } void allGather(void* sendbuf, int size) override { MPI_Allgather(MPI_IN_PLACE, 0, MPI_BYTE, sendbuf, size, MPI_BYTE, MPI_COMM_WORLD); } diff --git a/test/mp_unit/mp_unit_tests.hpp b/test/mp_unit/mp_unit_tests.hpp index d079f711..54994273 100644 --- a/test/mp_unit/mp_unit_tests.hpp +++ b/test/mp_unit/mp_unit_tests.hpp @@ -56,6 +56,8 @@ class BootstrapTest : public MultiProcessTest { void bootstrapTestSendRecv(std::shared_ptr bootstrap); + void bootstrapTestIpcDomain(std::shared_ptr bootstrap); + void bootstrapTestAll(std::shared_ptr bootstrap); // Each test case should finish within 30 seconds. diff --git a/test/mp_unit/port_channel_tests.cu b/test/mp_unit/port_channel_tests.cu index 2bcd7a04..7cf8def0 100644 --- a/test/mp_unit/port_channel_tests.cu +++ b/test/mp_unit/port_channel_tests.cu @@ -667,22 +667,19 @@ __global__ void kernelPortChannelAtomicAddConcurrent(int64_t* localBuff, int nTr static constexpr int kMaxQps = 4; __constant__ DeviceHandle gMultiQpPortChans[kMaxQps]; -// Multi-QP bandwidth kernel: barrier on QP 0 only, then putWithSignal on all QPs. -// Only one signal/wait pair is needed for sync between two GPU kernels. +// Multi-QP bandwidth kernel: one thread per QP, putWithSignal per QP, parallel waits. __global__ void kernelMultiQpBandwidth(int nElemPerChan, int nIters, int numQps) { - if (threadIdx.x != 0) return; + int q = threadIdx.x; + if (q >= numQps) return; for (int i = 0; i < nIters; i++) { - // Barrier on QP 0 only — syncs both ranks - gMultiQpPortChans[0].signal(); - gMultiQpPortChans[0].wait(); - // Data transfer: put on all QPs simultaneously - for (int q = 0; q < numQps; q++) { - gMultiQpPortChans[q].putWithSignal(0, nElemPerChan * sizeof(int)); - } - // Wait for all remote data arrivals - for (int q = 0; q < numQps; q++) { - gMultiQpPortChans[q].wait(); + if (q == 0) { + gMultiQpPortChans[0].signal(); + gMultiQpPortChans[0].wait(); } + __syncthreads(); + gMultiQpPortChans[q].putWithSignal(0, nElemPerChan * sizeof(int)); + gMultiQpPortChans[q].wait(); + __syncthreads(); } } @@ -840,15 +837,15 @@ void PortChannelOneToOneTest::testMultiQpBandwidth(IbMode ibMode, int numQps) { for (int nElemPerChan : {256, 16 * 1024, 256 * 1024, 1024 * 1024, 4 * 1024 * 1024, 16 * 1024 * 1024, 32 * 1024 * 1024}) { - int nIters = 10000; + int nIters = 200; // Warm-up - kernelMultiQpBandwidth<<<1, 1>>>(nElemPerChan, 10, numQps); + kernelMultiQpBandwidth<<<1, numQps>>>(nElemPerChan, 10, numQps); MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); communicator->bootstrap()->barrier(); // Measure mscclpp::Timer timer; - kernelMultiQpBandwidth<<<1, 1>>>(nElemPerChan, nIters, numQps); + kernelMultiQpBandwidth<<<1, numQps>>>(nElemPerChan, nIters, numQps); MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); double elapsedUs = timer.elapsed(); communicator->bootstrap()->barrier(); @@ -873,20 +870,46 @@ void PortChannelOneToOneTest::testMultiQpBandwidth(IbMode ibMode, int numQps) { for (auto& m : remoteMems) registeredMemories.push_back(m); } +PERF_TEST(PortChannelOneToOneTest, SingleQpBandwidthIbHostMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::Host); + testMultiQpBandwidth(IbMode::Host, /*numQps=*/1); +} + +PERF_TEST(PortChannelOneToOneTest, TwoQpBandwidthIbHostMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::Host); + testMultiQpBandwidth(IbMode::Host, /*numQps=*/2); +} + PERF_TEST(PortChannelOneToOneTest, MultiQpBandwidthIbHostMode) { REQUIRE_IBVERBS; REQUIRE_GDR_FOR_IB_MODE(IbMode::Host); - for (int numQps : {1, 2, 4}) { - testMultiQpBandwidth(IbMode::Host, numQps); - } + testMultiQpBandwidth(IbMode::Host, /*numQps=*/4); +} + +PERF_TEST(PortChannelOneToOneTest, SingleQpBandwidthIbHostNoAtomicMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); + testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/1); +} + +PERF_TEST(PortChannelOneToOneTest, TwoQpBandwidthIbHostNoAtomicMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); + testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/2); +} + +PERF_TEST(PortChannelOneToOneTest, ThreeQpBandwidthIbHostNoAtomicMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); + testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/3); } PERF_TEST(PortChannelOneToOneTest, MultiQpBandwidthIbHostNoAtomicMode) { REQUIRE_IBVERBS; REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); - for (int numQps : {1, 2, 4}) { - testMultiQpBandwidth(IbMode::HostNoAtomic, numQps); - } + testMultiQpBandwidth(IbMode::HostNoAtomic, /*numQps=*/4); } // Multi-QP flush-stress kernel: one thread per QP, all calling putWithSignalAndFlush @@ -911,7 +934,7 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps) if (gEnv->rank >= numRanksToUse) return; const int rank = communicator->bootstrap()->getRank(); - const int maxElemPerChan = 64 * 1024; + const int maxElemPerChan = 8 * 1024 * 1024; std::vector> sendBuffs; std::vector localMems; @@ -930,8 +953,8 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps) const std::string qpLabel = std::to_string(numQps) + " QP" + (numQps > 1 ? "s" : ""); - for (int nElemPerChan : {256, 4 * 1024, 64 * 1024}) { - int nIters = 2000; + for (int nElemPerChan : {256, 4 * 1024, 64 * 1024, 256 * 1024, 1024 * 1024, 4 * 1024 * 1024, 8 * 1024 * 1024}) { + int nIters = (nElemPerChan >= 256 * 1024) ? 200 : 2000; kernelMultiQpFlushStress<<<1, numQps>>>(nElemPerChan, 10, numQps); MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); communicator->bootstrap()->barrier(); @@ -948,8 +971,10 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps) int bytesPerChan = nElemPerChan * (int)sizeof(int); std::string sizeLabel = (bytesPerChan >= 1024) ? (std::to_string(bytesPerChan / 1024) + " KB") : (std::to_string(bytesPerChan) + " B"); + double aggGbps = ((double)bytesPerChan * numQps) / usPerIter * 1e-3; // bytes/us = MB/s × 1e-3 = GB/s ::mscclpp::test::reportPerfResult(sizeLabel + " (" + qpLabel + ") per-iter", usPerIter, "us"); ::mscclpp::test::reportPerfResult(sizeLabel + " (" + qpLabel + ") per-iter/QP", usPerIterPerQp, "us"); + ::mscclpp::test::reportPerfResult(sizeLabel + " (" + qpLabel + ") aggregate", aggGbps, "GB/s"); } } @@ -959,20 +984,40 @@ void PortChannelOneToOneTest::testMultiQpFlushStress(IbMode ibMode, int numQps) for (auto& m : remoteMems) registeredMemories.push_back(m); } +PERF_TEST(PortChannelOneToOneTest, SingleQpFlushStressIbHostMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::Host); + testMultiQpFlushStress(IbMode::Host, /*numQps=*/1); +} + +PERF_TEST(PortChannelOneToOneTest, TwoQpFlushStressIbHostMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::Host); + testMultiQpFlushStress(IbMode::Host, /*numQps=*/2); +} + PERF_TEST(PortChannelOneToOneTest, MultiQpFlushStressIbHostMode) { REQUIRE_IBVERBS; REQUIRE_GDR_FOR_IB_MODE(IbMode::Host); - for (int numQps : {1, 2, 4}) { - testMultiQpFlushStress(IbMode::Host, numQps); - } + testMultiQpFlushStress(IbMode::Host, /*numQps=*/4); +} + +PERF_TEST(PortChannelOneToOneTest, SingleQpFlushStressIbHostNoAtomicMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); + testMultiQpFlushStress(IbMode::HostNoAtomic, /*numQps=*/1); +} + +PERF_TEST(PortChannelOneToOneTest, TwoQpFlushStressIbHostNoAtomicMode) { + REQUIRE_IBVERBS; + REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); + testMultiQpFlushStress(IbMode::HostNoAtomic, /*numQps=*/2); } PERF_TEST(PortChannelOneToOneTest, MultiQpFlushStressIbHostNoAtomicMode) { REQUIRE_IBVERBS; REQUIRE_GDR_FOR_IB_MODE(IbMode::HostNoAtomic); - for (int numQps : {1, 2, 4}) { - testMultiQpFlushStress(IbMode::HostNoAtomic, numQps); - } + testMultiQpFlushStress(IbMode::HostNoAtomic, /*numQps=*/4); } // Same-channel concurrent-flush kernel: N GPU threads on the same PortChannel each call diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt index a345effc..0d075ac8 100644 --- a/test/unit/CMakeLists.txt +++ b/test/unit/CMakeLists.txt @@ -14,5 +14,6 @@ target_sources(unit_tests PRIVATE utils_tests.cc utils_internal_tests.cc compile_tests.cu + gpu_data_types_tests.cu local_channel_tests.cu ) diff --git a/test/unit/gpu_data_types_tests.cu b/test/unit/gpu_data_types_tests.cu new file mode 100644 index 00000000..5f91c684 --- /dev/null +++ b/test/unit/gpu_data_types_tests.cu @@ -0,0 +1,175 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include + +#include "../framework.hpp" + +namespace { + +constexpr int kConversionPaths = 3; + +template +std::array makeArray(Args... args) { + return {static_cast(args)...}; +} + +__device__ uint32_t floatToBitsDevice(float value) { + union { + float f; + uint32_t u; + } cvt = {value}; + return cvt.u; +} + +uint32_t floatToBitsHost(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +__global__ void kernelE4m3b15TypeConvert(const float* input, int encodeCases, const uint8_t* raw, int decodeCases, + uint8_t* encoded, uint32_t* decodedBits) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + for (int offset = 0; offset < encodeCases; offset += 4) { + mscclpp::f32x4 inputX4; + for (int i = 0; i < 4; ++i) { + inputX4.data[i] = input[offset + i]; + } + + mscclpp::f8_e4m3b15x4 encodedX4 = mscclpp::to(inputX4); + for (int i = 0; i < 4; ++i) { + encoded[offset + i] = encodedX4.data[i].__x; + } + + for (int pair = 0; pair < 2; ++pair) { + mscclpp::f32x2 inputX2; + inputX2.data[0] = input[offset + pair * 2]; + inputX2.data[1] = input[offset + pair * 2 + 1]; + mscclpp::f8_e4m3b15x2 encodedX2 = mscclpp::to(inputX2); + encoded[encodeCases + offset + pair * 2] = encodedX2.data[0].__x; + encoded[encodeCases + offset + pair * 2 + 1] = encodedX2.data[1].__x; + } + } + + for (int i = 0; i < encodeCases; ++i) { + encoded[2 * encodeCases + i] = __fp8_e4m3b15(input[i]).__x; + } + + for (int offset = 0; offset < decodeCases; offset += 4) { + mscclpp::f8_e4m3b15x4 rawX4; + for (int i = 0; i < 4; ++i) { + rawX4.data[i] = __fp8_e4m3b15::fromRaw(raw[offset + i]); + } + + mscclpp::f32x4 decodedX4 = mscclpp::to(rawX4); + for (int i = 0; i < 4; ++i) { + decodedBits[offset + i] = floatToBitsDevice(decodedX4.data[i]); + } + + for (int pair = 0; pair < 2; ++pair) { + mscclpp::f8_e4m3b15x2 rawX2; + rawX2.data[0] = __fp8_e4m3b15::fromRaw(raw[offset + pair * 2]); + rawX2.data[1] = __fp8_e4m3b15::fromRaw(raw[offset + pair * 2 + 1]); + mscclpp::f32x2 decodedX2 = mscclpp::to(rawX2); + decodedBits[decodeCases + offset + pair * 2] = floatToBitsDevice(decodedX2.data[0]); + decodedBits[decodeCases + offset + pair * 2 + 1] = floatToBitsDevice(decodedX2.data[1]); + } + } + + for (int i = 0; i < decodeCases; ++i) { + decodedBits[2 * decodeCases + i] = floatToBitsDevice(float(__fp8_e4m3b15::fromRaw(raw[i]))); + } +} + +} // namespace + +TEST(GpuDataTypesTest, E4m3b15TypeConvert) { + const float inf = std::numeric_limits::infinity(); + const float nan = std::numeric_limits::quiet_NaN(); + const float maxFloat = std::numeric_limits::max(); + + // Each input value maps to the byte at the same index in expectedEncoded. The fp8_e4m3b15 format has no + // NaN/Inf encoding, so NaN, Inf, and overflow inputs saturate to +/-1.875 (max byte 0x7f/0xff). + const auto input = makeArray(0.0f, -0.0f, // +/-0 + 0x1.0p-19f, -0x1.0p-19f, // +/-2^-19: underflows to signed 0 + 0x1.0p-18f, -0x1.0p-18f, // +/-2^-18: rounds to min subnormal + 0x1.0p-17f, -0x1.0p-17f, // +/-2^-17: min subnormal + 0x1.0p-14f, -0x1.0p-14f, // +/-2^-14: min normal + 0x1.0fcp-2f, -0x1.0fcp-2f, // Boundary rounds down in magnitude + 0x1.0fep-2f, -0x1.0fep-2f, // Boundary rounds up in magnitude + 0x1.cfep-2f, -0x1.cfep-2f, // Boundary rounds to +/-0.46875 + 0x1.cp0f, -0x1.cp0f, // +/-1.75: max finite + 2.0f, -2.0f, // Overflow saturation + inf, -inf, // +/-Inf saturation + nan, -maxFloat); // NaN / large negative saturation + + const auto expectedEncoded = makeArray(0x00, 0x80, // +/-0 + 0x00, 0x80, // Underflow to signed zero + 0x01, 0x81, // Round to min signed subnormal + 0x01, 0x81, // Min signed subnormal + 0x08, 0x88, // Min signed normal + 0x68, 0xe8, // Boundary rounds to +/-0.25 + 0x69, 0xe9, // Boundary rounds to +/-0.28125 + 0x6f, 0xef, // Boundary rounds to +/-0.46875 + 0x7e, 0xfe, // Max finite at fp16 grid (1.75) + 0x7f, 0xff, // Overflow saturation (1.875) + 0x7f, 0xff, // Inf saturation (1.875) + 0x7f, 0xff); // NaN / large negative saturation (1.875) + + // Raw bytes to decode, with expectedDecoded giving the exact float value at the same index. + const auto raw = makeArray(0x00, 0x80, // +/-0 + 0x01, 0x81, // +/-2^-17: min subnormal + 0x08, 0x88, // +/-2^-14: min normal + 0x68, 0xe8, // +/-0.25 + 0x69, 0xe9, // +/-0.28125 + 0x7e, 0xfe); // +/-1.75: max finite + const auto expectedDecoded = makeArray(0.0f, -0.0f, // +/-0 + 0x1.0p-17f, -0x1.0p-17f, // +/-2^-17: min subnormal + 0x1.0p-14f, -0x1.0p-14f, // +/-2^-14: min normal + 0x1.0p-2f, -0x1.0p-2f, // +/-0.25 + 0x1.2p-2f, -0x1.2p-2f, // +/-0.28125 + 0x1.cp0f, -0x1.cp0f); // +/-1.75: max finite + + ASSERT_EQ(input.size(), expectedEncoded.size()); + ASSERT_EQ(raw.size(), expectedDecoded.size()); + ASSERT_EQ(input.size() % 4, size_t(0)); + ASSERT_EQ(raw.size() % 4, size_t(0)); + + auto inputDev = mscclpp::detail::gpuCallocShared(input.size()); + auto rawDev = mscclpp::detail::gpuCallocShared(raw.size()); + auto encodedDev = mscclpp::detail::gpuCallocShared(input.size() * kConversionPaths); + auto decodedBitsDev = mscclpp::detail::gpuCallocShared(raw.size() * kConversionPaths); + + mscclpp::gpuMemcpy(inputDev.get(), input.data(), input.size(), cudaMemcpyHostToDevice); + mscclpp::gpuMemcpy(rawDev.get(), raw.data(), raw.size(), cudaMemcpyHostToDevice); + + kernelE4m3b15TypeConvert<<<1, 1>>>(inputDev.get(), static_cast(input.size()), rawDev.get(), + static_cast(raw.size()), encodedDev.get(), decodedBitsDev.get()); + MSCCLPP_CUDATHROW(cudaGetLastError()); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector encoded(input.size() * kConversionPaths); + std::vector decodedBits(raw.size() * kConversionPaths); + mscclpp::gpuMemcpy(encoded.data(), encodedDev.get(), encoded.size(), cudaMemcpyDeviceToHost); + mscclpp::gpuMemcpy(decodedBits.data(), decodedBitsDev.get(), decodedBits.size(), cudaMemcpyDeviceToHost); + + for (int path = 0; path < kConversionPaths; ++path) { + for (size_t i = 0; i < input.size(); ++i) { + EXPECT_EQ(static_cast(encoded[path * input.size() + i]), static_cast(expectedEncoded[i])); + } + } + + for (int path = 0; path < kConversionPaths; ++path) { + for (size_t i = 0; i < raw.size(); ++i) { + EXPECT_EQ(decodedBits[path * raw.size() + i], floatToBitsHost(expectedDecoded[i])); + } + } +} diff --git a/tools/npkit/npkit_trace_generator.py b/tools/npkit/npkit_trace_generator.py index 294516e6..52f899dc 100644 --- a/tools/npkit/npkit_trace_generator.py +++ b/tools/npkit/npkit_trace_generator.py @@ -39,6 +39,8 @@ def parse_npkit_event_header(npkit_event_header_path): "PIPELINE", "SEM_RELEASE", "SEM_ACQUIRE", + "MULTI_STORE", + "MULTI_STORE_PKT", ] executor_op_to_offset = {} for executor_op in executor_ops: