#!/usr/bin/env bash
set -euo pipefail

world_size="${1:-${H3_WORLD_SIZE:-}}"
if [[ -z "$world_size" ]]; then
  world_size="$(python -c 'import torch; print(torch.cuda.device_count())')"
fi
if [[ "$world_size" -lt 1 ]]; then
  echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
  exit 1
fi
backend="${2:-sdpa}"
sequence="${3:-20000}"
output="${H3_BENCHMARK_OUTPUT:-/output/h3-baselines/ulysses-${world_size}gpu-${backend}-${sequence}t.json}"

torchrun \
  --standalone \
  --nnodes=1 \
  --nproc-per-node="$world_size" \
  tools/benchmark_ulysses.py \
  --sequence "$sequence" \
  --backend "$backend" \
  --output "$output"