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

available="$(python -c 'import torch; print(torch.cuda.device_count())')"
if [[ "$available" -lt 1 ]]; then
  echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
  exit 1
fi

read -ra counts <<< "${H3_GPU_COUNTS:-1 2 4 6 8}"
read -ra modes <<< "${H3_DISTRIBUTED_MODES:-ulysses tensor}"
for count in "${counts[@]}"; do
  if [[ "$count" -gt "$available" ]]; then
    echo "Skipping ${count} GPUs; only ${available} are visible." >&2
    continue
  fi
  for mode in "${modes[@]}"; do
    tools/run_distributed_t2va.sh "$mode" "$count" "${H3_ATTENTION:-sdpa}"
  done
done