#!/bin/bash
set -e

NUM_RANKS=${NUM_RANKS:-2}
NUM_TOKENS=${NUM_TOKENS:-128}
HIDDEN=${HIDDEN:-4096}
INTERMEDIATE=${INTERMEDIATE:-2048}
NUM_EXPERTS=${NUM_EXPERTS:-16}
NUM_TOPK=${NUM_TOPK:-6}
OUTPUT_DIR=${OUTPUT_DIR:-./nsys-results}

mkdir -p "$OUTPUT_DIR"

TIMESTAMP=$(date +%Y%m%d_%H%M%S)
REPORT_NAME="nsys_sm90_megamoe_t${NUM_TOKENS}_r${NUM_RANKS}_${TIMESTAMP}"

echo "=== SM90 MegaMoE nsys profiling ==="
echo "  Ranks: $NUM_RANKS"
echo "  Tokens: $NUM_TOKENS"
echo "  Hidden: $HIDDEN"
echo "  Intermediate: $INTERMEDIATE"
echo "  Experts: $NUM_EXPERTS"
echo "  TopK: $NUM_TOPK"
echo "  Output: $OUTPUT_DIR/$REPORT_NAME"
echo ""

NSYS_FLAGS=(
    --trace=cuda,nvtx
    --cuda-memory-usage=true
    --force-overwrite=true
    --output="$OUTPUT_DIR/$REPORT_NAME"
    --stats=true
)

TORCH_FLAGS=(
    --nproc_per_node="$NUM_RANKS"
)

TEST_FLAGS=(
    --num-tokens="$NUM_TOKENS"
    --hidden="$HIDDEN"
    --intermediate-hidden="$INTERMEDIATE"
    --num-experts="$NUM_EXPERTS"
    --num-topk="$NUM_TOPK"
)

nsys profile "${NSYS_FLAGS[@]}" \
    torchrun "${TORCH_FLAGS[@]}" \
    tests/test_mega_moe_sm90.py "${TEST_FLAGS[@]}"

echo ""
echo "=== Done. Report: $OUTPUT_DIR/$REPORT_NAME.nsys-rep ==="