diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index b56bc285aa..49738a429b 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -1,40 +1,40 @@ # CollectiveX Sweep — one structured run instead of thousands of dispatches. # -# Shape (mirrors the InferenceX CI tracker): setup -> sweep (a MATRIX job = "a job with other jobs -# in it") -> aggregate (the collector "at the end"). The matrix unit is a SHARD = one allocation that -# sweeps many cases sharing (sku, backend, mode, resource) — generate_matrix's own grouping, chunked -# so no cell exceeds the job budget. Each cell emits a handful of per-case JSONs; the aggregate job -# collects every shard into ONE line-delimited file (results/aggregate/*.ndjson) so there aren't -# thousands of individual result files. Run once per backend (deepep / uccl / flashinfer / -# deepep-hybrid / nccl-ep, + deepep_v2) for full parity. +# Shape: setup resolves the suites into a shard matrix, then one GPU cell runs per +# shard, executes the benchmark, and uploads the emitted neutral artifacts as-is. +# This repo only schedules, executes, and uploads — it does not validate, promote, +# rank, recommend, or decide what a frontend displays. This file is registered +# on the default branch, so its collectivex branch revision can be dispatched with --ref. name: CollectiveX Sweep +permissions: + actions: read + contents: read on: workflow_dispatch: inputs: backend: - description: EP library to sweep (deepep matrix is remapped onto the others, capability-filtered) + description: "EP library to sweep — 'all' runs every EP backend in one matrix" type: choice - default: deepep - options: [deepep, uccl, flashinfer, deepep-hybrid, nccl-ep] - deepep_v2: - description: DeepEP V2 from-source kernels (kernel_gen=v2; deepep backend only) - type: boolean - default: false + default: all + options: [all, deepep, deepep-v2, uccl, deepep-hybrid, mori, nccl-ep] suites: description: "'all' or comma-list of suite names" type: string default: all only_sku: - description: Restrict to one SKU (h100-dgxc|h200|b300|b200-dgxc|gb200|gb300|mi355x); blank = all + description: Restrict to one GHA runner pool (h100-dgxc|h200-dgxc|b300|b200-dgxc|gb200|gb300|mi300x|mi355x); blank = all type: string default: '' - max_cases: - description: Max cases per shard cell (chunk larger shards) + exclude_skus: + description: Comma-list of runner pools to drop from the matrix (partial run, e.g. b300); disjoint from only_sku; blank = none type: string - default: '14' - + default: '' + ep_sizes: + description: Keep only shards whose expert-parallel degree is in this comma-list (8 keeps EP8 only, dropping EP16 so GB SKUs co-schedule with 8-GPU SKUs; blank = all) + type: string + default: '' concurrency: - group: cx-sweep-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.deepep_v2 }}-${{ inputs.only_sku }} + group: cx-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.only_sku }} cancel-in-progress: false jobs: @@ -44,28 +44,93 @@ jobs: outputs: matrix: ${{ steps.gen.outputs.matrix }} n: ${{ steps.gen.outputs.n }} + max_parallel: ${{ steps.gen.outputs.max_parallel }} steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } - - run: pip install --quiet pyyaml + with: { clean: true, persist-credentials: false } + - name: Install matrix dependencies + # PyYAML is used by sweep_matrix.py to read the suites config. + run: python3 -m pip install --quiet PyYAML==6.0.2 - id: gen working-directory: experimental/CollectiveX + env: + INPUT_BACKEND: ${{ inputs.backend }} + INPUT_SUITES: ${{ inputs.suites }} + INPUT_ONLY_SKU: ${{ inputs.only_sku }} + INPUT_EXCLUDE_SKUS: ${{ inputs.exclude_skus }} + INPUT_EP_SIZES: ${{ inputs.ep_sizes }} + run: | + set -euo pipefail + args=(--suites "$INPUT_SUITES") + case "$INPUT_BACKEND" in + all) args+=(--backends all) ;; + *) args+=(--backend "$INPUT_BACKEND") ;; + esac + [ -n "$INPUT_ONLY_SKU" ] && args+=(--only-sku "$INPUT_ONLY_SKU") + [ -n "$INPUT_EXCLUDE_SKUS" ] && args+=(--exclude-skus "$INPUT_EXCLUDE_SKUS") + [ -n "$INPUT_EP_SIZES" ] && args+=(--ep-sizes "$INPUT_EP_SIZES") + python3 sweep_matrix.py "${args[@]}" --out matrix_full.json >/dev/null + SLIM=$(python3 -c "import json;m=json.load(open('matrix_full.json'));print(json.dumps({'include':[{k:v for k,v in x.items() if k!='case_ids'} for x in m['include']]}))") + { + echo "matrix=$SLIM" + echo "n=$(python3 -c "import json;print(len(json.load(open('matrix_full.json'))['include']))")" + echo "source_backends=$(python3 -c "import json;m=json.load(open('matrix_full.json'));print(' '.join(sorted({x['backend'] for x in m['include']} & {'deepep-v2','deepep-hybrid'})))")" + echo "max_parallel=$(python3 -c "import json;m=json.load(open('matrix_full.json'));w=max(x['execution_weight'] for x in m['include']);print(max(1,min(10,4096//w)))")" + } >> "$GITHUB_OUTPUT" + python3 -c "import json;m=json.load(open('matrix_full.json'));print('execution-cells:',len(m['include']))" + - name: Prepare pinned backend source archive + if: ${{ steps.gen.outputs.source_backends != '' }} + working-directory: experimental/CollectiveX + env: + SOURCE_BACKENDS: ${{ steps.gen.outputs.source_backends }} + COLLECTIVEX_EXECUTION_ID: ${{ github.run_id }}_${{ github.run_attempt }}_sources run: | set -euo pipefail - ov=""; [ "${{ inputs.backend }}" != "deepep" ] && ov="--backend ${{ inputs.backend }}" - v2=""; [ "${{ inputs.deepep_v2 }}" = "true" ] && v2="--deepep-v2" - os=""; [ -n "${{ inputs.only_sku }}" ] && os="--only-sku ${{ inputs.only_sku }}" - # full matrix (with cases) -> artifact for the cells; slim (no cases) -> the strategy output. - python3 sweep_matrix.py --suites "${{ inputs.suites }}" --max-cases "${{ inputs.max_cases }}" $ov $v2 $os --out matrix_full.json >/dev/null - SLIM=$(python3 -c "import json;m=json.load(open('matrix_full.json'));print(json.dumps({'include':[{k:v for k,v in x.items() if k!='cases'} for x in m['include']]}))") - echo "matrix=$SLIM" >> "$GITHUB_OUTPUT" - echo "n=$(python3 -c "import json;print(len(json.load(open('matrix_full.json'))['include']))")" >> "$GITHUB_OUTPUT" - python3 -c "import json;m=json.load(open('matrix_full.json'));print('shard-cells:',len(m['include']),'cases:',sum(x['n'] for x in m['include']))" + source runtime/common.sh + work="$RUNNER_TEMP/collectivex-backend-sources" + archive="$RUNNER_TEMP/collectivex-backend-sources.tar" + rm -rf -- "$work" "$archive" + umask 077 + mkdir -m 700 "$work" + mkdir -p "$work/experimental/CollectiveX" + read -r -a backends <<< "$SOURCE_BACKENDS" + [ "${#backends[@]}" -gt 0 ] + source_root="$work/experimental/CollectiveX/.cx_sources" + members=() + for backend in "${backends[@]}"; do + cx_prepare_backend_source "$work" "$backend" + source_path="$(cx_backend_source_path "$source_root" "$backend")" + source_basename="${source_path#"$source_root/"}" + [ -n "$source_basename" ] \ + && [ "$source_path" = "$source_root/$source_basename" ] \ + && [[ "$source_basename" != */* ]] + members+=(".cx_sources/$source_basename") + done + cx_cleanup_private_logs 0 + tar --sort=name --mtime='@1' --owner=0 --group=0 --numeric-owner \ + --no-recursion -C "$work/experimental/CollectiveX" \ + -cf "$archive" .cx_sources + for member in "${members[@]}"; do + tar --sort=name --mtime='@1' --owner=0 --group=0 --numeric-owner \ + -C "$work/experimental/CollectiveX" -rf "$archive" "$member" + done + sha256sum "$archive" + rm -rf -- "$work" + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + if: ${{ steps.gen.outputs.source_backends != '' }} + with: + name: cxbackend-sources-${{ github.run_id }} + path: ${{ runner.temp }}/collectivex-backend-sources.tar + if-no-files-found: error + overwrite: true + retention-days: 14 - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: cxsweep-matrix-${{ github.run_id }} path: experimental/CollectiveX/matrix_full.json if-no-files-found: error + overwrite: true + retention-days: 14 # ---- sweep: ONE matrix cell per shard (the parent job with child jobs) ---- sweep: @@ -73,82 +138,386 @@ jobs: if: ${{ fromJSON(needs.setup.outputs.n) > 0 }} strategy: fail-fast: false - max-parallel: 10 # don't saturate the ~20-runner fleet; cells queue as slots free + max-parallel: ${{ fromJSON(needs.setup.outputs.max_parallel) }} matrix: ${{ fromJSON(needs.setup.outputs.matrix) }} - # h200 label spans two clusters; pin to the validated dgxc pool (mirrors collectivex-experimental). - runs-on: ${{ matrix.sku == 'h200' && 'h200-dgxc' || matrix.sku }} + runs-on: ${{ matrix.sku }} timeout-minutes: 350 env: CX_BENCH: ${{ matrix.backend }} - CX_DEEPEP_V2: ${{ matrix.deepep_v2 && '1' || '' }} CX_NODES: ${{ matrix.nodes }} - CX_SHARD_FILE: results/.shard_${{ matrix.id }}.json + CX_GPUS_PER_NODE: ${{ matrix.gpus_per_node }} + CX_SCALE_UP_DOMAIN: ${{ matrix.scale_up_domain }} + CX_SHARD_FILE: .shards/${{ matrix.id }}.json + CX_SHARD_SKU: ${{ matrix.sku }} + COLLECTIVEX_CANONICAL_GHA: '1' COLLECTIVEX_SOURCE_SHA: ${{ github.sha }} - CX_NODELIST: ${{ matrix.sku == 'mi355x' && 'mia1-p01-g10,mia1-p01-g15' || '' }} - CX_STAGE_DIR: ${{ matrix.sku == 'gb200' && '/mnt/lustre01/users-public/sa-shared/cx-stage' || '' }} + COLLECTIVEX_ARTIFACT_NAME: ${{ format('cxshard-{0}-{1}-{2}', matrix.id, github.run_id, github.run_attempt) }} + # Consolidated shards run one bounded build-group in one Slurm allocation. + # MI300X's compute partition has a 180-minute ceiling; the other production + # pools accept 300 minutes. Allocations release as soon as the shard finishes. + CX_TIME: ${{ matrix.sku == 'mi300x' && '180' || '300' }} + COLLECTIVEX_EXECUTION_ID: ${{ github.run_id }}_${{ github.run_attempt }}_${{ matrix.id }} + CX_JOB_PARENT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || '/tmp' }} + CX_JOB_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}', github.run_id, github.run_attempt, matrix.id) }} + CX_SOURCE_ROOT: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}/source', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}/source', github.run_id, github.run_attempt, matrix.id) }} + HOME: ${{ matrix.launcher == 'mi-amds' && format('/tmp/inferencex-collectivex-parent-{0}-{1}-{2}/inferencex-collectivex-{0}-{1}-{2}/home', github.run_id, github.run_attempt, matrix.id) || format('/tmp/inferencex-collectivex-{0}-{1}-{2}/home', github.run_id, github.run_attempt, matrix.id) }} steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } + - name: Provision host Python + # The self-hosted runners ship a bare python3 with no system pip. The host + # launcher and summarize.py invoke bare `python3` and need PyYAML/numpy, so + # build a venv and put it on PATH for every later step. venv + ensurepip are + # available even where system pip is not. + env: + HOME: ${{ runner.temp }} + run: | + set -euo pipefail + venv="$RUNNER_TEMP/cx-hostpy" + python3 -m venv --clear "$venv" + "$venv/bin/python3" -m pip install --quiet --no-cache-dir --upgrade pip + "$venv/bin/python3" -m pip install --quiet --no-cache-dir \ + PyYAML==6.0.2 "numpy>=1.26,<3" + echo "$venv/bin" >> "$GITHUB_PATH" + - name: Prepare isolated source + id: source + env: + COLLECTIVEX_REPOSITORY: ${{ github.repository }} + run: | + set -euo pipefail + python3 - <<'PY' + import os + import re + import shutil + import stat + import time + + pattern = re.compile(r"inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+") + link_pattern = re.compile( + r"/tmp/inferencex-collectivex-parent-[0-9]+-[0-9]+-[A-Za-z0-9._-]+" + ) + parent = os.environ["CX_JOB_PARENT"] + scan_parent = parent + if parent != "/tmp": + if not link_pattern.fullmatch(parent): + raise SystemExit("CollectiveX isolated parent is invalid") + shared_parent = os.path.join(os.getcwd(), ".collectivex-jobs") + try: + shared_metadata = os.lstat(shared_parent) + except FileNotFoundError: + os.mkdir(shared_parent, mode=0o700) + shared_metadata = os.lstat(shared_parent) + if ( + not stat.S_ISDIR(shared_metadata.st_mode) + or shared_metadata.st_uid != os.getuid() + or stat.S_IMODE(shared_metadata.st_mode) != 0o700 + ): + raise SystemExit("CollectiveX shared parent is unsafe") + if os.path.lexists(parent): + raise SystemExit("CollectiveX isolated parent already exists") + os.symlink(shared_parent, parent) + scan_parent = shared_parent + cutoff = time.time() - 86400 + for entry in os.scandir(scan_parent): + if not pattern.fullmatch(entry.name): + continue + try: + metadata = entry.stat(follow_symlinks=False) + except FileNotFoundError: + continue + if ( + not stat.S_ISDIR(metadata.st_mode) + or metadata.st_uid != os.getuid() + or stat.S_IMODE(metadata.st_mode) != 0o700 + or metadata.st_mtime >= cutoff + ): + continue + marked = False + for marker_name in ("cleanup-safe", "cleanup-unsafe"): + try: + marker = os.stat( + os.path.join(entry.path, marker_name), follow_symlinks=False + ) + except FileNotFoundError: + continue + marked = ( + stat.S_ISREG(marker.st_mode) + and marker.st_uid == os.getuid() + and stat.S_IMODE(marker.st_mode) == 0o600 + ) + if marked: + break + if marked: + shutil.rmtree(entry.path) + PY + [ "${CX_JOB_ROOT%/*}" = "$CX_JOB_PARENT" ] \ + && [[ "${CX_JOB_ROOT##*/}" =~ ^inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX isolated root is invalid" >&2; exit 1; } + [ "$CX_SOURCE_ROOT" = "$CX_JOB_ROOT/source" ] \ + || { echo "CollectiveX source root is invalid" >&2; exit 1; } + if [ -e "$CX_JOB_ROOT" ] || [ -L "$CX_JOB_ROOT" ]; then + echo "CollectiveX isolated root already exists" >&2 + exit 1 + fi + umask 077 + mkdir -m 700 -- "$CX_JOB_ROOT" + trap 'rc=$?; if [ "$rc" != 0 ]; then rm -rf -- "$CX_JOB_ROOT"; [ "$CX_JOB_PARENT" = /tmp ] || rm -f -- "$CX_JOB_PARENT"; fi; exit "$rc"' EXIT + mkdir -m 700 -- "$HOME" "$CX_JOB_ROOT/control" "$CX_JOB_ROOT/artifact" "$CX_SOURCE_ROOT" + : > "$CX_JOB_ROOT/cleanup-safe" + if ! { + GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null git init -q "$CX_SOURCE_ROOT" + GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null \ + git -C "$CX_SOURCE_ROOT" remote add origin \ + "https://github.com/${COLLECTIVEX_REPOSITORY}.git" + GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null \ + git -C "$CX_SOURCE_ROOT" -c credential.helper= -c protocol.version=2 \ + fetch -q --no-tags --depth=1 origin "$COLLECTIVEX_SOURCE_SHA" + GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null \ + git -C "$CX_SOURCE_ROOT" -c advice.detachedHead=false \ + checkout -q --detach FETCH_HEAD + [ "$(git -C "$CX_SOURCE_ROOT" rev-parse HEAD)" = "$COLLECTIVEX_SOURCE_SHA" ] + } /dev/null 2>&1; then + echo "CollectiveX source preparation failed" >&2 + exit 1 + fi + [ "$(stat -c '%a' "$CX_JOB_ROOT")" = 700 ] \ + || { echo "CollectiveX isolated root has unsafe permissions" >&2; exit 1; } + echo 'prepared=true' >> "$GITHUB_OUTPUT" + trap - EXIT - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 with: name: cxsweep-matrix-${{ github.run_id }} - path: experimental/CollectiveX - - name: Extract this shard's cases (stdlib only — no runner deps) - working-directory: experimental/CollectiveX + path: ${{ env.CX_JOB_ROOT }}/control + - name: Download pinned backend source archive + if: ${{ matrix.backend == 'deepep-v2' || matrix.backend == 'deepep-hybrid' }} + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: cxbackend-sources-${{ github.run_id }} + path: ${{ env.CX_JOB_ROOT }}/control + - name: Install pinned backend source seed + if: ${{ matrix.backend == 'deepep-v2' || matrix.backend == 'deepep-hybrid' }} + env: + EXPECTED_BACKEND: ${{ matrix.backend }} run: | set -euo pipefail - python3 -c " - import json - m=json.load(open('matrix_full.json')) - s=[x for x in m['include'] if x['id']=='${{ matrix.id }}'] - assert s, 'shard ${{ matrix.id }} not in matrix' - s=s[0] - json.dump({'id':s['id'],'sku':s['sku'],'backend':s['backend'],'nodes':s['nodes'],'deepep_v2':s['deepep_v2'],'cases':s['cases']}, open('results/.shard_${{ matrix.id }}.json','w')) - print('shard ${{ matrix.id }}:', len(s['cases']), 'cases') - " - - name: Sweep shard ${{ matrix.id }} (${{ matrix.n }} cases, one allocation) + archive="$CX_JOB_ROOT/control/collectivex-backend-sources.tar" + destination="$CX_SOURCE_ROOT/experimental/CollectiveX" + seed_root="$destination/.cx_sources" + [ -f "$archive" ] && [ ! -e "$seed_root" ] && [ ! -L "$seed_root" ] + chmod 600 -- "$archive" + source "$destination/runtime/common.sh" + source_path="$(cx_backend_source_path "$seed_root" "$EXPECTED_BACKEND")" + source_basename="${source_path#"$seed_root/"}" + [ -n "$source_basename" ] \ + && [ "$source_path" = "$seed_root/$source_basename" ] \ + && [[ "$source_basename" != */* ]] + python3 "$destination/source_archive.py" \ + "$archive" "$destination" "$source_basename" + cx_backend_source_is_valid "$EXPECTED_BACKEND" "$source_path" + printf 'CX_BACKEND_SOURCE_SEED_ROOT=%s\n' "$seed_root" >> "$GITHUB_ENV" + - name: Extract and validate this execution control + run: | + set -euo pipefail + cd "$CX_SOURCE_ROOT/experimental/CollectiveX" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + python3 sweep_matrix.py \ + --extract-from "$CX_JOB_ROOT/control/matrix_full.json" \ + --shard-id '${{ matrix.id }}' \ + --expect-sku '${{ matrix.sku }}' \ + --expect-backend '${{ matrix.backend }}' \ + --expect-nodes '${{ matrix.nodes }}' \ + --out '${{ env.CX_SHARD_FILE }}' >/dev/null + - name: Execute sweep cell ${{ matrix.id }} + id: sweep_shard env: - RUNNER_NAME: ${{ runner.name }} - run: bash "experimental/CollectiveX/launchers/launch_${RUNNER_NAME%%_*}.sh" + COLLECTIVEX_OPERATOR_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_OPERATOR_CONFIG_V1 }} + COLLECTIVEX_NETWORK_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_NETWORK_CONFIG_V1 }} + COLLECTIVEX_H100_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_H100_CONFIG_V1 }} + COLLECTIVEX_B300_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_B300_CONFIG_V1 }} + COLLECTIVEX_B200_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_B200_CONFIG_V1 }} + COLLECTIVEX_MI300_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_MI300_CONFIG_V1 }} + COLLECTIVEX_MI355_CONFIG_CONTENT: ${{ secrets.COLLECTIVEX_MI355_CONFIG_V1 }} + COLLECTIVEX_OPERATOR_AUDIT_SALT: ${{ secrets.COLLECTIVEX_AUDIT_SALT_V1 }} + COLLECTIVEX_OPERATOR_CONFIG_REQUIRED: '1' + run: | + set -euo pipefail + umask 077 + operator_config="$CX_JOB_ROOT/operator-config.json" + python3 - "$operator_config" <<'PY' + import json + import os + from pathlib import Path + import sys + + def pairs(items): + result = {} + for key, value in items: + if key in result: + raise ValueError("duplicate configuration key") + result[key] = value + return result + + def load(name): + return json.loads( + os.environ[name], object_pairs_hook=pairs, + parse_constant=lambda _: (_ for _ in ()).throw(ValueError()), + ) + + base = load("COLLECTIVEX_OPERATOR_CONFIG_CONTENT") + for overlay_name in ( + "COLLECTIVEX_NETWORK_CONFIG_CONTENT", "COLLECTIVEX_H100_CONFIG_CONTENT", + "COLLECTIVEX_B300_CONFIG_CONTENT", + "COLLECTIVEX_B200_CONFIG_CONTENT", "COLLECTIVEX_MI300_CONFIG_CONTENT", + "COLLECTIVEX_MI355_CONFIG_CONTENT", + ): + overlay_raw = os.environ.get(overlay_name, "") + if not overlay_raw: + continue + overlay = load(overlay_name) + private_fields = { + "socket_ifname", "rdma_devices", "ib_gid_index", "rdma_service_level", + "rdma_traffic_class", + } + if set(overlay) != {"schema_version", "runners"} or overlay["schema_version"] != 1: + raise ValueError("invalid network overlay envelope") + if not isinstance(overlay["runners"], dict) or not overlay["runners"]: + raise ValueError("invalid network overlay runners") + if overlay_name == "COLLECTIVEX_B200_CONFIG_CONTENT": + if set(overlay["runners"]) != {"b200-dgxc"}: + raise ValueError("invalid B200 overlay runner") + private_fields.update({"nodelist", "qos"}) + elif overlay_name == "COLLECTIVEX_H100_CONFIG_CONTENT": + if set(overlay["runners"]) != {"h100-dgxc"}: + raise ValueError("invalid H100 overlay runner") + private_fields.add("exclude_nodes") + elif overlay_name == "COLLECTIVEX_B300_CONFIG_CONTENT": + if set(overlay["runners"]) != {"b300"}: + raise ValueError("invalid B300 overlay runner") + private_fields.add("exclude_nodes") + elif overlay_name == "COLLECTIVEX_MI355_CONFIG_CONTENT": + if set(overlay["runners"]) != {"mi355x"}: + raise ValueError("invalid MI355 overlay runner") + elif overlay_name == "COLLECTIVEX_MI300_CONFIG_CONTENT": + if set(overlay["runners"]) != {"mi300x"}: + raise ValueError("invalid MI300 overlay runner") + private_fields.update({ + "partition", "squash_dir", "stage_dir", "exclude_nodes", + "nodelist", "lock_dir", + }) + for runner, fields in overlay["runners"].items(): + if not isinstance(fields, dict) or not fields: + raise ValueError("invalid network overlay runner") + if runner not in base.get("runners", {}): + if overlay_name != "COLLECTIVEX_MI300_CONFIG_CONTENT": + raise ValueError("invalid network overlay runner") + base.setdefault("runners", {})[runner] = {} + if set(fields) - private_fields: + raise ValueError("private overlay contains a disallowed field") + base["runners"][runner].update(fields) + payload = json.dumps(base, sort_keys=True, separators=(",", ":")) + "\n" + if len(payload.encode()) > 65536: + raise ValueError("merged operator configuration is too large") + path = Path(sys.argv[1]) + descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + stream.write(payload) + PY + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT COLLECTIVEX_NETWORK_CONFIG_CONTENT + unset COLLECTIVEX_H100_CONFIG_CONTENT COLLECTIVEX_B300_CONFIG_CONTENT + unset COLLECTIVEX_B200_CONFIG_CONTENT + unset COLLECTIVEX_MI300_CONFIG_CONTENT COLLECTIVEX_MI355_CONFIG_CONTENT + unset COLLECTIVEX_OPERATOR_CONFIG_REQUIRED + export COLLECTIVEX_OPERATOR_CONFIG="$operator_config" + export COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL=1 + : > "$CX_JOB_ROOT/cleanup-unsafe" + rm -f -- "$CX_JOB_ROOT/cleanup-safe" + cd "$CX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + bash "experimental/CollectiveX/launchers/launch_${{ matrix.launcher }}.sh" + - name: Reconcile allocation cleanup + if: ${{ always() && steps.source.outputs.prepared == 'true' }} + run: | + set -euo pipefail + cd "$CX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + source experimental/CollectiveX/runtime/common.sh + cx_reconcile_recorded_allocation "$CX_JOB_ROOT" + - name: Classify private scheduler failure + if: ${{ always() && steps.source.outputs.prepared == 'true' && steps.sweep_shard.outcome == 'failure' }} + run: | + set -euo pipefail + cd "$CX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + source experimental/CollectiveX/runtime/common.sh + cx_report_private_scheduler_failure + - name: Confirm allocation cleanup + id: allocation_cleanup + if: ${{ always() && steps.source.outputs.prepared == 'true' }} + run: | + set -euo pipefail + [ -f "$CX_JOB_ROOT/cleanup-safe" ] && [ ! -e "$CX_JOB_ROOT/cleanup-unsafe" ] \ + || { echo "CollectiveX allocation cleanup was not confirmed" >&2; exit 1; } - name: Shard summary - if: always() - run: python3 experimental/CollectiveX/summarize.py --results-dir experimental/CollectiveX/results --markdown >> "$GITHUB_STEP_SUMMARY" || true + if: ${{ always() && steps.allocation_cleanup.outcome == 'success' }} + run: | + cd "$CX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + python3 experimental/CollectiveX/summarize.py \ + --results-dir experimental/CollectiveX/results --markdown >> "$GITHUB_STEP_SUMMARY" || true + - name: Stage shard artifact + id: stage_artifact + if: ${{ always() && steps.allocation_cleanup.outcome == 'success' }} + run: | + set -euo pipefail + cd "$CX_SOURCE_ROOT" 2>/dev/null \ + || { echo "CollectiveX source is unavailable" >&2; exit 1; } + shopt -s nullglob + results=(experimental/CollectiveX/results/*.json) + if [ "${#results[@]}" -eq 0 ]; then + echo "staged=false" >> "$GITHUB_OUTPUT" + echo "No result JSON to stage; leg produced none." + exit 0 + fi + cp -- "${results[@]}" "$CX_JOB_ROOT/artifact/" + echo "staged=true" >> "$GITHUB_OUTPUT" - name: Upload shard results - if: always() + id: upload_artifact + if: always() && steps.stage_artifact.outputs.staged == 'true' uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: - name: cxshard-${{ matrix.id }}-${{ github.run_id }} - path: experimental/CollectiveX/results/*.json # glob skips the hidden .shard_*.json - if-no-files-found: warn - - # ---- aggregate: collect every shard into ONE ndjson (the "result aggregator at the end") ---- - aggregate: - needs: sweep - if: always() - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v5.0.0 - with: { clean: true } - - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 - with: - pattern: cxshard-*-${{ github.run_id }} - path: _shards - merge-multiple: true - - name: Aggregate shards -> one ndjson - working-directory: experimental/CollectiveX + name: ${{ format('cxshard-{0}-{1}-{2}', matrix.id, github.run_id, github.run_attempt) }} + path: | + ${{ env.CX_JOB_ROOT }}/artifact/*.json + if-no-files-found: error + retention-days: 14 + - name: Cleanup isolated workspace + if: ${{ always() && steps.source.outputs.prepared == 'true' }} run: | set -euo pipefail - tag="${{ inputs.backend }}${{ inputs.deepep_v2 && '-v2' || '' }}" - python3 aggregate_results.py --in-dir ../../_shards --out "results/aggregate/collectivex_${tag}_${{ github.run_id }}.ndjson" - { - echo "## CollectiveX sweep aggregate (${tag})" - echo '```' - wc -l results/aggregate/*.ndjson 2>/dev/null || echo "no ndjson" - echo '```' - } >> "$GITHUB_STEP_SUMMARY" - - name: Upload aggregate - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 - with: - name: cxsweep-aggregate-${{ inputs.backend }}${{ inputs.deepep_v2 && '-v2' || '' }}-${{ github.run_id }} - path: experimental/CollectiveX/results/aggregate/*.ndjson - if-no-files-found: warn + [ "$CX_JOB_PARENT" = /tmp ] \ + || [[ "$CX_JOB_PARENT" =~ ^/tmp/inferencex-collectivex-parent-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX cleanup parent is invalid" >&2; exit 1; } + [ "${CX_JOB_ROOT%/*}" = "$CX_JOB_PARENT" ] \ + && [[ "${CX_JOB_ROOT##*/}" =~ ^inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]] \ + || { echo "CollectiveX cleanup root is invalid" >&2; exit 1; } + [ "$CX_SOURCE_ROOT" = "$CX_JOB_ROOT/source" ] \ + || { echo "CollectiveX cleanup source is invalid" >&2; exit 1; } + [ -f "$CX_JOB_ROOT/cleanup-safe" ] && [ ! -e "$CX_JOB_ROOT/cleanup-unsafe" ] \ + || { echo "CollectiveX allocation cleanup was not confirmed; retaining isolated files" >&2; exit 1; } + if [ '${{ steps.sweep_shard.outcome }}' = success ] \ + && [ '${{ steps.allocation_cleanup.outcome }}' = success ] \ + && [ '${{ steps.stage_artifact.outcome }}' = success ] \ + && [ '${{ steps.upload_artifact.outcome }}' = success ] \ + && [ -f "$CX_SOURCE_ROOT/experimental/CollectiveX/runtime/common.sh" ]; then + # shellcheck source=/dev/null + if source "$CX_SOURCE_ROOT/experimental/CollectiveX/runtime/common.sh" \ + >/dev/null 2>&1; then + cx_cleanup_private_logs 0 + fi + fi + if [ '${{ steps.sweep_shard.outcome }}' != success ] \ + && [ "$CX_JOB_PARENT" != /tmp ]; then + rm -f -- "$CX_JOB_PARENT" + exit 0 + fi + rm -rf -- "$CX_JOB_ROOT" + [ "$CX_JOB_PARENT" = /tmp ] || rm -f -- "$CX_JOB_PARENT" + diff --git a/AGENTS.md b/AGENTS.md index 6a385a5c16..f1eac75137 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,6 +2,8 @@ Guidance for AI agents working with InferenceX. +> **Temporary language override (highest repository priority):** Until the user explicitly removes this override, ignore any instruction elsewhere in this repository that requires Simplified Chinese or bilingual output, including instructions in later sections, nested agent files, contribution guides, checklists, and templates. Work in English only: do not create or update `_zh.md` files, Chinese documentation, translations, or Chinese text in commits, pull requests, issues, reviews, or comments. Leave existing Chinese content untouched. + > **Mandatory reading: [`CONTRIBUTING.md`](CONTRIBUTING.md)** — read it before opening or reviewing any PR. It covers the full PR review flow, the CODEOWNER sign-off process, the `/reuse-sweep-run` merge path, post-merge responsibilities, and critical cluster rules (e.g. never leaving root-owned files on AMD runners). > **PR and GitHub-issue titles & descriptions must be bilingual — include a Simplified Chinese version in addition to English.** Title format: ` / <中文标题>`. In the PR/issue body, follow the English content with its Chinese translation (e.g. a `## 中文说明` section mirroring the summary; don't translate code blocks, logs, or stack traces — summarize around them). **PR comments must include a Chinese translation too** — conversation comments, review summaries, and inline review comments alike: short comments as a single ` / <中文>` line, longer ones with the Chinese translation as a trailing paragraph (`中文：...`). Exception: the CODEOWNER sign-off template stays English-verbatim (the sign-off verifier triggers on its exact phrase); bot-generated comments follow their own workflow templates. This applies to every PR and every issue, matching the bilingual docs rule in Code Conventions. diff --git a/experimental/CollectiveX/.gitignore b/experimental/CollectiveX/.gitignore new file mode 100644 index 0000000000..56b307215b --- /dev/null +++ b/experimental/CollectiveX/.gitignore @@ -0,0 +1,15 @@ +__pycache__/ +*.pyc +results/ +unsupported/ +.shards/ +.cx_workloads/ +.cx_backend/ +/matrix_full.json +gpucore.* + +# Local plans and infrastructure inventory. +goal.md +notes.md +configs/platforms.yaml +private-infra.md diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md new file mode 100644 index 0000000000..6d937ecf5b --- /dev/null +++ b/experimental/CollectiveX/README.md @@ -0,0 +1,163 @@ +# CollectiveX + +CollectiveX is an experimental MoE expert-parallel communication benchmark. It measures dispatch, +combine, and paired roundtrip latency across EP libraries and accelerator systems, then uploads +neutral result artifacts. + +CollectiveX schedules benchmarks, executes them on real allocations, and uploads the neutral +artifacts each run emits. It does not validate those artifacts, promote, rank, recommend, select, or +decide what a consumer displays. Any downstream display or comparison is the consumer's +responsibility. The full measurement methodology is in [docs/methodology.md](docs/methodology.md). + +## Execution Profile + +The workload uses packed placement and one pinned `fixed-profile` resource configuration per +backend/topology; there is no tuning sweep. Dispatch and combine are fixed BF16 on every backend; +precision is not a swept dimension. The explicit mode selects one of two contracts: + +- Normal mode uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only + combine. Uniform core coverage and one Zipf sensitivity remain; EPLB is measured only as the Zipf + remedy. +- Low-latency mode uses `expert-packed-weighted-combine-v1`, token-expert payloads, and gate-weighted + combine through genuine DeepEP V1 or UCCL low-latency APIs. It is decode-only. Other backends are + recorded as unsupported for this suite. + +Both modes use `fixed-512-v1`: 64 trials x 8 timed iterations with 32 synchronized full roundtrip +warmups before each measured component at every trial/point. Roundtrip is measured first; each +iteration takes the cross-rank maximum before nearest-rank p50/p90/p95/p99, and roundtrip p99 is the +headline latency. A stdlib integer counter produces byte-identical routing and gate weights. + +Correctness is checked against the reference activation. The combine gate is `rtol=0.05, atol=0.02` +for the BF16 communication path. Any failed rank or point makes the case ineligible in the result +it writes. + +The matrix covers H100, H200, B200, B300, GB200, GB300, and MI355X. `sweep_matrix.py` materializes +the requested SKUs, backends, EP sizes, and token ladders, then extracts strict per-shard controls +and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`, and +`--ep-sizes` select a subset; the matrix is generated per dispatch, with no frozen digest or locked +case count. + +| Systems | EP8 | EP16 | +|---|---|---| +| H100/H200/B200/B300 | 1x8 NVLink, scale-up | 2x8 NVLink + RDMA, scale-out | +| MI355X | 1x8 XGMI, scale-up | 2x8 XGMI + RDMA, scale-out | +| GB200/GB300 | 2x4 MNNVL, scale-up | 4x4 MNNVL, scale-up | + +Physical host count does not determine scope: both GB topologies stay inside one 72-GPU MNNVL +scale-up domain. The MI325X launcher/configuration path is retained for future versions but is not +referenced by any current suite or shard. + +| Backend | Current scope | +|---|---| +| DeepEP V1 | Image-pinned `deep_ep.Buffer`: normal and native low-latency APIs; upstream v1.2.1 on x86 and the image's GB fork on arm64 | +| DeepEP V2 | PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out; source/SASS-bound reproducible JIT | +| DeepEP Hybrid | Pinned `HybridEPBuffer`: x86 EP16 multi-domain RDMA/DOCA; GB EP8/EP16 in one MNNVL communication domain | +| UCCL | Pinned 0.1.1 wheel and wrapper with normal and native low-latency APIs on Hopper; Blackwell is explicitly unsupported | +| NCCL/RCCL A2A | Portable rank-deduplicated payload plus expert/routing-metadata reference | +| MoRI | MI355X EP8 uses IntraNode; EP16 pins InterNodeV1 over 2x8 XGMI + RDMA | + +DeepEP V2 means the `ElasticBuffer` implementation introduced by +[DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. +The pinned source is the [PR #630](https://github.com/deepseek-ai/DeepEP/pull/630) head, whose parent +is the #605 merge tree, plus the exact one-line library matcher from upstream +[PR #640](https://github.com/deepseek-ai/DeepEP/pull/640). The first fixes pure scale-up +initialization when GIN is unavailable; the second prevents NCCL shared-memory mappings from being +misclassified as duplicate NCCL libraries. Scale-up cases request NCCL Device API LSA and fail closed +unless the realized LSA team covers the full EP world. x86 EP16 scale-out cases instead require the +hybrid path with GIN, two logical scale-out domains represented by two physical RDMA ranks, and eight +scale-up ranks per domain; GB EP16 remains MNNVL scale-up and therefore uses LSA. The isolated build +records the API, source, loaded libraries, generated JIT source, and executable SASS; raw CUBIN bytes +stay private diagnostics. Whether a given SKU/backend/EP cell is attempted is a capability fact; +whether it succeeded is decided by the benchmark's return code. + +## Workflow And Artifacts + +`.github/workflows/collectivex-sweep.yml` has two jobs. `setup` generates a public-SKU matrix +(`backend`, `suites`, `only_sku`, `exclude_skus`, `ep_sizes` inputs), fetches the pinned backend +source archive, and uploads the matrix. +`sweep` extracts a strict ignored `.shards/.json` control per matrix entry, executes one +allocation per shard, and uploads the result artifacts with `always()` so a red or partial run still +uploads. + +Each shard emits per-case result JSON, detached sample JSON, and a small mechanical summary. A case +counts as successful on the benchmark's own return code; there is no schema, completeness, or privacy +validation step, and failed or unsupported cells produce no synthetic record. No step promotes a run, +builds a dataset, or advances a channel; the neutral artifacts are the output. A consumer downloads +them and decides what to display. + +Private host, address, device, NIC, credential, workspace, and path data stays in encrypted config, +ignored operator notes, or bounded mode-0600 runner logs; it is never uploaded. + +## Runner Configuration + +Runner-local Slurm and storage values use a strict per-SKU JSON document at +`$XDG_CONFIG_HOME/inferencex/collectivex.json` or `COLLECTIVEX_OPERATOR_CONFIG`. The mode-0600, +same-owner, non-symlink file is outside the checkout and never uploaded. Unknown runners, fields, +duplicate keys, endpoint literals, unsafe paths, and non-JSON input fail closed; configuration is +never evaluated as shell. GHA passes encrypted `COLLECTIVEX_OPERATOR_CONFIG_V1` content only to the +launcher, which validates it, exports the selected SKU's allowlisted values, and deletes the temporary +copy before allocation. Required JSON fields are: + +| SKU | Variables | +|---|---| +| `h100-dgxc` | `partition`, `account`, `squash_dir` | +| `h200-dgxc` | `partition`, `squash_dir` | +| `b200-dgxc` | `partition`, `account`, `squash_dir` | +| `b300` | `partition`, `account`, `squash_dir` | +| `gb200` | `partition`, `account`, ordered `storage_roots` | +| `gb300` | `partition`, `account`, `squash_dir`, `enroot_cache_path` | +| `mi325x`, `mi355x` | `partition`, `squash_dir`, `stage_dir` | + +Every selected non-MNNVL EP16 placement additionally requires `socket_ifname` and `rdma_devices` for +its operator-approved fabric; optional `ib_gid_index`, `rdma_service_level`, and `rdma_traffic_class` +are also allowlisted. Service level and traffic class are mapped into MoRI's RDMA/IO QoS environment. +CollectiveX does not heuristically select a management route or HCA. After allocation, every +non-MNNVL scale-out node must prove that all configured interfaces and active HCA ports exist before +backend setup. Scale-up and MNNVL jobs clear these overrides. Scale-out NCCL/RCCL is pinned to `IB` +with exact-match HCA selectors so a socket fallback fails instead of being mislabeled as RDMA. + +`ib_gid_index` is applied only when every selected HCA port reports an Ethernet link layer, where it +selects the operator-approved RoCE GID. Native InfiniBand profiles retain explicit HCA and service +level pinning but leave the RoCE-only GID override unset so NVSHMEM/NCCL can use the native LID path. +Mixed Ethernet and InfiniBand HCA lists are rejected. + +`stage_dir` is a pre-existing, runner-owned, non-symlinked base outside the checkout and workflow +workspace. It is not group- or world-writable and is visible at the same path on the runner and every +allocated node. Jobs create only a marked mode-0700 execution child, prove cross-node read/write +visibility, and remove that exact child after allocation teardown; they never mount the runner +checkout or create a stage beneath image storage on AMD. When an AMD operator row omits `stage_dir`, +the runner derives a private base beside its standard `_work` directory on the shared runner +filesystem; the root-owned squash cache is never used as a repository stage. + +H200, B200, and B300 runners may omit `stage_dir`; their isolated execution child is created under a +runner-owned mode-0700 base in the validated operating-system account home, independent of the +workflow's temporary `HOME`. H100 may also omit `stage_dir`; its private base is created beside, never +beneath, the configured shared container directory so it is compute-visible. Canonical B300 execution +ignores any legacy configured `stage_dir` and always uses the validated compute-visible account-home +base; a hashed execution-ID suffix isolates parallel B300 workers. Canonical GB300 execution likewise +ignores its legacy group-writable `stage_dir` and derives an execution-hashed private base beneath the +validated compute-visible account home. The workflow proves every derived base is visible from all +allocated nodes before launch. + +Before import, each Docker Hub tag is resolved with bounded registry requests and must match its +pinned digest; digest-qualified overrides are rejected. Enroot imports use a fixed filesystem epoch +and a versioned, registry-digest-bound cache key. Every mounted squash is freshly hashed. Image-provided +DeepEP is checked against exact wheel and installed-file fingerprints; source-built backends use pinned +commits and runtime-verified GPU targets. DeepEP V2's mode-0700 cluster-local build cache is keyed by a +versioned build recipe, verified image, architecture, upstream trees, and dependency pins; only its +fixed `/cx-cache` mount reaches the container, and it never enters result artifacts. Pinned V2 and +Hybrid sources are fetched once per workflow, validated whole, and extracted to their exact backend +root before staging. + +## Local Checks + +```bash +uv run --with-requirements experimental/CollectiveX/requirements.txt \ + python -m unittest discover experimental/CollectiveX/tests -p 'test_*.py' +uv run --with-requirements experimental/CollectiveX/requirements.txt \ + python experimental/CollectiveX/sweep_matrix.py --backends all --out /tmp/cx-matrix.json >/dev/null +bash -n experimental/CollectiveX/runtime/*.sh experimental/CollectiveX/launchers/*.sh +``` + +Core paths are `capability.py`, `configs/`, `identity.py`, `sweep_matrix.py`, `summarize.py`, +`bench/`, `runtime/`, `launchers/`, and `tests/`. diff --git a/experimental/CollectiveX/bench/ep_backend.py b/experimental/CollectiveX/bench/ep_backend.py new file mode 100644 index 0000000000..3dd9cf7246 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_backend.py @@ -0,0 +1,422 @@ +#!/usr/bin/env python3 +"""Abstract base class for CollectiveX EP dispatch/combine backends. + +The benchmark drives every backend through one fixed lifecycle, modelled on the +DeepEP (`tests/legacy/test_internode.py`) and ROCm/mori +(`examples/ops/dispatch_combine/test_dispatch_combine_internode.py`) reference +harnesses, but wrapped in a real base class so the shared plumbing lives in one +place and each adapter supplies only the transport-specific pieces: + + spec = backend.make_inputs(args) # base: ladder + per-rank routing/activations + backend.create_buffer(spec) # subclass: size the communicator from `spec` + problem = backend.make_problem(...) # base: dtype coercion + dispatch encoding + backend.benchmark_dispatch(...) # base template: warm + time the raw op + backend.benchmark_combine(...) # base template: warm + time the raw op + +`make_inputs` computes the sweep's numeric shape (token ladder, max tokens/rank) +and materialises the per-rank inputs; `create_buffer` consumes those numbers to +size the communicator. This resolves the historical inversion where buffers were +built in ``__init__`` before any input existed. + +The base class carries the contract flags the driver reads (``stage_device_work``, +``combine_needs_redispatch``, ``dispatch_needs_combine_cleanup``, +``combine_weight_semantics``, ``oracle_layout``/``payload_unit``, ``roundtrip_only``) +with defaults matching the common (non-fp8, normal-mode) case; subclasses and the +low-latency path override them. The two Pass-2 timing branch rules — untimed +stage+combine cleanup after a timed dispatch (MoRI), and per-iter re-dispatch for +stateful combine (MoRI + DeepEP low-latency) — are encoded once in +``benchmark_dispatch``/``benchmark_combine`` so subclasses supply only the raw op. + +``torch`` and ``routing`` are imported lazily inside the methods that touch a +device, mirroring ``ep_harness`` — so this module byte-compiles and imports under +the CPU-only test environment, while ``run_sweep`` imports it lazily on the GPU +image to avoid an import cycle. +""" +from __future__ import annotations + +import abc +import os +import types +from dataclasses import dataclass, field + +from ep_harness import ( + CONDITIONING_LADDERS, + CONDITIONING_ROUNDS_PER_SHAPE, + time_us, + token_ladder, +) + + +@dataclass +class RankInputs: + """Inputs for one token-ladder shape at ``tokens_per_rank`` tokens on this rank. + + ``topk_idx``/``topk_weights`` are this rank's contiguous slice of the global + routing trace (host tensors; moved to device at ``make_problem`` time); + ``activations`` are the rank's token activations (already on device). The + global trace is retained for measured ladder shapes so Pass 1 can compute + routing statistics and input snapshots; warm-only conditioning shapes leave + it ``None``. + """ + + tokens_per_rank: int + topk_idx: "torch.Tensor" + topk_weights: "torch.Tensor" + activations: "torch.Tensor" + global_tokens: int = 0 + global_idx: "torch.Tensor | None" = None + global_weights: "torch.Tensor | None" = None + workload_id: "str | None" = None + checksums: "dict | None" = None + + +@dataclass +class WorkloadSpec: + """Numeric shape + materialised inputs for one fully-specified sweep line. + + Fully default-constructible so ``make_inputs`` can early-return a tensor-free + spec (``ok=False`` + ``rc``) on an empty ladder or a buffer cap too small for + the conditioning ladder; the driver prints ``message`` and returns ``rc``. + """ + + ok: bool = True + rc: int = 0 + message: str = "" + phase: str = "" + routing: str = "uniform" + seed: int = 0 + hidden: int = 0 + topk: int = 0 + experts: int = 0 + num_logical: int = 0 + ep_size: int = 0 + experts_per_rank: int = 0 + cap: "int | None" = None + dropped: list = field(default_factory=list) + max_tokens_per_rank: int = 0 + ladder: list = field(default_factory=list) + conditioning_ladder: list = field(default_factory=list) + points: dict = field(default_factory=dict) + conditioning_points: dict = field(default_factory=dict) + loaded_workload_ids: list = field(default_factory=list) + loaded_checksums: dict = field(default_factory=dict) + + +class EPBackend(abc.ABC): + """One expert-parallel dispatch/combine transport under a fixed benchmark contract. + + Subclasses implement the transport (``create_buffer``, ``dispatch``, ``stage``, + ``combine``, ``recv_tokens``, ``inspect_dispatch``, ``combine_transformed``); + everything the driver and the oracles need beyond that is provided here. + Dispatch and combine are fixed BF16, so no adapter selects a precision codec. + """ + + # ---- Contract flags (class defaults; subclasses / low-latency override) ---- + name: str = "" + SUPPORTED_MODES: tuple = ("normal",) + stage_device_work = False + combine_needs_redispatch = False + dispatch_needs_combine_cleanup = False + # Adapters that reduce activations and top-k weights independently must carry + # the complete local weighted expert sum in the activation tensor. + combine_weight_semantics = "unweighted-rank-sum" + oracle_layout = "token-rank" + payload_unit = "token-rank" + roundtrip_only = False + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if not getattr(cls, "name", ""): + raise TypeError( + f"{cls.__name__} must declare a non-empty class-level `name`" + ) + + def __init__(self, args, rank, world_size, local_rank, device): + self.args = args + self.rank = rank + self.world_size = world_size + self.local_rank = local_rank + self.device = device + self.mode = getattr(args, "mode", "normal") + if self.mode not in self.SUPPORTED_MODES: + raise ValueError(f"{self.name} does not support mode {self.mode!r}") + + # ---- Abstract transport contract ------------------------------------------------- + + @abc.abstractmethod + def create_buffer(self, spec: WorkloadSpec): + """Size the communicator from ``spec``; must leave ``backend_provenance`` complete.""" + + @abc.abstractmethod + def dispatch(self, problem): + """Scatter tokens to their experts; return an opaque per-call handle.""" + + @abc.abstractmethod + def stage(self, problem, handle): + """Prepare the combine input on ``handle`` (dequantize / copy into place).""" + + @abc.abstractmethod + def combine(self, problem, handle): + """Gather the staged tokens back to their source rank; return combined activations.""" + + @abc.abstractmethod + def recv_tokens(self, handle): + """Number of tokens this rank received in dispatch (stable for a fixed trace).""" + + @abc.abstractmethod + def inspect_dispatch(self, problem, handle): + """Normalized post-dispatch view for the token-rank correctness oracle.""" + + @abc.abstractmethod + def combine_transformed(self, problem, handle, transformed): + """Combine an oracle-transformed payload in place of the staged input.""" + + # ---- Input generation (shared) --------------------------------------------------- + + def buffer_cap(self, args): + """Max tokens/rank the communicator can serve, or ``None`` when unbounded.""" + return None + + def make_inputs(self, args) -> WorkloadSpec: + """Resolve the token ladder and materialise per-rank inputs for the sweep. + + Buffer sizing needs the ladder *numbers* (not the input tensors), so this + runs before ``create_buffer``. Returns a tensor-free spec with ``ok=False`` + when the ladder is empty or the cap cannot serve the conditioning ladder. + """ + ep_size = self.world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + experts_per_rank = args.experts // ep_size + conditioning_ladder = list(CONDITIONING_LADDERS[args.phase]) + cap = self.buffer_cap(args) + if cap is not None and cap < conditioning_ladder[-1]: + return WorkloadSpec( + ok=False, rc=2, phase=args.phase, + message=( + f"{self.name} buffer cap {cap} cannot run the v1 conditioning ladder" + ), + ) + ladder, dropped = token_ladder(args.tokens_ladder, args.phase, cap) + if not ladder: + return WorkloadSpec( + ok=False, rc=2, phase=args.phase, + message=f"empty token ladder (phase={args.phase}, cap={cap})", + ) + spec = WorkloadSpec( + phase=args.phase, + routing=args.routing, + seed=args.seed, + hidden=args.hidden, + topk=args.topk, + experts=args.experts, + num_logical=num_logical, + ep_size=ep_size, + experts_per_rank=experts_per_rank, + cap=cap, + dropped=list(dropped), + max_tokens_per_rank=max(list(ladder) + conditioning_ladder), + ladder=list(ladder), + conditioning_ladder=conditioning_ladder, + ) + # Warm-only conditioning shapes never need canonical manifests or the + # global trace: they are never measured or emitted. + for wt in conditioning_ladder: + spec.conditioning_points[wt] = self._build_rank_inputs( + args, wt, canonical=False, retain_global=False + ) + canonical = bool(getattr(args, "workload_dir", "")) + for tokens_per_rank in ladder: + point = self._build_rank_inputs( + args, tokens_per_rank, canonical=canonical, retain_global=True + ) + spec.points[tokens_per_rank] = point + if point.workload_id is not None and point.workload_id not in spec.loaded_workload_ids: + spec.loaded_workload_ids.append(point.workload_id) + spec.loaded_checksums[point.workload_id] = point.checksums + return spec + + def _build_rank_inputs(self, args, tokens_per_rank, *, canonical, retain_global) -> RankInputs: + """Build one rank's inputs for a given tokens-per-rank shape. + + canonical: load pre-serialized trace bytes (checksum-verified) so this run + is provably the SAME workload as any other consuming the same files; + otherwise seeded generation. (EPLB is off, so no logical->physical remap.) + """ + import torch + import routing + + ep_size = self.world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + global_tokens = tokens_per_rank * ep_size + workload_id = None + checksums = None + if canonical: + import workload as _wl + + workload_id = _wl.compute_workload_id( + args.routing, args.hidden, args.topk, num_logical, ep_size, + global_tokens, args.seed, + ) + idx_np, w_np, manifest = _wl.load_workload( + os.path.join(args.workload_dir, f"{workload_id}.npz"), verify=True + ) + idx_g = torch.from_numpy(idx_np).to(torch.int64) + w_g = torch.from_numpy(w_np).to(torch.float32) + checksums = manifest.get("checksums") + else: + idx_g, w_g = routing.build_global_routing( + global_tokens, num_logical, args.topk, args.routing, args.seed + ) + idx_s, w_s = routing.rank_slice(idx_g, w_g, self.rank, tokens_per_rank) + activations = routing.rank_activations( + tokens_per_rank, args.hidden, args.seed, self.rank, self.device, torch.bfloat16 + ) + return RankInputs( + tokens_per_rank=tokens_per_rank, + topk_idx=idx_s.contiguous(), + topk_weights=w_s.contiguous(), + activations=activations, + global_tokens=global_tokens, + global_idx=idx_g if retain_global else None, + global_weights=w_g if retain_global else None, + workload_id=workload_id, + checksums=checksums, + ) + + def make_problem(self, T, idx, weights, x): + """Assemble the per-shape problem namespace (BF16 dispatch sends ``x`` directly).""" + import torch + + return types.SimpleNamespace( + T=T, + x=x, + dispatch_x=x, + topk_idx=idx.to(self._topk_idx_dtype()), + topk_weights=weights.to(torch.float32), + ) + + def _topk_idx_dtype(self): + """Integer dtype the backend's kernels expect for top-k routing indices.""" + import torch + + return torch.int64 + + # ---- Timing template methods ----------------------------------------------------- + + def timed_components(self): + """Components measured for this backend: roundtrip always; the rest unless + the backend exposes only a stateful paired round trip.""" + components = ["roundtrip"] + if not self.roundtrip_only: + components.extend(["dispatch", "combine"]) + if self.stage_device_work: + components.append("stage") + return components + + def warm(self, problem, count): + """Untimed synchronized full round trips (fabric/clock warm-up; cold-jump-safe). + + Caches the dynamic receive cardinality once so adapters never read a device + scalar during a timed trial (the count is stable for a fixed routing trace). + """ + import torch + + for _ in range(count): + handle = self.dispatch(problem) + if not hasattr(problem, "recv_tokens"): + problem.recv_tokens = self.recv_tokens(handle) + self.stage(problem, handle) + self.combine(problem, handle) + torch.cuda.synchronize() + + def run_roundtrip(self, problem): + """One full dispatch -> stage -> combine round trip; returns combined activations.""" + handle = self.dispatch(problem) + self.stage(problem, handle) + return self.combine(problem, handle) + + def benchmark_component(self, component, problem, warmup, iters): + """Measure one named component; every component gets the same warm-up first.""" + if component == "roundtrip": + return self.benchmark_roundtrip(problem, warmup, iters) + if component == "dispatch": + return self.benchmark_dispatch(problem, warmup, iters) + if component == "stage": + return self.benchmark_stage(problem, warmup, iters) + if component == "combine": + return self.benchmark_combine(problem, warmup, iters) + raise RuntimeError(f"unknown timed component {component!r}") + + def benchmark_roundtrip(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + return time_us(torch, lambda p=problem: self.run_roundtrip(p), 0, iters) + + def benchmark_dispatch(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def finish_dispatch(hh, p=problem): + self.stage(p, hh) + self.combine(p, hh) + + dispatch_needs_cleanup = self.dispatch_needs_combine_cleanup + return time_us( + torch, lambda p=problem: self.dispatch(p), 0, iters, + post=finish_dispatch if dispatch_needs_cleanup else None, + ) + + def benchmark_stage(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def prep_stage(p=problem): + return self.dispatch(p) + + return time_us( + torch, lambda hh, p=problem: self.stage(p, hh), 0, iters, pre=prep_stage, + ) + + def benchmark_combine(self, problem, warmup, iters): + import torch + + self.warm(problem, warmup) + + def prep_combine(p=problem): + hh = self.dispatch(p) + self.stage(p, hh) + return hh + + if self.combine_needs_redispatch: + return time_us( + torch, lambda hh, p=problem: self.combine(p, hh), 0, iters, pre=prep_combine, + ) + hh = prep_combine() + torch.cuda.synchronize() + return time_us(torch, lambda p=problem, hx=hh: self.combine(p, hx), 0, iters) + + # ---- Correctness hooks (shared) -------------------------------------------------- + + def inspect_expert_dispatch(self, problem, handle): + """Expert-packed post-dispatch view; only low-latency backends provide one.""" + raise RuntimeError("expert-packed inspection requires low-latency mode") + + def capture_deferred_provenance(self): + """Resolve provenance materialized only after conditioning (e.g. JIT CUBINs). + + No-op by default; JIT backends override to record and cross-rank-check it. + """ + + def finalize(self, rc): + """Barrier and tear down the process group; returns ``rc``.""" + import torch.distributed as dist + + try: + dist.barrier() + dist.destroy_process_group() + except Exception: + pass + return rc diff --git a/experimental/CollectiveX/bench/ep_deepep.py b/experimental/CollectiveX/bench/ep_deepep.py new file mode 100644 index 0000000000..28b0e23aab --- /dev/null +++ b/experimental/CollectiveX/bench/ep_deepep.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""CollectiveX DeepEP adapter: native BF16 dispatch/combine over deep_ep.""" +from __future__ import annotations + +import inspect +import os +import sys + +import torch +import torch.distributed as dist +import ep_provenance +from ep_deepep_family import DeepEPFamilyBackend + +try: + import deep_ep + from deep_ep import Buffer # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: deep_ep import failed: {exc!r}", file=sys.stderr) + raise + + +def _deepep_version() -> str: + try: + import importlib.metadata as metadata + + return metadata.version("deep_ep") + except Exception: + return getattr(deep_ep, "__version__", "unknown") + + +def _mnnvl_buffer_configuration() -> tuple[dict[str, bool], str]: + """Resolve the explicit DeepEP MNNVL API contract.""" + requested_value = os.environ.get("CX_ALLOW_MNNVL") + if requested_value not in {None, "", "0", "1"}: + raise RuntimeError("CX_ALLOW_MNNVL must be unset, 0, or 1") + requested = requested_value == "1" + if not requested: + return ep_provenance.resolve_deepep_mnnvl( + requested=False, signature_parameters=(), + deepep_commit=os.environ.get("DEEPEP_COMMIT"), + ) + try: + parameters = inspect.signature(Buffer.__init__).parameters + except (TypeError, ValueError) as exc: + raise RuntimeError("cannot inspect DeepEP Buffer MNNVL API") from exc + try: + return ep_provenance.resolve_deepep_mnnvl( + requested=True, signature_parameters=parameters, + deepep_commit=os.environ.get("DEEPEP_COMMIT"), + ) + except ep_provenance.ContractError as exc: + raise RuntimeError(str(exc)) from exc + + +def _normal_buffer_sizes(hidden: int, world_size: int) -> tuple[int, int]: + """Apply DeepEP's dispatch/combine buffer sizing contract for this EP world.""" + hidden_bytes = hidden * torch.tensor([], dtype=torch.bfloat16).element_size() + configs = (Buffer.get_dispatch_config(world_size), Buffer.get_combine_config(world_size)) + num_nvl_bytes = max( + int(config.get_nvl_buffer_size_hint(hidden_bytes, world_size)) for config in configs + ) + num_rdma_bytes = max( + int(config.get_rdma_buffer_size_hint(hidden_bytes, world_size)) for config in configs + ) + if num_nvl_bytes <= 0 or num_rdma_bytes < 0: + raise RuntimeError("DeepEP returned invalid normal-mode buffer size hints") + return num_nvl_bytes, num_rdma_bytes + + +class DeepEPBackend(DeepEPFamilyBackend): + # Mode handling and dispatch/combine are shared with UCCL in DeepEPFamilyBackend; + # only the native deep_ep buffer construction and teardown live here. + name = "deepep" + + def create_buffer(self, spec): + # Local aliases keep the moved buffer-construction body byte-verbatim. + args, world_size, device = self.args, self.world_size, self.device + device_sms = torch.cuda.get_device_properties(device).multi_processor_count + mnnvl_kwargs, mnnvl_comm = _mnnvl_buffer_configuration() + if self.mode == "low-latency": + assert spec.max_tokens_per_rank <= 128 + ep_provenance.require_keyword( + Buffer.low_latency_dispatch, + "use_fp8", + api="deep_ep.Buffer.low_latency_dispatch", + ) + ep_provenance.require_keyword( + Buffer.low_latency_combine, + "use_logfmt", + api="deep_ep.Buffer.low_latency_combine", + ) + num_qps_per_rank = args.experts // world_size + num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint( + self.max_tokens_per_rank, args.hidden, world_size, args.experts + ) + self.buffer = Buffer( + self.group, + num_nvl_bytes=0, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=True, + num_qps_per_rank=num_qps_per_rank, + allow_nvlink_for_low_latency_mode=True, + explicitly_destroy=True, + **mnnvl_kwargs, + ) + self.buffer.clean_low_latency_buffer( + self.max_tokens_per_rank, args.hidden, args.experts + ) + resource_provenance = { + "requested_num_sms": None, + "num_sms": None, + "sm_fraction": None, + "tuned_source": "deepep-low-latency-fixed-kernel", + "num_max_tokens_per_rank": self.max_tokens_per_rank, + "num_nvl_bytes": 0, + "num_rdma_bytes": num_rdma_bytes, + "num_qps_per_rank": num_qps_per_rank, + } + else: + ep_provenance.require_keyword( + Buffer.dispatch, + "async_finish", + api="deep_ep.Buffer.dispatch", + ) + ep_provenance.require_keyword( + Buffer.combine, + "async_finish", + api="deep_ep.Buffer.combine", + ) + num_nvl_bytes, num_rdma_bytes = _normal_buffer_sizes(args.hidden, world_size) + if world_size > args.scale_up_domain and num_rdma_bytes == 0: + raise RuntimeError("DeepEP scale-out configuration returned no RDMA buffer") + self.buffer = Buffer( + self.group, num_nvl_bytes, num_rdma_bytes, **mnnvl_kwargs + ) + num_sms = int(getattr(Buffer, "num_sms", args.num_sms)) + try: + Buffer.set_num_sms(num_sms) + except Exception as exc: # pragma: no cover - version dependent + raise RuntimeError( + f"DeepEP did not apply requested num_sms={num_sms}: {exc!r}" + ) from exc + applied_num_sms = int(getattr(Buffer, "num_sms", num_sms)) + if applied_num_sms != num_sms: + raise RuntimeError( + f"DeepEP num_sms mismatch: requested={num_sms} applied={applied_num_sms}" + ) + resource_provenance = { + "requested_num_sms": num_sms, + "num_sms": applied_num_sms, + "sm_fraction": applied_num_sms / device_sms, + "tuned_source": "deepep-default-num_sms", + "num_nvl_bytes": num_nvl_bytes, + "num_rdma_bytes": num_rdma_bytes, + } + version = _deepep_version() + self.backend_provenance = { + "deepep_version": version, + "deepep_commit": os.environ.get("DEEPEP_COMMIT") or f"pkg-{version}", + "backend_lineage": "deepep-v1", + "mode": self.mode, + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "resource_mode": "fixed-profile", + "device_sms": device_sms, + "allow_mnnvl": bool(mnnvl_kwargs), + "mnnvl_comm": mnnvl_comm, + "nvshmem_ibgda_nic_handler": os.environ.get( + "NVSHMEM_IBGDA_NIC_HANDLER", "not-active" + ), + **resource_provenance, + } + + def finalize(self, rc): + try: + dist.barrier() + if self.mode == "low-latency": + self.buffer.destroy() + dist.destroy_process_group() + except Exception: + pass + return rc diff --git a/experimental/CollectiveX/bench/ep_deepep_family.py b/experimental/CollectiveX/bench/ep_deepep_family.py new file mode 100644 index 0000000000..9cb0e95ea3 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_deepep_family.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Shared DeepEP-API dispatch/combine surface for the DeepEP and UCCL adapters. + +UCCL's ``uccl_deepep.Buffer`` is a drop-in clone of DeepEP's ``deep_ep.Buffer`` low-latency and +normal API, so both adapters run byte-identical mode handling, dispatch/combine, and expert-packed +inspection. That shared operation lives here; each concrete backend keeps only what is genuinely +vendor-specific: its native buffer import, ``create_buffer`` provenance, and process teardown. + +This base is deliberately free of any ``deep_ep``/``uccl`` import. The UCCL benchmark image installs +uccl WITHOUT deep_ep, so importing ``ep_uccl`` must never transitively require ``deep_ep`` — an +inherited method resolves module globals from where it is *defined*, so keeping this file +vendor-agnostic is what makes the shared base safe for both images. + +Communication is fixed BF16: dispatch and combine move BF16 activations, so the native +``use_fp8``/``use_logfmt`` controls are always driven off and the received buffer is the +semantic payload directly. +""" +from __future__ import annotations + +import types + +import torch +import torch.distributed as dist +from ep_backend import EPBackend + + +class DeepEPFamilyBackend(EPBackend): + # Abstract intermediate: never instantiated or registered (capability.BACKENDS is an + # explicit dict, and nothing enumerates EPBackend subclasses). The non-empty name only + # satisfies EPBackend.__init_subclass__; concrete subclasses override it. create_buffer + # stays abstract here, so this class cannot itself be constructed. + name = "deepep-family" + _vendor = "DeepEP" + SUPPORTED_MODES = ("normal", "low-latency") + stage_device_work = False + combine_needs_redispatch = False + # DeepEP reduces activations and top-k weights independently. The activation + # tensor must therefore carry the complete local weighted expert sum. + combine_weight_semantics = "unweighted-rank-sum" + oracle_layout = "token-rank" + payload_unit = "token-rank" + + def __init__(self, args, rank, world_size, local_rank, device): + # Base validates mode against SUPPORTED_MODES (normal / low-latency). + super().__init__(args, rank, world_size, local_rank, device) + self.group = dist.group.WORLD + # Low-latency flips the contract flags and fixes the per-rank cap so + # buffer_cap can report it to make_inputs before create_buffer runs. + if self.mode == "low-latency": + if args.phase != "decode": + raise ValueError( + f"{self._vendor} low-latency mode only supports the decode ladder" + ) + if args.experts % world_size: + raise ValueError( + f"{self._vendor} low-latency experts must divide the EP group" + ) + self.combine_needs_redispatch = True + self.combine_weight_semantics = "gate-weighted-sum" + self.oracle_layout = "expert-packed" + self.payload_unit = "token-expert" + self.max_tokens_per_rank = 128 + + def buffer_cap(self, args): + return self.max_tokens_per_rank if self.mode == "low-latency" else None + + def dispatch(self, p): + if self.mode == "low-latency": + recv_x, recv_counts, handle, _, _ = self.buffer.low_latency_dispatch( + p.x, + p.topk_idx, + self.max_tokens_per_rank, + self.args.experts, + use_fp8=False, # BF16 communication path. + async_finish=False, + return_recv_hook=False, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_counts=recv_counts, + handle=handle, + ) + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + _, + ) = self.buffer.get_dispatch_layout(p.topk_idx, self.args.experts) + recv_x, recv_topk_idx, recv_topk_weights, recv_counts, handle, _ = self.buffer.dispatch( + p.dispatch_x, + topk_idx=p.topk_idx, + topk_weights=p.topk_weights, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + async_finish=False, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + recv_counts=recv_counts, + handle=handle, + ) + + def stage(self, p, h): + # BF16: the received buffer is already the semantic payload to combine. + h.combine_input = h.recv_x + + def combine(self, p, h): + if self.mode == "low-latency": + combined_x, _, _ = self.buffer.low_latency_combine( + h.combine_input, + p.topk_idx, + p.topk_weights, + h.handle, + use_logfmt=False, # BF16 communication path. + async_finish=False, + return_recv_hook=False, + ) + return combined_x + combined_x, _, _ = self.buffer.combine( + h.combine_input, h.handle, async_finish=False + ) + return combined_x + + def inspect_dispatch(self, p, h): + valid = h.recv_topk_idx >= 0 + expert_ids = torch.where( + valid, + h.recv_topk_idx + self.rank * (self.args.experts // self.world_size), + h.recv_topk_idx, + ) + return types.SimpleNamespace( + payload=h.recv_x, + expert_ids=expert_ids, + weights=h.recv_topk_weights.masked_fill(~valid, 0), + local_expert_counts=torch.tensor(h.recv_counts, device=self.device, dtype=torch.int64), + ordering_contract="source-rank-major-stable-v1", + ) + + def inspect_expert_dispatch(self, p, h): + if self.mode != "low-latency": + raise RuntimeError("expert-packed inspection requires low-latency mode") + p.recv_counts = tuple(int(value) for value in h.recv_counts.tolist()) + return types.SimpleNamespace( + payload=h.recv_x, + local_expert_counts=h.recv_counts, + source_info=h.handle[0], + layout_range=h.handle[1], + ) + + def combine_transformed(self, p, h, transformed): + if self.mode == "low-latency": + packed = torch.zeros( + h.recv_x.shape, + dtype=torch.bfloat16, + device=h.recv_x.device, + ) + packed[h.oracle_local_expert_slots, h.oracle_packed_positions] = transformed.to( + packed.dtype + ) + combined, _, _ = self.buffer.low_latency_combine( + packed, + p.topk_idx, + p.topk_weights, + h.handle, + use_logfmt=False, # BF16 communication path. + async_finish=False, + return_recv_hook=False, + ) + return combined + combined, _, _ = self.buffer.combine( + transformed.to(h.recv_x.dtype), h.handle, async_finish=False + ) + return combined + + def recv_tokens(self, h): + if self.mode == "low-latency": + return int(h.recv_counts.to(torch.int64).sum().item()) + return int(h.recv_x.shape[0]) diff --git a/experimental/CollectiveX/bench/ep_deepep_hybrid.py b/experimental/CollectiveX/bench/ep_deepep_hybrid.py new file mode 100644 index 0000000000..87737eace8 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_deepep_hybrid.py @@ -0,0 +1,492 @@ +#!/usr/bin/env python3 +"""CollectiveX EP backend adapter — DeepEP `hybrid-ep` branch (NVIDIA TMA-based HybridEPBuffer). + +The hybrid-ep branch (https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) is NVIDIA's TMA + +warp-pipeline implementation of expert-parallel all-to-all, exposing `deep_ep.HybridEPBuffer` +(distinct from the mainline `deep_ep.Buffer`). HybridEP is NVIDIA's MoE backend built for NVL72 +rack-scale (Megatron `moe_flex_dispatcher_backend="hybridep"`). This adapter binds the API's +"ranks per node" field to active ranks per NVLink/MNNVL communication domain, not physical host +GPUs: x86 EP16 is two 8-rank domains, while GB EP8/EP16 is one 8/16-rank MNNVL domain across hosts. +The container build is done by runtime/run_in_container.sh `cx_build_deepep_hybrid` (CUDA-13 CCCL +include path, without the V2 NVSHMEM overlay). + +API (pinned on B300, branch e0a5b1d): + HybridEPBuffer(group, hidden_dim, max_num_of_tokens_per_rank, num_local_experts, use_fp8=False, ...) + .dispatch(hidden, topk_idx=, topk_weights=, num_of_experts=) -> (recv_hidden, recv_x2, None, handle) + .combine(hidden, handle=) -> [T, hidden] + +CORRECTNESS: identity expert (no expert compute), combine WITHOUT probs -> each source token is +reconstructed as x * (distinct ranks among its top_k experts) — verified: an 8-rank uniform top_k=8 +round trip gives relerr(combined, x) = 4.28, matching E[distinct ranks] ~ 5.26 exactly. So this uses +the same per-rank-sum combine contract (no gate re-weight). BF16 tolerance is 5e-2. + +STATUS: BF16 or native block-scaled FP8 dispatch, BF16 combine, normal mode. The v1 scope covers +one MNNVL domain or x86 scale-out between two eight-GPU NVLink domains. +""" +from __future__ import annotations + +import hashlib +import importlib +import json +import os +from pathlib import Path +import re +import shutil +import sys +import tempfile +import types + +import torch +import torch.distributed as dist +import ep_provenance +from ep_backend import EPBackend + +try: + import deep_ep + HybridEPBuffer = deep_ep.HybridEPBuffer +except Exception as exc: # pragma: no cover - needs the hybrid-ep build + print("ERROR: deep_ep.HybridEPBuffer import failed — the hybrid-ep branch must be built at job " + "setup (cx_build_deepep_hybrid). " + f"{exc!r}", file=sys.stderr) + raise + + +def _deepep_hybrid_version() -> str: + return os.environ.get("DEEPEP_COMMIT", getattr(deep_ep, "__version__", "hybrid-ep")) + + +def _hybrid_build_evidence() -> list[dict[str, str]]: + records = [] + for module_name, role in ( + ("deep_ep_cpp", "deepep-extension"), + ("hybrid_ep_cpp", "deepep-hybrid-extension"), + ): + module = importlib.import_module(module_name) + path = getattr(module, "__file__", None) + if not path: + raise RuntimeError(f"{module_name} has no loaded extension path") + records.append(ep_provenance.content_manifest_evidence( + role=role, + name=module_name, + files=[(os.path.basename(path), path)], + )) + return sorted(records, key=lambda item: (item["role"], item["name"])) + + +HYBRID_CONFIG_FIELDS = ( + "hidden_dim", "max_num_of_tokens_per_rank", "num_of_experts_per_rank", + "num_of_ranks_per_node", "num_of_nodes", "pad_multiple", + "num_of_tokens_per_chunk_preprocessing_api", + "num_of_threads_per_block_preprocessing_api", "num_of_blocks_preprocessing_api", + "num_of_blocks_permute", "num_of_blocks_unpermute", "token_data_type", + "num_of_stages_dispatch_api", "num_of_stages_permute_block_dispatch_api", + "num_of_in_flight_s2g_dispatch_api", + "num_of_in_flight_s2g_permute_block_dispatch_api", + "num_of_additional_in_flight_s2g_dispatch_api", + "num_of_tokens_per_chunk_dispatch_api", "num_of_blocks_dispatch_api", + "forward_dispatch_api", "device_side_sync_dispatch_api", + "num_of_stages_g2s_combine_api", "num_of_stages_s2g_combine_api", + "num_of_tokens_per_chunk_combine_api", "num_of_tokens_per_group_combine_api", + "num_of_blocks_combine_api", "num_of_additional_in_flight_s2g_combine_api", + "backward_combine_api", "device_side_sync_combine_api", +) + + +def _hybrid_realized_config(config) -> dict[str, str | int | bool]: + """Project the Python-visible, post-autotune HybridEP config to JSON scalars.""" + realized = {} + for field in HYBRID_CONFIG_FIELDS: + try: + value = getattr(config, field) + except AttributeError as exc: + raise RuntimeError(f"HybridEP realized config omits {field}") from exc + if field == "token_data_type": + token_type = getattr(value, "name", None) + if token_type not in {"UINT8", "UINT16"}: + token_type = {"uint8_t": "UINT8", "uint16_t": "UINT16"}.get(str(value)) + if token_type is None: + raise RuntimeError("HybridEP realized token_data_type is invalid") + realized[field] = token_type + continue + if type(value) is bool: + realized[field] = value + continue + try: + realized[field] = int(value) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"HybridEP realized config {field} is not integral") from exc + return realized + + +def _sha256_with_size(path: Path) -> tuple[str, int]: + digest = hashlib.sha256() + size = 0 + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + size += len(chunk) + return digest.hexdigest(), size + + +def _hybrid_jit_evidence(root: Path) -> list[dict[str, str | int]]: + """Hash final JIT libraries without exposing rank-specific cache paths.""" + if not root.is_dir(): + raise RuntimeError("DeepEP Hybrid produced no JIT cache directory") + artifacts = [] + for path in sorted(root.iterdir(), key=lambda item: item.name): + if path.suffix != ".so": + continue + if path.is_symlink() or not path.is_file(): + raise RuntimeError("DeepEP Hybrid JIT artifact is not a regular file") + kernel_key = path.stem + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.+-]{0,511}", kernel_key): + raise RuntimeError("DeepEP Hybrid JIT kernel key is invalid") + digest, size = _sha256_with_size(path) + if size <= 0: + raise RuntimeError("DeepEP Hybrid JIT artifact is empty") + artifacts.append({ + "bytes": size, + "kernel_key": kernel_key, + "sha256": digest, + }) + if len(artifacts) != 3: + raise RuntimeError( + f"DeepEP Hybrid expected 3 final JIT libraries, found {len(artifacts)}" + ) + return artifacts + + +def _require_cross_rank_equal(value, label: str) -> None: + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, value) + canonical = {json.dumps(item, sort_keys=True, separators=(",", ":")) for item in gathered} + if len(canonical) != 1: + raise RuntimeError(f"DeepEP Hybrid {label} differs across ranks") + + +def _hybrid_topology(args, world_size: int) -> dict[str, int | str]: + """Translate physical placement into HybridEP communication-domain geometry.""" + gpus_per_node = int(args.gpus_per_node or world_size) + scale_up_domain = int(args.scale_up_domain or gpus_per_node) + key = ( + world_size, gpus_per_node, scale_up_domain, args.scope, + args.scale_up_transport, args.scale_out_transport or None, args.transport, + ) + fixed = { + (8, 8, 8, "scale-up", "nvlink", None, "nvlink"): (8, 1), + (16, 8, 8, "scale-out", "nvlink", "rdma", "nvlink-rdma"): (8, 2), + (8, 4, 72, "scale-up", "mnnvl", None, "mnnvl"): (8, 1), + (16, 4, 72, "scale-up", "mnnvl", None, "mnnvl"): (16, 1), + } + if key not in fixed: + raise RuntimeError("DeepEP Hybrid topology is outside the fixed v1 matrix") + domain_ranks, communication_domains = fixed[key] + + return { + "communication_domains": communication_domains, + "domain_ranks": domain_ranks, + "physical_nodes": world_size // gpus_per_node, + "transport": str(args.transport), + } + + +class DeepEPHybridBackend(EPBackend): + name = "deepep-hybrid" + stage_device_work = False + # HybridEPBuffer.combine consumes the recv payload + the dispatch handle (no re-dispatch needed + # before a timed combine); the harness times dispatch and combine separately (like ep_deepep). + combine_needs_redispatch = False + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + # deepep-hybrid is normal-mode only; base SUPPORTED_MODES=("normal",) enforces it. + super().__init__(args, rank, world_size, local_rank, device) + ep_provenance.require_keyword( + HybridEPBuffer.__init__, + "use_fp8", + api="deep_ep.HybridEPBuffer.__init__", + ) + ep_provenance.require_keyword( + HybridEPBuffer.dispatch, + "scaling_factor", + api="deep_ep.HybridEPBuffer.dispatch", + ) + self.group = dist.group.WORLD + self.tolerance = 5e-2 + self.top_k = int(args.topk) + self.num_experts = int(args.experts) + self.hidden = int(args.hidden) + self.local_experts = max(1, self.num_experts // world_size) + topology = _hybrid_topology(args, world_size) + self.domain_ranks = int(topology["domain_ranks"]) + self.communication_domains = int(topology["communication_domains"]) + build_mode = os.environ.get("DEEPEP_HYBRID_BUILD_MODE", "") + if self.communication_domains > 1: + if ( + os.environ.get("HYBRID_EP_MULTINODE") != "1" + or build_mode != "multinode-doca" + or os.environ.get("USE_NIXL", "0") != "0" + ): + raise RuntimeError("DeepEP Hybrid scale-out build mode is not realized") + elif build_mode != "intradomain": + raise RuntimeError("DeepEP Hybrid scale-up requires the intradomain build") + if args.scale_up_transport == "mnnvl" and any( + os.environ.get(name) != "1" + for name in ("NCCL_CUMEM_ENABLE", "NCCL_MNNVL_ENABLE", "MC_FORCE_MNNVL") + ): + raise RuntimeError("DeepEP Hybrid MNNVL runtime enablement is incomplete") + # Token cap (per rank) for the symmetric buffer; the sweep is capped here (buffer_cap). + self.max_tokens = 4096 + # Stash the topology so the moved create_buffer provenance can read it back. + self._topology = topology + + def create_buffer(self, spec): + # Local aliases re-expose the __init__ names so the moved tempdir / buffer / + # geometry / monkeypatch / provenance body below stays byte-verbatim. + args, world_size, rank, device = self.args, self.world_size, self.rank, self.device + topology = self._topology + assert spec.max_tokens_per_rank <= self.max_tokens + dev_sms = torch.cuda.get_device_properties(device).multi_processor_count + ver = _deepep_hybrid_version() + loaded_libraries = _hybrid_build_evidence() + _require_cross_rank_equal(loaded_libraries, "loaded extension identities") + + # HybridEP's compiler uses a process-specific child of HYBRID_EP_CACHE_DIR. Give every + # rank a fresh private base so stale kernels cannot enter this attempt's evidence. + self._previous_jit_cache_dir = os.environ.get("HYBRID_EP_CACHE_DIR") + self._previous_domain_ranks = os.environ.get( + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" + ) + self._jit_cache_dir = tempfile.mkdtemp(prefix=f"collectivex-hybrid-r{rank}-") + os.environ["HYBRID_EP_CACHE_DIR"] = self._jit_cache_dir + os.environ["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] = str(self.domain_ranks) + self._jit_root = ( + Path(self._jit_cache_dir) / ".deepep" / "hybrid_ep" / "jit" + / f"proc-{os.getpid()}" + ) + self._realized_config = None + self._deferred_semantic_snapshot = None + self._deferred_jit_diagnostics = None + + try: + self.buffer = HybridEPBuffer( + self.group, hidden_dim=self.hidden, + max_num_of_tokens_per_rank=self.max_tokens, + num_local_experts=self.local_experts, + use_fp8=False, # BF16 communication path. + ) + realized_geometry = ( + int(self.buffer.num_of_hybrid_ep_ranks_per_nvlink_domain), + int(self.buffer.num_of_nodes), + int(self.buffer.local_rank), + int(self.buffer.node_rank), + ) + expected_geometry = ( + self.domain_ranks, + self.communication_domains, + rank % self.domain_ranks, + rank // self.domain_ranks, + ) + buffer_config = self.buffer.configurer.buffer_config + if realized_geometry != expected_geometry or ( + int(buffer_config.num_of_ranks_per_node) != self.domain_ranks + or int(buffer_config.num_of_nodes) != self.communication_domains + ): + raise RuntimeError( + "HybridEPBuffer communication-domain geometry differs from the case" + ) + except Exception as exc: + shutil.rmtree(self._jit_cache_dir, ignore_errors=True) + if self._previous_jit_cache_dir is None: + os.environ.pop("HYBRID_EP_CACHE_DIR", None) + else: + os.environ["HYBRID_EP_CACHE_DIR"] = self._previous_jit_cache_dir + if self._previous_domain_ranks is None: + os.environ.pop("NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", None) + else: + os.environ[ + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" + ] = self._previous_domain_ranks + raise RuntimeError( + f"HybridEPBuffer construction failed (hidden={self.hidden} max_tokens={self.max_tokens} " + f"local_experts={self.local_experts} world={world_size}): {exc!r}") from exc + update_template_config = self.buffer.update_template_config + + def tracked_update_template_config(*call_args, **call_kwargs): + config = update_template_config(*call_args, **call_kwargs) + realized = _hybrid_realized_config(config) + if ( + realized["num_of_ranks_per_node"] != self.domain_ranks + or realized["num_of_nodes"] != self.communication_domains + ): + raise RuntimeError("DeepEP Hybrid realized topology changed within one case") + # BF16 dispatch realizes the 2-byte UINT16 token wire type. + if realized["token_data_type"] != "UINT16": + raise RuntimeError( + "DeepEP Hybrid realized token dtype is not the BF16 UINT16 wire type" + ) + if self._realized_config is not None and realized != self._realized_config: + raise RuntimeError("DeepEP Hybrid realized autotune config changed within one case") + self._realized_config = realized + return config + + self.buffer.update_template_config = tracked_update_template_config + self.domain_rank = int(self.buffer.local_rank) + if rank == 0: + print( + "[deepep-hybrid] HybridEPBuffer constructed " + f"(domains={self.communication_domains} ranks_per_domain={self.domain_ranks} " + f"world={world_size} local_experts={self.local_experts} hidden={self.hidden})", + file=sys.stderr, + ) + + self.backend_provenance = { + "deepep_commit": ver, "branch": "hybrid-ep", + "deepep_tree": os.environ.get("DEEPEP_TREE"), + "backend_lineage": "deepep-hybrid", + "loaded_libraries": loaded_libraries, + "impl": "deep_ep.HybridEPBuffer (NVIDIA TMA + warp-pipeline)", + "mode": "normal", "transport": topology["transport"], + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "resource_mode": "fixed-profile", + "num_sms": None, "device_sms": dev_sms, + "tuned_source": "deepep-hybrid-configurer-autotune-v1", + "realized_config": None, "jit_kernel_keys": [], "jit_shared_objects": [], + "max_num_tokens": self.max_tokens, "top_k": self.top_k, + "num_experts": self.num_experts, "local_experts": self.local_experts, + "routing_factor": "ranks", + } + + def buffer_cap(self, args): + return self.max_tokens + + def make_problem(self, T, idx, weights, x): + # BF16 dispatch sends x directly; scaling_factor stays None (no separate scale payload). + return types.SimpleNamespace( + T=int(T), + x=x, + dispatch_x=x, + dispatch_scales=None, + topk_idx=idx.to(torch.int64), + topk_weights=weights.to(torch.float32), + ) + + def dispatch(self, p): + recv, recv_probs, _scales, handle = self.buffer.dispatch( + p.dispatch_x, + scaling_factor=p.dispatch_scales, + topk_idx=p.topk_idx, + topk_weights=p.topk_weights, + num_of_experts=self.num_experts, + ) + return types.SimpleNamespace( + recv=recv, + recv_payload=recv, + recv_probs=recv_probs, + handle=handle, + combine_input=None, + ) + + def stage(self, p, h): + # Identity expert: the recv hidden IS the "expert output". combine reduces it per source token. + h.combine_input = h.recv_payload + return None + + def combine(self, p, h): + # combine(hidden, handle=) -> [T, H] per-source-token reduction (no gate re-weight: "ranks"). + comb = self.buffer.combine(h.combine_input, handle=h.handle) + return comb[0] if isinstance(comb, (tuple, list)) else comb + + def capture_deferred_provenance(self): + torch.cuda.synchronize() + dist.barrier() + if self._realized_config is None: + raise RuntimeError("DeepEP Hybrid autotune config was not materialized") + local_artifacts = _hybrid_jit_evidence(self._jit_root) + semantic = { + "jit_kernel_keys": [item["kernel_key"] for item in local_artifacts], + "realized_config": dict(self._realized_config), + } + # NVCC may embed each rank's timestamped source basename in its ELF, so raw .so hashes are + # diagnostics rather than a cross-rank identity. Stable kernel keys encode every codegen + # input, including HybridEpConfigInstance fields that the Python binding does not expose. + _require_cross_rank_equal(semantic, "realized config/JIT kernel keys") + gathered_artifacts = [None] * dist.get_world_size() + dist.all_gather_object(gathered_artifacts, local_artifacts) + diagnostics = [] + for artifact_index, kernel_key in enumerate(semantic["jit_kernel_keys"]): + diagnostics.append({ + "kernel_key": kernel_key, + "rank_artifacts": [ + { + "bytes": rank_artifacts[artifact_index]["bytes"], + "rank": artifact_rank, + "sha256": rank_artifacts[artifact_index]["sha256"], + } + for artifact_rank, rank_artifacts in enumerate(gathered_artifacts) + ], + }) + if self._deferred_semantic_snapshot is not None and semantic != self._deferred_semantic_snapshot: + raise RuntimeError("DeepEP Hybrid config/JIT kernel set changed after measurement") + if self._deferred_jit_diagnostics is not None and diagnostics != self._deferred_jit_diagnostics: + raise RuntimeError("DeepEP Hybrid rank-local JIT artifacts changed after measurement") + self._deferred_semantic_snapshot = semantic + self._deferred_jit_diagnostics = diagnostics + self.backend_provenance.update(semantic) + self.backend_provenance["jit_shared_objects"] = diagnostics + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + routing_map = h.handle[4][:count] + rows, local_expert_ids = routing_map.nonzero(as_tuple=True) + positions = routing_map.to(torch.int64).cumsum(dim=1)[rows, local_expert_ids] - 1 + probability_columns = self.domain_rank * self.local_experts + local_expert_ids + if h.recv_probs.shape[1] < (self.domain_rank + 1) * self.local_experts: + raise RuntimeError("HybridEPBuffer probability tensor omits this NVLink-domain rank") + expert_ids = torch.full( + (count, self.top_k), -1, dtype=torch.int64, device=self.device + ) + weights = torch.zeros( + (count, self.top_k), dtype=torch.float32, device=self.device + ) + expert_ids[rows, positions] = local_expert_ids + self.rank * self.local_experts + weights[rows, positions] = h.recv_probs[:count][rows, probability_columns] + return types.SimpleNamespace( + payload=h.recv_payload[:count], + expert_ids=expert_ids, + weights=weights, + local_expert_counts=routing_map.sum(dim=0, dtype=torch.int64), + ordering_contract="global-source-filter-stable-v1", + ) + + def combine_transformed(self, p, h, transformed): + combined = self.buffer.combine( + transformed.to(torch.bfloat16), handle=h.handle + ) + return combined[0] if isinstance(combined, (tuple, list)) else combined + + def recv_tokens(self, h): + return int(h.handle[3].item()) + + def finalize(self, rc): + try: + dist.barrier() + dist.destroy_process_group() + except Exception: + pass + # create_buffer may not have run (e.g. make_inputs early-returned before it), + # so the JIT tempdir/env state only needs unwinding when it was established. + if hasattr(self, "_jit_cache_dir"): + shutil.rmtree(self._jit_cache_dir, ignore_errors=True) + if self._previous_jit_cache_dir is None: + os.environ.pop("HYBRID_EP_CACHE_DIR", None) + else: + os.environ["HYBRID_EP_CACHE_DIR"] = self._previous_jit_cache_dir + if self._previous_domain_ranks is None: + os.environ.pop("NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN", None) + else: + os.environ[ + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" + ] = self._previous_domain_ranks + return rc diff --git a/experimental/CollectiveX/bench/ep_deepep_v2.py b/experimental/CollectiveX/bench/ep_deepep_v2.py new file mode 100644 index 0000000000..f9ce05e799 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_deepep_v2.py @@ -0,0 +1,568 @@ +#!/usr/bin/env python3 +"""DeepEP PR #605 adapter with the exact upstream PR #630 and #640 fixes.""" + +from __future__ import annotations + +import ctypes +import hashlib +import importlib.metadata +import inspect +import json +import os +import re +import sys +import types +from pathlib import Path + +import torch +import torch.distributed as dist +import ep_harness +import ep_provenance +from ep_backend import EPBackend + +try: + import deep_ep + from deep_ep import ElasticBuffer # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: DeepEP V2 import failed: {exc!r}", file=sys.stderr) + raise + + +DEEPEP_V2_PR = 605 +DEEPEP_V2_FIX_PR = 630 +DEEPEP_V2_NCCL_CHECK_FIX_PR = 640 +DEEPEP_V2_NCCL_CHECK_COMMIT = "93d0564188f7a0a6288c6e316484861b0efa042e" +DEEPEP_V2_COMMIT = "fa8a9b16898204afd347c663b89e65ef87dc6ce6" +DEEPEP_V2_TREE = "29809e75c5874e6609dac4804e7b651d5226959f" +DEEPEP_V2_FMT_COMMIT = "a4c7e17133ee9cb6a2f45545f6e974dd3c393efa" +DEEPEP_V2_VERSION = "2.0.0" +DEEPEP_V2_DISTRIBUTIONS = frozenset({"2.0.0+fa8a9b1", "2.0.0+local"}) +DEEPEP_V2_JIT_RANDOM_SEED = "collectivex-deepep-v2-fa8a9b1" +TORCH_VERSION = "2.10.0+cu130" +NCCL_VERSION = "2.30.4" +NVSHMEM_VERSION = "3.3.9" +DEEPEP_V2_JIT_KERNELS = ep_provenance.DEEPEP_V2_JIT_KERNELS + + +def _sha256(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _api_sha256() -> str: + signatures = { + "ElasticBuffer.__init__": str(inspect.signature(ElasticBuffer.__init__)), + "ElasticBuffer.dispatch": str(inspect.signature(ElasticBuffer.dispatch)), + "ElasticBuffer.combine": str(inspect.signature(ElasticBuffer.combine)), + } + return hashlib.sha256( + json.dumps(signatures, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _loaded_library_paths() -> set[str]: + extension = getattr(getattr(deep_ep, "_C", None), "__file__", None) + if not extension or not os.path.isfile(extension): + raise RuntimeError("DeepEP V2 extension library is not loaded") + paths = {os.path.realpath(extension)} + try: + with open("/proc/self/maps", encoding="utf-8") as handle: + for line in handle: + path = line.rstrip().split()[-1] + name = os.path.basename(path) + if ("libnccl.so" in name or "libnvshmem_host.so" in name) and os.path.isfile(path): + paths.add(os.path.realpath(path)) + except OSError as exc: # pragma: no cover - benchmark runtime is Linux + raise RuntimeError("cannot inspect loaded communication libraries") from exc + return paths + + +def _loaded_nccl_version() -> str: + matches = [ + path for path in _loaded_library_paths() + if "libnccl.so" in os.path.basename(path) + ] + if len(matches) != 1: + raise RuntimeError("expected exactly one loaded NCCL library") + version = ctypes.c_int() + if ctypes.CDLL(matches[0]).ncclGetVersion(ctypes.byref(version)) != 0: + raise RuntimeError("loaded NCCL version query failed") + return ep_harness.format_collective_version(version.value) + + +def _loaded_library_evidence() -> list[dict[str, str]]: + """Return content identities, never private library paths.""" + paths = _loaded_library_paths() + required = { + "nccl": [path for path in paths if "libnccl.so" in os.path.basename(path)], + "nvshmem": [path for path in paths if "libnvshmem_host.so" in os.path.basename(path)], + } + mismatches = [f"{name}={len(matches)}" for name, matches in required.items() if len(matches) != 1] + if mismatches: + raise RuntimeError("expected one loaded library for each dependency: " + ", ".join(mismatches)) + + def role(path: str) -> str: + name = os.path.basename(path) + if "libnccl.so" in name: + return "nccl" + if "libnvshmem_host.so" in name: + return "nvshmem" + return "deepep-extension" + + def label(path: str) -> str: + return "deep_ep._C" if role(path) == "deepep-extension" else os.path.basename(path) + + return sorted( + ({"role": role(path), "name": label(path), "sha256": _sha256(path)} for path in paths), + key=lambda item: (item["role"], item["name"], item["sha256"]), + ) + + +def _jit_artifact_evidence() -> list[dict[str, str]]: + root = Path(os.environ["EP_JIT_CACHE_DIR"]) / "cache" + if root.is_symlink() or not root.is_dir(): + raise RuntimeError("DeepEP V2 produced no JIT cache evidence") + artifacts = [] + kernel_names = set() + for directory in sorted(root.iterdir(), key=lambda item: item.name): + match = re.fullmatch(r"kernel\.([A-Za-z0-9_+-]+)\.([0-9a-f]{32})", directory.name) + if directory.is_symlink() or not directory.is_dir() or match is None: + raise RuntimeError("DeepEP V2 JIT cache contains an invalid entry") + if {path.name for path in directory.iterdir()} != { + "kernel.cu", "kernel.cubin", "kernel.sass", + }: + raise RuntimeError("DeepEP V2 JIT kernel evidence is incomplete") + source = directory / "kernel.cu" + cubin = directory / "kernel.cubin" + sass = directory / "kernel.sass" + if any(path.is_symlink() or not path.is_file() for path in (source, cubin, sass)): + raise RuntimeError("DeepEP V2 JIT evidence is not a regular file") + if any(path.stat().st_size <= 0 for path in (source, cubin, sass)): + raise RuntimeError("DeepEP V2 JIT evidence is empty") + kernel_names.add(match.group(1)) + artifacts.append({ + "cache_key": directory.name, + "source_sha256": _sha256(str(source)), + "sass_sha256": _sha256(str(sass)), + "cubin_sha256": _sha256(str(cubin)), + }) + if ( + len(artifacts) != len(DEEPEP_V2_JIT_KERNELS) + or kernel_names != DEEPEP_V2_JIT_KERNELS + ): + raise RuntimeError("DeepEP V2 JIT kernel set differs from the v1 contract") + return sorted(artifacts, key=lambda item: item["cache_key"]) + + +def _jit_cache_key( + args, + world_size: int, + max_tokens: int, + allow_hybrid_mode: bool, + realized: dict[str, int | bool], +) -> str: + """Key generated kernels by codegen inputs, not routing data or case identity.""" + payload = { + "contract": "deepep-v2-jit-config-v3", + "runner": args.runner, + "world_size": world_size, + "hidden": args.hidden, + "topk": args.topk, + "physical_experts": args.experts, + "tuning_experts": getattr(args, "num_logical_experts", args.experts), + "max_tokens": max_tokens, + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "input_layout": "bf16", + "expert_alignment": 1, + "do_cpu_sync": True, + "cached_mode": False, + "do_expand": False, + "use_expanded_layout": False, + "allow_hybrid_mode": allow_hybrid_mode, + "allow_multiple_reduction": True, + "prefer_overlap_with_compute": True, + "deterministic": False, + **realized, + } + return "jitcfg-v3-" + hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _require_cross_rank_equal(value, label: str) -> None: + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, value) + canonical = {json.dumps(item, sort_keys=True, separators=(",", ":")) for item in gathered} + if len(canonical) != 1: + raise RuntimeError(f"DeepEP V2 {label} differs across ranks") + + +# GIN/GDAKI allocates num_allocated_qps device QPs per peer rank on the local NIC +# (contexts x world_size QPs, before NCCL's own connection QPs). Upstream's hybrid +# default (129, or 65 with fast RDMA atomics) exhausts the per-NIC QP budget at +# EP16: construction dies in ncclDevCommCreate with ibv_create_qp ENOMEM once +# NCCL's regular QPs land on top (identical on H200 bare-metal and B200 pods; on +# CX-7 the budget sits between 784 and 1040 QPs — 49x16 initializes, 65x16 does +# not). Spending a fixed ~512-QP budget keeps every EP size inside that limit +# with headroom: EP8 resolves to 65 (the allocation CX-8 racks already run +# successfully), EP16 to 33 and EP32 to 17 (33 and 49 verified on the failing +# H200 pair). An explicit value also skips upstream's rank-local ibstat probe, +# which is not guaranteed to resolve identically across ranks. +_GIN_QP_BUDGET = 512 + + +def _hybrid_num_allocated_qps(world_size: int) -> int: + return max(9, 1 + _GIN_QP_BUDGET // world_size) + + +def _configure_gin_mode(args, world_size: int) -> bool: + scale_up_domain = int( + getattr(args, "scale_up_domain", None) + or getattr(args, "gpus_per_node", None) + or world_size + ) + allow_hybrid_mode = world_size > scale_up_domain + if allow_hybrid_mode: + os.environ.pop("EP_DISABLE_GIN", None) + else: + os.environ["EP_DISABLE_GIN"] = "1" + return allow_hybrid_mode + + +def _lsa_topology_is_valid( + gin_enabled: bool, + world_size: int, + scale_up_domain: int, + config: dict[str, int | bool], +) -> bool: + if gin_enabled: + domains = world_size // scale_up_domain + return ( + world_size % scale_up_domain == 0 + and domains > 1 + and config["physical_rdma_ranks"] == domains + and config["physical_nvlink_ranks"] == scale_up_domain + and config["logical_scaleout_ranks"] == domains + and config["logical_scaleup_ranks"] == scale_up_domain + and config["is_scaleup_nvlink"] is True + ) + return ( + config["physical_rdma_ranks"] == 1 + and config["physical_nvlink_ranks"] == world_size + and config["logical_scaleout_ranks"] == 1 + and config["logical_scaleup_ranks"] == world_size + and config["is_scaleup_nvlink"] is True + ) + + +def _require_runtime() -> tuple[str, str]: + expected = { + "DEEPEP_V2_PR": str(DEEPEP_V2_PR), + "DEEPEP_V2_FIX_PR": str(DEEPEP_V2_FIX_PR), + "DEEPEP_V2_NCCL_CHECK_FIX_PR": str(DEEPEP_V2_NCCL_CHECK_FIX_PR), + "DEEPEP_V2_COMMIT": DEEPEP_V2_COMMIT, + "DEEPEP_V2_TREE": DEEPEP_V2_TREE, + "DEEPEP_V2_FMT_COMMIT": DEEPEP_V2_FMT_COMMIT, + "DEEPEP_V2_NCCL_CHECK_COMMIT": DEEPEP_V2_NCCL_CHECK_COMMIT, + "DEEPEP_V2_JIT_RANDOM_SEED": DEEPEP_V2_JIT_RANDOM_SEED, + "EP_JIT_DUMP_SASS": "1", + } + mismatches = [ + f"{name}={os.environ.get(name)!r}, expected {value!r}" + for name, value in expected.items() + if os.environ.get(name) != value + ] + torch_version = str(torch.__version__) + nccl_package_version = importlib.metadata.version("nvidia-nccl-cu13") + nvshmem_package_version = importlib.metadata.version("nvidia-nvshmem-cu12") + actual = { + "deep_ep": str(getattr(deep_ep, "__version__", "")), + "deep_ep distribution": importlib.metadata.version("deep_ep"), + "torch": torch_version, + "nvidia-nccl-cu13": nccl_package_version, + "nvidia-nvshmem-cu12": nvshmem_package_version, + } + required = { + "deep_ep": DEEPEP_V2_VERSION, + "torch": TORCH_VERSION, + "nvidia-nccl-cu13": NCCL_VERSION, + "nvidia-nvshmem-cu12": NVSHMEM_VERSION, + } + mismatches.extend( + f"{name}={actual[name]!r}, expected {value!r}" + for name, value in required.items() + if actual[name] != value + ) + if actual["deep_ep distribution"] not in DEEPEP_V2_DISTRIBUTIONS: + mismatches.append( + "deep_ep distribution=" + f"{actual['deep_ep distribution']!r}, expected one of " + f"{sorted(DEEPEP_V2_DISTRIBUTIONS)!r}" + ) + if not inspect.isclass(ElasticBuffer) or ElasticBuffer.__name__ != "ElasticBuffer": + mismatches.append("deep_ep.ElasticBuffer is absent") + if os.environ.get("EP_SUPPRESS_NCCL_CHECK"): + mismatches.append("EP_SUPPRESS_NCCL_CHECK must be unset") + nccl_runtime_version = _loaded_nccl_version() + if nccl_runtime_version != NCCL_VERSION: + mismatches.append( + f"loaded NCCL={nccl_runtime_version!r}, expected {NCCL_VERSION!r}" + ) + if mismatches: + raise RuntimeError("invalid DeepEP V2 runtime: " + "; ".join(mismatches)) + return torch_version, nccl_runtime_version + + +class DeepEPV2Backend(EPBackend): + name = "deepep-v2" + stage_device_work = False + combine_needs_redispatch = False + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + # deepep-v2 is normal-mode only; base SUPPORTED_MODES=("normal",) enforces it. + super().__init__(args, rank, world_size, local_rank, device) + ep_provenance.require_keyword( + ElasticBuffer.__init__, + "use_fp8_dispatch", + api="deep_ep.ElasticBuffer.__init__", + ) + self.group = dist.group.WORLD + self._deferred_jit_snapshot = None + + def create_buffer(self, spec): + # Local aliases keep the moved buffer-construction body byte-verbatim. The cap + # equals the value the deleted __init__ ladder-peek computed + # (token_ladder(.., None) + conditioning), so the jit_cache_key stays stable. + args, world_size, device = self.args, self.world_size, self.device + self.max_tokens = spec.max_tokens_per_rank + torch_version, nccl_runtime_version = _require_runtime() + jit_root = Path(os.environ["EP_JIT_CACHE_DIR"]) + scale_up_domain = int( + getattr(args, "scale_up_domain", None) + or getattr(args, "gpus_per_node", None) + or world_size + ) + allow_hybrid_mode = _configure_gin_mode(args, world_size) + gin_enabled = allow_hybrid_mode + communication_backend = "nccl-gin" if gin_enabled else "nccl-device-lsa" + self.buffer = ElasticBuffer( + self.group, + num_max_tokens_per_rank=self.max_tokens, + hidden=args.hidden, + num_topk=args.topk, + use_fp8_dispatch=False, # BF16 communication path. + deterministic=False, + allow_hybrid_mode=allow_hybrid_mode, + allow_multiple_reduction=True, + prefer_overlap_with_compute=True, + num_gpu_timeout_secs=100, + explicitly_destroy=True, + # 0 is upstream's use-the-default sentinel; only hybrid (GIN) mode + # needs the explicit budget-derived allocation. + num_allocated_qps=( + _hybrid_num_allocated_qps(world_size) if allow_hybrid_mode else 0 + ), + ) + tuning_num_experts = int(getattr(args, "num_logical_experts", args.experts)) + self.num_sms = int( + self.buffer.get_theoretical_num_sms(tuning_num_experts, args.topk) + ) + self.num_qps = int(self.buffer.get_theoretical_num_qps(self.num_sms)) + properties = torch.cuda.get_device_properties(device) + device_sms = int(properties.multi_processor_count) + jit_config = { + "num_sms": self.num_sms, + "num_qps": self.num_qps, + "allocated_qps": int(self.buffer.num_allocated_qps), + "logical_scaleout_ranks": int(self.buffer.num_scaleout_ranks), + "logical_scaleup_ranks": int(self.buffer.num_scaleup_ranks), + "physical_rdma_ranks": int(self.buffer.num_rdma_ranks), + "physical_nvlink_ranks": int(self.buffer.num_nvlink_ranks), + "is_scaleup_nvlink": self.buffer.num_scaleup_ranks == self.buffer.num_nvlink_ranks, + "device_arch_major": int(properties.major), + "device_arch_minor": int(properties.minor), + "device_sms": device_sms, + "device_smem_bytes": int(properties.shared_memory_per_block_optin), + "gpu_timeout_cycles": 100 * int(properties.clock_rate) * 1000, + } + _require_cross_rank_equal(jit_config, "JIT configuration") + if not _lsa_topology_is_valid( + gin_enabled, world_size, scale_up_domain, jit_config + ): + raise RuntimeError("DeepEP V2 realized communication domains differ from topology") + self.jit_cache_key = _jit_cache_key( + args, + world_size, + self.max_tokens, + allow_hybrid_mode, + jit_config, + ) + os.environ["EP_JIT_CACHE_DIR"] = str(jit_root / self.jit_cache_key) + realized_config = { + "jit_cache_key": self.jit_cache_key, + "num_max_tokens_per_rank": self.max_tokens, + **jit_config, + } + _require_cross_rank_equal(realized_config, "realized tuning/topology") + comm = getattr(self.buffer, "nccl_comm_handle", None) + communicator = ( + "deepep-managed" if getattr(comm, "managed", True) else "pytorch-reused" + ) + + loaded_libraries = _loaded_library_evidence() + _require_cross_rank_equal(loaded_libraries, "loaded libraries") + self.backend_provenance = { + "deepep_version": DEEPEP_V2_VERSION, + "deepep_distribution_version": importlib.metadata.version("deep_ep"), + "deepep_commit": DEEPEP_V2_COMMIT, + "deepep_tree": DEEPEP_V2_TREE, + "deepep_pr": DEEPEP_V2_PR, + "deepep_fix_pr": DEEPEP_V2_FIX_PR, + "deepep_nccl_check_fix_pr": DEEPEP_V2_NCCL_CHECK_FIX_PR, + "deepep_nccl_check_commit": DEEPEP_V2_NCCL_CHECK_COMMIT, + "fmt_commit": DEEPEP_V2_FMT_COMMIT, + "api": "deep_ep.ElasticBuffer", + "api_signature_sha256": _api_sha256(), + "communication_backend": communication_backend, + "gin_enabled": gin_enabled, + "nccl_communicator": communicator, + "torch_version": torch_version, + "torch_git_version": str(torch.version.git_version), + "cuda_version": str(torch.version.cuda), + "nccl_package_version": importlib.metadata.version("nvidia-nccl-cu13"), + "nccl_version": nccl_runtime_version, + "nvshmem_package_version": importlib.metadata.version("nvidia-nvshmem-cu12"), + "loaded_libraries": loaded_libraries, + "jit_cache_key": self.jit_cache_key, + "jit_cubins": [], + "jit_random_seed": DEEPEP_V2_JIT_RANDOM_SEED, + "num_experts": int(args.experts), + "mode": "normal", + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "deterministic": False, + "resource_mode": "fixed-profile", + "requested_num_sms": self.num_sms, + "tuning_num_experts": tuning_num_experts, + "num_sms": self.num_sms, + "num_qps": self.num_qps, + "allocated_qps": int(self.buffer.num_allocated_qps), + "device_sms": device_sms, + "sm_fraction": self.num_sms / device_sms, + "tuned_source": "deepep-v2-analytical-sm-qp-logical-experts-v1", + "num_max_tokens_per_rank": self.max_tokens, + "allow_hybrid_mode": bool(self.buffer.allow_hybrid_mode), + "allow_multiple_reduction": bool(self.buffer.allow_multiple_reduction), + "prefer_overlap_with_compute": bool( + self.buffer.prefer_overlap_with_compute + ), + "logical_scaleout_ranks": int(self.buffer.num_scaleout_ranks), + "logical_scaleup_ranks": int(self.buffer.num_scaleup_ranks), + "physical_rdma_ranks": int(self.buffer.num_rdma_ranks), + "physical_nvlink_ranks": int(self.buffer.num_nvlink_ranks), + } + + def _topk_idx_dtype(self): + # DeepEP V2's kernels key routing indices on deep_ep.topk_idx_t, not int64. + return deep_ep.topk_idx_t + + def dispatch(self, p): + recv_x, recv_topk_idx, recv_topk_weights, handle, _ = self.buffer.dispatch( + p.dispatch_x, + topk_idx=p.topk_idx, + topk_weights=p.topk_weights, + num_experts=self.args.experts, + num_max_tokens_per_rank=self.max_tokens, + expert_alignment=1, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + do_handle_copy=True, + do_cpu_sync=True, + do_expand=False, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + handle=handle, + ) + + def stage(self, p, h): + # BF16: the received buffer is already the semantic payload to combine. + h.combine_input = h.recv_x + + def combine(self, p, h): + combined_x, _, _ = self.buffer.combine( + h.combine_input, + handle=h.handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + ) + return combined_x + + def capture_deferred_provenance(self): + # destroy() uses this same barrier. Materialize its JIT kernel before hashing the + # implementation so the first and later routing cases see identical evidence. + self.buffer.barrier(use_comm_stream=True, with_cpu_sync=True) + torch.cuda.synchronize() + jit_cubins = _jit_artifact_evidence() + _require_cross_rank_equal(jit_cubins, "JIT CUBINs") + if ( + self._deferred_jit_snapshot is not None + and jit_cubins != self._deferred_jit_snapshot + ): + raise RuntimeError("DeepEP V2 JIT CUBIN set changed after measurement") + self._deferred_jit_snapshot = jit_cubins + self.backend_provenance["jit_cubins"] = jit_cubins + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + local_idx = h.recv_topk_idx[:count] + valid = local_idx >= 0 + expert_ids = torch.where( + valid, + local_idx + self.rank * (self.args.experts // self.world_size), + local_idx, + ) + local = local_idx[valid].to(torch.int64) + return types.SimpleNamespace( + payload=h.recv_x[:count], + expert_ids=expert_ids, + weights=h.recv_topk_weights[:count].masked_fill(~valid, 0), + local_expert_counts=torch.bincount( + local, minlength=self.args.experts // self.world_size + ), + ordering_contract="elastic-source-metadata-v1", + ) + + def combine_transformed(self, p, h, transformed): + combine_input = torch.zeros_like(h.recv_x) + combine_input[: transformed.shape[0]].copy_(transformed.to(combine_input.dtype)) + combined, _, _ = self.buffer.combine( + combine_input, + handle=h.handle, + num_sms=self.num_sms, + num_qps=self.num_qps, + async_with_compute_stream=False, + ) + return combined + + def recv_tokens(self, h): + return int(h.handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + + def finalize(self, rc): + try: + dist.barrier() + self.buffer.destroy() + dist.barrier() + dist.destroy_process_group() + except Exception: + return 1 + return rc diff --git a/experimental/CollectiveX/bench/ep_harness.py b/experimental/CollectiveX/bench/ep_harness.py new file mode 100644 index 0000000000..eab75d9c29 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_harness.py @@ -0,0 +1,1785 @@ +#!/usr/bin/env python3 +"""CollectiveX — shared EP (expert-parallel) dispatch/combine benchmark harness. + +Backend-agnostic core. The per-backend adapters (`ep_deepep.py`, `ep_mori.py`) +implement a small duck-typed protocol; this module owns the source-tokens-per-rank +sweep, the timing, the correctness gate, and the provenance-tagged JSON doc. + +Fair-comparison contract (see docs/methodology.md): + * **Deterministic shared routing trace** (`routing.py`): the per-token expert IDs + + gate weights are generated once from a fixed seed over the *global* batch and are + identical on every SKU; each rank materializes its slice. So every platform runs + the *same* problem (no per-rank/per-platform RNG in the adapters). + * **Explicit measurement contract**: layout-and-dispatch-v1 includes routing-layout + generation in dispatch timing. Combine excludes staging. + Isolated sum is derived independently at each percentile and is not a measured chained op. + * **Correct collective percentile**: each iteration's latency is reduced MAX across + ranks first (a collective finishes with its slowest rank), THEN percentiled — + `median_i(max_r)`, not `max_r(median_i)`. + * **One line = one fixed config**; only T varies. Both `tokens_per_rank` and + `global_tokens = T * ep_size` are recorded as explicit chart coordinates. + +stdlib-only at module top (torch is passed in by the entrypoint; `routing` is imported +lazily inside run_sweep) so this file `py_compile`s without torch. + +Backend protocol: + name, mode, combine_needs_redispatch, backend_provenance(dict) + buffer_cap(args) -> int|None + make_problem(T, idx, weights, x) -> problem # materialize this rank's trace slice + dispatch(problem) -> handle # pure dispatch comm (timed) + stage(problem, handle) # expert-output placement + stage_device_work # true only when stage launches device work + combine(problem, handle) -> tensor # pure combine comm (timed) + inspect_dispatch(problem, handle) -> view # normalized payload/expert/weight metadata + combine_transformed(problem, handle, tensor) -> tensor + recv_tokens(handle) -> int # realized tokens received this rank + finalize(rc) -> int|NoReturn +""" +from __future__ import annotations + +import argparse +import datetime as _dt +import hashlib +import json +import math +import os +import types + +import ep_provenance +import identity +import workload as workload_contract + +# Raw v1 result emitted by one benchmark case. Publication uses a separate contract. +SCHEMA_VERSION = 1 + +# Every comparison-grade EP point uses the same literal timing profile on every SKU/backend. +# Eight timed iterations keep each MoRI burst well below its sustained-iteration wedge, 64 trials +# provide 512 observations per operation, and 32 warmups meet Blackwell's measured clock-ramp floor. +SAMPLING_CONTRACT = identity.V1_CASE_PROFILE["sampling_contract"] +TIMED_SAMPLES_PER_POINT = 512 +TIMED_ITERS_PER_TRIAL = 8 +TRIALS_PER_POINT = 64 +WARMUP_ITERS_PER_TRIAL = 32 +WARMUP_SEMANTICS = "full-roundtrip-before-each-component-trial-point-v1" +QUALIFICATION_RUNS = 3 +ROUTING_SEED = 67 +ROUTING_GENERATOR = workload_contract.GENERATOR_VERSION +ACTIVATION_PROFILE = "canonical-counter-source-v4" +ACTIVATION_GENERATOR = workload_contract.ACTIVATION_GENERATOR +PLACEMENT = "packed" +COMPONENT_ORDER_CONTRACT = "qualification-hash-rotated-components-v1" +LOW_LATENCY_MODE = "low-latency" +LOW_LATENCY_MAX_TOKENS_PER_RANK = 128 +LOW_LATENCY_MEASUREMENT_CONTRACT = "expert-packed-weighted-combine-v1" +LOW_LATENCY_COMPONENT_ORDER_CONTRACT = "qualification-hash-rotated-components-v1" +LOW_LATENCY_ORACLE_CONTRACT = "expert-assignment-transform-v1" +LOW_LATENCY_CORRECTNESS_SCOPE = "expert-assignment-and-weighted-combine" + +# Phase-default sweeps — token-size regimes, NOT distinct kernels (both run normal +# mode; "decode"/"prefill" name the small/large-token regime). Powers of two for a +# clean log x-axis; clamped to the backend buffer ceiling (MoRI's registerable heap). +DECODE_LADDER = [1, 2, 4, 8, 16, 32, 64, 128] +PREFILL_LADDER = [128, 256, 512, 1024, 2048, 4096] +# Conditioning replays a fixed phase ramp before each measured shape to settle +# clocks and routing state; these rounds are never timed or emitted. The ladders +# and round count are the frozen v1 measurement contract (see conditioning_contract). +CONDITIONING_LADDERS = { + "decode": [1, 2, 4, 8, 16, 32, 64, 128], + "prefill": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512], +} +CONDITIONING_ROUNDS_PER_SHAPE = 8 +CONDITIONING_CONTRACT = identity.V1_CASE_PROFILE["conditioning_contract"] +ORACLE_CONTRACT = identity.V1_CASE_PROFILE["oracle_contract"] +# Dispatch and combine are fixed BF16, so the combine oracle uses one frozen gate. +ORACLE_RTOL = 5e-2 +ORACLE_ATOL = 2e-2 + +EPLB_REDUNDANT_EXPERTS = 32 +EPLB_REFERENCE_TOKENS_PER_RANK = 2048 +EPLB_PLANNER = "greedy-rank-major-v1" +V1_PROFILE = { + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "mode": "normal", + "measurement_contract": "layout-and-dispatch-v1", + "resource_mode": "fixed-profile", + "placement": PLACEMENT, + "activation_profile": ACTIVATION_PROFILE, + "activation_generator": ACTIVATION_GENERATOR, + "routing_generator": ROUTING_GENERATOR, + "component_order_contract": COMPONENT_ORDER_CONTRACT, + "conditioning_contract": CONDITIONING_CONTRACT, + "eplb_reference_tokens_per_rank": EPLB_REFERENCE_TOKENS_PER_RANK, + "eplb_redundant_experts": EPLB_REDUNDANT_EXPERTS, + "eplb_planner": EPLB_PLANNER, + # DeepEP/UCCL use this only as the fallback when their tuned default is not exported. + "num_sms": 24, +} + + +def logical_byte_provenance(logical_copies: int, hidden: int) -> dict[str, int | str]: + """Return comparable logical BF16 activation bytes for one direction. + + Dispatch and combine both move BF16 (2 bytes/value) with no separate scale + payload, so ``scale_bytes`` is always zero. + """ + if logical_copies < 0 or hidden < 0: + raise ValueError("logical byte dimensions must be non-negative") + activation_data_bytes = logical_copies * hidden * 2 + return { + "accounting_contract": "activation-data-plus-scales-v1", + "activation_data_bytes": activation_data_bytes, + "scale_bytes": 0, + "total_logical_bytes": activation_data_bytes, + } + +def format_collective_version(raw) -> str: + """Normalize PyTorch's tuple or packed NCCL/RCCL version representation.""" + if isinstance(raw, int): + if raw < 10_000: + return f"{raw // 1000}.{raw // 100 % 10}.{raw % 100}" + return f"{raw // 10_000}.{raw // 100 % 100}.{raw % 100}" + if isinstance(raw, (tuple, list)): + return ".".join(map(str, raw)) + return str(raw) if raw not in (None, "") else "unknown" + + +def add_common_args(ap: argparse.ArgumentParser) -> None: + """Add the varying v1 inputs; fixed profile values are not CLI axes.""" + ap.set_defaults(**V1_PROFILE) + ap.add_argument("--mode", default="normal", choices=["normal", LOW_LATENCY_MODE]) + ap.add_argument("--phase", default="decode", choices=["decode", "prefill"], + help="token-size regime: decode (small T) / prefill (large T) — picks the default ladder") + ap.add_argument("--tokens-ladder", default="", + help="space/comma-separated source-tokens-per-rank sweep; blank = phase default") + ap.add_argument("--hidden", type=int, default=7168) + ap.add_argument("--topk", type=int, default=8) + ap.add_argument("--experts", type=int, default=256, help="TOTAL experts (fixed across EP degrees)") + ap.add_argument("--routing", default="uniform", choices=["uniform"]) + # Canonical workloads consume pre-generated trace bytes instead of the + # seeded runtime generator, so a result is provably the SAME workload as another machine's + # (checksum match). Points at a dir of .npz/.manifest.json (make_workloads.py). + ap.add_argument("--workload-dir", default="", + help="dir of canonical workload traces; empty = seeded runtime generation (dev)") + ap.add_argument("--case-id", default="") + ap.add_argument("--suite", default="") + ap.add_argument("--workload-name", default="") + ap.add_argument("--seed", type=int, default=ROUTING_SEED) + ap.add_argument( + "--qualification-index", + type=int, + choices=range(1, QUALIFICATION_RUNS + 1), + default=os.environ.get("CX_QUALIFICATION_INDEX", "1"), + help="one-based qualification repeat used for deterministic measurement ordering", + ) + ap.add_argument( + "--version", + type=int, + default=os.environ.get("CX_VERSION", "1"), + help="iterable benchmark version copied verbatim into the emitted result", + ) + # 32: B300/Blackwell needs ~30 untimed iters to reach steady-state GPU clocks + + # establish NVLink/NVSHMEM connections — at warmup=8 its dispatch read ~1787us + # (cold), at warmup>=30 it settles to ~85us (faster than H100, reproducible within + # ~2.5%). H100/MI355X reach steady state much sooner; the extra iters are harmless. + ap.add_argument("--warmup", type=int, default=WARMUP_ITERS_PER_TRIAL, + help=f"untimed full roundtrips before each trial/point; fixed by " + f"{SAMPLING_CONTRACT} to {WARMUP_ITERS_PER_TRIAL}") + ap.add_argument("--iters", type=int, default=TIMED_ITERS_PER_TRIAL, + help=f"timed iterations per trial; fixed by {SAMPLING_CONTRACT} to " + f"{TIMED_ITERS_PER_TRIAL}") + ap.add_argument("--trials", type=int, default=TRIALS_PER_POINT, + help=f"timed trials; fixed by {SAMPLING_CONTRACT} to {TRIALS_PER_POINT}") + # provenance / output + ap.add_argument("--runner", required=True) + ap.add_argument("--topology-class", required=True) + ap.add_argument("--transport", default="") + ap.add_argument("--scope", required=True, choices=["scale-up", "scale-out"]) + ap.add_argument("--scale-up-transport", required=True) + ap.add_argument("--scale-out-transport", default="") + # gpus-per-node=0 means one node containing the whole EP group. + ap.add_argument("--gpus-per-node", type=int, default=0) + ap.add_argument("--scale-up-domain", type=int, default=0, help="0 = gpus_per_node*ep (one domain)") + ap.add_argument("--timestamp") + ap.add_argument("--out", required=True) + + +def token_ladder(spec: str, phase: str, cap: int | None) -> tuple[list[int], list[int]]: + """Return (ladder, dropped): explicit spec else the phase default; positive ints; + clamped to `cap` with dropped points reported (never silently truncated).""" + if spec and spec.strip(): + want = [int(t) for t in spec.replace(",", " ").split() if t] + else: + want = DECODE_LADDER if phase == "decode" else PREFILL_LADDER + want = sorted({t for t in want if t > 0}) + if cap is not None: + return [t for t in want if t <= cap], [t for t in want if t > cap] + return want, [] + + +def sampling_contract_error(iters: int, trials: int, warmup: int) -> str | None: + """Return a user-facing error unless the exact cross-SKU timing profile is used.""" + expected = (TIMED_ITERS_PER_TRIAL, TRIALS_PER_POINT, WARMUP_ITERS_PER_TRIAL) + observed = (iters, trials, warmup) + if observed != expected: + return (f"{SAMPLING_CONTRACT} requires exactly iters:trials:warmup=" + f"{expected[0]}:{expected[1]}:{expected[2]} on every SKU/backend; got " + f"{observed[0]}:{observed[1]}:{observed[2]} " + f"({iters * trials if iters > 0 and trials > 0 else 'invalid'} timed samples)") + return None + + +def qualification_order( + values: list, qualification_index: int, trial_index: int, *, identity_key: str = "" +) -> list: + """Return a deterministic, position-balanced order for one qualification trial. + + Official runs bind the base permutation to the case identity. The cyclic schedule then gives + every value every position equally often over 64 trials while qualification repeats start at + different offsets. Keeping the empty-key behavior stable preserves local diagnostic fixtures. + """ + if not values or len(values) != len(set(values)): + raise ValueError("qualification order requires non-empty unique values") + if qualification_index not in range(1, QUALIFICATION_RUNS + 1): + raise ValueError(f"qualification_index must be in 1..{QUALIFICATION_RUNS}") + if type(trial_index) is not int or trial_index < 0: + raise ValueError("trial_index must be a non-negative integer") + if not isinstance(identity_key, str): + raise ValueError("qualification identity_key must be a string") + base_values = list(values) + if identity_key: + base_values.sort( + key=lambda value: hashlib.sha256( + f"{identity_key}\0{qualification_index}\0{value}".encode("utf-8") + ).digest() + ) + position = trial_index + qualification_index - 1 + cycle, offset = divmod(position, len(values)) + base = base_values if cycle % 2 == 0 else list(reversed(base_values)) + return base[offset:] + base[:offset] + + +def sampled_component_evidence(trials: list[list[float]]) -> dict: + """Validate and copy private 64x8 trial blocks without flattening their boundaries.""" + if not trials: + return {"availability": "unavailable", "sample_count": 0, "trials": None} + if len(trials) != TRIALS_PER_POINT: + raise ValueError( + f"measured component needs {TRIALS_PER_POINT} trial blocks; got {len(trials)}" + ) + normalized: list[list[float]] = [] + for trial in trials: + if len(trial) != TIMED_ITERS_PER_TRIAL: + raise ValueError( + f"measured trial needs {TIMED_ITERS_PER_TRIAL} samples; got {len(trial)}" + ) + block = [] + for sample in trial: + if isinstance(sample, bool) or not isinstance(sample, (int, float)): + raise ValueError("measured samples must be numeric") + value = float(sample) + if not math.isfinite(value) or value < 0: + raise ValueError("measured samples must be finite and non-negative") + block.append(value) + normalized.append(block) + count = sum(map(len, normalized)) + if count != TIMED_SAMPLES_PER_POINT: + raise ValueError( + f"measured component needs {TIMED_SAMPLES_PER_POINT} samples; got {count}" + ) + return {"availability": "measured", "sample_count": count, "trials": normalized} + + +def _stats_vec(xs: list[int]) -> dict: + """min/mean/max/CV (+ empty count) of a per-rank count vector — self-describing source-token + or load summary without dumping the full vector.""" + n = len(xs) or 1 + mean = sum(xs) / n + var = sum((x - mean) ** 2 for x in xs) / n + cv = (var ** 0.5 / mean) if mean > 0 else 0.0 + return {"min": min(xs) if xs else 0, "mean": round(mean, 3), + "max": max(xs) if xs else 0, "cv": round(cv, 4), + "empty_ranks": sum(1 for x in xs if x == 0), "total": sum(xs), "ranks": n} + + +def percentile(xs: list[float], q: float) -> float: + if not xs: + return float("nan") + s = sorted(xs) + i = max(0, min(len(s) - 1, math.ceil(q / 100.0 * len(s)) - 1)) + return s[i] + + +def _sha256_json(value) -> str: + payload = json.dumps( + value, allow_nan=False, ensure_ascii=False, sort_keys=True, separators=(",", ":") + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def _write_bytes_atomic(path: str, payload: bytes) -> tuple[str, int]: + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + temporary = f"{path}.tmp-{os.getpid()}" + try: + with open(temporary, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + finally: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + return hashlib.sha256(payload).hexdigest(), len(payload) + + +def _write_json_atomic(path: str, value) -> tuple[str, int]: + payload = ( + json.dumps(value, allow_nan=False, ensure_ascii=False, indent=2) + "\n" + ).encode() + return _write_bytes_atomic(path, payload) + + +def time_us(torch, fn, warmup: int, iters: int, pre=None, post=None) -> list[float]: + """Per-iteration CUDA-event latencies (µs) for THIS rank. + + Without `pre`: times `fn()`. With `pre`: runs `pre()` UNTIMED each iteration (sync + before the start event so its GPU work can't bleed in), then times `fn(pre_result)`. + `post(result)` runs after the end event and synchronization, so stateful backends can + consume/reset a timed operation without charging that cleanup to its latency. Returns + the raw per-iteration series; the caller reduces across ranks per iteration before + percentiling. + """ + def sample(): + arg = pre() if pre is not None else None + if pre is not None: + torch.cuda.synchronize() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + result = fn(arg) if pre is not None else fn() + e.record() + torch.cuda.synchronize() + elapsed = s.elapsed_time(e) * 1000.0 # ms -> us + if post is not None: + post(result) + torch.cuda.synchronize() + return elapsed + + for _ in range(max(0, warmup)): + if pre is not None: + a = pre() + torch.cuda.synchronize() + fn(a) + else: + fn() + # sync EACH warmup iteration, not just once after the loop: the measured-roundtrip fn + # interleaves dispatch+combine on a backend's persistent comm buffer, so back-to-back + # un-synced warmup iterations let iter N+1's dispatch race iter N's combine (CUDA abort + # on a rank -> NCCL-watchdog SIGABRT). Cheap (warmup is small); timed samples already sync. + torch.cuda.synchronize() + return [sample() for _ in range(iters)] + + +def kernel_generation(backend) -> str: + """Return the adapter's explicit kernel family when one exists.""" + declared = getattr(backend, "kernel_generation", None) + if declared: + return declared + return { + "deepep": "v1", + "deepep-v2": "v2-elastic-buffer", + "deepep-hybrid": "hybrid", + }.get(backend.name, "n-a") + + +def _reduce_vec(torch, dist, device, vals, op): + t = torch.tensor(vals, device=device, dtype=torch.float64) + dist.all_reduce(t, op=op) + return [float(x) for x in t.tolist()] + + +def _reduce_int(torch, dist, device, v: int, op) -> int: + t = torch.tensor([int(v)], device=device, dtype=torch.int64) + dist.all_reduce(t, op=op) + return int(t.item()) + + +def _same_hash_across_ranks(torch, dist, device, digest: str) -> bool: + parts = [int(digest[offset:offset + 8], 16) for offset in range(0, 64, 8)] + low = torch.tensor(parts, device=device, dtype=torch.int64) + high = low.clone() + dist.all_reduce(low, op=dist.ReduceOp.MIN) + dist.all_reduce(high, op=dist.ReduceOp.MAX) + return bool(torch.equal(low, high)) + + +def _tensor_sha256(*tensors) -> str: + digest = hashlib.sha256() + for tensor in tensors: + digest.update(tensor.detach().contiguous().cpu().numpy().tobytes()) + return digest.hexdigest() + + +def _normalized_expert_metadata(torch, expert_ids, weights): + """Sort each row by global expert ID while keeping -1 sentinels last.""" + valid = expert_ids >= 0 + keys = torch.where(valid, expert_ids.to(torch.int64), torch.full_like(expert_ids, 1 << 30)) + order = torch.argsort(keys, dim=1, stable=True) + sorted_ids = torch.gather(expert_ids.to(torch.int64), 1, order) + sorted_weights = torch.gather(weights.to(torch.float32), 1, order) + sorted_valid = sorted_ids >= 0 + return ( + torch.where(sorted_valid, sorted_ids, torch.full_like(sorted_ids, -1)), + sorted_weights.masked_fill(~sorted_valid, 0), + ) + + +def expert_packed_slot_map( + counts, + src_info, + layout_range, + *, + tokens_per_rank: int, + experts_per_rank: int, + world_size: int, +) -> list[tuple[int, int, int]]: + """Decode and validate DeepEP's expert-packed receive metadata. + + ``src_info`` stores a source-local token index. The source rank is carried by + the corresponding packed ``layout_range`` interval, so neither field is + independently sufficient to identify a source token. + """ + if tokens_per_rank <= 0 or experts_per_rank <= 0 or world_size <= 0: + raise ValueError("expert-packed dimensions must be positive") + if len(counts) != experts_per_rank: + raise ValueError("expert-packed count shape differs from local experts") + if len(src_info) != experts_per_rank or len(layout_range) != experts_per_rank: + raise ValueError("expert-packed metadata shape differs from local experts") + + mask = (1 << 32) - 1 + slots: list[tuple[int, int, int]] = [] + pairs: set[tuple[int, int]] = set() + for local_expert in range(experts_per_rank): + count = counts[local_expert] + if type(count) is not int or count < 0: + raise ValueError("expert-packed receive count is invalid") + if len(layout_range[local_expert]) != world_size: + raise ValueError("expert-packed layout rank dimension is invalid") + if len(src_info[local_expert]) < count: + raise ValueError("expert-packed source metadata is truncated") + + covered = [False] * count + for source_rank, encoded in enumerate(layout_range[local_expert]): + if type(encoded) is not int or encoded < 0: + raise ValueError("expert-packed layout range is invalid") + begin, span = encoded >> 32, encoded & mask + if begin > count or begin + span > count: + raise ValueError("expert-packed layout range exceeds valid slots") + for packed_position in range(begin, begin + span): + if covered[packed_position]: + raise ValueError("expert-packed layout ranges overlap") + covered[packed_position] = True + local_source = src_info[local_expert][packed_position] + if ( + type(local_source) is not int + or local_source < 0 + or local_source >= tokens_per_rank + ): + raise ValueError("expert-packed source token index is invalid") + source_id = source_rank * tokens_per_rank + local_source + pair = (source_id, local_expert) + if pair in pairs: + raise ValueError("expert-packed source/expert assignment is duplicated") + pairs.add(pair) + slots.append((local_expert, packed_position, source_id)) + if not all(covered): + raise ValueError("expert-packed layout ranges omit valid receive slots") + return slots + + +def expert_packed_dispatch_view( + torch, + packed_payload, + packed_counts, + packed_src_info, + packed_layout_range, + *, + rank: int, + tokens_per_rank: int, + experts_per_rank: int, + world_size: int, +): + """Return the valid expert-packed rows with exact global source identities.""" + if packed_payload.ndim != 3: + raise ValueError("expert-packed payload must have shape [experts, slots, hidden]") + if packed_payload.shape[0] != experts_per_rank: + raise ValueError("expert-packed payload expert dimension is invalid") + if tuple(packed_counts.shape) != (experts_per_rank,): + raise ValueError("expert-packed count tensor shape is invalid") + if tuple(packed_src_info.shape[:1]) != (experts_per_rank,): + raise ValueError("expert-packed source tensor shape is invalid") + if tuple(packed_layout_range.shape) != (experts_per_rank, world_size): + raise ValueError("expert-packed layout tensor shape is invalid") + if packed_src_info.ndim != 2 or packed_src_info.shape[1] < packed_payload.shape[1]: + raise ValueError("expert-packed source tensor capacity is invalid") + + counts = [int(value) for value in packed_counts.detach().cpu().tolist()] + if any(count > packed_payload.shape[1] for count in counts): + raise ValueError("expert-packed receive count exceeds payload capacity") + slots = expert_packed_slot_map( + counts, + packed_src_info.detach().cpu().tolist(), + packed_layout_range.detach().cpu().tolist(), + tokens_per_rank=tokens_per_rank, + experts_per_rank=experts_per_rank, + world_size=world_size, + ) + device = packed_payload.device + local_expert_slots = torch.tensor( + [slot[0] for slot in slots], device=device, dtype=torch.int64 + ) + packed_positions = torch.tensor( + [slot[1] for slot in slots], device=device, dtype=torch.int64 + ) + source_ids = torch.tensor( + [slot[2] for slot in slots], device=device, dtype=torch.int64 + ) + expert_ids = local_expert_slots + rank * experts_per_rank + payload = packed_payload[local_expert_slots, packed_positions] + return types.SimpleNamespace( + payload=payload, + source_ids=source_ids, + expert_ids=expert_ids, + local_expert_counts=packed_counts.to(torch.int64), + local_expert_slots=local_expert_slots, + packed_positions=packed_positions, + ordering_contract="expert-major/layout-addressed-packed-slot-v1", + ) + + +def _expert_transform(torch, payload, expert_ids, weights, combine_weight_semantics): + """Build one local expert aggregate for the v1 unweighted combine contract.""" + if combine_weight_semantics != "unweighted-rank-sum": + raise ValueError("v1 requires unweighted rank-sum combine") + valid = expert_ids >= 0 + expert = expert_ids.clamp(min=0).to(torch.int64) + gate = weights.to(torch.float32).masked_fill(~valid, 0) + scale = ((expert * 17 + 5) % 31 + 1).to(torch.float32) / 32 + offset_a = (((expert * 29 + 7) % 37) - 18).to(torch.float32) / 64 + offset_b = (((expert * 43 + 11) % 41) - 20).to(torch.float32) / 128 + scale_sum = (gate * scale).sum(dim=1, keepdim=True) + offset_a_sum = (gate * offset_a).sum(dim=1, keepdim=True) + offset_b_sum = (gate * offset_b).sum(dim=1, keepdim=True) + columns = torch.arange(payload.shape[1], device=payload.device, dtype=torch.int64) + pattern = (((columns * 13) % 17) - 8).to(torch.float32) / 8 + transformed = ( + payload.float() * scale_sum + offset_a_sum + offset_b_sum * pattern.unsqueeze(0) + ) + return transformed.to(payload.dtype) + + +def _expert_transform_expanded(torch, payload, expert_ids): + """Apply the oracle transform to one row per token/expert assignment.""" + expert = expert_ids.to(torch.int64) + scale = (((expert * 17 + 5) % 31 + 1).to(torch.float32) / 32).unsqueeze(1) + offset_a = ((((expert * 29 + 7) % 37) - 18).to(torch.float32) / 64).unsqueeze(1) + offset_b = ((((expert * 43 + 11) % 41) - 20).to(torch.float32) / 128).unsqueeze(1) + columns = torch.arange(payload.shape[1], device=payload.device, dtype=torch.int64) + pattern = (((columns * 13) % 17) - 8).to(torch.float32) / 8 + transformed = payload.float() * scale + offset_a + offset_b * pattern.unsqueeze(0) + return transformed.to(payload.dtype) + + +def _expected_transformed_combine(torch, problem): + """Independently derive sum_i gate_i * expert_i(x) for each source token.""" + semantic_x = getattr(problem, "oracle_x", problem.x) + expected = torch.zeros_like(semantic_x, dtype=torch.float32) + expert_ids = problem.topk_idx.to(torch.int64) + weights = problem.topk_weights.to(torch.float32) + columns = torch.arange(semantic_x.shape[1], device=semantic_x.device, dtype=torch.int64) + pattern = (((columns * 13) % 17) - 8).to(torch.float32) / 8 + for slot in range(expert_ids.shape[1]): + expert = expert_ids[:, slot] + gate = weights[:, slot].unsqueeze(1) + scale = (((expert * 17 + 5) % 31 + 1).to(torch.float32) / 32).unsqueeze(1) + offset_a = ((((expert * 29 + 7) % 37) - 18).to(torch.float32) / 64).unsqueeze(1) + offset_b = ((((expert * 43 + 11) % 41) - 20).to(torch.float32) / 128).unsqueeze(1) + expert_output = semantic_x.float() * scale + offset_a + offset_b * pattern.unsqueeze(0) + expected.add_(gate * expert_output) + return expected + + +def _run_expert_packed_oracle( + torch, + routing, + backend, + problem, + global_idx, + global_weights, + rank: int, + experts_per_rank: int, + seed: int, +): + """Verify an expert-packed dispatch and native gate-weighted combine.""" + contract = LOW_LATENCY_ORACLE_CONTRACT + oracle_atol, oracle_rtol = ORACLE_ATOL, ORACLE_RTOL + handle = backend.dispatch(problem) + torch.cuda.synchronize() + try: + packed = backend.inspect_expert_dispatch(problem, handle) + view = expert_packed_dispatch_view( + torch, + packed.payload, + packed.local_expert_counts, + packed.source_info, + packed.layout_range, + rank=rank, + tokens_per_rank=problem.T, + experts_per_rank=experts_per_rank, + world_size=backend.world_size, + ) + decoded_source_ids = routing.decode_source_ids(view.payload, seed) + except Exception as inspection_error: + try: + problem.recv_tokens = backend.recv_tokens(handle) + backend.stage(problem, handle) + backend.combine(problem, handle) + torch.cuda.synchronize() + except Exception as cleanup_error: + raise inspection_error from cleanup_error + return { + "contract": contract, + "passed": False, + "ordering_contract": "adapter-inspection-failed", + "order_sha256": None, + "dispatch_sha256": None, + "combine_weight_semantics": getattr( + backend, "combine_weight_semantics", "undeclared" + ), + "receive_count": 0, + "atol": oracle_atol, + "max_absolute_error": None, + "max_elementwise_relative_error": None, + "max_relative_error": None, + "max_weight_error": None, + "rtol": oracle_rtol, + "checks": { + "combine_values": False, + "counts": False, + "metadata": False, + "multiplicity": False, + "payload": False, + "source_set": False, + "weights": False, + }, + } + + device = problem.x.device + world_size = backend.world_size + total_experts = experts_per_rank * world_size + global_idx_device = global_idx.to(device=device, dtype=torch.int64) + global_weights_device = global_weights.to(device=device, dtype=torch.float32) + source_grid = torch.arange( + global_idx_device.shape[0], device=device, dtype=torch.int64 + ).unsqueeze(1).expand_as(global_idx_device) + local_mask = (global_idx_device // experts_per_rank) == rank + expected_sources = source_grid[local_mask] + expected_experts = global_idx_device[local_mask] + expected_pair_weights = global_weights_device[local_mask] + + receive_count = int(view.payload.shape[0]) + shape_ok = ( + view.payload.ndim == 2 + and view.source_ids.shape == (receive_count,) + and view.expert_ids.shape == (receive_count,) + and view.local_expert_counts.shape == (experts_per_rank,) + ) + source_range = bool( + receive_count == 0 + or ( + (view.source_ids >= 0) + & (view.source_ids < global_idx_device.shape[0]) + ).all().item() + ) + expected_payload = ( + routing.activations_for_source_ids( + view.source_ids, problem.x.shape[1], seed, problem.x.dtype + ) + if source_range + else torch.empty_like(view.payload) + ) + payload_ok = bool( + source_range + and torch.equal(decoded_source_ids.to(torch.int64), view.source_ids) + and torch.equal(view.payload, expected_payload) + ) + + actual_keys = view.source_ids * total_experts + view.expert_ids + expected_keys = expected_sources * total_experts + expected_experts + actual_order = torch.argsort(actual_keys, stable=True) + expected_order = torch.argsort(expected_keys, stable=True) + canonical_sources = view.source_ids.index_select(0, actual_order) + canonical_experts = view.expert_ids.index_select(0, actual_order) + canonical_expected_weights = expected_pair_weights.index_select(0, expected_order) + expected_local_idx = global_idx_device[ + rank * problem.T:(rank + 1) * problem.T + ] + metadata_ok = bool( + shape_ok + and torch.equal(problem.topk_idx.to(torch.int64), expected_local_idx) + and torch.equal( + actual_keys.index_select(0, actual_order), + expected_keys.index_select(0, expected_order), + ) + ) + expected_counts = torch.bincount( + expected_experts - rank * experts_per_rank, minlength=experts_per_rank + ) + counts_ok = torch.equal( + view.local_expert_counts.to(torch.int64), expected_counts.to(torch.int64) + ) + actual_multiplicity = torch.bincount( + view.source_ids, minlength=global_idx_device.shape[0] + ) + expected_multiplicity = torch.bincount( + expected_sources, minlength=global_idx_device.shape[0] + ) + multiplicity_ok = torch.equal(actual_multiplicity, expected_multiplicity) + source_set_ok = torch.equal( + torch.sort(torch.unique(view.source_ids)).values, + torch.sort(torch.unique(expected_sources)).values, + ) + + expected_local_weights = global_weights_device[ + rank * problem.T:(rank + 1) * problem.T + ] + if problem.topk_weights.shape == expected_local_weights.shape: + max_weight_error = ( + float((problem.topk_weights.float() - expected_local_weights).abs().max().item()) + if expected_local_weights.numel() + else 0.0 + ) + else: + max_weight_error = None + weights_ok = max_weight_error == 0.0 + ordering_contract = f"canonical-source-expert-v1/{view.ordering_contract}" + order_sha256 = _tensor_sha256(canonical_sources, canonical_experts) + dispatch_sha256 = _tensor_sha256( + canonical_sources, canonical_experts, canonical_expected_weights + ) + + handle.oracle_local_expert_slots = view.local_expert_slots + handle.oracle_packed_positions = view.packed_positions + problem.recv_tokens = receive_count + transformed = _expert_transform_expanded(torch, view.payload, view.expert_ids) + view.combine_input = transformed + combined = backend.combine_transformed(problem, handle, transformed) + torch.cuda.synchronize() + expected_combined = _expected_transformed_combine(torch, problem) + if combined.shape == expected_combined.shape and combined.numel(): + absolute_error = (combined.float() - expected_combined).abs() + max_absolute_error = float(absolute_error.max().item()) + max_relative_error = max_absolute_error / ( + float(expected_combined.abs().max().item()) + 1e-6 + ) + max_elementwise_relative_error = float( + (absolute_error / expected_combined.abs().clamp_min(oracle_atol)).max().item() + ) + combine_values_ok = bool(torch.allclose( + combined.float(), expected_combined, rtol=oracle_rtol, atol=oracle_atol + )) + elif combined.shape == expected_combined.shape: + max_absolute_error = 0.0 + max_elementwise_relative_error = 0.0 + max_relative_error = 0.0 + combine_values_ok = True + else: + max_absolute_error = None + max_elementwise_relative_error = None + max_relative_error = None + combine_values_ok = False + tolerance = oracle_rtol + checks = { + "combine_values": combine_values_ok, + "counts": counts_ok, + "metadata": metadata_ok, + "multiplicity": multiplicity_ok, + "payload": payload_ok, + "source_set": source_set_ok, + "weights": weights_ok, + } + return { + "contract": contract, + "passed": bool( + all(checks.values()) + and ordering_contract + and max_relative_error is not None + and max_relative_error < tolerance + ), + "atol": oracle_atol, + "combine_weight_semantics": backend.combine_weight_semantics, + "ordering_contract": ordering_contract, + "order_sha256": order_sha256, + "dispatch_sha256": dispatch_sha256, + "receive_count": receive_count, + "max_absolute_error": max_absolute_error, + "max_elementwise_relative_error": max_elementwise_relative_error, + "max_relative_error": max_relative_error, + "max_weight_error": max_weight_error, + "rtol": oracle_rtol, + "checks": checks, + } + + +def _run_expert_oracle( + torch, + routing, + backend, + problem, + global_idx, + global_weights, + rank: int, + experts_per_rank: int, + seed: int, +): + """Verify one real dispatch/transform/combine without entering a timed region.""" + if getattr(backend, "oracle_layout", "token-rank") == "expert-packed": + return _run_expert_packed_oracle( + torch, + routing, + backend, + problem, + global_idx, + global_weights, + rank, + experts_per_rank, + seed, + ) + oracle_atol, oracle_rtol = ORACLE_ATOL, ORACLE_RTOL + handle = backend.dispatch(problem) + torch.cuda.synchronize() + try: + view = backend.inspect_dispatch(problem, handle) + source_ids = routing.decode_source_ids(view.payload, seed) + except Exception as inspection_error: + try: + problem.recv_tokens = backend.recv_tokens(handle) + backend.stage(problem, handle) + backend.combine(problem, handle) + torch.cuda.synchronize() + except Exception as cleanup_error: + raise inspection_error from cleanup_error + return { + "contract": ORACLE_CONTRACT, + "passed": False, + "ordering_contract": "adapter-inspection-failed", + "order_sha256": None, + "dispatch_sha256": None, + "combine_weight_semantics": getattr( + backend, "combine_weight_semantics", "undeclared" + ), + "receive_count": 0, + "atol": oracle_atol, + "max_absolute_error": None, + "max_elementwise_relative_error": None, + "max_relative_error": None, + "max_weight_error": None, + "rtol": oracle_rtol, + "checks": { + "combine_values": False, + "counts": False, + "metadata": False, + "multiplicity": False, + "payload": False, + "source_set": False, + "weights": False, + }, + } + + receive_count = int(view.payload.shape[0]) + shape_ok = ( + view.payload.ndim == 2 + and view.expert_ids.shape == (receive_count, problem.topk_idx.shape[1]) + and view.weights.shape == view.expert_ids.shape + ) + source_range = bool( + receive_count == 0 + or ((source_ids >= 0) & (source_ids < global_idx.shape[0])).all().item() + ) + if source_range: + expected_idx = global_idx.to(problem.x.device).index_select(0, source_ids) + expected_weights = global_weights.to(problem.x.device).index_select(0, source_ids) + local = (expected_idx // experts_per_rank) == rank + expected_ids = torch.where(local, expected_idx, torch.full_like(expected_idx, -1)) + expected_weights = expected_weights.masked_fill(~local, 0) + expected_payload = routing.activations_for_source_ids( + source_ids, problem.x.shape[1], seed, problem.x.dtype + ) + else: + expected_ids = torch.full_like(view.expert_ids, -1) + expected_weights = torch.zeros_like(view.weights) + expected_payload = torch.empty_like(view.payload) + actual_ids, actual_weights = _normalized_expert_metadata( + torch, view.expert_ids, view.weights + ) + expected_ids, expected_weights = _normalized_expert_metadata( + torch, expected_ids, expected_weights + ) + expected_sources = ( + ((global_idx // experts_per_rank) == rank).any(dim=1).nonzero(as_tuple=True)[0] + ).to(problem.x.device) + source_set_ok = ( + source_range + and source_ids.numel() == torch.unique(source_ids).numel() + and torch.equal(torch.sort(source_ids).values, expected_sources) + ) + payload_ok = source_range and torch.equal(view.payload, expected_payload) + metadata_ok = shape_ok and torch.equal(actual_ids, expected_ids) + max_weight_error = ( + float((actual_weights - expected_weights).abs().max().item()) + if actual_weights.numel() + else 0.0 + ) + weights_ok = max_weight_error == 0.0 + valid_expected = expected_ids >= 0 + expected_local = expected_ids[valid_expected] - rank * experts_per_rank + expected_counts = torch.bincount(expected_local, minlength=experts_per_rank) + counts_ok = torch.equal( + view.local_expert_counts.to(torch.int64), expected_counts.to(torch.int64) + ) + multiplicity_ok = torch.equal( + (actual_ids >= 0).sum(dim=1), (expected_ids >= 0).sum(dim=1) + ) + # Receive-slot assignment may use atomics and is not a semantic EP guarantee. Compare + # pre/post dispatch evidence in canonical source-token order without changing the native path. + canonical_order = torch.argsort(source_ids.to(torch.int64), stable=True) + canonical_sources = source_ids.to(torch.int64).index_select(0, canonical_order) + canonical_ids = actual_ids.to(torch.int64).index_select(0, canonical_order) + canonical_weights = actual_weights.index_select(0, canonical_order) + ordering_contract = f"canonical-source-id-v1/{view.ordering_contract}" + order_sha256 = _tensor_sha256(canonical_sources) + dispatch_sha256 = _tensor_sha256( + canonical_sources, canonical_ids, canonical_weights + ) + + problem.recv_tokens = receive_count + combine_weight_semantics = backend.combine_weight_semantics + transformed = _expert_transform( + torch, view.payload, actual_ids, actual_weights, combine_weight_semantics + ) + view.combine_input = transformed + combined = backend.combine_transformed(problem, handle, transformed) + torch.cuda.synchronize() + expected_combined = _expected_transformed_combine(torch, problem) + if combined.shape == expected_combined.shape and combined.numel(): + absolute_error = (combined.float() - expected_combined).abs() + max_absolute_error = float(absolute_error.max().item()) + max_relative_error = max_absolute_error / ( + float(expected_combined.abs().max().item()) + 1e-6 + ) + max_elementwise_relative_error = float( + (absolute_error / expected_combined.abs().clamp_min(oracle_atol)).max().item() + ) + combine_values_ok = bool(torch.allclose( + combined.float(), expected_combined, rtol=oracle_rtol, atol=oracle_atol + )) + elif combined.shape == expected_combined.shape: + max_absolute_error = 0.0 + max_elementwise_relative_error = 0.0 + max_relative_error = 0.0 + combine_values_ok = True + else: + max_absolute_error = None + max_elementwise_relative_error = None + max_relative_error = None + combine_values_ok = False + tolerance = oracle_rtol + checks = { + "combine_values": combine_values_ok, + "counts": counts_ok, + "metadata": metadata_ok, + "multiplicity": multiplicity_ok, + "payload": payload_ok, + "source_set": source_set_ok, + "weights": weights_ok, + } + return { + "contract": ORACLE_CONTRACT, + "passed": bool( + all(checks.values()) + and ordering_contract + and max_relative_error is not None + and max_relative_error < tolerance + ), + "atol": oracle_atol, + "combine_weight_semantics": combine_weight_semantics, + "ordering_contract": ordering_contract, + "order_sha256": order_sha256, + "dispatch_sha256": dispatch_sha256, + "receive_count": receive_count, + "max_absolute_error": max_absolute_error, + "max_elementwise_relative_error": max_elementwise_relative_error, + "max_relative_error": max_relative_error, + "max_weight_error": max_weight_error, + "rtol": oracle_rtol, + "checks": checks, + } + + +def _histogram(xs: list[float], nbins: int = 40) -> dict: + """Compact equal-width summary of the exact private cross-rank-max samples.""" + if not xs: + return {"n": 0} + lo, hi = min(xs), max(xs) + if hi <= lo: + return {"n": len(xs), "min": lo, "max": hi, "bins": nbins, "counts": [len(xs)]} + counts = [0] * nbins + span = hi - lo + for x in xs: + b = min(nbins - 1, int((x - lo) / span * nbins)) + counts[b] += 1 + return {"n": len(xs), "min": round(lo, 3), "max": round(hi, 3), "bins": nbins, "counts": counts} + + +def run_sweep(args, backend, torch, dist, device, rank: int, world_size: int) -> int: + """Drive the source-tokens-per-rank sweep for one fully-specified line.""" + mode = getattr(args, "mode", "normal") + try: + case_profile = identity.profile_for_case({"mode": mode}) + except identity.IdentityError as exc: + if rank == 0: + print(f"ERROR: {exc}") + return 2 + sampling_error = sampling_contract_error(args.iters, args.trials, args.warmup) + if sampling_error: + if rank == 0: + print(f"ERROR: {sampling_error}") + return 2 + import routing # torch-based; imported lazily so the module byte-compiles without torch + import eplb # stdlib planner + torch remap (the EPLB transform) + + ep_size = world_size + num_logical = getattr(args, "num_logical_experts", args.experts) + if args.experts % ep_size != 0: + if rank == 0: + print(f"ERROR: experts ({args.experts}) must divide ep_size ({ep_size})") + return 2 + experts_per_rank = args.experts // ep_size + if getattr(backend, "mode", None) != mode: + if rank == 0: + print(f"ERROR: backend mode {getattr(backend, 'mode', None)!r} != {mode!r}") + return 2 + expected_weight_semantics = ( + "gate-weighted-sum" + if case_profile["combine_semantics"] == "gate-weighted" + else "unweighted-rank-sum" + ) + if getattr(backend, "combine_weight_semantics", None) != expected_weight_semantics: + if rank == 0: + print( + f"ERROR: {mode} requires combine semantics {expected_weight_semantics}" + ) + return 2 + if mode == LOW_LATENCY_MODE and ( + args.phase != "decode" + or getattr(backend, "oracle_layout", None) != "expert-packed" + or getattr(backend, "payload_unit", None) != "token-expert" + ): + if rank == 0: + print("ERROR: low-latency requires decode expert-packed token-expert execution") + return 2 + + spec = backend.make_inputs(args) + if not spec.ok: + if rank == 0: + print(f"ERROR: {spec.message}") + return spec.rc + cap = spec.cap + conditioning_ladder = spec.conditioning_ladder + ladder, dropped = spec.ladder, spec.dropped + if rank == 0 and dropped: + print(f"NOTE: dropped tokens/rank {dropped} — exceed {backend.name} buffer cap {cap} " + f"(hidden={args.hidden}); not silently truncated.") + loaded_workload_ids = spec.loaded_workload_ids + loaded_checksums = spec.loaded_checksums + canonical = bool(getattr(args, "workload_dir", "")) + MAX, MIN, SUM = dist.ReduceOp.MAX, dist.ReduceOp.MIN, dist.ReduceOp.SUM + + # EPLB is off in v1 (uniform placement): the plan/calibration stay None so the + # disabled eplb_record and the identity trace-remap paths downstream take their + # inert branch. `eplb` remains imported for those None-guarded call sites. + eplb_plan = None + eplb_calibration = None + + # Size the communicator from the resolved numeric shape. Buffer sizing needs the + # ladder numbers (not the input tensors), so it runs after make_inputs rather than + # in __init__ — resolving the historical build-buffers-before-inputs inversion. + # make_inputs already materialized the per-rank routing/activation inputs (canonical + # bytes when a workload_dir was staged, else seeded generation) on `spec`. + backend.create_buffer(spec) + + for wt in conditioning_ladder: + # Fabric/clock warm-up BEFORE any timed point (review: H200 had an anomalous + # cold first point and a 40% decode-vs-prefill mismatch at the shared T=128). + # make_inputs already materialized these warm-only per-rank inputs; ramping + # through the small ladder shapes untimed warms clocks/fabric for everyone and + # is cold-jump-safe for MoRI. Warm-only shapes carry no canonical manifest: + # they are never measured or emitted. + cwp = spec.conditioning_points[wt] + wp = backend.make_problem( + wt, cwp.topk_idx.to(device), cwp.topk_weights.to(device), cwp.activations + ) + backend.warm(wp, CONDITIONING_ROUNDS_PER_SHAPE) + torch.cuda.synchronize() + dist.barrier() + # Setup may materialize deferred provenance such as DeepEP V2 JIT CUBINs. + # Resolve it after conditioning but before correctness or timed measurements. + backend.capture_deferred_provenance() + provenance_issues = ep_provenance.backend_provenance_issues( + backend.name, backend.backend_provenance + ) + if provenance_issues: + if rank == 0: + print( + f"ERROR: unpinned provenance {provenance_issues} " + f"in {backend.backend_provenance}" + ) + return 4 + # ---- Pass 1: build each deterministic problem and run the expert oracle. ---- + problems, gate, gts, global_traces, input_snapshots = {}, {}, {}, {}, {} + routing_hashes = set() + for T in ladder: + counts = [T] * ep_size + gt = T * ep_size + gts[T] = gt + point = spec.points[T] + idx_g, w_g = point.global_idx, point.global_weights + rstats = routing.routing_stats(idx_g, args.experts, experts_per_rank, weights=w_g) + gpn = args.gpus_per_node or ep_size + rstats["locality"] = routing.routing_locality(idx_g, experts_per_rank, ep_size, max(1, T), + gpn, args.scale_up_domain or None) + rstats["source_token_stats"] = _stats_vec(counts) + routing_hashes.add(rstats["routing_hash"]) + my_cnt = T + problem = backend.make_problem( + my_cnt, point.topk_idx.to(device), point.topk_weights.to(device), point.activations + ) + input_snapshots[T] = ( + problem.x.clone(), problem.topk_idx.clone(), problem.topk_weights.clone() + ) + oracle = _run_expert_oracle( + torch, routing, backend, problem, idx_g, w_g, rank, experts_per_rank, + args.seed, + ) + before_x, before_idx, before_weights = input_snapshots[T] + pre_input_unchanged = ( + torch.equal(problem.x, before_x) + and torch.equal(problem.topk_idx, before_idx) + and torch.equal(problem.topk_weights, before_weights) + ) + problems[T] = problem + global_traces[T] = (idx_g, w_g) + gate[T] = { + "rstats": rstats, + "recv_local": oracle["receive_count"], + "max_rel": oracle["max_relative_error"] or 0.0, + "local_ok": int(oracle["passed"]), + "oracle_pre": oracle, + "pre_input_unchanged": pre_input_unchanged, + } + + # ---- Pass 2: every backend uses the same ascending point order and conditioning ramp. + # Per-iteration cross-rank MAX samples are pooled across trials. ---- + disp_pool = {T: [] for T in ladder} # pooled per-iteration cross-rank MAX (dispatch) + stage_pool = {T: [] for T in ladder} # measured only when stage launches device work + comb_pool = {T: [] for T in ladder} # ... combine + rt_pool = {T: [] for T in ladder} # independently measured round trip + disp_trials = {T: [] for T in ladder} + stage_trials = {T: [] for T in ladder} + comb_trials = {T: [] for T in ladder} + rt_trials = {T: [] for T in ladder} + order_identity = args.case_id or _sha256_json({ + "backend": backend.name, + "ep_size": ep_size, + "mode": mode, + "phase": args.phase, + "runner": args.runner, + "suite": args.suite, + }) + for trial_index in range(args.trials): + order = qualification_order( + list(ladder), args.qualification_index, trial_index, + identity_key=f"{order_identity}:tokens", + ) + for T in order: + problem = problems[T] + # timed_components() encodes the roundtrip-only vs full-component contract + # (and whether stage launches device work) once, in the base class. + component_order = qualification_order( + backend.timed_components(), + args.qualification_index, + trial_index, + identity_key=f"{order_identity}:components:{T}", + ) + measured = {name: [] for name in ("dispatch", "stage", "combine", "roundtrip")} + for component_name in component_order: + # The base template gives every component the same synchronized + # full-roundtrip warm-up before its timed trial and encodes the two + # branch rules (dispatch cleanup, combine re-dispatch) internally. + measured[component_name] = backend.benchmark_component( + component_name, problem, args.warmup, args.iters + ) + disp_iters = measured["dispatch"] + stage_iters = measured["stage"] + comb_iters = measured["combine"] + rt_iters = measured["roundtrip"] + # per-iteration cross-rank MAX (the distributed-op latency per iter), pooled. + if disp_iters: + reduced_dispatch = _reduce_vec(torch, dist, device, disp_iters, MAX) + reduced_combine = _reduce_vec(torch, dist, device, comb_iters, MAX) + disp_trials[T].append(reduced_dispatch) + comb_trials[T].append(reduced_combine) + disp_pool[T] += reduced_dispatch + comb_pool[T] += reduced_combine + if stage_iters: + reduced_stage = _reduce_vec(torch, dist, device, stage_iters, MAX) + stage_trials[T].append(reduced_stage) + stage_pool[T] += reduced_stage + reduced_roundtrip = _reduce_vec(torch, dist, device, rt_iters, MAX) + rt_trials[T].append(reduced_roundtrip) + rt_pool[T] += reduced_roundtrip + + # ---- Pass 3: prove timed inputs were immutable and repeat the full oracle. ---- + for T in ladder: + problem = problems[T] + before_x, before_idx, before_weights = input_snapshots[T] + input_unchanged = gate[T]["pre_input_unchanged"] and ( + torch.equal(problem.x, before_x) + and torch.equal(problem.topk_idx, before_idx) + and torch.equal(problem.topk_weights, before_weights) + ) + idx_g, w_g = global_traces[T] + post = _run_expert_oracle( + torch, routing, backend, problem, idx_g, w_g, rank, experts_per_rank, + args.seed, + ) + pre = gate[T]["oracle_pre"] + order_stable = ( + pre["ordering_contract"] == post["ordering_contract"] + and pre["order_sha256"] == post["order_sha256"] + and pre["dispatch_sha256"] == post["dispatch_sha256"] + ) + gate[T].update({ + "input_unchanged": input_unchanged, + "local_ok": int(pre["passed"] and post["passed"] and input_unchanged and order_stable), + "max_rel": max(pre["max_relative_error"] or 0.0, post["max_relative_error"] or 0.0), + "oracle_post": post, + "order_stable": order_stable, + }) + + # ---- Pass 4: percentiles (p50/p90/p95/p99, nearest-rank) from pooled samples + bytes + row ---- + def pcts(xs): + return ({"p50": percentile(xs, 50), "p90": percentile(xs, 90), + "p95": percentile(xs, 95), "p99": percentile(xs, 99)} if xs else None) + + def component(percentiles, count, *, derived=False): + if percentiles is None: + return {"availability": "unavailable", "origin": None, + "percentiles_us": None, "sample_count": 0} + return { + "availability": "derived" if derived else "measured", + "origin": "derived-percentile-sum" if derived else "measured", + "percentiles_us": percentiles, + "sample_count": 0 if derived else count, + } + rows = [] + all_anomalies = [] + thr_rt = 3.0 + for T in ladder: + gt = gts[T] + g = gate[T] + rstats = g["rstats"] + d, s, c, rt = disp_pool[T], stage_pool[T], comb_pool[T], rt_pool[T] + dp, sp, cp, rtp = pcts(d), pcts(s), pcts(c), pcts(rt) + # isolated_sum = SUM of the isolated dispatch+stage+combine percentiles. Stage contributes + # zero when it is explicitly not applicable. This is NOT a measured chained operation + # (can't reveal shared sync / launch amortization / overlap) — do NOT use for throughput + # or SLO capacity. The MEASURED round trip (rtp) is the real chained latency. + isum = ( + {key: dp[key] + (sp[key] if sp is not None else 0.0) + cp[key] for key in dp} + if dp and cp else None + ) + recv_total = _reduce_int(torch, dist, device, g["recv_local"], SUM) + recv_max = _reduce_int(torch, dist, device, g["recv_local"], MAX) + recv_min = _reduce_int(torch, dist, device, g["recv_local"], MIN) + global_ok = _reduce_int(torch, dist, device, g["local_ok"], MIN) + max_rel = _reduce_vec(torch, dist, device, [g["max_rel"]], MAX)[0] + point_ok = bool(global_ok) and recv_total > 0 + rank_evidence = [None] * world_size + dist.all_gather_object( + rank_evidence, + { + "input_unchanged": g["input_unchanged"], + "order_stable": g["order_stable"], + "post_timing": g["oracle_post"], + "pre_timing": g["oracle_pre"], + "rank": rank, + }, + ) + # Canonical LOGICAL payload byte contracts (from the routing trace, NOT backend recv + # tensors): token-rank = one copy per unique (token,dest-rank); token-expert = one copy + # per routed (token,expert). routed_copies = token-rank copies; gt*topk = token-expert. + token_rank_copies = rstats["routed_copies"] + logical_copies = ( + sum(rstats["expert_assignments_per_rank"]) + if case_profile["payload_unit"] == "token-expert" + else token_rank_copies + ) + H = args.hidden + throughput = { + percentile_name: gt / (latency_us * 1e-6) + for percentile_name, latency_us in rtp.items() + } + dispatch_bytes = logical_byte_provenance(logical_copies, H) + combine_bytes = logical_byte_provenance(logical_copies, H) + stage_bytes = { + "accounting_contract": "activation-data-plus-scales-v1", + "activation_data_bytes": 0, + "scale_bytes": 0, + "total_logical_bytes": 0, + } + roundtrip_bytes = { + "accounting_contract": "activation-data-plus-scales-v1", + **{ + field: dispatch_bytes[field] + combine_bytes[field] + for field in ( + "activation_data_bytes", "scale_bytes", "total_logical_bytes" + ) + }, + } + # Contract-level anomalies are attached to the row and rolled into validity. + # roundtrip_gt_isolated_sum: measured RT p99 >> Σ(isolated dispatch+combine) p99. + # roundtrip_lt_component_floor: measured RT p50 < max(dispatch,combine) p50 — a chained + # op can't finish faster than its slowest required component (sync semantics violated). + row_anoms = [] + if isum and isum["p99"] > 0 and rtp["p99"] > thr_rt * isum["p99"]: + row_anoms.append({"type": "roundtrip_gt_isolated_sum", "T": T, + "roundtrip_p99": round(rtp["p99"], 2), "isolated_sum_p99": round(isum["p99"], 2), + "ratio": round(rtp["p99"] / isum["p99"], 2), "threshold": thr_rt}) + floor = ( + max(dp["p50"], cp["p50"], sp["p50"] if sp is not None else 0.0) + if dp and cp else None + ) + if floor and rtp["p50"] > 0 and rtp["p50"] < 0.95 * floor: + row_anoms.append({"type": "roundtrip_lt_component_floor", "T": T, + "roundtrip_p50": round(rtp["p50"], 2), "component_floor_p50": round(floor, 2)}) + all_anomalies.extend(row_anoms) + rows.append({ + "anomalies": row_anoms, + "components": { + "combine": component(cp, len(c)), + "dispatch": component(dp, len(d)), + "isolated_sum": component(isum, 0, derived=True), + "roundtrip": component(rtp, len(rt)), + "stage": component(sp, len(s)), + }, + "correctness": { + "contract": case_profile["oracle_contract"], + "max_relative_error": max_rel, + "passed": point_ok, + "rank_evidence": rank_evidence, + "scope": case_profile["correctness_scope"], + }, + "global_tokens": gt, + "byte_provenance": { + "combine": combine_bytes, + "dispatch": dispatch_bytes, + "roundtrip": roundtrip_bytes, + "stage": stage_bytes, + }, + "receive": { + "max": recv_max, + "mean": recv_total / world_size, + "min": recv_min, + "total": recv_total, + }, + "routing": { + "empty_expert_count": rstats["empty_expert_count"], + "empty_rank_count": rstats["empty_rank_count"], + "expert_assignment_rank_cv": rstats["expert_assignment_rank_cv"], + "expert_assignments_per_rank": rstats["expert_assignments_per_rank"], + "expert_load_cv": rstats["expert_load_cv"], + "expert_load_max": rstats["expert_load_max"], + "expert_load_mean": rstats["expert_load_mean"], + "expert_load_min": rstats["expert_load_min"], + "fanout_histogram": rstats["fanout_hist"], + "fanout_max": rstats["fanout_max"], + "fanout_mean": rstats["fanout_mean"], + "fanout_min": rstats["fanout_min"], + "hash": rstats["routing_hash"], + "hotspot_ratio": rstats["hotspot_ratio"], + "locality": rstats.get("locality"), + "payload_copies_per_rank": rstats["payload_copies_per_rank"], + "payload_rank_cv": rstats["payload_rank_cv"], + "routed_copies": rstats["routed_copies"], + "source_token_stats": rstats.get("source_token_stats"), + }, + "sample_histograms": { + "dispatch": _histogram(d) if d else None, + "stage": _histogram(s) if s else None, + "combine": _histogram(c) if c else None, + "roundtrip": _histogram(rt), + }, + "token_rate_at_latency_percentile": throughput, + "tokens_per_rank": T, + }) + if rank == 0: + component_log = (f"disp p50/p99={dp['p50']:7.1f}/{dp['p99']:7.1f} " + f"comb {cp['p50']:6.1f}/{cp['p99']:6.1f} " if dp and cp + else "components=unavailable ") + print(f" T={T:<5} {component_log}" + f"RT p50/p99={rtp['p50']:7.1f}/{rtp['p99']:7.1f}us n={len(rt)} fanout={rstats['fanout_mean']:.2f} " + f"recv[min/mean/max]={recv_min}/{recv_total // world_size}/{recv_max} " + f"correct={point_ok}") + + # Cross-rank workload-identity proof: every rank must have built the SAME global routing + # (one hash per T here); confirm all ranks agree by hashing the per-T hash set and + # MIN/MAX-reducing it — a mismatch means NVIDIA and AMD did NOT run identical routing. + trace_sig = hashlib.sha256("|".join(sorted(routing_hashes)).encode()).hexdigest() + routing_consistent = _same_hash_across_ranks(torch, dist, device, trace_sig) + + # Capture again after correctness and timing so no lazily generated kernel can escape + # the implementation identity recorded in the artifact. + backend.capture_deferred_provenance() + + # status=valid requires correctness AND a proven-identical routing trace across ranks. + all_ok = bool(rows) and all(r["correctness"]["passed"] for r in rows) and routing_consistent + + # Adapters never self-label official; status is derived from these gates. + prov = backend.backend_provenance + provenance_complete = ep_provenance.provenance_complete( + prov, + backend.name, + getattr(args, "git_run", None), + image_digest=getattr(args, "image_digest", None), + image_verified=getattr(args, "image_digest_verified", False), + squash_sha256=getattr(args, "squash_sha256", None), + ) + resource_profile = ep_provenance.project_resource_profile(prov) + resource_conformance = resource_profile["conformance_class"] + # record the canonical workload identity consumed (one trace per T -> set of ids/checksums). + if canonical and loaded_workload_ids: + args.workload_id = identity.workload_id( + { + "members": [ + {"checksums": loaded_checksums[member], "workload_id": member} + for member in sorted(loaded_workload_ids) + ] + } + ) + args.workload_members = sorted(loaded_workload_ids) + args.workload_checksums = loaded_checksums + canonical_workload = bool(getattr(args, "workload_id", None)) + activation_identity = workload_contract.compute_activation_identity(args.seed, args.hidden) + # EPLB identity covers replica placement, not only counts. + eplb_mapping_hash = None + if eplb_plan is not None: + eplb_mapping_hash = eplb.mapping_hash(eplb_plan) + anomaly_free = len(all_anomalies) == 0 + validity = { + "execution_status": "complete" if rows else "failed", + "semantic_correctness": ( + "pass" if rows and all(r["correctness"]["passed"] for r in rows) else "fail" + ), + "workload_identity": "consistent-across-ranks" if routing_consistent else "inconsistent", + "workload_source": "canonical-serialized" if canonical_workload else "seeded-runtime", + "measurement_conformance": "conformant", # run_ep gate rejects nonconformant pre-run + "sampling_conformance": "conformant", # fixed-512-v1 gate rejects any other profile + "resource_conformance": resource_conformance, + "provenance_complete": provenance_complete, + # anomaly-free unless a contract-level timing anomaly fired (then diagnostic, see above). + "anomaly_free": anomaly_free, + } + + shape = { # FIXED line identity (no T, no per-backend resource knobs) + "hidden": args.hidden, "topk": args.topk, "experts": args.experts, + "experts_per_rank": experts_per_rank, + "routing": args.routing, "eplb": bool(eplb_plan), "num_logical_experts": num_logical, + # V2 is reserved for the PR #605 ElasticBuffer adapter; package versions never imply it. + "kernel_gen": kernel_generation(backend), + "activation_profile": ACTIVATION_PROFILE, + } + generated_at = args.timestamp or _dt.datetime.now().astimezone().isoformat() + realized_placement = getattr(args, "realized_placement", None) + nodes = ( + realized_placement["nodes"] + if realized_placement is not None + else int(os.environ.get("SLURM_NNODES", "1")) + ) + scheduled_case = { + "backend": backend.name, + "canonical": canonical, + "eplb": bool(eplb_plan), + "ep": ep_size, + "experts": num_logical, + "gpus_per_node": args.gpus_per_node or ep_size, + "hidden": args.hidden, + "ladder": " ".join(map(str, ladder)), + "mode": mode, + "nodes": nodes, + "phase": args.phase, + "routing": args.routing, + "samples_per_point": TIMED_SAMPLES_PER_POINT, + "scale_up_domain": args.scale_up_domain or (args.gpus_per_node or ep_size), + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "scope": args.scope, + "suite": args.suite or "manual", + "timing": f"{args.iters}:{args.trials}:{args.warmup}", + "topk": args.topk, + "topology_class": args.topology_class, + "transport": args.transport, + "warmup_semantics": WARMUP_SEMANTICS, + "workload": args.workload_name or "manual", + } + case_factors = { + "case": scheduled_case, + "profile": case_profile, + "sku": args.runner, + } + computed_case_id = identity.case_id_from_factors(case_factors) + if args.case_id and args.case_id != computed_case_id: + raise ValueError( + f"scheduled case ID does not match realized factors: {args.case_id} != {computed_case_id}" + ) + case_identifier = args.case_id or computed_case_id + git_run = getattr(args, "git_run", None) or {} + allocation_factors = { + "artifact": git_run.get("artifact"), + "execution_id": getattr(args, "allocation_execution_id", None), + "job": git_run.get("job"), + "repo": git_run.get("repo"), + "run_attempt": git_run.get("run_attempt"), + "run_id": git_run.get("run_id"), + "runner": args.runner, + "source_sha": git_run.get("source_sha"), + } + try: + attempt_ordinal = int(os.environ.get("CX_ATTEMPT_ID", "1")) + except ValueError: + attempt_ordinal = 0 + if attempt_ordinal <= 0: + raise ValueError("CX_ATTEMPT_ID must be a positive integer") + attempt_identifier = identity.attempt_id( + case=case_identifier, ordinal=attempt_ordinal + ) + runtime_fingerprint = getattr(args, "runtime_fingerprint", None) or {} + + sample_points = [] + for row in rows: + token_count = row["tokens_per_rank"] + + sample_point = { + "components": { + "combine": sampled_component_evidence(comb_trials[token_count]), + "dispatch": sampled_component_evidence(disp_trials[token_count]), + "roundtrip": sampled_component_evidence(rt_trials[token_count]), + "stage": sampled_component_evidence(stage_trials[token_count]), + }, + "tokens_per_rank": token_count, + } + sample_sha256 = _sha256_json(sample_point) + point_identifier = identity.point_id( + case=case_identifier, tokens_per_rank=token_count + ) + sample_point.update( + { + "point_id": point_identifier, + "sample_sha256": sample_sha256, + } + ) + sample_points.append(sample_point) + row.update({ + "point_id": point_identifier, + "sample_sha256": sample_sha256, + }) + + samples_path = args.out[:-5] + ".samples.json" if args.out.endswith(".json") else args.out + ".samples.json" + samples_document = { + "attempt_id": attempt_identifier, + "case_id": case_identifier, + "format": "collectivex.samples.v1", + "points": sample_points, + "sampling": { + "iterations_per_trial": args.iters, + "reduction": case_profile["rank_reduction"], + "trials": args.trials, + }, + "schema_version": 1, + } + samples_payload = ep_provenance.canonical_json_bytes(samples_document) + samples_sha256 = hashlib.sha256(samples_payload).hexdigest() + samples_bytes = len(samples_payload) + sample_artifact = { + "bytes": samples_bytes, + "format": "collectivex.samples.v1", + "path": os.path.basename(samples_path), + "sha256": samples_sha256, + } + headline = next((r for r in rows if r["tokens_per_rank"] == 64), rows[len(rows) // 2]) + eplb_record = ( + { + "calibration_token_offset": eplb_calibration["token_offset"], + "calibration_trace_sha256": eplb_calibration["trace_sha256"], + "calibration_window": eplb_calibration["window"], + "calibration_workload_id": eplb_calibration["workload_id"], + "enabled": True, + "imbalance_after": eplb_plan["imbalance_after"], + "imbalance_before": eplb_plan["imbalance_before"], + "mapping_hash": eplb_mapping_hash, + "max_replicas": eplb_plan["max_replicas"], + "num_logical_experts": num_logical, + "num_physical_experts": args.experts, + "num_redundant": args.experts - num_logical, + "planner": EPLB_PLANNER, + "reference_tokens_per_rank": EPLB_REFERENCE_TOKENS_PER_RANK, + "replicated_experts": eplb_plan["replicated_experts"], + } + if eplb_plan + else { + "calibration_token_offset": None, + "calibration_trace_sha256": None, + "calibration_window": None, + "calibration_workload_id": None, + "enabled": False, + "imbalance_after": None, + "imbalance_before": None, + "mapping_hash": None, + "max_replicas": None, + "num_logical_experts": num_logical, + "num_physical_experts": args.experts, + "num_redundant": 0, + "planner": None, + "reference_tokens_per_rank": None, + "replicated_experts": 0, + } + ) + doc = { + "format": "collectivex.ep.v1", + "schema_version": SCHEMA_VERSION, + "version": args.version, + "record_type": "case-attempt", + "generated_at": generated_at, + "identity": { + "allocation_factors": allocation_factors, + "attempt_id": attempt_identifier, + "attempt_ordinal": attempt_ordinal, + "case_factors": case_factors, + "case_id": case_identifier, + }, + "case": { + "attempt_ordinal": attempt_ordinal, + "backend": backend.name, + "eplb": eplb_record, + "ep_size": ep_size, + "mode": mode, + "phase": args.phase, + "resource_mode": "fixed-profile", + "runner": args.runner, + "shape": shape, + "suite": args.suite or "manual", + "workload_name": args.workload_name or "manual", + }, + "workload": { + "activation_generator": ACTIVATION_GENERATOR, + "activation_identity": activation_identity, + "activation_profile": ACTIVATION_PROFILE, + "cross_rank_consistent": routing_consistent, + "manifest_checksums": getattr(args, "workload_checksums", None), + "members": getattr(args, "workload_members", None), + "routing_generator": ROUTING_GENERATOR, + "source": validity["workload_source"], + "trace_hashes": sorted(routing_hashes), + "trace_signature": trace_sig, + "workload_id": getattr(args, "workload_id", None), + }, + "measurement": { + "component_order_contract": case_profile["component_order_contract"], + "conditioning": { + "contract": CONDITIONING_CONTRACT, + "ladder": conditioning_ladder, + "roundtrips_per_shape": CONDITIONING_ROUNDS_PER_SHAPE, + }, + "contract": case_profile["contract"], + "rows": rows, + "sampling": { + "contract": SAMPLING_CONTRACT, + "iterations_per_trial": args.iters, + "percentile_method": case_profile["percentile_method"], + "reduction": case_profile["rank_reduction"], + "samples_per_component": TIMED_SAMPLES_PER_POINT, + "trials": args.trials, + "warmup_iterations": args.warmup, + "warmup_semantics": WARMUP_SEMANTICS, + }, + "source_allocation": "even", + }, + "implementation": { + "kernel_generation": kernel_generation(backend), + "name": backend.name, + "provenance": backend.backend_provenance, + "resource_profile": resource_profile, + }, + "topology": { + "device_count": getattr(args, "runtime_device_count", None), + "device_product": getattr(args, "runtime_device_product", None), + "gpus_per_node": args.gpus_per_node or ep_size, + "nodes": nodes, + "placement": "packed", + "realized_placement": realized_placement, + "scale_up_domain": args.scale_up_domain or (args.gpus_per_node or ep_size), + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "scope": args.scope, + "topology_class": args.topology_class, + "transport": args.transport, + "world_size": world_size, + }, + "runtime_fingerprint": runtime_fingerprint, + "provenance": { + "command": getattr(args, "reproduction_command", ""), + "distributed_launcher": getattr(args, "distributed_launcher", None), + "git_run": getattr(args, "git_run", None), + "image": { + "arch": getattr(args, "image_arch", None), + "digest": getattr(args, "image_digest", "") or None, + "digest_verified": getattr(args, "image_digest_verified", False), + "reference": getattr(args, "image", "") or None, + "squash_sha256": getattr(args, "squash_sha256", None), + }, + "redaction": "sanitized-v1", + }, + "sample_artifact": sample_artifact, + "outcome": { + "reasons": [] if all_ok else ["semantic correctness or routing identity failed"], + "status": "success" if all_ok else "invalid", + "validity": validity, + }, + } + if rank == 0: + _write_bytes_atomic(samples_path, samples_payload) + _write_json_atomic(args.out, doc) + dispatch_percentiles = headline["components"]["dispatch"]["percentiles_us"] + dispatch_p99 = dispatch_percentiles["p99"] if dispatch_percentiles else None + component_summary = (f"disp_p99={dispatch_p99:.1f}us " + if dispatch_p99 is not None + else "components=unavailable ") + print(f"{backend.name} ep-dispatch-combine [{args.phase}/{mode}/{case_profile['contract']}]: " + f"status={doc['outcome']['status']} {len(rows)} pts, routing_consistent={routing_consistent}, " + f"headline T={headline['tokens_per_rank']} {component_summary}" + f"-> {args.out}") + # A complete invalid document is still a successfully captured terminal outcome. Launchers + # inspect its status to fail the case without conflating it with an execution failure. + return 0 diff --git a/experimental/CollectiveX/bench/ep_mori.py b/experimental/CollectiveX/bench/ep_mori.py new file mode 100644 index 0000000000..ffe174d336 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_mori.py @@ -0,0 +1,421 @@ +#!/usr/bin/env python3 +"""CollectiveX MoRI adapter: native BF16 dispatch/combine over mori.ops.""" +from __future__ import annotations + +import os +from pathlib import Path +import re +import sys +import types + +# MoRI registers the whole symmetric heap at import time. The pinned upstream +# inter-node benchmark uses 6 GiB for its InterNodeV1 staging and signal buffers. +os.environ["MORI_SHMEM_HEAP_SIZE"] = "6G" + +import torch +import torch.distributed as dist +from ep_backend import EPBackend + +try: + import mori # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: mori import failed: {exc!r}", file=sys.stderr) + raise + + +def _project_local_metadata(torch_module, raw_expert_ids, raw_weights, rank, experts_per_rank): + local_start = rank * experts_per_rank + local = (raw_expert_ids >= local_start) & ( + raw_expert_ids < local_start + experts_per_rank + ) + expert_ids = torch_module.where( + local, raw_expert_ids, torch_module.full_like(raw_expert_ids, -1) + ) + weights = torch_module.where(local, raw_weights, torch_module.zeros_like(raw_weights)) + return expert_ids, weights, raw_expert_ids[local] - local_start + + +def _mori_source_commit() -> str: + module_path = Path(mori.__file__).resolve() + for root in module_path.parents: + head = root / ".git" / "HEAD" + if not head.is_symlink() and head.is_file() and head.stat().st_size <= 128: + value = head.read_text(encoding="ascii").strip() + if re.fullmatch(r"[0-9a-f]{40}", value): + return value + raise RuntimeError("MoRI image source is not pinned to a detached commit") + raise RuntimeError("MoRI image source revision is unavailable") + + +def _mori_input_capacity(args, world_size: int) -> int: + """Reserve enough rows for the worst legal Zipf destination rank.""" + base = 512 + if getattr(args, "routing", "uniform") != "zipf": + return base + spec = str(getattr(args, "tokens_ladder", "")).replace(",", " ").split() + if spec: + largest = max(int(token) for token in spec) + else: + largest = 128 if getattr(args, "phase", "decode") == "decode" else 4096 + return max(base, largest * world_size) + + +class MoRIBackend(EPBackend): + name = "mori" + stage_device_work = False + combine_needs_redispatch = True + dispatch_needs_combine_cleanup = True + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + # mori is normal-mode only; base SUPPORTED_MODES=("normal",) enforces it. + super().__init__(args, rank, world_size, local_rank, device) + # The runner pins the AMD GPU architecture MoRI is built against. This is a + # hardware-lineage contract, independent of the (fixed BF16) communication path. + runner = str(getattr(args, "runner", "")) + if runner.startswith("mi355x"): + expected_arch = "gfx950" + elif runner.startswith(("mi300x", "mi325x")): + expected_arch = "gfx942" + else: + raise RuntimeError( + f"MoRI has no pinned GPU architecture for runner {runner!r}" + ) + + self.ep_size = world_size + self.experts_per_rank = args.experts // self.ep_size + device_properties = torch.cuda.get_device_properties(device) + device_cus = device_properties.multi_processor_count + realized_arch = str(getattr(device_properties, "gcnArchName", "")).split(":", 1)[0] + if realized_arch != expected_arch: + raise RuntimeError( + f"MoRI runner {runner!r} realized architecture {realized_arch!r}, " + f"expected {expected_arch!r}" + ) + gpus_per_node = int(args.gpus_per_node or world_size) + scale_up_domain = int(args.scale_up_domain or gpus_per_node) + scale_out = world_size > scale_up_domain + if ( + gpus_per_node <= 0 + or scale_up_domain <= 0 + or world_size % gpus_per_node + or world_size % scale_up_domain + ): + raise RuntimeError("MoRI placement is not divisible into complete domains") + if scale_out != (args.scope == "scale-out"): + raise RuntimeError("MoRI requested scope differs from the EP topology") + if not scale_out and ( + world_size != 8 + or gpus_per_node != 8 + or scale_up_domain != 8 + or args.scale_up_transport != "xgmi" + or args.scale_out_transport + or args.transport != "xgmi" + ): + raise RuntimeError("MoRI scale-up is pinned to EP8 over one XGMI domain") + if scale_out and ( + world_size != 16 + or gpus_per_node != 8 + or scale_up_domain != 8 + or args.scale_up_transport != "xgmi" + or args.scale_out_transport != "rdma" + or args.transport != "xgmi-rdma" + ): + raise RuntimeError( + "MoRI InterNodeV1 is pinned to EP16 over two 8-GPU XGMI/RDMA nodes" + ) + self.block_num = self._block_target = 80 + self.rdma_block_num = 0 + self.num_qps = 1 + self._block_floored = False + self._tuned_source = "default-80" + self.dispatch_warps = 16 + self.combine_warps = 8 + + # MI355X uses the direct intranode kernel. MI325X uses MoRI's split + # AsyncLL send/receive kernel as its normal-mode XGMI transport. + kernel_request = os.environ.get("CX_MORI_KERNEL_TYPE", "intranode").strip().lower() + self._kernel_type = None + self._kernel_type_label = "IntraNode" + self._async_ll = False + self._inter_node = False + if kernel_request in ("asyncll", "async_ll", "async-ll"): + if scale_out: + raise RuntimeError("MoRI EP16 must use InterNodeV1, not AsyncLL") + kernel_enum = getattr(mori.ops, "EpDispatchCombineKernelType", None) + if kernel_enum is None or not hasattr(kernel_enum, "AsyncLL"): + raise RuntimeError( + "CX_MORI_KERNEL_TYPE=asyncll requires " + "EpDispatchCombineKernelType.AsyncLL" + ) + self._kernel_type = kernel_enum.AsyncLL + self._kernel_type_label = "AsyncLL" + self._async_ll = True + self.block_num = self._block_target = 64 + self.dispatch_warps = self.combine_warps = 8 + self._tuned_source = "upstream-asyncll-64x8-external-input" + elif kernel_request in ("internode-v1", "internode_v1", "internodev1"): + if not scale_out: + raise RuntimeError("MoRI InterNodeV1 is valid only for scale-out EP16") + kernel_enum = getattr(mori.ops, "EpDispatchCombineKernelType", None) + if kernel_enum is None or not hasattr(kernel_enum, "InterNodeV1"): + raise RuntimeError( + "CX_MORI_KERNEL_TYPE=internode-v1 requires " + "EpDispatchCombineKernelType.InterNodeV1" + ) + self._kernel_type = kernel_enum.InterNodeV1 + self._kernel_type_label = "InterNodeV1" + self._inter_node = True + self.block_num = self._block_target = 96 + self.rdma_block_num = 64 + self.dispatch_warps = self.combine_warps = 8 + self._tuned_source = "upstream-internode-v1-96-64x8-qps1" + elif kernel_request not in ("intranode", "intra_node", "intra-node", ""): + raise RuntimeError( + f"unknown CX_MORI_KERNEL_TYPE={kernel_request!r} " + "(expected intranode|asyncll|internode-v1)" + ) + elif scale_out: + raise RuntimeError("MoRI scale-out EP16 requires CX_MORI_KERNEL_TYPE=internode-v1") + self.kernel_generation = ( + "inter-node-v1" if self._inter_node + else "async-ll" if self._async_ll + else "intranode" + ) + self._external_input = self._async_ll or self._inter_node + # Registered-input MoRI copies expert output into a device-side symmetric buffer. External + # input kernels consume the dispatch output directly, so their stage is not applicable. + self.stage_device_work = not self._external_input + # Stash the __init__-only locals the moved create_buffer body reads back. + self._gpus_per_node = gpus_per_node + + def create_buffer(self, spec): + # Local aliases re-expose the __init__ names so the moved shmem-init / + # config / op / provenance body below stays byte-verbatim. + args, world_size, rank, device = self.args, self.world_size, self.rank, self.device + gpus_per_node = self._gpus_per_node + device_cus = torch.cuda.get_device_properties(device).multi_processor_count + + world_group = torch.distributed.group.WORLD + torch._C._distributed_c10d._register_process_group("default", world_group) + mori.shmem.shmem_torch_process_group_init("default") + realized_qps = int(mori.shmem.shmem_num_qp_per_pe()) + if realized_qps < self.num_qps: + raise RuntimeError( + f"MoRI realized {realized_qps} QPs per PE; {self.num_qps} required" + ) + + self._cap = self.buffer_cap(args) + assert spec.max_tokens_per_rank <= self._cap + # BF16 communication path: no FP8 dispatch dtype, no scale payload, no combine quantization. + dispatch_dtype = torch.bfloat16 + config_kwargs = { + "data_type": dispatch_dtype, + "rank": rank, + "world_size": world_size, + "hidden_dim": args.hidden, + "scale_dim": 0, + "scale_type_size": 1, + "max_token_type_size": ( + torch.tensor([], dtype=torch.bfloat16).element_size() + if self._inter_node + else torch.tensor([], dtype=torch.float32).element_size() + ), + "max_num_inp_token_per_rank": max(512, self._cap), + "num_experts_per_rank": self.experts_per_rank, + "num_experts_per_token": args.topk, + "use_external_inp_buf": self._external_input, + "quant_type": "none", + } + if self._kernel_type is not None: + config_kwargs["kernel_type"] = self._kernel_type + if self._async_ll: + config_kwargs["max_total_recv_tokens"] = 0 + if self._async_ll or self._inter_node: + config_kwargs["block_num"] = self.block_num + config_kwargs["warp_num_per_block"] = self.dispatch_warps + if self._inter_node: + config_kwargs.update({ + "gpu_per_node": gpus_per_node, + "rdma_block_num": self.rdma_block_num, + "num_qp_per_pe": self.num_qps, + }) + self.config = mori.ops.EpDispatchCombineConfig(**config_kwargs) + expected_config = { + "data_type": dispatch_dtype, + "scale_dim": 0, + "scale_type_size": 1, + "use_external_inp_buf": self._external_input, + "quant_type": config_kwargs["quant_type"], + } + if self._async_ll or self._inter_node: + expected_config.update({ + "block_num": self.block_num, + "warp_num_per_block": self.dispatch_warps, + }) + if self._inter_node: + expected_config.update({ + "gpu_per_node": 8, + "rdma_block_num": 64, + "num_qp_per_pe": 1, + }) + if any(getattr(self.config, key, None) != value for key, value in expected_config.items()): + raise RuntimeError("MoRI requested launch/topology configuration was not realized") + # The newer pinned MoRI revision can otherwise replace explicit values + # with token-dependent tuning rules from the image. + os.environ["MORI_EP_LAUNCH_CONFIG_MODE"] = "MANUAL" + self.op = mori.ops.EpDispatchCombineOp(self.config) + if getattr(self.op, "launch_config_mode", None) != "MANUAL": + raise RuntimeError("MoRI explicit launch configuration was not applied") + + expected_mori_commit = os.environ.get("MORI_COMMIT") + mori_commit = _mori_source_commit() + if expected_mori_commit and mori_commit != expected_mori_commit: + raise RuntimeError("MoRI image source revision differs from canonical provenance") + self.backend_provenance = { + "mori_commit": mori_commit, + "api": ( + "mori.ops.EpDispatchCombineOp/external-input" + if self._external_input + else "mori.ops.EpDispatchCombineOp/registered-input" + ), + "mode": "normal", + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "kernel_type": self._kernel_type_label, + "enable_sdma": os.environ.get("MORI_ENABLE_SDMA"), + "heap_size": os.environ.get("MORI_SHMEM_HEAP_SIZE"), + "max_num_inp_token_per_rank": max(512, self._cap), + "max_total_recv_tokens": config_kwargs.get("max_total_recv_tokens"), + "gpus_per_node": gpus_per_node, + "rdma_block_num": self.rdma_block_num, + "use_external_inp_buf": self._external_input, + "num_qps": self.num_qps, + "resource_mode": "fixed-profile", + "block_num": self.block_num, + "block_num_target": self._block_target, + "block_num_floored": self._block_floored, + "dispatch_warps": self.dispatch_warps, + "combine_warps": self.combine_warps, + "device_cus": device_cus, + "sm_fraction": None if self._async_ll else self.block_num / device_cus, + "tuned_source": self._tuned_source, + } + + def buffer_cap(self, args): + return _mori_input_capacity(args, self.world_size) + + def make_problem(self, T, idx, weights, x): + indices = idx.to(torch.int32) + gate_weights = weights.to(torch.float32) + return types.SimpleNamespace( + T=T, + x=x, + dispatch_x=x, + topk_idx=indices, + topk_weights=gate_weights, + indices=indices, + weights=gate_weights, + scales=torch.empty((T, 0), dtype=torch.uint8, device=self.device), + ) + + def dispatch(self, p): + dispatch_output, dispatch_weights, _scales, dispatch_indices, recv_num = ( + self.op.dispatch( + p.dispatch_x, + p.weights, + p.scales, + p.indices, + block_num=self.block_num, + rdma_block_num=self.rdma_block_num, + warp_per_block=self.dispatch_warps, + ) + ) + if self._async_ll: + self.op.dispatch_recv(warp_per_block=self.dispatch_warps) + return types.SimpleNamespace( + dispatch_output=dispatch_output, + dispatch_weights=dispatch_weights, + dispatch_indices=dispatch_indices, + recv_num=recv_num[0], + combine_input=None, + ) + + def stage(self, p, h): + rows = getattr(p, "recv_tokens", None) + if not isinstance(rows, int) or rows < 0 or rows > h.dispatch_output.size(0): + raise RuntimeError("MoRI receive count was not validated before staging") + h.combine_input = h.dispatch_output + if self._external_input: + return None + buffer = self.op.get_registered_combine_input_buffer( + torch.bfloat16, hidden_dim=h.combine_input.size(1) + ) + buffer[:rows, :].copy_(h.combine_input[:rows, :]) + h.combine_input = buffer + + def combine(self, p, h): + combine_indices = p.indices if self._async_ll else h.dispatch_indices + combined, _weights = self.op.combine( + h.combine_input, + None, + combine_indices, + block_num=self.block_num, + rdma_block_num=self.rdma_block_num, + warp_per_block=self.combine_warps, + ) + if self._async_ll: + self.op.combine_recv(warp_per_block=self.combine_warps) + return combined[:p.T] + + def inspect_dispatch(self, p, h): + count = self.recv_tokens(h) + if h.dispatch_weights is None: + raise RuntimeError("MoRI dispatch did not expose gate weights") + if count < 0 or any( + tensor.ndim == 0 or count > tensor.size(0) + for tensor in (h.dispatch_output, h.dispatch_indices, h.dispatch_weights) + ): + raise RuntimeError("MoRI receive count exceeds dispatch metadata") + raw_expert_ids = h.dispatch_indices[:count].to(torch.int64) + expert_ids, weights, local_expert_ids = _project_local_metadata( + torch, + raw_expert_ids, + h.dispatch_weights[:count].to(torch.float32), + self.rank, + self.experts_per_rank, + ) + return types.SimpleNamespace( + payload=h.dispatch_output[:count], + expert_ids=expert_ids, + weights=weights, + local_expert_counts=torch.bincount( + local_expert_ids, minlength=self.experts_per_rank + ), + ordering_contract="mori-global-topk-masked-v1", + ) + + def combine_transformed(self, p, h, transformed): + h.combine_input = transformed.to(torch.bfloat16) + rows = getattr(p, "recv_tokens", None) + if not isinstance(rows, int) or rows < 0 or rows > h.combine_input.size(0): + raise RuntimeError("MoRI receive count was not validated before transformed combine") + if not self._external_input: + buffer = self.op.get_registered_combine_input_buffer( + torch.bfloat16, hidden_dim=h.combine_input.size(1) + ) + buffer[:rows, :].copy_(h.combine_input[:rows, :]) + h.combine_input = buffer + return self.combine(p, h) + + def recv_tokens(self, h): + return int(h.recv_num.item()) + + def finalize(self, rc): + try: + dist.barrier() + except Exception: + pass + sys.stdout.flush() + sys.stderr.flush() + os._exit(rc if 0 <= rc <= 255 else 1) diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py new file mode 100644 index 0000000000..309d7125bf --- /dev/null +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -0,0 +1,194 @@ +"""CollectiveX NCCL all-to-all expert-parallel reference backend. + +The canonical "token-shuffle" EP built on torch.distributed's NCCL ``all_to_all_single``. Like the +DeepEP-family APIs, dispatch sends one hidden-state copy to each distinct destination rank, even when +multiple selected experts live on that rank. Combine reverses the shuffle and sums those rank copies. + +Why this exists alongside DeepEP/UCCL/MoRI: it is the portable collective reference baseline for the +same rank-deduplicated payload and routing metadata. It keeps the library comparison anchored to the +platform collective stack without claiming the custom fused kernels use the same transport algorithm. + +Scope: BF16, normal mode, layout-and-dispatch-v1. The timed dispatch includes layout, count exchange, +payload, rank-masked expert indices, gate weights, and source-token metadata; combine returns only +the activation payload. RCCL exposes the same API. The v1 AMD matrix uses this backend at EP8 and EP16. +""" + +import os +import re +import types + +import torch +import torch.distributed as dist +import ep_provenance +from ep_backend import EPBackend + + +def _runtime_collective(args, torch_module) -> tuple[str, str]: + expected = "rccl" if torch_module.version.hip else "nccl" + fingerprint = getattr(args, "runtime_fingerprint", None) + collective = fingerprint.get("collective_library") if isinstance(fingerprint, dict) else None + if ( + not isinstance(collective, dict) + or collective.get("kind") != expected + or not isinstance(collective.get("version"), str) + or not re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+", collective["version"]) + ): + raise RuntimeError("loaded collective runtime identity is unavailable") + return expected, collective["version"] + + +class NCCLBackend(EPBackend): + name = "nccl-ep" + stage_device_work = False + combine_needs_redispatch = False # dispatch saves the permutation + splits + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + # Base validates mode against SUPPORTED_MODES=("normal",) — this backend is + # bf16 normal only, so an unsupported mode raises there. + super().__init__(args, rank, world_size, local_rank, device) + self.experts = args.experts + if args.experts % world_size: + raise ValueError(f"experts({args.experts}) must divide world_size({world_size})") + self.experts_per_rank = args.experts // world_size + self.tolerance = 5e-2 # bf16 round-trip + _library, _version = _runtime_collective(args, torch) + if args.scale_out_transport: + hcas = os.environ.get("NCCL_IB_HCA", "") + network_selection = os.environ.get("NCCL_NET") + if _library == "nccl" and network_selection != "IB": + raise RuntimeError("scale-out NCCL network mode is not IB") + if _library == "rccl" and network_selection: + raise RuntimeError("scale-out RCCL must auto-select its HCA-pinned network") + if not re.fullmatch( + r"=[A-Za-z][A-Za-z0-9_.-]{0,31}(?::[1-9][0-9]*)?" + r"(?:,[A-Za-z][A-Za-z0-9_.-]{0,31}(?::[1-9][0-9]*)?)*", + hcas, + ): + entries = hcas.removeprefix("=").split(",") if hcas else [] + shaped = all( + re.fullmatch( + r"[A-Za-z][A-Za-z0-9_.-]{0,31}(?::[1-9][0-9]*)?", entry + ) + for entry in entries + ) + raise RuntimeError( + "scale-out collective HCA allowlist is invalid " + f"(present={bool(hcas)}, exact={hcas.startswith('=')}, " + f"entries={len(entries)}, shaped={shaped})" + ) + self.kernel_generation = ep_provenance.collective_kernel_generation(_library) + self.backend_provenance = { + "backend": f"{_library}-all2all", + "backend_lineage": _library, + "collective_library": _library, + "nccl_version": _version, + "transport": f"{_library}-all_to_all_single", + "resource_mode": "fixed-profile", + "num_sms": None, + "device_sms": torch.cuda.get_device_properties(device).multi_processor_count, + "tuned_source": "nccl-collective", + "reference_semantics": "rank-deduplicated-payload-plus-routing-metadata-v2", + "routing_metadata": "expert-index-gate-weight-source-token", + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + } + + def create_buffer(self, spec): + # No fixed pre-allocated communicator: all-to-all sizes itself per step, and + # all provenance is resolved in __init__. Nothing to size from `spec`. + return None + + def dispatch(self, p): + ws = self.world_size + x = p.x # [T, H] bf16 + idx = p.topk_idx # [T, topk] + T, H = int(x.shape[0]), int(x.shape[1]) + dev = x.device + # DeepEP dispatches one token per destination rank, not one copy per expert. Build the same + # rank-deduplicated routing map so NCCL traffic and combine semantics are comparable. + destinations = (idx // self.experts_per_rank).clamp_(0, ws - 1) + present = torch.zeros((T, ws), dtype=torch.bool, device=dev) + present.scatter_(1, destinations, True) + flat_token, flat_dest = present.nonzero(as_tuple=True) + # Group rank copies by destination (stable -> deterministic, invertible permutation). + order = torch.argsort(flat_dest, stable=True) + ordered_token = flat_token.index_select(0, order) + ordered_dest = flat_dest.index_select(0, order) + send_counts = torch.bincount(flat_dest, minlength=ws) # [ws] + send_x = x.index_select(0, ordered_token).contiguous() + send_topk_idx = idx.index_select(0, ordered_token).contiguous() + expert_start = ordered_dest.unsqueeze(1) * self.experts_per_rank + local_mask = ((send_topk_idx >= expert_start) + & (send_topk_idx < expert_start + self.experts_per_rank)) + send_topk_idx = torch.where( + local_mask, send_topk_idx - expert_start, torch.full_like(send_topk_idx, -1) + ) + send_topk_weights = p.topk_weights.index_select(0, ordered_token).contiguous() + send_topk_weights.masked_fill_(~local_mask, 0) + send_src_metadata = (ordered_token.to(torch.int64) | (self.rank << 32)).contiguous() + # Exchange per-rank counts so every rank can size its receive buffer. + recv_counts = torch.empty_like(send_counts) + dist.all_to_all_single(recv_counts, send_counts) + sc = send_counts.tolist() + rc = recv_counts.tolist() + total_recv = int(sum(rc)) + recv_x = torch.empty((total_recv, H), dtype=x.dtype, device=dev) + recv_topk_idx = torch.empty((total_recv, int(idx.shape[1])), dtype=idx.dtype, device=dev) + recv_topk_weights = torch.empty((total_recv, int(idx.shape[1])), + dtype=p.topk_weights.dtype, device=dev) + recv_src_metadata = torch.empty((total_recv,), dtype=torch.int64, device=dev) + # Dispatch the uneven per-rank splits over the configured collective transport. + dist.all_to_all_single(recv_x, send_x, rc, sc) + dist.all_to_all_single(recv_topk_idx, send_topk_idx, rc, sc) + dist.all_to_all_single(recv_topk_weights, send_topk_weights, rc, sc) + dist.all_to_all_single(recv_src_metadata, send_src_metadata, rc, sc) + return types.SimpleNamespace( + recv_x=recv_x, combine_input=None, order=order, flat_token=flat_token, + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, recv_src_rank=recv_src_metadata >> 32, + recv_src_token=recv_src_metadata & ((1 << 32) - 1), send_counts=sc, recv_counts=rc, + T=T, H=H, topk=int(idx.shape[1]), total_recv=total_recv) + + def stage(self, p, h): + # No expert compute: the expert "output" is the received tokens as-is (the round-trip identity). + h.combine_input = h.recv_x + return None + + def combine(self, p, h): + # Reverse all-to-all: ship expert outputs back to their origin ranks (swap the split lists). + send_back = torch.empty((int(h.order.shape[0]), h.H), dtype=h.combine_input.dtype, + device=h.combine_input.device) + dist.all_to_all_single(send_back, h.combine_input.contiguous(), + h.send_counts, h.recv_counts) + # send_back is in send (sorted) order; invert the argsort to token-copy order. + copies = torch.empty_like(send_back) + copies[h.order] = send_back + # Sum one copy per destination rank under this reference's explicit unweighted contract. + out = torch.zeros((h.T, h.H), dtype=torch.float32, device=send_back.device) + out.index_add_(0, h.flat_token, copies.float()) + return out.to(p.x.dtype) + + def inspect_dispatch(self, p, h): + valid = h.recv_topk_idx >= 0 + expert_ids = torch.where( + valid, + h.recv_topk_idx + self.rank * self.experts_per_rank, + h.recv_topk_idx, + ) + return types.SimpleNamespace( + payload=h.recv_x, + expert_ids=expert_ids, + weights=h.recv_topk_weights.masked_fill(~valid, 0), + local_expert_counts=torch.bincount( + h.recv_topk_idx[valid], minlength=self.experts_per_rank + ), + ordering_contract="source-rank-major-stable-v1", + ) + + def combine_transformed(self, p, h, transformed): + h.combine_input = transformed.to(h.recv_x.dtype) + return self.combine(p, h) + + def recv_tokens(self, h): + return int(h.total_recv) diff --git a/experimental/CollectiveX/bench/ep_provenance.py b/experimental/CollectiveX/bench/ep_provenance.py new file mode 100644 index 0000000000..06c56259de --- /dev/null +++ b/experimental/CollectiveX/bench/ep_provenance.py @@ -0,0 +1,665 @@ +#!/usr/bin/env python3 +"""Backend implementation-provenance evidence and resource projection. + +The EP benchmark emitter builds and self-checks the implementation-provenance +evidence it embeds in each raw attempt document: which backend libraries were +loaded, that a deepep-v2 attempt carries internally consistent JIT cubins, that +a hybrid attempt realized its kernels, and so on. Those are cross-field facts a +JSON Schema cannot express, so they live in Python here, beside the executable +bench modules that build and self-check them. +""" +from __future__ import annotations + +import hashlib +import inspect +import json +import math +import os +import re +from pathlib import Path, PurePosixPath +from typing import Any, Iterable + +class ContractError(ValueError): + """A provenance payload differs from the CollectiveX emitter contract.""" + + +# Git run-identity fields an emitted attempt carries when produced under CI. +GIT_RUN_FIELDS = { + "artifact", "job", "ref", "repo", "run_attempt", "run_id", "source_sha", +} + + +def _finite_tree(value: Any, path: str = "$") -> None: + if isinstance(value, float) and not math.isfinite(value): + raise ContractError(f"{path} contains a non-finite number") + if isinstance(value, list): + for index, item in enumerate(value): + _finite_tree(item, f"{path}[{index}]") + elif isinstance(value, dict): + for key, item in value.items(): + _finite_tree(item, f"{path}.{key}") + + +def canonical_json_bytes(value: Any) -> bytes: + """Canonical finite JSON bytes for checksums and immutable artifacts.""" + _finite_tree(value) + try: + return json.dumps( + value, allow_nan=False, ensure_ascii=False, sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise ContractError(f"value is not canonical JSON: {exc}") from exc + + +# Backend implementation-provenance constants: the required evidence fields per +# backend and the pinned native versions a self-check compares against. They +# describe executable code, so they live beside the emitter that consumes them +# rather than in the neutral document/delivery validator. +DEEPEP_V2_JIT_KERNELS = frozenset({ + "barrier", "combine", "combine_reduce_epilogue", "dispatch", + "dispatch_copy_epilogue", +}) +DEEPEP_V2_V1_PROVENANCE = { + "deepep_version": "2.0.0", + "deepep_distribution_version": "2.0.0+fa8a9b1", + "deepep_commit": "fa8a9b16898204afd347c663b89e65ef87dc6ce6", + "deepep_tree": "29809e75c5874e6609dac4804e7b651d5226959f", + "deepep_pr": 605, + "deepep_fix_pr": 630, + "deepep_nccl_check_fix_pr": 640, + "deepep_nccl_check_commit": "93d0564188f7a0a6288c6e316484861b0efa042e", + "fmt_commit": "a4c7e17133ee9cb6a2f45545f6e974dd3c393efa", + "torch_version": "2.10.0+cu130", + "nccl_package_version": "2.30.4", + "nccl_version": "2.30.4", + "nvshmem_package_version": "3.3.9", +} +DEEPEP_V2_DISTRIBUTION_VERSIONS = frozenset({ + "2.0.0+fa8a9b1", "2.0.0+local", +}) +UCCL_DEPENDENCY_VERSIONS = { + "intervaltree": "3.1.0", + "nvidia-cuda-runtime-cu12": "12.9.79", + "sortedcontainers": "2.4.0", +} +REQUIRED_BACKEND_PROVENANCE = { + "deepep": ( + "deepep_version", "deepep_commit", "backend_lineage", "allow_mnnvl", + "mnnvl_comm", "mode", "num_nvl_bytes", "num_rdma_bytes", + "nvshmem_ibgda_nic_handler", + ), + "deepep-v2": ( + *DEEPEP_V2_V1_PROVENANCE, "api_signature_sha256", "loaded_libraries", + "jit_cubins", "jit_random_seed", "deterministic", "num_experts", + "tuning_num_experts", "allow_hybrid_mode", "gin_enabled", + "communication_backend", + ), + "deepep-hybrid": ( + "deepep_commit", "deepep_tree", "branch", "backend_lineage", + "loaded_libraries", "realized_config", "jit_kernel_keys", "jit_shared_objects", + ), + "uccl": ( + "uccl_version", "uccl_commit", "uccl_wrapper_commit", "backend_lineage", + "loaded_libraries", "uccl_dependency_versions", "mode", "num_nvl_bytes", + "num_rdma_bytes", + ), + "mori": ("mori_commit",), + "nccl-ep": ("nccl_version", "collective_library", "backend_lineage"), +} + + +def require_keyword(callable_object, keyword: str, *, api: str) -> None: + """Fail closed when a pinned native API does not expose a required control. + + A runtime-API contract guard: it verifies the loaded kernel build still accepts + the keyword the adapter drives (e.g. the BF16 ``use_fp8=False`` control), so a + version mismatch surfaces here instead of as an opaque call-site TypeError. + """ + try: + parameters = inspect.signature(callable_object).parameters + except (TypeError, ValueError) as exc: + raise ContractError(f"cannot inspect required native API {api}") from exc + if keyword not in parameters: + raise ContractError(f"required native API {api} omits {keyword!r}") + + +def resolve_deepep_mnnvl( + *, requested: bool, signature_parameters: Iterable[str], deepep_commit: str | None +) -> tuple[dict[str, bool], str]: + """Resolve one explicit DeepEP MNNVL API mode without signature fallbacks.""" + if not requested: + return {}, "not-requested" + if "allow_mnnvl" in set(signature_parameters): + return {"allow_mnnvl": True}, "explicit-allow-mnnvl" + raise ContractError( + f"requested DeepEP MNNVL is unsupported by commit {deepep_commit or 'unknown'}" + ) + + +def collective_kernel_generation(collective_library: Any) -> str: + """Return the public NCCL/RCCL implementation lineage.""" + if collective_library not in {"nccl", "rccl"}: + raise ContractError("reference collective library must be nccl or rccl") + return collective_library + + +def project_resource_profile(provenance: dict[str, Any]) -> dict[str, Any]: + """Project backend provenance into the canonical cross-backend resource vocabulary.""" + device_units = provenance.get("device_sms") or provenance.get("device_cus") + if provenance.get("num_sms") is not None: + kind, configured = "sm", provenance["num_sms"] + elif ( + provenance.get("block_num") is not None + and provenance.get("kernel_type") != "AsyncLL" + ): + kind, configured = "cu_block", provenance["block_num"] + else: + kind, configured = None, None + achieved = configured / device_units if configured and device_units else None + fixed = "fixed-kernel" in str(provenance.get("tuned_source", "")) + source = str(provenance.get("tuned_source", "")) + num_nvl_bytes = provenance.get("num_nvl_bytes") + num_rdma_bytes = provenance.get("num_rdma_bytes") + persistent_bytes = ( + (num_nvl_bytes or 0) + (num_rdma_bytes or 0) + if num_nvl_bytes is not None or num_rdma_bytes is not None + else provenance.get("heap_size") + ) + return { + "achieved_fraction": round(achieved, 4) if achieved else None, + "comm_units_kind": kind, + "configured_units": configured, + "conformance_class": ( + "not-applicable" if fixed else "backend-default" if "default" in source + else "pinned-upstream" + ), + "device_units": device_units, + "fixed_kernel": fixed, + "nonconforming": False, + "pareto_eligible": False, + "persistent_bytes": persistent_bytes, + "qps_per_rank": provenance.get("num_qps_per_rank"), + "requested_fraction": None, + "resource_class": "fixed-kernel" if fixed else "fixed-profile", + "target_achieved_within_tol": None, + "tolerance": 0.10, + "tuned_source": provenance.get("tuned_source"), + "warps_combine": provenance.get("combine_warps"), + "warps_dispatch": provenance.get("dispatch_warps"), + } + + +def _resolved_provenance_value(field: str, value: Any) -> bool: + if value is None or isinstance(value, (dict, list, tuple, set)) and not value: + return False + text = str(value).strip().lower() + if not text or text in {"unknown", "none", "null", "n/a", "?", "capture-failed"}: + return False + if "capture-failed" in text: + return False + if field.endswith("_commit") and ( + text in {"main", "hybrid-ep", "uccl", "pkg-uccl"} + or text.endswith(("-unknown", "-none", "-main", "-hybrid-ep")) + ): + return False + return True + + +def _content_evidence_is_valid(value: Any, required_roles: set[str]) -> bool: + if not isinstance(value, list) or not value: + return False + records: set[tuple[str, str]] = set() + roles: set[str] = set() + for item in value: + if not isinstance(item, dict) or set(item) != {"name", "role", "sha256"}: + return False + name, role, digest = item["name"], item["role"], item["sha256"] + if ( + not isinstance(name, str) + or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.+-]{0,159}", name) + or not isinstance(role, str) + or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.+-]{0,127}", role) + or not isinstance(digest, str) + or not re.fullmatch(r"[0-9a-f]{64}", digest) + or (role, name) in records + ): + return False + records.add((role, name)) + roles.add(role) + return required_roles <= roles + + +def _deepep_v2_jit_cubins_are_valid(value: Any) -> bool: + if not isinstance(value, list) or len(value) != len(DEEPEP_V2_JIT_KERNELS): + return False + cache_keys = [] + kernel_names = set() + for item in value: + if not isinstance(item, dict) or set(item) != { + "cache_key", "cubin_sha256", "sass_sha256", "source_sha256", + }: + return False + cache_key = item["cache_key"] + match = ( + re.fullmatch(r"kernel\.([A-Za-z0-9_+-]+)\.[0-9a-f]{32}", cache_key) + if isinstance(cache_key, str) + else None + ) + if ( + match is None + or any( + not isinstance(item[field], str) + or not re.fullmatch(r"[0-9a-f]{64}", item[field]) + for field in ("cubin_sha256", "sass_sha256", "source_sha256") + ) + ): + return False + cache_keys.append(cache_key) + kernel_names.add(match.group(1)) + return ( + cache_keys == sorted(set(cache_keys)) + and kernel_names == DEEPEP_V2_JIT_KERNELS + ) + + +HYBRID_REALIZED_CONFIG_FIELDS = { + "hidden_dim", "max_num_of_tokens_per_rank", "num_of_experts_per_rank", + "num_of_ranks_per_node", "num_of_nodes", "pad_multiple", + "num_of_tokens_per_chunk_preprocessing_api", + "num_of_threads_per_block_preprocessing_api", "num_of_blocks_preprocessing_api", + "num_of_blocks_permute", "num_of_blocks_unpermute", "token_data_type", + "num_of_stages_dispatch_api", "num_of_stages_permute_block_dispatch_api", + "num_of_in_flight_s2g_dispatch_api", + "num_of_in_flight_s2g_permute_block_dispatch_api", + "num_of_additional_in_flight_s2g_dispatch_api", + "num_of_tokens_per_chunk_dispatch_api", "num_of_blocks_dispatch_api", + "forward_dispatch_api", "device_side_sync_dispatch_api", + "num_of_stages_g2s_combine_api", "num_of_stages_s2g_combine_api", + "num_of_tokens_per_chunk_combine_api", "num_of_tokens_per_group_combine_api", + "num_of_blocks_combine_api", "num_of_additional_in_flight_s2g_combine_api", + "backward_combine_api", "device_side_sync_combine_api", +} +HYBRID_REALIZED_BOOL_FIELDS = { + "forward_dispatch_api", "device_side_sync_dispatch_api", "backward_combine_api", + "device_side_sync_combine_api", +} + + +def _hybrid_realized_config_is_valid(value: Any) -> bool: + if not isinstance(value, dict) or set(value) != HYBRID_REALIZED_CONFIG_FIELDS: + return False + for field, field_value in value.items(): + if field in HYBRID_REALIZED_BOOL_FIELDS: + if type(field_value) is not bool: + return False + elif field == "token_data_type": + if field_value not in {"UINT8", "UINT16"}: + return False + elif type(field_value) is not int or field_value < 0: + return False + return all(value[field] > 0 for field in ( + "hidden_dim", "max_num_of_tokens_per_rank", "num_of_experts_per_rank", + "num_of_ranks_per_node", "num_of_nodes", + )) + + +def _hybrid_kernel_keys_are_valid(value: Any) -> bool: + return ( + isinstance(value, list) + and len(value) == 3 + and len(set(value)) == 3 + and value == sorted(value) + and all( + isinstance(key, str) + and re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.+-]{0,511}", key) + for key in value + ) + ) + + +def _hybrid_jit_evidence_is_valid(value: Any, kernel_keys: Any) -> bool: + if not _hybrid_kernel_keys_are_valid(kernel_keys) or not isinstance(value, list): + return False + if len(value) != len(kernel_keys): + return False + rank_sets = [] + for expected_key, item in zip(kernel_keys, value): + if not isinstance(item, dict) or set(item) != {"kernel_key", "rank_artifacts"}: + return False + rank_artifacts = item["rank_artifacts"] + if item["kernel_key"] != expected_key or not isinstance(rank_artifacts, list): + return False + ranks = [] + for artifact in rank_artifacts: + if not isinstance(artifact, dict) or set(artifact) != {"bytes", "rank", "sha256"}: + return False + rank, digest, size = artifact["rank"], artifact["sha256"], artifact["bytes"] + if ( + type(rank) is not int + or rank < 0 + or not isinstance(digest, str) + or not re.fullmatch(r"[0-9a-f]{64}", digest) + or type(size) is not int + or size <= 0 + ): + return False + ranks.append(rank) + if not ranks or ranks != list(range(len(ranks))): + return False + rank_sets.append(ranks) + return all(ranks == rank_sets[0] for ranks in rank_sets) + + +def backend_provenance_issues(backend: str, provenance: dict[str, Any]) -> list[str]: + unknown = [ + field for field, value in provenance.items() + if isinstance(value, str) and value.strip().lower() == "unknown" + ] + unresolved = [ + field for field in REQUIRED_BACKEND_PROVENANCE.get(backend, ()) + if not _resolved_provenance_value(field, provenance.get(field)) + ] + if backend == "deepep": + mode = provenance.get("mnnvl_comm") + allow = provenance.get("allow_mnnvl") + valid_modes = { + "not-requested": False, + "explicit-allow-mnnvl": True, + } + if type(allow) is not bool or valid_modes.get(mode) is not allow: + unresolved.append("mnnvl_comm") + if provenance.get("backend_lineage") != "deepep-v1": + unresolved.append("backend_lineage") + if provenance.get("nvshmem_ibgda_nic_handler") not in { + "cpu", "gpu", "not-active", + }: + unresolved.append("nvshmem_ibgda_nic_handler") + if backend in {"deepep", "uccl"}: + mode = provenance.get("mode") + num_nvl_bytes = provenance.get("num_nvl_bytes") + num_rdma_bytes = provenance.get("num_rdma_bytes") + if mode not in {"normal", "low-latency"}: + unresolved.append("mode") + if type(num_nvl_bytes) is not int or num_nvl_bytes < 0: + unresolved.append("num_nvl_bytes") + if type(num_rdma_bytes) is not int or num_rdma_bytes < 0: + unresolved.append("num_rdma_bytes") + if mode == "normal" and (type(num_nvl_bytes) is not int or num_nvl_bytes <= 0): + unresolved.append("num_nvl_bytes") + if mode == "low-latency": + if num_nvl_bytes != 0: + unresolved.append("num_nvl_bytes") + if type(num_rdma_bytes) is not int or num_rdma_bytes <= 0: + unresolved.append("num_rdma_bytes") + if ( + type(provenance.get("num_max_tokens_per_rank")) is not int + or provenance["num_max_tokens_per_rank"] <= 0 + ): + unresolved.append("num_max_tokens_per_rank") + if backend == "deepep" and ( + type(provenance.get("num_qps_per_rank")) is not int + or provenance["num_qps_per_rank"] <= 0 + ): + unresolved.append("num_qps_per_rank") + if backend == "deepep-v2": + for field in ("num_experts", "tuning_num_experts"): + if type(provenance.get(field)) is not int or provenance[field] <= 0: + unresolved.append(field) + if not _deepep_v2_jit_cubins_are_valid(provenance.get("jit_cubins")): + unresolved.append("jit_cubins") + if provenance.get("jit_random_seed") != "collectivex-deepep-v2-fa8a9b1": + unresolved.append("jit_random_seed") + unresolved.extend( + field for field, expected in DEEPEP_V2_V1_PROVENANCE.items() + if field != "deepep_distribution_version" + and provenance.get(field) != expected + ) + if provenance.get("deepep_distribution_version") not in ( + DEEPEP_V2_DISTRIBUTION_VERSIONS + ): + unresolved.append("deepep_distribution_version") + policy = ( + provenance.get("allow_hybrid_mode"), + provenance.get("gin_enabled"), + provenance.get("communication_backend"), + ) + if policy not in { + (False, False, "nccl-device-lsa"), + (True, True, "nccl-gin"), + }: + unresolved.extend( + ("allow_hybrid_mode", "gin_enabled", "communication_backend") + ) + content_roles = { + "deepep-v2": {"deepep-extension", "nccl", "nvshmem"}, + "deepep-hybrid": {"deepep-extension", "deepep-hybrid-extension"}, + "uccl": { + "uccl-distribution", "uccl-wrapper", "intervaltree-distribution", + "sortedcontainers-distribution", "cuda-runtime", + }, + }.get(backend) + if content_roles is not None and not _content_evidence_is_valid( + provenance.get("loaded_libraries"), content_roles + ): + unresolved.append("loaded_libraries") + if backend in {"deepep-v2", "deepep-hybrid"} and not re.fullmatch( + r"[0-9a-f]{40}", str(provenance.get("deepep_tree", "")) + ): + unresolved.append("deepep_tree") + if backend == "deepep-hybrid" and provenance.get("backend_lineage") != "deepep-hybrid": + unresolved.append("backend_lineage") + if backend == "deepep-hybrid": + if not _hybrid_realized_config_is_valid(provenance.get("realized_config")): + unresolved.append("realized_config") + if not _hybrid_kernel_keys_are_valid(provenance.get("jit_kernel_keys")): + unresolved.append("jit_kernel_keys") + if not _hybrid_jit_evidence_is_valid( + provenance.get("jit_shared_objects"), provenance.get("jit_kernel_keys") + ): + unresolved.append("jit_shared_objects") + if backend == "uccl" and provenance.get("backend_lineage") != "uccl": + unresolved.append("backend_lineage") + if backend == "uccl" and provenance.get("uccl_dependency_versions") != ( + UCCL_DEPENDENCY_VERSIONS + ): + unresolved.append("uccl_dependency_versions") + if backend == "nccl-ep": + collective = provenance.get("collective_library") + if collective not in {"nccl", "rccl"}: + unresolved.append("collective_library") + if provenance.get("backend_lineage") != collective: + unresolved.append("backend_lineage") + if backend == "mori" and provenance.get("kernel_type") == "InterNodeV1": + expected = { + "block_num": 96, + "rdma_block_num": 64, + "dispatch_warps": 8, + "combine_warps": 8, + "num_qps": 1, + "use_external_inp_buf": True, + "gpus_per_node": 8, + } + unresolved.extend( + field for field, value in expected.items() + if provenance.get(field) != value + ) + for field, minimum in ( + ("num_nvl_bytes", 0), ("num_rdma_bytes", 0), + ("num_qps_per_rank", 1), + ): + if field in provenance and ( + type(provenance[field]) is not int or provenance[field] < minimum + ): + unresolved.append(field) + if "rdma_block_num" in provenance and ( + type(provenance["rdma_block_num"]) is not int + or provenance["rdma_block_num"] < 0 + ): + unresolved.append("rdma_block_num") + if "use_external_inp_buf" in provenance and type( + provenance["use_external_inp_buf"] + ) is not bool: + unresolved.append("use_external_inp_buf") + return sorted(set(unknown + unresolved)) + + +def provenance_complete( + provenance: dict[str, Any], backend: str, git_run: dict[str, Any] | None, + *, image_digest: Any, image_verified: Any, squash_sha256: Any, +) -> bool: + """Return whether backend provenance and run identity are fully resolved.""" + image = str(image_digest or "") + squash = str(squash_sha256 or "") + return ( + not backend_provenance_issues(backend, provenance) + and image_verified is True + and bool(re.fullmatch(r"sha256:[0-9a-f]{64}", image)) + and bool(re.fullmatch(r"[0-9a-f]{64}", squash)) + and isinstance(git_run, dict) + and all(git_run.get(field) for field in GIT_RUN_FIELDS) + ) + + +def content_manifest_evidence( + *, role: str, name: str, files: Iterable[tuple[str, str | os.PathLike[str]]] +) -> dict[str, str]: + """Hash a labeled file set without exposing any host path in provenance.""" + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.+-]{0,127}", role): + raise ContractError("content evidence role is invalid") + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.+-]{0,159}", name): + raise ContractError("content evidence name is invalid") + manifest: list[dict[str, Any]] = [] + labels: set[str] = set() + for label, raw_path in files: + logical = PurePosixPath(label) + if ( + not label + or logical.is_absolute() + or ".." in logical.parts + or label in labels + or any(ord(character) < 0x20 or ord(character) > 0x7E for character in label) + ): + raise ContractError("content evidence label is invalid or duplicated") + path = Path(raw_path) + if not path.is_file(): + raise ContractError("content evidence source is not a file") + digest = hashlib.sha256() + size = 0 + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + size += len(chunk) + labels.add(label) + manifest.append({"bytes": size, "label": label, "sha256": digest.hexdigest()}) + if not manifest: + raise ContractError("content evidence cannot be empty") + digest = hashlib.sha256( + canonical_json_bytes(sorted(manifest, key=lambda item: item["label"])) + ).hexdigest() + return {"name": name, "role": role, "sha256": digest} + + +# Source-built backend libraries whose raw build tree is projected to a stable +# public source identity (private cache keys / host paths are dropped). +SOURCE_BUILT_LIBRARY_ROLES = frozenset({ + "deepep-extension", "deepep-hybrid-extension", +}) + + +def series_provenance(provenance: dict[str, Any]) -> dict[str, Any]: + """Project stable semantic build identity while dropping private binaries and host paths.""" + projected = { + key: value for key, value in provenance.items() + if key not in {"jit_cache_key", "jit_shared_objects", "path", "sm_fraction"} + } + libraries = provenance.get("loaded_libraries") + if isinstance(libraries, list): + projected["loaded_libraries"] = [ + { + "name": item.get("name"), + "role": item.get("role"), + "source_tree": provenance.get("deepep_tree"), + } + if isinstance(item, dict) and item.get("role") in SOURCE_BUILT_LIBRARY_ROLES + else item + for item in libraries + ] + jit_cubins = provenance.get("jit_cubins") + if isinstance(jit_cubins, list): + projected["jit_cubins"] = [ + { + "cache_key": item.get("cache_key"), + "sass_sha256": item.get("sass_sha256"), + "source_sha256": item.get("source_sha256"), + } + if isinstance(item, dict) + else item + for item in jit_cubins + ] + return projected + + +def _sha256_json(value: Any) -> str: + """Canonical sha256 of a finite JSON value.""" + return hashlib.sha256(canonical_json_bytes(value)).hexdigest() + + +def backend_version(provenance: dict[str, Any]) -> str | None: + """Return the canonical public backend version from implementation provenance.""" + for field in ( + "deepep_version", "uccl_version", "nccl_version", + "mori_commit", "deepep_commit", + ): + value = provenance.get(field) + if value is not None and str(value).strip(): + return str(value)[:160] + return None + + +def public_series_config( + *, kernel_generation: Any, provenance: dict[str, Any], + resource_profile: dict[str, Any], resource_mode: Any, device_product: Any, +) -> dict[str, Any]: + """Project raw implementation facts into the exact public configuration fields.""" + generation = None if kernel_generation == "n-a" else kernel_generation + profile = "profile-" + _sha256_json(resource_profile)[:16] + return { + "backend": { + "generation": generation, + "version": backend_version(provenance), + }, + "resource": { + "mode": resource_mode, + "profile": profile, + "comm_units_kind": resource_profile.get("comm_units_kind"), + "configured_units": resource_profile.get("configured_units"), + }, + "system": {"label": str(device_product)[:160]}, + } + + +def public_series_config_sha256(config: dict[str, Any]) -> str: + """Commit the canonical public configuration projection into series identity.""" + return _sha256_json(config) + + +def routing_implementation_control_sha256(implementation: dict[str, Any]) -> str: + """Bind routing cohorts to the same static build/generator and non-treatment configuration.""" + provenance = implementation.get("provenance") + if not isinstance(provenance, dict): + raise ContractError("implementation provenance is unavailable") + semantic = series_provenance(provenance) + treatment_fields = { + "jit_cache_key", "jit_cubins", "jit_kernel_keys", "jit_shared_objects", + "local_experts", "num_experts", "path", "realized_config", "sm_fraction", + } + return _sha256_json({ + "kernel_generation": implementation.get("kernel_generation"), + "name": implementation.get("name"), + "provenance": { + key: value for key, value in semantic.items() + if key not in treatment_fields + }, + "resource_profile": implementation.get("resource_profile"), + }) diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py new file mode 100644 index 0000000000..55d07bdce2 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +"""CollectiveX UCCL adapter: native BF16 dispatch/combine over uccl_deepep.""" +from __future__ import annotations + +import importlib.metadata as metadata +import json +import os +from pathlib import Path +from pathlib import PurePosixPath +import sys + +import torch +import torch.distributed as dist +import ep_provenance +from ep_deepep_family import DeepEPFamilyBackend + +try: + import uccl + import uccl_deepep + from uccl_deepep import Buffer # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: uccl.ep import failed: {exc!r}", file=sys.stderr) + raise + + +def _uccl_version() -> str: + try: + return metadata.version("uccl") + except Exception: + return getattr(uccl, "__version__", "unknown") + + +def _uccl_dependency_versions() -> dict[str, str]: + versions = { + package: metadata.version(package) + for package in ep_provenance.UCCL_DEPENDENCY_VERSIONS + } + if versions != ep_provenance.UCCL_DEPENDENCY_VERSIONS: + raise RuntimeError( + "UCCL runtime dependency versions differ from the v1 contract" + ) + return versions + + +def _is_uccl_runtime_payload(name: str) -> bool: + path = PurePosixPath(name) + return ( + bool(path.parts) + and path.parts[0] in {"uccl", "uccl.libs"} + and "__pycache__" not in path.parts + and path.suffix != ".pyc" + ) + + +def _python_dependency_evidence(package: str, version: str) -> dict[str, str]: + distribution = metadata.distribution(package) + runtime_files = [] + for entry in distribution.files or (): + logical = PurePosixPath(entry.as_posix()) + path = Path(distribution.locate_file(entry)) + if ( + logical.parts + and logical.parts[0] == package + and "__pycache__" not in logical.parts + and logical.suffix != ".pyc" + and path.is_file() + ): + runtime_files.append((entry.as_posix(), path)) + return ep_provenance.content_manifest_evidence( + role=f"{package}-distribution", + name=f"{package}-{version}", + files=runtime_files, + ) + + +def _loaded_libcudart_evidence( + version: str, maps_path: Path = Path("/proc/self/maps") +) -> dict[str, str]: + distribution = metadata.distribution("nvidia-cuda-runtime-cu12") + candidates = { + Path(distribution.locate_file(entry)).resolve() + for entry in distribution.files or () + if PurePosixPath(entry.as_posix()).name.startswith("libcudart.so") + and Path(distribution.locate_file(entry)).is_file() + } + candidate_names = {path.name for path in candidates} + if not candidates or not candidate_names: + raise RuntimeError("pinned CUDA runtime distribution has no libcudart payload") + + loaded: set[Path] = set() + try: + mappings = maps_path.read_text().splitlines() + except OSError as exc: + raise RuntimeError("cannot inspect mapped UCCL runtime libraries") from exc + for mapping in mappings: + columns = mapping.split(maxsplit=5) + if len(columns) != 6: + continue + raw_path = columns[5] + deleted = raw_path.endswith(" (deleted)") + if deleted: + raw_path = raw_path.removesuffix(" (deleted)") + mapped = Path(raw_path) + if mapped.name not in candidate_names: + continue + if deleted or not mapped.is_file(): + raise RuntimeError( + "mapped libcudart is unavailable for content verification" + ) + resolved = mapped.resolve() + if resolved not in candidates: + raise RuntimeError( + "mapped libcudart is not owned by the pinned CUDA runtime package" + ) + loaded.add(resolved) + if len(loaded) != 1: + raise RuntimeError( + "expected exactly one mapped libcudart from the pinned CUDA runtime" + ) + return ep_provenance.content_manifest_evidence( + role="cuda-runtime", + name=f"nvidia-cuda-runtime-cu12-{version}", + files=[("libcudart.so", loaded.pop())], + ) + + +def _uccl_build_evidence( + version: str, dependency_versions: dict[str, str] +) -> list[dict[str, str]]: + distribution = metadata.distribution("uccl") + distribution_files = [ + (entry.as_posix(), distribution.locate_file(entry)) + for entry in distribution.files or () + if _is_uccl_runtime_payload(entry.as_posix()) + and Path(distribution.locate_file(entry)).is_file() + ] + wrapper_root = Path(uccl_deepep.__file__).resolve().parent + wrapper_files = [ + (path.relative_to(wrapper_root).as_posix(), path) + for path in wrapper_root.rglob("*.py") + if path.is_file() + ] + return [ + ep_provenance.content_manifest_evidence( + role="uccl-distribution", + name=f"uccl-{version}", + files=distribution_files, + ), + ep_provenance.content_manifest_evidence( + role="uccl-wrapper", + name="uccl-deepep-wrapper", + files=wrapper_files, + ), + _python_dependency_evidence("intervaltree", dependency_versions["intervaltree"]), + _python_dependency_evidence( + "sortedcontainers", dependency_versions["sortedcontainers"] + ), + _loaded_libcudart_evidence(dependency_versions["nvidia-cuda-runtime-cu12"]), + ] + + +def _require_cross_rank_equal(value, label: str) -> None: + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, value) + canonical = {json.dumps(item, sort_keys=True, separators=(",", ":")) for item in gathered} + if len(canonical) != 1: + raise RuntimeError(f"UCCL {label} differs across ranks") + + +def _normal_buffer_sizes(hidden: int, world_size: int) -> tuple[int, int]: + """Apply the wrapped DeepEP dispatch/combine sizing contract for this EP world.""" + hidden_bytes = hidden * torch.tensor([], dtype=torch.bfloat16).element_size() + configs = (Buffer.get_dispatch_config(world_size), Buffer.get_combine_config(world_size)) + num_nvl_bytes = max( + int(config.get_nvl_buffer_size_hint(hidden_bytes, world_size)) for config in configs + ) + num_rdma_bytes = max( + int(config.get_rdma_buffer_size_hint(hidden_bytes, world_size)) for config in configs + ) + if num_nvl_bytes <= 0 or num_rdma_bytes < 0: + raise RuntimeError("UCCL returned invalid normal-mode buffer size hints") + return num_nvl_bytes, num_rdma_bytes + + +class UCCLBackend(DeepEPFamilyBackend): + # uccl_deepep.Buffer is a DeepEP-API clone, so mode handling and dispatch/combine are + # shared in DeepEPFamilyBackend; only the native uccl_deepep buffer construction, its + # content-provenance, and process teardown are UCCL-specific. + name = "uccl" + _vendor = "UCCL" + + def create_buffer(self, spec): + # Local aliases keep the moved buffer-construction body byte-verbatim. + args, world_size, device = self.args, self.world_size, self.device + device_sms = torch.cuda.get_device_properties(device).multi_processor_count + if self.mode == "low-latency": + assert spec.max_tokens_per_rank <= 128 + ep_provenance.require_keyword( + Buffer.low_latency_dispatch, + "use_fp8", + api="uccl_deepep.Buffer.low_latency_dispatch", + ) + ep_provenance.require_keyword( + Buffer.low_latency_combine, + "use_logfmt", + api="uccl_deepep.Buffer.low_latency_combine", + ) + num_qps_per_rank = args.experts // world_size + num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint( + self.max_tokens_per_rank, args.hidden, world_size, args.experts + ) + self.buffer = Buffer( + self.group, + num_nvl_bytes=0, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=True, + num_qps_per_rank=num_qps_per_rank, + allow_nvlink_for_low_latency_mode=True, + ) + self.buffer.clean_low_latency_buffer( + self.max_tokens_per_rank, args.hidden, args.experts + ) + resource_provenance = { + "requested_num_sms": None, + "num_sms": None, + "sm_fraction": None, + "tuned_source": "uccl-low-latency-fixed-kernel", + "num_max_tokens_per_rank": self.max_tokens_per_rank, + "num_nvl_bytes": 0, + "num_rdma_bytes": num_rdma_bytes, + } + else: + ep_provenance.require_keyword( + Buffer.dispatch, + "async_finish", + api="uccl_deepep.Buffer.dispatch", + ) + ep_provenance.require_keyword( + Buffer.combine, + "async_finish", + api="uccl_deepep.Buffer.combine", + ) + num_nvl_bytes, num_rdma_bytes = _normal_buffer_sizes(args.hidden, world_size) + if world_size > args.scale_up_domain and num_rdma_bytes == 0: + raise RuntimeError("UCCL scale-out configuration returned no RDMA buffer") + self.buffer = Buffer(self.group, num_nvl_bytes, num_rdma_bytes) + num_sms = int(getattr(Buffer, "num_sms", args.num_sms)) + try: + Buffer.set_num_sms(num_sms) + except Exception as exc: # pragma: no cover - version dependent + raise RuntimeError( + f"UCCL did not apply requested num_sms={num_sms}: {exc!r}" + ) from exc + applied_num_sms = int(getattr(Buffer, "num_sms", num_sms)) + if applied_num_sms != num_sms: + raise RuntimeError( + f"UCCL num_sms mismatch: requested={num_sms} applied={applied_num_sms}" + ) + resource_provenance = { + "requested_num_sms": num_sms, + "num_sms": applied_num_sms, + "sm_fraction": applied_num_sms / device_sms, + "tuned_source": "uccl-default-num_sms", + "num_nvl_bytes": num_nvl_bytes, + "num_rdma_bytes": num_rdma_bytes, + } + version = _uccl_version() + dependency_versions = _uccl_dependency_versions() + loaded_libraries = _uccl_build_evidence(version, dependency_versions) + _require_cross_rank_equal(loaded_libraries, "installed content identities") + self.backend_provenance = { + "uccl_version": version, + "uccl_commit": os.environ.get("UCCL_COMMIT") or f"pkg-{version}", + "uccl_wrapper_commit": os.environ.get("UCCL_WRAPPER_COMMIT"), + "backend_lineage": "uccl", + "uccl_dependency_versions": dependency_versions, + "loaded_libraries": loaded_libraries, + "mode": self.mode, + "dispatch_dtype": "bf16", + "combine_dtype": "bf16", + "resource_mode": "fixed-profile", + "device_sms": device_sms, + **resource_provenance, + } + + def finalize(self, rc): + # UCCL's proxy teardown can crash after results are written; preserve the real rc. + try: + dist.barrier() + except Exception: + pass + sys.stdout.flush() + sys.stderr.flush() + os._exit(rc if 0 <= rc <= 255 else 1) diff --git a/experimental/CollectiveX/bench/eplb.py b/experimental/CollectiveX/bench/eplb.py new file mode 100644 index 0000000000..b1479da9f1 --- /dev/null +++ b/experimental/CollectiveX/bench/eplb.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""CollectiveX — EPLB (Expert-Parallel Load Balancer), the DeepSeek-style remedy for +skewed (zipf) expert load. + +Under skewed routing, the ranks hosting hot logical experts receive far more token-copies +than the rest; dispatch/combine latency is gated by that busiest rank (the cross-rank MAX +the harness measures), so the whole collective stalls on it. EPLB REPLICATES hot experts +onto extra physical slots and PLACES the slots so every rank carries ~equal load. + +This module is backend-agnostic: it is purely a transform of the deterministic routing +trace. The trick that keeps every adapter unchanged — DeepEP/MoRI both route expert i to +rank `i // experts_per_rank` (contiguous block placement) — is to number the physical slots +RANK-MAJOR (rank r owns physical ids [r*spp, (r+1)*spp)), so the standard contiguous mapping +reproduces EPLB's balanced placement. The harness then runs with `experts = num_physical` +and the remapped (physical) trace; nothing else changes. + + num_physical = num_logical + redundant (redundant rounded up to a multiple of ep_size) + build_plan(): greedy replicate-by-load + equal-cardinality balanced packing onto ep_size ranks + remap_idx(): each token's logical targets -> physical replicas, spread by global token id + +Pure-Python planner (no torch) so it unit-tests on a login node; remap_idx needs torch. +""" +from __future__ import annotations + +import hashlib +import json + + +def physical_count(num_logical: int, num_redundant: int, ep_size: int) -> int: + """num_logical + redundant, with redundant rounded UP to a multiple of ep_size so the + physical experts divide evenly across ranks (symmetric dispatch).""" + r = ((max(0, num_redundant) + ep_size - 1) // ep_size) * ep_size + return num_logical + r + + +def _contiguous_rank_load(logical_load, ep_size): + """Per-rank received load WITHOUT EPLB: logical experts placed contiguously + (experts_per_rank = num_logical/ep_size), so rank r carries its block's total.""" + n = len(logical_load) + per = n // ep_size + return [sum(logical_load[r * per:(r + 1) * per]) for r in range(ep_size)] + + +def build_plan(logical_load, num_physical: int, ep_size: int) -> dict: + """logical_load: list[float] length num_logical (token-copies per logical expert). + Returns the replication+placement plan (all pure-Python lists) + before/after balance.""" + num_logical = len(logical_load) + assert num_physical >= num_logical, "num_physical must be >= num_logical" + assert num_physical % ep_size == 0, "num_physical must divide ep_size" + assert num_logical % ep_size == 0, "num_logical must divide ep_size" + spp = num_physical // ep_size # physical slots per rank (fixed) + + # 1) Replica allocation — start one slot per logical expert, then hand each redundant + # slot to the expert with the highest CURRENT per-replica load (greedy min-max). + replicas = [1] * num_logical + for _ in range(num_physical - num_logical): + best, best_lps = 0, -1.0 + for e in range(num_logical): + lps = logical_load[e] / replicas[e] + if lps > best_lps: + best, best_lps = e, lps + replicas[best] += 1 + + # 2) Slots = (per-replica load, logical expert), one per replica. + slots = [] + for e in range(num_logical): + lps = logical_load[e] / replicas[e] + slots.extend((lps, e) for _ in range(replicas[e])) + + # 3) Balanced packing into ep_size bins of EQUAL cardinality (spp each), minimizing the + # max per-rank load: heaviest slot first -> least-loaded rank that still has capacity. + slots.sort(reverse=True) + rank_slots = [[] for _ in range(ep_size)] + rank_load = [0.0] * ep_size + for lps, e in slots: + r = min((r for r in range(ep_size) if len(rank_slots[r]) < spp), + key=lambda r: rank_load[r]) + rank_slots[r].append(e) + rank_load[r] += lps + + # 4) Rank-major physical numbering -> contiguous placement == this balanced placement. + phys2log, rank_of_phys = [], [] + for r in range(ep_size): + for e in rank_slots[r]: + phys2log.append(e) + rank_of_phys.append(r) + log2phys = [[] for _ in range(num_logical)] + for pid, e in enumerate(phys2log): + log2phys[e].append(pid) + + before = _contiguous_rank_load(logical_load, ep_size) + total = sum(logical_load) or 1.0 + mean = total / ep_size + return { + "num_logical": num_logical, "num_physical": num_physical, "ep_size": ep_size, + "slots_per_rank": spp, "replicas": replicas, "max_replicas": max(replicas), + "phys2log": phys2log, "rank_of_phys": rank_of_phys, "log2phys": log2phys, + "rank_load_after": rank_load, "rank_load_before": before, + # imbalance = busiest rank / mean (1.0 = perfect). This is the number EPLB cuts. + "imbalance_before": max(before) / mean, "imbalance_after": max(rank_load) / mean, + "replicated_experts": sum(1 for r in replicas if r > 1), + } + + +def mapping_hash(plan: dict) -> str: + """Hash the placement fields that fully determine the logical-to-physical remap.""" + payload = { + "phys2log": plan["phys2log"], + "rank_of_phys": plan["rank_of_phys"], + "replicas": plan["replicas"], + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + + +def remap_rows(indices: list[list[int]], plan: dict) -> list[list[int]]: + """Pure-Python equivalent of remap_idx for contract verification.""" + replicas = plan["log2phys"] + return [ + [replicas[expert][token % len(replicas[expert])] for expert in row] + for token, row in enumerate(indices) + ] + + +def remap_idx(idx_logical, plan): + """idx_logical: torch [gt, topk] int64 logical-expert ids (global trace). + Returns idx_physical [gt, topk]: each token's logical target -> one of that expert's + physical replicas, SPREAD by global token id (row) so a hot expert's tokens fan out + across its replicas (= across ranks). Replicas of distinct logical experts are disjoint, + so a token's top-k physical ids stay distinct (dispatch invariant preserved).""" + import torch + replicas = plan["replicas"] + num_logical = len(replicas) + max_rc = plan["max_replicas"] + rc = torch.tensor(replicas, dtype=torch.int64) + # padded [num_logical, max_rc] table of physical ids (pad with replica 0; never indexed + # past rc[e] because the replica index is taken mod rc[e]). + padded = torch.zeros(num_logical, max_rc, dtype=torch.int64) + for e, phys in enumerate(plan["log2phys"]): + for k in range(max_rc): + padded[e, k] = phys[k] if k < len(phys) else phys[0] + gt = idx_logical.shape[0] + rows = torch.arange(gt, dtype=torch.int64).unsqueeze(1) # [gt,1] global token id + e = idx_logical.to(torch.int64) # [gt,topk] + ridx = rows % rc[e] # [gt,topk] replica index + return padded[e, ridx] # [gt,topk] physical ids + + +# --------------------------------------------------------------------------- self-test +if __name__ == "__main__": + # Synthetic zipf load (popularity ∝ 1/(e+1)) — the case EPLB targets. No torch needed. + import sys + NUM_LOGICAL, EP, REDUNDANT = 256, 8, 32 + load = [1.0 / (e + 1) for e in range(NUM_LOGICAL)] + nphys = physical_count(NUM_LOGICAL, REDUNDANT, EP) + plan = build_plan(load, nphys, EP) + print(f"num_logical={NUM_LOGICAL} ep={EP} num_physical={nphys} slots/rank={plan['slots_per_rank']}") + print(f"replicated experts={plan['replicated_experts']} max_replicas={plan['max_replicas']} " + f"(hottest expert 0 replicas={plan['replicas'][0]})") + print(f"per-rank load BEFORE (contiguous): {[round(x,3) for x in plan['rank_load_before']]}") + print(f"per-rank load AFTER (EPLB): {[round(x,3) for x in plan['rank_load_after']]}") + print(f"imbalance (max/mean) BEFORE={plan['imbalance_before']:.2f}x AFTER={plan['imbalance_after']:.2f}x") + # Gates: equal slot cardinality, every logical expert placed, big imbalance cut. + assert all(plan["replicas"][e] >= 1 for e in range(NUM_LOGICAL)) + assert sum(plan["replicas"]) == nphys + assert len(plan["phys2log"]) == nphys + assert all(len(plan["log2phys"][e]) == plan["replicas"][e] for e in range(NUM_LOGICAL)) + # rank-major numbering => contiguous block per rank => rank_of_phys is non-decreasing + assert plan["rank_of_phys"] == sorted(plan["rank_of_phys"]) + assert plan["imbalance_after"] < plan["imbalance_before"], "EPLB must reduce imbalance" + assert plan["imbalance_after"] < 1.30, f"EPLB should get within ~30% of perfect, got {plan['imbalance_after']:.2f}" + # remap (if torch present): distinctness + balanced receive on a sampled zipf trace. + try: + import torch + g = torch.Generator().manual_seed(0) + p = torch.tensor(load) + p = (p / p.sum()).expand(4096, NUM_LOGICAL) + idx_l = torch.multinomial(p, 8, replacement=False, generator=g).to(torch.int64) + idx_p = remap_idx(idx_l, plan) + assert idx_p.shape == idx_l.shape + # top-k physical ids distinct per token + assert all(len(set(row.tolist())) == 8 for row in idx_p), "physical top-k must stay distinct" + spp = plan["slots_per_rank"] + recv_before = [0] * EP + recv_after = [0] * EP + per_log = NUM_LOGICAL // EP + for row_l, row_p in zip(idx_l.tolist(), idx_p.tolist()): + for e in row_l: + recv_before[e // per_log] += 1 + for pid in row_p: + recv_after[pid // spp] += 1 + ib = max(recv_before) / (sum(recv_before) / EP) + ia = max(recv_after) / (sum(recv_after) / EP) + print(f"sampled-trace receive imbalance BEFORE={ib:.2f}x AFTER={ia:.2f}x") + assert ia < ib and ia < 1.35, "remap must balance per-rank receive load" + print("remap self-test: OK") + except ImportError: + print("(torch absent — skipped remap self-test; planner gates passed)") + print("EPLB self-test: PASS") + sys.exit(0) diff --git a/experimental/CollectiveX/bench/make_workloads.py b/experimental/CollectiveX/bench/make_workloads.py new file mode 100644 index 0000000000..f92d20c923 --- /dev/null +++ b/experimental/CollectiveX/bench/make_workloads.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Generate canonical serialized workloads. Runs the stdlib counter generator for +each (routing, global_tokens) in a ladder and writes .npz + .manifest.json into a +dir that runs then consume via `run_ep.py --workload-dir`. One trace is emitted per global-token +count because global token count is part of workload identity. + + python3 bench/make_workloads.py --out-dir /path/to/cx_workloads \\ + --routing uniform --ep 8 --hidden 7168 --topk 8 --experts 256 --seed 67 \\ + --tokens-ladder "1 2 4 8 16 32 64 128 256 512" + +Or by the named workload in configs/workloads.yaml. Explicit dimension flags still override it: + + python3 bench/make_workloads.py --out-dir /path/to/cx_workloads --workload deepseek-v3 --routing uniform --ep 8 + +--id-only prints the content-bound workload_id per ladder point without torch/numpy: + + python3 bench/make_workloads.py --workload deepseek-v3 --ep 8 --id-only + +Generate every routing the suites need by running once per --routing. Idempotent (same id => same +file). The dir is the cross-hardware artifact: copy it to each cluster so all consume identical bytes. +""" +from __future__ import annotations + +import argparse +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import workload as wl # noqa: E402 + +# Repo root holds configs/ (this file is in tests/). Used only for --workload name resolution. +_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def resolve_manifest(name): + """Look a workload name up in configs/workloads.yaml and return (hidden, topk, experts). + Searches synthetic + model_derived; expert count = `experts` or (for model-derived) `routed_experts`. + Raises SystemExit with the known names if the manifest is absent. Pure PyYAML + stdlib.""" + import yaml + path = os.path.join(_REPO, "configs", "workloads.yaml") + with open(path) as handle: + cfg = yaml.safe_load(handle) + known = [] + for section in ("synthetic", "model_derived"): + sec = cfg.get(section) or {} + known += list(sec) + m = sec.get(name) + if m is None: + continue + experts = m.get("experts", m.get("routed_experts")) + if m.get("hidden") is None or m.get("topk") is None or experts is None: + raise SystemExit(f"workload '{name}' is missing hidden/topk/experts in {path}") + return int(m["hidden"]), int(m["topk"]), int(experts) + raise SystemExit(f"unknown --workload '{name}'; known: {sorted(known)}") + + +def main() -> int: + ap = argparse.ArgumentParser(description="Generate canonical CollectiveX workloads") + ap.add_argument("--out-dir", help="required unless --id-only") + ap.add_argument("--workload", help="named manifest in configs/workloads.yaml (sets hidden/topk/experts)") + ap.add_argument("--routing", default="uniform", choices=["uniform"]) + ap.add_argument("--ep", type=int, required=True, help="ep_size (global_tokens = T * ep)") + ap.add_argument("--hidden", type=int, help="override (default 7168, or the --workload's hidden)") + ap.add_argument("--topk", type=int, help="override (default 8, or the --workload's topk)") + ap.add_argument("--experts", type=int, help="override (default 256, or the --workload's experts)") + ap.add_argument("--seed", type=int, default=67) + ap.add_argument("--tokens-ladder", default="1 2 4 8 16 32 64 128 256 512") + ap.add_argument("--id-only", action="store_true", + help="print content-bound workload_id per point without torch/numpy") + a = ap.parse_args() + + # Resolve dims: a named --workload supplies defaults; explicit --hidden/--topk/--experts override + # per field. With neither, fall back to the v1 DeepSeek dimensions (7168/8/256). + base_h, base_t, base_e = (7168, 8, 256) + if a.workload: + base_h, base_t, base_e = resolve_manifest(a.workload) + hidden = a.hidden if a.hidden is not None else base_h + topk = a.topk if a.topk is not None else base_t + experts = a.experts if a.experts is not None else base_e + + if not a.id_only and not a.out_dir: + ap.error("--out-dir is required unless --id-only") + + raw_ladder = [int(token) for token in a.tokens_ladder.replace(",", " ").split()] + if (a.ep <= 0 or min(hidden, topk, experts) <= 0 or topk > experts or experts % a.ep + or not raw_ladder or any(token <= 0 for token in raw_ladder) + or len(raw_ladder) != len(set(raw_ladder))): + ap.error("shape, EP, and token ladder must be positive, divisible, and unique") + ladder = sorted(raw_ladder) + epr = experts // a.ep + label = f"workload={a.workload} " if a.workload else "" + + if a.id_only: + # The stdlib counter generator derives the same content-bound ID on every runtime. + made = [] + for T in ladder: + gt = T * a.ep + wid = wl.compute_workload_id(a.routing, hidden, topk, experts, a.ep, gt, a.seed) + made.append((T, gt, wid)) + print(f" T={T:<5} gt={gt:<6} routing={a.routing} -> {wid}") + print(f"{label}id-only: {len(made)} workload_id(s) " + f"(hidden={hidden} topk={topk} experts={experts} ep={a.ep} routing={a.routing} seed={a.seed})") + return 0 + + os.makedirs(a.out_dir, exist_ok=True) + made = [] + for T in ladder: + gt = T * a.ep + idx, w, man = wl.build_workload(hidden, topk, experts, a.routing, gt, a.seed, epr) + wid = wl.save_workload(a.out_dir, idx, w, man) + made.append((T, gt, wid)) + print(f" T={T:<5} gt={gt:<6} routing={a.routing} -> {wid} " + f"(trace sha {man['checksums']['trace'][:12]})") + print(f"{label}wrote {len(made)} canonical workloads to {a.out_dir} (routing={a.routing}, ep={a.ep})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/bench/routing.py b/experimental/CollectiveX/bench/routing.py new file mode 100644 index 0000000000..08d2cddae2 --- /dev/null +++ b/experimental/CollectiveX/bench/routing.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""CollectiveX — deterministic, platform-independent MoE routing trace. + +Fair-comparison fix #1: routing (per-token expert IDs + gate weights) is generated +ONCE from a fixed seed over the *global* token batch, indexed by global token id, and +is identical on every SKU for the same (seed, routing, global_tokens, experts, top-k). +Each rank materializes its slice `[rank*T,(rank+1)*T)`. Activations +are per-rank (same rank ⇒ same x on any platform), so a given global token id has +identical activation everywhere without materializing a global activation tensor. + +The v1 suite uses a single routing distribution: + + * uniform — top-k distinct experts drawn uniformly per token. The DEFAULT. + Expected fan-out for top-k=8, 256 experts, EP8 (32 experts/rank) ≈ + 8·(1 − C(224,8)/C(256,8)) ≈ 5.3 ranks/token. Load ~ Poisson. + +Always publish the realized fan-out so the workload is never misread again +(`routing_stats`). +""" +from __future__ import annotations + +import hashlib + +import torch + +ACTIVATION_GENERATOR = "collectivex-activation-counter-v4" +SOURCE_ID_BITS = 32 +SOURCE_CHECKSUM_BITS = 16 +SOURCE_ID_COLUMNS = SOURCE_ID_BITS + SOURCE_CHECKSUM_BITS +SOURCE_ID_CONTRACT = "bounded-sign-bit-source-v1" + + +def build_global_routing( + global_tokens: int, + experts: int, + topk: int, + routing: str, + seed: int, + *, + token_offset: int = 0, +): + """Return one byte-stable counter-generated routing window on CPU.""" + import workload + + indices, weights = workload.canonical_routing_rows( + int(global_tokens), + int(experts), + int(topk), + routing, + int(seed), + token_offset=token_offset, + ) + return ( + torch.tensor(indices, dtype=torch.int64), + torch.tensor(weights, dtype=torch.float32), + ) + + +def rank_slice(idx, weights, rank: int, tokens_per_rank: int): + lo = rank * tokens_per_rank + return idx[lo:lo + tokens_per_rank].contiguous(), weights[lo:lo + tokens_per_rank].contiguous() + + +def rank_activations(tokens: int, hidden: int, seed: int, rank: int, device, + dtype=torch.bfloat16): + """Exact counter-derived inputs with a quantization-safe source-token prefix.""" + source = torch.arange(tokens, device=device, dtype=torch.int64) + rank * tokens + return activations_for_source_ids(source, hidden, seed, dtype) + + +def activations_for_source_ids(source, hidden: int, seed: int, dtype=torch.bfloat16): + """Materialize canonical activations for arbitrary global source-token IDs.""" + if hidden < SOURCE_ID_COLUMNS: + raise ValueError(f"hidden must be at least {SOURCE_ID_COLUMNS}") + source = source.to(torch.int64) + column = torch.arange(hidden, device=source.device, dtype=torch.int64) + values = (source[:, None] * 131 + column[None, :] * 17 + int(seed) * 19) % 257 - 128 + output = values.to(dtype).mul_(1 / 64) + if bool((source < 0).any().item()) or bool((source >= (1 << SOURCE_ID_BITS)).any().item()): + raise ValueError("source token ID is outside the bounded identity contract") + source_columns = torch.arange(SOURCE_ID_BITS, device=source.device, dtype=torch.int64) + source_bits = ((source[:, None] >> source_columns[None, :]) & 1) * 2 - 1 + checksum = (source * 0x9E37 + int(seed) * 0xA24B) & ((1 << SOURCE_CHECKSUM_BITS) - 1) + checksum_columns = torch.arange( + SOURCE_CHECKSUM_BITS, device=source.device, dtype=torch.int64 + ) + checksum_bits = ((checksum[:, None] >> checksum_columns[None, :]) & 1) * 2 - 1 + # Magnitude one sits inside the ordinary [-2, 2] activation range, so the identity cannot set + # an FP8 block scale. Decode depends only on sign and remains stable after dequantization. + output[:, :SOURCE_ID_BITS] = source_bits.to(dtype) + output[:, SOURCE_ID_BITS:SOURCE_ID_COLUMNS] = checksum_bits.to(dtype) + return output + + +def decode_source_ids(payload, seed: int): + """Decode and validate source IDs carried by rank_activations.""" + if payload.ndim != 2 or payload.shape[1] < SOURCE_ID_COLUMNS: + raise ValueError("received payload cannot carry the source-token prefix") + prefix = payload[:, :SOURCE_ID_COLUMNS].float() + if not bool(torch.isfinite(prefix).all().item()) or bool((prefix.abs() < 0.25).any().item()): + raise ValueError("received source-token prefix is not quantization-stable") + bits = prefix >= 0 + powers = 1 << torch.arange(SOURCE_ID_BITS, device=payload.device, dtype=torch.int64) + source = (bits[:, :SOURCE_ID_BITS].to(torch.int64) * powers).sum(dim=1) + checksum_powers = 1 << torch.arange( + SOURCE_CHECKSUM_BITS, device=payload.device, dtype=torch.int64 + ) + observed_checksum = ( + bits[:, SOURCE_ID_BITS:SOURCE_ID_COLUMNS].to(torch.int64) * checksum_powers + ).sum(dim=1) + checksum = (source * 0x9E37 + int(seed) * 0xA24B) & ( + (1 << SOURCE_CHECKSUM_BITS) - 1 + ) + if not torch.equal(checksum, observed_checksum): + raise ValueError("received source-token checksum differs") + return source + + +def routing_locality(idx, experts_per_rank: int, ep_size: int, tokens_per_rank: int, + gpus_per_node: int, scale_up_domain: int = None) -> dict: + """Locality of rank-deduplicated payload copies under packed placement.""" + import torch as _t + gt = idx.shape[0] + assignments = (idx // experts_per_rank).clamp(max=ep_size - 1) + destinations = _t.zeros((gt, ep_size), dtype=_t.bool) + destinations.scatter_(1, assignments, True) + token, dest = destinations.nonzero(as_tuple=True) + src = (token // max(1, tokens_per_rank)).clamp(max=ep_size - 1) + sud = scale_up_domain or (gpus_per_node * ep_size) # default: all one domain + phys = _t.arange(ep_size, dtype=_t.int64) + pd, ps = phys[dest], phys[src] + local = (dest == src) + same_node = (pd // gpus_per_node) == (ps // gpus_per_node) + same_dom = (pd // sud) == (ps // sud) + n = dest.numel() + return { + "placement": "packed", + "local_rank_fraction": float(local.float().mean()), + "same_node_fraction": float(same_node.float().mean()), + "same_scaleup_domain_fraction": float(same_dom.float().mean()), + "cross_node_fraction": float((~same_node).float().mean()), + "cross_domain_fraction": float((~same_dom).float().mean()), + "gpus_per_node": gpus_per_node, "scale_up_domain": sud, "copies": int(n), + } + + +def routing_stats(idx, experts: int, experts_per_rank: int, weights=None) -> dict: + """Realized routing properties for the GLOBAL trace — published per point so the + fan-out / load can never be silently misread. idx is the global [gt, topk] tensor; + weights the matching [gt, topk] gate weights (hashed too for workload identity). + """ + ep = max(1, experts // max(1, experts_per_rank)) + ranks = (idx // experts_per_rank) # [gt, topk] destination rank per assignment + # unique destination ranks per token (fan-out) + onehot = torch.zeros(idx.shape[0], ep, dtype=torch.bool) + onehot.scatter_(1, ranks.clamp(max=ep - 1), True) + fanout = onehot.sum(dim=1) # [gt] + hist = torch.bincount(fanout, minlength=ep + 1)[1:ep + 1].tolist() # counts for fan-out 1..ep + load = torch.bincount(idx.reshape(-1), minlength=experts).float() + # Keep expert assignments (compute load) separate from rank-deduplicated payload copies + # (network load). Conflating them overstates traffic when two experts share a rank. + assignment_load = torch.bincount( + ranks.reshape(-1).clamp(max=ep - 1), minlength=ep + ).float() + payload_load = onehot.sum(dim=0).float() + # One-number imbalance summaries so a row is self-describing for the distribution-sensitivity + # suite (no need to read the full histograms): CV = std/mean of the load; hotspot_ratio = + # worst expert load over the mean. Zipf should be more concentrated than uniform. + def _cv(t): + m = float(t.mean()) + return float(t.std(unbiased=False) / m) if m > 0 else 0.0 + expert_load_cv = _cv(load) + assignment_rank_cv = _cv(assignment_load) + payload_rank_cv = _cv(payload_load) + hotspot_ratio = float(load.max() / load.mean()) if float(load.mean()) > 0 else 0.0 + # Empty experts capture compute skew; empty destination ranks capture network skew. + empty_expert_count = int((load == 0).sum()) + empty_rank_count = int((payload_load == 0).sum()) + # SHA-256 workload identity over both topk_idx and gate weights: a chart + # point's routing is provably identical across SKUs only if both hashes match. + idx_bytes = idx.to(torch.int32).cpu().numpy().tobytes() + idx_hash = hashlib.sha256(idx_bytes).hexdigest() + if weights is not None: + w_bytes = weights.to(torch.float32).cpu().numpy().tobytes() + w_hash = hashlib.sha256(w_bytes).hexdigest() + routing_hash = hashlib.sha256(idx_bytes + w_bytes).hexdigest() + else: + w_hash, routing_hash = None, idx_hash + return { + "fanout_mean": float(fanout.float().mean()), + "fanout_min": int(fanout.min()), "fanout_max": int(fanout.max()), + "fanout_hist": hist, # index k-1 = #tokens with fan-out k + "expert_assignments_per_rank": [int(x) for x in assignment_load.tolist()], + "payload_copies_per_rank": [int(x) for x in payload_load.tolist()], + "routed_copies": int(fanout.sum()), # total (token, dest-rank) pairs + "expert_load_min": int(load.min()), "expert_load_max": int(load.max()), + "expert_load_mean": float(load.mean()), "expert_load_cv": expert_load_cv, + "expert_assignment_rank_cv": assignment_rank_cv, + "payload_rank_cv": payload_rank_cv, "hotspot_ratio": hotspot_ratio, + "empty_expert_count": empty_expert_count, "empty_rank_count": empty_rank_count, + "routing_hash": routing_hash, "idx_hash": idx_hash, "weights_hash": w_hash, + } + + +# --------------------------------------------------------------------------- self-test +if __name__ == "__main__": + import sys + E, TOPK, EPR, GT = 256, 8, 32, 4096 + ui, _ = build_global_routing(GT, E, TOPK, "uniform", 67) + assert all(len(set(row.tolist())) == TOPK for row in ui[:16]) + uniform = routing_stats(ui, E, EPR) + assert uniform["hotspot_ratio"] >= 1.0 + dev = torch.device("cpu") + first = rank_activations(8, 256, 67, 0, dev, dtype=torch.float32) + second = rank_activations(8, 256, 67, 0, dev, dtype=torch.float32) + assert torch.equal(first, second) and torch.isfinite(first).all() + print("routing self-test: PASS") + sys.exit(0) diff --git a/experimental/CollectiveX/bench/run_ep.py b/experimental/CollectiveX/bench/run_ep.py new file mode 100644 index 0000000000..ae43f4dbc5 --- /dev/null +++ b/experimental/CollectiveX/bench/run_ep.py @@ -0,0 +1,562 @@ +#!/usr/bin/env python3 +"""CollectiveX v1 EP benchmark entrypoint for torchrun or rank environments.""" + +from __future__ import annotations + +import argparse +import ctypes +import hashlib +import hmac +import json +import os +import platform +import re +import shlex +import socket +import subprocess +import sys + +# Make the sibling bench/ modules importable when run as `bench/run_ep.py` under +# torchrun (it executes the file as __main__, not as a package). +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path[:0] = [HERE, os.path.dirname(HERE)] + +import ep_harness # noqa: E402 (stdlib-only; safe before torch) +import identity # noqa: E402 + + +ALLOCATION_STRATUM_CONTRACT = "collectivex-allocation-stratum-v1" +PRIVATE_FABRIC_ENV = { + "ib_gid_index": "CX_IB_GID_INDEX", + "rdma_devices": "CX_RDMA_DEVICES", + "rdma_service_level": "CX_RDMA_SERVICE_LEVEL", + "rdma_traffic_class": "CX_RDMA_TRAFFIC_CLASS", + "socket_ifname": "CX_SOCKET_IFNAME", +} + + +def _numeric_version(command: list[str]) -> str | None: + try: + result = subprocess.run( + command, capture_output=True, check=False, text=True, timeout=10 + ) + except (OSError, subprocess.TimeoutExpired): + return None + if result.returncode != 0: + return None + match = re.search(r"\b[0-9]+(?:\.[0-9]+){1,3}\b", result.stdout) + return match.group(0) if match else None + + +def _loaded_collective_version() -> str | None: + try: + with open("/proc/self/maps", encoding="utf-8") as handle: + paths = { + os.path.realpath(line.rstrip().split()[-1]) + for line in handle + if any(name in line for name in ("libnccl.so", "librccl.so")) + and os.path.isfile(line.rstrip().split()[-1]) + } + if len(paths) != 1: + return None + version = ctypes.c_int() + library = ctypes.CDLL(paths.pop()) + if library.ncclGetVersion(ctypes.byref(version)) != 0: + return None + return ep_harness.format_collective_version(version.value) + except (AttributeError, OSError): + return None + + +def _runtime_fingerprint( + torch, device, *, machine: str, vendor: str, arch: str +) -> dict: + """Return strict runtime facts without hosts, addresses, UUIDs, or paths.""" + properties = torch.cuda.get_device_properties(device) + if vendor == "nvidia": + driver = _numeric_version( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"] + ) + runtime_kind, runtime_version, collective_kind = ( + "cuda", + torch.version.cuda, + "nccl", + ) + else: + driver = _numeric_version(["rocm-smi", "--showdriverversion"]) + runtime_kind, runtime_version, collective_kind = ( + "hip", + torch.version.hip, + "rccl", + ) + return { + "accelerator_runtime": {"kind": runtime_kind, "version": runtime_version}, + "collective_library": { + "kind": collective_kind, + "version": _loaded_collective_version(), + }, + "device": { + "arch": arch, + "compute_units": int(properties.multi_processor_count), + "memory_bytes": int(properties.total_memory), + "product": torch.cuda.get_device_name(device), + "warp_size": int(properties.warp_size), + }, + "driver_version": driver, + "framework": {"kind": "torch", "version": str(torch.__version__)}, + "machine": machine, + "python_version": platform.python_version(), + "vendor": vendor, + } + + +def _summarize_realized_placement( + records: list[tuple[str, int]], + *, + expected_nodes: int, + expected_gpus_per_node: int, + expected_world_size: int, +) -> dict: + """Validate private host/rank records and return only publication-safe aggregates.""" + if expected_nodes < 1 or expected_gpus_per_node < 1: + raise ValueError("requested placement dimensions must be positive") + if expected_nodes * expected_gpus_per_node != expected_world_size: + raise ValueError("requested nodes x GPUs per node differs from world size") + if len(records) != expected_world_size: + raise ValueError("realized rank count differs from world size") + + by_host: dict[str, list[int]] = {} + for host, local_rank in records: + if not isinstance(host, str) or not host or type(local_rank) is not int: + raise ValueError("realized placement record has invalid types") + by_host.setdefault(host, []).append(local_rank) + + counts = sorted(len(local_ranks) for local_ranks in by_host.values()) + complete_local_ranks = all( + sorted(local_ranks) == list(range(expected_gpus_per_node)) + for local_ranks in by_host.values() + ) + unique_pairs = len(set(records)) == len(records) + if len(by_host) != expected_nodes: + raise ValueError( + f"realized node count {len(by_host)} differs from requested {expected_nodes}" + ) + if counts != [expected_gpus_per_node] * expected_nodes: + raise ValueError("realized ranks per node differ from requested GPUs per node") + if not complete_local_ranks or not unique_pairs: + raise ValueError("realized local ranks are incomplete or duplicated") + return { + "gpus_per_node": expected_gpus_per_node, + "nodes": expected_nodes, + "ranks_per_node": expected_gpus_per_node, + "unique_local_ranks": True, + "valid": True, + } + + +def _common_runtime_fingerprint(records: list[dict]) -> dict: + """Return the shared sanitized fingerprint, rejecting heterogeneous ranks.""" + if not records: + raise ValueError("runtime fingerprint evidence is empty") + canonical = { + json.dumps(record, allow_nan=False, sort_keys=True, separators=(",", ":")) + for record in records + } + if len(canonical) != 1: + raise ValueError("runtime fingerprint differs across distributed ranks") + return records[0] + + +def _all_gather_json( + value, + *, + torch_module, + dist_module, + device, + world_size: int, +) -> list: + """Gather bounded JSON values using device tensors, not Python object collectives.""" + payload = json.dumps( + value, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + if not payload or len(payload) > 1_048_576: + raise ValueError("distributed consensus payload size is invalid") + size = torch_module.tensor([len(payload)], dtype=torch_module.int64, device=device) + sizes = [torch_module.zeros_like(size) for _ in range(world_size)] + dist_module.all_gather(sizes, size) + lengths = [int(item.item()) for item in sizes] + if any(length < 1 or length > 1_048_576 for length in lengths): + raise ValueError("distributed consensus size evidence is invalid") + width = max(lengths) + encoded = torch_module.zeros(width, dtype=torch_module.uint8, device=device) + encoded[: len(payload)] = torch_module.tensor( + list(payload), dtype=torch_module.uint8, device=device + ) + gathered = [torch_module.empty_like(encoded) for _ in range(world_size)] + dist_module.all_gather(gathered, encoded) + records = [] + for tensor, length in zip(gathered, lengths, strict=True): + raw = bytes(tensor[:length].cpu().tolist()) + records.append(json.loads(raw.decode("utf-8"))) + return records + + +def _allocation_stratum_sha256( + physical_hosts: list[str], + *, + audit_salt: str | None, + fabric_selectors: dict[str, str | None], + required: bool, +) -> str | None: + """Commit private allocation/fabric identity without exposing its inputs.""" + if audit_salt in (None, ""): + if required: + raise ValueError("canonical execution requires a private allocation audit salt") + return None + if not isinstance(audit_salt, str) or not re.fullmatch(r"[0-9a-f]{64}", audit_salt): + raise ValueError("allocation audit salt is invalid") + if set(fabric_selectors) != set(PRIVATE_FABRIC_ENV): + raise ValueError("private fabric selector set differs from the stratum contract") + for value in fabric_selectors.values(): + if value is not None and ( + not isinstance(value, str) + or not value + or len(value) > 512 + or any(ord(char) < 32 or ord(char) == 127 for char in value) + ): + raise ValueError("private fabric selector is invalid") + if not physical_hosts or any( + not isinstance(host, str) + or not host + or len(host) > 255 + or any(ord(char) < 32 or ord(char) == 127 for char in host) + for host in physical_hosts + ): + raise ValueError("physical allocation host evidence is invalid") + payload = json.dumps( + { + "contract": ALLOCATION_STRATUM_CONTRACT, + "fabric_selectors": fabric_selectors, + "physical_hosts": sorted(set(physical_hosts)), + }, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hmac.new(bytes.fromhex(audit_salt), payload, hashlib.sha256).hexdigest() + + +def _common_allocation_stratum( + records: list[str | None], *, required: bool +) -> str | None: + """Require every distributed rank to derive the same private stratum.""" + if not records or any( + value is not None + and (not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value)) + for value in records + ): + raise ValueError("allocation stratum evidence is invalid") + distinct = set(records) + if len(distinct) != 1: + raise ValueError("allocation stratum differs across distributed ranks") + value = records[0] + if required and value is None: + raise ValueError("canonical execution requires an allocation stratum") + return value + + +def main() -> int: + ap = argparse.ArgumentParser(description="CollectiveX EP dispatch/combine sweep") + ap.add_argument( + "--backend", + required=True, + choices=[ + "deepep", + "deepep-v2", + "deepep-hybrid", + "mori", + "uccl", + "nccl-ep", + ], + ) + ep_harness.add_common_args(ap) + args = ap.parse_args() + + if args.mode == ep_harness.LOW_LATENCY_MODE: + if args.backend not in {"deepep", "uccl"}: + print( + "ERROR: low-latency mode is supported only by deepep and uccl", + file=sys.stderr, + ) + return 2 + if args.phase != "decode": + print("ERROR: low-latency mode requires --phase decode", file=sys.stderr) + return 2 + if args.case_id and not identity.is_case_id(args.case_id): + print(f"ERROR: invalid native case ID {args.case_id!r}", file=sys.stderr) + return 2 + if args.case_id and args.seed != ep_harness.ROUTING_SEED: + print( + f"ERROR: scheduled v1 cases require seed={ep_harness.ROUTING_SEED}; got {args.seed}", + file=sys.stderr, + ) + return 2 + if args.qualification_index not in range(1, ep_harness.QUALIFICATION_RUNS + 1): + print( + f"ERROR: qualification index must be in 1..{ep_harness.QUALIFICATION_RUNS}", + file=sys.stderr, + ) + return 2 + + sampling_error = ep_harness.sampling_contract_error( + args.iters, args.trials, args.warmup + ) + if sampling_error: + print(f"ERROR: {sampling_error}", file=sys.stderr) + return 2 + + try: + import torch + import torch.distributed as dist + except Exception as exc: # pragma: no cover + print(f"ERROR: torch unavailable: {exc!r}", file=sys.stderr) + return 3 + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "12355") + + import capability + + sku = capability.PLATFORMS.get(args.runner) + if sku is None: + print(f"ERROR: unknown runner identity {args.runner!r}", file=sys.stderr) + return 5 + machine = {"x86_64": "amd64", "aarch64": "arm64"}.get( + platform.machine(), platform.machine() + ) + props = torch.cuda.get_device_properties(device) + if torch.version.hip: + vendor = "amd" + accelerator = str(getattr(props, "gcnArchName", "")).split(":", 1)[0] + else: + vendor = "nvidia" + major, minor = torch.cuda.get_device_capability(device) + accelerator = f"sm{major}{minor}" + device_name = torch.cuda.get_device_name(device) + device_count = torch.cuda.device_count() + identity_issues = capability.runtime_identity_issues( + args.runner, + vendor=vendor, + arch=accelerator, + machine=machine, + device_name=device_name, + device_count=device_count, + world_size=world_size, + ) + if identity_issues: + print( + f"ERROR: runtime identity does not match {args.runner}: " + + "; ".join(identity_issues), + file=sys.stderr, + ) + return 5 + observed_gpus_per_node = args.gpus_per_node or device_count + if observed_gpus_per_node != sku["gpus_per_node"]: + print( + f"ERROR: {args.runner} requires {sku['gpus_per_node']} GPUs per node", + file=sys.stderr, + ) + return 5 + if world_size % observed_gpus_per_node: + print("ERROR: distributed world is not divisible by GPUs per node", file=sys.stderr) + return 5 + observed_nodes = world_size // observed_gpus_per_node + topology = capability.topology_for(args.runner, world_size) + observed_topology = { + "nodes": observed_nodes, + "gpus_per_node": observed_gpus_per_node, + "scale_up_domain": args.scale_up_domain or observed_gpus_per_node, + "scope": args.scope, + "scale_up_transport": args.scale_up_transport, + "scale_out_transport": args.scale_out_transport or None, + "transport": args.transport, + "topology_class": args.topology_class, + } + if topology is None or any( + observed_topology[field] != topology[field] for field in observed_topology + ): + print( + f"ERROR: runtime topology does not match {args.runner} EP{world_size}", + file=sys.stderr, + ) + return 5 + schedulable, reason = capability.resolve( + args.runner, + args.backend, + ep=world_size, + nodes=observed_nodes, + routing=args.routing, + eplb=False, + mode=args.mode, + ) + if not schedulable: + print(f"ERROR: scheduled case is unsupported: {reason}", file=sys.stderr) + return 5 + args.runtime_device_product = device_name + args.runtime_device_count = device_count + args.allocation_execution_id = os.environ.get("COLLECTIVEX_EXECUTION_ID") + + # Reproduction provenance (recorded in the artifact). Rack launchers provide ranks directly + # through srun, while single-node launchers use torchrun; do not claim torchrun for both. + if os.environ.get("TORCHELASTIC_RUN_ID"): + args.distributed_launcher = "torchrun" + prefix = f"torchrun --nproc_per_node={world_size}" + else: + args.distributed_launcher = "rank-environment" + prefix = f"RANK={rank} WORLD_SIZE={world_size} LOCAL_RANK={local_rank} python3" + args.reproduction_command = f"{prefix} bench/run_ep.py {shlex.join(sys.argv[1:])}" + args.image = os.environ.get("COLLECTIVEX_IMAGE", "") + args.image_digest = os.environ.get("COLLECTIVEX_IMAGE_DIGEST", "") + args.image_digest_verified = ( + os.environ.get("COLLECTIVEX_IMAGE_DIGEST_VERIFIED") == "1" + ) + # Container architecture and local squash hash for Enroot/Pyxis. + args.image_arch = machine + args.squash_sha256 = os.environ.get("COLLECTIVEX_SQUASH_SHA256") + # GitHub provenance: repo, run ID, attempt, ref, source SHA, job, + # artifact. A result is only publication-'official' when these are present (validity gate). + _run = { + "run_id": os.environ.get("GITHUB_RUN_ID"), + "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), + "ref": os.environ.get("GITHUB_REF_NAME") or os.environ.get("GITHUB_REF"), + "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") + or os.environ.get("GITHUB_SHA"), + "repo": os.environ.get("GITHUB_REPOSITORY"), + "job": os.environ.get("GITHUB_JOB"), + "artifact": os.environ.get("COLLECTIVEX_ARTIFACT_NAME"), + } + if any(_run.values()): + args.git_run = _run + else: + args.git_run = None + + # Import the backend class only after torch initializes. The selected mode is an + # explicit case dimension; adapters do not infer it from the token ladder. + if args.backend == "mori": + from ep_mori import MoRIBackend as Backend + elif args.backend == "nccl-ep": + from ep_nccl import NCCLBackend as Backend + elif args.backend == "uccl": + from ep_uccl import UCCLBackend as Backend + elif args.backend == "deepep-hybrid": + from ep_deepep_hybrid import DeepEPHybridBackend as Backend + elif args.backend == "deepep-v2": + from ep_deepep_v2 import DeepEPV2Backend as Backend + else: + from ep_deepep import DeepEPBackend as Backend + + # MoRI registers the default GPU process group with its SHMEM runtime. Keep that + # group device-only so scale-out does not also depend on a host Gloo fabric. + if not dist.is_initialized(): + if args.backend == "mori": + dist.init_process_group( + backend="nccl", + rank=rank, + world_size=world_size, + device_id=device, + ) + elif args.backend == "deepep-v2": + # PR #605 reuses PyTorch's NCCL communicator through ``_comm_ptr``. Supplying + # device_id eagerly forms it before ElasticBuffer construction. + dist.init_process_group("nccl", device_id=device) + else: + dist.init_process_group("nccl") + + args.runtime_fingerprint = _runtime_fingerprint( + torch, device, machine=machine, vendor=vendor, arch=accelerator + ) + + gpus_per_node = args.gpus_per_node or sku["gpus_per_node"] + try: + expected_nodes = int( + os.environ.get("SLURM_NNODES", str(world_size // gpus_per_node)) + ) + except ValueError as exc: + raise ValueError("SLURM_NNODES must be a positive integer") from exc + realized_records = _all_gather_json( + [socket.gethostname(), local_rank, args.runtime_fingerprint], + torch_module=torch, + dist_module=dist, + device=device, + world_size=world_size, + ) + args.realized_placement = _summarize_realized_placement( + [(record[0], record[1]) for record in realized_records], + expected_nodes=expected_nodes, + expected_gpus_per_node=gpus_per_node, + expected_world_size=world_size, + ) + args.runtime_fingerprint = _common_runtime_fingerprint( + [record[2] for record in realized_records] + ) + canonical = bool(args.workload_dir) + local_stratum = _allocation_stratum_sha256( + [record[0] for record in realized_records], + audit_salt=os.environ.get("CX_AUDIT_SALT"), + fabric_selectors={ + field: os.environ.get(environment) or None + for field, environment in PRIVATE_FABRIC_ENV.items() + }, + required=canonical, + ) + stratum_records = _all_gather_json( + local_stratum, + torch_module=torch, + dist_module=dist, + device=device, + world_size=world_size, + ) + args.allocation_stratum_sha256 = _common_allocation_stratum( + stratum_records, required=canonical + ) + + # Construct + run inside a try so a backend exception (esp. a new adapter on GPU) prints its + # FULL traceback to STDOUT — torchrun captures per-rank stdout but only summarizes stderr, so an + # uncaught exception is otherwise invisible in CI. Print on every rank (prefixed) then re-raise. + try: + backend = Backend(args, rank, world_size, local_rank, device) + if rank == 0: + print( + f"[run_ep] backend={args.backend} phase={args.phase} mode={args.mode} " + f"world={world_size} ep_size={world_size} hidden={args.hidden} " + f"topk={args.topk} experts={args.experts} dtype=bf16 " + f"routing={args.routing} seed={args.seed} " + f"qualification_index={args.qualification_index}" + ) + rc = ep_harness.run_sweep(args, backend, torch, dist, device, rank, world_size) + except Exception: + import traceback + + print( + f"[run_ep][rank{rank}] backend={args.backend} FAILED:\n" + + traceback.format_exc(), + flush=True, + ) + raise + # finalize() handles backend-specific teardown: DeepEP returns rc cleanly; + # MoRI hard-exits past its post-shmem_finalize teardown assertion. + return backend.finalize(rc) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/bench/workload.py b/experimental/CollectiveX/bench/workload.py new file mode 100644 index 0000000000..bb9b3f18c3 --- /dev/null +++ b/experimental/CollectiveX/bench/workload.py @@ -0,0 +1,403 @@ +#!/usr/bin/env python3 +"""Canonical, byte-stable CollectiveX routing workloads. + +A *canonical workload* is a routing trace generated ONCE, serialized to a platform-independent +file, and referenced by an immutable `workload_id`. Every promoted benchmark point consumes the +SAME serialized bytes, so "did NVIDIA and AMD run the identical workload?" is answered by a +checksum match, not by trusting that two machines re-ran the same seeded generator. + +Layout on disk (one workload = two files, basename = workload_id): + /.npz topk_idx [gt,topk] int32, topk_weights [gt,topk] float32 + /.manifest.json dims, routing profile, generator version, seed, SHA-256s + +Routing and gate weights come from a stdlib integer counter, not a framework RNG. The same +parameters therefore produce the same int32/float32 bytes across PyTorch and accelerator images. +""" +from __future__ import annotations + +from array import array +import hashlib +import json +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import identity # noqa: E402 + +WORKLOAD_SCHEMA_VERSION = 1 +# Bump when the counter or byte encoding changes. The workload ID binds parameters and trace bytes. +GENERATOR_VERSION = "collectivex-routing-counter-v3" +GATE_WEIGHT_FORMAT = "counter-u16-normalized-f32" +ACTIVATION_GENERATOR = "collectivex-activation-counter-v4" +EPLB_CALIBRATION_WINDOW = "collectivex-eplb-calibration-window-v1" +EPLB_CALIBRATION_TOKEN_OFFSET = 1 << 32 +_MASK64 = (1 << 64) - 1 + + +def _sha256(b: bytes) -> str: + return hashlib.sha256(b).hexdigest() + + +def _mix64(value: int) -> int: + value = (value + 0x9E3779B97F4A7C15) & _MASK64 + value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & _MASK64 + value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & _MASK64 + return value ^ (value >> 31) + + +def _counter(seed: int, token: int, slot: int, attempt: int, stream: int) -> int: + value = ( + (seed & _MASK64) + ^ (((token + 1) * 0xD2B74407B1CE6E93) & _MASK64) + ^ (((slot + 1) * 0xCA5A826395121157) & _MASK64) + ^ (((attempt + 1) * 0x9E3779B185EBCA87) & _MASK64) + ^ (((stream + 1) * 0xA24BAED4963EE407) & _MASK64) + ) + return _mix64(value) + + +def canonical_routing_rows( + global_tokens: int, + experts: int, + topk: int, + routing: str, + seed: int, + *, + token_offset: int = 0, +) -> tuple[list[list[int]], list[list[float]]]: + """Generate a deterministic routing window from exact integer counters.""" + if routing != "uniform": + raise ValueError(f"unknown routing {routing!r} (uniform)") + if global_tokens <= 0 or experts <= 0 or topk <= 0 or topk > experts: + raise ValueError("global_tokens/experts/topk must be positive and topk <= experts") + if type(token_offset) is not int or token_offset < 0: + raise ValueError("token_offset must be a non-negative integer") + + indices: list[list[int]] = [] + weights: list[list[float]] = [] + for local_token in range(global_tokens): + token = token_offset + local_token + selected: list[int] = [] + used: set[int] = set() + for slot in range(topk): + attempt = 0 + while True: + value = _counter(seed, token, slot, attempt, 0) + expert = value % experts + if expert not in used: + used.add(expert) + selected.append(expert) + break + attempt += 1 + if attempt > experts * 16: + raise RuntimeError("counter routing could not select distinct experts") + raw = [1 + _counter(seed, token, slot, 0, 1) % 65535 for slot in range(topk)] + denominator = float(sum(raw)) + indices.append(selected) + weights.append([value / denominator for value in raw]) + return indices, weights + + +def _canonical_bytes( + indices: list[list[int]], weights: list[list[float]] +) -> tuple[bytes, bytes]: + idx = array("i", (value for row in indices for value in row)) + gate = array("f", (value for row in weights for value in row)) + if idx.itemsize != 4 or gate.itemsize != 4: + raise RuntimeError("canonical workload requires 32-bit int and float arrays") + if sys.byteorder != "little": + idx.byteswap() + gate.byteswap() + return idx.tobytes(), gate.tobytes() + + +def trace_checksums( + indices: list[list[int]], weights: list[list[float]] +) -> dict[str, str]: + """Return the manifest hashes for exact logical or remapped routing rows.""" + idx_bytes, weight_bytes = _canonical_bytes(indices, weights) + return { + "topk_idx": _sha256(idx_bytes), + "topk_weights": _sha256(weight_bytes), + "trace": _sha256(idx_bytes + weight_bytes), + } + + +def canonical_member( + routing: str, + hidden: int, + topk: int, + experts: int, + ep_size: int, + tokens_per_rank: int, + seed: int, + *, + token_offset: int = 0, +) -> tuple[str, dict[str, str], list[list[int]], list[list[float]]]: + """Derive one canonical manifest member and retain its rows for proof checks.""" + global_tokens = ep_size * tokens_per_rank + indices, weights = canonical_routing_rows( + global_tokens, + experts, + topk, + routing, + seed, + token_offset=token_offset, + ) + checksums = trace_checksums(indices, weights) + member = compute_workload_id( + routing, + hidden, + topk, + experts, + ep_size, + global_tokens, + seed, + trace_checksum=checksums["trace"], + token_offset=token_offset, + ) + return member, checksums, indices, weights + + +def canonical_eplb_calibration_member( + routing: str, + hidden: int, + topk: int, + experts: int, + ep_size: int, + tokens_per_rank: int, + seed: int, +) -> tuple[str, dict[str, str], list[list[int]], list[list[float]]]: + """Return the EPLB calibration trace from a disjoint global-token window.""" + return canonical_member( + routing, + hidden, + topk, + experts, + ep_size, + tokens_per_rank, + seed, + token_offset=EPLB_CALIBRATION_TOKEN_OFFSET, + ) + + +def compute_workload_id(routing: str, hidden: int, topk: int, experts: int, + ep_size: int, global_tokens: int, seed: int, + generator: str = GENERATOR_VERSION, + trace_checksum: str | None = None, + token_offset: int = 0) -> str: + """Deterministic ID over parameters and canonical trace bytes.""" + if generator != GENERATOR_VERSION: + raise ValueError(f"unsupported workload generator {generator!r}") + if type(token_offset) is not int or token_offset < 0: + raise ValueError("token_offset must be a non-negative integer") + if trace_checksum is None: + indices, weights = canonical_routing_rows( + global_tokens, + experts, + topk, + routing, + seed, + token_offset=token_offset, + ) + idx_bytes, weight_bytes = _canonical_bytes(indices, weights) + trace_checksum = _sha256(idx_bytes + weight_bytes) + key = { + "generator": generator, "routing": routing, "hidden": hidden, "topk": topk, + "experts": experts, "ep_size": ep_size, "global_tokens": global_tokens, + "seed": seed, "trace_sha256": trace_checksum, + "activation_generator": ACTIVATION_GENERATOR, + "activation_identity": compute_activation_identity(seed, hidden), + } + if token_offset: + key.update({ + "routing_window": EPLB_CALIBRATION_WINDOW, + "token_offset": token_offset, + }) + return identity.workload_id(key) + + +def compute_activation_identity(seed, hidden, generator=ACTIVATION_GENERATOR) -> str: + """Identity of the exact counter-derived activation generator.""" + key = f"counter|seed={seed}|hidden={hidden}|gen={generator}" + return _sha256(key.encode()) + + +def build_manifest(routing, hidden, topk, experts, global_tokens, seed, experts_per_rank, + idx_np, weights_np): + """Assemble the manifest dict from the (numpy) trace arrays. Pure numpy/stdlib.""" + if experts % experts_per_rank: + raise ValueError("experts must be divisible by experts_per_rank") + idx_bytes = idx_np.astype(" str: + import numpy as np + os.makedirs(out_dir, exist_ok=True) + wid = manifest["workload_id"] + np.savez_compressed(os.path.join(out_dir, f"{wid}.npz"), + topk_idx=idx_np.astype(np.int32), topk_weights=weights_np.astype(np.float32)) + with open(os.path.join(out_dir, f"{wid}.manifest.json"), "w") as fh: + json.dump(manifest, fh, indent=2, sort_keys=True) + return wid + + +def load_workload(npz_path, verify=True): + """Load a canonical trace (numpy + stdlib only). Returns (idx_np, weights_np, manifest). + Raises ValueError if verify=True and the on-disk bytes don't match the manifest checksums.""" + import numpy as np + base = npz_path[:-4] if npz_path.endswith(".npz") else npz_path + with open(base + ".manifest.json") as fh: + manifest = json.load(fh) + if manifest.get("workload_id") != os.path.basename(base): + raise ValueError(f"workload manifest ID does not match filename for {base}") + with np.load(base + ".npz", allow_pickle=False) as archive: + if set(archive.files) != {"topk_idx", "topk_weights"}: + raise ValueError(f"workload archive fields differ for {base}") + idx_np = np.ascontiguousarray(archive["topk_idx"]) + w_np = np.ascontiguousarray(archive["topk_weights"]) + if verify: + ok, reason = verify_workload(manifest, idx_np, w_np) + if not ok: + raise ValueError(f"workload checksum mismatch for {base}: {reason}") + return idx_np, w_np, manifest + + +def verify_workload(manifest, idx_np, weights_np): + """Recompute checksums and compare to the manifest. Returns (ok, reason).""" + import numpy as np + expected_fields = { + "schema_version", "workload_id", "generator_version", "gate_weight_format", "dims", + "routing_profile", "seed", "checksums", "activation_profile", "activation_generator", + "activation_identity", + } + if not isinstance(manifest, dict) or set(manifest) != expected_fields: + return False, "manifest fields differ from the v1 contract" + if (manifest["schema_version"] != WORKLOAD_SCHEMA_VERSION + or manifest["generator_version"] != GENERATOR_VERSION + or manifest["gate_weight_format"] != GATE_WEIGHT_FORMAT + or manifest["routing_profile"] != "uniform"): + return False, "manifest version or generator is unsupported" + if (isinstance(manifest["seed"], bool) or not isinstance(manifest["seed"], int) + or not identity.is_workload_id(manifest["workload_id"])): + return False, "manifest seed or workload ID is invalid" + dims = manifest["dims"] + dim_fields = {"hidden", "topk", "experts", "ep_size", "tokens_per_rank", + "global_tokens", "experts_per_rank"} + if not isinstance(dims, dict) or set(dims) != dim_fields: + return False, "manifest dimensions are invalid" + if any(isinstance(dims[key], bool) or not isinstance(dims[key], int) or dims[key] <= 0 + for key in dim_fields): + return False, "manifest dimensions must be positive integers" + if (dims["experts"] != dims["ep_size"] * dims["experts_per_rank"] + or dims["global_tokens"] != dims["ep_size"] * dims["tokens_per_rank"]): + return False, "manifest EP dimensions are inconsistent" + shape = (dims["global_tokens"], dims["topk"]) + if (idx_np.dtype != np.int32 or weights_np.dtype != np.float32 + or idx_np.shape != shape or weights_np.shape != shape + or not idx_np.flags.c_contiguous or not weights_np.flags.c_contiguous): + return False, "workload array dtype, shape, or layout is invalid" + if (np.any(idx_np < 0) or np.any(idx_np >= dims["experts"]) + or np.any(np.diff(np.sort(idx_np, axis=1), axis=1) == 0)): + return False, "expert indices are out of range or repeated" + if (not np.isfinite(weights_np).all() or np.any(weights_np < 0) + or not np.allclose(weights_np.sum(axis=1), 1.0, rtol=1e-5, atol=1e-6)): + return False, "gate weights are invalid" + if (manifest["activation_profile"] != "canonical-counter-source-v3" + or manifest["activation_generator"] != ACTIVATION_GENERATOR + or manifest["activation_identity"] + != compute_activation_identity( + manifest["seed"], dims["hidden"], manifest["activation_generator"] + )): + return False, "activation identity is invalid" + ib = idx_np.astype(" must fail + idx2[0, 0] = (int(idx2[0, 0]) + 1) % 256 + bad, _ = verify_workload(man2, idx2, w2) + assert not bad, "verify must catch tampering" + print(f"save/load/verify roundtrip OK (workload_id={wid})") + except ImportError: + print("(numpy unavailable — skipped serialization roundtrip; id logic passed)") + print("workload self-test: PASS") + sys.exit(0) diff --git a/experimental/CollectiveX/capability.py b/experimental/CollectiveX/capability.py new file mode 100644 index 0000000000..22e2f3f0c6 --- /dev/null +++ b/experimental/CollectiveX/capability.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +"""Public runner and backend capability registry for CollectiveX.""" + +from __future__ import annotations + +import re +from typing import Any + + +DEEPEP_V2_COMMIT = "fa8a9b16898204afd347c663b89e65ef87dc6ce6" +DEEPEP_V2_SKU_CAPABILITIES = { + "h100-dgxc": { + "schedulable": False, + "basis": "current-runner-nccl-device-api-symmetric-memory-unavailable", + }, + "h200-dgxc": {"schedulable": True, "basis": "upstream-sm90-requirement"}, + "b200-dgxc": {"schedulable": True, "basis": "upstream-sm100-result"}, + "gb200": {"schedulable": True, "basis": "upstream-sm100-result"}, + "b300": {"schedulable": True, "basis": "pinned-pr605-pr630-sm103-maps-sm100f"}, + "gb300": {"schedulable": True, "basis": "pinned-pr605-pr630-sm103-maps-sm100f"}, + "mi325x": {"schedulable": False, "basis": "nvidia-only"}, + "mi300x": {"schedulable": False, "basis": "nvidia-only"}, + "mi355x": {"schedulable": False, "basis": "nvidia-only"}, +} + + +def _topologies( + product: str, *, gpus_per_node: int, scale_up_domain: int, scale_up_transport: str +) -> dict[int, dict[str, Any]]: + scale_up_class = ( + f"{product}-nvl72-mnnvl" + if scale_up_transport == "mnnvl" + else f"{product}-xgmi" + if scale_up_transport == "xgmi" + else f"{product}-{scale_up_transport}-island" + ) + return { + 8: { + "nodes": 8 // gpus_per_node, + "gpus_per_node": gpus_per_node, + "scale_up_domain": scale_up_domain, + "scope": "scale-up", + "scale_up_transport": scale_up_transport, + "scale_out_transport": None, + "transport": scale_up_transport, + "topology_class": scale_up_class, + }, + 16: { + "nodes": 16 // gpus_per_node, + "gpus_per_node": gpus_per_node, + "scale_up_domain": scale_up_domain, + "scope": "scale-up" if scale_up_domain >= 16 else "scale-out", + "scale_up_transport": scale_up_transport, + "scale_out_transport": None if scale_up_domain >= 16 else "rdma", + "transport": ( + scale_up_transport + if scale_up_domain >= 16 + else f"{scale_up_transport}-rdma" + ), + "topology_class": ( + scale_up_class + if scale_up_domain >= 16 + else f"{product}-{scale_up_transport}-rdma" + ), + }, + } + + +def _platform( + *, vendor: str, arch: str, machine: str, product: str, gpus_per_node: int, + scale_up_domain: int, scale_up_transport: str, launcher: str, +) -> dict[str, Any]: + topologies = _topologies( + product, + gpus_per_node=gpus_per_node, + scale_up_domain=scale_up_domain, + scale_up_transport=scale_up_transport, + ) + ep8 = topologies[8] + return { + "vendor": vendor, + "arch": arch, + "machine": machine, + "product": product, + # EP8 defaults remain while downstream readers migrate to per-EP records. + "transport": ep8["transport"], + "topology_class": ep8["topology_class"], + "gpus_per_node": gpus_per_node, + "scale_up_domain": scale_up_domain, + "ep_degrees": tuple(topologies), + "topologies": topologies, + "launcher": launcher, + } + + +PLATFORMS = { + "h100-dgxc": _platform( + vendor="nvidia", arch="sm90", machine="amd64", product="h100", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="nvlink", + launcher="single-slurm", + ), + "h200-dgxc": _platform( + vendor="nvidia", arch="sm90", machine="amd64", product="h200", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="nvlink", + launcher="single-slurm", + ), + "b200-dgxc": _platform( + vendor="nvidia", arch="sm100", machine="amd64", product="b200", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="nvlink", + launcher="single-slurm", + ), + "b300": _platform( + vendor="nvidia", arch="sm103", machine="amd64", product="b300", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="nvlink", + launcher="single-slurm", + ), + "gb200": _platform( + vendor="nvidia", arch="sm100", machine="arm64", product="gb200", + gpus_per_node=4, scale_up_domain=72, scale_up_transport="mnnvl", + launcher="gb-nv", + ), + "gb300": _platform( + vendor="nvidia", arch="sm103", machine="arm64", product="gb300", + gpus_per_node=4, scale_up_domain=72, scale_up_transport="mnnvl", + launcher="gb-nv", + ), + "mi325x": _platform( + vendor="amd", arch="gfx942", machine="amd64", product="mi325x", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="xgmi", + launcher="mi-amds", + ), + "mi300x": _platform( + vendor="amd", arch="gfx942", machine="amd64", product="mi300x", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="xgmi", + launcher="mi-amds", + ), + "mi355x": _platform( + vendor="amd", arch="gfx950", machine="amd64", product="mi355x", + gpus_per_node=8, scale_up_domain=8, scale_up_transport="xgmi", + launcher="mi-amds", + ), +} + +BACKENDS = { + "deepep": {"vendors": {"nvidia"}}, + "deepep-v2": { + "vendors": {"nvidia"}, + "implementation": "deep_ep.ElasticBuffer", + "source": "deepseek-ai/DeepEP#605+#630+#640", + "commit": DEEPEP_V2_COMMIT, + "communication_backend": "nccl-device-lsa", + "torch": "2.10.0+cu130", + "nccl": "2.30.4", + "sku_capabilities": DEEPEP_V2_SKU_CAPABILITIES, + }, + "uccl": { + "vendors": {"nvidia"}, + "machines": {"amd64"}, + "excluded_skus": {"b200-dgxc", "b300"}, + }, + "deepep-hybrid": {"vendors": {"nvidia"}}, + "mori": {"vendors": {"amd"}}, + "nccl-ep": {"vendors": {"nvidia", "amd"}}, +} +SWEEP_BACKENDS = tuple(BACKENDS) + +# Publication-quality topology exceptions apply after ordinary backend support +# checks. They describe the currently usable benchmark fabric, not a library +# implementation limit, and can be removed when the named topology is repaired. +TOPOLOGY_CELL_OVERRIDES: dict[tuple[str, int], str] = {} + +# Backend-specific topology limits require repeated native execution evidence. +# Keep these narrower than platform overrides so working reference paths remain +# measurable on the same fabric. +BACKEND_TOPOLOGY_CELL_OVERRIDES: dict[tuple[str, str, int], str] = { + ("b200-dgxc", "deepep", 16): ( + "DeepEP V1 EP16 requires unavailable GPU doorbell mapping or GDRCopy on B200" + ), + ("b200-dgxc", "deepep-hybrid", 16): ( + "DeepEP Hybrid EP16 requires unavailable GPU-to-NIC doorbell/UAR mappings on B200" + ), + ("b300", "deepep", 16): ( + "DeepEP V1 EP16 requires GDRCopy /dev/gdrdrv for NVSHMEM-IBGDA, " + "unprovisioned on B300 hosts" + ), + ("b300", "deepep-v2", 16): ( + "DeepEP V2 EP16 requires GDRCopy /dev/gdrdrv for NVSHMEM-IBGDA, " + "unprovisioned on B300 hosts" + ), + ("mi355x", "mori", 8): ( + "Pinned MoRI backend construction does not complete on MI355X EP8" + ), + ("mi355x", "mori", 16): ( + "Pinned MoRI backend construction does not complete on MI355X EP16" + ), + ("mi300x", "mori", 16): ( + "Pinned MoRI distributed initialization does not complete on MI300X EP16" + ), +} + + +def runtime_identity_issues( + sku: str, *, vendor: str, arch: str, machine: str, device_name: str, + device_count: int, world_size: int, +) -> list[str]: + """Validate public product identity on every rank without private device identifiers.""" + platform = PLATFORMS.get(sku) + if platform is None: + return [f"unknown runner identity {sku!r}"] + issues = [] + for field, observed in (("vendor", vendor), ("arch", arch), ("machine", machine)): + if observed != platform[field]: + issues.append(f"{field}={observed!r}, expected {platform[field]!r}") + products = set(re.findall(r"[a-z]+\d+[a-z]*", device_name.lower())) + if platform["product"] not in products: + issues.append(f"device product {device_name!r} does not identify {platform['product']}") + if device_count != platform["gpus_per_node"]: + issues.append( + f"visible GPUs={device_count}, expected {platform['gpus_per_node']} per node" + ) + if world_size not in platform["ep_degrees"]: + issues.append(f"EP{world_size} is not registered for {sku}") + return issues + + +def topology_for(sku: str, ep: int) -> dict[str, Any] | None: + """Return the exact public topology registered for one SKU/EP cell.""" + platform = PLATFORMS.get(sku) + if platform is None: + return None + return platform["topologies"].get(ep) + + +def _resolve_base( + sku: str, + backend: str, + *, + ep: int | None = None, + nodes: int | None = None, + routing: str = "uniform", + eplb: bool = False, + mode: str = "normal", +) -> tuple[bool, str]: + """Resolve the base BF16 capability for one SKU/backend/EP cell.""" + platform, implementation = PLATFORMS.get(sku), BACKENDS.get(backend) + if platform is None: + return False, f"unknown GHA runner label {sku!r}" + if implementation is None: + return False, f"unknown backend {backend!r}" + if mode not in {"normal", "low-latency"}: + return False, f"unknown benchmark mode {mode!r}" + if mode == "low-latency" and backend not in {"deepep", "uccl"}: + return False, f"{backend} has no distinct low-latency API" + if ep is None: + if nodes is None: + ep = platform["ep_degrees"][0] + else: + matches = [ + degree for degree, topology in platform["topologies"].items() + if topology["nodes"] == nodes + ] + if len(matches) != 1: + return False, f"{sku} does not register a unique {nodes}-node EP degree" + ep = matches[0] + topology = topology_for(sku, ep) + if topology is None or (nodes is not None and nodes != topology["nodes"]): + return False, f"{sku} does not register EP{ep} on {nodes} nodes" + if routing != "uniform" or eplb: + return False, "core routing is uniform; EPLB is unavailable" + if platform["vendor"] not in implementation["vendors"]: + return False, f"{backend} does not support {platform['vendor']}" + sku_capability = implementation.get("sku_capabilities", {}).get(sku) + if sku_capability is not None and not sku_capability["schedulable"]: + return False, f"{backend} is unsupported on {sku}: {sku_capability['basis']}" + if platform["machine"] not in implementation.get("machines", {platform["machine"]}): + return False, f"{backend} does not support {platform['machine']}" + if sku in implementation.get("excluded_skus", set()): + return False, f"{backend} is unavailable on {sku}" + backend_topology_override = BACKEND_TOPOLOGY_CELL_OVERRIDES.get( + (sku, backend, ep) + ) + if backend_topology_override is not None: + return False, backend_topology_override + topology_override = TOPOLOGY_CELL_OVERRIDES.get((sku, ep)) + if topology_override is not None: + return False, topology_override + return True, "ok" + + +def resolve_disposition( + sku: str, + backend: str, + *, + ep: int | None = None, + nodes: int | None = None, + routing: str = "uniform", + eplb: bool = False, + mode: str = "normal", +) -> tuple[str, str]: + """Resolve a BF16 cell to its capability disposition.""" + base_ok, base_detail = _resolve_base( + sku, + backend, + ep=ep, + nodes=nodes, + routing=routing, + eplb=eplb, + mode=mode, + ) + return ("supported", "ok") if base_ok else ("unsupported", base_detail) + + +def resolve( + sku: str, + backend: str, + *, + ep: int | None = None, + nodes: int | None = None, + routing: str = "uniform", + eplb: bool = False, + mode: str = "normal", +) -> tuple[bool, str]: + """Return whether one fixed-profile case can run on a public GHA runner label.""" + disposition, detail = resolve_disposition( + sku, + backend, + ep=ep, + nodes=nodes, + routing=routing, + eplb=eplb, + mode=mode, + ) + return disposition == "supported", detail diff --git a/experimental/CollectiveX/configs/suites.yaml b/experimental/CollectiveX/configs/suites.yaml new file mode 100644 index 0000000000..882c6779b0 --- /dev/null +++ b/experimental/CollectiveX/configs/suites.yaml @@ -0,0 +1,27 @@ +# CollectiveX EP comparison suites. +schema_version: 1 +# Iterable benchmark version. Bump when a methodology change makes new results +# incomparable with prior ones; it is copied verbatim into the matrix, every +# emitted case result, every extracted shard control, and the run summary so the +# frontend can group and pick runs by it. This is the only place it is set. +version: 1 + +suites: + ep-core: + mode: normal + workloads: [deepseek-v3] + platforms: [h100-dgxc, h200-dgxc, b300, b200-dgxc, gb300, gb200, mi300x, mi355x] + ep_degrees: [8, 16] + routings: [uniform] + phases: [decode, prefill] + token_points_prefill: [256, 512] + + ep-low-latency: + mode: low-latency + backends: [deepep, uccl] + workloads: [deepseek-v3] + platforms: [h100-dgxc, h200-dgxc, b300, b200-dgxc, gb300, gb200, mi300x, mi355x] + ep_degrees: [8, 16] + routings: [uniform] + phases: [decode] + token_points_decode: [1, 2, 4, 8, 16, 32, 64, 128] diff --git a/experimental/CollectiveX/configs/workloads.yaml b/experimental/CollectiveX/configs/workloads.yaml new file mode 100644 index 0000000000..ba750b9100 --- /dev/null +++ b/experimental/CollectiveX/configs/workloads.yaml @@ -0,0 +1,9 @@ +# CollectiveX EP canonical workload and phase metadata. +schema_version: 1 + +model_derived: + deepseek-v3: + hidden: 7168 + topk: 8 + routed_experts: 256 + verified_against: "deepseek-ai/DeepSeek-V3@e815299b0bcbac849fa540c768ef21845365c9eb/config.json" diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md new file mode 100644 index 0000000000..a4306fc0c8 --- /dev/null +++ b/experimental/CollectiveX/docs/methodology.md @@ -0,0 +1,243 @@ +# CollectiveX EP Benchmark Methodology + +CollectiveX schedules expert-parallel (EP) communication benchmarks, executes them on real +accelerator allocations, and uploads the neutral artifacts each run emits. It does **not** validate +those artifacts, promote, rank, recommend, select, hide, or decide what any consumer displays. The +frontend reads the neutral matrix, result, summary, and catalog artifacts and makes its own coverage +and display decisions. This document describes how a case is scheduled, measured, checked, and +recorded — not a publication or qualification contract. + +## Product Boundary + +CollectiveX is a communication microbenchmark for: + +- comparing EP libraries on one chip/topology; +- comparing EP latency and logical payload bandwidth across systems under the same workload; and +- surfacing unsupported, failed, invalid, and unstable cases rather than hiding them. + +It does not predict serving throughput without a separate correlation study. + +## Matrix + +The implemented workload is `deepseek-v3`: hidden 7168, top-k 8, 256 routed experts, packed +placement, and one pinned fixed resource profile per backend/topology. Dispatch and combine are +fixed BF16 on every backend; precision is not a swept dimension. Each case explicitly selects normal +`layout-and-dispatch-v1` or low-latency `expert-packed-weighted-combine-v1` semantics. + +- `ep-core`: uniform routing; decode T=1..128 powers of two; prefill T=256/512. +- `ep-low-latency`: DeepEP V1 / UCCL native low-latency APIs; uniform decode T=1..128 powers of + two; other backends are recorded as unsupported rather than fabricating a low-latency path. + +`sweep_matrix.py` materializes the requested SKUs, backends, EP sizes, and token ladders into a +matrix document, then extracts strict per-shard controls. `--only-sku`, `--exclude-skus`, and +`--ep-sizes` select a subset; a subset produces a smaller matrix, not a different contract. The +matrix is generated per dispatch; there is no frozen matrix digest or locked case count. + +| Systems | EP8 | EP16 | +|---|---|---| +| H100/H200/B200/B300 | 1x8 NVLink, scale-up | 2x8 NVLink + RDMA, scale-out | +| MI355X | 1x8 XGMI, scale-up | 2x8 XGMI + RDMA, scale-out | +| GB200/GB300 | 2x4 MNNVL, scale-up | 4x4 MNNVL, scale-up | + +Physical host count does not define scope. Both GB cells remain inside one 72-GPU MNNVL scale-up +domain. The MI325X launcher/configuration path is retained for future versions but is not referenced +by any current suite or shard. + +Unsupported combinations are explicitly classified in the matrix, not silently skipped coverage. DeepEP V2 is the +`ElasticBuffer` introduced by PR #605, pinned with upstream PR #630's minimal pure-scale-up fix and +the exact upstream PR #640 library matcher that excludes NCCL shared-memory mappings. Scale-up cases +request NCCL Device API LSA and fail closed unless the realized LSA team covers the full EP world. +x86 EP16 scale-out uses the hybrid path with GIN and requires two logical scale-out domains +represented by two physical RDMA ranks, with eight scale-up ranks per domain. GB EP16 remains MNNVL +scale-up and uses LSA. MoRI EP8 uses MI355X IntraNode in normal mode; EP16 uses pinned InterNodeV1 +over 2x8 XGMI + RDMA with 96 blocks, 64 RDMA blocks, 8 warps, one QP per PE, and external input. +MoRI's AsyncLL transport is not the low-latency suite contract and is never labeled as such. Whether +a given SKU/backend/EP cell is attempted is a capability fact; whether it succeeded is decided only +by the emitted artifact. + +## Workload Identity + +One canonical workload is generated over the global token batch and sliced by source rank. Expert +indices and gate weights are serialized. Activations use a versioned integer counter formula whose +BF16 values are exact across runtimes; its full identity is bound into the manifest. The manifest +also binds shape/EP coordinates and oracle version. SHA-256 covers canonical bytes and parameters; +library RNG regeneration is not proof of identity. + +Routing traffic distinguishes: + +- token-expert assignments, which determine expert compute load; and +- rank-deduplicated token payload copies, which determine EP activation traffic. + +Adapters may not generate routing or reinterpret one quantity as the other. + +## Measurement + +Normal mode uses `layout-and-dispatch-v1`: dispatch timing includes layout plus communication, and +combine returns activation payload through an unweighted rank-sum path. Low-latency mode uses +`expert-packed-weighted-combine-v1`: native DeepEP V1 / UCCL APIs dispatch token-expert assignments +and perform gate-weighted combine. Expert-output staging is outside isolated combine timing and +inside the measured paired roundtrip. Each component declares availability, origin, start/end states, +stage scope, and sample count. A paired-only API reports null isolated components. `isolated_sum` is +derived. Normal and low-latency evidence describe different measurement contracts and are not +directly comparable; the artifact records the mode so a reader can keep them separate. + +Every measured component uses `fixed-512-v1`: + +- 64 trials x 8 timed iterations = 512 observations; +- 32 synchronized full dispatch-stage-combine warmups before each available measured component at + every trial/point; +- roundtrip first, then isolated dispatch and combine, with a fixed per-phase conditioning ladder; and +- per-iteration maximum latency across ranks before nearest-rank p50/p90/p95/p99. + +Measured roundtrip p99 is the headline latency. Retries remain separate attempts; a later success +does not erase earlier failures. Decode and prefill identify the serving regime represented by one +MoE-layer collective; they do not change the timed primitive at an otherwise identical shape. + +The NCCL/RCCL reference is an end-to-end Python adapter, not a bare fabric primitive. Its dispatch +boundary includes layout, count exchange, a device-to-host split synchronization, fresh receive +allocation, and four payload/metadata all-to-all calls; activation-only combine adds one all-to-all +plus scatter/reduction. Its p99 therefore measures the complete reference-adapter boundary and can be +host/scheduler-sensitive. It is useful for portable system controls but must not be labeled fabric, +link, bus, or single-collective latency. + +The versioned conditioning and EPLB planner contracts (reference trace, redundant count, and +placement/remap version) are part of scheduled and evidence identity. + +Logical payload bandwidth is: + +`logical_payload_bytes / measured_latency_seconds` + +Normal-mode payload bytes use rank-deduplicated token-rank activations; low-latency bytes use +token-expert assignments. Both add required scale bytes at the named boundary and exclude expert +metadata, padding, and backend buffer capacity. Algorithm bandwidth, bus bandwidth, wire +utilization, and physical-link utilization are not emitted without a defined primitive model or +transport counters. Logical bandwidth must never be labeled physical bandwidth. Payload and token +rates are named `rate_at_latency_percentile`: bytes or tokens divided by the matching latency +percentile. They are lower-tail service rates at p99 latency, not p99 percentiles of an inverted +rate distribution. + +## Correctness + +An implementation-independent oracle uses an expert-specific deterministic transform so wrong expert +routing cannot pass an identity roundtrip. For every rank and point it verifies: + +1. destination rank/expert, source token, multiplicity, gate weight, and receive counts; +2. dispatched payload and metadata before timing; +3. combined output before timing; +4. unchanged semantic inputs through all timed samples; and +5. dispatched payload/metadata and combined output again after timing. + +Normal-mode adapters use activation-only, unweighted rank-sum combine. The oracle builds each rank's +gate-weighted expert aggregate before combine, independently derives `sum(gate * expert(token))`, and +checks the dispatch metadata and transformed output. Low-latency adapters separately verify the +expert-packed source/expert assignment, native gate weights, and gate-weighted combined output. Both +contracts compare against the reference activation. The combine gate is `rtol=0.05, atol=0.02` for +the BF16 communication path. This threshold is a correctness gate, not an estimate of transport +error. Any failed rank or point makes the case ineligible in the result it writes. +Pre/post dispatch evidence is hashed in canonical source-token order. Native receive slots may be +assigned nondeterministically, so physical receive order is not treated as a correctness property. + +## Result Artifact + +One raw case document uses `format: "collectivex.ep.v1"`, carries `schema_version: 1`, rejects +unknown fields, and contains: + +- `case`: stable case ID, suite, and coordinate; +- `workload`: canonical identity and logical MoE shape; +- `measurement`: sampling, component states, timing, and byte accounting; +- `implementation`: instantiated class/API, pinned source, loaded libraries, and resources; +- `topology`: requested and realized SKU, devices, placement, scale-up domain, and transport; +- `provenance`: source SHA, image/squash hashes, allocation, run, and attempt; +- `rows`: point latency, byte accounting, token rate, correctness, load, fanout, and anomaly + evidence; and +- `outcome`: `success`, `failed`, `invalid`, `unsupported`, with `diagnostic` and reasons. + +Exact per-point samples are emitted as detached `collectivex.samples.v1` documents referenced by path +and SHA-256, so the raw document stays compact. Each dispatched case writes its raw result document; +unsupported or never-run cells produce no synthetic record. Private +environment details (hosts, addresses, device selectors, credentials, workspace paths) remain in +local mode-0600 logs and ignored operator notes and never enter an emitted artifact. + +## Identity + +Identifiers are readable factor strings with a short digest suffix, not opaque hashes: + +- `case_id`: `{sku}-{backend}-{workload}-{mode}-{phase}-ep{ep}-{routing}[-eplb]` followed by a + 12-hex digest of the full case factors (SKU, profile, and case), so the body is human-readable + while the suffix makes the ID collision-proof and a tamper-evident join key; +- `attempt_id`: `case_id` plus `-a{ordinal:02d}`; and +- `point_id`: `case_id` plus `-t{tokens_per_rank}`. + +Content SHA-256 is retained only where it identifies separately-stored bytes or executable code: the +workload manifest (`workload_id` = `cxwork-v1-{sha256}`), detached sample files (`sample_sha256`), +source commits, container images and squash layers, and deferred/JIT-generated source. DeepEP V2 +uses a fixed NVCC random seed and binds final cache keys plus generated-source and executable-SASS +hashes; raw CUBIN bytes remain private diagnostics and are stripped before they enter an artifact. +Hybrid binds its realized auto-tuned config and complete kernel-key set while retaining rank-local +shared-object hashes as private diagnostics. Locally built extension hashes are diagnostic; their +pinned source trees, build recipe, runtime, and dependencies remain bound to the case factors. + +These IDs let a consumer group matched configurations and separate distinct ones. The backend does +not itself compute cohorts, controlled comparisons, sensitivity pairs, eligibility, or +recommendations — a reader decides which cases to surface and how to compare them. + +## Execution Isolation + +Every non-MNNVL scale-out case uses operator-pinned socket and RDMA selectors. The launcher rejects +missing or partial profiles, then probes every allocated node for the configured interface, active +HCA port, and configured GID before backend initialization. It never substitutes a default route, +inherited runner environment, or transport fallback. Scale-up and MNNVL cases clear the profile; +scale-out NCCL/RCCL forces `NCCL_NET=IB` and exact HCA matching. Selector values remain in encrypted +config and mode-0600 private logs. + +Repository staging uses a pre-existing, runner-owned, group/world non-writable shared base outside +the checkout and workflow workspace. The parent process resolves the exact execution child before +copying, claims it with a runner-owned marker, and verifies that all allocated nodes can read and +write the same bytes. Cleanup waits for confirmed allocation teardown and removes only that child, +including a safely identified partial claim. The same-run V2/Hybrid source archive is fully validated +under fixed member and expanded-size bounds, and only the selected pinned root is extracted; a +symlink is accepted only when it is a relative leaf pointing to a regular member inside the same +backend root, followed by exact Git tree/submodule validation. + +H200, B200, and B300 may derive that private base beneath the validated operating-system account home +when it is compute-visible. H100 instead derives a sibling of its shared container directory, never a +child of image storage. The launcher still proves cross-node visibility before any benchmark starts. +Canonical B300 execution ignores the legacy operator `stage_dir` field and always derives the base +from the validated shared account home. Its UID-mapped Actions shell may accept that exact base when +its owner matches the private parent owner; explicit stages and all other runners retain the strict +effective-UID ownership rule. A hashed execution-ID suffix isolates parallel B300 workers without +exposing private runner identity. The current NFS export may realize a newly created hashed base as +UID 0; only that creation path is accepted, while a pre-existing root-owned base is rejected. +Canonical GB300 execution likewise ignores its legacy group-writable `stage_dir` and derives an +execution-hashed private base beneath the validated compute-visible account home. + +## Image Pinning And Build Isolation + +Container tags are checked against pinned registry digests. Enroot imports use a fixed +`SOURCE_DATE_EPOCH` and versioned cache generation; every mounted squash is freshly hashed into series +identity. Image-provided DeepEP is also checked against exact per-architecture wheel and installed-file +fingerprints, so a stale cache cannot inherit the pinned source identity. Source-built DeepEP V2 uses +a separate mode-0700 cluster-local cache mounted only as `/cx-cache`. Its content key binds a +versioned build recipe, verified image digest, CPU/GPU architecture, upstream source trees, and pinned +build dependencies. The cache is never an artifact; per-execution source/results stages remain +isolated and disposable, and marker plus runtime probes fail closed before reuse. The runner UID is +inside the trusted cluster boundary: this cache guards against stale or accidental mutation, not +hostile same-UID jobs. Only an unpublished partial build may be reset automatically; a cache that +fails integrity or runtime checks is left intact and rejected so a concurrent allocation cannot lose +files it is using. + +## Neutral Artifact Delivery + +There is no results server, attached store, or managed object store. Each shard runs one allocation, +emits per-case result JSON plus detached sample JSON and a small mechanical summary, and uploads them +as GitHub artifacts with `always()` so a red or partial run still uploads. A case counts as successful +on the benchmark's own return code; there is no schema, completeness, detached-sample, or privacy +validation step before upload, and failed or unsupported cells produce no synthetic record. + +No step promotes a run, builds a dataset, or advances a channel; the artifacts are the output. Any +downstream display or comparison is the consumer's responsibility. + +## Legacy Data + +Historical numeric schemas 3-5 are outside this benchmark's artifacts. They remain historical +diagnostic evidence and are not produced or consumed by the current sweep. diff --git a/experimental/CollectiveX/identity.py b/experimental/CollectiveX/identity.py new file mode 100644 index 0000000000..2fa34bf8d4 --- /dev/null +++ b/experimental/CollectiveX/identity.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +"""Canonical cross-runtime identities for CollectiveX. + +A case identity is a readable, factor-derived join between a scheduled case and +its emitted result: a human-readable body (SKU, backend, workload, mode, phase, +EP size, routing) plus a short digest of the exact case factors. The suffix keeps +the identity collision-proof and tamper-evident without hiding what the case is. +Attempt and point identities are readable derivations of the case identity. + +Content SHA-256 survives only where it identifies separately-stored bytes: the +workload manifest whose exact routing/gate data a result consumes. Detached +sample files, source commits, container images, and generated kernels carry their +own SHA-256 fields in the artifacts; this module does not mint identifiers for +catalogs, series, evidence, allocations, or any other publisher-side grouping. +""" +from __future__ import annotations + +import hashlib +import json +import re +from typing import Any + +MAX_SAFE_INTEGER = (1 << 53) - 1 +WORKLOAD_PREFIX = "cxwork-v1-" +_CASE_SUFFIX_LEN = 12 +_NON_SLUG = re.compile(r"[^a-z0-9]+") +_CASE_ID = re.compile(r"^[a-z0-9][a-z0-9_.-]*-[0-9a-f]{%d}$" % _CASE_SUFFIX_LEN) +_WORKLOAD_ID = re.compile(r"^cxwork-v1-[0-9a-f]{64}$") + +V1_NORMAL_CASE_PROFILE = { + "activation_generator": "collectivex-activation-counter-v4", + "activation_profile": "canonical-counter-source-v4", + "combine_dtype": "bf16", + "combine_semantics": "activation-only", + "component_order_contract": "qualification-hash-rotated-components-v1", + "conditioning_contract": "fixed-phase-ramp-8-roundtrips-v1", + "contract": "layout-and-dispatch-v1", + "correctness_scope": "dispatch-metadata-and-transformed-combine", + "dtype": "bf16", + "eplb_planner": "greedy-rank-major-v1", + "eplb_redundant_experts": 32, + "eplb_reference_tokens_per_rank": 2048, + "mode": "normal", + "oracle_contract": "expert-specific-transform-v1", + "payload_unit": "token-rank", + "placement": "packed", + "percentile_method": "nearest-rank", + "rank_reduction": "cross-rank-max-per-iteration", + "resource_mode": "fixed-profile", + "routing_generator": "collectivex-routing-counter-v3", + "sampling_contract": "fixed-512-v1", + "seed": 67, + "source_identity_contract": "bounded-sign-bit-source-v1", +} + +V1_LOW_LATENCY_CASE_PROFILE = { + **V1_NORMAL_CASE_PROFILE, + "combine_semantics": "gate-weighted", + "contract": "expert-packed-weighted-combine-v1", + "correctness_scope": "expert-assignment-and-weighted-combine", + "mode": "low-latency", + "oracle_contract": "expert-assignment-transform-v1", + "payload_unit": "token-expert", +} + +# Compatibility alias for normal-mode callers. New scheduling and validation +# must select a profile from the explicit case mode. +V1_CASE_PROFILE = V1_NORMAL_CASE_PROFILE +V1_CASE_PROFILES = { + "normal": V1_NORMAL_CASE_PROFILE, + "low-latency": V1_LOW_LATENCY_CASE_PROFILE, +} + + +class IdentityError(ValueError): + """An identity payload cannot be represented consistently across runtimes.""" + + +def case_profile(mode: str) -> dict[str, Any]: + """Return the immutable measurement profile for one scheduled mode.""" + try: + return V1_CASE_PROFILES[mode] + except KeyError as exc: + raise IdentityError(f"unknown CollectiveX case mode {mode!r}") from exc + + +def profile_for_case(case: dict[str, Any]) -> dict[str, Any]: + """Resolve a scheduled case's explicit mode to its identity profile. + + Dispatch and combine are fixed BF16 benchmark facts, so a case's identity + profile is fully determined by its measurement mode. + """ + mode = case.get("mode") + if not isinstance(mode, str): + raise IdentityError("scheduled case mode is missing") + return case_profile(mode) + + +def _validate(value: Any, path: str = "$") -> None: + if value is None or isinstance(value, bool): + return + if isinstance(value, str): + if any(ord(character) < 0x20 or ord(character) > 0x7E for character in value): + raise IdentityError(f"{path}: string must contain printable ASCII only") + return + if type(value) is int: + if abs(value) > MAX_SAFE_INTEGER: + raise IdentityError(f"{path}: integer exceeds the cross-runtime safe range") + return + if isinstance(value, list): + for index, item in enumerate(value): + _validate(item, f"{path}[{index}]") + return + if isinstance(value, dict): + for key, item in value.items(): + if not isinstance(key, str): + raise IdentityError(f"{path}: object key is not a string") + if any(ord(character) < 0x20 or ord(character) > 0x7E for character in key): + raise IdentityError(f"{path}: object key must contain printable ASCII only") + _validate(item, f"{path}.{key}") + return + raise IdentityError(f"{path}: unsupported identity value {type(value).__name__}") + + +def canonical_bytes(value: Any) -> bytes: + """Return compact UTF-8 JSON after enforcing the portable value subset.""" + _validate(value) + return json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + +def _content_sha256(value: Any) -> str: + return hashlib.sha256(canonical_bytes(value)).hexdigest() + + +def _slug(value: Any) -> str: + text = _NON_SLUG.sub("-", str(value).strip().lower()).strip("-") + if not text: + raise IdentityError("case factor has no printable identity component") + return text + + +def _case_body(sku: str, case: dict[str, Any]) -> str: + parts = [ + _slug(sku), + _slug(case["backend"]), + _slug(case.get("workload") or "manual"), + _slug(case["mode"]), + _slug(case["phase"]), + "ep%d" % int(case["ep"]), + _slug(case["routing"]), + ] + if case.get("eplb"): + parts.append("eplb") + return "-".join(parts) + + +def case_id(*, sku: str, profile: dict[str, Any], case: dict[str, Any]) -> str: + """Readable case identity: factor body plus a digest of the exact factors.""" + factors = {"case": case, "profile": profile, "sku": sku} + return f"{_case_body(sku, case)}-{_content_sha256(factors)[:_CASE_SUFFIX_LEN]}" + + +def case_id_from_factors(factors: dict[str, Any]) -> str: + """Recompute a case identity from an emitted `case_factors` object.""" + return case_id(sku=factors["sku"], profile=factors["profile"], case=factors["case"]) + + +def attempt_id(*, case: str, ordinal: int) -> str: + if not isinstance(ordinal, int) or isinstance(ordinal, bool) or ordinal < 1: + raise IdentityError("attempt ordinal must be a positive integer") + return f"{case}-a{ordinal:02d}" + + +def point_id(*, case: str, tokens_per_rank: int) -> str: + if not isinstance(tokens_per_rank, int) or isinstance(tokens_per_rank, bool) or tokens_per_rank < 1: + raise IdentityError("tokens_per_rank must be a positive integer") + return f"{case}-t{tokens_per_rank}" + + +def is_case_id(value: Any) -> bool: + return bool(isinstance(value, str) and _CASE_ID.fullmatch(value)) + + +def workload_id(value: dict[str, Any]) -> str: + """Content identity of a workload manifest's canonical routing/gate data.""" + return WORKLOAD_PREFIX + _content_sha256( + {"kind": "workload", "value": value, "version": 1} + ) + + +def is_workload_id(value: Any) -> bool: + return bool(isinstance(value, str) and _WORKLOAD_ID.fullmatch(value)) diff --git a/experimental/CollectiveX/launchers/launch_gb-nv.sh b/experimental/CollectiveX/launchers/launch_gb-nv.sh new file mode 100644 index 0000000000..16371719ee --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_gb-nv.sh @@ -0,0 +1,111 @@ +#!/usr/bin/env bash +# CollectiveX shared GB200/GB300 NVL72 (aarch64) launcher. +# shellcheck disable=SC2016,SC2034 +# +# EP8/EP16 use one Slurm task per GPU across two or four trays in the same +# MNNVL scale-up domain. +set -euo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CX_DIR="$(cd "$HERE/.." && pwd)"; REPO_ROOT="$(cd "$CX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +PRODUCT="${CX_SHARD_SKU:-${CX_GB_PRODUCT:-${CX_PUBLIC_RUNNER:-}}}" +case "$PRODUCT" in + gb200|gb300) ;; + *) cx_die "set CX_SHARD_SKU or CX_PUBLIC_RUNNER to gb200 or gb300" ;; +esac +RUNNER="$PRODUCT" +export CX_RUNNER="$RUNNER" CX_BENCH="${CX_BENCH:-deepep}" +export CX_IMAGE_PLATFORM=linux/arm64 +JOB_ID="" +cx_install_launcher_fail_safe +cx_set_failure_stage setup +cx_load_operator_config +cx_lock_canonical_gha_env "$RUNNER" +NODES="${CX_NODES:-2}"; GPN="${CX_GPUS_PER_NODE:-4}" +SCALE_UP_DOMAIN="${CX_SCALE_UP_DOMAIN:-72}" +EXPECTED_WORLD=$((NODES * GPN)) +NGPUS="${CX_NGPUS:-$EXPECTED_WORLD}" +if [ "$PRODUCT" = gb200 ]; then default_time=30; else default_time=90; fi +TIME_MIN="${CX_TIME:-$default_time}" +[ "$NODES" = 2 ] || [ "$NODES" = 4 ] \ + || cx_die "$PRODUCT v1 supports two or four four-GPU trays" +[ "$GPN" = 4 ] || cx_die "$PRODUCT requires four GPUs per tray" +[ "$SCALE_UP_DOMAIN" = 72 ] || cx_die "$PRODUCT requires the NVL72 scale-up domain" +[ "$NGPUS" = "$EXPECTED_WORLD" ] \ + || cx_die "$PRODUCT world size must equal nodes x GPUs per tray" +cx_apply_timing_profile +IMAGE="${CX_IMAGE:-$(cx_default_image "$PRODUCT")}" +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +export CX_RUNNER="$RUNNER" CX_TS="$TS" CX_TOPO="${PRODUCT}-nvl72-mnnvl" +export CX_SCOPE=scale-up CX_TRANSPORT=mnnvl CX_SCALE_UP_TRANSPORT=mnnvl +export CX_NODES="$NODES" CX_GPUS_PER_NODE="$GPN" CX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +export CX_NGPUS="$NGPUS" +unset CX_SCALE_OUT_TRANSPORT +case "$CX_BENCH" in + deepep|deepep-v2|deepep-hybrid|nccl-ep) ;; + *) cx_die "unsupported $PRODUCT EP backend: $CX_BENCH" ;; +esac +cx_validate_shard_control "$CX_DIR" +cx_load_network_control_mode "$CX_DIR" || cx_die "cannot resolve network control mode" +cx_require_vars CX_PARTITION CX_ACCOUNT CX_SQUASH_DIR CX_STAGE_DIR +[ "$PRODUCT" != gb300 ] || cx_require_vars CX_ENROOT_CACHE_PATH +PARTITION="$CX_PARTITION"; ACCOUNT="$CX_ACCOUNT"; SQUASH_DIR="$CX_SQUASH_DIR" +[ -z "${CX_ENROOT_CACHE_PATH:-}" ] || export ENROOT_CACHE_PATH="$CX_ENROOT_CACHE_PATH" +export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 MC_FORCE_MNNVL=1 +cx_apply_network_profile "$NODES" "$CX_TRANSPORT" + +cx_log "$PRODUCT nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$CX_BENCH" +[ "${CX_DRYRUN:-0}" = 1 ] && { cx_log "DRYRUN"; exit 0; } +cx_set_failure_stage registry-verification +cx_verify_registry_image "$IMAGE" +cx_set_failure_stage repository-stage +MOUNT_SRC="$(cx_stage_path "$REPO_ROOT" "$CX_STAGE_DIR")" +cx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +cx_prepare_runtime_marker "$MOUNT_SRC" +CONTAINER_MOUNTS="$MOUNT_SRC:/ix" +if [ "$CX_BENCH" = deepep-v2 ] || [ "$CX_BENCH" = deepep-hybrid ]; then + cx_set_failure_stage backend-setup + cx_prepare_backend_source "$MOUNT_SRC" "$CX_BENCH" \ + || cx_die "cannot stage the pinned backend source" + export CX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.cx_sources +fi +if [ "$CX_BENCH" = deepep-v2 ]; then + cx_prepare_backend_cache "$CX_SQUASH_DIR" \ + || cx_die "cannot prepare the isolated backend cache" + CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$CX_PREPARED_BACKEND_CACHE:/cx-cache" + export CX_BACKEND_CACHE_ROOT=/cx-cache +fi + +cx_set_failure_stage scheduler-allocation +command -v salloc >/dev/null || cx_die "salloc not found" +cx_salloc_jobid --partition="$PARTITION" --account="$ACCOUNT" --nodes="$NODES" \ + --gres=gpu:"$GPN" --ntasks-per-node="$GPN" --exclusive --mem=0 --cpus-per-task=35 \ + --time="$TIME_MIN" +[ -n "$JOB_ID" ] || cx_die "no JOB_ID from salloc" +cx_set_failure_stage container-import +SQUASH_FILE="$(cx_ensure_squash_on_job "$JOB_ID" "$SQUASH_DIR" "$IMAGE")" +cx_set_failure_stage container-hash +cx_export_squash_identity "$SQUASH_FILE" +cx_preflight_allocation "$JOB_ID" "$NODES" "$MOUNT_SRC" "$SQUASH_FILE" \ + "${CX_SHARD_FILE:-}" + +# Keep the loader policy here because it is platform/container specific and +# security tests evaluate this literal independently. +SOURCE_BACKEND_ENV='case "${SLURM_NODEID:-}" in ""|*[!0-9]*) exit 66;; esac; env_file="/ix/experimental/CollectiveX/.cx_backend/env/node-${SLURM_NODEID}.sh"; env_root="${env_file%/*}"; [ -d "$env_root" ] && [ ! -L "$env_root" ] || exit 66; case "$(stat -c "%a" "$env_root")" in 700|[1-7]700) ;; *) exit 66;; esac; [ -f "$env_file" ] && [ -r "$env_file" ] && [ ! -L "$env_file" ] && [ "$(stat -c "%u:%a" "$env_file")" = "$(stat -c "%u" "$env_root"):600" ] || exit 66; . "$env_file" || exit 66' +BACKEND_PROBE="$SOURCE_BACKEND_ENV"'; case "$CX_BENCH" in deepep) python3 -c "from deep_ep import Buffer";; deepep-v2) python3 -c "import deep_ep; assert hasattr(deep_ep, '\''ElasticBuffer'\'')";; deepep-hybrid) python3 -c "import deep_ep; assert hasattr(deep_ep, '\''HybridEPBuffer'\'')";; nccl-ep) python3 -c "import torch";; esac' +WRAP="${SOURCE_BACKEND_ENV}"$'\n'"$(cx_slurm_rank_wrapper)" +CX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) +[ "$CX_BENCH" != deepep ] || export CX_ALLOW_MNNVL=1 +run_rc=0 +cx_set_failure_stage container-launch +cx_run_distributed_shard || run_rc=$? + +cx_adopt_runtime_stage "$MOUNT_SRC" +collect_rc=0 +cx_collect_results "$MOUNT_SRC" "$REPO_ROOT" || collect_rc=$? +[ "$run_rc" != 0 ] || [ "$collect_rc" = 0 ] || cx_set_failure_stage artifact-collection +final_rc="$run_rc" +[ "$final_rc" != 0 ] || final_rc="$collect_rc" +exit "$final_rc" diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh new file mode 100644 index 0000000000..ff3552aa34 --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -0,0 +1,183 @@ +#!/usr/bin/env bash +# CollectiveX shared AMD Slurm launcher (one or two nodes). +# shellcheck disable=SC2016,SC2034 +set -euo pipefail + +HERE="$(cd -P -- "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +CX_DIR="$(cd "$HERE/.." && pwd)" +REPO_ROOT="$(cd "$CX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +RUNNER="${CX_SHARD_SKU:-${CX_PUBLIC_RUNNER:-}}" +case "$RUNNER" in + mi300x|mi325x) CPUS_PER_TASK=256; DEVICE_MOUNTS=",/dev/kfd:/dev/kfd,/dev/dri:/dev/dri" ;; + mi355x) CPUS_PER_TASK=128; DEVICE_MOUNTS="" ;; + *) cx_die "set CX_SHARD_SKU or CX_PUBLIC_RUNNER to a registered AMD SKU" ;; +esac +export CX_RUNNER="$RUNNER" CX_BENCH="${CX_BENCH:-mori}" +export CX_IMAGE_PLATFORM=linux/amd64 +JOB_ID="" +cx_install_launcher_fail_safe +cx_set_failure_stage setup +cx_load_operator_config +cx_lock_canonical_gha_env "$RUNNER" + +NODES="${CX_NODES:-1}"; GPN="${CX_GPUS_PER_NODE:-8}" +SCALE_UP_DOMAIN="${CX_SCALE_UP_DOMAIN:-8}" +EXPECTED_WORLD=$((NODES * GPN)) +NGPUS="${CX_NGPUS:-$EXPECTED_WORLD}" +TIME_MIN="${CX_TIME:-60}" +EXCLUDE_NODES="${CX_EXCLUDE_NODES:-}" +NODELIST="${CX_NODELIST:-}" +MOUNT_DIR=/ix +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +[ "$NODES" = 1 ] || [ "$NODES" = 2 ] \ + || cx_die "$RUNNER supports one or two nodes" +[ "$GPN" = 8 ] || cx_die "$RUNNER requires eight GPUs per node" +[ "$SCALE_UP_DOMAIN" = 8 ] || cx_die "$RUNNER requires an eight-GPU scale-up domain" +[ "$NGPUS" = "$EXPECTED_WORLD" ] \ + || cx_die "$RUNNER world size must equal nodes x GPUs per node" +case "$CX_BENCH" in + mori|nccl-ep) ;; + *) cx_die "unsupported AMD EP backend: $CX_BENCH" ;; +esac + +if [ "$RUNNER" = mi300x ] || [ "$RUNNER" = mi325x ]; then + export MORI_DISABLE_AUTO_XGMI="${MORI_DISABLE_AUTO_XGMI:-0}" + export MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-1}" + export MORI_APP_LOG_LEVEL="${MORI_APP_LOG_LEVEL:-info}" + export MORI_SHMEM_LOG_LEVEL="${MORI_SHMEM_LOG_LEVEL:-info}" + export MORI_IO_LOG_LEVEL="${MORI_IO_LOG_LEVEL:-info}" + [ "$CX_BENCH" != mori ] \ + || export CX_IMAGE="${CX_IMAGE:-$CX_IMAGE_AMD_MORI_MI325}" +fi +if [ "$CX_BENCH" = mori ]; then + if [ "$NODES" -gt 1 ]; then + export CX_MORI_KERNEL_TYPE=internode-v1 + elif [ "$RUNNER" = mi300x ] || [ "$RUNNER" = mi325x ]; then + export CX_MORI_KERNEL_TYPE="${CX_MORI_KERNEL_TYPE:-asyncll}" + else + export CX_MORI_KERNEL_TYPE="${CX_MORI_KERNEL_TYPE:-intranode}" + fi +fi +IMAGE="${CX_IMAGE:-$(cx_default_image "$RUNNER")}" +export CX_RUNNER="$RUNNER" CX_NGPUS="$NGPUS" CX_NODES="$NODES" +export CX_GPUS_PER_NODE="$GPN" CX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" CX_TS="$TS" +export CX_SCALE_UP_TRANSPORT=xgmi +if [ "$NODES" -gt 1 ]; then + export CX_SCOPE=scale-out CX_SCALE_OUT_TRANSPORT=rdma + export CX_TRANSPORT=xgmi-rdma CX_TOPO="${RUNNER}-xgmi-rdma" +else + export CX_SCOPE=scale-up CX_TRANSPORT=xgmi CX_TOPO="${RUNNER}-xgmi" + unset CX_SCALE_OUT_TRANSPORT +fi +export CX_RUN_TIMEOUT="${CX_RUN_TIMEOUT:-1800}" +cx_validate_shard_control "$CX_DIR" +cx_load_network_control_mode "$CX_DIR" || cx_die "cannot resolve network control mode" +cx_apply_network_profile "$NODES" "$CX_TRANSPORT" +cx_require_vars CX_PARTITION CX_SQUASH_DIR CX_STAGE_DIR +PARTITION="$CX_PARTITION"; SQUASH_DIR="$CX_SQUASH_DIR" + +cx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$CX_BENCH" +cx_set_failure_stage repository-stage +MOUNT_SRC="$(cx_stage_path "$REPO_ROOT" "$CX_STAGE_DIR")" +cx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +cx_prepare_runtime_marker "$MOUNT_SRC" +[ "${CX_DRYRUN:-0}" != 1 ] || { cx_log "CX_DRYRUN=1 - not allocating"; exit 0; } +cx_set_failure_stage registry-verification +cx_verify_registry_image "$IMAGE" +cx_set_failure_stage scheduler-allocation +command -v salloc >/dev/null || cx_die "salloc not found on this runner" + +allocation=(--partition="$PARTITION" --nodes="$NODES" --gres=gpu:"$GPN" + --time="$TIME_MIN") +if [ "$RUNNER" = mi355x ]; then + allocation+=(--exclusive) +fi +if [ "$NODES" = 1 ]; then + allocation+=(--cpus-per-task="$CPUS_PER_TASK") +else + allocation+=(--ntasks-per-node="$GPN" --cpus-per-task="$((CPUS_PER_TASK / GPN))") +fi +excluded_nodes="$EXCLUDE_NODES" +for allocation_attempt in 1 2 3; do + attempt_allocation=("${allocation[@]}") + if [ -n "$NODELIST" ]; then + cx_log "using configured node pin" + attempt_allocation+=(--nodelist="$NODELIST") + elif [ -n "$excluded_nodes" ]; then + attempt_allocation+=(--exclude="$excluded_nodes") + fi + export CX_SALLOC_ATTEMPT="$allocation_attempt" + export CX_NETWORK_VALIDATION_ATTEMPT="$allocation_attempt" + cx_salloc_jobid "${attempt_allocation[@]}" + [ -n "$JOB_ID" ] || cx_die "could not resolve allocated JOB_ID from salloc" + cx_set_failure_stage setup + cx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$CX_TRANSPORT" + cx_set_failure_stage container-import + if SQUASH_FILE="$(cx_ensure_squash_on_job \ + "$JOB_ID" "$SQUASH_DIR" "$IMAGE" "${CX_LOCK_DIR:-}")"; then + break + fi + if [ -n "$NODELIST" ] || [ "$allocation_attempt" = 3 ]; then + cx_die "allocated nodes failed container import" + fi + rejected_nodes="$(cx_allocation_nodes_csv "$JOB_ID")" \ + || cx_die "cannot identify nodes from a rejected allocation" + cx_log "allocated nodes failed container import; retrying elsewhere" + cx_cancel_job "$JOB_ID" || cx_die "cannot release a rejected allocation" + cx_clear_allocation_jobid || cx_die "cannot reset rejected allocation state" + JOB_ID="" + [ -z "$excluded_nodes" ] || excluded_nodes+=, + excluded_nodes+="$rejected_nodes" +done +unset CX_SALLOC_ATTEMPT CX_NETWORK_VALIDATION_ATTEMPT +cx_set_failure_stage container-hash +import_log="$(cx_private_log_path image-hash)" +if ! COLLECTIVEX_SQUASH_SHA256="$( + srun --jobid="$JOB_ID" --nodes=1 --ntasks=1 --chdir=/tmp \ + --export="$(cx_host_exports)" sha256sum "$SQUASH_FILE" \ + 2>>"$import_log" | awk 'NR==1 {print $1}' +)"; then + cx_fail_stage container-hash "$import_log" +fi +[[ "$COLLECTIVEX_SQUASH_SHA256" =~ ^[0-9a-f]{64}$ ]] \ + || cx_fail_stage container-hash "$import_log" +export COLLECTIVEX_SQUASH_SHA256 +cx_preflight_allocation "$JOB_ID" "$NODES" "$MOUNT_SRC" "$SQUASH_FILE" \ + "${CX_SHARD_FILE:-}" +CONTAINER_MOUNTS="$MOUNT_SRC:$MOUNT_DIR$DEVICE_MOUNTS" + +if [ "$NODES" = 1 ]; then + run_rc=0 + cx_set_failure_stage container-launch + runtime_log="$(cx_private_log_path runtime)" + srun --jobid="$JOB_ID" --chdir=/tmp --container-image="$SQUASH_FILE" \ + --container-mounts="$CONTAINER_MOUNTS" --container-writable --container-remap-root \ + --no-container-mount-home --container-workdir="$MOUNT_DIR/experimental/CollectiveX" \ + --no-container-entrypoint --export="$(cx_container_exports)" \ + bash "$MOUNT_DIR/experimental/CollectiveX/runtime/run_in_container.sh" \ + >"$runtime_log" 2>&1 || run_rc=$? +else + SOURCE_BACKEND_ENV='case "${SLURM_NODEID:-}" in ""|*[!0-9]*) exit 66;; esac; env_file="/ix/experimental/CollectiveX/.cx_backend/env/node-${SLURM_NODEID}.sh"; env_root="${env_file%/*}"; [ -d "$env_root" ] && [ ! -L "$env_root" ] || exit 66; case "$(stat -c "%a" "$env_root")" in 700|[1-7]700) ;; *) exit 66;; esac; [ -f "$env_file" ] && [ -r "$env_file" ] && [ ! -L "$env_file" ] && [ "$(stat -c "%u:%a" "$env_file")" = "$(stat -c "%u" "$env_root"):600" ] || exit 66; . "$env_file" || exit 66' + BACKEND_PROBE="$SOURCE_BACKEND_ENV"'; case "$CX_BENCH" in mori) python3 -c "import mori";; nccl-ep) python3 -c "import torch";; esac' + WRAP="${SOURCE_BACKEND_ENV}"$'\n'"$(cx_slurm_rank_wrapper)" + CX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) + run_rc=0 + cx_set_failure_stage container-launch + cx_run_distributed_shard || run_rc=$? +fi + +cx_adopt_runtime_stage "$MOUNT_SRC" +if [ "$NODES" = 1 ] && [ "$run_rc" != 0 ]; then + cx_fail_stage "$CX_FAILSAFE_MODE" "$runtime_log" || true +fi +collect_rc=0 +cx_collect_results "$MOUNT_SRC" "$REPO_ROOT" || collect_rc=$? +[ "$run_rc" != 0 ] || [ "$collect_rc" = 0 ] || cx_set_failure_stage artifact-collection +final_rc="$run_rc" +[ "$final_rc" != 0 ] || final_rc="$collect_rc" +rm -f "$MOUNT_SRC"/experimental/CollectiveX/gpucore.* 2>/dev/null || true +cx_log "done - result artifacts collected" +exit "$final_rc" diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh new file mode 100644 index 0000000000..485b647734 --- /dev/null +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -0,0 +1,202 @@ +#!/usr/bin/env bash +# CollectiveX shared standard NVIDIA Slurm launcher (one or two nodes). +# shellcheck disable=SC2016,SC2034 +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CX_DIR="$(cd "$HERE/.." && pwd)" +REPO_ROOT="$(cd "$CX_DIR/../.." && pwd)" +# shellcheck source=../runtime/common.sh +source "$HERE/../runtime/common.sh" + +RUNNER="${CX_SHARD_SKU:-${CX_PUBLIC_RUNNER:-}}" +ALLOC_EXTRA=(); SRUN_EXTRA=(); LOCAL_IMPORT=0 +case "$RUNNER" in + h100-dgxc) PRODUCT=h100; TOPO=h100-nvlink-island; DEFAULT_TIME=45; REQUIRE_ACCOUNT=1 ;; + h200-dgxc) + PRODUCT=h200; TOPO=h200-nvlink-island; DEFAULT_TIME=45; REQUIRE_ACCOUNT=0 + SRUN_EXTRA=(--container-remap-root) + ;; + b200-dgxc) + PRODUCT=b200; TOPO=b200-nvlink-island; DEFAULT_TIME=30; REQUIRE_ACCOUNT=1 + ALLOC_EXTRA=(--mem=0) + ;; + b300) + PRODUCT=b300; TOPO=b300-nvlink-island; DEFAULT_TIME=45; REQUIRE_ACCOUNT=1 + # Do not restore ALLOC_EXTRA=(-N 1 --mem=0); it blocks two-node B300 jobs. + ALLOC_EXTRA=(--mem=0) + SRUN_EXTRA=(--mpi=none --container-remap-root) + LOCAL_IMPORT=1 + ;; + *) cx_die "set CX_SHARD_SKU or CX_PUBLIC_RUNNER to a registered NVIDIA SKU" ;; +esac +export CX_RUNNER="$RUNNER" CX_BENCH="${CX_BENCH:-deepep}" +export CX_IMAGE_PLATFORM=linux/amd64 +JOB_ID="" +cx_install_launcher_fail_safe +cx_set_failure_stage setup +cx_load_operator_config +cx_lock_canonical_gha_env "$RUNNER" + +NODES="${CX_NODES:-1}"; GPN="${CX_GPUS_PER_NODE:-8}" +SCALE_UP_DOMAIN="${CX_SCALE_UP_DOMAIN:-8}" +EXPECTED_WORLD=$((NODES * GPN)) +NGPUS="${CX_NGPUS:-$EXPECTED_WORLD}" +TIME_MIN="${CX_TIME:-$DEFAULT_TIME}" +IMAGE="${CX_IMAGE:-$(cx_default_image "$PRODUCT")}" +TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" +[ "$NODES" = 1 ] || [ "$NODES" = 2 ] \ + || cx_die "$RUNNER supports one or two nodes" +[ "$GPN" = 8 ] || cx_die "$RUNNER requires eight GPUs per node" +[ "$SCALE_UP_DOMAIN" = 8 ] || cx_die "$RUNNER requires an eight-GPU scale-up domain" +[ "$NGPUS" = "$EXPECTED_WORLD" ] \ + || cx_die "$RUNNER world size must equal nodes x GPUs per node" +case "$CX_BENCH" in + deepep|deepep-v2|deepep-hybrid|uccl|nccl-ep) ;; + *) cx_die "unsupported $RUNNER EP backend: $CX_BENCH" ;; +esac + +export CX_RUNNER="$RUNNER" CX_NGPUS="$NGPUS" CX_NODES="$NODES" +export CX_GPUS_PER_NODE="$GPN" CX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" +export CX_TS="$TS" CX_SCALE_UP_TRANSPORT=nvlink +if [ "$NODES" -gt 1 ]; then + export CX_SCOPE=scale-out CX_SCALE_OUT_TRANSPORT=rdma + export CX_TRANSPORT=nvlink-rdma CX_TOPO="${PRODUCT}-nvlink-rdma" +else + export CX_SCOPE=scale-up CX_TRANSPORT=nvlink CX_TOPO="$TOPO" + unset CX_SCALE_OUT_TRANSPORT +fi +export CX_NCCL_HOME="${CX_NCCL_HOME:-/usr}" NCCL_CUMEM_ENABLE=1 +cx_validate_shard_control "$CX_DIR" +cx_load_network_control_mode "$CX_DIR" || cx_die "cannot resolve network control mode" +cx_apply_network_profile "$NODES" "$CX_TRANSPORT" +cx_require_vars CX_PARTITION CX_SQUASH_DIR +[ "$REQUIRE_ACCOUNT" = 0 ] || cx_require_vars CX_ACCOUNT +[ "$RUNNER" != b300 ] || cx_require_vars CX_STAGE_DIR + +cx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$CX_BENCH" +[ "${CX_DRYRUN:-0}" != 1 ] || { cx_log "CX_DRYRUN=1 - not allocating"; exit 0; } +cx_set_failure_stage registry-verification +cx_verify_registry_image "$IMAGE" +SQUASH_FILE="" +cx_set_failure_stage repository-stage +MOUNT_SRC="$(cx_stage_path "$REPO_ROOT" "${CX_STAGE_DIR:-}")" +cx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +cx_prepare_runtime_marker "$MOUNT_SRC" +CONTAINER_MOUNTS="$MOUNT_SRC:/ix" +if [ "$CX_BENCH" = deepep-v2 ] || [ "$CX_BENCH" = deepep-hybrid ]; then + cx_set_failure_stage backend-setup + cx_prepare_backend_source "$MOUNT_SRC" "$CX_BENCH" \ + || cx_die "cannot stage the pinned backend source" + export CX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.cx_sources +fi +if [ "$CX_BENCH" = deepep-v2 ]; then + cx_prepare_backend_cache "$CX_SQUASH_DIR" \ + || cx_die "cannot prepare the isolated backend cache" + CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$CX_PREPARED_BACKEND_CACHE:/cx-cache" + export CX_BACKEND_CACHE_ROOT=/cx-cache +fi + +cx_set_failure_stage scheduler-allocation +command -v salloc >/dev/null || cx_die "salloc not found on this runner" +allocation=(--partition="$CX_PARTITION" --nodes="$NODES" --gres=gpu:"$GPN" --exclusive + --time="$TIME_MIN" "${ALLOC_EXTRA[@]}") +[ "$NODES" = 1 ] || allocation+=(--ntasks-per-node="$GPN") +[ -z "${CX_ACCOUNT:-}" ] || allocation+=(--account="$CX_ACCOUNT") +[ -z "${CX_QOS:-}" ] || allocation+=(--qos="$CX_QOS") +[ -z "${CX_NODELIST:-}" ] || allocation+=(--nodelist="$CX_NODELIST") +excluded_nodes="${CX_EXCLUDE_NODES:-}" +for allocation_attempt in 1 2 3; do + validation_failure="" + attempt_allocation=("${allocation[@]}") + [ -z "$excluded_nodes" ] || attempt_allocation+=(--exclude="$excluded_nodes") + export CX_SALLOC_ATTEMPT="$allocation_attempt" + export CX_NETWORK_VALIDATION_ATTEMPT="$allocation_attempt" + cx_salloc_jobid "${attempt_allocation[@]}" + [ -n "$JOB_ID" ] || cx_die "could not resolve allocated JOB_ID from salloc" + cx_set_failure_stage setup + if ! cx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$CX_TRANSPORT" 0; then + validation_failure=network + elif [ "$RUNNER" = b300 ] \ + && ! cx_validate_cuda_context_on_job "$JOB_ID" "$NODES" "$GPN"; then + validation_failure=cuda-context + else + break + fi + retryable=0 + [ "$RUNNER:$validation_failure" != h100-dgxc:network ] || retryable=1 + [ "$RUNNER:$validation_failure" != b300:cuda-context ] || retryable=1 + if [ "$retryable" = 0 ] || [ "$allocation_attempt" = 3 ]; then + if [ "$validation_failure" = network ]; then + cx_fail_stage setup "$CX_NETWORK_PROFILE_LOG" || true + cx_die "allocated nodes failed the network profile" + fi + cx_fail_stage setup "$CX_CUDA_CONTEXT_LOG" || true + cx_die "allocated nodes failed accelerator context validation" + fi + rejected_nodes="$(cx_allocation_nodes_csv "$JOB_ID")" \ + || cx_die "cannot identify nodes from a rejected allocation" + cx_log "allocated nodes failed $validation_failure validation; retrying elsewhere" + cx_cancel_job "$JOB_ID" || cx_die "cannot release a rejected allocation" + cx_clear_allocation_jobid || cx_die "cannot reset rejected allocation state" + JOB_ID="" + [ -z "$excluded_nodes" ] || excluded_nodes+=, + excluded_nodes+="$rejected_nodes" +done +unset CX_SALLOC_ATTEMPT CX_NETWORK_VALIDATION_ATTEMPT +if [ "$RUNNER" = b200-dgxc ] && [ "$CX_BENCH" = deepep ] && [ "$NODES" -gt 1 ]; then + gdrcopy_log="$(cx_private_log_path gdrcopy-preflight)" + if ! srun --jobid="$JOB_ID" --nodes="$NODES" --ntasks="$NODES" --ntasks-per-node=1 \ + --chdir=/tmp --export="$(cx_host_exports)" test -c /dev/gdrdrv \ + >"$gdrcopy_log" 2>&1; then + cx_log "ERROR: allocated B200 nodes do not expose the required GDRCopy device" + cx_fail_stage setup "$gdrcopy_log" || true + cx_die "B200 DeepEP CPU NIC handling requires GDRCopy" + fi + CONTAINER_MOUNTS="$CONTAINER_MOUNTS,/dev/gdrdrv:/dev/gdrdrv" +fi +if [ "$LOCAL_IMPORT" = 1 ]; then + cx_set_failure_stage container-import + SQUASH_FILE="$(CX_ENROOT_LOCAL_IMPORT=1 cx_ensure_squash "$CX_SQUASH_DIR" "$IMAGE")" + cx_set_failure_stage container-hash + cx_export_squash_identity "$SQUASH_FILE" +else + cx_set_failure_stage container-import + SQUASH_FILE="$(cx_ensure_squash_on_job "$JOB_ID" "$CX_SQUASH_DIR" "$IMAGE")" + cx_set_failure_stage container-hash + cx_export_squash_identity "$SQUASH_FILE" +fi +cx_preflight_allocation "$JOB_ID" "$NODES" "$MOUNT_SRC" "$SQUASH_FILE" \ + "${CX_SHARD_FILE:-}" + +if [ "$NODES" = 1 ]; then + run_rc=0 + cx_set_failure_stage container-launch + runtime_log="$(cx_private_log_path runtime)" + srun --jobid="$JOB_ID" --container-image="$SQUASH_FILE" \ + --container-mounts="$CONTAINER_MOUNTS" --no-container-mount-home \ + --container-workdir=/ix/experimental/CollectiveX --no-container-entrypoint \ + "${SRUN_EXTRA[@]}" --export="$(cx_container_exports)" \ + bash /ix/experimental/CollectiveX/runtime/run_in_container.sh \ + >"$runtime_log" 2>&1 || run_rc=$? +else + SOURCE_BACKEND_ENV='case "${SLURM_NODEID:-}" in ""|*[!0-9]*) exit 66;; esac; env_file="/ix/experimental/CollectiveX/.cx_backend/env/node-${SLURM_NODEID}.sh"; env_root="${env_file%/*}"; [ -d "$env_root" ] && [ ! -L "$env_root" ] || exit 66; case "$(stat -c "%a" "$env_root")" in 700|[1-7]700) ;; *) exit 66;; esac; [ -f "$env_file" ] && [ -r "$env_file" ] && [ ! -L "$env_file" ] && [ "$(stat -c "%u:%a" "$env_file")" = "$(stat -c "%u" "$env_root"):600" ] || exit 66; . "$env_file" || exit 66' + BACKEND_PROBE="$SOURCE_BACKEND_ENV"'; case "$CX_BENCH" in deepep) python3 -c "from deep_ep import Buffer";; deepep-v2) python3 -c "import deep_ep; assert hasattr(deep_ep, '\''ElasticBuffer'\'')";; deepep-hybrid) python3 -c "import deep_ep; assert hasattr(deep_ep, '\''HybridEPBuffer'\'')";; uccl) python3 -c "import torch; from uccl_deepep import Buffer";; nccl-ep) python3 -c "import torch";; esac' + WRAP="${SOURCE_BACKEND_ENV}"$'\n'"$(cx_slurm_rank_wrapper)" + CX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable "${SRUN_EXTRA[@]}") + run_rc=0 + cx_set_failure_stage container-launch + cx_run_distributed_shard || run_rc=$? +fi + +cx_adopt_runtime_stage "$MOUNT_SRC" +if [ "$NODES" = 1 ] && [ "$run_rc" != 0 ]; then + cx_fail_stage "$CX_FAILSAFE_MODE" "$runtime_log" || true +fi +collect_rc=0 +cx_collect_results "$MOUNT_SRC" "$REPO_ROOT" || collect_rc=$? +[ "$run_rc" != 0 ] || [ "$collect_rc" = 0 ] || cx_set_failure_stage artifact-collection +final_rc="$run_rc" +[ "$final_rc" != 0 ] || final_rc="$collect_rc" +cx_log "done - result artifacts collected" +exit "$final_rc" diff --git a/experimental/CollectiveX/requirements.txt b/experimental/CollectiveX/requirements.txt new file mode 100644 index 0000000000..23ed3544c0 --- /dev/null +++ b/experimental/CollectiveX/requirements.txt @@ -0,0 +1,5 @@ +# Host-side matrix generation. GPU libraries are supplied by benchmark images. +PyYAML==6.0.2 + +# Canonical workload serialization. +numpy>=1.26,<3 diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh new file mode 100644 index 0000000000..5679c10cb7 --- /dev/null +++ b/experimental/CollectiveX/runtime/common.sh @@ -0,0 +1,3046 @@ +# shellcheck shell=bash +# CollectiveX — shared launcher helpers (sourced, not executed). +# +# Cluster-generic scaffolding only (Slurm/container/build/staging); no +# model-serving. Logging goes to stderr so functions can `echo` a single +# result on stdout. + +CX_SQUASH_FORMAT_VERSION="repro-v1" +CX_SQUASH_SOURCE_DATE_EPOCH=1 +CX_DEEPEP_V2_COMMIT="fa8a9b16898204afd347c663b89e65ef87dc6ce6" # pragma: allowlist secret +CX_DEEPEP_V2_TREE="29809e75c5874e6609dac4804e7b651d5226959f" # pragma: allowlist secret +CX_DEEPEP_V2_FMT_COMMIT="a4c7e17133ee9cb6a2f45545f6e974dd3c393efa" # pragma: allowlist secret +# Consumed by run_in_container.sh after this helper is sourced. +# shellcheck disable=SC2034 +CX_DEEPEP_V2_NCCL_CHECK_COMMIT="93d0564188f7a0a6288c6e316484861b0efa042e" # pragma: allowlist secret +CX_DEEPEP_V2_INIT_SHA256="f090bbacc38c7d5dba29f07fd38a918eb820b6adac6f76903fe56273060e4870" +CX_DEEPEP_HYBRID_COMMIT="e0a5b1d9848ab3e7b4a67842bf06f067bfac67f8" # pragma: allowlist secret +CX_DEEPEP_HYBRID_TREE="d77aeab7f1bb52b615666fe178d26ced41fae08e" # pragma: allowlist secret +CX_DEEPEP_HYBRID_NCCL_COMMIT="1e0c869c39bb33f1034cb9920bd2a8a8406f04a3" # pragma: allowlist secret +unset COLLECTIVEX_OPERATOR_CONFIG_LOADED COLLECTIVEX_EPHEMERAL_CONFIG_PATH + +cx_log() { printf '[collectivex] %s\n' "$*" >&2; } +cx_die() { printf '[collectivex] FATAL: %s\n' "$*" >&2; exit 1; } + +# Public failure telemetry is a closed vocabulary. Raw scheduler, container, +# host, and filesystem diagnostics stay in the mode-0600 private logs. +cx_set_failure_stage() { + local stage="$1" + case "$stage" in + setup|repository-stage|registry-verification|scheduler-allocation|container-import) ;; + container-hash|container-launch|backend-setup|execution|artifact-collection) ;; + *) cx_die "invalid launcher failure stage" ;; + esac + export CX_FAILSAFE_MODE="$stage" +} + +cx_fail_stage() { + local stage="$1" log_path="${2:-}" diagnostic="unknown" + cx_set_failure_stage "$stage" + if [ -n "$log_path" ] && [ -f "$log_path" ]; then + if grep -aEqi 'no space left|disk quota|quota exceeded' "$log_path"; then + diagnostic="storage-capacity" + elif grep -aEqi 'permission denied|operation not permitted|read-only file system|source mount (creation|ownership validation|permission inspection|permission normalization|permission validation) failed' "$log_path"; then + diagnostic="storage-permission" + elif grep -aEqi 'outside one realized LSA domain|lsa(Size| team| domain).*(mismatch|invalid|expected)|ranks.*not in (one|the same) nvlink.domain' "$log_path" \ + || { [ "${CX_BENCH:-}" = deepep-v2 ] \ + && grep -aEqi 'nccl[.]cu:(111|112)([^0-9]|$)' "$log_path"; }; then + diagnostic="accelerator-topology" + elif grep -aEqi 'cuda driver version is insufficient|call requires newer driver|cudaErrorCallRequiresNewerDriver|CUDA_ERROR_SYSTEM_DRIVER_MISMATCH|unsupported toolchain' "$log_path"; then + diagnostic="accelerator-driver" + elif grep -aEqi 'cudaErrorDevicesUnavailable|CUDA_ERROR_DEVICE_UNAVAILABLE|CUDA-capable device\(s\) is/are busy or unavailable|primary context retain failed' "$log_path"; then + diagnostic="accelerator-unavailable" + elif grep -aEqi 'ncclDevCommCreate|ncclCommWindowRegister|ncclGetLsa(Device)?Pointer|Communicator does not support symmetric memory|Symmetric memory is not supported' "$log_path" \ + || { [ "${CX_BENCH:-}" = deepep-v2 ] \ + && grep -aEqi 'nccl[.]cu:(106|127|128|129|135)([^0-9]|$)' "$log_path"; }; then + diagnostic="nccl-device-api" + elif grep -aEqi 'NVCC (PTX )?compilation failed|cuobjdump failed|invalid device (kernel )?image|no kernel image is available' "$log_path"; then + diagnostic="jit-toolchain" + elif grep -aEqi 'cuda out of memory|CUDA_ERROR_OUT_OF_MEMORY|out of memory.*cuda' "$log_path"; then + diagnostic="accelerator-memory" + elif grep -aEqi 'does not match its pinned image contract|requires the exact pinned|version mismatch' "$log_path"; then + diagnostic="backend-version" + elif grep -aEqi 'nvshmem is unavailable|build-tool installation failed' "$log_path"; then + diagnostic="backend-dependency" + elif grep -aEqi 'revision fetch failed|submodule fetch failed|package installation failed|staged source is invalid|source (pin resolution|seed validation|seed copy|checkout creation|publication validation|existing source validation) failed' "$log_path"; then + diagnostic="backend-source" + elif grep -aEqi 'backend preparation failed|backend import failed|build (failed|is incomplete)|cache (mount identity )?validation failed|import failed' "$log_path"; then + diagnostic="backend-build" + elif grep -aEqi 'failed to mount|squashfs|enroot|pyxis|mount.*invalid argument|invalid argument.*mount' "$log_path"; then + diagnostic="container-runtime" + elif grep -aEqi 'command not found|not found on this runner|git lookup failed' "$log_path"; then + diagnostic="missing-runtime" + elif grep -aEqi 'too many requests|rate.?limit' "$log_path"; then + diagnostic="registry-rate-limit" + elif grep -aEqi 'ncclRemoteError|remote process exited|connection closed by peer' "$log_path"; then + diagnostic="collective-remote" + elif grep -aEqi 'ncclSystemError|unhandled system error' "$log_path"; then + diagnostic="collective-system" + elif grep -aEqi 'ncclInternalError|internal check failed' "$log_path"; then + diagnostic="collective-internal" + elif grep -aEqi 'ncclInvalidUsage|invalid usage' "$log_path"; then + if grep -aEqi 'dist[.]init_process_group|init_process_group[(]' "$log_path"; then + diagnostic="collective-init-invalid-usage" + elif grep -aEqi 'dist[.]all_gather_object|all_gather_object[(]' "$log_path"; then + diagnostic="collective-consensus-invalid-usage" + elif grep -aEqi 'dist[.]all_to_all_single|all_to_all_single[(]' "$log_path"; then + diagnostic="collective-alltoall-invalid-usage" + else + diagnostic="collective-invalid-usage" + fi + elif grep -aEqi 'timed out|operation timeout|wait timeout after|watchdog.*timeout|timeout: sending signal|connection reset|could not resolve|TLS|certificate' "$log_path"; then + diagnostic="network-or-timeout" + elif grep -aEqi 'salloc:|srun:.*(unable to create step|step creation|invalid partition|invalid account)|unable to create step|job allocation' "$log_path"; then + diagnostic="scheduler" + elif grep -aEqi 'ModuleNotFoundError|ImportError:' "$log_path"; then + diagnostic="python-import" + elif grep -aEqi 'AttributeError:|TypeError:.*(unexpected|argument|operand)|has no attribute' "$log_path"; then + diagnostic="backend-api" + elif grep -aEqi 'AssertionError:' "$log_path"; then + diagnostic="python-assertion" + elif grep -aEqi 'RuntimeError:' "$log_path"; then + diagnostic="python-runtime" + elif grep -aEqi 'ValueError:' "$log_path" \ + && grep -aEqi 'ep_harness[.]py' "$log_path"; then + diagnostic="harness-value" + elif grep -aEqi 'ValueError:' "$log_path" \ + && grep -aEqi 'workload[.]py|make_workloads[.]py' "$log_path"; then + diagnostic="workload-value" + elif grep -aEqi 'ValueError:' "$log_path" \ + && grep -aEqi 'run_ep[.]py' "$log_path"; then + diagnostic="runner-value" + elif grep -aEqi 'ValueError:' "$log_path" \ + && grep -aEqi 'ep_deepep[.]py' "$log_path"; then + diagnostic="deepep-adapter-value" + elif grep -aEqi 'ValueError:' "$log_path" \ + && grep -aEqi '/(torch|numpy)/|site-packages/(torch|numpy)' "$log_path"; then + diagnostic="dependency-value" + elif grep -aEqi 'ValueError:' "$log_path"; then + diagnostic="python-value" + elif grep -aEqi 'KeyError:' "$log_path"; then + diagnostic="python-key" + elif grep -aEqi '(FileNotFoundError|PermissionError|IsADirectoryError|NotADirectoryError|OSError):' "$log_path"; then + diagnostic="python-os" + elif grep -aEqi '(NotImplemented|System)Error:' "$log_path"; then + diagnostic="python-system" + elif grep -aEqi 'DistBackendError:' "$log_path"; then + diagnostic="collective-backend" + elif grep -aEqi 'CalledProcessError:' "$log_path"; then + diagnostic="python-subprocess" + elif grep -aEqi 'Traceback \(most recent call last\)' "$log_path"; then + diagnostic="python-exception" + elif grep -aEqi 'SHARD done: [0-9]+/[0-9]+ case\(s\) failed|WARN: .* run failed rc=|completed with invalid semantic evidence' "$log_path"; then + diagnostic="benchmark-case-failure" + elif [ -s "$log_path" ]; then + diagnostic="unclassified" + else + diagnostic="empty-log" + fi + fi + cx_log "ERROR: failure-class=$stage diagnostic=$diagnostic" + # Surface the failing stage's captured output verbatim on stdout for debugging. + # Bounded by line count only (anti-flood, not redaction): this is public CI + # output, so the tail carries whatever the underlying tool emitted, unmasked. + if [ -n "$log_path" ] && [ -f "$log_path" ] && [ -s "$log_path" ]; then + local tail_lines="${CX_LOG_TAIL_LINES:-100}" + [[ "$tail_lines" =~ ^[0-9]+$ ]] || tail_lines=100 + cx_log "--- $stage log tail (last $tail_lines lines, verbatim) ---" + tail -n "$tail_lines" "$log_path" >&2 || true + cx_log "--- end $stage log tail ---" + fi + return 1 +} + +cx_job_root_is_safe() { + local root="$1" + if [[ "$root" =~ ^/tmp/inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+$ ]]; then + : + elif [[ "$root" =~ ^/tmp/inferencex-collectivex-parent-([0-9]+)-([0-9]+)-([A-Za-z0-9._-]+)/inferencex-collectivex-([0-9]+)-([0-9]+)-([A-Za-z0-9._-]+)$ ]]; then + [ "${BASH_REMATCH[1]}" = "${BASH_REMATCH[4]}" ] \ + && [ "${BASH_REMATCH[2]}" = "${BASH_REMATCH[5]}" ] \ + && [ "${BASH_REMATCH[3]}" = "${BASH_REMATCH[6]}" ] || return 1 + else + return 1 + fi + [ -d "$root" ] && [ ! -L "$root" ] \ + && [ "$(stat -c '%u:%a' "$root" 2>/dev/null)" = "$(id -u):700" ] +} + +# Runner-local deployment settings are strict JSON kept outside the checkout. +# Only the selected runner's allowlisted values are exported; the document is +# never sourced or evaluated as shell. +cx_load_operator_config() { + [ -n "${COLLECTIVEX_OPERATOR_CONFIG_LOADED:-}" ] \ + && [ "$COLLECTIVEX_OPERATOR_CONFIG_LOADED" = "$$" ] && return 0 + local config_path generated=0 parsed_path config_log key value + local audit_salt_override validation_code + audit_salt_override="${COLLECTIVEX_OPERATOR_AUDIT_SALT:-}" + unset COLLECTIVEX_OPERATOR_AUDIT_SALT + unset CX_PARTITION CX_ACCOUNT CX_QOS CX_SQUASH_DIR CX_STAGE_DIR CX_ENROOT_CACHE_PATH + unset ENROOT_CACHE_PATH + unset CX_EXCLUDE_NODES CX_NODELIST CX_LOCK_DIR CX_MASTER_PORT + unset CX_SOCKET_IFNAME CX_RDMA_DEVICES CX_IB_GID_INDEX CX_RDMA_SERVICE_LEVEL + unset CX_RDMA_TRAFFIC_CLASS + unset CX_AUDIT_SALT + unset MASTER_ADDR MASTER_PORT RANK WORLD_SIZE LOCAL_RANK LOCAL_WORLD_SIZE + config_path="${COLLECTIVEX_OPERATOR_CONFIG:-${XDG_CONFIG_HOME:-${HOME}/.config}/inferencex/collectivex.json}" + if [ -n "${COLLECTIVEX_OPERATOR_CONFIG_CONTENT:-}" ]; then + umask 077 + if cx_job_root_is_safe "${CX_JOB_ROOT:-}"; then + config_path="$CX_JOB_ROOT/operator-config.json" + (set -C; : > "$config_path") 2>/dev/null \ + || cx_die "cannot create ephemeral runner configuration" + else + config_path="$(mktemp /tmp/inferencex-collectivex-config.XXXXXX)" \ + || cx_die "cannot create ephemeral runner configuration" + fi + COLLECTIVEX_EPHEMERAL_CONFIG_PATH="$config_path" + generated=1 + if ! printf '%s' "$COLLECTIVEX_OPERATOR_CONFIG_CONTENT" > "$config_path"; then + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT + rm -f -- "$config_path" + unset COLLECTIVEX_EPHEMERAL_CONFIG_PATH + cx_die "cannot materialize runner configuration" + fi + elif [ "${COLLECTIVEX_OPERATOR_CONFIG_REQUIRED:-0}" = 1 ]; then + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT + cx_die "runner configuration is unavailable" + fi + unset COLLECTIVEX_OPERATOR_CONFIG_CONTENT COLLECTIVEX_OPERATOR_CONFIG_REQUIRED + if [ ! -e "$config_path" ]; then + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" != 1 ] \ + || cx_die "runner configuration is unavailable" + COLLECTIVEX_OPERATOR_CONFIG_LOADED="$$" + return 0 + fi + umask 077 + parsed_path="$(mktemp /tmp/inferencex-collectivex-parsed.XXXXXX)" || { + [ "$generated" = 0 ] || rm -f -- "$config_path" + cx_die "cannot parse runner configuration" + } + config_log="$(cx_private_log_path operator-config)" + if ! python3 - "$config_path" "${CX_RUNNER:-${CX_SHARD_SKU:-${CX_PUBLIC_RUNNER:-}}}" \ + "${COLLECTIVEX_CANONICAL_GHA:-0}" "$audit_salt_override" \ + > "$parsed_path" 2> "$config_log" <<'PY' +import json +import os +import posixpath +import re +import stat +import sys + +RUNNERS = { + "h100-dgxc", "h200-dgxc", "b200-dgxc", "b300", + "gb200", "gb300", "mi300x", "mi325x", "mi355x", +} +FIELDS = { + "partition": "CX_PARTITION", + "account": "CX_ACCOUNT", + "qos": "CX_QOS", + "squash_dir": "CX_SQUASH_DIR", + "stage_dir": "CX_STAGE_DIR", + "enroot_cache_path": "CX_ENROOT_CACHE_PATH", + "exclude_nodes": "CX_EXCLUDE_NODES", + "nodelist": "CX_NODELIST", + "lock_dir": "CX_LOCK_DIR", + "socket_ifname": "CX_SOCKET_IFNAME", + "rdma_devices": "CX_RDMA_DEVICES", + "ib_gid_index": "CX_IB_GID_INDEX", + "rdma_service_level": "CX_RDMA_SERVICE_LEVEL", + "rdma_traffic_class": "CX_RDMA_TRAFFIC_CLASS", +} +NETWORK_FIELDS = { + "socket_ifname", "rdma_devices", "ib_gid_index", "rdma_service_level", + "rdma_traffic_class", +} +REQUIRED = { + "h100-dgxc": {"partition", "account", "squash_dir"}, + "h200-dgxc": {"partition", "squash_dir"}, + "b200-dgxc": {"partition", "account", "squash_dir"}, + "b300": {"partition", "account", "squash_dir"}, + "gb200": {"partition", "account", "storage_roots"}, + "gb300": {"partition", "account", "squash_dir", "enroot_cache_path"}, + "mi300x": {"partition", "squash_dir"}, + "mi325x": {"partition", "squash_dir"}, + "mi355x": {"partition", "squash_dir"}, +} +ALLOWED = { + "h100-dgxc": REQUIRED["h100-dgxc"] | {"exclude_nodes", "stage_dir"} | NETWORK_FIELDS, + "h200-dgxc": REQUIRED["h200-dgxc"] | {"account", "exclude_nodes", "stage_dir"} | NETWORK_FIELDS, + "b200-dgxc": REQUIRED["b200-dgxc"] | { + "exclude_nodes", "nodelist", "stage_dir", "qos", + } | NETWORK_FIELDS, + "b300": REQUIRED["b300"] | {"exclude_nodes", "stage_dir"} | NETWORK_FIELDS, + "gb200": REQUIRED["gb200"] | NETWORK_FIELDS, + "gb300": REQUIRED["gb300"] | {"stage_dir"} | NETWORK_FIELDS, + "mi300x": REQUIRED["mi300x"] | {"exclude_nodes", "nodelist", "stage_dir", "lock_dir"} | NETWORK_FIELDS, + "mi325x": REQUIRED["mi325x"] | {"exclude_nodes", "nodelist", "stage_dir", "lock_dir"} | NETWORK_FIELDS, + "mi355x": REQUIRED["mi355x"] | {"exclude_nodes", "nodelist", "stage_dir", "lock_dir"} | NETWORK_FIELDS, +} +TOKEN = re.compile(r"^[A-Za-z0-9_.\[\],-]+$") +PATH = re.compile(r"^/[A-Za-z0-9._/+\-]+$") +IPV4 = re.compile(r"(? 65536 + ): + raise ValueError + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(path, flags) + try: + opened = os.fstat(descriptor) + if (opened.st_dev, opened.st_ino) != (metadata.st_dev, metadata.st_ino): + raise ValueError + payload = b"" + while len(payload) <= 65536: + chunk = os.read(descriptor, 65537 - len(payload)) + if not chunk: + break + payload += chunk + document = json.loads( + payload.decode("utf-8"), + object_pairs_hook=pairs, + parse_constant=lambda _: (_ for _ in ()).throw(ValueError()), + ) + finally: + os.close(descriptor) + if ( + set(document) not in ( + {"schema_version", "runners"}, + {"schema_version", "audit_salt", "runners"}, + ) + or type(document["schema_version"]) is not int + or document["schema_version"] != 1 + ): + raise ValueError + audit_salt = document.get("audit_salt") + if ( + (audit_salt is not None and ( + not isinstance(audit_salt, str) or not AUDIT_SALT.fullmatch(audit_salt) + )) + or (audit_override and not AUDIT_SALT.fullmatch(audit_override)) + or (audit_salt is not None and audit_override and audit_salt != audit_override) + ): + raise ValueError + audit_salt = audit_salt or audit_override or None + if audit_required == "1" and audit_salt is None: + raise ValueError + runners = document["runners"] + if ( + not isinstance(runners, dict) or not runners or set(runners) - RUNNERS + or runner not in runners + ): + raise ValueError + selected = None + for name, config in runners.items(): + if not isinstance(config, dict): + raise ValueError + if name == runner: + missing = sorted(REQUIRED[name] - set(config)) + if missing: + print( + "validation-missing-required-" + "-".join(missing), + file=sys.stderr, + ) + raise SystemExit(1) + if set(config) - ALLOWED[name]: + raise ValueError + for field, value in config.items(): + if field == "storage_roots": + if ( + not isinstance(value, list) or not 1 <= len(value) <= 16 + or len(value) != len(set(value)) or not all(valid_path(item) for item in value) + ): + raise ValueError + elif field == "socket_ifname": + if not isinstance(value, str) or not INTERFACES.fullmatch(value): + raise ValueError + elif field == "rdma_devices": + if not isinstance(value, str) or not RDMA_DEVICES.fullmatch(value): + raise ValueError + elif field == "ib_gid_index": + config[field] = bounded_integer(value, 255) + elif field == "rdma_service_level": + config[field] = bounded_integer(value, 15) + elif field == "rdma_traffic_class": + config[field] = bounded_integer(value, 255) + elif field.endswith(("_dir", "_path")): + if not valid_path(value): + raise ValueError + elif ( + not isinstance(value, str) or not value or len(value) > 512 + or not TOKEN.fullmatch(value) or IPV4.search(value) + ): + raise ValueError + if name == runner: + selected = dict(config) + if selected is None: + raise ValueError + roots = selected.pop("storage_roots", None) + if roots is not None: + for root in roots: + squash = posixpath.join(root, "collectivex", "containers") + stage = posixpath.join(root, "collectivex", "stage") + probes = [] + try: + for directory in (squash, stage): + os.makedirs(directory, mode=0o700, exist_ok=True) + probe = posixpath.join(directory, f".write-probe-{os.getpid()}") + fd = os.open(probe, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + os.close(fd) + probes.append(probe) + selected.update(squash_dir=squash, stage_dir=stage) + break + except OSError: + pass + finally: + for probe in probes: + try: + os.unlink(probe) + except OSError: + pass + else: + raise ValueError + if audit_salt is not None: + sys.stdout.buffer.write(b"CX_AUDIT_SALT\0" + audit_salt.encode() + b"\0") + for field, value in selected.items(): + key = FIELDS[field] + sys.stdout.buffer.write( + key.encode() + b"\0" + str(value).encode() + b"\0" + ) +except (KeyError, OSError, TypeError, UnicodeError, ValueError): + import traceback + + location = traceback.extract_tb(sys.exc_info()[2])[-1].lineno + print(f"validation-line-{location}", file=sys.stderr) + raise SystemExit(1) +PY + then + validation_code="$(head -n 1 "$config_log" 2>/dev/null || true)" + rm -f -- "$parsed_path" + [ "$generated" = 0 ] || rm -f -- "$config_path" + unset COLLECTIVEX_EPHEMERAL_CONFIG_PATH + unset COLLECTIVEX_OPERATOR_CONFIG COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL + [[ "$validation_code" =~ ^validation-(line-[0-9]+|missing-required-[a-z0-9_-]+)$ ]] \ + || validation_code="validation-unknown" + cx_die "runner-local configuration failed ($validation_code)" + fi + while IFS= read -r -d '' key && IFS= read -r -d '' value; do + printf -v "$key" '%s' "$value" + export "${key?}" + done < "$parsed_path" + rm -f -- "$parsed_path" + if [ "$generated" = 1 ] || [ "${COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL:-0}" = 1 ]; then + rm -f -- "$config_path" || cx_die "cannot remove ephemeral runner configuration" + fi + unset COLLECTIVEX_EPHEMERAL_CONFIG_PATH + unset COLLECTIVEX_OPERATOR_CONFIG COLLECTIVEX_OPERATOR_CONFIG_EPHEMERAL + COLLECTIVEX_OPERATOR_CONFIG_LOADED="$$" +} + +cx_private_log_path() { + local label="$1" tag="${COLLECTIVEX_EXECUTION_ID:-manual_$$}" path + path="$(python3 - "$tag" "$label" <<'PY' 2>/dev/null +import os +import re +import shutil +import stat +import sys +import time + +tag, label = sys.argv[1:] +if not all(re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", value) for value in (tag, label)): + raise SystemExit(1) +root = f"/tmp/inferencex-collectivex-{os.getuid()}" +job_root = os.environ.get("CX_JOB_ROOT", "") +job_parent = os.environ.get("CX_JOB_PARENT", "") +if ( + os.environ.get("COLLECTIVEX_CANONICAL_GHA") == "1" + and job_parent + and job_parent != "/tmp" +): + if ( + not os.path.isabs(job_root) + or os.path.dirname(job_root) != job_parent + or not re.fullmatch( + r"inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+", + os.path.basename(job_root), + ) + ): + raise SystemExit(1) + control = os.path.join(job_root, "control") + control_metadata = os.stat(control, follow_symlinks=False) + if ( + not stat.S_ISDIR(control_metadata.st_mode) + or control_metadata.st_uid != os.getuid() + or stat.S_IMODE(control_metadata.st_mode) != 0o700 + ): + raise SystemExit(1) + root = os.path.join(control, "private-logs") +old_umask = os.umask(0o077) +flags = os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0) +try: + try: + os.mkdir(root, 0o700) + except FileExistsError: + pass + root_fd = os.open(root, flags) + try: + metadata = os.fstat(root_fd) + if metadata.st_uid != os.getuid() or stat.S_IMODE(metadata.st_mode) != 0o700: + raise OSError("unsafe root") + cutoff = time.time() - 86400 + for entry in os.scandir(root): + try: + if ( + entry.name != tag and entry.is_dir(follow_symlinks=False) + and entry.stat(follow_symlinks=False).st_mtime < cutoff + ): + shutil.rmtree(entry.path) + except OSError: + pass + try: + os.mkdir(tag, 0o700, dir_fd=root_fd) + except FileExistsError: + pass + directory_fd = os.open(tag, flags, dir_fd=root_fd) + try: + metadata = os.fstat(directory_fd) + if metadata.st_uid != os.getuid() or stat.S_IMODE(metadata.st_mode) != 0o700: + raise OSError("unsafe directory") + log_flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + log_fd = os.open(f"{label}.log", log_flags, 0o600, dir_fd=directory_fd) + os.close(log_fd) + finally: + os.close(directory_fd) + finally: + os.close(root_fd) +finally: + os.umask(old_umask) +print(f"{root}/{tag}/{label}.log", end="") +PY +)" || cx_die "cannot create private runtime log" + printf '%s' "$path" +} + +# Manual successes delete diagnostics immediately. Canonical workflow logs survive +# until artifact upload succeeds; failed logs remain private for debugging, and a +# later run prunes abandoned directories older than 24 hours. +cx_cleanup_private_logs() { + local rc="$1" tag="${COLLECTIVEX_EXECUTION_ID:-manual_$$}" + [ "$rc" = 0 ] || return 0 + python3 - "$tag" <<'PY' >/dev/null 2>&1 || true +import os +import re +import shutil +import stat +import sys + +tag = sys.argv[1] +if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", tag): + raise SystemExit(1) +root = f"/tmp/inferencex-collectivex-{os.getuid()}" +job_root = os.environ.get("CX_JOB_ROOT", "") +job_parent = os.environ.get("CX_JOB_PARENT", "") +if ( + os.environ.get("COLLECTIVEX_CANONICAL_GHA") == "1" + and job_parent + and job_parent != "/tmp" +): + if ( + not os.path.isabs(job_root) + or os.path.dirname(job_root) != job_parent + or not re.fullmatch( + r"inferencex-collectivex-[0-9]+-[0-9]+-[A-Za-z0-9._-]+", + os.path.basename(job_root), + ) + ): + raise SystemExit(1) + control = os.path.join(job_root, "control") + control_metadata = os.stat(control, follow_symlinks=False) + if ( + not stat.S_ISDIR(control_metadata.st_mode) + or control_metadata.st_uid != os.getuid() + or stat.S_IMODE(control_metadata.st_mode) != 0o700 + ): + raise SystemExit(1) + root = os.path.join(control, "private-logs") +flags = os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0) +root_fd = os.open(root, flags) +try: + metadata = os.fstat(root_fd) + if metadata.st_uid != os.getuid() or stat.S_IMODE(metadata.st_mode) != 0o700: + raise SystemExit(1) +finally: + os.close(root_fd) +path = os.path.join(root, tag) +if os.path.isdir(path) and not os.path.islink(path): + shutil.rmtree(path) +PY +} + +cx_report_private_scheduler_failure() { + local root="${CX_JOB_ROOT:-}" tag="${COLLECTIVEX_EXECUTION_ID:-}" diagnostic + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" = 1 ] \ + && [ -n "$root" ] && [ -n "$tag" ] \ + && [[ "$tag" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || return 0 + diagnostic="$(python3 - "$root" "$tag" <<'PY' 2>/dev/null || true +import os +import re +import stat +import sys + +root, tag = sys.argv[1:] +directory = os.path.join(root, "control", "private-logs", tag) +try: + metadata = os.stat(directory, follow_symlinks=False) +except OSError: + print("missing") + raise SystemExit +if ( + not stat.S_ISDIR(metadata.st_mode) + or metadata.st_uid != os.getuid() + or stat.S_IMODE(metadata.st_mode) != 0o700 +): + print("unsafe") + raise SystemExit +payload = b"" +for name in sorted(os.listdir(directory)): + if not re.fullmatch(r"scheduler-allocation(?:-a[23])?[.]log", name): + continue + path = os.path.join(directory, name) + item = os.stat(path, follow_symlinks=False) + if ( + not stat.S_ISREG(item.st_mode) + or item.st_uid != os.getuid() + or stat.S_IMODE(item.st_mode) != 0o600 + or item.st_size > 65536 + ): + print("unsafe") + raise SystemExit + with open(path, "rb") as stream: + payload += stream.read(65537) +text = payload.decode("utf-8", "replace") +if not text: + result = "empty" +elif re.search(r"invalid (?:account|partition|qos|reservation)|association.*not permitted|access denied", text, re.I): + result = "policy" +elif re.search(r"pending job allocation|job .* pending|waiting for resource", text, re.I): + result = "pending" +elif re.search(r"requested node configuration is not available|nodes required.*not available|resources? unavailable", text, re.I): + result = "capacity" +elif re.search(r"unable to contact slurm controller|communication connection failure|socket timed out|slurmctld.*(?:down|unreachable)", text, re.I): + result = "controller" +elif re.search(r"invalid generic resource|invalid gres|invalid node count|invalid cpu count|memory specification can not be satisfied|requested.*configuration.*invalid|requested time limit is invalid|time limit.*(?:invalid|exceed)", text, re.I): + result = "resource-request" +elif re.search(r"job violates accounting[/ ]qos policy|maximum.*jobs|association.*limit|qos.*limit", text, re.I): + result = "account-limit" +elif re.search(r"invalid credential|authentication failure|authentication error", text, re.I): + result = "authentication" +elif re.search(r"job submit plugin|job_submit", text, re.I): + result = "submit-plugin" +elif re.search(r"unrecognized option|unknown option|invalid option", text, re.I): + result = "option" +elif re.search(r"allocation (?:revoked|cancelled)|job .* cancelled", text, re.I): + result = "revoked" +elif re.search(r"timed out|timeout", text, re.I): + result = "timeout" +elif re.search(r"no space left|disk quota|quota exceeded", text, re.I): + result = "storage-capacity" +else: + result = "unclassified" +print(result) +PY +)" + case "$diagnostic" in + missing|unsafe|empty|policy|pending|capacity|controller|resource-request|account-limit) ;; + authentication|submit-plugin|option|revoked|timeout|storage-capacity|unclassified) ;; + *) diagnostic=unclassified ;; + esac + cx_log "ERROR: scheduler-diagnostic=$diagnostic" +} + +# Explicit Slurm export boundary. Operator config, runner credentials, HOME, +# workspace paths, and unrelated service secrets never enter the container. +cx_container_exports() { + printf '%s' 'COLLECTIVEX_SOURCE_SHA,COLLECTIVEX_ARTIFACT_NAME,COLLECTIVEX_EXECUTION_ID,COLLECTIVEX_CONTROL_SHA256,COLLECTIVEX_IMAGE,COLLECTIVEX_IMAGE_DIGEST,COLLECTIVEX_IMAGE_DIGEST_VERIFIED,COLLECTIVEX_SQUASH_SHA256,GITHUB_REF_NAME,GITHUB_REF,GITHUB_REPOSITORY,GITHUB_JOB,GITHUB_RUN_ID,GITHUB_RUN_ATTEMPT,GITHUB_SHA,CX_RUNNER,CX_BENCH,CX_NODES,CX_GPUS_PER_NODE,CX_SCALE_UP_DOMAIN,CX_SHARD_FILE,CX_SHARD_SKU,CX_NGPUS,CX_TS,CX_TOPO,CX_SCOPE,CX_TRANSPORT,CX_SCALE_UP_TRANSPORT,CX_SCALE_OUT_TRANSPORT,CX_MODE,CX_PHASE,CX_ROUTING,CX_EPLB,CX_CASE_ID,CX_SUITE,CX_WORKLOAD_NAME,CX_QUALIFICATION_INDEX,CX_VERSION,CX_HIDDEN,CX_TOPK,CX_EXPERTS,CX_TOKENS_LADDER,CX_CANONICAL,CX_ITERS,CX_TRIALS,CX_WARMUP,CX_SAMPLES_PER_POINT,CX_WARMUP_SEMANTICS,CX_SEED,CX_RUN_TIMEOUT,CX_NCCL_HOME,CX_ALLOW_MNNVL,CX_ATTEMPT_ID,CX_RUNTIME_MARKER,CX_MORI_KERNEL_TYPE,CX_WORKLOAD_DIR,CX_BACKEND_CACHE_ROOT,CX_BACKEND_CACHE_SENTINEL_SHA256,CX_BACKEND_SOURCE_ROOT,CX_AUDIT_SALT,CX_SOCKET_IFNAME,CX_RDMA_DEVICES,CX_IB_GID_INDEX,CX_RDMA_SERVICE_LEVEL,CX_RDMA_TRAFFIC_CLASS,CX_RDMA_LINK_LAYER,MASTER_ADDR,MASTER_PORT,RANK,WORLD_SIZE,LOCAL_RANK,LOCAL_WORLD_SIZE,NCCL_NET,NCCL_SOCKET_IFNAME,GLOO_SOCKET_IFNAME,NCCL_IB_HCA,NCCL_IB_GID_INDEX,NCCL_IB_SL,NVSHMEM_DISABLE_IB,NVSHMEM_ENABLE_NIC_PE_MAPPING,NVSHMEM_HCA_LIST,NVSHMEM_IB_GID_INDEX,NVSHMEM_IB_SL,NVSHMEM_IB_ENABLE_IBGDA,NVSHMEM_IBGDA_NIC_HANDLER,EP_NIC_NAME,EP_OVERRIDE_RDMA_SL,UCCL_SOCKET_IFNAME,UCCL_IB_GID_INDEX,UCCL_IB_SL,MORI_RDMA_DEVICES,MORI_RDMA_TC,MORI_IO_TC,MORI_RDMA_SL,MORI_IO_SL,HYBRID_EP_MULTINODE,USE_NIXL,RDMA_CORE_HOME,DEEPEP_HYBRID_BUILD_MODE,NCCL_CUMEM_ENABLE,NCCL_MNNVL_ENABLE,MC_FORCE_MNNVL,MORI_DISABLE_AUTO_XGMI,MORI_ENABLE_SDMA,MORI_APP_LOG_LEVEL,MORI_SHMEM_LOG_LEVEL,MORI_IO_LOG_LEVEL,MORI_COMMIT' +} + +# Host-side utility steps need only the basic login paths. They never receive +# the complete Actions or runner environment. +cx_host_exports() { + printf '%s' 'HOME,PATH,USER,XDG_CACHE_HOME,ENROOT_CACHE_PATH' +} + +cx_prepare_runtime_marker() { + local mount_src="$1" tag="${COLLECTIVEX_EXECUTION_ID:-${CX_TS:-}}" marker + [[ "$tag" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || cx_die "cannot create runtime stage marker" + marker=".shards/runtime-stage-${tag}.txt" + mkdir -p "$mount_src/experimental/CollectiveX/.shards" >/dev/null 2>&1 \ + || cx_die "cannot create runtime stage marker" + rm -f -- "$mount_src/experimental/CollectiveX/$marker" >/dev/null 2>&1 \ + || cx_die "cannot reset runtime stage marker" + export CX_RUNTIME_MARKER="$marker" +} + +cx_write_runtime_stage() { + local stage="$1" marker="${CX_RUNTIME_MARKER:-}" + [ -n "$marker" ] || return 0 + [[ "$marker" =~ ^\.shards/runtime-stage-[A-Za-z0-9][A-Za-z0-9._-]*\.txt$ ]] \ + || return 1 + case "$stage" in backend-setup|execution) ;; *) return 1 ;; esac + printf '%s\n' "$stage" > "$marker" +} + +cx_adopt_runtime_stage() { + local mount_src="$1" marker="${CX_RUNTIME_MARKER:-}" stage="" + [ -n "$marker" ] || return 0 + if [[ "$marker" =~ ^\.shards/runtime-stage-[A-Za-z0-9][A-Za-z0-9._-]*\.txt$ ]] \ + && [ -f "$mount_src/experimental/CollectiveX/$marker" ]; then + IFS= read -r stage < "$mount_src/experimental/CollectiveX/$marker" || true + rm -f -- "$mount_src/experimental/CollectiveX/$marker" >/dev/null 2>&1 || true + case "$stage" in + backend-setup|execution) cx_set_failure_stage "$stage" ;; + esac + fi +} + +cx_require_vars() { + local name + local -a missing=() + for name in "$@"; do + [ -n "${!name:-}" ] || missing+=("$name") + done + [ "${#missing[@]}" -eq 0 ] || cx_die \ + "missing runner-local configuration: ${missing[*]} (set them in COLLECTIVEX_OPERATOR_CONFIG)" +} + +cx_bool_enabled() { + local normalized + normalized="$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')" + case "$normalized" in + 1|true|yes) return 0 ;; + *) return 1 ;; + esac +} + +cx_require_record_safe() { + local value + for value in "$@"; do + case "$value" in + *'|'*|*$'\n'*|*$'\r'*) cx_die "manual case field contains a record delimiter" ;; + esac + done +} + +cx_nccl_hca_device_name() { + local selector="${1#=}" + printf '%s' "${selector%%:*}" +} + +cx_export_gid_index_for_link_layer() { + local link_layer="$1" scaleout="$2" + unset NVSHMEM_IB_GID_INDEX NCCL_IB_GID_INDEX UCCL_IB_GID_INDEX + [ -n "${CX_IB_GID_INDEX:-}" ] || return 0 + case "$link_layer" in + roce) + export NVSHMEM_IB_GID_INDEX="$CX_IB_GID_INDEX" + if [ "$scaleout" = 1 ]; then + export NCCL_IB_GID_INDEX="$CX_IB_GID_INDEX" + export UCCL_IB_GID_INDEX="$CX_IB_GID_INDEX" + fi + ;; + infiniband) ;; + *) cx_die "unsupported RDMA link layer" ;; + esac +} + +# Convert private, runner-local network selectors into the public library +# variables needed inside the container. Values are interface/HCA identifiers, +# never addresses; the rendezvous hostname is derived from the allocation. +cx_apply_network_profile() { + local nodes="$1" transport="$2" + local selector rdma_name rdma_names="" ep_nic="" + local scaleout=0 single_node_rdma=0 + local -a selectors + [[ "$nodes" =~ ^[1-9][0-9]*$ ]] || cx_die "invalid network placement" + unset NCCL_NET NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME NCCL_IB_HCA + unset NCCL_IB_GID_INDEX NCCL_IB_SL + unset NVSHMEM_DISABLE_IB NVSHMEM_ENABLE_NIC_PE_MAPPING + unset NVSHMEM_HCA_LIST NVSHMEM_IB_GID_INDEX NVSHMEM_IB_SL + unset NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_IBGDA_NIC_HANDLER + unset NVSHMEM_HCA_PE_MAPPING NVSHMEM_REMOTE_TRANSPORT + unset EP_NIC_NAME EP_OVERRIDE_RDMA_SL + unset UCCL_SOCKET_IFNAME UCCL_IB_GID_INDEX UCCL_IB_SL MORI_RDMA_DEVICES + unset MORI_RDMA_TC MORI_IO_TC MORI_RDMA_SL MORI_IO_SL + if [ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]; then + scaleout=1 + elif [ "${CX_SHARD_SKU:-}" = b300 ] && [ "$nodes" = 1 ] \ + && [ "${CX_BENCH:-}" = deepep ] && [ "${CX_MODE:-}" = low-latency ]; then + # DeepEP V1 low-latency kernels use pure RDMA even for EP8 within one + # NVLink domain. Keep NCCL on NVLink, but bind NVSHMEM to the private HCA. + single_node_rdma=1 + elif [ "${CX_SHARD_SKU:-}" = b300 ] && [ "$nodes" = 1 ]; then + export NVSHMEM_DISABLE_IB=1 + fi + [ "$scaleout" = 1 ] || [ "$single_node_rdma" = 1 ] || return 0 + [ -n "${CX_RDMA_DEVICES:-}" ] \ + || cx_die "RDMA execution requires a private device selector" + if [ "$scaleout" = 1 ] && [ -n "${CX_SOCKET_IFNAME:-}" ]; then + [[ "$CX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(,[A-Za-z][A-Za-z0-9_.-]{0,31})*$ ]] \ + || cx_die "invalid private socket interface selector" + export NCCL_SOCKET_IFNAME="$CX_SOCKET_IFNAME" GLOO_SOCKET_IFNAME="$CX_SOCKET_IFNAME" + export UCCL_SOCKET_IFNAME="$CX_SOCKET_IFNAME" + fi + if [ -n "${CX_RDMA_DEVICES:-}" ]; then + [[ "$CX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] \ + || cx_die "invalid private RDMA device selector" + IFS=, read -r -a selectors <<< "$CX_RDMA_DEVICES" + for selector in "${selectors[@]}"; do + rdma_name="${selector%%:*}" + rdma_names="${rdma_names}${rdma_names:+,}${rdma_name}" + [ -n "$ep_nic" ] || ep_nic="$rdma_name" + done + export NVSHMEM_HCA_LIST="$CX_RDMA_DEVICES" + export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 + if [ "$scaleout" = 1 ]; then + if [ "${CX_SHARD_SKU:-}" = mi300x ] || [ "${CX_SHARD_SKU:-}" = mi325x ] \ + || [ "${CX_SHARD_SKU:-}" = mi355x ]; then + unset NCCL_NET + else + export NCCL_NET=IB + fi + export NCCL_IB_HCA="=$CX_RDMA_DEVICES" + export MORI_RDMA_DEVICES="$rdma_names" EP_NIC_NAME="$ep_nic" + fi + fi + if [ -n "${CX_IB_GID_INDEX:-}" ]; then + [[ "$CX_IB_GID_INDEX" =~ ^[0-9]+$ ]] && [ "$CX_IB_GID_INDEX" -le 255 ] \ + || cx_die "invalid private IB GID index" + fi + if [ -n "${CX_RDMA_SERVICE_LEVEL:-}" ]; then + [[ "$CX_RDMA_SERVICE_LEVEL" =~ ^[0-9]+$ ]] && [ "$CX_RDMA_SERVICE_LEVEL" -le 15 ] \ + || cx_die "invalid private RDMA service level" + export NVSHMEM_IB_SL="$CX_RDMA_SERVICE_LEVEL" + if [ "$scaleout" = 1 ]; then + export NCCL_IB_SL="$CX_RDMA_SERVICE_LEVEL" UCCL_IB_SL="$CX_RDMA_SERVICE_LEVEL" + export EP_OVERRIDE_RDMA_SL="$CX_RDMA_SERVICE_LEVEL" + export MORI_RDMA_SL="$CX_RDMA_SERVICE_LEVEL" MORI_IO_SL="$CX_RDMA_SERVICE_LEVEL" + fi + fi + if [ -n "${CX_RDMA_TRAFFIC_CLASS:-}" ]; then + [[ "$CX_RDMA_TRAFFIC_CLASS" =~ ^[0-9]+$ ]] && [ "$CX_RDMA_TRAFFIC_CLASS" -le 255 ] \ + || cx_die "invalid private RDMA traffic class" + [ "$scaleout" = 1 ] \ + && export MORI_RDMA_TC="$CX_RDMA_TRAFFIC_CLASS" MORI_IO_TC="$CX_RDMA_TRAFFIC_CLASS" + fi + local nic_handler=gpu + if [ "${CX_SHARD_SKU:-}" = b200-dgxc ] && [ "${CX_BENCH:-}" = deepep ]; then + nic_handler=cpu + fi + export NVSHMEM_IB_ENABLE_IBGDA=1 NVSHMEM_IBGDA_NIC_HANDLER="$nic_handler" + if [ -n "${CX_RDMA_LINK_LAYER:-}" ]; then + case "$CX_RDMA_LINK_LAYER" in + roce|infiniband) ;; + *) cx_die "invalid validated RDMA link layer" ;; + esac + cx_export_gid_index_for_link_layer "$CX_RDMA_LINK_LAYER" "$scaleout" + fi +} + +# Slurm may remove NCCL's leading exact-match marker while propagating an +# inherited environment. Reconstruct it from the validated private selector at +# the container boundary instead of accepting a prefix-matched HCA list. +cx_restore_exact_hca_selector() { + if [ "${CX_NODES:-1}" -le 1 ] || [ "${CX_TRANSPORT:-}" = mnnvl ]; then + return 0 + fi + [ -n "${CX_RDMA_DEVICES:-}" ] \ + || { cx_log "ERROR: scale-out RDMA selector is unavailable"; return 1; } + [[ "$CX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] \ + || { cx_log "ERROR: invalid scale-out RDMA selector"; return 1; } + export NCCL_IB_HCA="=$CX_RDMA_DEVICES" +} + +cx_default_route_interface() { + python3 - <<'PY' +from pathlib import Path +for line in Path('/proc/net/route').read_text().splitlines()[1:]: + fields = line.split() + if len(fields) >= 4 and fields[1] == '00000000' and int(fields[3], 16) & 1: + print(fields[0], end=''); break +PY +} + +# Prove that the operator-pinned scale-out fabric exists on every allocated +# node before image import or backend initialization. Selector values and node +# diagnostics stay in the runner-private log. +cx_validate_network_profile_on_job() { + local job_id="$1" nodes="$2" transport="$3" report_failure="${4:-1}" + local log_label=network-profile log rc=0 scaleout=0 marker_count link_layer + local single_node_rdma=0 + if [ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]; then + scaleout=1 + elif [ "${CX_SHARD_SKU:-}" = b300 ] && [ "$nodes" = 1 ] \ + && [ "${CX_BENCH:-}" = deepep ] && [ "${CX_MODE:-}" = low-latency ]; then + single_node_rdma=1 + fi + [ "$scaleout" = 1 ] || [ "$single_node_rdma" = 1 ] || return 0 + [[ "$job_id" =~ ^[1-9][0-9]*$ && "$nodes" =~ ^[1-9][0-9]*$ ]] \ + || return 1 + [ -n "${CX_RDMA_DEVICES:-}" ] || return 1 + case "${CX_NETWORK_VALIDATION_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${CX_NETWORK_VALIDATION_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(cx_private_log_path "$log_label")" || return 1 + CX_NETWORK_PROFILE_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp --input=all \ + --export="$(cx_host_exports),CX_SOCKET_IFNAME,CX_RDMA_DEVICES,CX_IB_GID_INDEX" \ + bash -s > "$log" 2>&1 <<'BASH' || rc=$? +set -euo pipefail +if [ -z "${CX_SOCKET_IFNAME:-}" ]; then + CX_SOCKET_IFNAME="$(python3 - <<'PY' +from pathlib import Path + +for line in Path('/proc/net/route').read_text().splitlines()[1:]: + fields = line.split() + if len(fields) >= 4 and fields[1] == '00000000' and int(fields[3], 16) & 1: + print(fields[0], end='') + break +PY +)" + [ -n "$CX_SOCKET_IFNAME" ] \ + || { printf '[collectivex-private] socket-interface-1=default-route-missing\n'; exit 1; } +fi +[[ "$CX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(,[A-Za-z][A-Za-z0-9_.-]{0,31})*$ ]] +printf '[collectivex-private] socket-interface-selected=%s\n' "$CX_SOCKET_IFNAME" +[[ "$CX_RDMA_DEVICES" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?(,[A-Za-z][A-Za-z0-9_.-]{0,31}(:[1-9][0-9]*)?)*$ ]] +if [ -n "${CX_IB_GID_INDEX:-}" ]; then + [[ "$CX_IB_GID_INDEX" =~ ^[0-9]+$ ]] && [ "$CX_IB_GID_INDEX" -le 255 ] +fi +if [ -n "${CX_SOCKET_IFNAME:-}" ]; then + IFS=, read -r -a interfaces <<< "$CX_SOCKET_IFNAME" + socket_ordinal=0 + for interface in "${interfaces[@]}"; do + socket_ordinal=$((socket_ordinal + 1)) + [ -d "/sys/class/net/$interface" ] \ + || { printf '[collectivex-private] socket-interface-%s=missing\n' "$socket_ordinal"; exit 1; } + state="$(cat "/sys/class/net/$interface/operstate")" + [ "$state" = up ] || [ "$state" = unknown ] \ + || { printf '[collectivex-private] socket-interface-%s=down\n' "$socket_ordinal"; exit 1; } + done +fi +check_port() { + local port_path="$1" ordinal="$2" state gid link_layer + [ -d "$port_path" ] \ + || { printf '[collectivex-private] rdma-port-%s=missing\n' "$ordinal"; return 1; } + read -r state _ < "$port_path/state" + [ "$state" = 4: ] \ + || { printf '[collectivex-private] rdma-port-%s=inactive\n' "$ordinal"; return 1; } + if [ -n "${CX_IB_GID_INDEX:-}" ]; then + [ -r "$port_path/gids/$CX_IB_GID_INDEX" ] \ + || { printf '[collectivex-private] rdma-port-%s=gid-missing\n' "$ordinal"; return 1; } + gid="$(tr -d ':0[:space:]' < "$port_path/gids/$CX_IB_GID_INDEX")" + [ -n "$gid" ] \ + || { printf '[collectivex-private] rdma-port-%s=gid-empty\n' "$ordinal"; return 1; } + fi + [ -r "$port_path/link_layer" ] \ + || { printf '[collectivex-private] rdma-port-%s=link-layer-missing\n' "$ordinal"; return 1; } + link_layer="$(< "$port_path/link_layer")" + case "$link_layer" in + Ethernet) link_layer=roce ;; + InfiniBand) link_layer=infiniband ;; + *) printf '[collectivex-private] rdma-port-%s=link-layer-invalid\n' "$ordinal"; return 1 ;; + esac + [ -z "${profile:-}" ] || [ "$profile" = "$link_layer" ] \ + || { printf '[collectivex-private] rdma-port-%s=link-layer-mixed\n' "$ordinal"; return 1; } + profile="$link_layer" +} +profile="" +IFS=, read -r -a devices <<< "$CX_RDMA_DEVICES" +ordinal=0 +for selector in "${devices[@]}"; do + ordinal=$((ordinal + 1)) + device="${selector%%:*}" + configured_port="" + [ "$selector" = "$device" ] || configured_port="${selector#*:}" + ports="/sys/class/infiniband/$device/ports" + [ -d "$ports" ] \ + || { printf '[collectivex-private] rdma-device-%s=missing\n' "$ordinal"; exit 1; } + if [ -n "$configured_port" ]; then + check_port "$ports/$configured_port" "$ordinal" + else + active=0 + for port_path in "$ports"/*; do + if check_port "$port_path" "$ordinal"; then + active=1 + fi + done + [ "$active" = 1 ] + fi +done +[ -n "$profile" ] +printf '[collectivex-private] rdma-link-layer=%s\n' "$profile" +BASH + if [ "$rc" != 0 ]; then + marker="$(grep -aoE '(socket-interface|rdma-(device|port))-[0-9]+=(missing|down|inactive|default-route-missing|gid-missing|gid-empty|link-layer-missing|link-layer-invalid|link-layer-mixed)' "$log" \ + | tail -n 1 || true)" + [ -z "$marker" ] || cx_log "ERROR: network-profile-$marker" + [ "$report_failure" = 0 ] || cx_fail_stage setup "$log" || true + return "$rc" + fi + socket_ifname="$( + sed -nE 's/^\[collectivex-private\] socket-interface-selected=([A-Za-z][A-Za-z0-9_.-]{0,31})$/\1/p' "$log" \ + | sort -u + )" + marker_count="$(grep -Ec '^\[collectivex-private\] socket-interface-selected=' "$log")" + socket_unique_count="$(printf '%s\n' "$socket_ifname" | sed '/^$/d' | wc -l | tr -d ' ')" + if [ "$socket_unique_count" -lt 1 ] || [ "$marker_count" != "$nodes" ]; then + cx_log "ERROR: network-profile-socket-markers=$marker_count/$nodes unique=$socket_unique_count" + return 1 + fi + if [ "$socket_unique_count" = 1 ]; then + export CX_SOCKET_IFNAME="$socket_ifname" + else + unset CX_SOCKET_IFNAME + fi + link_layer="$( + sed -nE 's/^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$/\1/p' "$log" \ + | sort -u + )" + marker_count="$(grep -Ec '^\[collectivex-private\] rdma-link-layer=(roce|infiniband)$' "$log")" + case "$marker_count:$link_layer" in + "$nodes":roce|"$nodes":infiniband) ;; + *) [ "$report_failure" = 0 ] || cx_fail_stage setup "$log" || true; return 1 ;; + esac + export CX_RDMA_LINK_LAYER="$link_layer" + cx_export_gid_index_for_link_layer "$link_layer" "$scaleout" +} + +cx_allocation_nodes_csv() { + local job_id="$1" nodelist node output="" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + nodelist="$(squeue -h -j "$job_id" -o %N 2>/dev/null)" || return 1 + [[ "$nodelist" =~ ^[][A-Za-z0-9._,-]+$ ]] || return 1 + while IFS= read -r node; do + [ -n "$node" ] || continue + [[ "$node" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] || return 1 + [ -z "$output" ] || output+=, + output+="$node" + done < <(scontrol show hostnames "$nodelist" 2>/dev/null) + [ -n "$output" ] || return 1 + printf '%s' "$output" +} + +cx_resolve_slurm_rendezvous() { + local job_id="$1" master_addr master_port socket_ifname="${CX_SOCKET_IFNAME:-}" + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || cx_die "invalid rendezvous allocation" + # Query relative node zero directly so MASTER_ADDR always hosts global rank 0. + # Prefer the address on the already validated cross-node socket interface; + # a short hostname may resolve onto a management network that ranks cannot use. + if [[ "$socket_ifname" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]]; then + master_addr="$(srun --jobid="$job_id" --nodes=1 --ntasks=1 --relative=0 \ + --chdir=/tmp --export="$(cx_host_exports)" bash -s -- "$socket_ifname" \ + 2>/dev/null <<'BASH' | head -n1 +set -euo pipefail +ip -o -4 address show dev "$1" scope global \ + | awk 'NR == 1 {split($4, address, "/"); print address[1]}' +BASH +)" + [[ "$master_addr" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] \ + || cx_die "could not resolve the allocated primary interface" + else + master_addr="$(srun --jobid="$job_id" --nodes=1 --ntasks=1 --relative=0 \ + --chdir=/tmp --export="$(cx_host_exports)" hostname -s 2>/dev/null | head -n1)" + [[ "$master_addr" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || cx_die "could not resolve the allocated primary node" + fi + master_port="${CX_MASTER_PORT:-29551}" + [[ "$master_port" =~ ^[1-9][0-9]*$ ]] && [ "$master_port" -le 65535 ] \ + || cx_die "invalid distributed rendezvous port" + export MASTER_ADDR="$master_addr" MASTER_PORT="$master_port" +} + +# Printed into `bash -c` for one Slurm task per GPU. Every rank derives its +# identity from Slurm rather than accepting caller-supplied rank values. +cx_slurm_rank_wrapper() { + cat <<'BASH' +case "${SLURM_PROCID:-}:${SLURM_NTASKS:-}:${SLURM_LOCALID:-}:${SLURM_NODEID:-}" in + *[!0-9:]*|:*|*::*|*:) exit 67 ;; +esac +[ "$SLURM_NTASKS" = "$CX_NGPUS" ] || exit 67 +[ "$SLURM_LOCALID" -lt "$CX_GPUS_PER_NODE" ] || exit 67 +. /ix/experimental/CollectiveX/runtime/common.sh || exit 68 +if [ "${CX_NODES:-1}" -gt 1 ] && [ "${CX_TRANSPORT:-}" != mnnvl ]; then + if [ -z "${CX_SOCKET_IFNAME:-}" ]; then + CX_SOCKET_IFNAME="$(cx_default_route_interface)" || exit 68 + [[ "$CX_SOCKET_IFNAME" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]] || exit 68 + export CX_SOCKET_IFNAME + fi + cx_apply_network_profile "$CX_NODES" "$CX_TRANSPORT" || exit 68 +fi +cx_write_runtime_stage execution || exit 68 +export RANK="$SLURM_PROCID" WORLD_SIZE="$SLURM_NTASKS" +export LOCAL_RANK="$SLURM_LOCALID" LOCAL_WORLD_SIZE="$CX_GPUS_PER_NODE" +exec python3 bench/run_ep.py "$@" +BASH +} + +# A set shard path is an execution contract, never a hint. Validate it before +# staging/allocation and again in-container so a missing or stale control file +# cannot silently fall back to a manual single-case run. +cx_validate_shard_control() { + local cx_root="$1" shard="${CX_SHARD_FILE:-}" path expected_sku control_sha256 + [ -n "$shard" ] || return 0 + expected_sku="${CX_SHARD_SKU:-}" + [ -n "$expected_sku" ] || cx_die "CX_SHARD_SKU is required with CX_SHARD_FILE" + [ -n "${CX_BENCH:-}" ] || cx_die "CX_BENCH is required with CX_SHARD_FILE" + [[ "${CX_NODES:-}" =~ ^[1-9][0-9]*$ ]] \ + || cx_die "positive CX_NODES is required with CX_SHARD_FILE" + path="$shard" + [ -f "$path" ] || path="${cx_root%/}/$shard" + [ -f "$path" ] || cx_die "shard control does not exist" + [ -s "$path" ] || cx_die "shard control is empty" + python3 "${cx_root%/}/sweep_matrix.py" \ + --validate-control "$path" --expect-sku "$expected_sku" \ + --expect-backend "$CX_BENCH" --expect-nodes "$CX_NODES" >/dev/null 2>&1 \ + || cx_die "invalid shard control" + control_sha256="$(sha256sum "$path" | awk '{print $1}')" + [[ "$control_sha256" =~ ^[0-9a-f]{64}$ ]] \ + || cx_die "cannot hash shard control" + export COLLECTIVEX_CONTROL_SHA256="$control_sha256" +} + +# Load only the case mode needed to choose the allocation/network profile. A +# consolidated shard may contain normal and low-latency cases; selecting the +# strictest mode here makes allocation preflight prove every required device. +# The in-container dispatcher reapplies the profile for each individual case. +cx_load_network_control_mode() { + local cx_root="$1" shard="${CX_SHARD_FILE:-}" path mode + [ -n "$shard" ] || return 0 + path="$shard" + [ -f "$path" ] || path="${cx_root%/}/$shard" + [ -f "$path" ] || return 1 + mode="$(python3 - "$path" <<'PY' +import json +import sys + +with open(sys.argv[1]) as stream: + cases = json.load(stream)["cases"] +modes = {case["mode"] for case in cases} +if not modes or modes - {"normal", "low-latency"}: + raise SystemExit("invalid shard mode") +print("low-latency" if "low-latency" in modes else "normal") +PY +)" || return 1 + case "$mode" in + normal|low-latency) export CX_MODE="$mode" ;; + *) return 1 ;; + esac +} + +cx_apply_timing_profile() { + [ -n "${CX_TIMING:-}" ] || return 0 + local iters trials warmup extra + IFS=: read -r iters trials warmup extra <<< "$CX_TIMING" + [[ "$iters" =~ ^[1-9][0-9]*$ && "$trials" =~ ^[1-9][0-9]*$ \ + && "$warmup" =~ ^[1-9][0-9]*$ && -z "$extra" ]] \ + || cx_die "CX_TIMING must be positive iters:trials:warmup" + export CX_ITERS="$iters" CX_TRIALS="$trials" CX_WARMUP="$warmup" +} + +# Use an opaque, execution-bound name so a missing grant message can be +# reconciled without exposing runner or shard details in public logs. +cx_scheduler_job_name() { + local execution_id="${COLLECTIVEX_EXECUTION_ID:-manual-$$}" digest + digest="$(printf '%s' "$execution_id" | sha256sum | awk '{print $1}')" \ + || return 1 + [[ "$digest" =~ ^[0-9a-f]{64}$ ]] || return 1 + printf 'cx-%s' "${digest:0:24}" +} + +# Return 0 after recovering one allocation ID, 2 after three successful empty +# observations, and 1 for every ambiguous or failed lookup. Callers inspect the +# state variables rather than the status because all missing-ID paths still fail. +cx_reconcile_salloc_jobid() { + local job_name="$1" scheduler_user queue_output line delay attempt + local -a ids=() + scheduler_user="$(id -un 2>/dev/null)" || return 1 + [[ "$scheduler_user" =~ ^[A-Za-z0-9_.-]+$ \ + && "$job_name" =~ ^cx-[0-9a-f]{24}$ ]] || return 1 + for attempt in 1 2 3; do + ids=() + if ! queue_output="$( + squeue -h --user="$scheduler_user" --name="$job_name" -o %A 2>/dev/null + )"; then + return 1 + fi + while IFS= read -r line; do + [[ "$line" =~ ^[[:space:]]*$ ]] && continue + if [[ "$line" =~ ^[[:space:]]*([1-9][0-9]*)[[:space:]]*$ ]]; then + ids+=("${BASH_REMATCH[1]}") + else + return 1 + fi + done <<< "$queue_output" + if [ "${#ids[@]}" -eq 1 ]; then + JOB_ID="${ids[0]}" + CX_ALLOCATION_UNCERTAIN=0 + return 0 + fi + [ "${#ids[@]}" -eq 0 ] || return 1 + if [ "$attempt" -eq 3 ]; then + CX_ALLOCATION_UNCERTAIN=0 + return 2 + fi + delay=$((1 << (attempt - 1))) + sleep "$delay" || return 1 + done + return 1 +} + +cx_verify_salloc_jobid() { + local job_id="$1" queue_output line count=0 + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + queue_output="$(squeue -h -j "$job_id" -o %A 2>/dev/null)" || return 1 + while IFS= read -r line; do + [[ "$line" =~ ^[[:space:]]*$ ]] && continue + [[ "$line" =~ ^[[:space:]]*${job_id}[[:space:]]*$ ]] || return 1 + count=$((count + 1)) + done <<< "$queue_output" + [ "$count" -eq 1 ] +} + +# Allocate via salloc's stable grant message and assign JOB_ID in this shell. +# Raw scheduler output remains in the bounded private execution log. +cx_salloc_jobid() { + local log_label=scheduler-allocation log job_id job_name argument salloc_rc=0 + case "${CX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${CX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + if ! log="$(cx_private_log_path "$log_label")"; then + cx_log "ERROR: failure-class=scheduler-allocation diagnostic=private-log" + return 1 + fi + for argument in "$@"; do + case "$argument" in + --job-name|--job-name=*|-J|-J*) + cx_log "ERROR: scheduler job names are managed by CollectiveX" + return 1 + ;; + esac + done + if ! job_name="$(cx_scheduler_job_name)"; then + cx_log "ERROR: failure-class=scheduler-allocation diagnostic=job-name" + return 1 + fi + CX_ALLOCATION_UNCERTAIN=1 + # salloc has no portable --parsable option. Parse the stable grant message + # used by the production launchers, while also accepting a bare ID from + # site wrappers. Contain shell-function wrappers that call exit so the + # launcher can still reconcile and cancel an allocation. + cx_log "scheduler-request=submit" + (salloc "$@" --job-name="$job_name" --no-shell) > "$log" 2>&1 || salloc_rc=$? + if ! job_id="$(sed -nE \ + -e 's/^([0-9]+)(;[^[:space:]]+)?$/\1/p; t found' \ + -e 's/.*Granted job allocation ([0-9]+).*/\1/p; t found' \ + -e 'b' -e ':found' -e 'q' "$log")"; then + cx_log "ERROR: failure-class=scheduler-allocation diagnostic=grant-parse" + cx_reconcile_salloc_jobid "$job_name" || true + [ -z "$JOB_ID" ] || cx_record_allocation_jobid "$JOB_ID" || true + return 1 + fi + if [ -n "$job_id" ]; then + [[ "$job_id" =~ ^[0-9]+$ ]] || return 1 + JOB_ID="$job_id" + CX_ALLOCATION_UNCERTAIN=0 + fi + if [ "$salloc_rc" != 0 ]; then + if [ -n "$JOB_ID" ] && cx_verify_salloc_jobid "$JOB_ID"; then + cx_log "scheduler-request=verified-grant" + cx_record_allocation_jobid "$JOB_ID" || return 1 + return 0 + fi + cx_log "ERROR: scheduler-request=rejected" + if [ "$salloc_rc" -ge 128 ] && [ -z "$JOB_ID" ]; then + cx_fail_stage scheduler-allocation "$log" + return 1 + fi + [ -n "$JOB_ID" ] || cx_reconcile_salloc_jobid "$job_name" || true + [ -z "$JOB_ID" ] || cx_record_allocation_jobid "$JOB_ID" || true + cx_fail_stage scheduler-allocation "$log" + return 1 + fi + if [ -z "$JOB_ID" ]; then + cx_log "ERROR: scheduler-request=missing-grant" + cx_reconcile_salloc_jobid "$job_name" || true + cx_fail_stage scheduler-allocation "$log" + return 1 + fi + cx_record_allocation_jobid "$JOB_ID" || return 1 +} + +cx_record_allocation_jobid() { + local job_id="$1" root="${CX_JOB_ROOT:-}" path temporary + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || return 1 + [ -n "$root" ] || return 0 + cx_job_root_is_safe "$root" || return 1 + path="$root/allocation-job-id" + temporary="$(mktemp "$root/.allocation-job-id.XXXXXX")" || return 1 + chmod 600 "$temporary" || { rm -f -- "$temporary"; return 1; } + printf '%s\n' "$job_id" > "$temporary" \ + || { rm -f -- "$temporary"; return 1; } + mv -f -- "$temporary" "$path" || { rm -f -- "$temporary"; return 1; } +} + +cx_clear_allocation_jobid() { + local root="${CX_JOB_ROOT:-}" path + [ -n "$root" ] || return 0 + cx_job_root_is_safe "$root" || return 1 + path="$root/allocation-job-id" + [ ! -e "$path" ] || { + [ -f "$path" ] && [ ! -L "$path" ] \ + && [ "$(stat -c '%u:%a' "$path" 2>/dev/null)" = "$(id -u):600" ] || return 1 + rm -f -- "$path" + } +} + +cx_cancel_job() { + local job_id="$1" active delay + [[ "$job_id" =~ ^[0-9]+$ ]] || return 1 + scancel "$job_id" >/dev/null 2>&1 || true + for delay in 1 2 4 8 16 32 64; do + if ! active="$(squeue -h -j "$job_id" -o %A 2>/dev/null)"; then + sleep "$delay" + continue + fi + [ -n "$active" ] || return 0 + sleep "$delay" + done + cx_log "ERROR: scheduled allocation did not terminate during cleanup" + return 1 +} + +cx_write_cleanup_guard() { + local state="$1" root="${CX_JOB_ROOT:-}" safe unsafe + cx_job_root_is_safe "$root" || return 0 + safe="$root/cleanup-safe" + unsafe="$root/cleanup-unsafe" + umask 077 + case "$state" in + safe) : > "$safe" && rm -f -- "$unsafe" ;; + unsafe) rm -f -- "$safe" && : > "$unsafe" ;; + *) return 1 ;; + esac +} + +# A workflow cancellation may kill a foreground Slurm step before Bash can run +# the launcher trap. Reconcile the mode-0600 allocation record from an always() +# workflow step before isolated source cleanup is allowed. +cx_reconcile_recorded_allocation() { + local root="$1" path job_id + cx_job_root_is_safe "$root" || return 1 + export CX_JOB_ROOT="$root" + path="$root/allocation-job-id" + if [ ! -e "$path" ]; then + [ -f "$root/cleanup-safe" ] && [ ! -e "$root/cleanup-unsafe" ] + return + fi + [ -f "$path" ] && [ ! -L "$path" ] \ + && [ "$(stat -c '%u:%a' "$path" 2>/dev/null)" = "$(id -u):600" ] \ + || { cx_write_cleanup_guard unsafe || true; return 1; } + IFS= read -r job_id < "$path" || { + cx_write_cleanup_guard unsafe || true + return 1 + } + [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || { + cx_write_cleanup_guard unsafe || true + return 1 + } + if cx_cancel_job "$job_id" && cx_clear_allocation_jobid; then + cx_write_cleanup_guard safe + return + fi + cx_write_cleanup_guard unsafe || true + return 1 +} + +# Single multi-arch container for ALL NVIDIA SKUs: tag `v0.5.11-cu130` is an OCI +# image index covering linux/amd64 (B200) + linux/arm64 (GB200); enroot import +# pulls the matching arch. (cu130 = CUDA 13, system nccl.h in /usr/include, torch 2.9.x.) +# Import remains tag-based because Enroot cannot reliably import a digest-qualified +# Docker Hub reference non-interactively. The registry digest is resolved and checked +# immediately before import, then recorded as verified provenance. +CX_IMAGE_MULTIARCH_DIGEST="sha256:061fb71f838e82000a1768c159654d526c2f17ebe751c21e7fc48ca53c8ef975" +# (v0.5.12-cu130 was rejected: its 62 layers overflow enroot's overlay-based +# squash creation on these nodes — "failed to mount overlay ... Invalid argument". +# v0.5.11-cu130 imports cleanly.) +# Runtime setup verifies the image-bundled DeepEP build for the detected GPU target. +CX_IMAGE_MULTIARCH="lmsysorg/sglang:v0.5.11-cu130" + +# AMD (ROCm/CDNA): separate single-arch images bundle MoRI. +CX_IMAGE_AMD_MORI="rocm/sgl-dev:sglang-0.5.9-rocm720-mi35x-mori-0227-2" +CX_IMAGE_AMD_MORI_DIGEST="sha256:24c3b30d64475937abbb6498e3b29528649adcb836dde7a468979f767809b0e8" +CX_MORI_COMMIT_MI355="99bc0a3a6e7a70aacc6372cd9a4275ccfb4de567" # pragma: allowlist secret +CX_IMAGE_AMD_MORI_MI325="rocm/sgl-dev:sglang-0.5.14-rocm720-mi35x-mori-0701" +CX_IMAGE_AMD_MORI_MI325_DIGEST="sha256:ea42375343c2ef8f73b3bdb9e1b7b435556e3ca92aba5e3f74ada29ba217fabc" +CX_MORI_COMMIT_MI325="bf99bdf18fc69887a346913ca01c315c2aa9bd4c" # pragma: allowlist secret +cx_default_image() { + case "$1" in + mi300x*|mi325x*) echo "$CX_IMAGE_AMD_MORI_MI325" ;; + mi355x*) echo "$CX_IMAGE_AMD_MORI" ;; + b200*|gb200*|b300*|gb300*|h100*|h200*) echo "$CX_IMAGE_MULTIARCH" ;; + *) cx_die "no default image for runner prefix: $1" ;; + esac +} + +cx_resolve_registry_digest() { + local image="$1" repository reference token digest registry + if [[ "$image" == *@* ]]; then + cx_die "digest-qualified image overrides are unsupported; configure a tag and pinned digest" + fi + registry="${image%%/*}" + if [[ "$image" == */* && ( "$registry" == *.* || "$registry" == *:* || "$registry" = localhost ) ]]; then + case "$registry" in + docker.io|registry-1.docker.io) image="${image#*/}" ;; + *) cx_die "only Docker Hub images are supported by the registry verifier" ;; + esac + fi + repository="${image%:*}" + reference="${image##*:}" + [ "$repository" != "$image" ] || { repository="$image"; reference=latest; } + [ -n "$repository" ] && [ -n "$reference" ] \ + || cx_die "configured image reference is malformed" + [[ "$repository" == */* ]] || repository="library/$repository" + token="$(curl -fsSLG --connect-timeout 10 --max-time 30 --retry 2 \ + --retry-delay 1 --retry-all-errors 'https://auth.docker.io/token' \ + --data-urlencode 'service=registry.docker.io' \ + --data-urlencode "scope=repository:${repository}:pull" \ + | python3 -c 'import json,sys; print(json.load(sys.stdin)["token"])')" \ + || cx_die "cannot authenticate to the image registry" + digest="$(curl -fsSI --connect-timeout 10 --max-time 30 --retry 2 \ + --retry-delay 1 --retry-all-errors \ + -H "Authorization: Bearer $token" \ + -H 'Accept: application/vnd.oci.image.index.v1+json, application/vnd.oci.image.manifest.v1+json, application/vnd.docker.distribution.manifest.list.v2+json, application/vnd.docker.distribution.manifest.v2+json' \ + "https://registry-1.docker.io/v2/${repository}/manifests/${reference}" \ + | tr -d '\r' | awk 'tolower($1)=="docker-content-digest:" {print $2; exit}')" \ + || cx_die "cannot resolve the configured image digest" + [[ "$digest" =~ ^sha256:[0-9a-f]{64}$ ]] \ + || cx_die "registry returned an invalid image digest" + printf '%s' "$digest" +} + +cx_verify_registry_image() { + local image="$1" expected actual + expected="${CX_IMAGE_DIGEST:-$(cx_default_image_digest "$image")}" + [[ "$expected" =~ ^sha256:[0-9a-f]{64}$ ]] \ + || cx_die "a pinned digest is required for the configured image" + actual="$(cx_resolve_registry_digest "$image")" + [ "$actual" = "$expected" ] \ + || cx_die "configured image tag no longer matches its pinned digest" + export COLLECTIVEX_IMAGE="$image" COLLECTIVEX_IMAGE_DIGEST="$actual" + export COLLECTIVEX_IMAGE_DIGEST_VERIFIED=1 +} + +cx_default_image_digest() { + case "$1" in + "$CX_IMAGE_MULTIARCH") printf '%s' "$CX_IMAGE_MULTIARCH_DIGEST" ;; + "$CX_IMAGE_AMD_MORI") printf '%s' "$CX_IMAGE_AMD_MORI_DIGEST" ;; + "$CX_IMAGE_AMD_MORI_MI325") printf '%s' "$CX_IMAGE_AMD_MORI_MI325_DIGEST" ;; + esac +} + +# Create a per-UID cache under validated cluster-local storage. Only the fixed +# /cx-cache mount enters the container; the operator host path does not. +cx_prepare_backend_cache() { + local stage_parent="$1" cache info sentinel_sha256 + unset CX_PREPARED_BACKEND_CACHE CX_BACKEND_CACHE_SENTINEL_SHA256 + info="$(python3 - "$stage_parent" <<'PY' +import hashlib +import os +import secrets +import stat +import sys + +configured_parent = sys.argv[1] +try: + if ( + not os.path.isabs(configured_parent) + or "\n" in configured_parent + or "\r" in configured_parent + ): + raise OSError + parent = os.path.realpath(configured_parent) + if not os.path.isdir(parent): + raise OSError + flags = os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0) + parent_fd = os.open(parent, flags) + try: + probe_name = f".collectivex-owner-probe-{os.getpid()}-{secrets.token_hex(8)}" + os.mkdir(probe_name, 0o700, dir_fd=parent_fd) + try: + probe_fd = os.open(probe_name, flags, dir_fd=parent_fd) + try: + probe = os.fstat(probe_fd) + if stat.S_IMODE(probe.st_mode) & 0o777 != 0o700: + raise OSError + realized_owner = probe.st_uid + finally: + os.close(probe_fd) + finally: + os.rmdir(probe_name, dir_fd=parent_fd) + for generation in (3, 4): + name = f".collectivex-backend-cache-v{generation}-{os.getuid()}" + try: + os.mkdir(name, 0o700, dir_fd=parent_fd) + except FileExistsError: + pass + try: + cache_fd = os.open(name, flags, dir_fd=parent_fd) + try: + metadata = os.fstat(cache_fd) + if ( + metadata.st_uid != realized_owner + or stat.S_IMODE(metadata.st_mode) & 0o777 != 0o700 + ): + raise OSError + sentinel_name = ".collectivex-mount-sentinel-v1" + temporary_name = ( + f"{sentinel_name}.tmp.{os.getpid()}.{secrets.token_hex(8)}" + ) + create_flags = ( + os.O_WRONLY | os.O_CREAT | os.O_EXCL + | getattr(os, "O_NOFOLLOW", 0) + ) + payload = secrets.token_bytes(32) + temporary_fd = os.open( + temporary_name, create_flags, 0o600, dir_fd=cache_fd + ) + try: + try: + view = memoryview(payload) + try: + while view: + written = os.write(temporary_fd, view) + if written <= 0: + raise OSError + view = view[written:] + os.fsync(temporary_fd) + finally: + view.release() + finally: + os.close(temporary_fd) + try: + os.link( + temporary_name, + sentinel_name, + src_dir_fd=cache_fd, + dst_dir_fd=cache_fd, + follow_symlinks=False, + ) + except FileExistsError: + pass + finally: + try: + os.unlink(temporary_name, dir_fd=cache_fd) + except FileNotFoundError: + pass + sentinel_fd = os.open( + sentinel_name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + dir_fd=cache_fd, + ) + try: + sentinel = os.fstat(sentinel_fd) + payload = os.read(sentinel_fd, 33) + if ( + not stat.S_ISREG(sentinel.st_mode) + or sentinel.st_uid != realized_owner + or stat.S_IMODE(sentinel.st_mode) & 0o777 != 0o600 + or sentinel.st_size != 32 + or len(payload) != 32 + ): + raise OSError + sentinel_sha256 = hashlib.sha256(payload).hexdigest() + finally: + os.close(sentinel_fd) + finally: + os.close(cache_fd) + except OSError: + if generation == 3: + continue + raise + break + finally: + os.close(parent_fd) +except OSError: + raise SystemExit(1) +print(sentinel_sha256, os.path.join(parent, name), end="") +PY +)" || return 1 + sentinel_sha256="${info%% *}" + cache="${info#* }" + [ "$cache" != "$info" ] && [[ "$sentinel_sha256" =~ ^[0-9a-f]{64}$ ]] \ + && [[ "$cache" = /* ]] || return 1 + export CX_PREPARED_BACKEND_CACHE="$cache" + export CX_BACKEND_CACHE_SENTINEL_SHA256="$sentinel_sha256" +} + +cx_verify_backend_cache_mount() { + python3 - "${CX_BACKEND_CACHE_ROOT:-}" \ + "${CX_BACKEND_CACHE_SENTINEL_SHA256:-}" <<'PY' +import hashlib +import os +import re +import stat +import sys + +root, expected = sys.argv[1:] +try: + if ( + not os.path.isabs(root) + or os.path.realpath(root) != root + or re.fullmatch(r"[0-9a-f]{64}", expected) is None + ): + raise OSError + flags = os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0) + root_fd = os.open(root, flags) + try: + root_item = os.fstat(root_fd) + if ( + not stat.S_ISDIR(root_item.st_mode) + or stat.S_IMODE(root_item.st_mode) & 0o777 != 0o700 + ): + raise OSError + sentinel_fd = os.open( + ".collectivex-mount-sentinel-v1", + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + dir_fd=root_fd, + ) + try: + sentinel = os.fstat(sentinel_fd) + payload = os.read(sentinel_fd, 33) + if ( + not stat.S_ISREG(sentinel.st_mode) + or sentinel.st_uid != root_item.st_uid + or stat.S_IMODE(sentinel.st_mode) & 0o777 != 0o600 + or sentinel.st_size != 32 + or len(payload) != 32 + or hashlib.sha256(payload).hexdigest() != expected + ): + raise OSError + finally: + os.close(sentinel_fd) + finally: + os.close(root_fd) +except OSError: + raise SystemExit(1) +PY +} + +cx_git() { + GIT_CONFIG_NOSYSTEM=1 GIT_CONFIG_GLOBAL=/dev/null GIT_TERMINAL_PROMPT=0 \ + git -c credential.helper= "$@" +} + +cx_git_in_tree() { + local directory="$1" canonical + shift + [[ "$directory" = /* ]] && [ -d "$directory" ] && [ ! -L "$directory" ] \ + || return 1 + [[ "$directory" != *'*'* && "$directory" != *$'\n'* && "$directory" != *$'\r'* ]] \ + || return 1 + canonical="$(cd -P -- "$directory" && pwd -P)" || return 1 + cx_git -c "safe.directory=$canonical" -C "$canonical" "$@" +} + +cx_fetch_revision() { + local repository="$1" revision="$2" destination="$3" attempt + for attempt in 1 2 3; do + rm -rf -- "$destination" + if cx_git init -q "$destination" \ + && cx_git_in_tree "$destination" remote add origin "$repository" \ + && cx_git_in_tree "$destination" fetch -q --no-tags --depth 1 origin "$revision" \ + && cx_git_in_tree "$destination" -c advice.detachedHead=false \ + checkout -q --detach FETCH_HEAD \ + && [ "$(cx_git_in_tree "$destination" rev-parse HEAD)" = "$revision" ]; then + return 0 + fi + [ "$attempt" = 3 ] || sleep $((attempt * 5)) + done + return 1 +} + +cx_backend_source_pin() { + case "$1" in + deepep-v2) + printf '%s|%s|%s' \ + "$CX_DEEPEP_V2_COMMIT" "$CX_DEEPEP_V2_TREE" "$CX_DEEPEP_V2_FMT_COMMIT" + ;; + deepep-hybrid) + printf '%s|%s||%s' "$CX_DEEPEP_HYBRID_COMMIT" "$CX_DEEPEP_HYBRID_TREE" \ + "$CX_DEEPEP_HYBRID_NCCL_COMMIT" + ;; + *) return 1 ;; + esac +} + +cx_backend_source_path() { + local root="$1" backend="$2" revision tree fmt nccl pin + pin="$(cx_backend_source_pin "$backend")" || return 1 + IFS='|' read -r revision tree fmt nccl <<< "$pin" + printf '%s/%s-%s' "$root" "$backend" "$revision" +} + +cx_backend_source_is_valid() { + local backend="$1" source="$2" revision tree fmt nccl pin status ignored + pin="$(cx_backend_source_pin "$backend")" || return 1 + IFS='|' read -r revision tree fmt nccl <<< "$pin" + [ -d "$source" ] && [ ! -L "$source" ] \ + && [ "$(cx_git_in_tree "$source" rev-parse HEAD 2>/dev/null)" = "$revision" ] \ + && [ "$(cx_git_in_tree "$source" rev-parse 'HEAD^{tree}' 2>/dev/null)" = "$tree" ] \ + || return 1 + status="$(GIT_OPTIONAL_LOCKS=0 cx_git_in_tree "$source" \ + status --porcelain --untracked-files=all --ignore-submodules=none \ + 2>/dev/null)" || return 1 + if [ "$backend" = deepep-v2 ]; then + [ "$status" = " M deep_ep/__init__.py" ] \ + && [ "$(sha256sum "$source/deep_ep/__init__.py" | awk '{print $1}')" = \ + "$CX_DEEPEP_V2_INIT_SHA256" ] \ + || return 1 + else + [ -z "$status" ] || return 1 + fi + ignored="$(cx_git_in_tree "$source" ls-files --others --ignored --exclude-standard \ + 2>/dev/null)" || return 1 + [ -z "$ignored" ] || return 1 + [ -z "$fmt" ] \ + || [ "$(cx_git_in_tree "$source/third-party/fmt" rev-parse HEAD 2>/dev/null)" = "$fmt" ] \ + || return 1 + [ -z "$nccl" ] \ + || [ "$(cx_git_in_tree "$source/third-party/nccl" rev-parse HEAD 2>/dev/null)" = "$nccl" ] +} + +cx_apply_deepep_v2_nccl_check_fix() { + local source="$1" + python3 - "$source/deep_ep/__init__.py" "$CX_DEEPEP_V2_INIT_SHA256" <<'PY' +import hashlib +from pathlib import Path +import sys + +path = Path(sys.argv[1]) +expected = sys.argv[2] +old = "for so in [line.strip().split(' ')[-1] for line in f if 'nccl' in line]:" +new = "for so in [line.strip().split(' ')[-1] for line in f if 'libnccl' in line]:" +payload = path.read_text(encoding="utf-8") +if payload.count(old) != 1 or new in payload: + raise SystemExit(1) +path.write_text(payload.replace(old, new), encoding="utf-8") +if hashlib.sha256(path.read_bytes()).hexdigest() != expected: + raise SystemExit(1) +PY +} + +cx_extension_pair_sha256() { + python3 - "$1" "$2" "$3" <<'PY' +import hashlib +import os +from pathlib import Path +import stat +import sys + +root = Path(sys.argv[1]) +digest = hashlib.sha256() +try: + if root.is_symlink() or not root.is_dir(): + raise OSError + for pattern in sys.argv[2:]: + matches = list(root.glob(pattern)) + if len(matches) != 1 or matches[0].is_symlink(): + raise OSError + path = matches[0] + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + try: + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode): + raise OSError + file_digest = hashlib.sha256() + with os.fdopen(descriptor, "rb", closefd=False) as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + file_digest.update(chunk) + digest.update(path.name.encode("utf-8") + b"\0") + digest.update(str(metadata.st_size).encode("ascii") + b"\0") + digest.update(file_digest.digest()) + finally: + os.close(descriptor) +except (OSError, UnicodeError): + raise SystemExit(1) +print(digest.hexdigest(), end="") +PY +} + +# Acquire source before compute allocation, preferring the verified same-run GHA seed. +_cx_prepare_backend_source() { + local mount_src="$1" backend="$2" root source temporary revision tree fmt nccl pin + local root_mode stage_mode root_owner stage_owner + local seed_root="${CX_BACKEND_SOURCE_SEED_ROOT:-}" seed seed_mode + root="$mount_src/experimental/CollectiveX/.cx_sources" + CX_BACKEND_SOURCE_STEP="source mount creation" + if [ ! -e "$root" ] && [ ! -L "$root" ]; then + mkdir -m 700 -- "$root" || return 1 + fi + CX_BACKEND_SOURCE_STEP="source mount ownership validation" + [ -d "$mount_src" ] && [ ! -L "$mount_src" ] \ + && [ -d "$root" ] && [ ! -L "$root" ] || return 1 + stage_owner="$(stat -c '%u' "$mount_src" 2>/dev/null)" || return 1 + root_owner="$(stat -c '%u' "$root" 2>/dev/null)" || return 1 + [ "$root_owner" = "$stage_owner" ] || return 1 + stage_mode="$(stat -c '%a' "$mount_src" 2>/dev/null)" || return 1 + case "$stage_mode" in 700|[1-7]700) ;; *) return 1 ;; esac + # Shared stage parents may retain harmless special bits despite mkdir -m. + CX_BACKEND_SOURCE_STEP="source mount permission inspection" + root_mode="$(stat -c '%a' "$root" 2>/dev/null)" || return 1 + case "$root_mode" in + 700|[1-7]700) ;; + *) + CX_BACKEND_SOURCE_STEP="source mount permission normalization" + chmod 700 "$root" || return 1 + CX_BACKEND_SOURCE_STEP="source mount permission validation" + root_mode="$(stat -c '%a' "$root" 2>/dev/null)" || return 1 + case "$root_mode" in 700|[1-7]700) ;; *) return 1 ;; esac + ;; + esac + CX_BACKEND_SOURCE_STEP="git lookup" + command -v git >/dev/null || return 1 + CX_BACKEND_SOURCE_STEP="source pin resolution" + source="$(cx_backend_source_path "$root" "$backend")" || return 1 + if [ -e "$source" ] || [ -L "$source" ]; then + CX_BACKEND_SOURCE_STEP="existing source validation" + cx_backend_source_is_valid "$backend" "$source" + return + fi + if [ -n "$seed_root" ]; then + CX_BACKEND_SOURCE_STEP="source seed validation" + [[ "$seed_root" = /* ]] && [ -d "$seed_root" ] && [ ! -L "$seed_root" ] \ + || return 1 + seed_mode="$(stat -c '%a' "$seed_root" 2>/dev/null)" || return 1 + case "$seed_mode" in 700|[1-7]700) ;; *) return 1 ;; esac + seed="$(cx_backend_source_path "$seed_root" "$backend")" || return 1 + cx_backend_source_is_valid "$backend" "$seed" || return 1 + CX_BACKEND_SOURCE_STEP="source seed copy" + temporary="$(mktemp -d "$root/.${backend}.XXXXXX")" || return 1 + if ! cp -R -- "$seed/." "$temporary/" \ + || ! cx_backend_source_is_valid "$backend" "$temporary" \ + || ! mv -- "$temporary" "$source"; then + rm -rf -- "$temporary" + return 1 + fi + return + fi + if [ "${COLLECTIVEX_CANONICAL_GHA:-0}" = 1 ]; then + CX_BACKEND_SOURCE_STEP="source seed validation" + return 1 + fi + CX_BACKEND_SOURCE_STEP="source checkout creation" + temporary="$(mktemp -d "$root/.${backend}.XXXXXX")" || return 1 + CX_BACKEND_SOURCE_STEP="source pin resolution" + pin="$(cx_backend_source_pin "$backend")" || { + rm -rf -- "$temporary" + return 1 + } + IFS='|' read -r revision tree fmt nccl <<< "$pin" + CX_BACKEND_SOURCE_STEP="revision fetch" + if ! cx_fetch_revision \ + https://github.com/deepseek-ai/DeepEP "$revision" "$temporary"; then + rm -rf -- "$temporary" + return 1 + fi + CX_BACKEND_SOURCE_STEP="submodule fetch" + if [ -n "$fmt" ] && ! cx_git_in_tree "$temporary" \ + -c "safe.directory=$temporary/third-party/fmt" \ + submodule update -q --init --depth 1 third-party/fmt; then + rm -rf -- "$temporary" + return 1 + fi + if [ -n "$nccl" ] && ! cx_git_in_tree "$temporary" \ + -c "safe.directory=$temporary/third-party/nccl" \ + submodule update -q --init --depth 1 third-party/nccl; then + rm -rf -- "$temporary" + return 1 + fi + if [ "$backend" = deepep-v2 ]; then + CX_BACKEND_SOURCE_STEP="upstream NCCL check fix" + cx_apply_deepep_v2_nccl_check_fix "$temporary" || { + rm -rf -- "$temporary" + return 1 + } + fi + CX_BACKEND_SOURCE_STEP="source publication validation" + if ! cx_backend_source_is_valid "$backend" "$temporary" \ + || ! mv -- "$temporary" "$source"; then + rm -rf -- "$temporary" + return 1 + fi +} + +cx_prepare_backend_source() { + local log backend="$2" CX_BACKEND_SOURCE_STEP="initialization" + log="$(cx_private_log_path "backend-source-$backend")" || return 1 + if _cx_prepare_backend_source "$@" > "$log" 2>&1; then + return 0 + fi + printf '%s failed\n' "$CX_BACKEND_SOURCE_STEP" >> "$log" + cx_log "ERROR: backend-source-step=${CX_BACKEND_SOURCE_STEP// /-}" + cx_fail_stage backend-setup "$log" +} + +cx_materialize_backend_source() { + local backend="$1" destination="$2" source parent temporary + [ -n "${CX_BACKEND_SOURCE_ROOT:-}" ] || return 1 + source="$(cx_backend_source_path "$CX_BACKEND_SOURCE_ROOT" "$backend")" || return 1 + cx_backend_source_is_valid "$backend" "$source" || return 1 + parent="${destination%/*}" + [ "$parent" != "$destination" ] && [ -d "$parent" ] && [ ! -L "$parent" ] \ + || return 1 + temporary="$(mktemp -d "$parent/.collectivex-source.XXXXXX")" || return 1 + if ! cp -R -- "$source/." "$temporary/" \ + || ! cx_backend_source_is_valid "$backend" "$temporary"; then + rm -rf -- "$temporary" + return 1 + fi + if ! rm -rf -- "$destination" || ! mv -- "$temporary" "$destination"; then + rm -rf -- "$temporary" + return 1 + fi + if ! cx_backend_source_is_valid "$backend" "$destination"; then + rm -rf -- "$destination" + return 1 + fi + return 0 +} + +cx_prepare_implicit_stage_base() { + python3 - "${1:-}" "${2:-}" <<'PY' +import hashlib +import os +from pathlib import Path +import pwd +import stat +import sys + +def reject(reason): + print(f"[collectivex] FATAL: implicit-stage-validation={reason}", file=sys.stderr) + raise SystemExit(1) + +try: + configured_home = Path(sys.argv[1] or pwd.getpwuid(os.getuid()).pw_dir) +except (KeyError, OSError): + reject("account-home") +if not configured_home.is_absolute(): + reject("path-shape") +home = Path(os.path.realpath(configured_home)) +try: + metadata = os.stat(home, follow_symlinks=False) +except OSError: + reject("home-stat") +if not stat.S_ISDIR(metadata.st_mode): + reject("home-type") +if metadata.st_uid != os.getuid(): + reject("home-owner") +if stat.S_IMODE(metadata.st_mode) & stat.S_IWGRP: + reject("home-group-writable") +if stat.S_IMODE(metadata.st_mode) & stat.S_IWOTH: + reject("home-world-writable") +home_owner = metadata.st_uid +try: + isolation_key = sys.argv[2] + suffix = "" + if isolation_key: + suffix = "-" + hashlib.sha256(isolation_key.encode("utf-8")).hexdigest()[:16] + current = home / f".inferencex-collectivex-stage{suffix}" + created = False + try: + os.mkdir(current, mode=0o700) + created = True + except FileExistsError: + pass + metadata = os.stat(current, follow_symlinks=False) + if not stat.S_ISDIR(metadata.st_mode): + reject("child-type") + if metadata.st_uid not in {os.getuid(), home_owner} and not ( + isolation_key and created and metadata.st_uid == 0 + ): + reject("child-owner") + if Path(os.path.realpath(current)) != current: + reject("child-symlink") + if stat.S_IMODE(metadata.st_mode) & (stat.S_IWGRP | stat.S_IWOTH): + reject("child-writable") + os.chmod(current, 0o700) +except OSError: + reject("child-access") +print(current, end="") +PY +} + +cx_prepare_runner_shared_stage_base() { + local runner_temp="${RUNNER_TEMP:-}" runner_root + case "$runner_temp" in + /*/_work/_temp) runner_root="${runner_temp%/_work/_temp}" ;; + *) cx_die "canonical AMD execution requires a standard shared runner temp" ;; + esac + [ -n "$runner_root" ] && [ "$runner_root" != "$runner_temp" ] \ + || cx_die "canonical AMD execution requires a shared runner root" + cx_prepare_implicit_stage_base "$runner_root" +} + +cx_lock_canonical_gha_env() { + local runner="$1" expected_nodes expected_gpn expected_world trusted_lock_dir="" + local trusted_stage_dir="" + local trusted_qos="" + local trusted_socket_ifname="" trusted_rdma_devices="" + local trusted_ib_gid_index="" trusted_rdma_service_level="" + local trusted_rdma_traffic_class="" + local trusted_audit_salt="" + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" = 1 ] || return 0 + [ "${GITHUB_ACTIONS:-}" = true ] \ + || cx_die "canonical CollectiveX execution requires GitHub Actions" + [ -n "${CX_SHARD_FILE:-}" ] && [ "${CX_SHARD_SKU:-}" = "$runner" ] \ + || cx_die "canonical CollectiveX execution requires a matched shard" + [[ "${GITHUB_RUN_ID:-}" =~ ^[1-9][0-9]*$ \ + && "${GITHUB_RUN_ATTEMPT:-}" =~ ^[1-9][0-9]*$ \ + && "${COLLECTIVEX_SOURCE_SHA:-}" =~ ^[0-9a-f]{40,64}$ ]] \ + || cx_die "canonical CollectiveX workflow identity is incomplete" + + # cx_load_operator_config clears inherited values before setting this process marker. + # Preserve only values parsed from that private strict document. + if [ "${COLLECTIVEX_OPERATOR_CONFIG_LOADED:-}" = "$$" ]; then + trusted_lock_dir="${CX_LOCK_DIR:-}" + trusted_stage_dir="${CX_STAGE_DIR:-}" + trusted_qos="${CX_QOS:-}" + trusted_socket_ifname="${CX_SOCKET_IFNAME:-}" + trusted_rdma_devices="${CX_RDMA_DEVICES:-}" + trusted_ib_gid_index="${CX_IB_GID_INDEX:-}" + trusted_rdma_service_level="${CX_RDMA_SERVICE_LEVEL:-}" + trusted_rdma_traffic_class="${CX_RDMA_TRAFFIC_CLASS:-}" + trusted_audit_salt="${CX_AUDIT_SALT:-}" + fi + # The legacy B300 operator row contains a root-owned stage path. B300's + # compute-visible account home is the canonical source for its private base. + case "$runner" in b300|gb300) trusted_stage_dir="" ;; esac + unset CX_NCCL_HOME CX_MASTER_PORT CX_MORI_KERNEL_TYPE CX_LOCK_DIR CX_STAGE_DIR CX_QOS + unset CX_STAGE_PARENT_OWNER_OK + unset MASTER_ADDR MASTER_PORT RANK WORLD_SIZE LOCAL_RANK LOCAL_WORLD_SIZE + unset CX_SOCKET_IFNAME CX_RDMA_DEVICES CX_IB_GID_INDEX CX_RDMA_SERVICE_LEVEL + unset CX_RDMA_TRAFFIC_CLASS + unset CX_AUDIT_SALT + unset NCCL_NET NCCL_SOCKET_IFNAME GLOO_SOCKET_IFNAME NCCL_IB_HCA + unset NCCL_IB_GID_INDEX NCCL_IB_SL + unset NVSHMEM_DISABLE_IB NVSHMEM_ENABLE_NIC_PE_MAPPING + unset NVSHMEM_HCA_LIST NVSHMEM_IB_GID_INDEX NVSHMEM_IB_SL + unset NVSHMEM_IB_ENABLE_IBGDA NVSHMEM_IBGDA_NIC_HANDLER + unset NVSHMEM_HCA_PE_MAPPING NVSHMEM_REMOTE_TRANSPORT + unset EP_NIC_NAME EP_OVERRIDE_RDMA_SL + unset UCCL_SOCKET_IFNAME UCCL_IB_GID_INDEX UCCL_IB_SL MORI_RDMA_DEVICES + unset MORI_RDMA_TC MORI_IO_TC MORI_RDMA_SL MORI_IO_SL + unset HYBRID_EP_MULTINODE USE_NIXL RDMA_CORE_HOME DEEPEP_HYBRID_BUILD_MODE + unset MORI_COMMIT MORI_DISABLE_AUTO_XGMI MORI_ENABLE_SDMA + unset MORI_APP_LOG_LEVEL MORI_SHMEM_LOG_LEVEL MORI_IO_LOG_LEVEL + unset NCCL_CUMEM_ENABLE NCCL_MNNVL_ENABLE MC_FORCE_MNNVL + unset CX_BACKEND_CACHE_ROOT CX_BACKEND_CACHE_SENTINEL_SHA256 + unset CX_PREPARED_BACKEND_CACHE CX_BACKEND_SOURCE_ROOT + + [ -n "${CX_SQUASH_DIR:-}" ] \ + || cx_die "canonical CollectiveX execution requires shared container storage" + if [ -z "$trusted_stage_dir" ]; then + case "$runner" in + h100-dgxc) + trusted_stage_dir="$(cx_prepare_implicit_stage_base "${CX_SQUASH_DIR%/*}")" \ + || cx_die "canonical CollectiveX execution cannot create an isolated shared stage directory" + ;; + b300) + trusted_stage_dir="$(cx_prepare_implicit_stage_base "" \ + "${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-}}")" \ + || cx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + gb300) + trusted_stage_dir="$(cx_prepare_implicit_stage_base "" \ + "${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-}}")" \ + || cx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + h200-dgxc|b200-dgxc) + trusted_stage_dir="$(cx_prepare_implicit_stage_base)" \ + || cx_die "canonical CollectiveX execution cannot create an isolated stage directory" + ;; + mi300x|mi325x|mi355x) + # AMD self-hosted runners and compute nodes share the runner filesystem, + # while the image cache may be root-owned. Derive a runner-owned base + # outside _work instead of weakening stage ownership validation. + trusted_stage_dir="$(cx_prepare_runner_shared_stage_base)" \ + || cx_die "canonical AMD execution cannot create an isolated shared stage directory" + ;; + *) cx_die "canonical CollectiveX execution requires a configured shared stage directory" ;; + esac + elif [ "$runner" = mi300x ]; then + # The MI300X runner home is a shared-filesystem symlink. Resolve the + # operator-selected base once; cx_stage_path still validates the canonical + # directory's ownership, permissions, overlap, and per-run child path. + trusted_stage_dir="$(python3 -c \ + 'import os,sys; p=os.path.realpath(sys.argv[1]); assert os.path.isdir(p); print(p,end="")' \ + "$trusted_stage_dir")" \ + || cx_die "canonical MI300X execution cannot resolve the shared stage directory" + fi + [[ "$trusted_audit_salt" =~ ^[0-9a-f]{64}$ ]] \ + || cx_die "canonical CollectiveX execution requires a private audit salt" + if [ "$runner" = b300 ]; then + CX_STAGE_PARENT_OWNER_OK=1 + fi + + case "$runner" in + h100-dgxc|h200-dgxc|b200-dgxc|b300) + expected_nodes="${CX_NODES:-}"; expected_gpn=8 + [ "$expected_nodes" = 1 ] || [ "$expected_nodes" = 2 ] \ + || cx_die "canonical NVIDIA execution requires one or two nodes" + CX_IMAGE="$CX_IMAGE_MULTIARCH" + CX_IMAGE_DIGEST="$CX_IMAGE_MULTIARCH_DIGEST" + CX_NCCL_HOME=/usr + ;; + gb200|gb300) + expected_nodes="${CX_NODES:-}"; expected_gpn=4 + [ "$expected_nodes" = 2 ] || [ "$expected_nodes" = 4 ] \ + || cx_die "canonical GB execution requires two or four trays" + CX_IMAGE="$CX_IMAGE_MULTIARCH" + CX_IMAGE_DIGEST="$CX_IMAGE_MULTIARCH_DIGEST" + CX_NCCL_HOME=/usr + CX_MASTER_PORT=29551 + ;; + mi300x|mi325x) + expected_nodes="${CX_NODES:-}"; expected_gpn=8 + [ "$expected_nodes" = 1 ] || [ "$expected_nodes" = 2 ] \ + || cx_die "canonical AMD execution requires one or two nodes" + CX_IMAGE="$CX_IMAGE_AMD_MORI_MI325" + CX_IMAGE_DIGEST="$CX_IMAGE_AMD_MORI_MI325_DIGEST" + if [ "$expected_nodes" = 2 ]; then + CX_MORI_KERNEL_TYPE=internode-v1 + else + CX_MORI_KERNEL_TYPE=asyncll + fi + MORI_COMMIT="$CX_MORI_COMMIT_MI325" + MORI_DISABLE_AUTO_XGMI=0 + MORI_ENABLE_SDMA=1 + MORI_APP_LOG_LEVEL=info + MORI_SHMEM_LOG_LEVEL=info + MORI_IO_LOG_LEVEL=info + ;; + mi355x) + expected_nodes="${CX_NODES:-}"; expected_gpn=8 + [ "$expected_nodes" = 1 ] || [ "$expected_nodes" = 2 ] \ + || cx_die "canonical AMD execution requires one or two nodes" + CX_IMAGE="$CX_IMAGE_AMD_MORI" + CX_IMAGE_DIGEST="$CX_IMAGE_AMD_MORI_DIGEST" + if [ "$expected_nodes" = 2 ]; then + CX_MORI_KERNEL_TYPE=internode-v1 + else + CX_MORI_KERNEL_TYPE=intranode + fi + MORI_COMMIT="$CX_MORI_COMMIT_MI355" + ;; + *) cx_die "canonical CollectiveX runner is not registered" ;; + esac + case "$runner:$trusted_lock_dir" in + mi300x:?*|mi325x:?*|mi355x:?*) export CX_LOCK_DIR="$trusted_lock_dir" ;; + esac + CX_STAGE_DIR="$trusted_stage_dir" + [ -z "$trusted_qos" ] || export CX_QOS="$trusted_qos" + [ -z "$trusted_socket_ifname" ] \ + || export CX_SOCKET_IFNAME="$trusted_socket_ifname" + [ -z "$trusted_rdma_devices" ] \ + || export CX_RDMA_DEVICES="$trusted_rdma_devices" + [ -z "$trusted_ib_gid_index" ] \ + || export CX_IB_GID_INDEX="$trusted_ib_gid_index" + [ -z "$trusted_rdma_service_level" ] \ + || export CX_RDMA_SERVICE_LEVEL="$trusted_rdma_service_level" + [ -z "$trusted_rdma_traffic_class" ] \ + || export CX_RDMA_TRAFFIC_CLASS="$trusted_rdma_traffic_class" + CX_AUDIT_SALT="$trusted_audit_salt" + export CX_STAGE_DIR CX_AUDIT_SALT + [ "${CX_NODES:-}" = "$expected_nodes" ] \ + && [ "${CX_GPUS_PER_NODE:-}" = "$expected_gpn" ] \ + || cx_die "canonical CollectiveX placement differs from the shard" + expected_world=$((expected_nodes * expected_gpn)) + CX_NGPUS="$expected_world" + CX_SEED=67 + case "$runner" in mi300x|mi325x|mi355x) CX_RUN_TIMEOUT=1800 ;; *) CX_RUN_TIMEOUT=900 ;; esac + unset CX_PUBLIC_RUNNER CX_GB_PRODUCT CX_DRYRUN CX_TIMING CX_ALLOW_MNNVL + unset CX_ENROOT_LOCAL_IMPORT COLLECTIVEX_IMAGE COLLECTIVEX_IMAGE_DIGEST + unset COLLECTIVEX_IMAGE_DIGEST_VERIFIED COLLECTIVEX_SQUASH_SHA256 + export CX_IMAGE CX_IMAGE_DIGEST CX_NGPUS CX_SEED CX_RUN_TIMEOUT + case "$runner" in + h100-dgxc|h200-dgxc|b200-dgxc|b300) export CX_NCCL_HOME ;; + gb200|gb300) export CX_NCCL_HOME CX_MASTER_PORT ;; + mi300x|mi325x) + export CX_MORI_KERNEL_TYPE MORI_COMMIT MORI_DISABLE_AUTO_XGMI MORI_ENABLE_SDMA + export MORI_APP_LOG_LEVEL MORI_SHMEM_LOG_LEVEL MORI_IO_LOG_LEVEL + ;; + mi355x) export CX_MORI_KERNEL_TYPE MORI_COMMIT ;; + esac +} + +cx_reverify_registry_image() { + local image="$1" actual + [[ "${COLLECTIVEX_IMAGE_DIGEST:-}" =~ ^sha256:[0-9a-f]{64}$ ]] \ + && [ "${COLLECTIVEX_IMAGE_DIGEST_VERIFIED:-0}" = 1 ] || return 1 + actual="$(cx_resolve_registry_digest "$image")" || return 1 + [ "$actual" = "$COLLECTIVEX_IMAGE_DIGEST" ] || { + cx_log "ERROR: configured image tag changed during container import" + return 1 + } +} + +cx_export_squash_identity() { + local image="$1" digest log + log="$(cx_private_log_path container-hash)" + digest="$(sha256sum "$image" 2>> "$log" | awk '{print $1}')" + [[ "$digest" =~ ^[0-9a-f]{64}$ ]] \ + || { cx_fail_stage container-hash "$log"; return 1; } + export COLLECTIVEX_SQUASH_SHA256="$digest" +} + +cx_squash_path() { + local squash_dir="$1" image="$2" key platform + [[ "${COLLECTIVEX_IMAGE_DIGEST:-}" =~ ^sha256:[0-9a-f]{64}$ ]] \ + || return 1 + case "${CX_IMAGE_PLATFORM:-}" in + linux/amd64) platform="" ;; + linux/arm64) platform="_linux_arm64" ;; + *) return 1 ;; + esac + key="${CX_SQUASH_FORMAT_VERSION}${platform}_${COLLECTIVEX_IMAGE_DIGEST#sha256:}_$( + printf '%s' "$image" | sed 's#[/:@#]#_#g' + )" + printf '%s' "$squash_dir/${key}.sqsh" +} + +# cx_ensure_squash -> echoes the squash file path. +# Imports via Enroot only if a valid squash is not already present, under a lock. +cx_ensure_squash() { + local squash_dir="$1" image="$2" key sq locks lock_fd log + local enroot_local="" import_rc=0 machine + log="$(cx_private_log_path container-import)" + machine="$(uname -m)" + case "${CX_IMAGE_PLATFORM:-}:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) cx_fail_stage container-import "$log"; return 1 ;; + esac + mkdir -p "$squash_dir" 2>> "$log" \ + || { cx_fail_stage container-import "$log"; return 1; } + sq="$(cx_squash_path "$squash_dir" "$image")" \ + || { cx_fail_stage container-import "$log"; return 1; } + key="${sq##*/}" + key="${key%.sqsh}" + locks="$squash_dir/.locks" + mkdir -p "$locks" 2>> "$log" \ + || { cx_fail_stage container-import "$log"; return 1; } + { exec {lock_fd}>"$locks/${key}.lock"; } 2>> "$log" \ + || { cx_fail_stage container-import "$log"; return 1; } + flock -w 900 "$lock_fd" 2>> "$log" \ + || { cx_fail_stage container-import "$log"; return 1; } + if unsquashfs -l "$sq" >/dev/null 2>&1; then + cx_log "container squash ready" + else + cx_log "importing configured container image" + rm -f "$sq" 2>> "$log" \ + || { cx_fail_stage container-import "$log"; return 1; } + # > "$log" 2>&1 || import_rc=$? + rm -rf -- "$enroot_local" >/dev/null 2>&1 || true + [ "$import_rc" = 0 ] \ + || { cx_fail_stage container-import "$log"; return 1; } + else + SOURCE_DATE_EPOCH="$CX_SQUASH_SOURCE_DATE_EPOCH" \ + enroot import -o "$sq" "docker://$image" > "$log" 2>&1 \ + || { cx_fail_stage container-import "$log"; return 1; } + fi + unsquashfs -l "$sq" >> "$log" 2>&1 \ + || { cx_fail_stage container-import "$log"; return 1; } + fi + if ! cx_reverify_registry_image "$image" >> "$log" 2>&1; then + flock -u "$lock_fd" >/dev/null 2>&1 || true + exec {lock_fd}>&- + cx_fail_stage container-import "$log" + return 1 + fi + flock -u "$lock_fd" + exec {lock_fd}>&- + echo "$sq" +} + +# Import on an allocated compute node so multiarch tags resolve for the target +# architecture. The squash directory must be shared with the submit host. +cx_ensure_squash_on_job() { + local job_id="$1" squash_dir="$2" image="$3" lock_dir="${4:-}" sq key lock + local log_label=container-import log + [[ "$job_id" =~ ^[0-9]+$ ]] || return 1 + case "${CX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${CX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + sq="$(cx_squash_path "$squash_dir" "$image")" || return 1 + key="${sq##*/}" + key="${key%.sqsh}" + [ -n "$lock_dir" ] || lock_dir="$squash_dir/.locks" + lock="$lock_dir/${key}.lock" + log="$(cx_private_log_path "$log_label")" + if ! srun --jobid="$job_id" --nodes=1 --ntasks=1 --chdir=/tmp \ + --export="$(cx_host_exports)" \ + bash -s -- "$sq" "$lock" "$image" "$CX_SQUASH_SOURCE_DATE_EPOCH" \ + "$CX_IMAGE_PLATFORM" \ + > "$log" 2>&1 <<'BASH' +set -euo pipefail +sq="$1"; lock="$2"; image="$3"; source_date_epoch="$4"; platform="$5" +machine="$(uname -m)" +case "$platform:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) exit 13 ;; +esac +compute_home="$(mktemp -d /tmp/inferencex-collectivex-home.XXXXXX)" +trap 'rm -rf -- "$compute_home"' EXIT +export HOME="$compute_home" XDG_CACHE_HOME="$compute_home/.cache" +export ENROOT_TEMP_PATH="$compute_home/enroot-tmp" +export ENROOT_CACHE_PATH="$compute_home/enroot-cache" +export ENROOT_DATA_PATH="$compute_home/enroot-data" +export ENROOT_RUNTIME_PATH="$compute_home/enroot-run" +mkdir -p "$(dirname "$sq")" "$(dirname "$lock")" \ + "$ENROOT_TEMP_PATH" "$ENROOT_CACHE_PATH" "$ENROOT_DATA_PATH" "$ENROOT_RUNTIME_PATH" +exec 9>"$lock" +flock -w 900 9 +if unsquashfs -l "$sq" >/dev/null 2>&1; then + echo 'container squash ready' +else + rm -f -- "$sq" + SOURCE_DATE_EPOCH="$source_date_epoch" \ + enroot import -o "$sq" "docker://$image" /dev/null 2>&1 +fi +BASH + then + cx_fail_stage container-import "$log" + return 1 + fi + if ! cx_reverify_registry_image "$image" >> "$log" 2>&1; then + cx_fail_stage container-import "$log" + return 1 + fi + printf '%s' "$sq" +} + +cx_preflight_allocation() { + local job_id="$1" nodes="$2" mount_src="$3" squash="$4" shard="${5:-}" + local log rc=0 runtime shard_path="" probe_root probe_token index + runtime="$mount_src/experimental/CollectiveX/runtime/run_in_container.sh" + [ -z "$shard" ] || shard_path="$mount_src/experimental/CollectiveX/$shard" + log="$(cx_private_log_path allocation-preflight)" + probe_root="$mount_src/.collectivex-preflight" + probe_token="$probe_root/source" + if [ -e "$probe_root" ] || [ -L "$probe_root" ] \ + || ! mkdir -m 700 "$probe_root"; then + cx_fail_stage repository-stage "$log" + return 1 + fi + if ! printf '%s\n' "${COLLECTIVEX_EXECUTION_ID:-manual-$$}" > "$probe_token" \ + || ! chmod 600 "$probe_token"; then + chmod 700 "$probe_root" >/dev/null 2>&1 || true + rm -rf -- "$probe_root" >/dev/null 2>&1 || true + cx_fail_stage repository-stage "$log" + return 1 + fi + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp --input=all \ + --export="$(cx_host_exports)" bash -s -- "$runtime" "$shard_path" "$squash" \ + "$CX_IMAGE_PLATFORM" "$probe_root" \ + > "$log" 2>&1 <<'BASH' || rc=$? +set -euo pipefail +machine="$(uname -m)" +case "$4:$machine" in + linux/amd64:x86_64|linux/amd64:amd64|linux/arm64:aarch64|linux/arm64:arm64) ;; + *) exit 13 ;; +esac +test -r "$1" || exit 10 +[ -z "$2" ] || test -r "$2" || exit 11 +test -r "$3" || exit 12 +unsquashfs -s "$3" >/dev/null 2>&1 || exit 12 +case "${SLURM_NODEID:-}" in ""|*[!0-9]*) exit 10 ;; esac +[ -d "$5" ] && [ ! -L "$5" ] && [ -r "$5/source" ] || exit 10 +(set -C; cat "$5/source" > "$5/node-$SLURM_NODEID") || exit 10 +cmp -s -- "$5/source" "$5/node-$SLURM_NODEID" || exit 10 +BASH + if [ "$rc" = 0 ]; then + for ((index = 0; index < nodes; index++)); do + if ! cmp -s -- "$probe_token" "$probe_root/node-$index"; then + rc=10 + break + fi + done + fi + if [ -d "$probe_root" ] && [ ! -L "$probe_root" ]; then + chmod 700 "$probe_root" >/dev/null 2>&1 || rc=10 + fi + rm -rf -- "$probe_root" >/dev/null 2>&1 || rc=10 + [ "$rc" = 0 ] && return 0 + case "$rc" in + 10|11) cx_fail_stage repository-stage "$log" ;; + 12) cx_fail_stage container-hash "$log" ;; + *) cx_fail_stage container-launch "$log" ;; + esac + return 1 +} + +# A clean nvidia-smi inventory does not prove that a prior cancelled workload +# released every CUDA context. Retaining each primary context catches poisoned +# allocations before a full shard spends time failing every case. +cx_validate_cuda_context_on_job() { + local job_id="$1" nodes="$2" gpus_per_node="$3" log_label=cuda-context log + case "${CX_SALLOC_ATTEMPT:-1}" in + 1) ;; + 2|3) log_label+="-a${CX_SALLOC_ATTEMPT}" ;; + *) return 1 ;; + esac + log="$(cx_private_log_path "$log_label")" + CX_CUDA_CONTEXT_LOG="$log" + srun --jobid="$job_id" --nodes="$nodes" --ntasks="$nodes" --ntasks-per-node=1 \ + --gres=gpu:"$gpus_per_node" --chdir=/tmp --input=all \ + --export="$(cx_host_exports)" python3 - "$gpus_per_node" \ + >"$log" 2>&1 <<'PY' +import ctypes +import socket +import sys + +expected = int(sys.argv[1]) +cuda = ctypes.CDLL("libcuda.so.1") +cuda.cuInit.argtypes = [ctypes.c_uint] +cuda.cuInit.restype = ctypes.c_int +cuda.cuDeviceGetCount.argtypes = [ctypes.POINTER(ctypes.c_int)] +cuda.cuDeviceGetCount.restype = ctypes.c_int +cuda.cuDeviceGet.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int] +cuda.cuDeviceGet.restype = ctypes.c_int +cuda.cuDevicePrimaryCtxRetain.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), ctypes.c_int, +] +cuda.cuDevicePrimaryCtxRetain.restype = ctypes.c_int +cuda.cuDevicePrimaryCtxRelease.argtypes = [ctypes.c_int] +cuda.cuDevicePrimaryCtxRelease.restype = ctypes.c_int + + +def check(result: int, operation: str) -> None: + if result != 0: + raise RuntimeError(f"{operation} failed with CUDA result {result}") + + +check(cuda.cuInit(0), "CUDA initialization") +count = ctypes.c_int() +check(cuda.cuDeviceGetCount(ctypes.byref(count)), "CUDA device count") +if count.value != expected: + raise RuntimeError( + f"CUDA device count mismatch on {socket.gethostname()}: " + f"expected {expected}, found {count.value}" + ) + +devices: list[int] = [] +try: + for ordinal in range(expected): + device = ctypes.c_int() + context = ctypes.c_void_p() + check(cuda.cuDeviceGet(ctypes.byref(device), ordinal), "CUDA device lookup") + result = cuda.cuDevicePrimaryCtxRetain(ctypes.byref(context), device.value) + if result != 0: + raise RuntimeError( + f"CUDA primary context retain failed on {socket.gethostname()} " + f"device {ordinal} with CUDA result {result}" + ) + devices.append(device.value) +finally: + for device in reversed(devices): + check(cuda.cuDevicePrimaryCtxRelease(device), "CUDA primary context release") +PY +} + +# Resolve the exact per-execution child before any copy starts, so the parent +# EXIT trap can remove an interrupted partial stage. The configured base must +# already exist on compute-visible storage and must not traverse symlinks. +cx_stage_path() { + local repo_root="$1" stage_base="${2:-}" tag safe_tag stage_path + tag="${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-manual-$$}}" + [[ "$tag" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] \ + || cx_die "invalid staging execution identity" + safe_tag="$(printf '%s' "$tag" | tr -c 'A-Za-z0-9._-' '_')" + if [ -z "$stage_base" ] || [ "$stage_base" = "$repo_root" ]; then + [ -n "${CX_SQUASH_DIR:-}" ] \ + || cx_die "CollectiveX staging requires CX_STAGE_DIR or CX_SQUASH_DIR" + stage_base="$CX_SQUASH_DIR" + stage_path="${stage_base%/}/.collectivex-stage-$safe_tag" + else + stage_path="${stage_base%/}/job_$safe_tag" + fi + python3 - "$repo_root" "$stage_base" "$stage_path" \ + "${CX_JOB_ROOT:-}" "${GITHUB_WORKSPACE:-}" \ + "${CX_STAGE_PARENT_OWNER_OK:-0}" <<'PY' +import os +import stat +import sys + +repo, base, child, job_root, workspace, allow_parent_owner = sys.argv[1:] +def reject(reason): + print(f"[collectivex] FATAL: stage-base-validation={reason}", file=sys.stderr) + raise SystemExit(1) + +try: + if ( + not os.path.isabs(repo) + or os.path.realpath(repo) != repo + or not os.path.isabs(base) + or os.path.realpath(base) != base + or not os.path.isabs(child) + or os.path.dirname(child) != base.rstrip("/") + ): + reject("path-shape") + if os.path.lexists(child): + reject("child-exists") + try: + metadata = os.stat(base, follow_symlinks=False) + except OSError: + reject("base-stat") + excluded = [repo] + excluded.extend(path for path in (job_root, workspace) if path) + for path in excluded: + resolved = os.path.realpath(path) + if os.path.commonpath((base, resolved)) == resolved: + reject("overlap") + if not stat.S_ISDIR(metadata.st_mode): + reject("not-directory") + if metadata.st_uid != os.getuid(): + if allow_parent_owner != "1": + reject("owner") + parent = os.path.dirname(base.rstrip("/")) + parent_metadata = os.stat(parent, follow_symlinks=False) + if ( + not stat.S_ISDIR(parent_metadata.st_mode) + or metadata.st_uid not in {parent_metadata.st_uid, 0} + or stat.S_IMODE(parent_metadata.st_mode) & (stat.S_IWGRP | stat.S_IWOTH) + ): + reject("parent-owner") + if stat.S_IMODE(metadata.st_mode) & stat.S_IWGRP: + reject("group-writable") + if stat.S_IMODE(metadata.st_mode) & stat.S_IWOTH: + reject("world-writable") + if not os.access(base, os.W_OK | os.X_OK): + reject("access") +except ValueError: + reject("path-shape") +print(child, end="") +PY +} + +# Stage only the public benchmark tree into a pre-resolved, private execution +# child. A runner-owned marker makes recursive cleanup an explicit capability. +cx_stage_repo() { + local repo_root="$1" stage_dir="$2" expected log tag marker + cx_validate_shard_control "$repo_root/experimental/CollectiveX" + expected="$(cx_stage_path "$repo_root" "${CX_STAGE_DIR:-}")" \ + || cx_die "configured stage base is unavailable or unsafe" + [ "$stage_dir" = "$expected" ] \ + || cx_die "execution stage differs from the configured stage base" + tag="${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-manual-$$}}" + if [ -e "$stage_dir" ] || [ -L "$stage_dir" ]; then + cx_die "refusing to reuse a pre-existing execution stage" + fi + mkdir -m 700 "$stage_dir" 2>/dev/null \ + || cx_die "cannot create the configured stage directory" + chmod 700 "$stage_dir" 2>/dev/null \ + || cx_die "cannot protect the configured stage directory" + marker="$stage_dir/.collectivex-stage-v1" + umask 077 + (set -C; printf 'collectivex-stage-v1\n%s\n' "$tag" > "$marker") 2>/dev/null \ + || cx_die "cannot claim the configured stage directory" + chmod 600 "$marker" 2>/dev/null \ + || cx_die "cannot protect the configured stage directory" + mkdir -m 700 "$stage_dir/experimental" 2>/dev/null \ + || cx_die "cannot create the configured stage directory" + cx_log "staging CollectiveX on compute-visible storage" + log="$(cx_private_log_path repository-stage)" + if ! python3 - "$repo_root/experimental/CollectiveX" \ + "$stage_dir/experimental/CollectiveX" > "$log" 2>&1 <<'PY' +import os +from pathlib import Path +import shutil +import sys + +def report_error(kind, value, trace): + error_number = getattr(value, "errno", 0) + print(f"collectivex-stage-copy-error={kind.__name__}:{error_number or 0}", file=sys.stderr) + +sys.excepthook = report_error +source, target = map(Path, sys.argv[1:]) +excluded = { + Path("__pycache__"), Path("results"), Path(".cx_workloads"), + Path(".cx_backend"), Path(".cx_sources"), Path("configs/platforms.yaml"), + Path("private-infra.md"), Path("goal.md"), Path("notes.md"), +} +for root, directories, files in os.walk(source, followlinks=False): + root_path = Path(root) + relative_root = root_path.relative_to(source) + directories[:] = [ + name for name in directories + if relative_root / name not in excluded and not (root_path / name).is_symlink() + ] + destination = target / relative_root + destination.mkdir(mode=0o700, parents=True, exist_ok=True) + for name in files: + relative = relative_root / name + if relative in excluded: + continue + source_file = root_path / name + if source_file.is_symlink() or not source_file.is_file(): + raise RuntimeError("unsupported source entry") + with source_file.open("rb") as input_file, (destination / name).open("xb") as output_file: + shutil.copyfileobj(input_file, output_file) +PY + then + copy_error="$(grep -aoE 'collectivex-stage-copy-error=[A-Za-z]+:[0-9]+' "$log" \ + | tail -n 1 || true)" + [ -z "$copy_error" ] || cx_log "ERROR: repository-stage-$copy_error" + rm -rf -- "$stage_dir" >/dev/null 2>&1 \ + || cx_log "ERROR: cannot remove the incomplete execution stage" + cx_fail_stage repository-stage "$log" || true + return 1 + fi +} + +# cx_collect_results +# When the run used a staged (compute-visible) mount, copy result JSONs back to +# the original checkout's results/ so the workflow's upload-artifact (which reads +# the checkout, not the stage dir) finds them. No-op when no staging was used. +cx_collect_results() { + local mount_src="$1" repo_root="$2" dst log + local -a files + [ "$mount_src" = "$repo_root" ] && return 0 + log="$(cx_private_log_path "artifact-collection-$$-${RANDOM}")" + dst="$repo_root/experimental/CollectiveX/results" + mkdir -p "$dst" 2>> "$log" \ + || { cx_log "ERROR: cannot create checkout result directory"; return 1; } + shopt -s nullglob + files=("$mount_src/experimental/CollectiveX/results/"*.json) + shopt -u nullglob + [ "${#files[@]}" -gt 0 ] || { cx_log "ERROR: staged run produced no result JSON"; return 1; } + cp -- "${files[@]}" "$dst/" >> "$log" 2>&1 \ + || { cx_log "ERROR: staged result collection failed"; return 1; } + cx_log "collected staged results for artifact validation" +} + +cx_cleanup_stage() { + local mount_src="$1" repo_root="$2" base="${CX_STAGE_DIR:-}" tag safe_tag expected + tag="${COLLECTIVEX_EXECUTION_ID:-${GITHUB_RUN_ID:-manual-$$}}" + safe_tag="$(printf '%s' "$tag" | tr -c 'A-Za-z0-9._-' '_')" + [ "$mount_src" != "$repo_root" ] || return 0 + if [ -n "$base" ] && [ "$base" != "$repo_root" ]; then + expected="${base%/}/job_$safe_tag" + else + [ -n "${CX_SQUASH_DIR:-}" ] \ + || { cx_log "ERROR: cannot identify the generated stage directory"; return 1; } + expected="${CX_SQUASH_DIR%/}/.collectivex-stage-$safe_tag" + fi + if [ "$mount_src" != "$expected" ] || [ "$mount_src" = / ] \ + || { [ -n "$base" ] && [ "$mount_src" = "$base" ]; }; then + cx_log "ERROR: refusing to remove an unrecognized stage directory" + return 1 + fi + if ! python3 - "$mount_src" "$tag" "${CX_STAGE_PARENT_OWNER_OK:-0}" <<'PY' +import os +from pathlib import Path +import stat +import sys + +root = Path(sys.argv[1]) +expected = f"collectivex-stage-v1\n{sys.argv[2]}\n" +allow_uid_mapped_root = sys.argv[3] == "1" +try: + metadata = os.stat(root, follow_symlinks=False) + marker = root / ".collectivex-stage-v1" + owner_ok = metadata.st_uid == os.getuid() + if not owner_ok and allow_uid_mapped_root and metadata.st_uid == 0: + base = root.parent + private_parent = base.parent + base_metadata = os.stat(base, follow_symlinks=False) + parent_metadata = os.stat(private_parent, follow_symlinks=False) + owner_ok = ( + stat.S_ISDIR(base_metadata.st_mode) + and base_metadata.st_uid == 0 + and stat.S_IMODE(base_metadata.st_mode) == 0o700 + and stat.S_ISDIR(parent_metadata.st_mode) + and parent_metadata.st_uid != 0 + and not stat.S_IMODE(parent_metadata.st_mode) & (stat.S_IWGRP | stat.S_IWOTH) + ) + if ( + not stat.S_ISDIR(metadata.st_mode) + or not owner_ok + or (stat.S_IMODE(metadata.st_mode) & 0o777) != 0o700 + ): + raise OSError + entries = list(root.iterdir()) + if marker.exists(): + marker_metadata = os.stat(marker, follow_symlinks=False) + if ( + not stat.S_ISREG(marker_metadata.st_mode) + or marker_metadata.st_uid != metadata.st_uid + or stat.S_IMODE(marker_metadata.st_mode) != 0o600 + ): + raise OSError + marker_content = marker.read_text() + if marker_content != expected and entries != [marker]: + raise OSError + elif entries: + raise OSError +except (OSError, UnicodeError): + raise SystemExit(1) +PY + then + cx_log "ERROR: refusing to remove an unowned stage directory" + return 1 + fi + rm -rf -- "$mount_src" >/dev/null 2>&1 || { + cx_log "ERROR: cannot remove generated stage directory" + return 1 + } + cx_log "removed generated per-execution stage directory" +} + +# Run one validated shard with one Slurm task per GPU. Launchers provide only +# allocation/container policy through globals and CX_DISTRIBUTED_CONTAINER_ARGS. +# shellcheck disable=SC2153 +cx_run_distributed_shard() { + local build_log build_rc cases_file expected_cases ci=0 failed_cases=0 + local ph mode routing eplb hidden topk experts ladder suite workload + local canonical case_id ep timing case_iters case_trials case_warmup case_stem + local scope scale_up_transport scale_out_transport transport topology_class nodes gpn domain + local workload_dir workload_ladder workload_log stage_rc attempt_tag out + local runtime_log run_rc summary_log + local -a container_args workload_args ep_args + [ "${NODES:-0}" -gt 1 ] && [ "${NGPUS:-0}" = "$((NODES * GPN))" ] \ + || cx_die "invalid distributed launcher placement" + [ -n "${JOB_ID:-}" ] && [ -n "${SQUASH_FILE:-}" ] \ + && [ -n "${CONTAINER_MOUNTS:-}" ] || cx_die "distributed launcher is incomplete" + [ -n "${SOURCE_BACKEND_ENV:-}" ] && [ -n "${BACKEND_PROBE:-}" ] \ + && [ -n "${WRAP:-}" ] || cx_die "distributed rank wrapper is incomplete" + + cx_resolve_slurm_rendezvous "$JOB_ID" + mkdir -p "$MOUNT_SRC/experimental/CollectiveX/results" + container_args=(--container-mounts="$CONTAINER_MOUNTS" --no-container-mount-home + --container-workdir=/ix/experimental/CollectiveX --no-container-entrypoint) + if declare -p CX_DISTRIBUTED_CONTAINER_ARGS >/dev/null 2>&1; then + container_args+=("${CX_DISTRIBUTED_CONTAINER_ARGS[@]}") + fi + local container_name="cxep_${JOB_ID}" + + cx_log "distributed backend preparation: bench=$CX_BENCH nodes=$NODES" + cx_set_failure_stage backend-setup + build_log="$(cx_private_log_path backend-prepare)" + set +e + srun --jobid="$JOB_ID" --nodes="$NODES" --ntasks-per-node=1 --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" --export="$(cx_container_exports),CX_BUILD_ONLY=1" \ + bash /ix/experimental/CollectiveX/runtime/run_in_container.sh \ + "$build_log" 2>&1 + build_rc=$? + if [ "$build_rc" = 0 ]; then + srun --jobid="$JOB_ID" --nodes="$NODES" --ntasks-per-node=1 --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" \ + --export="$(cx_container_exports)" bash -c "$BACKEND_PROBE" \ + >"$build_log" 2>&1 + build_rc=$? + fi + set -e + if [ "$build_rc" != 0 ]; then + cx_fail_stage backend-setup "$build_log" || true + return "$build_rc" + fi + cx_set_failure_stage execution + + cases_file="$(mktemp)" || return 1 + local shard="${CX_SHARD_FILE:-}" + [ -z "$shard" ] || [ -f "$shard" ] || shard="$CX_DIR/$shard" + # Iterable benchmark version is a shard-level scalar; export it once so the + # harness copies it verbatim into every emitted result for this shard. + if [ -n "$shard" ] && [ -f "$shard" ]; then + CX_VERSION="$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['version'])" "$shard")" + export CX_VERSION + fi + if [ -n "$shard" ]; then + if [ ! -f "$shard" ] || ! python3 - "$shard" > "$cases_file" <<'PY' +import json +import sys + +with open(sys.argv[1]) as handle: + cases = json.load(handle)["cases"] +for case in cases: + get = lambda key, default="": str(case.get(key) or default) + fields = ( + get("phase", "decode"), get("mode", "normal"), get("routing", "uniform"), + "1" if case.get("eplb") else "", get("hidden", "7168"), + get("topk", "8"), get("experts", "256"), get("ladder"), + get("suite"), get("workload"), + "1" if case.get("canonical") else "", get("case_id"), get("ep"), + get("timing", "8:64:32"), get("nodes"), get("gpus_per_node"), + get("scale_up_domain"), get("scope"), get("scale_up_transport"), + get("scale_out_transport"), get("transport"), get("topology_class"), + ) + print("|".join(fields)) +PY + then + rm -f "$cases_file" + cx_die "could not enumerate validated shard cases" + fi + else + local phases="${CX_PHASE:-decode}" phase + [ "$phases" = both ] && phases="decode prefill" + cx_require_record_safe "$phases" "${CX_MODE:-normal}" "${CX_ROUTING:-uniform}" \ + "${CX_EPLB:-}" "${CX_HIDDEN:-7168}" "${CX_TOPK:-8}" "${CX_EXPERTS:-256}" \ + "${CX_TOKENS_LADDER:-}" "${CX_SUITE:-}" "${CX_WORKLOAD_NAME:-}" \ + "${CX_CANONICAL:-}" "${CX_CASE_ID:-}" \ + "${CX_ITERS:-8}" "${CX_TRIALS:-64}" "${CX_WARMUP:-32}" \ + "${CX_SCOPE:-scale-up}" \ + "${CX_SCALE_UP_TRANSPORT:-unknown}" "${CX_SCALE_OUT_TRANSPORT:-}" \ + "${CX_TRANSPORT:-unknown}" "${CX_TOPO:-manual}" + for phase in $phases; do + (IFS='|'; printf '%s\n' "$phase|${CX_MODE:-normal}|${CX_ROUTING:-uniform}|${CX_EPLB:-}|${CX_HIDDEN:-7168}|${CX_TOPK:-8}|${CX_EXPERTS:-256}|${CX_TOKENS_LADDER:-}|${CX_SUITE:-}|${CX_WORKLOAD_NAME:-}|${CX_CANONICAL:-}|${CX_CASE_ID:-}|$NGPUS|${CX_ITERS:-8}:${CX_TRIALS:-64}:${CX_WARMUP:-32}|$NODES|$GPN|$SCALE_UP_DOMAIN|${CX_SCOPE:-scale-up}|${CX_SCALE_UP_TRANSPORT:-unknown}|${CX_SCALE_OUT_TRANSPORT:-}|${CX_TRANSPORT:-unknown}|${CX_TOPO:-manual}") + done > "$cases_file" + fi + expected_cases="$(wc -l < "$cases_file" | tr -d ' ')" + [ "$expected_cases" -gt 0 ] \ + || { rm -f "$cases_file"; cx_die "distributed case list is empty"; } + + while IFS='|' read -r ph mode routing eplb hidden topk experts ladder suite workload \ + canonical case_id ep timing nodes gpn domain scope scale_up_transport \ + scale_out_transport transport topology_class; do + [ -n "$ph" ] || continue + ci=$((ci + 1)) + case_stem="${RUNNER}_${CX_BENCH}_${ph}_${TS}-c$(printf '%03d' "$ci")" + IFS=: read -r case_iters case_trials case_warmup <<< "${timing:-8:64:32}" + case_iters="${case_iters:-8}" + case_trials="${case_trials:-64}" + case_warmup="${case_warmup:-32}" + ep="${ep:-$NGPUS}" + export CX_MODE="$mode" CX_PHASE="$ph" CX_CASE_ID="$case_id" CX_SUITE="$suite" + export CX_WORKLOAD_NAME="$workload" + export CX_CANONICAL="$canonical" CX_EP="$ep" + export CX_ROUTING="$routing" CX_EPLB="$eplb" CX_TOKENS_LADDER="$ladder" + export CX_HIDDEN="$hidden" CX_TOPK="$topk" CX_EXPERTS="$experts" + export CX_NODES="$nodes" CX_GPUS_PER_NODE="$gpn" CX_SCALE_UP_DOMAIN="$domain" + export CX_SCOPE="$scope" CX_SCALE_UP_TRANSPORT="$scale_up_transport" + export CX_SCALE_OUT_TRANSPORT="$scale_out_transport" + export CX_TRANSPORT="$transport" CX_TOPO="$topology_class" + export CX_ITERS="$case_iters" CX_TRIALS="$case_trials" CX_WARMUP="$case_warmup" + export CX_SAMPLES_PER_POINT="$((case_iters * case_trials))" + export CX_WARMUP_SEMANTICS="full-roundtrip-before-each-component-trial-point-v1" + cx_apply_network_profile "$NODES" "$transport" + cx_log "EP${NGPUS}[$ci] id=${case_id:-manual} $mode/$ph $CX_BENCH" + if [ "$ep" != "$NGPUS" ] || [ "$nodes" != "$NODES" ] || [ "$gpn" != "$GPN" ] \ + || [ "$domain" != "$SCALE_UP_DOMAIN" ]; then + cx_log "ERROR: EP${NGPUS}[$ci] topology mismatch (ep=$ep nodes=$nodes gpn=$gpn domain=$domain); skipping case" + failed_cases=$((failed_cases + 1)) + continue + fi + + workload_dir="" + if cx_bool_enabled "$canonical"; then + workload_dir=".cx_workloads/c$(printf '%03d' "$ci")" + workload_ladder="$ladder" + [ -n "$workload_ladder" ] \ + || workload_ladder="1 2 4 8 16 32 64 128 256 512 1024 2048 4096" + workload_args=(python3 bench/make_workloads.py --out-dir "$workload_dir" + --routing "$routing" --ep "$ep" --hidden "$hidden" --topk "$topk" + --experts "$experts" --seed "${CX_SEED:-67}" --tokens-ladder "$workload_ladder") + workload_log="$(cx_private_log_path "workload-c$(printf '%03d' "$ci")")" + set +e + srun --jobid="$JOB_ID" --nodes=1 --ntasks=1 --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" \ + --export="$(cx_container_exports)" "${workload_args[@]}" \ + "$workload_log" 2>&1 + stage_rc=$? + set -e + if [ "$stage_rc" != 0 ]; then + cx_log "ERROR: EP${NGPUS}[$ci] workload staging failed (rc=$stage_rc)" + failed_cases=$((failed_cases + 1)) + continue + fi + fi + + ep_args=(--backend "$CX_BENCH" --mode "$mode" --phase "$ph" --routing "$routing" + --gpus-per-node "$gpn" --scale-up-domain "$domain" --scope "$scope" + --scale-up-transport "$scale_up_transport" --scale-out-transport "$scale_out_transport" + --tokens-ladder "$ladder" --hidden "$hidden" --topk "$topk" --experts "$experts" + --warmup "$case_warmup" --iters "$case_iters" --trials "$case_trials" + --seed "${CX_SEED:-67}" --runner "$RUNNER" --topology-class "$topology_class" + --transport "$transport" --case-id "$case_id" --suite "$suite" + --workload-name "$workload" + --qualification-index "${CX_QUALIFICATION_INDEX:-1}" --version "${CX_VERSION:-1}") + cx_bool_enabled "$eplb" && ep_args+=(--eplb) + [ -z "$workload_dir" ] || ep_args+=(--workload-dir "$workload_dir") + export CX_ATTEMPT_ID=1 + attempt_tag=a01 + out="results/${case_stem}_${attempt_tag}.json" + runtime_log="$(cx_private_log_path "runtime-c$(printf '%03d' "$ci")-$attempt_tag")" + set +e + timeout -k 30 "${CX_RUN_TIMEOUT:-900}" srun --jobid="$JOB_ID" --nodes="$NODES" \ + --ntasks="$NGPUS" --ntasks-per-node="$GPN" --chdir=/tmp \ + --container-name="$container_name" --container-image="$SQUASH_FILE" \ + "${container_args[@]}" \ + --export="$(cx_container_exports)" \ + bash -c "$WRAP" _ "${ep_args[@]}" --out "$out" \ + "$runtime_log" 2>&1 + run_rc=$? + set -e + # Terminal-outcome emission and result-document gating were removed with + # contracts.py; a case now counts as run purely on the distributed command's + # return code. The rank-zero result the harness wrote (if any) is left in place + # for the summary renderer, which validates nothing. + if [ "$run_rc" != 0 ]; then + cx_fail_stage execution "$runtime_log" || true + failed_cases=$((failed_cases + 1)) + fi + done < "$cases_file" + rm -f "$cases_file" + [ "$ci" -eq "$expected_cases" ] \ + || cx_die "enumerated $expected_cases cases but executed $ci" + if [ "$failed_cases" -ne 0 ]; then + summary_log="$(cx_private_log_path shard-summary)" + printf 'SHARD done: %s/%s case(s) failed\n' "$failed_cases" "$expected_cases" \ + > "$summary_log" + cx_fail_stage execution "$summary_log" || true + return 1 + fi + return 0 +} + +# Remove this allocation's persistent pyxis container before the allocation is +# released. The cluster runs pyxis with container_scope=global, so a named +# --container-writable container (the distributed path's cxep_) survives +# job teardown and its unpacked rootfs — tens of GB per node — would otherwise +# accumulate on every allocated node's local image store until it fills and the +# next writable extraction fails with ENOSPC. Best-effort and bounded: teardown +# must never hang or fail on this. Single-node legs use an unnamed, ephemeral +# container that pyxis reclaims on its own, so only NODES>1 needs removal. +cx_remove_distributed_container() { + local job_id="$1" nodes="${2:-1}" + [ -n "$job_id" ] || return 0 + [ "$nodes" -gt 1 ] 2>/dev/null || return 0 + timeout 120 srun --jobid="$job_id" --nodes="$nodes" --ntasks-per-node=1 \ + --chdir=/tmp enroot remove -f "pyxis_cxep_${job_id}" \ + /dev/null 2>&1 || true +} + +cx_launcher_cleanup() { + local rc="$1" stage_root="${MOUNT_SRC:-}" source_root allocation_stopped=1 + source_root="${stage_root:-${REPO_ROOT:-}}" + trap - EXIT HUP INT TERM + if [ -n "${COLLECTIVEX_EPHEMERAL_CONFIG_PATH:-}" ]; then + rm -f -- "$COLLECTIVEX_EPHEMERAL_CONFIG_PATH" >/dev/null 2>&1 || true + unset COLLECTIVEX_EPHEMERAL_CONFIG_PATH + fi + if [ -n "${JOB_ID:-}" ]; then + cx_remove_distributed_container "$JOB_ID" "${NODES:-1}" + if ! cx_cancel_job "$JOB_ID"; then + allocation_stopped=0 + [ "$rc" != 0 ] || rc=1 + elif ! cx_clear_allocation_jobid; then + allocation_stopped=0 + [ "$rc" != 0 ] || rc=1 + fi + elif [ "${CX_ALLOCATION_UNCERTAIN:-0}" = 1 ]; then + allocation_stopped=0 + [ "$rc" != 0 ] || rc=1 + fi + if [ "$allocation_stopped" = 1 ]; then + cx_write_cleanup_guard safe || true + else + cx_write_cleanup_guard unsafe || true + fi + [ "$allocation_stopped" = 1 ] || source_root="${REPO_ROOT:-$source_root}" + if [ "$rc" != 0 ] \ + && [ -n "${REPO_ROOT:-}" ] && [ -n "${CX_BENCH:-}" ]; then + cx_log "ERROR: terminal-failure-class=${CX_FAILSAFE_MODE:-setup}" + [ -d "$source_root/experimental/CollectiveX" ] || source_root="$REPO_ROOT" + [ "$source_root" = "$REPO_ROOT" ] \ + || cx_collect_results "$source_root" "$REPO_ROOT" || true + fi + if [ "$allocation_stopped" = 1 ] && [ -n "${REPO_ROOT:-}" ] \ + && [ -n "$stage_root" ] && [ "$stage_root" != "$REPO_ROOT" ]; then + if ! cx_cleanup_stage "$stage_root" "$REPO_ROOT"; then + [ "$rc" != 0 ] || rc=1 + fi + fi + [ "${COLLECTIVEX_CANONICAL_GHA:-0}" = 1 ] || cx_cleanup_private_logs "$rc" + exit "$rc" +} + +cx_install_launcher_fail_safe() { + CX_ALLOCATION_UNCERTAIN=0 + trap 'cx_launcher_cleanup "$?"' EXIT + trap 'cx_launcher_cleanup 129' HUP + trap 'cx_launcher_cleanup 130' INT + trap 'cx_launcher_cleanup 143' TERM +} diff --git a/experimental/CollectiveX/runtime/run_in_container.sh b/experimental/CollectiveX/runtime/run_in_container.sh new file mode 100644 index 0000000000..b2d3f01982 --- /dev/null +++ b/experimental/CollectiveX/runtime/run_in_container.sh @@ -0,0 +1,1074 @@ +#!/usr/bin/env bash +# CollectiveX — generic in-container benchmark dispatcher (single-node). +# +# Runs INSIDE the container under `srun` for single-node shards. The GB EP8 launcher invokes +# run_ep.py directly across nodes. The SKU adapter handles allocation/container/transport-env; +# this script selects one EP backend from CX_BENCH and writes result JSON under results/. +# +# Required env (exported by the adapter): CX_RUNNER CX_NGPUS CX_TS CX_TOPO +# Selector: CX_BENCH = deepep | deepep-v2 | mori | uccl | nccl-ep | deepep-hybrid +# EP knobs passed to bench/run_ep.py: +# CX_PHASE = decode | prefill | both (default decode) <- picks the token sweep +# CX_TOKENS_LADDER (space/comma sep; blank = phase default) +# CX_HIDDEN CX_TOPK CX_EXPERTS CX_ROUTING CX_SEED CX_ITERS +set -euo pipefail + +cd /ix/experimental/CollectiveX +# shellcheck source=../runtime/common.sh +source runtime/common.sh +mkdir -p results +cx_write_runtime_stage backend-setup || cx_die "cannot record runtime stage" + +: "${CX_RUNNER:?CX_RUNNER not set}" +: "${CX_NGPUS:?CX_NGPUS not set}" +: "${CX_TS:?CX_TS not set}" +: "${CX_TOPO:?CX_TOPO not set}" +CX_BENCH="${CX_BENCH:-deepep}" +CX_TRANSPORT="${CX_TRANSPORT:-}" + +cx_apply_timing_profile + +cx_log "in-container: runner=$CX_RUNNER ngpus=$CX_NGPUS bench=$CX_BENCH topo=$CX_TOPO" + +# Blank ladders use the phase default in bench/run_ep.py. +cx_ep_ladder() { + printf '%s' "${CX_TOKENS_LADDER:-}" +} + +# Canonical workload staging. Every SKU/backend generates identical canonical array bytes and +# content IDs in-container; the NPZ container bytes themselves are not an identity boundary. When CX_CANONICAL=1 +# (and CX_WORKLOAD_DIR not already provided) we generate routing traces for the run's ladder +# into a NON-results dir (.cx_workloads/ — so the *.manifest.json never pollute the results glob) and +# point run_ep at it. Raw attempts remain diagnostic until the publisher validates full coverage. +cx_stage_canonical() { + cx_bool_enabled "${CX_CANONICAL:-0}" || return 0 + [ -n "${CX_WORKLOAD_DIR:-}" ] && return 0 + local dir="$PWD/.cx_workloads" + local ladder; ladder="$(cx_ep_ladder)" + # cover both phase ladders when none is given, so either phase finds its files. + [ -z "$ladder" ] && ladder="1 2 4 8 16 32 64 128 256 512 1024 2048 4096" + cx_log "staging canonical workloads (routing=${CX_ROUTING:-uniform} ep=$CX_NGPUS ladder='$ladder')" + python3 bench/make_workloads.py --out-dir "$dir" --routing "${CX_ROUTING:-uniform}" \ + --ep "$CX_NGPUS" --hidden "${CX_HIDDEN:-7168}" --topk "${CX_TOPK:-8}" \ + --experts "${CX_EXPERTS:-256}" --seed "${CX_SEED:-67}" --tokens-ladder "$ladder" \ + || { cx_log "ERROR: canonical workload staging failed"; return 1; } + export CX_WORKLOAD_DIR="$dir" + cx_log "canonical workloads staged at $dir" +} + +# run_ep_suite +# One bench/run_ep.py invocation per phase (decode/prefill/both); dispatch and +# combine are timed separately inside it. One JSON per (backend, phase). +run_ep_suite() { + local backend="$1" phase phases ladder failure_kind rc=0 rc_run + ladder="$(cx_ep_ladder)" + phases="${CX_PHASE:-decode}" + [ "$phases" = "both" ] && phases="decode prefill" + if ! cx_stage_canonical; then + cx_log "ERROR: $backend canonical workload staging failed" + return 1 + fi + for phase in $phases; do + cx_log "ep backend=$backend phase=$phase ngpus=$CX_NGPUS ladder='${ladder:-}'" + local out="results/${CX_RUNNER}_${backend}_${phase}_${CX_TS}.json" + local -a EPARGS=(--backend "$backend" --mode "${CX_MODE:-normal}" --phase "$phase" + --tokens-ladder "$ladder" + --hidden "${CX_HIDDEN:-7168}" --topk "${CX_TOPK:-8}" --experts "${CX_EXPERTS:-256}" + --routing "${CX_ROUTING:-uniform}" --seed "${CX_SEED:-67}" --iters "${CX_ITERS:-8}" + --trials "${CX_TRIALS:-64}" --warmup "${CX_WARMUP:-32}" + --gpus-per-node "${CX_GPUS_PER_NODE:-0}" --scale-up-domain "${CX_SCALE_UP_DOMAIN:-0}" + --scope "${CX_SCOPE:-scale-up}" --scale-up-transport "${CX_SCALE_UP_TRANSPORT:-unknown}" + --scale-out-transport "${CX_SCALE_OUT_TRANSPORT:-}" + --case-id "${CX_CASE_ID:-}" --suite "${CX_SUITE:-}" --workload-name "${CX_WORKLOAD_NAME:-}" + --qualification-index "${CX_QUALIFICATION_INDEX:-1}" --version "${CX_VERSION:-1}" + --runner "$CX_RUNNER" --topology-class "$CX_TOPO" --transport "$CX_TRANSPORT" + --out "$out") + [ -n "${CX_WORKLOAD_DIR:-}" ] && EPARGS+=(--workload-dir "$CX_WORKLOAD_DIR") + cx_write_runtime_stage execution || cx_die "cannot record runtime stage" + if timeout -k 30 "${CX_RUN_TIMEOUT:-900}" \ + torchrun --nproc_per_node="$CX_NGPUS" bench/run_ep.py "${EPARGS[@]}"; then + rc_run=0 + else + rc_run=$? + fi + # Result-document gating and terminal-outcome emission were removed with + # contracts.py; success is now the run_ep.py return code alone. Any result + # document it wrote is left in place for the summary renderer, which validates + # nothing. + if [ "$rc_run" != 0 ]; then + failure_kind=failed + [ "$rc_run" != 124 ] && [ "$rc_run" != 137 ] || failure_kind="timed out" + if [ "$failure_kind" = "timed out" ]; then + cx_log "WARN: $backend $phase run timed out rc=$rc_run (limit=${CX_RUN_TIMEOUT:-900}s)" + else + cx_log "WARN: $backend $phase run failed rc=$rc_run" + fi + rc=1 + fi + done + return "$rc" +} + +# Resolve and verify the actual CUDA target before compiling source kernels. +cx_cuda_arch() { + local expected detected + case "$CX_RUNNER" in + h100*|h200*) expected="9.0" ;; + b200*|gb200*) expected="10.0" ;; + b300*|gb300*) expected="10.3" ;; + *) cx_log "ERROR: no CUDA target registered for $CX_RUNNER"; return 1 ;; + esac + detected="$(python3 - <<'PY' +import torch + +major, minor = torch.cuda.get_device_capability() +print(f"{major}.{minor}") +PY +)" || return 1 + [ "$detected" = "$expected" ] || { + cx_log "ERROR: $CX_RUNNER expected CUDA target $expected, detected $detected" + return 1 + } + printf '%s' "$detected" +} + +cx_nvidia_package_root() { + local package="$1" component="$2" + python3 - "$package" "$component" <<'PY' +from importlib import metadata +from pathlib import Path, PurePosixPath +import sys + +package, component = sys.argv[1:] +try: + distribution = metadata.distribution(package) + prefix = f"nvidia/{component}/" + entries = [str(entry).replace("\\", "/") for entry in distribution.files or ()] + if not any(entry.startswith(prefix) for entry in entries): + raise ValueError + root = Path(distribution.locate_file(PurePosixPath("nvidia") / component)).resolve() + if not root.is_dir(): + raise ValueError +except (metadata.PackageNotFoundError, OSError, TypeError, ValueError): + raise SystemExit(1) +print(root, end="") +PY +} + +cx_prepare_cuda_cccl() { + local cccl="" candidate cuda_home nvcc + nvcc="$(command -v nvcc)" \ + || { cx_log "ERROR: CUDA nvcc is unavailable"; return 1; } + nvcc="$(readlink -f -- "$nvcc")" \ + || { cx_log "ERROR: CUDA nvcc cannot be resolved"; return 1; } + case "$nvcc" in + */bin/nvcc) cuda_home="${nvcc%/bin/nvcc}" ;; + *) cx_log "ERROR: CUDA nvcc has an unexpected path"; return 1 ;; + esac + [ -x "$cuda_home/bin/nvcc" ] && [ -d "$cuda_home/include" ] \ + && [ -d "$cuda_home/lib64" ] \ + || { cx_log "ERROR: CUDA toolkit root is incomplete"; return 1; } + for candidate in "$cuda_home"/targets/*/include/cccl; do + if [ -d "$candidate" ]; then + cccl="$candidate" + break + fi + done + [ -n "$cccl" ] || { cx_log "ERROR: CUDA CCCL headers are unavailable"; return 1; } + export CUDA_HOME="$cuda_home" CX_CUDA_CCCL="$cccl" + export CPATH="$cccl:${CPATH:-}" + export NVCC_PREPEND_FLAGS="-I$cccl ${NVCC_PREPEND_FLAGS:-}" +} + +cx_prepare_deepep_toolchain() { + local packaged overlay path root temporary + packaged="$(cx_nvidia_package_root nvidia-nvshmem-cu12 nvshmem)" \ + || { cx_log "ERROR: nvidia.nvshmem is unavailable"; return 1; } + root="$(cx_deepep_v2_root)" || return 1 + overlay="$root/nvshmem-overlay" + if ! ( + umask 077 + exec 8>"$root/nvshmem-overlay.lock" || exit 1 + flock 8 || exit 1 + if [ ! -d "$overlay" ]; then + temporary="$root/.nvshmem-overlay.$$" + rm -rf "$temporary" || exit 1 + mkdir -p "$temporary/lib" || exit 1 + ln -s "$packaged/include" "$temporary/include" || exit 1 + for path in "$packaged"/lib/*; do + ln -s "$path" "$temporary/lib/${path##*/}" || exit 1 + done + [ ! -e "$packaged/lib/libnvshmem_host.so.3" ] \ + || ln -sf "$packaged/lib/libnvshmem_host.so.3" \ + "$temporary/lib/libnvshmem_host.so" || exit 1 + mv "$temporary" "$overlay" || exit 1 + fi + [ ! -L "$overlay" ] \ + && [ "$(readlink -f "$overlay/include")" = "$(readlink -f "$packaged/include")" ] \ + && [ -e "$overlay/lib/libnvshmem_host.so" ] \ + && [ -e "$overlay/lib/libnvshmem_device.a" ] + ); then + cx_log "ERROR: DeepEP V2 NVSHMEM overlay is invalid" + return 1 + fi + NVSHMEM_DIR="$overlay" + export NVSHMEM_DIR + cx_prepare_cuda_cccl || return 1 + export LD_LIBRARY_PATH="$NVSHMEM_DIR/lib:${LD_LIBRARY_PATH:-}" +} + +cx_probe_deepep() { + local expected_record_sha256 expected_version expected_wheel_sha256 + if [ "${COLLECTIVEX_IMAGE:-}" != "$CX_IMAGE_MULTIARCH" ] \ + || [ "${COLLECTIVEX_IMAGE_DIGEST:-}" != "$CX_IMAGE_MULTIARCH_DIGEST" ] \ + || [ "${COLLECTIVEX_IMAGE_DIGEST_VERIFIED:-0}" != 1 ]; then + cx_log "ERROR: DeepEP V1 requires the exact pinned multi-architecture image" + return 1 + fi + cx_cuda_arch >/dev/null || return 1 + case "$CX_RUNNER" in + gb200|gb300) + expected_version="1.1.0+814e508" + expected_wheel_sha256="784dabec0877b6cf72619b7e93eda7e2f365648487bd37fc3ff6960e53669313" + expected_record_sha256="2671cff7baf8c2c214ff4bac721af875d513130670bec57601998bd1aae82882" + DEEPEP_COMMIT="814e508537c6ffc775d59f6f1b9ba43f3a65968c" + ;; + *) + expected_version="1.2.1" + expected_wheel_sha256="7c02c29306ea0fe2dd474618e72e0f310f260187a9c0700a656d2f6964e8c307" + expected_record_sha256="6548e9c504a12b2471af4b7f4d9546321210a57a456b5dc55bd4a8dad0f932ac" + DEEPEP_COMMIT="9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee" + ;; + esac + export DEEPEP_COMMIT + python3 - "$expected_version" "$expected_wheel_sha256" "$expected_record_sha256" <<'PY' || { +import base64 +import csv +import hashlib +import importlib.metadata as metadata +import io +import json +from pathlib import Path +import sys + +import deep_ep +from deep_ep import Buffer + +distribution = metadata.distribution("deep_ep") +assert distribution.version == sys.argv[1] +assert Buffer.__name__ == "Buffer" +recorded_files = { + Path(distribution.locate_file(entry)).resolve() for entry in distribution.files or () +} +buffer_module = sys.modules.get(Buffer.__module__) +assert Path(deep_ep.__file__).resolve() in recorded_files +assert buffer_module is not None and Path(buffer_module.__file__).resolve() in recorded_files +direct_url = json.loads(distribution.read_text("direct_url.json")) +assert direct_url["archive_info"]["hashes"]["sha256"] == sys.argv[2] +record_entry = next( + entry for entry in distribution.files or () + if str(entry).endswith(".dist-info/RECORD") +) +record = distribution.locate_file(record_entry).read_bytes() +assert hashlib.sha256(record).hexdigest() == sys.argv[3] +for path, encoded_digest, size in csv.reader(io.StringIO(record.decode())): + if not encoded_digest: + continue + algorithm, expected = encoded_digest.split("=", 1) + assert algorithm == "sha256" + payload = distribution.locate_file(path).read_bytes() + observed = base64.urlsafe_b64encode(hashlib.sha256(payload).digest()).decode().rstrip("=") + assert observed == expected + assert not size or len(payload) == int(size) +PY + cx_log "ERROR: container DeepEP build does not match its pinned image contract" + return 1 + } + cx_log "DeepEP image build ready ($DEEPEP_COMMIT)" +} + +# DeepEP V2 is PR #605's ElasticBuffer implementation with upstream PR #630's pure scale-up +# initialization fix and PR #640's exact libnccl mapping check. Canonical launchers stage the +# pinned source and mount a private cluster-local build cache at /cx-cache. +cx_deepep_v2_root() { + local arch cpu base identity key image_digest + arch="$(cx_cuda_arch)" || return 1 + cpu="$(uname -m)" + [[ "$cpu" =~ ^[A-Za-z0-9._-]+$ ]] || return 1 + base="${CX_BACKEND_CACHE_ROOT:-}" + [[ "$base" = /* ]] || return 1 + image_digest="${COLLECTIVEX_IMAGE_DIGEST:-manual-unverified}" + [[ "$image_digest" = manual-unverified || "$image_digest" =~ ^sha256:[0-9a-f]{64}$ ]] \ + || return 1 + # Bump the recipe generation whenever the build procedure changes. Benchmark-only + # source revisions must reuse the same immutable environment instead of leaking GBs. + identity="deepep-v2-cache-v3|$cpu|sm${arch/./}|image=$image_digest|recipe=aot-persistent-nvshmem-active-cuda-maxjobs16-v3|$CX_DEEPEP_V2_COMMIT|$CX_DEEPEP_V2_TREE|$CX_DEEPEP_V2_FMT_COMMIT|$CX_DEEPEP_V2_NCCL_CHECK_COMMIT|pip=26.1.2|setuptools=82.0.1|wheel=0.47.0|ninja=1.13.0|numpy=2.2.6|torch=2.10.0+cu130|nccl=2.30.4|nvshmem=3.3.9|max-jobs=16" + key="$(printf '%s' "$identity" | sha256sum | awk '{print $1}')" + [[ "$key" =~ ^[0-9a-f]{64}$ ]] || return 1 + printf '%s/deepep-v2-%s' "$base" "$key" +} + +cx_activate_deepep_v2() { + local root venv execution_id + root="$(cx_deepep_v2_root)" || return 1 + venv="$root/venv" + [ -x "$venv/bin/python" ] \ + || { cx_log "ERROR: DeepEP V2 venv interpreter is unavailable"; return 1; } + export VIRTUAL_ENV="$venv" + export PATH="$venv/bin:${PATH#"$venv/bin:"}" + EP_NCCL_ROOT_DIR="$(cx_nvidia_package_root nvidia-nccl-cu13 nccl)" \ + || { cx_log "ERROR: DeepEP V2 NCCL package root is unavailable"; return 1; } + EP_NVSHMEM_ROOT_DIR="$(cx_nvidia_package_root nvidia-nvshmem-cu12 nvshmem)" \ + || { cx_log "ERROR: DeepEP V2 NVSHMEM package root is unavailable"; return 1; } + export EP_NCCL_ROOT_DIR EP_NVSHMEM_ROOT_DIR + export LD_LIBRARY_PATH="$EP_NCCL_ROOT_DIR/lib:$EP_NVSHMEM_ROOT_DIR/lib:${LD_LIBRARY_PATH:-}" + execution_id="${COLLECTIVEX_EXECUTION_ID:-manual}" + [[ "$execution_id" =~ ^[A-Za-z0-9._-]+$ ]] \ + || { cx_log "ERROR: DeepEP V2 execution identity is invalid"; return 1; } + # JIT CUBINs are per-execution evidence and must be node-local. A shared NFS cache lets + # ranks on different nodes race the same compiler output and can trip compiler.hpp asserts. + # The identical absolute path still lets ranks on one node share their cold build. + export EP_JIT_CACHE_DIR="/tmp/collectivex-deepep-v2-jit-$execution_id" + export EP_REUSE_NCCL_COMM=1 + export DEEPEP_V2_PR=605 DEEPEP_V2_FIX_PR=630 DEEPEP_V2_NCCL_CHECK_FIX_PR=640 + DEEPEP_V2_COMMIT="$CX_DEEPEP_V2_COMMIT" + DEEPEP_V2_TREE="$CX_DEEPEP_V2_TREE" + DEEPEP_V2_FMT_COMMIT="$CX_DEEPEP_V2_FMT_COMMIT" + DEEPEP_V2_NCCL_CHECK_COMMIT="$CX_DEEPEP_V2_NCCL_CHECK_COMMIT" + export DEEPEP_V2_COMMIT DEEPEP_V2_TREE DEEPEP_V2_FMT_COMMIT + export DEEPEP_V2_NCCL_CHECK_COMMIT + [ ! -L "$EP_JIT_CACHE_DIR" ] \ + || { cx_log "ERROR: DeepEP V2 JIT cache path is unsafe"; return 1; } + if ! mkdir -p "$EP_JIT_CACHE_DIR" || ! chmod 700 "$EP_JIT_CACHE_DIR"; then + cx_log "ERROR: DeepEP V2 JIT cache is unavailable" + return 1 + fi + unset EP_SUPPRESS_NCCL_CHECK +} + +cx_enable_deepep_v2_jit_reproducibility() { + local seed="collectivex-deepep-v2-fa8a9b1" cccl + [ -n "${CUDA_HOME:-}" ] \ + || { cx_log "ERROR: active CUDA toolkit is unavailable"; return 1; } + cccl="${CX_CUDA_CCCL:-}" + case "$cccl" in + "$CUDA_HOME"/targets/*/include/cccl) ;; + *) cx_log "ERROR: CUDA CCCL headers differ from the active toolkit"; return 1 ;; + esac + [ -d "$cccl" ] || { cx_log "ERROR: CUDA CCCL headers are unavailable"; return 1; } + CPATH="$cccl" + NVCC_PREPEND_FLAGS="--frandom-seed=$seed -I$cccl" + DEEPEP_V2_JIT_RANDOM_SEED="$seed" + EP_JIT_DUMP_SASS=1 + unset EP_JIT_DEBUG EP_JIT_DUMP_ASM EP_JIT_DUMP_PTX EP_JIT_WITH_LINEINFO + unset EP_JIT_PTXAS_VERBOSE EP_JIT_PRINT_COMPILER_COMMAND EP_JIT_NVCC_COMPILER + unset EP_JIT_CPP_STANDARD EP_JIT_PTXAS_CHECK EP_GIN_GDAKI_DEBUG EP_NUM_TOPK_IDX_BITS + export CPATH DEEPEP_V2_JIT_RANDOM_SEED EP_JIT_DUMP_SASS NVCC_PREPEND_FLAGS +} + +cx_probe_deepep_v2() { + python3 - <<'PY' +import ctypes +import importlib.metadata as metadata +import inspect +import os + +import torch + +assert torch.__version__ == "2.10.0+cu130", torch.__version__ +assert metadata.version("nvidia-nccl-cu13") == "2.30.4" +assert metadata.version("nvidia-nvshmem-cu12") == "3.3.9" +assert metadata.version("numpy") == "2.2.6" + +import deep_ep +assert deep_ep.__version__ == "2.0.0", deep_ep.__version__ +assert metadata.version("deep_ep") in {"2.0.0+fa8a9b1", "2.0.0+local"}, metadata.version("deep_ep") +assert inspect.isclass(deep_ep.ElasticBuffer) +assert deep_ep.ElasticBuffer.__name__ == "ElasticBuffer" +assert os.environ.get("EP_SUPPRESS_NCCL_CHECK") is None +with open("/proc/self/maps", encoding="utf-8") as handle: + loaded_nccl = { + os.path.realpath(line.rstrip().split()[-1]) + for line in handle + if "libnccl.so" in line and os.path.isfile(line.rstrip().split()[-1]) + } +assert len(loaded_nccl) == 1 +runtime_version = ctypes.c_int() +assert ctypes.CDLL(loaded_nccl.pop()).ncclGetVersion(ctypes.byref(runtime_version)) == 0 +assert runtime_version.value == 23004, runtime_version.value +PY +} + +cx_deepep_v2_content_sha256() { + python3 - <<'PY' +import hashlib +from importlib import metadata +import os +from pathlib import Path, PurePosixPath +import stat + +distribution = metadata.distribution("deep_ep") +entries = sorted(distribution.files or (), key=lambda entry: entry.as_posix()) +if not entries: + raise SystemExit(1) +venv_path = Path(os.environ["VIRTUAL_ENV"]).absolute() +if venv_path.is_symlink() or not venv_path.is_dir(): + raise SystemExit(1) +venv = venv_path.resolve(strict=True) +digest = hashlib.sha256() +extension = False +for entry in entries: + relative = PurePosixPath(entry.as_posix()) + if ( + relative.is_absolute() + or ".." in relative.parts + or not relative.parts + or not ( + relative.parts[0] == "deep_ep" + or relative.parts[0].startswith("deep_ep-") + and relative.parts[0].endswith(".dist-info") + ) + ): + raise SystemExit(1) + path = Path(distribution.locate_file(entry)).absolute() + resolved = path.resolve(strict=True) + try: + path.relative_to(venv_path) + resolved.relative_to(venv) + except ValueError: + raise SystemExit(1) + parent = path.parent + while parent != venv_path: + if parent.is_symlink(): + raise SystemExit(1) + parent = parent.parent + item = os.lstat(path) + if not stat.S_ISREG(item.st_mode): + raise SystemExit(1) + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + try: + opened = os.fstat(descriptor) + if (opened.st_dev, opened.st_ino) != (item.st_dev, item.st_ino): + raise SystemExit(1) + file_digest = hashlib.sha256() + while chunk := os.read(descriptor, 1024 * 1024): + file_digest.update(chunk) + finally: + os.close(descriptor) + name = relative.as_posix() + extension |= name.startswith("deep_ep/") and name.endswith(".so") + digest.update(name.encode()) + digest.update(b"\0") + digest.update(str(item.st_size).encode()) + digest.update(b"\0") + digest.update(file_digest.digest()) +if not extension: + raise SystemExit(1) +print(digest.hexdigest(), end="") +PY +} + +cx_deepep_v2_marker_content_sha256() { + local root="$1" marker="$2" revision="$3" tree="$4" fmt_revision="$5" cache_key="$6" + python3 - "$root" "$marker" "$revision" "$tree" "$fmt_revision" "$cache_key" <<'PY' +import os +import re +import stat +import sys + +root, marker, revision, tree, fmt_revision, cache_key = sys.argv[1:] +try: + root_item = os.lstat(root) + marker_item = os.lstat(marker) + children = [os.lstat(os.path.join(root, name)) for name in ("source", "venv")] + if ( + not stat.S_ISDIR(root_item.st_mode) + or stat.S_IMODE(root_item.st_mode) & 0o777 != 0o700 + or not stat.S_ISREG(marker_item.st_mode) + or marker_item.st_uid != root_item.st_uid + or stat.S_IMODE(marker_item.st_mode) & 0o777 != 0o600 + or marker_item.st_size > 1024 + or any( + not stat.S_ISDIR(child.st_mode) + or child.st_uid != root_item.st_uid + or stat.S_IMODE(child.st_mode) & 0o022 + for child in children + ) + ): + raise OSError + descriptor = os.open(marker, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + try: + opened = os.fstat(descriptor) + if (opened.st_dev, opened.st_ino) != (marker_item.st_dev, marker_item.st_ino): + raise OSError + payload = os.read(descriptor, 1025) + finally: + os.close(descriptor) + lines = payload.decode("ascii").splitlines() + if lines[:4] != [revision, tree, fmt_revision, cache_key] or len(lines) != 5: + raise ValueError + if not re.fullmatch(r"[0-9a-f]{64}", lines[4]): + raise ValueError +except (OSError, UnicodeError, ValueError): + raise SystemExit(1) +print(lines[4], end="") +PY +} + +cx_deepep_v2_cache_is_valid() { + local root="$1" marker="$2" revision="$3" tree="$4" fmt_revision="$5" cache_key="$6" + local expected_content actual_content + expected_content="$( + cx_deepep_v2_marker_content_sha256 \ + "$root" "$marker" "$revision" "$tree" "$fmt_revision" "$cache_key" + )" || return 1 + [ -d "$root/source" ] && [ ! -L "$root/source" ] \ + && [ "$(cx_git_in_tree "$root/source" rev-parse 'HEAD^{tree}' 2>/dev/null)" = "$tree" ] \ + && [ "$(cx_git_in_tree "$root/source/third-party/fmt" rev-parse HEAD 2>/dev/null)" = "$fmt_revision" ] \ + || return 1 + cx_activate_deepep_v2 || return 1 + actual_content="$(cx_deepep_v2_content_sha256)" || return 1 + [ "$actual_content" = "$expected_content" ] +} + +cx_build_deepep_v2() { + local root venv source marker marker_tmp lock_path arch cache_key cache_ready content_sha256 + local revision="fa8a9b16898204afd347c663b89e65ef87dc6ce6" + local tree="29809e75c5874e6609dac4804e7b651d5226959f" + local fmt_revision="a4c7e17133ee9cb6a2f45545f6e974dd3c393efa" + cx_verify_backend_cache_mount \ + || { cx_log "ERROR: DeepEP V2 cache mount identity validation failed"; return 1; } + arch="$(cx_cuda_arch)" || return 1 + root="$(cx_deepep_v2_root)" || return 1 + cache_key="${root##*/deepep-v2-}" + [[ "$cache_key" =~ ^[0-9a-f]{64}$ ]] || return 1 + venv="$root/venv"; source="$root/source"; marker="$root/.collectivex-complete" + lock_path="${root}.lock" + command -v flock >/dev/null || { cx_log "ERROR: flock is required for DeepEP V2"; return 1; } + mkdir -p "${root%/*}" || return 1 + cx_log "DeepEP V2: preparing PR #605 with upstream PR #630 and #640 fixes ($revision)" + if ! ( + [ ! -L "$lock_path" ] \ + || { cx_log "ERROR: DeepEP V2 cache lock is unsafe"; exit 1; } + (umask 077; : >> "$lock_path") && chmod 600 "$lock_path" \ + || { cx_log "ERROR: DeepEP V2 cache-lock-create failed"; exit 1; } + exec 9<>"$lock_path" \ + || { cx_log "ERROR: DeepEP V2 cache-lock-open failed"; exit 1; } + flock 9 \ + || { cx_log "ERROR: DeepEP V2 cache-lock-acquire failed"; exit 1; } + cache_ready=0 + if [ -e "$marker" ] || [ -L "$marker" ]; then + if ( + cx_deepep_v2_cache_is_valid \ + "$root" "$marker" "$revision" "$tree" "$fmt_revision" "$cache_key" + ); then + cache_ready=1 + else + cx_log "ERROR: published DeepEP V2 cache failed integrity validation; refusing reset" + exit 1 + fi + fi + if [ "$cache_ready" != 1 ]; then + if [ -e "$root" ] || [ -L "$root" ]; then + rm -rf "$root" \ + || { cx_log "ERROR: incomplete DeepEP V2 cache-reset failed"; exit 1; } + fi + mkdir -m 700 "$root" \ + || { cx_log "ERROR: DeepEP V2 cache-create failed"; exit 1; } + python3 -m venv "$venv" \ + || { cx_log "ERROR: DeepEP V2 venv creation failed"; exit 1; } + "$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input \ + "pip==26.1.2" "setuptools==82.0.1" "wheel==0.47.0" "ninja==1.13.0" \ + "numpy==2.2.6" "nvidia-nvshmem-cu12==3.3.9" >&2 2>&1 \ + || { cx_log "ERROR: DeepEP V2 build-tool installation failed"; exit 1; } + "$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input \ + --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple "torch==2.10.0" >&2 2>&1 \ + || { cx_log "ERROR: torch 2.10.0+cu130 installation failed"; exit 1; } + # Torch pins NCCL 2.28.9; the PR #605 ElasticBuffer implementation requires 2.30.4. + "$venv/bin/python" -m pip install -q --disable-pip-version-check --no-input \ + --force-reinstall --no-deps "nvidia-nccl-cu13==2.30.4" >&2 2>&1 \ + || { cx_log "ERROR: NCCL 2.30.4 installation failed"; exit 1; } + cx_activate_deepep_v2 \ + || { cx_log "ERROR: DeepEP V2 environment activation failed"; exit 1; } + cx_prepare_deepep_toolchain \ + || { cx_log "ERROR: DeepEP V2 toolchain preparation failed"; exit 1; } + EP_NVSHMEM_ROOT_DIR="$NVSHMEM_DIR" + export EP_NVSHMEM_ROOT_DIR + cx_materialize_backend_source deepep-v2 "$source" \ + || { cx_log "ERROR: DeepEP V2 staged source is invalid"; exit 1; } + (cd "$source" && SOURCE_DATE_EPOCH="$(cx_git_in_tree "$source" show -s --format=%ct HEAD)" \ + TORCH_CUDA_ARCH_LIST="$arch" MAX_JOBS=16 \ + python3 -m pip install -q --no-build-isolation --no-deps --force-reinstall .) >&2 2>&1 \ + || { cx_log "ERROR: DeepEP V2 build failed"; exit 1; } + cx_probe_deepep_v2 \ + || { cx_log "ERROR: DeepEP V2 ElasticBuffer/runtime probe failed"; exit 1; } + content_sha256="$(cx_deepep_v2_content_sha256)" \ + || { cx_log "ERROR: DeepEP V2 installed-content hashing failed"; exit 1; } + marker_tmp="$(mktemp "$root/.collectivex-complete.tmp.XXXXXX")" \ + || { cx_log "ERROR: DeepEP V2 cache-marker-create failed"; exit 1; } + chmod 600 "$marker_tmp" \ + || { cx_log "ERROR: DeepEP V2 cache-marker-permission failed"; exit 1; } + printf '%s\n%s\n%s\n%s\n%s\n' \ + "$revision" "$tree" "$fmt_revision" "$cache_key" "$content_sha256" > "$marker_tmp" \ + || { cx_log "ERROR: DeepEP V2 cache-marker-write failed"; exit 1; } + mv -f -- "$marker_tmp" "$marker" \ + || { cx_log "ERROR: DeepEP V2 cache-marker-publish failed"; exit 1; } + fi + cx_deepep_v2_cache_is_valid \ + "$root" "$marker" "$revision" "$tree" "$fmt_revision" "$cache_key" \ + || { cx_log "ERROR: DeepEP V2 cache validation failed"; exit 1; } + ); then + cx_log "ERROR: shared DeepEP V2 environment is incomplete" + return 1 + fi + cx_activate_deepep_v2 || return 1 + cx_prepare_deepep_toolchain || return 1 + cx_enable_deepep_v2_jit_reproducibility || return 1 + EP_NVSHMEM_ROOT_DIR="$NVSHMEM_DIR" + export EP_NVSHMEM_ROOT_DIR + cx_probe_deepep_v2 || { cx_log "ERROR: DeepEP V2 shared runtime probe failed"; return 1; } + cx_log "DeepEP V2 ready ($DEEPEP_V2_COMMIT, ElasticBuffer, NCCL Device API; LSA/Gin selected by adapter)" +} + +# Build the pinned DeepEP `hybrid-ep` implementation. MNNVL remains one scale-up +# domain; true x86 scale-out uses the upstream DOCA/RDMA build explicitly. +cx_configure_deepep_hybrid_build() { + local interface device rdma_name + local -a interfaces devices + unset HYBRID_EP_MULTINODE USE_NIXL RDMA_CORE_HOME DEEPEP_HYBRID_BUILD_MODE + if [ "${CX_NODES:-1}" -le 1 ] || [ "${CX_TRANSPORT:-}" = mnnvl ]; then + export DEEPEP_HYBRID_BUILD_MODE=intradomain + return 0 + fi + [ "$(uname -m)" = x86_64 ] \ + || { cx_log "ERROR: hybrid-ep scale-out is registered only on x86_64"; return 1; } + [ -n "${GLOO_SOCKET_IFNAME:-}" ] && [ -n "${NCCL_IB_HCA:-}" ] \ + || { cx_log "ERROR: hybrid-ep scale-out network selectors are unavailable"; return 1; } + IFS=, read -r -a interfaces <<< "$GLOO_SOCKET_IFNAME" + for interface in "${interfaces[@]}"; do + [ -d "/sys/class/net/$interface" ] \ + || { cx_log "ERROR: configured hybrid-ep socket interface is absent"; return 1; } + done + IFS=, read -r -a devices <<< "$NCCL_IB_HCA" + for device in "${devices[@]}"; do + rdma_name="$(cx_nccl_hca_device_name "$device")" + [ -d "/sys/class/infiniband/$rdma_name" ] \ + || { cx_log "ERROR: configured hybrid-ep RDMA device is absent"; return 1; } + done + command -v make >/dev/null \ + || { cx_log "ERROR: make is required for hybrid-ep scale-out"; return 1; } + [ -r /usr/include/infiniband/verbs.h ] && [ -r /usr/include/infiniband/mlx5dv.h ] \ + || { cx_log "ERROR: pinned hybrid-ep RDMA headers are unavailable"; return 1; } + python3 - <<'PY' >/dev/null 2>&1 || { +import ctypes.util +import sys +sys.exit(0 if all(ctypes.util.find_library(name) for name in ("ibverbs", "mlx5")) else 1) +PY + cx_log "ERROR: pinned hybrid-ep RDMA libraries are unavailable" + return 1 + } + export HYBRID_EP_MULTINODE=1 USE_NIXL=0 RDMA_CORE_HOME=/usr + export DEEPEP_HYBRID_BUILD_MODE=multinode-doca +} + +cx_deepep_hybrid_marker_content_sha256() { + python3 - "$1" "$2" "$3" "$4" "${5:-}" <<'PY' +import os +import re +import stat +import sys + +root, marker, revision, tree, build_mode = sys.argv[1:] +try: + root_item = os.lstat(root) + marker_item = os.lstat(marker) + if ( + not stat.S_ISDIR(root_item.st_mode) + or stat.S_IMODE(root_item.st_mode) & 0o777 != 0o700 + or not stat.S_ISREG(marker_item.st_mode) + or marker_item.st_uid != root_item.st_uid + or stat.S_IMODE(marker_item.st_mode) & 0o777 != 0o600 + or marker_item.st_size > 512 + ): + raise OSError + descriptor = os.open(marker, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + try: + opened = os.fstat(descriptor) + if (opened.st_dev, opened.st_ino) != (marker_item.st_dev, marker_item.st_ino): + raise OSError + payload = os.read(descriptor, 513) + finally: + os.close(descriptor) + lines = payload.decode("ascii").splitlines() + expected = [revision, tree, build_mode] if build_mode else [revision, tree] + if len(lines) != len(expected) + 1 or lines[:-1] != expected: + raise ValueError + if not re.fullmatch(r"[0-9a-f]{64}", lines[-1]): + raise ValueError +except (OSError, UnicodeError, ValueError): + raise SystemExit(1) +print(lines[-1], end="") +PY +} + +cx_deepep_hybrid_cache_is_valid() { + local root="$1" marker="$2" revision="$3" tree="$4" build_mode="${5:-}" + local expected actual status extra + expected="$(cx_deepep_hybrid_marker_content_sha256 \ + "$root" "$marker" "$revision" "$tree" "$build_mode")" || return 1 + [ "$(cx_git_in_tree "$root" rev-parse HEAD 2>/dev/null)" = "$revision" ] \ + && [ "$(cx_git_in_tree "$root" rev-parse 'HEAD^{tree}' 2>/dev/null)" = "$tree" ] \ + || return 1 + status="$(cx_git_in_tree "$root" status --porcelain --untracked-files=no \ + --ignore-submodules=none 2>/dev/null)" || return 1 + [ -z "$status" ] || return 1 + extra="$(cx_git_in_tree "$root" ls-files --others --exclude-standard -- \ + 'deep_ep/*.py' 'deep_ep/*.so' 2>/dev/null)" || return 1 + [ -z "$extra" ] || return 1 + extra="$(cx_git_in_tree "$root" ls-files --others --ignored --exclude-standard -- \ + 'deep_ep/*.py' 'deep_ep/*.so' 2>/dev/null)" || return 1 + [ -z "$extra" ] || return 1 + actual="$(cx_extension_pair_sha256 "$root" 'deep_ep_cpp*.so' 'hybrid_ep_cpp*.so')" \ + || return 1 + [ "$actual" = "$expected" ] +} + +cx_build_deepep_hybrid() { + local arch revision="$CX_DEEPEP_HYBRID_COMMIT" tree="$CX_DEEPEP_HYBRID_TREE" + local build_root marker marker_tmp lock_path content_sha256 cache_ready build_mode + export DEEPEP_COMMIT="$revision" DEEPEP_TREE="$tree" + arch="$(cx_cuda_arch)" || return 1 + cx_configure_deepep_hybrid_build || return 1 + build_mode="$DEEPEP_HYBRID_BUILD_MODE" + build_root="$PWD/.cx_backend/deepep-hybrid-${arch/./}-${build_mode}" + marker="$build_root/.collectivex-complete" + lock_path="${build_root}.lock" + cx_log "DeepEP hybrid-ep: building $revision for CUDA target $arch" + unset NVSHMEM_DIR + cx_prepare_cuda_cccl || return 1 + command -v flock >/dev/null || { cx_log "ERROR: flock is required for hybrid-ep"; return 1; } + mkdir -p "$PWD/.cx_backend" || return 1 + if ! ( + [ ! -L "$lock_path" ] || exit 1 + (umask 077; : >> "$lock_path") && chmod 600 "$lock_path" || exit 1 + exec 9<>"$lock_path" || exit 1 + flock 9 || exit 1 + cache_ready=0 + if [ -e "$marker" ] || [ -L "$marker" ]; then + cx_deepep_hybrid_cache_is_valid \ + "$build_root" "$marker" "$revision" "$tree" "$build_mode" \ + || exit 1 + cache_ready=1 + fi + if [ "$cache_ready" != 1 ]; then + cx_materialize_backend_source deepep-hybrid "$build_root" \ + || { cx_log "ERROR: hybrid-ep staged source is invalid"; exit 1; } + if [ "$build_mode" = multinode-doca ]; then + [ "$(cx_git_in_tree "$build_root/third-party/nccl" rev-parse HEAD 2>/dev/null)" \ + = "$CX_DEEPEP_HYBRID_NCCL_COMMIT" ] \ + || { cx_log "ERROR: pinned hybrid-ep NCCL transport source is absent"; exit 1; } + fi + (cd "$build_root" && \ + SOURCE_DATE_EPOCH="$(cx_git_in_tree "$build_root" show -s --format=%ct HEAD)" \ + TORCH_CUDA_ARCH_LIST="$arch" MAX_JOBS=16 \ + python3 setup.py build_ext --inplace) >&2 2>&1 \ + || { cx_log "ERROR: hybrid-ep build failed"; exit 1; } + content_sha256="$(cx_extension_pair_sha256 \ + "$build_root" 'deep_ep_cpp*.so' 'hybrid_ep_cpp*.so')" || exit 1 + marker_tmp="$(mktemp "$build_root/.collectivex-complete.tmp.XXXXXX")" || exit 1 + chmod 600 "$marker_tmp" || exit 1 + printf '%s\n%s\n%s\n%s\n' \ + "$revision" "$tree" "$build_mode" "$content_sha256" > "$marker_tmp" \ + || exit 1 + mv -f -- "$marker_tmp" "$marker" || exit 1 + fi + cx_deepep_hybrid_cache_is_valid \ + "$build_root" "$marker" "$revision" "$tree" "$build_mode" + ); then + cx_log "ERROR: shared hybrid-ep build is incomplete" + return 1 + fi + export PYTHONPATH="$build_root:${PYTHONPATH:-}" + python3 -c "import deep_ep; assert hasattr(deep_ep,'HybridEPBuffer'); print('built hybrid-ep deep_ep', getattr(deep_ep,'__version__','?'))" >&2 \ + || { cx_log "ERROR: hybrid-ep import / HybridEPBuffer missing after build"; return 1; } + cx_log "DeepEP hybrid-ep ready ($DEEPEP_COMMIT, mode=$build_mode)" +} + +# UCCL EP (uccl.ep.Buffer is a DeepEP-API clone). The prebuilt wheel is cu12; on a cu13 +# image its kernels need a cu12 CUDA runtime on LD_LIBRARY_PATH (probe-confirmed). PEP-668 +# images need PIP_BREAK_SYSTEM_PACKAGES. Best-effort; failure to import fails loudly. +cx_build_uccl() { + if [ -f /tmp/.cx_built_uccl ]; then + cx_log "UCCL EP already prepared this allocation — skip rebuild" + python3 -c "import torch; from uccl_deepep import Buffer" 2>/dev/null || return 1 + return 0 + fi + local version="0.1.1" tag="v0.1.1" node_id wrapper_root + local wheel_sha256="390c1320918972206546e44d79b132988f2818ec07e23afcd0595f7183916cec" + cx_log "UCCL EP: installing uccl==$version + cu12 runtime shim" + export PIP_BREAK_SYSTEM_PACKAGES=1 + pip install -q --no-deps "sortedcontainers==2.4.0" "intervaltree==3.1.0" >&2 2>&1 \ + || { cx_log "ERROR: UCCL support dependency installation failed"; return 1; } + printf 'uccl==%s --hash=sha256:%s\n' "$version" "$wheel_sha256" \ + | pip install -q --no-deps --only-binary=:all: --require-hashes -r /dev/stdin >&2 2>&1 \ + || { cx_log "ERROR: pip install uccl==$version failed"; return 1; } + pip install -q --no-deps "nvidia-cuda-runtime-cu12==12.9.79" >&2 2>&1 \ + || { cx_log "ERROR: CUDA 12 runtime shim install failed"; return 1; } + local cu12lib + cu12lib="$(python3 -c "import nvidia.cuda_runtime as m, os; print(os.path.join(os.path.dirname(m.__file__),'lib'))" 2>/dev/null)" + [ -n "$cu12lib" ] && export LD_LIBRARY_PATH="$cu12lib:${LD_LIBRARY_PATH:-}" + local installed + installed="$(python3 -c 'import importlib.metadata as m; print(m.version("uccl"))')" \ + || { cx_log "ERROR: cannot read installed UCCL version"; return 1; } + [ "$installed" = "$version" ] \ + || { cx_log "ERROR: expected UCCL $version, installed $installed"; return 1; } + UCCL_COMMIT="pkg-$installed" + export UCCL_COMMIT + # import torch FIRST: uccl.ep's C extension links libc10.so (torch), which is only on the loader + # path once torch is imported (rpath). The adapter (ep_uccl.py) imports torch before uccl.ep too. + python3 -c "import torch; from uccl.ep import Buffer; print('uccl.ep ready')" >&2 \ + || { cx_log "ERROR: uccl.ep import failed (cu12 runtime on LD_LIBRARY_PATH?)"; return 1; } + # Vendor UCCL's DeepEP-API wrapper (ep/deep_ep_wrapper/deep_ep) under a NON-conflicting name + # (uccl_deepep) so it doesn't shadow the container's real deep_ep. Its Buffer(group, num_nvl_bytes, + # ...) takes a torch ProcessGroup (matching DeepEP + ep_uccl.py's calls) and runs the full + # proxy/IPC-handle/runtime.sync bootstrap that the low-level uccl.ep.Buffer(rank,num_ranks) lacks. + case "${SLURM_NODEID:-0}" in ""|*[!0-9]*) return 1 ;; esac + node_id="${SLURM_NODEID:-0}" + wrapper_root="$PWD/.cx_backend/uccl-wrapper-node-$node_id" + case "$wrapper_root" in /ix/experimental/CollectiveX/.cx_backend/uccl-wrapper-node-*) ;; *) return 1 ;; esac + rm -rf /tmp/uccl_src "$wrapper_root" + # Pin the wrapper to the SAME tag as the installed wheel (pkg-0.1.1 -> v0.1.1): the wrapper's + # dispatch calls into uccl.ep (get_rdma_buffer etc.), so a main-branch wrapper vs a 0.1.1 wheel + # mismatches signatures. Match them. + if git clone --depth 1 --branch "$tag" https://github.com/uccl-project/uccl /tmp/uccl_src >&2 2>&1 \ + && [ "$(git -C /tmp/uccl_src rev-parse HEAD)" = "73ee4f12ba71717d6de34ba06806e1baaabe3f42" ] \ + && [ -d /tmp/uccl_src/ep/deep_ep_wrapper/deep_ep ]; then + mkdir -p "$wrapper_root/uccl_deepep" + chmod 700 "$PWD/.cx_backend" "$wrapper_root" "$wrapper_root/uccl_deepep" + cp /tmp/uccl_src/ep/deep_ep_wrapper/deep_ep/*.py "$wrapper_root/uccl_deepep/" 2>/dev/null + export PYTHONPATH="$wrapper_root:${PYTHONPATH:-}" + python3 -c "import torch; from uccl_deepep import Buffer; print('uccl_deepep wrapper ready')" >&2 \ + || { cx_log "ERROR: uccl_deepep wrapper import failed"; return 1; } + export CX_UCCL_WRAPPER=1 + export UCCL_WRAPPER_COMMIT="73ee4f12ba71717d6de34ba06806e1baaabe3f42" + else + cx_log "ERROR: uccl deep_ep_wrapper not available" + return 1 + fi + : > /tmp/.cx_built_uccl + cx_log "UCCL EP ready ($UCCL_COMMIT, wrapper=${CX_UCCL_WRAPPER:-0})" +} + +# Rack build and rank steps may enter different container instances. Persist each node's +# loader/import path and build identity on the shared staged mount, then require it from every rank. +cx_persist_backend_env() { + local root="$PWD/.cx_backend/env" node_id="${SLURM_NODEID:-0}" path temporary name + local -a names=(PATH VIRTUAL_ENV LD_LIBRARY_PATH PYTHONPATH CUDA_HOME CPATH NVCC_PREPEND_FLAGS + NVSHMEM_DIR DEEPEP_COMMIT DEEPEP_TREE + EP_NCCL_ROOT_DIR EP_NVSHMEM_ROOT_DIR EP_JIT_CACHE_DIR EP_REUSE_NCCL_COMM + EP_JIT_DUMP_SASS + DEEPEP_V2_PR DEEPEP_V2_FIX_PR DEEPEP_V2_NCCL_CHECK_FIX_PR DEEPEP_V2_COMMIT + DEEPEP_V2_TREE DEEPEP_V2_FMT_COMMIT DEEPEP_V2_NCCL_CHECK_COMMIT + DEEPEP_V2_JIT_RANDOM_SEED + HYBRID_EP_MULTINODE USE_NIXL RDMA_CORE_HOME DEEPEP_HYBRID_BUILD_MODE + UCCL_COMMIT UCCL_WRAPPER_COMMIT CX_UCCL_WRAPPER) + [[ "$node_id" =~ ^[0-9]+$ ]] || return 1 + mkdir -p "$root" || return 1 + chmod 700 "$root" || return 1 + temporary="$(mktemp "$root/.node-${node_id}.XXXXXX")" || return 1 + chmod 600 "$temporary" || { rm -f "$temporary"; return 1; } + for name in "${names[@]}"; do + if declare -p "$name" >/dev/null 2>&1; then + printf 'export %s=%q\n' "$name" "${!name}" >> "$temporary" \ + || { rm -f "$temporary"; return 1; } + fi + done + path="$root/node-${node_id}.sh" + mv -f -- "$temporary" "$path" || { rm -f "$temporary"; return 1; } +} + +# Validate private scale-out selectors on every allocated compute node before a +# backend can initialize or build transport code. +cx_probe_scaleout_network() { + local interface device rdma_name + local -a interfaces devices + if [ "${CX_NODES:-1}" -le 1 ] || [ "${CX_TRANSPORT:-}" = mnnvl ]; then + return 0 + fi + cx_restore_exact_hca_selector || return 1 + [ -n "${GLOO_SOCKET_IFNAME:-}" ] && [ -n "${NCCL_IB_HCA:-}" ] \ + || { cx_log "ERROR: scale-out network selectors are unavailable"; return 1; } + IFS=, read -r -a interfaces <<< "$GLOO_SOCKET_IFNAME" + for interface in "${interfaces[@]}"; do + [ -d "/sys/class/net/$interface" ] \ + || { cx_log "ERROR: configured scale-out socket interface is absent"; return 1; } + done + IFS=, read -r -a devices <<< "$NCCL_IB_HCA" + for device in "${devices[@]}"; do + rdma_name="$(cx_nccl_hca_device_name "$device")" + [ -d "/sys/class/infiniband/$rdma_name" ] \ + || { cx_log "ERROR: configured scale-out RDMA device is absent"; return 1; } + done +} + +# Prepare and probe one backend without running a benchmark. The same hook is used +# by normal in-container runs and by rack launchers' persistent build-only step. +cx_prepare_backend() { + local backend="${1:-}" + case "$backend" in + deepep) + cx_probe_deepep || return 1 + ;; + deepep-v2) + cx_build_deepep_v2 || return 1 + ;; + deepep-hybrid) + cx_build_deepep_hybrid || return 1 + ;; + uccl) + cx_build_uccl || return 1 + ;; + mori) + python3 -c "import mori" \ + || { cx_log "ERROR: MoRI backend import failed"; return 1; } + ;; + nccl-ep) + ;; + *) + cx_log "ERROR: unknown backend preparation request" + return 1 + ;; + esac +} + +prepare_backend_or_record() { + local backend="$1" + cx_write_runtime_stage backend-setup || return 1 + if cx_prepare_backend "$backend"; then + return 0 + fi + cx_log "WARN: $backend preparation failed" + return 1 +} + +# dispatch_bench runs the CURRENT CX_BENCH (+ CX_* config env) once. The sweep workflow runs many +# of these per allocation (SHARD mode below), reusing this single container + its built backend. +dispatch_bench() { + case "$CX_BENCH" in + nccl-ep) + run_ep_suite "$CX_BENCH" + ;; + deepep|deepep-v2|deepep-hybrid|mori|uccl) + prepare_backend_or_record "$CX_BENCH" && run_ep_suite "$CX_BENCH" + ;; + *) + cx_die "unknown CX_BENCH=$CX_BENCH (want deepep|deepep-v2|mori|uccl|nccl-ep|deepep-hybrid)" + ;; + esac +} + +rc=0 +cx_validate_shard_control "$PWD" +cx_load_network_control_mode "$PWD" || cx_die "cannot resolve network control mode" +cx_apply_network_profile "${CX_NODES:-1}" "${CX_TRANSPORT:-}" +# Build-only mode: rack launchers run the shared backend preparation hook once per +# node inside a persistent named container, then direct rank processes reuse it. +if [ -n "${CX_BUILD_ONLY:-}" ]; then + if cx_probe_scaleout_network && cx_prepare_backend "${CX_BENCH:-}"; then + cx_persist_backend_env || rc=1 + else + rc=1 + fi + cx_log "backend preparation: bench=${CX_BENCH:-unknown} rc=$rc" + exit "$rc" +fi +if [ -n "${CX_SHARD_FILE:-}" ]; then + # SHARD/SWEEP mode (collectivex-sweep.yml): run EVERY case of this shard in THIS one allocation. + # All cases share (sku, backend, nodes), so backend preparation is paid once and cached. + ncases="$(python3 -c "import json;print(len(json.load(open('$CX_SHARD_FILE'))['cases']))")" + # The iterable benchmark version is a shard-level scalar (identical for every case); + # export it once so run_ep copies it verbatim into each emitted result. + CX_VERSION="$(python3 -c "import json;print(json.load(open('$CX_SHARD_FILE'))['version'])")" + export CX_VERSION + cx_log "SHARD mode: $ncases case(s) in one allocation (shard=$CX_SHARD_FILE, version=$CX_VERSION)" + _cx_ts_base="$CX_TS" # per-case CX_TS suffix below keeps each case's result file UNIQUE (else + # cases sharing backend+phase overwrite each other at the same timestamp). + ci=0 + failed_cases=0 + while [ "$ci" -lt "$ncases" ]; do + CX_TS="${_cx_ts_base}-c$(printf '%03d' "$ci")" + export CX_TS + # Map varying case fields plus the frozen v1 defaults into CX_* env. + _exports="$(python3 - "$CX_SHARD_FILE" "$ci" <<'PY' +import json, sys, shlex +c = json.load(open(sys.argv[1]))["cases"][int(sys.argv[2])] +def g(k, d=""): + v = c.get(k, d); return "" if v is None else str(v) +env = { + "CX_BENCH": g("backend"), + "CX_MODE": g("mode", "normal"), + "CX_ROUTING": g("routing", "uniform"), "CX_PHASE": g("phase", "decode"), + "CX_EP": g("ep", "1"), + "CX_EPLB": "1" if c.get("eplb") else "", + "CX_CASE_ID": g("case_id"), "CX_SUITE": g("suite"), "CX_WORKLOAD_NAME": g("workload"), + "CX_HIDDEN": g("hidden"), "CX_TOPK": g("topk"), "CX_EXPERTS": g("experts"), + "CX_TOKENS_LADDER": g("ladder"), "CX_CANONICAL": ("1" if c.get("canonical") else ""), + "CX_NODES": g("nodes"), "CX_GPUS_PER_NODE": g("gpus_per_node"), + "CX_SCALE_UP_DOMAIN": g("scale_up_domain"), "CX_SCOPE": g("scope"), + "CX_SCALE_UP_TRANSPORT": g("scale_up_transport"), + "CX_SCALE_OUT_TRANSPORT": g("scale_out_transport"), + "CX_TRANSPORT": g("transport"), "CX_TOPO": g("topology_class"), + "CX_SAMPLES_PER_POINT": g("samples_per_point"), + "CX_WARMUP_SEMANTICS": g("warmup_semantics"), +} +lines = [f"export {k}={shlex.quote(v)}" for k, v in env.items()] +# Per-case timing "iters:trials:warmup" (fixed-512-v1 requires 8:64:32 everywhere); +# cases without one must fall back to the harness defaults, so UNSET rather than export-empty +# (an empty CX_ITERS would defeat the 8-iter default and break the run_ep argparse; NOTE no +# apostrophes in this heredoc — bash command-substitution scanning chokes on unbalanced quotes). +timing = g("timing") +if timing: + parts = (timing.split(":") + ["", "", ""])[:3] + for k, v in zip(("CX_ITERS", "CX_TRIALS", "CX_WARMUP"), parts): + if v: + lines.append(f"export {k}={shlex.quote(v)}") +else: + lines.append("unset CX_ITERS CX_TRIALS CX_WARMUP 2>/dev/null || true") +print("\n".join(lines)) +PY +)" + eval "$_exports" + cx_apply_network_profile "$CX_NODES" "$CX_TRANSPORT" + # Each case has its OWN routing/dims -> its own canonical workload manifest. cx_stage_canonical + # short-circuits when CX_WORKLOAD_DIR is already set, so without this unset the first case's + # staged dir is reused for the rest and run_ep.py can't find the later cases' manifests + # (FileNotFoundError .cx_workloads/.manifest.json). Unset so every case re-stages its own. + unset CX_WORKLOAD_DIR 2>/dev/null || true + cx_log " [$((ci+1))/$ncases] $CX_BENCH $CX_MODE/$CX_PHASE routing=$CX_ROUTING eplb=${CX_EPLB:-0}" + _cx_case_ts="$CX_TS" + CX_TS="${_cx_case_ts}-a01" + export CX_ATTEMPT_ID=1 CX_TS + dispatch_bench || { + failed_cases=$((failed_cases+1)) + cx_log " [$((ci+1))/$ncases] $CX_BENCH case FAILED; failed-case record preserved" + } + export CX_TS="$_cx_case_ts" + ci=$((ci + 1)) + done + if [ "${failed_cases:-0}" -gt 0 ]; then + cx_log "SHARD done: $failed_cases/$ncases case(s) failed" + rc=1 + fi + # The base timestamp matches every per-case file, so the final summary covers the whole shard. + export CX_TS="$_cx_ts_base" +else + _cx_single_ts="$CX_TS" + CX_TS="${_cx_single_ts}-a01" + export CX_ATTEMPT_ID=1 CX_TS + dispatch_bench || rc=1 +fi + +# Summary table for the log; also fails the job if no valid results were produced. +python3 summarize.py --results-dir results --runner "$CX_RUNNER" --ts "$CX_TS" || rc=1 +exit "$rc" diff --git a/experimental/CollectiveX/source_archive.py b/experimental/CollectiveX/source_archive.py new file mode 100644 index 0000000000..c027490a65 --- /dev/null +++ b/experimental/CollectiveX/source_archive.py @@ -0,0 +1,349 @@ +#!/usr/bin/env python3 +"""Validate and extract one pinned backend from a shared source tar.""" +from __future__ import annotations + +import argparse +import os +from pathlib import Path, PurePosixPath +import stat +import tarfile +from typing import Optional, Sequence + + +PathParts = tuple[str, ...] +_DIRECTORY_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC +_FILE_FLAGS = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW | os.O_CLOEXEC +MAX_ARCHIVE_MEMBERS = 20_000 +MAX_MEMBER_BYTES = 512 * 1024 * 1024 +MAX_EXPANDED_BYTES = 2 * 1024 * 1024 * 1024 +MAX_ARCHIVE_BYTES = 4 * 1024 * 1024 * 1024 +MAX_ARCHIVE_HEADERS = 40_000 +MAX_EXTENSION_BYTES = 64 * 1024 * 1024 +MAX_EXTENSION_MEMBER_BYTES = 1024 * 1024 +MAX_EXTENSION_CHAIN = 8 +_TAR_BLOCK = 512 +_EXTENSION_TYPES = {b"L", b"K", b"x", b"g", b"X"} + + +class SourceArchiveError(ValueError): + """The backend source archive cannot be extracted safely.""" + + +def _tar_size(field: bytes) -> int: + if field[0] in (0o200, 0o377): + value = int.from_bytes(field[1:], "big") + if field[0] == 0o377: + value -= 256 ** (len(field) - 1) + return value + try: + text = field.split(b"\0", 1)[0].decode("ascii").strip() + return int(text or "0", 8) + except (UnicodeDecodeError, ValueError) as exc: + raise SourceArchiveError("archive contains an invalid size field") from exc + + +def _preflight_archive(descriptor: int, archive_size: int) -> None: + if archive_size <= 0 or archive_size > MAX_ARCHIVE_BYTES: + raise SourceArchiveError("backend source archive exceeds the raw size limit") + offset = headers = extension_bytes = extension_chain = 0 + while offset < archive_size: + header = os.pread(descriptor, _TAR_BLOCK, offset) + if len(header) != _TAR_BLOCK: + raise SourceArchiveError("archive header is truncated") + if not any(header): + return + headers += 1 + if headers > MAX_ARCHIVE_HEADERS: + raise SourceArchiveError("archive has too many physical headers") + size = _tar_size(header[124:136]) + if size < 0: + raise SourceArchiveError("archive contains a negative payload size") + type_flag = header[156:157] + if type_flag in _EXTENSION_TYPES: + extension_chain += 1 + extension_bytes += size + if ( + extension_chain > MAX_EXTENSION_CHAIN + or size > MAX_EXTENSION_MEMBER_BYTES + or extension_bytes > MAX_EXTENSION_BYTES + ): + raise SourceArchiveError("archive extension metadata exceeds its limit") + if type_flag in {b"x", b"g", b"X"}: + payload = os.pread(descriptor, size, offset + _TAR_BLOCK) + if len(payload) != size: + raise SourceArchiveError("archive extension metadata is truncated") + if b"GNU.sparse." in payload: + raise SourceArchiveError("archive contains sparse extension metadata") + else: + extension_chain = 0 + if type_flag == b"S": + raise SourceArchiveError("archive contains a sparse member") + blocks = (size + _TAR_BLOCK - 1) // _TAR_BLOCK + offset += _TAR_BLOCK + blocks * _TAR_BLOCK + if offset > archive_size: + raise SourceArchiveError("archive payload is truncated") + + +def _member_parts(name: str) -> PathParts: + if not name or "\\" in name or "\0" in name: + raise SourceArchiveError("archive contains a noncanonical member path") + path = PurePosixPath(name) + if ( + path.is_absolute() + or path.as_posix() != name + or not path.parts + or path.parts[0] != ".cx_sources" + or any(part in {"", ".", ".."} for part in path.parts) + ): + raise SourceArchiveError("archive contains a noncanonical member path") + return path.parts + + +def _root_parts(root_basename: str) -> PathParts: + path = PurePosixPath(root_basename) + if ( + not root_basename + or "\\" in root_basename + or "\0" in root_basename + or path.is_absolute() + or path.as_posix() != root_basename + or len(path.parts) != 1 + or path.parts[0] in {"", ".", ".."} + ): + raise SourceArchiveError("invalid backend source root") + return (".cx_sources", root_basename) + + +def _read_members(archive: tarfile.TarFile) -> list[tarfile.TarInfo]: + members: list[tarfile.TarInfo] = [] + for member in archive: + if len(members) >= MAX_ARCHIVE_MEMBERS: + raise SourceArchiveError("archive has an invalid member count") + members.append(member) + return members + + +def _validate_members( + members: list[tarfile.TarInfo], selected_root: PathParts +) -> dict[PathParts, tarfile.TarInfo]: + if not members or len(members) > MAX_ARCHIVE_MEMBERS: + raise SourceArchiveError("archive has an invalid member count") + entries: dict[PathParts, tarfile.TarInfo] = {} + expanded_bytes = 0 + for member in members: + parts = _member_parts(member.name) + if parts in entries: + raise SourceArchiveError("archive contains duplicate member paths") + if member.sparse is not None: + raise SourceArchiveError("archive contains a sparse member") + if member.isdir(): + if member.size != 0: + raise SourceArchiveError("archive contains an invalid directory") + elif member.isfile(): + if member.size < 0 or member.size > MAX_MEMBER_BYTES: + raise SourceArchiveError("archive member exceeds the size limit") + expanded_bytes += member.size + if expanded_bytes > MAX_EXPANDED_BYTES: + raise SourceArchiveError("archive exceeds the expanded size limit") + elif member.issym(): + if member.size != 0: + raise SourceArchiveError("archive contains an invalid symbolic link") + else: + raise SourceArchiveError("archive contains a non-file member") + entries[parts] = member + + source_parent = entries.get((".cx_sources",)) + selected = entries.get(selected_root) + if source_parent is None or not source_parent.isdir(): + raise SourceArchiveError("archive is missing its source directory") + if selected is None or not selected.isdir(): + raise SourceArchiveError("archive is missing the selected backend source") + + for parts in entries: + for depth in range(1, len(parts)): + parent = entries.get(parts[:depth]) + if parent is None or not parent.isdir(): + raise SourceArchiveError("archive member has an unsafe parent") + + for parts, member in entries.items(): + if not member.issym(): + continue + target_name = member.linkname + target_path = PurePosixPath(target_name) + if ( + not target_name + or "\\" in target_name + or "\0" in target_name + or target_path.is_absolute() + or target_path.as_posix() != target_name + ): + raise SourceArchiveError("archive contains an unsafe symbolic link") + target = list(parts[:-1]) + for component in target_path.parts: + if component == "..": + if len(target) <= 2: + raise SourceArchiveError("symbolic link escapes its backend source") + target.pop() + else: + target.append(component) + resolved = tuple(target) + if resolved[:2] != parts[:2]: + raise SourceArchiveError("symbolic link crosses backend sources") + target_member = entries.get(resolved) + if target_member is None or not target_member.isfile(): + raise SourceArchiveError("symbolic link target is not a regular archive file") + return entries + + +def _open_directory(root_fd: int, parts: PathParts) -> int: + descriptor = os.dup(root_fd) + try: + for part in parts: + child = os.open(part, _DIRECTORY_FLAGS, dir_fd=descriptor) + os.close(descriptor) + descriptor = child + return descriptor + except BaseException: + os.close(descriptor) + raise + + +def _create_directory(root_fd: int, parts: PathParts) -> None: + parent_fd = _open_directory(root_fd, parts[:-1]) + try: + os.mkdir(parts[-1], mode=0o700, dir_fd=parent_fd) + finally: + os.close(parent_fd) + + +def _extract_file( + archive: tarfile.TarFile, root_fd: int, parts: PathParts, member: tarfile.TarInfo +) -> None: + parent_fd = _open_directory(root_fd, parts[:-1]) + descriptor = -1 + source = None + try: + mode = 0o700 if member.mode & 0o111 else 0o600 + descriptor = os.open(parts[-1], _FILE_FLAGS, mode, dir_fd=parent_fd) + source = archive.extractfile(member) + if source is None: + raise SourceArchiveError("archive file has no readable payload") + remaining = member.size + while remaining: + chunk = source.read(min(1024 * 1024, remaining)) + if not chunk: + raise SourceArchiveError("archive file payload is truncated") + view = memoryview(chunk) + while view: + written = os.write(descriptor, view) + view = view[written:] + remaining -= len(chunk) + os.fchmod(descriptor, mode) + finally: + if source is not None: + source.close() + if descriptor >= 0: + os.close(descriptor) + os.close(parent_fd) + + +def _extract_symlink(root_fd: int, parts: PathParts, member: tarfile.TarInfo) -> None: + parent_fd = _open_directory(root_fd, parts[:-1]) + try: + os.symlink(member.linkname, parts[-1], dir_fd=parent_fd) + finally: + os.close(parent_fd) + + +def _extract_selected( + archive: tarfile.TarFile, + destination_fd: int, + entries: dict[PathParts, tarfile.TarInfo], + selected_root: PathParts, +) -> None: + try: + os.stat(".cx_sources", dir_fd=destination_fd, follow_symlinks=False) + except FileNotFoundError: + pass + else: + raise SourceArchiveError("backend source output already exists") + + selected = { + parts: member + for parts, member in entries.items() + if parts[: len(selected_root)] == selected_root + } + _create_directory(destination_fd, (".cx_sources",)) + directories = sorted( + (parts for parts, member in selected.items() if member.isdir()), + key=lambda parts: (len(parts), parts), + ) + for parts in directories: + _create_directory(destination_fd, parts) + for parts, member in sorted(selected.items()): + if member.isfile(): + _extract_file(archive, destination_fd, parts, member) + for parts, member in sorted(selected.items()): + if member.issym(): + _extract_symlink(destination_fd, parts, member) + + +def extract_source_archive( + archive_path: Path, destination: Path, root_basename: str +) -> None: + """Validate the complete tar, then safely extract one backend source root.""" + selected_root = _root_parts(root_basename) + archive_fd = os.open(archive_path, os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC) + try: + metadata = os.fstat(archive_fd) + if ( + not stat.S_ISREG(metadata.st_mode) + or metadata.st_uid != os.getuid() + or stat.S_IMODE(metadata.st_mode) & 0o022 + ): + raise SourceArchiveError("backend source archive has unsafe metadata") + _preflight_archive(archive_fd, metadata.st_size) + with os.fdopen(os.dup(archive_fd), "rb") as stream: + try: + with tarfile.open(fileobj=stream, mode="r:") as archive: + entries = _validate_members(_read_members(archive), selected_root) + destination_fd = os.open(destination, _DIRECTORY_FLAGS) + try: + destination_metadata = os.fstat(destination_fd) + if ( + destination_metadata.st_uid != os.getuid() + or stat.S_IMODE(destination_metadata.st_mode) != 0o700 + ): + raise SourceArchiveError("backend source destination is unsafe") + previous_umask = os.umask(0o077) + try: + _extract_selected( + archive, destination_fd, entries, selected_root + ) + finally: + os.umask(previous_umask) + finally: + os.close(destination_fd) + except RecursionError as exc: + raise SourceArchiveError("archive extension metadata is recursive") from exc + finally: + os.close(archive_fd) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + description="Safely install one pinned backend source archive" + ) + parser.add_argument("archive", type=Path) + parser.add_argument("destination", type=Path) + parser.add_argument("root_basename") + args = parser.parse_args(argv) + try: + extract_source_archive(args.archive, args.destination, args.root_basename) + except (OSError, SourceArchiveError, tarfile.TarError) as exc: + parser.error(f"backend source archive rejected: {exc}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/summarize.py b/experimental/CollectiveX/summarize.py new file mode 100644 index 0000000000..91197da20d --- /dev/null +++ b/experimental/CollectiveX/summarize.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +"""Render a small native-v1 shard summary and gate on a successful case.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +# Emitted document format this summary reads. This is a best-effort renderer over +# whatever raw attempts a shard produced; it validates nothing. +RAW_FORMAT = "collectivex.ep.v1" + + +def load_results(directory: str, runner: str | None, timestamp: str | None) -> list[dict]: + documents: list[dict] = [] + for path in sorted(Path(directory).glob("*.json")): + if runner and not path.name.startswith(f"{runner}_"): + continue + if timestamp and timestamp not in path.name: + continue + try: + with path.open() as handle: + document = json.load(handle) + except (OSError, ValueError): + continue + if isinstance(document, dict) and document.get("format") == RAW_FORMAT: + documents.append(document) + return documents + + +def _identity(document: dict) -> tuple[str, str, str, str, bool, int]: + case = document["case"] + routing = case["shape"]["routing"] + eplb = case["eplb"]["enabled"] + sku = document["identity"]["case_factors"]["sku"] + return ( + sku, case["suite"], routing, case["phase"], eplb, + case.get("ep_size", case.get("ep", 0)), + ) + + +def _headline(document: dict) -> tuple[int | str, float | str, float | str]: + rows = document["measurement"]["rows"] + row = next((item for item in rows if item["tokens_per_rank"] == 64), rows[len(rows) // 2]) + latency = row["components"]["roundtrip"]["percentiles_us"] + return row["tokens_per_rank"], latency["p50"], latency["p99"] + + +def render(documents: list[dict], markdown: bool) -> str: + documents = sorted(documents, key=_identity) + if markdown: + lines = [ + "## CollectiveX EP results", "", + "| ver | sku | backend | suite | phase | routing | ep | outcome | T* | p50 us | p99 us |", + "|--:|---|---|---|---|---|--:|---|--:|--:|--:|", + ] + for document in documents: + sku, suite, routing, phase, eplb, ep = _identity(document) + backend = document["case"]["backend"] + token, p50, p99 = _headline(document) + lines.append( + f"| {document['version']} | {sku} | `{backend}` | {suite} | {phase} | " + f"{routing}{'+eplb' if eplb else ''} | {ep} | " + f"{document['outcome']['status']} | {token} | {p50} | {p99} |" + ) + if not documents: + lines.append("\n> No valid native outcome documents found.") + return "\n".join(lines) + lines = ["CollectiveX EP results", "======================"] + for document in documents: + sku, suite, routing, phase, eplb, ep = _identity(document) + backend = document["case"]["backend"] + token, _, p99 = _headline(document) + lines.append( + f" v{document['version']} {sku:<10} {backend:<16} {suite:<13} {phase:<7} " + f"{routing}{'+eplb' if eplb else ''} ep{ep} " + f"{document['outcome']['status']} T={token} roundtrip_p99_us={p99}" + ) + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Summarize CollectiveX native v1 outcomes") + parser.add_argument("--results-dir", default="results") + parser.add_argument("--runner") + parser.add_argument("--ts") + parser.add_argument("--markdown", action="store_true") + args = parser.parse_args() + documents = load_results(args.results_dir, args.runner, args.ts) + print(render(documents, args.markdown)) + if args.markdown: + return 0 + return 0 if any( + document["outcome"]["status"] == "success" + for document in documents + ) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py new file mode 100644 index 0000000000..4a8ce4878d --- /dev/null +++ b/experimental/CollectiveX/sweep_matrix.py @@ -0,0 +1,998 @@ +#!/usr/bin/env python3 +"""Resolve CollectiveX suites and extract validated execution shards. + +Mode changes measurement semantics and therefore participates in case identity. +Dispatch and combine are fixed BF16 benchmark facts, so a case's coordinates are +suite/workload/backend/topology only; the matrix schedules, it never ranks. +""" +from __future__ import annotations + +import argparse +import itertools +import json +import os +from pathlib import Path +import re +import sys +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE / "bench")) + +try: # Shard extraction on GPU runners is intentionally stdlib-only. + import yaml # type: ignore +except ModuleNotFoundError: # pragma: no cover - exercised by the workflow environment + yaml = None + +import capability as cap # noqa: E402 +import ep_harness # noqa: E402 +import identity # noqa: E402 + + +EP_TIMING_PROFILE = ( + f"{ep_harness.TIMED_ITERS_PER_TRIAL}:{ep_harness.TRIALS_PER_POINT}:" + f"{ep_harness.WARMUP_ITERS_PER_TRIAL}" +) +V1_WORKLOAD = ("deepseek-v3", 7168, 8, 256) +V1_SUITE_CONTRACTS = { + "ep-core": { + "mode": "normal", + "coordinates": { + ("normal", "decode", "uniform", False), + ("normal", "prefill", "uniform", False), + }, + "ladders": { + "decode": tuple(ep_harness.DECODE_LADDER), + "prefill": (256, 512), + }, + }, + "ep-low-latency": { + "mode": "low-latency", + "backends": {"deepep", "uccl"}, + "coordinates": {("low-latency", "decode", "uniform", False)}, + "ladders": {"decode": tuple(ep_harness.DECODE_LADDER)}, + }, +} +IDENTIFIER = re.compile(r"[a-z0-9][a-z0-9.-]*") +SUITE_FIELDS = { + "backends", "ep_degrees", "eplb", "mode", "phases", "platforms", + "routings", "token_points", "token_points_decode", "token_points_prefill", + "workloads", +} +SUITE_REQUIRED = { + "ep_degrees", "mode", "phases", "platforms", "routings", "workloads", +} +TOPOLOGY_FIELDS = ( + "nodes", "gpus_per_node", "scale_up_domain", "scope", "scale_up_transport", + "scale_out_transport", "transport", "topology_class", +) + + +class MatrixError(ValueError): + """A matrix or shard-control document violates the execution contract.""" + + +if yaml is not None: + class _UniqueKeyLoader(yaml.SafeLoader): + pass + + def _unique_mapping(loader: Any, node: Any, deep: bool = False) -> dict[Any, Any]: + result: dict[Any, Any] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + if key in result: + raise SystemExit(f"duplicate YAML key {key!r} at line {key_node.start_mark.line + 1}") + result[key] = loader.construct_object(value_node, deep=deep) + return result + + _UniqueKeyLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _unique_mapping + ) + + +def _load(name: str) -> dict[str, Any]: + if yaml is None: + raise SystemExit("matrix generation requires PyYAML; shard extraction does not") + try: + with (HERE / "configs" / name).open() as fh: + document = yaml.load(fh, Loader=_UniqueKeyLoader) + except yaml.YAMLError as exc: + raise SystemExit(f"configs/{name} is not valid YAML: {exc}") from exc + if not isinstance(document, dict): + raise SystemExit(f"configs/{name} must contain a YAML object") + return document + + +def _workload_registry(workloads: dict[str, Any]) -> dict[str, dict[str, Any]]: + return { + name: cfg + for section in ("synthetic", "model_derived") + for name, cfg in (workloads.get(section) or {}).items() + } + + +def _fields(value: Any, path: str, allowed: set[str], required: set[str]) -> dict[str, Any]: + if not isinstance(value, dict): + raise SystemExit(f"{path} must be an object") + if any(not isinstance(key, str) for key in value): + raise SystemExit(f"{path} field names must be strings") + unknown, missing = set(value) - allowed, required - set(value) + if unknown or missing: + raise SystemExit(f"{path} fields: unknown={sorted(unknown)}, missing={sorted(missing)}") + return value + + +def _list(value: Any, path: str, item_type: type, allowed: set[Any] | None = None) -> list[Any]: + if (not isinstance(value, list) or not value + or any(type(item) is not item_type for item in value) + or len(value) != len(set(value)) + or (allowed is not None and any(item not in allowed for item in value))): + raise SystemExit(f"{path} must be a non-empty unique list of valid {item_type.__name__}s") + return value + + +def validate_config_documents( + suites_document: dict[str, Any], workloads: dict[str, Any] +) -> None: + """Reject configuration that is ambiguous, unused, or outside the v1 grid.""" + _fields( + suites_document, "configs/suites.yaml", + {"schema_version", "version", "suites"}, {"schema_version", "version", "suites"}, + ) + _fields( + workloads, "configs/workloads.yaml", + {"schema_version", "synthetic", "model_derived"}, {"schema_version"}, + ) + if type(suites_document["schema_version"]) is not int or suites_document["schema_version"] != 1: + raise SystemExit("configs/suites.yaml schema_version must be integer 1") + # The iterable benchmark version is independent of the JSON schema_version: it + # is any positive integer the operator bumps to mark a new, incomparable grid. + if type(suites_document["version"]) is not int or suites_document["version"] < 1: + raise SystemExit("configs/suites.yaml version must be a positive integer") + if type(workloads["schema_version"]) is not int or workloads["schema_version"] != 1: + raise SystemExit("configs/workloads.yaml schema_version must be integer 1") + registry: dict[str, dict[str, Any]] = {} + for section, expert_field in ( + ("synthetic", "experts"), + ("model_derived", "routed_experts"), + ): + entries = workloads.get(section, {}) + if not isinstance(entries, dict): + raise SystemExit(f"workloads.{section} must be an object") + for name, value in entries.items(): + if not isinstance(name, str) or not IDENTIFIER.fullmatch(name) or name in registry: + raise SystemExit(f"workloads.{section} has invalid or duplicate name {name!r}") + fields = {"hidden", "topk", expert_field, "verified_against"} + config = _fields(value, f"workload {name}", fields, fields - {"verified_against"}) + dimensions = [config[key] for key in ("hidden", "topk", expert_field)] + if any(type(item) is not int or item <= 0 for item in dimensions): + raise SystemExit(f"workload {name} dimensions must be positive integers") + if dimensions[1] > dimensions[2]: + raise SystemExit(f"workload {name}.topk exceeds its expert count") + source = config.get("verified_against") + if source is not None and (not isinstance(source, str) or not source.strip()): + raise SystemExit(f"workload {name}.verified_against must be a non-empty string") + registry[name] = config + if not registry: + raise SystemExit("configs/workloads.yaml must define at least one workload") + + suites = suites_document["suites"] + if not isinstance(suites, dict) or not suites: + raise SystemExit("configs/suites.yaml suites must be a non-empty object") + referenced: set[str] = set() + for name, value in suites.items(): + if not isinstance(name, str) or not IDENTIFIER.fullmatch(name): + raise SystemExit(f"invalid suite name {name!r}") + suite = _fields(value, f"suite {name}", SUITE_FIELDS, SUITE_REQUIRED) + contract = V1_SUITE_CONTRACTS.get(name) + if contract is None: + raise SystemExit(f"suite {name} is outside the frozen v1 catalog") + mode = suite["mode"] + if mode not in identity.V1_CASE_PROFILES or mode != contract["mode"]: + raise SystemExit(f"suite {name}.mode differs from the frozen v1 catalog") + suite_backends = _list( + suite.get("backends", list(cap.SWEEP_BACKENDS)), + f"suite {name}.backends", + str, + set(cap.SWEEP_BACKENDS), + ) + expected_backends = contract.get("backends") + if expected_backends is not None and set(suite_backends) != expected_backends: + raise SystemExit(f"suite {name}.backends differs from the frozen v1 catalog") + if expected_backends is None and "backends" in suite: + raise SystemExit(f"suite {name}.backends must be omitted") + suite_workloads = _list(suite["workloads"], f"suite {name}.workloads", str) + unknown = sorted(set(suite_workloads) - set(registry)) + if unknown: + raise SystemExit(f"suite {name}: unknown workloads {unknown}") + referenced.update(suite_workloads) + platforms = _list( + suite["platforms"], f"suite {name}.platforms", str, set(cap.PLATFORMS) + ) + phases = _list(suite["phases"], f"suite {name}.phases", str, {"decode", "prefill"}) + routings = _list(suite["routings"], f"suite {name}.routings", str, {"uniform"}) + eplb = _list(suite.get("eplb", [False]), f"suite {name}.eplb", bool) + if True in eplb: + raise SystemExit(f"suite {name}: EPLB is unavailable for v1 uniform routing") + degrees = _list(suite["ep_degrees"], f"suite {name}.ep_degrees", int) + if degrees != [8, 16]: + raise SystemExit(f"suite {name}.ep_degrees must be exactly [8, 16]") + for platform in platforms: + if not set(degrees).issubset(cap.PLATFORMS[platform]["ep_degrees"]): + raise SystemExit(f"suite {name}: invalid EP degree for {platform}") + for phase in {"decode", "prefill"} - set(phases): + if f"token_points_{phase}" in suite: + raise SystemExit(f"suite {name}.token_points_{phase} is unreachable") + if "token_points" in suite and all( + f"token_points_{phase}" in suite for phase in phases + ): + raise SystemExit(f"suite {name}.token_points is unreachable") + for phase in phases: + _ladder(suite, phase) + coordinates = { + (mode, phase, routing, enabled) + for phase, routing, enabled in itertools.product(phases, routings, eplb) + } + if coordinates != contract["coordinates"] or any( + tuple(map(int, _ladder(suite, phase).split())) != contract["ladders"][phase] + for phase in phases + ): + raise SystemExit(f"suite {name} coordinates differ from the frozen v1 catalog") + unused = sorted(set(registry) - referenced) + if unused: + raise SystemExit(f"unreferenced workloads: {unused}") + + +def _dims(workloads: dict[str, Any], name: str) -> tuple[int, int, int]: + config = _workload_registry(workloads)[name] + values = ( + config.get("hidden"), + config.get("topk"), + config.get("experts", config.get("routed_experts")), + ) + return values # type: ignore[return-value] + + +def _ladder(suite: dict[str, Any], phase: str) -> str: + points = suite.get(f"token_points_{phase}", suite.get("token_points")) + if points is None: + points = ep_harness.DECODE_LADDER if phase == "decode" else ep_harness.PREFILL_LADDER + if (not isinstance(points, list) or not points + or any(isinstance(point, bool) or not isinstance(point, int) or point <= 0 + for point in points) + or points != sorted(set(points))): + raise SystemExit(f"invalid {phase} token ladder: {points!r}") + return " ".join(map(str, points)) + + +def _v1_requested_ladder(case: dict[str, Any]) -> str: + """Bind extracted controls to the frozen v1 suite and workload catalog.""" + suite = V1_SUITE_CONTRACTS.get(case.get("suite")) + coordinate = ( + case.get("mode"), case.get("phase"), case.get("routing"), case.get("eplb") + ) + if ( + suite is None + or coordinate not in suite["coordinates"] + or ( + case.get("workload"), case.get("hidden"), case.get("topk"), case.get("experts") + ) != V1_WORKLOAD + ): + raise MatrixError("case differs from the frozen v1 suite/workload catalog") + return " ".join(map(str, suite["ladders"][case["phase"]])) + + +def _expected_disposition( + sku: str, case: dict[str, Any] +) -> tuple[str, str | None, str | None]: + requested_ladder = _v1_requested_ladder(case) + disposition, detail = cap.resolve_disposition( + sku, case["backend"], ep=case["ep"], nodes=case["nodes"], + routing=case["routing"], eplb=case["eplb"], mode=case["mode"], + ) + if disposition == "supported": + if case["ladder"] != requested_ladder: + raise MatrixError("case ladder differs from the frozen v1 suite catalog") + return "runnable", None, None + if case["ladder"] != requested_ladder: + raise MatrixError("unsupported case ladder differs from the frozen v1 suite catalog") + if disposition == "unsupported": + return "unsupported", "backend-platform-unsupported", detail + raise MatrixError(f"invalid capability disposition {disposition!r}") + + +def _case_id(sku: str, case: dict[str, Any]) -> str: + return identity.case_id( + sku=sku, profile=identity.profile_for_case(case), case=case + ) + + +def _semantic_points(sku: str, case: dict[str, Any]) -> list[str]: + execution = { + key: value for key, value in case.items() + if key not in {"canonical", "case_id", "ladder", "suite", "workload"} + } + return [ + json.dumps( + {"sku": sku, "tokens_per_rank": int(point), **execution}, + sort_keys=True, + separators=(",", ":"), + ) + for point in case["ladder"].split() + ] + + +def _select_backends(backend: str, backends: str) -> list[str]: + available = list(cap.SWEEP_BACKENDS) + if backend and backends: + raise SystemExit("--backend and --backends are mutually exclusive") + if backends: + names = available if backends == "all" else [ + value.strip() for value in backends.split(",") if value.strip() + ] + else: + names = [backend or "deepep"] + unknown = sorted(set(names) - set(available)) + if unknown: + raise SystemExit(f"unknown backend values {unknown}; have {available}") + if len(names) != len(set(names)): + raise SystemExit("backend selection contains duplicates") + return names + + +def resolve_matrix( + suites: str = "all", + backend: str = "", + backends: str = "", + only_sku: str = "", + exclude_skus: str = "", + ep_sizes: str = "", + max_cases: int = 128, +) -> dict[str, Any]: + """Resolve suite configuration into allocation-sized workflow shards.""" + if max_cases <= 0: + raise SystemExit("--max-cases must be positive") + # --ep-sizes narrows the matrix to specific expert-parallel degrees at dispatch + # time: "8" keeps every EP8 shard and drops EP16, so a comprehensive run can + # co-schedule the 8-GPU SKUs' single-node EP8 with the GB SKUs' two-node EP8 + # without dispatching any EP16 leg. Blank keeps every degree. The resulting + # matrix is a partial subset that only omits cases; it never reclassifies them. + selected_eps: set[int] = set() + for value in (part.strip() for part in ep_sizes.split(",")): + if not value: + continue + if not value.isdigit() or int(value) <= 0: + raise SystemExit(f"invalid --ep-sizes {ep_sizes!r}; expected positive integers") + selected_eps.add(int(value)) + if only_sku and only_sku not in cap.PLATFORMS: + raise SystemExit(f"unknown --only-sku {only_sku!r}; have {sorted(cap.PLATFORMS)}") + # --exclude-skus narrows the matrix to a subset by dropping whole runner pools + # — e.g. exclude a SKU whose cluster is unavailable. It only omits cases. + excluded = {value.strip() for value in exclude_skus.split(",") if value.strip()} + unknown_excluded = sorted(excluded - set(cap.PLATFORMS)) + if unknown_excluded: + raise SystemExit( + f"unknown --exclude-skus {unknown_excluded}; have {sorted(cap.PLATFORMS)}" + ) + if only_sku and only_sku in excluded: + raise SystemExit("--only-sku and --exclude-skus select disjoint pools") + + workloads = _load("workloads.yaml") + suites_document = _load("suites.yaml") + validate_config_documents(suites_document, workloads) + registry = suites_document["suites"] + select_all = suites == "all" + names = ( + list(registry) + if select_all + else [value.strip() for value in suites.split(",") if value.strip()] + ) + if not names or len(names) != len(set(names)): + raise SystemExit("suite selection must be non-empty and unique") + unknown = sorted(set(names) - set(registry)) + if unknown: + raise SystemExit(f"unknown suites {unknown}; have {sorted(registry)}") + targets = _select_backends(backend, backends) + + shards: dict[tuple[str, str, int], list[dict[str, Any]]] = {} + requested_cases: list[dict[str, Any]] = [] + scheduled: set[str] = set() + for suite_name in names: + suite = registry[suite_name] + mode = suite["mode"] + phases = suite["phases"] + routings = suite["routings"] + eplb_values = suite.get("eplb", [False]) + suite_backends = set(suite.get("backends", cap.SWEEP_BACKENDS)) + suite_targets = [target for target in targets if target in suite_backends] + if not suite_targets: + continue + for platform_name in suite["platforms"]: + if only_sku and platform_name != only_sku: + continue + if platform_name in excluded: + continue + ep_degrees = suite["ep_degrees"] + for workload, ep, phase, routing, eplb, target in itertools.product( + suite["workloads"], ep_degrees, phases, routings, eplb_values, + suite_targets, + ): + if selected_eps and ep not in selected_eps: + continue + topology = cap.topology_for(platform_name, ep) + if topology is None: + raise SystemExit( + f"suite {suite_name}: {platform_name} EP{ep} is not registered" + ) + nodes = int(topology["nodes"]) + capability_disposition, capability_detail = cap.resolve_disposition( + platform_name, + target, + ep=ep, + nodes=nodes, + routing=routing, + eplb=bool(eplb), + mode=mode, + ) + hidden, topk, experts = _dims(workloads, workload) + + def add_case( + case_ladder: str, + disposition: str, + reason: str | None, + detail: str | None, + ) -> None: + case: dict[str, Any] = { + "suite": suite_name, + "workload": workload, + "backend": target, + "routing": routing, + "phase": phase, + "ep": ep, + "eplb": eplb, + "hidden": hidden, + "topk": topk, + "experts": experts, + "samples_per_point": ep_harness.TIMED_SAMPLES_PER_POINT, + "warmup_semantics": ep_harness.WARMUP_SEMANTICS, + "ladder": case_ladder, + "mode": mode, + "timing": EP_TIMING_PROFILE, + "canonical": True, + **{field: topology[field] for field in TOPOLOGY_FIELDS}, + } + for signature in _semantic_points(platform_name, case): + if signature in scheduled: + raise SystemExit( + f"suite {suite_name}: duplicate semantic point for {platform_name}" + ) + scheduled.add(signature) + case["case_id"] = _case_id(platform_name, case) + requested_cases.append( + { + "sku": platform_name, + "case": case, + "disposition": disposition, + "reason": reason, + "detail": detail, + } + ) + if disposition == "runnable": + shards.setdefault((platform_name, target, nodes), []).append(case) + + requested_ladder = _ladder(suite, phase) + if capability_disposition == "unsupported": + add_case( + requested_ladder, + "unsupported", + "backend-platform-unsupported", + capability_detail, + ) + continue + if capability_disposition != "supported": + raise SystemExit( + f"suite {suite_name}: invalid capability disposition " + f"{capability_disposition!r}" + ) + add_case(requested_ladder, "runnable", None, None) + + shards_by_sku: dict[str, list[dict[str, Any]]] = {} + for (sku, target, nodes), cases in sorted(shards.items()): + chunk_size = max_cases + for offset in range(0, len(cases), chunk_size): + chunk = cases[offset:offset + chunk_size] + part = offset // chunk_size + shard_id = f"{sku}-{target}-n{nodes}" + if len(cases) > chunk_size: + shard_id += f"-p{part}" + shards_by_sku.setdefault(sku, []).append({ + "id": shard_id, + "sku": sku, + "backend": target, + "launcher": cap.PLATFORMS[sku]["launcher"], + **{field: chunk[0][field] for field in TOPOLOGY_FIELDS}, + "n": len(chunk), + "execution_weight": execution_weight(chunk), + "case_ids": [case["case_id"] for case in chunk], + }) + include = [ + shards_by_sku[sku][round_index] + for round_index in range(max(map(len, shards_by_sku.values()), default=0)) + for sku in sorted(shards_by_sku) + if round_index < len(shards_by_sku[sku]) + ] + return { + "format": "collectivex.matrix.v1", + "schema_version": 1, + "version": suites_document["version"], + "requested_cases": requested_cases, + "include": include, + } + + +def _strict_json_load(path: Path) -> Any: + def reject_constant(value: str) -> None: + raise MatrixError(f"non-finite JSON number {value}") + + def reject_duplicates(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise MatrixError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + if not path.is_file(): + raise MatrixError(f"matrix does not exist: {path}") + if path.stat().st_size == 0: + raise MatrixError(f"matrix is empty: {path}") + try: + with path.open() as fh: + return json.load( + fh, parse_constant=reject_constant, object_pairs_hook=reject_duplicates + ) + except (OSError, json.JSONDecodeError) as exc: + raise MatrixError(f"matrix is not valid JSON: {exc}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if type(value) is not int: + raise MatrixError(f"{field} must be a positive integer") + if value <= 0: + raise MatrixError(f"{field} must be a positive integer") + return value + + +def execution_weight(cases: list[dict[str, Any]]) -> int: + """Return deterministic GPU-point work used to bound workflow parallelism.""" + if not isinstance(cases, list) or not cases: + raise MatrixError("execution weight requires at least one case") + weight = 0 + for case in cases: + ep = _positive_int(case.get("ep"), "execution-weight.ep") + ladder = case.get("ladder") + if not isinstance(ladder, str) or not ladder.split(): + raise MatrixError("execution weight requires a token ladder") + weight += ep * len(ladder.split()) + return weight + + +def validate_shard_control( + shard: dict[str, Any], + *, + sku: str, + backend: str, + nodes: int, + require_runnable: bool = True, +) -> None: + """Validate one shard against the workflow cell that requested it.""" + if not isinstance(shard, dict): + raise MatrixError("shard must be a JSON object") + if sku not in cap.PLATFORMS or backend not in cap.SWEEP_BACKENDS: + raise MatrixError("shard platform/backend is not registered") + top_fields = { + "schema_version", "version", "id", "sku", "backend", "nodes", "n", "cases", + "execution_weight", + } + if ( + set(shard) != top_fields + or type(shard.get("schema_version")) is not int + or shard["schema_version"] != 1 + ): + raise MatrixError("shard fields or schema version differ from v1 contract") + if type(shard.get("version")) is not int or shard["version"] < 1: + raise MatrixError("shard benchmark version must be a positive integer") + if not isinstance(shard.get("id"), str) or not IDENTIFIER.fullmatch(shard["id"]): + raise MatrixError("shard has invalid id") + for field, expected in (("sku", sku), ("backend", backend)): + if shard.get(field) != expected: + raise MatrixError( + f"shard {field} mismatch: expected {expected!r}, got {shard.get(field)!r}" + ) + if _positive_int(shard.get("nodes"), "shard.nodes") != nodes: + raise MatrixError( + f"shard nodes mismatch: expected {nodes}, got {shard.get('nodes')!r}" + ) + cases = shard.get("cases") + if not isinstance(cases, list) or not cases: + raise MatrixError("shard must contain at least one case") + if _positive_int(shard.get("n"), "shard.n") != len(cases): + raise MatrixError("shard.n does not match the number of cases") + seen: set[str] = set() + required = { + "case_id", "suite", "workload", "backend", "routing", "mode", "phase", "ep", + "eplb", "hidden", "topk", "experts", "samples_per_point", "warmup_semantics", + "ladder", "timing", "canonical", + } | set(TOPOLOGY_FIELDS) + for index, case in enumerate(cases): + if not isinstance(case, dict): + raise MatrixError(f"case {index} must be a JSON object") + fields = set(case) + if fields != required: + raise MatrixError( + f"case {index} fields differ from v1 contract: " + f"missing={sorted(required - fields)}, extra={sorted(fields - required)}" + ) + case_id = case["case_id"] + if not identity.is_case_id(case_id): + raise MatrixError(f"case {index} has invalid case_id") + if case_id in seen: + raise MatrixError(f"duplicate case_id {case_id}") + seen.add(case_id) + string_fields = [ + "suite", "workload", "backend", "mode", "routing", "phase", + "warmup_semantics", "ladder", "timing", + ] + for field in string_fields: + if not isinstance(case[field], str) or not case[field]: + raise MatrixError(f"case {index}.{field} must be a non-empty string") + identifier_fields = ["suite", "workload", "backend", "routing", "phase"] + for field in identifier_fields: + if not IDENTIFIER.fullmatch(case[field]): + raise MatrixError(f"case {index}.{field} is not a safe identifier") + case_identity = {key: value for key, value in case.items() if key != "case_id"} + if case_id != _case_id(sku, case_identity): + raise MatrixError(f"case {index} case_id does not match its contents") + if case["backend"] != backend: + raise MatrixError(f"case {index} backend does not match shard") + if case["mode"] not in identity.V1_CASE_PROFILES: + raise MatrixError(f"case {index} mode is invalid") + if _positive_int(case["nodes"], f"case {index}.nodes") != nodes: + raise MatrixError(f"case {index} nodes does not match shard") + ep = _positive_int(case["ep"], f"case {index}.ep") + gpus_per_node = _positive_int( + case["gpus_per_node"], f"case {index}.gpus_per_node" + ) + topology = cap.topology_for(sku, ep) + if topology is None or any(case[field] != topology[field] for field in TOPOLOGY_FIELDS): + raise MatrixError(f"case {index} differs from the platform registry") + if ep != nodes * gpus_per_node: + raise MatrixError(f"case {index} ep does not equal nodes * gpus_per_node") + if case["samples_per_point"] != ep_harness.TIMED_SAMPLES_PER_POINT: + raise MatrixError(f"case {index} violates fixed-512-v1") + if case["timing"] != EP_TIMING_PROFILE: + raise MatrixError(f"case {index} has invalid timing profile") + if case["warmup_semantics"] != ep_harness.WARMUP_SEMANTICS: + raise MatrixError(f"case {index} has invalid warmup semantics") + if case["phase"] not in {"decode", "prefill"}: + raise MatrixError(f"case {index} has invalid phase") + if case["routing"] != "uniform": + raise MatrixError(f"case {index} has invalid routing") + if not isinstance(case["eplb"], bool) or case["eplb"]: + raise MatrixError(f"case {index} has invalid EPLB setting") + if not isinstance(case["canonical"], bool) or not case["canonical"]: + raise MatrixError(f"case {index} must use a canonical workload") + for field in ("ep", "nodes", "gpus_per_node", "hidden", "topk", "experts", + "samples_per_point", "scale_up_domain"): + if isinstance(case[field], bool) or not isinstance(case[field], int): + raise MatrixError(f"case {index}.{field} must be an integer") + _positive_int(case[field], f"case {index}.{field}") + scale_up_domain = _positive_int( + case["scale_up_domain"], f"case {index}.scale_up_domain" + ) + expected_scope = "scale-up" if ep <= scale_up_domain else "scale-out" + if case["scope"] != expected_scope or ( + expected_scope == "scale-out" and ep % scale_up_domain + ): + raise MatrixError(f"case {index} has invalid scale-up/scale-out geometry") + try: + ladder = [int(value) for value in case["ladder"].split()] + except (AttributeError, ValueError) as exc: + raise MatrixError(f"case {index} has invalid token ladder") from exc + if (not ladder or any(value <= 0 for value in ladder) + or ladder != sorted(set(ladder)) + or case["ladder"] != " ".join(map(str, ladder))): + raise MatrixError(f"case {index} has invalid token ladder") + if require_runnable: + disposition, reason, _ = _expected_disposition(sku, case) + if disposition != "runnable": + raise MatrixError(f"case {index} violates capability registry: {reason}") + else: + _v1_requested_ladder(case) + if _positive_int( + shard.get("execution_weight"), "shard.execution_weight" + ) != execution_weight(cases): + raise MatrixError("shard execution_weight differs from its cases") + + +def validate_matrix_document(document: Any) -> dict[str, Any]: + """Validate the complete requested grid and its runnable shard partition.""" + if not isinstance(document, dict) or set(document) != { + "format", "schema_version", "version", "requested_cases", "include" + }: + raise MatrixError("matrix fields differ from the v1 contract") + if ( + document["format"] != "collectivex.matrix.v1" + or type(document["schema_version"]) is not int + or document["schema_version"] != 1 + ): + raise MatrixError("matrix format/schema differs from v1") + if type(document["version"]) is not int or document["version"] < 1: + raise MatrixError("matrix benchmark version must be a positive integer") + requested = document["requested_cases"] + include = document["include"] + if not isinstance(requested, list) or not requested: + raise MatrixError("matrix.requested_cases must be non-empty") + if not isinstance(include, list): + raise MatrixError("matrix.include must be an array") + + cases_by_id: dict[str, dict[str, Any]] = {} + runnable_ids: set[str] = set() + semantic_points: set[str] = set() + for index, value in enumerate(requested): + path = f"matrix.requested_cases[{index}]" + if not isinstance(value, dict) or set(value) != { + "sku", "case", "disposition", "reason", "detail" + }: + raise MatrixError(f"{path} fields differ from the v1 contract") + sku = value["sku"] + case = value["case"] + disposition = value["disposition"] + if sku not in cap.PLATFORMS: + raise MatrixError(f"{path}.sku is unknown") + if disposition not in {"runnable", "unsupported"}: + raise MatrixError(f"{path}.disposition is invalid") + if disposition == "runnable": + if value["reason"] is not None or value["detail"] is not None: + raise MatrixError(f"{path} runnable cases cannot have a reason") + else: + if ( + not isinstance(value["reason"], str) + or not IDENTIFIER.fullmatch(value["reason"]) + or not isinstance(value["detail"], str) + or not value["detail"] + ): + raise MatrixError(f"{path} unsupported cases need a public reason and detail") + if not isinstance(case, dict): + raise MatrixError(f"{path}.case must be an object") + backend = case.get("backend") + nodes = case.get("nodes") + if not isinstance(backend, str) or type(nodes) is not int: + raise MatrixError(f"{path}.case backend/nodes are invalid") + requested_case_plan = [case] + validate_shard_control( + { + "schema_version": 1, + "version": document["version"], + "id": "requested-case", + "sku": sku, + "backend": backend, + "nodes": nodes, + "n": 1, + "execution_weight": execution_weight(requested_case_plan), + "cases": requested_case_plan, + }, + sku=sku, + backend=backend, + nodes=nodes, + require_runnable=disposition == "runnable", + ) + case_id = case["case_id"] + if case_id in cases_by_id: + raise MatrixError(f"duplicate requested case_id {case_id}") + for signature in _semantic_points(sku, case): + if signature in semantic_points: + raise MatrixError(f"{path} duplicates a semantic token point") + semantic_points.add(signature) + cases_by_id[case_id] = value + expected = _expected_disposition(sku, case) + if (disposition, value["reason"], value["detail"]) != expected: + raise MatrixError(f"{path} disposition differs from the frozen v1 catalog") + if disposition == "runnable": + runnable_ids.add(case_id) + + shard_ids: set[str] = set() + assigned: list[str] = [] + for index, shard in enumerate(include): + path = f"matrix.include[{index}]" + expected = { + "id", "sku", "backend", "launcher", "n", "execution_weight", "case_ids", + } | set(TOPOLOGY_FIELDS) + if not isinstance(shard, dict) or set(shard) != expected: + raise MatrixError(f"{path} fields differ from the v1 contract") + shard_id = shard["id"] + if not isinstance(shard_id, str) or not IDENTIFIER.fullmatch(shard_id): + raise MatrixError(f"{path}.id is invalid") + if shard_id in shard_ids: + raise MatrixError(f"duplicate shard id {shard_id}") + shard_ids.add(shard_id) + sku = shard["sku"] + if sku not in cap.PLATFORMS: + raise MatrixError(f"{path}.sku is unknown") + platform = cap.PLATFORMS[sku] + if shard["launcher"] != platform["launcher"]: + raise MatrixError(f"{path}.launcher differs from the platform registry") + case_ids = shard["case_ids"] + if not isinstance(case_ids, list) or not case_ids or len(case_ids) != len(set(case_ids)): + raise MatrixError(f"{path}.case_ids must be a non-empty unique array") + if _positive_int(shard["n"], f"{path}.n") != len(case_ids): + raise MatrixError(f"{path}.n differs from case_ids") + nodes = _positive_int(shard["nodes"], f"{path}.nodes") + for case_id in case_ids: + wrapper = cases_by_id.get(case_id) + if wrapper is None or wrapper["disposition"] != "runnable": + raise MatrixError(f"{path} references a missing or unsupported case") + case = wrapper["case"] + if ( + wrapper["sku"] != sku + or case["backend"] != shard["backend"] + or case["nodes"] != nodes + or any(shard[field] != case[field] for field in TOPOLOGY_FIELDS) + ): + raise MatrixError(f"{path} case does not match shard coordinates") + assigned.append(case_id) + if shard["execution_weight"] != execution_weight( + [cases_by_id[case_id]["case"] for case_id in case_ids] + ): + raise MatrixError(f"{path}.execution_weight differs from its cases") + if len(assigned) != len(set(assigned)): + raise MatrixError("a runnable case is assigned to more than one shard") + if set(assigned) != runnable_ids: + raise MatrixError("runnable requested cases and shard assignments differ") + return document + + +def extract_shard( + matrix_path: str | os.PathLike[str], + shard_id: str, + output_path: str | os.PathLike[str], + *, + sku: str, + backend: str, + nodes: int, +) -> dict[str, Any]: + """Extract one strictly matched shard control file, writing it atomically.""" + document = validate_matrix_document(_strict_json_load(Path(matrix_path))) + include = document["include"] + matches = [item for item in include if isinstance(item, dict) and item.get("id") == shard_id] + if len(matches) != 1: + raise MatrixError(f"expected exactly one shard {shard_id!r}, found {len(matches)}") + source = matches[0] + requested = { + item["case"]["case_id"]: item + for item in document["requested_cases"] + } + cases = [requested[case_id]["case"] for case_id in source["case_ids"]] + control = { + "schema_version": 1, + "version": document["version"], + "id": source.get("id"), + "sku": source.get("sku"), + "backend": source.get("backend"), + "nodes": source.get("nodes"), + "n": source.get("n"), + "execution_weight": source.get("execution_weight"), + "cases": cases, + } + validate_shard_control(control, sku=sku, backend=backend, nodes=nodes) + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_name(f".{output.name}.tmp-{os.getpid()}") + try: + with temporary.open("w") as fh: + json.dump(control, fh, sort_keys=True, separators=(",", ":")) + fh.write("\n") + os.replace(temporary, output) + finally: + temporary.unlink(missing_ok=True) + return control + + +def main() -> int: + parser = argparse.ArgumentParser(description="CollectiveX matrix resolver") + parser.add_argument("--suites", default="all", help="'all' or comma-list of suites") + parser.add_argument("--backend", default="", help="select one EP backend") + parser.add_argument("--backends", default="", help="'all' or comma-list of EP backends") + parser.add_argument("--only-sku", default="") + parser.add_argument( + "--exclude-skus", + default="", + help="comma-list of runner pools to drop (partial matrix); disjoint from --only-sku", + ) + parser.add_argument( + "--ep-sizes", + default="", + help="comma-list of expert-parallel degrees to keep (e.g. 8 drops EP16); blank = all", + ) + parser.add_argument("--max-cases", type=int, default=128) + parser.add_argument("--extract-from", default="", metavar="MATRIX") + parser.add_argument("--validate-control", default="", metavar="SHARD") + parser.add_argument("--shard-id", default="") + parser.add_argument("--expect-sku", default="") + parser.add_argument("--expect-backend", default="") + parser.add_argument("--expect-nodes", type=int, default=0) + parser.add_argument("--out", default="") + args = parser.parse_args() + + if args.validate_control: + if not all((args.expect_sku, args.expect_backend, args.expect_nodes)): + parser.error( + "control validation requires --expect-sku, --expect-backend, and --expect-nodes" + ) + try: + control = _strict_json_load(Path(args.validate_control)) + validate_shard_control( + control, + sku=args.expect_sku, + backend=args.expect_backend, + nodes=args.expect_nodes, + ) + except MatrixError as exc: + parser.error(str(exc)) + print(f"validated {control.get('id')}: {control['n']} cases", file=sys.stderr) + return 0 + + if args.extract_from: + if not all((args.shard_id, args.expect_sku, args.expect_backend, args.expect_nodes, args.out)): + parser.error( + "shard extraction requires --shard-id, --expect-sku, --expect-backend, " + "--expect-nodes, and --out" + ) + try: + control = extract_shard( + args.extract_from, + args.shard_id, + args.out, + sku=args.expect_sku, + backend=args.expect_backend, + nodes=args.expect_nodes, + ) + except MatrixError as exc: + parser.error(str(exc)) + print(f"extracted {control['id']}: {control['n']} cases", file=sys.stderr) + print(json.dumps(control, separators=(",", ":"))) + return 0 + + matrix = resolve_matrix( + suites=args.suites, + backend=args.backend, + backends=args.backends, + only_sku=args.only_sku, + exclude_skus=args.exclude_skus, + ep_sizes=args.ep_sizes, + max_cases=args.max_cases, + ) + try: + validate_matrix_document(matrix) + except MatrixError as exc: + parser.error(str(exc)) + if args.out: + with open(args.out, "w") as fh: + json.dump(matrix, fh, sort_keys=True, separators=(",", ":")) + fh.write("\n") + runnable = sum( + item["disposition"] == "runnable" for item in matrix["requested_cases"] + ) + unsupported = len(matrix["requested_cases"]) - runnable + print( + f"resolved {len(matrix['include'])} shard-cells, " + f"{runnable} runnable and {unsupported} unsupported cases", + file=sys.stderr, + ) + print(json.dumps(matrix)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/experimental/CollectiveX/tests/test_ep_backend.py b/experimental/CollectiveX/tests/test_ep_backend.py new file mode 100644 index 0000000000..503ac3e1c7 --- /dev/null +++ b/experimental/CollectiveX/tests/test_ep_backend.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +"""Small torch-free smoke tests for the shared EP backend lifecycle.""" +from __future__ import annotations + +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT), str(ROOT / "bench")] + +import ep_backend # noqa: E402 +from ep_backend import EPBackend, RankInputs # noqa: E402 + + +def args(**updates): + values = dict( + experts=8, phase="decode", tokens_ladder="", routing="uniform", seed=0, + hidden=16, topk=2, mode="normal", workload_dir="", + ) + values.update(updates) + return types.SimpleNamespace(**values) + + +class FakeBackend(EPBackend): + name = "fake" + + def __init__(self, options, *, cap=None, world_size=1): + super().__init__(options, 0, world_size, 0, "cpu") + self.cap = cap + self.calls: list[str] = [] + + def create_buffer(self, spec): + return None + + def dispatch(self, problem): + self.calls.append("dispatch") + return object() + + def stage(self, problem, handle): + self.calls.append("stage") + + def combine(self, problem, handle): + self.calls.append("combine") + + def recv_tokens(self, handle): + return 0 + + def inspect_dispatch(self, problem, handle): + return None + + def combine_transformed(self, problem, handle, transformed): + return None + + def buffer_cap(self, options): + return self.cap + + def _build_rank_inputs(self, options, tokens, *, canonical, retain_global): + return RankInputs( + tokens_per_rank=tokens, topk_idx=None, topk_weights=None, + activations=None, workload_id=f"workload-{tokens}" if canonical else None, + checksums={"tokens": tokens}, + ) + + +class BackendTests(unittest.TestCase): + def test_input_plan_sizes_for_measured_and_conditioning_ladders(self): + backend = FakeBackend(args(tokens_ladder="8 16", workload_dir="/tmp"), world_size=2) + spec = backend.make_inputs(backend.args) + self.assertTrue(spec.ok) + self.assertEqual(spec.ladder, [8, 16]) + self.assertEqual(spec.max_tokens_per_rank, 128) + self.assertEqual((spec.ep_size, spec.experts_per_rank), (2, 4)) + self.assertEqual(sorted(spec.points), [8, 16]) + + def test_invalid_ladder_or_small_buffer_fails_before_execution(self): + for backend, message in ( + (FakeBackend(args(tokens_ladder="0")), "empty token ladder"), + (FakeBackend(args(), cap=64), "buffer cap 64"), + ): + with self.subTest(message=message): + spec = backend.make_inputs(backend.args) + self.assertEqual(spec.rc, 2) + self.assertIn(message, spec.message) + + def test_timed_components_follow_backend_contract(self): + backend = FakeBackend(args()) + self.assertEqual(backend.timed_components(), ["roundtrip", "dispatch", "combine"]) + backend.stage_device_work = True + self.assertEqual( + backend.timed_components(), ["roundtrip", "dispatch", "combine", "stage"] + ) + backend.roundtrip_only = True + self.assertEqual(backend.timed_components(), ["roundtrip"]) + + def test_dispatch_cleanup_is_outside_timed_call(self): + backend = FakeBackend(args()) + backend.dispatch_needs_combine_cleanup = True + captured = {} + + def fake_time(_torch, operation, _warmup, _iters, **kwargs): + handle = operation() + kwargs["post"](handle) + captured.update(kwargs) + return [1.0] + + with mock.patch.dict(sys.modules, {"torch": types.SimpleNamespace()}), mock.patch.object( + ep_backend, "time_us", side_effect=fake_time + ): + backend.benchmark_dispatch(object(), 0, 1) + self.assertIn("post", captured) + self.assertEqual(backend.calls, ["dispatch", "stage", "combine"]) + + def test_mode_is_fail_closed(self): + with self.assertRaises(ValueError): + FakeBackend(args(mode="low-latency")) + + +if __name__ == "__main__": + unittest.main() diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py new file mode 100644 index 0000000000..490d580ede --- /dev/null +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Matrix, subset, and shard-extraction tests.""" +from __future__ import annotations + +import json +import sys +import tempfile +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import sweep_matrix # noqa: E402 + + +def matrix(**options): + return sweep_matrix.validate_matrix_document( + sweep_matrix.resolve_matrix(max_cases=128, **options) + ) + + +class MatrixTests(unittest.TestCase): + def test_shard_extraction_is_deterministic_and_preserves_cases(self): + document = matrix(backend="deepep", only_sku="h100-dgxc") + cell = document["include"][0] + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + source = root / "matrix.json" + source.write_text(json.dumps(document, sort_keys=True)) + outputs = [ + sweep_matrix.extract_shard( + source, cell["id"], root / f"shard-{index}.json", + sku=cell["sku"], backend=cell["backend"], nodes=cell["nodes"], + ) + for index in range(2) + ] + self.assertEqual(outputs[0], outputs[1]) + self.assertEqual([case["case_id"] for case in outputs[0]["cases"]], cell["case_ids"]) + + def test_sku_and_ep_filters_only_remove_cases(self): + full = matrix(suites="all", backends="all") + for options, keep in ( + ({"exclude_skus": "b300"}, lambda item: item["sku"] != "b300"), + ({"ep_sizes": "8"}, lambda item: item["case"]["ep"] == 8), + ): + partial = matrix(suites="all", backends="all", **options) + expected = { + item["case"]["case_id"]: item for item in full["requested_cases"] if keep(item) + } + actual = {item["case"]["case_id"]: item for item in partial["requested_cases"]} + self.assertEqual(actual, expected) + + def test_invalid_filters_fail_closed(self): + for options in ( + {"exclude_skus": "unknown"}, + {"only_sku": "b300", "exclude_skus": "b300"}, + {"ep_sizes": "0"}, + {"ep_sizes": "eight"}, + ): + with self.subTest(options=options), self.assertRaises(SystemExit): + sweep_matrix.resolve_matrix(**options) + + +if __name__ == "__main__": + unittest.main()