Skip to content

Commit 4bdc125

Browse files
functionstackxclaude
functionstackx
and
claude
authored
Add OperatorX root files, scripts/, and testlists/ to experimental/operatorx (#2236)
Copies the remaining OperatorX repo contents (at 3afea73) missed by the initial import: .gitignore, CLUSTERS.md, README.md, containers.toml, pyproject.toml, uv.lock, scripts/, and testlists/. Only results/ is excluded. Package layout stays flat. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent 886a0db commit 4bdc125

28 files changed

Lines changed: 215599 additions & 0 deletions

‎experimental/operatorx/.gitignore‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Python
2+
__pycache__/
3+
*.py[cod]
4+
*.egg-info/
5+
.eggs/
6+
build/
7+
dist/
8+
.venv/
9+
venv/
10+
11+
# Runtime / caches
12+
logs/
13+
*.log
14+
.pytest_cache/
15+
.triton/
16+
triton_cache/
17+
18+
# OS
19+
.DS_Store

‎experimental/operatorx/CLUSTERS.md‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
# Clusters
2+
3+
Reference for the clusters operatorx targets: hardware, how the runner is
4+
launched per platform, and the per-platform quirks. Access details
5+
(hostnames, credentials, checkout paths) are deployment-specific and are not
6+
included here — fill them in for your own environment.
7+
8+
Platform → cluster routing lives in `operatorx/clusters.py` (`CLUSTER_PLATFORMS`);
9+
SLURM/env defaults live in `scripts/submit_run.py` (`DEFAULT_CLUSTER`).
10+
11+
## Conventions
12+
13+
- `submit_run.py` is **SLURM-only** — it submits `sbatch`/`srun` jobs and is
14+
used on the NVIDIA and AMD clusters. On TPU / Trainium hosts there is no
15+
SLURM; run `python -m operatorx` directly inside the appropriate venv/image.
16+
- Point `OPERATORX_SQUASH_DIR` at your container squash directory, and set
17+
`OPERATORX_PARTITION` / `OPERATORX_ACCOUNT` / `OPERATORX_QOS` to your
18+
cluster's SLURM values (see the env-var table at the bottom).
19+
- Set `OPERATORX_JOB_NAME` to a recognizable SLURM job name for your runs.
20+
21+
## At a glance
22+
23+
| Cluster id | Platform | Hardware | Scheduler |
24+
|---------------|----------|---------------------------------|------------|
25+
| `b200_dgx_8x` | nvidia | 8× B200 SXM (DGX-style) | SLURM |
26+
| `b300_hgx_8x` | nvidia | 8× B300 (HGX-style) | SLURM |
27+
| `b200_nvl72` | nvidia | B200 NVL72 (routing only) | SLURM |
28+
| `mi355x_8x` | amd | 8× MI355 OAM per node | SLURM |
29+
| `v6e_1x` | tpu | 1× TPU v6e (1×1) | run direct |
30+
| `v6e_4x` | tpu | 4× TPU v6e (2×2) | run direct |
31+
| `v6e_pod` | tpu | TPU v6e pod slice | run direct |
32+
| `trn3_1x` | trainium | 1 LNC of a trn3 host | run direct |
33+
| `trn3_8x` | trainium | trn3 subset | run direct |
34+
| `trn3_16x` | trainium | full trn3 host (16 devices) | run direct |
35+
36+
## NVIDIA
37+
38+
### b200 — DGX-style, 8× B200 SXM (`b200_dgx_8x`)
39+
40+
| Setting | Value |
41+
|----------|----------------------------------------------------------|
42+
| Hardware | 8× B200 SXM per node |
43+
| GRES | e.g. `gpu:nvidia_b200:8` |
44+
| Defaults | `submit_run.py` defaults target this cluster (`OPERATORX_CLUSTER=b200_dgx_8x`, partition/squash dir from the script/env) |
45+
46+
```bash
47+
OPERATORX_JOB_NAME=<job-name> python3 scripts/submit_run.py nvidia
48+
```
49+
50+
### b300 — HGX-style, 8× B300 (`b300_hgx_8x`)
51+
52+
| Setting | Value |
53+
|-------------------|--------------------------------------------------------------|
54+
| Hardware | 8× B300 per node |
55+
| OS note | If the login node runs Python 3.10, `submit_run.py` falls back to `tomli` for TOML parsing |
56+
| Excluded backends | DeepEP — known IBGDA hangs on this fabric, so exclude it |
57+
58+
Override the SLURM knobs for your cluster and drop DeepEP:
59+
60+
```bash
61+
OPERATORX_CLUSTER=b300_hgx_8x \
62+
OPERATORX_PARTITION=<partition> \
63+
OPERATORX_ACCOUNT=<account> \
64+
OPERATORX_QOS=<qos> \
65+
OPERATORX_SQUASH_DIR=<squash-dir> \
66+
OPERATORX_BACKENDS=torch,deepgemm,flashinfer,sglang \
67+
OPERATORX_JOB_NAME=<job-name> \
68+
python3 scripts/submit_run.py nvidia
69+
```
70+
71+
### `b200_nvl72`
72+
73+
Present in `CLUSTER_PLATFORMS` (routes to nvidia) as a placeholder; no
74+
hardware-specific guidance yet.
75+
76+
## AMD — `mi355x_8x`
77+
78+
| Setting | Value |
79+
|-----------|--------------------------------------------------------------|
80+
| Hardware | 8× MI355 OAM per node |
81+
| GRES | e.g. `gpu:amd_instinct_mi355_oam:8` |
82+
| Backends | `containers.toml` `amd.torch` / `amd.aiter` → `vllm-openai-rocm` |
83+
84+
```bash
85+
OPERATORX_CLUSTER=mi355x_8x \
86+
OPERATORX_PARTITION=<partition> \
87+
OPERATORX_SQUASH_DIR=<squash-dir> \
88+
OPERATORX_JOB_NAME=<job-name> \
89+
python3 scripts/submit_run.py amd
90+
```
91+
92+
## TPU
93+
94+
No SLURM — run `python -m operatorx` directly on the TPU VM. Access a VM with
95+
`gcloud compute tpus tpu-vm ssh <vm-name> --zone=<zone>`.
96+
97+
| Cluster id | Topology | Chips | ws range | Use for |
98+
|------------|----------|-------|----------|-----------------------------------------------------|
99+
| `v6e_1x` | 1×1 | 1 | 1 | Single-chip tests (gemm, moe_gemm). No collectives. |
100+
| `v6e_4x` | 2×2 | 4 | 1–4 | Collectives + MoE-EP up to ws=4. |
101+
| `v6e_pod` | pod | pod | — | Multi-host v6e pod slice. |
102+
103+
```bash
104+
OPERATORX_CLUSTER=v6e_4x python -m operatorx
105+
```
106+
107+
`moe_forward` on TPU emits `unsupported` unless Google's MaxText is installed
108+
on the VM (the `jax` backend works without it):
109+
110+
```bash
111+
git clone https://github.com/AI-Hypercomputer/maxtext ~/maxtext
112+
pip install -e ~/maxtext
113+
```
114+
115+
## Trainium — `trn3_16x`
116+
117+
No SLURM — run `python -m operatorx` directly on the instance.
118+
119+
| Setting | Value |
120+
|---------------------|-------------------------------------------------------------------|
121+
| Hardware | 16 Neuron devices × 4 cores (64 physical NCs, 8 LNC=2 logical units) |
122+
| Per-device HBM | 144 GiB |
123+
| Cluster ids | `trn3_16x` (full host); `trn3_8x`, `trn3_1x` are subset configs |
124+
| Device topology cmd | `/opt/aws/neuron/bin/neuron-ls` (may not be on `PATH`) |
125+
126+
Neuron venvs (standard AWS Neuron paths under `/opt/`):
127+
128+
- `aws_neuronx_venv_pytorch_2_9` — base PyTorch + NeuronX
129+
- `aws_neuronx_venv_pytorch_2_9_nxd_inference` — adds NeuronX Distributed Inference
130+
- `aws_neuronx_venv_pytorch_2_9_nxd_training` — adds NXD training
131+
- `aws_neuronx_venv_pytorch_inference_vllm_0_16` — vLLM build
132+
- `aws_neuronx_venv_jax_0_7` — Jax (not used by operatorx)
133+
134+
```bash
135+
source /opt/aws_neuronx_venv_pytorch_2_9/bin/activate
136+
OPERATORX_CLUSTER=trn3_16x python -m operatorx
137+
```
138+
139+
The Trainium runner forces `NEURON_LOGICAL_NC_CONFIG=2` and appends
140+
`--logical-nc-config=2` to `NEURON_CC_FLAGS`
141+
(`operatorx/runners/trainium/runner.py`). `world_size=1` on this platform =
142+
one LNC=2 unit (1 HBM bank + 2 paired NC-v4 cores).
143+
144+
## Cluster id → platform (`operatorx/clusters.py`)
145+
146+
```python
147+
CLUSTER_PLATFORMS = {
148+
"b200_dgx_8x": "nvidia",
149+
"b300_hgx_8x": "nvidia",
150+
"b200_nvl72": "nvidia",
151+
"mi355x_8x": "amd",
152+
"v6e_1x": "tpu",
153+
"v6e_4x": "tpu",
154+
"v6e_pod": "tpu",
155+
"trn3_1x": "trainium",
156+
"trn3_8x": "trainium",
157+
"trn3_16x": "trainium",
158+
}
159+
```
160+
161+
## `submit_run.py` env vars
162+
163+
| Var | Default | Notes |
164+
|------------------------|-------------------------------|---------------------------------------------------------------------|
165+
| `OPERATORX_CLUSTER` | per-platform (see script) | Routes the runner via `CLUSTER_PLATFORMS` |
166+
| `OPERATORX_PARTITION` | site default | SLURM partition |
167+
| `OPERATORX_ACCOUNT` | (omitted) | SLURM `--account` |
168+
| `OPERATORX_QOS` | (omitted) | SLURM `--qos` |
169+
| `OPERATORX_SQUASH_DIR` | site default | Where `<safe_image>.sqsh` lives (used by `srun --container-image=`) |
170+
| `OPERATORX_BACKENDS` | all backends for the platform | CSV allowlist |
171+
| `OPERATORX_JOB_NAME` | `benchmark` | SLURM job name |
172+
| `OPERATORX_TESTLISTS` | all under `testlists/` | CSV of testlist stems to run |
173+
| `OPERATORX_TIME_MIN` | `30` | `--time` (minutes) |
174+
175+
`WORLD_SIZES = [1, 2, 4, 8]` — single-node only; multi-node NCCL IB bring-up
176+
currently hangs on the B200/B300 fabrics. `MASTER_ADDR` is derived by parsing
177+
`SLURM_NODELIST`.

‎experimental/operatorx/README.md‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# operatorx
2+
3+
Multi-platform inference operator benchmark suite. Times one op at a time
4+
(gemm, attention, moe, collectives, ...) on NVIDIA / AMD / TPU / Trainium and
5+
emits one JSON per run under `results/<platform>/<cluster>/`.
6+
7+
See `CLUSTERS.md` for how to reach each cluster and the per-host quirks.
8+
9+
## Running on a cluster
10+
11+
`scripts/submit_run.py <platform>` fans out one sbatch job per
12+
`(container_image, world_size)` pair to your local SLURM. Each job runs
13+
`python -m operatorx` inside the container.
14+
15+
### b200 (DGX-style, 8x B200 SXM)
16+
17+
```bash
18+
ssh tailscale-b200
19+
cd /home/sa-shared/harrison/oss-inference-tracker/operatorx
20+
python3 scripts/submit_run.py nvidia
21+
```
22+
23+
Defaults that apply: `OPERATORX_CLUSTER=b200_dgx_8x`,
24+
`OPERATORX_PARTITION=gpu-2`, `OPERATORX_SQUASH_DIR=/home/sa-shared/containers`.
25+
26+
### b300 (HGX-style, 8x B300)
27+
28+
The b300 cluster needs a non-default partition + account + qos, has its own
29+
squash dir, and DeepEP has known IBGDA issues here so we exclude it.
30+
31+
```bash
32+
ssh tailscale-b300
33+
cd /data/home/sa-shared/harrison/oss-inference-tracker/operatorx
34+
OPERATORX_CLUSTER=b300_hgx_8x \
35+
OPERATORX_PARTITION=batch_1 \
36+
OPERATORX_ACCOUNT=benchmark \
37+
OPERATORX_QOS=batch_1_qos \
38+
OPERATORX_SQUASH_DIR=/data/home/sa-shared/harrison/containers \
39+
OPERATORX_BACKENDS=torch,deepgemm,flashinfer,sglang \
40+
python3 scripts/submit_run.py nvidia
41+
```
42+
43+
### TPU / Trainium
44+
45+
These hosts have no SLURM, so `submit_run.py` (which submits `sbatch` jobs)
46+
does not apply — run the benchmark directly on the VM/instance:
47+
48+
```bash
49+
OPERATORX_CLUSTER=v6e_4x python -m operatorx # TPU (default tpu cluster)
50+
OPERATORX_CLUSTER=trn3_16x python -m operatorx # Trainium (default trainium cluster)
51+
```
52+
53+
The TPU `maxtext` backend depends on Google's MaxText library. Install it
54+
once per TPU VM (the `jax` backend works without it):
55+
56+
```bash
57+
git clone https://github.com/AI-Hypercomputer/maxtext ~/maxtext
58+
pip install -e ~/maxtext
59+
```
60+
61+
If MaxText isn't installed, `moe_forward` on TPU emits `unsupported` rows
62+
rather than running our old single-device dense fallback.
63+
64+
## Env vars honored by `submit_run.py`
65+
66+
| Var | Default | Notes |
67+
|-----|---------|-------|
68+
| `OPERATORX_CLUSTER` | per-platform (see script) | Cluster id used to route the runner (see `operatorx.clusters.CLUSTER_PLATFORMS`). |
69+
| `OPERATORX_PARTITION` | `gpu-2` | SLURM partition. |
70+
| `OPERATORX_ACCOUNT` | (omitted) | SLURM `--account`. |
71+
| `OPERATORX_QOS` | (omitted) | SLURM `--qos`. |
72+
| `OPERATORX_SQUASH_DIR` | `/home/sa-shared/containers` | Where `<safe_image>.sqsh` lives. |
73+
| `OPERATORX_BACKENDS` | all backends for the platform | CSV allowlist. |
74+
| `OPERATORX_JOB_NAME` | `benchmark` | SLURM job name. Use `h-benchmark` for benchmark runs (see `CLUSTERS.md`). |
75+
76+
`WORLD_SIZES` in the script is `[1, 2, 4, 8]` — single-node only (ws>8 is
77+
disabled: multi-node NCCL IB bring-up currently hangs on b200/b300).
78+
`MASTER_ADDR` is derived by parsing `SLURM_NODELIST`.
79+
80+
## Adding a backend / op
81+
82+
- Backend impl: `operatorx/runners/<platform>/backends/<name>.py`, exports an
83+
`IMPLS = [BackendImpl(...)]` list.
84+
- Op spec: `operatorx/ops/<name>.py`, calls `register(OpSpec(..., flops=, bytes=))`.
85+
- Add the container image to `containers.toml`, then on each cluster run
86+
`python3 scripts/pull_containers.py nvidia` to enroot-import it into
87+
`$OPERATORX_SQUASH_DIR`.
88+
- Add shapes to `testlists/<name>.json`.

‎experimental/operatorx/containers.toml‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
[nvidia.torch]
2+
image = "nvcr.io/nvidia/pytorch:26.03-py3"
3+
4+
[nvidia.flashinfer]
5+
image = "lmsysorg/sglang:latest"
6+
7+
[nvidia.deepgemm]
8+
image = "lmsysorg/sglang:latest"
9+
10+
[nvidia.sglang]
11+
image = "lmsysorg/sglang:latest"
12+
13+
[nvidia.deepep]
14+
image = "lmsysorg/sglang:latest"
15+
16+
[nvidia.flashinfer_comm]
17+
image = "lmsysorg/sglang:latest"
18+
19+
[nvidia.sglang_comm]
20+
image = "lmsysorg/sglang:latest"
21+
22+
[amd.torch]
23+
image = "vllm/vllm-openai-rocm:v0.19.0"
24+
25+
[amd.aiter]
26+
image = "vllm/vllm-openai-rocm:v0.19.0"
27+
28+
[tpu.jax]
29+
image = "us-central1-docker.pkg.dev/jax-image-build/jax-tpu/jax-tpu:latest"
30+
31+
[trainium.neuronx]
32+
image = "public.ecr.aws/neuron/pytorch-training-neuronx:2.5.1-neuronx-py310-sdk2.21.0-ubuntu22.04"

‎experimental/operatorx/pyproject.toml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
[build-system]
2+
requires = ["setuptools>=68"]
3+
build-backend = "setuptools.build_meta"
4+
5+
[project]
6+
name = "operatorx"
7+
version = "0.1.0"
8+
requires-python = ">=3.10"
9+
dependencies = []
10+
11+
[tool.setuptools.packages.find]
12+
include = ["operatorx*"]

0 commit comments

Comments
 (0)