#!/bin/bash

rm output -rf
mkdir output

DATA_DIR=/data
MODELS_DEEPSEEK_V2=(
  "${DATA_DIR}/vllm/models/LLM-Research/deepseek-v2"
)

MODELS=(${MODELS_DEEPSEEK_V2[@]})

# 定义变量
use_ray=0
use_eager=0
use_pp=0
use_kernel_analysis=0
# context parameter
input_sizes=(1024)
output_sizes=(1)
# batch_sizes=(1 2 4 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40)
batch_sizes=(1 4 8 16 32)

# decoder parameter
# input_sizes=(1)
# output_sizes=(128)
# batch_sizes=(1 2 4 8 16 32 64 128 256 512 1024 1280 1536 1600 1616 1632 1648 1652 1656 1660 1661 1662 1663 1664 1728 1792 2048)
# batch_sizes=(1 4 8 16 32 64 128 256 512 1024 2048)

tp_sizes=(8)
moe_ep_sizes=(8 -1)
pp_sizes=(1)

if [ $use_pp -gt 0 ]; then
  tp_sizes=(1)
  moe_ep_sizes=(-1)
  pp_sizes=(8)
  BENCHMARK_CMD=benchmarks/benchmark_throughput.py
  benchmark_option="--backend vllm --num-prompts 1000 --output-json output_throughput.csv --async-engine"
else
  BENCHMARK_CMD=benchmarks/benchmark_latency.py
  benchmark_option="--num-iters-warmup 1 --num-iters 3 --only_average"
fi

max_position_embeddings=163840

#export MLU_VISIBLE_DEVICES=4,5,6,7
export EXPERT_PARALLEL_EN=true
export VLLM_LATENCY_DEBUG=true
export VLLM_GRAPH_DEBUG=false
# export VLLM_DUMP_MLU_INFO=true
export OUTPUT_CSV_PATH=/data/solution-sdk/kangpengtao/tmp/deepseek/output.csv

ray_option=""
if [ $use_ray -gt 0 ]; then
    ray_option="--distributed-executor-backend ray --ray-workers-use-nsight"
fi

record_option=""
if [ $use_kernel_analysis -gt 0 ]; then
    # ref: https://wiki.cambricon.com/pages/viewpage.action?pageId=434445235
    export CNPERF_KERNEL_ANALYSIS=1
    record_option="--pmu --capture_range=cnpx --cnpx_include kangpengtao --cnpx_exclude kangpengtao_exec --events tp_core__write_bytes,tp_core__read_bytes,tp_memcore__write_bytes,tp_memcore__read_bytes,tp_core__lt_cycles,tp_core__csimd_pre_cycles,tp_core__csimd_post_cycles"
    use_eager=1
fi

eager_option=""
if [ $use_eager -gt 0 ]; then
    eager_option="--enforce-eager"
fi

# 遍历所有组合
for HF_MODEL in "${MODELS[@]}"; do
    quantization_option=""
    if [[ "${HF_MODEL}" == *"sq_per_token_per_channel"* ]]; then
        quantization_option="--quantization=smoothquant"
    fi
    for tp_size in "${tp_sizes[@]}"; do
        for moe_ep_size in "${moe_ep_sizes[@]}"; do
            for pp_size in "${pp_sizes[@]}"; do
                for input_size in "${input_sizes[@]}"; do
                    for output_size in "${output_sizes[@]}"; do
                        for batch_size in "${batch_sizes[@]}"; do
                            max_seq_len_to_capture=$(expr $input_size \+ $output_size)
                            max_num_batched_tokens=$(expr $batch_size \* $input_size)
                            max_model_len=$max_seq_len_to_capture 
                            if [ $max_model_len -gt $max_position_embeddings ]; then
                                continue
                            fi
                            # max_num_seqs=256
                            # if [ $max_num_seqs -lt $batch_size ]; then
                            #     max_num_seqs=$batch_size
                            # fi
                            max_num_seqs=$batch_size
                            if [ $max_model_len -gt $max_num_batched_tokens ]; then
                                max_num_batched_tokens=$max_model_len
                            fi
                            if [ $max_num_seqs -gt $max_num_batched_tokens ]; then
                                max_num_batched_tokens=$max_num_seqs
                            fi
    
                            pp_option="--pipeline-parallel-size ${pp_size}"
                            tp_option="-tp ${tp_size}"
                            ep_option="--moe-ep-size ${moe_ep_size}"
                            batch_size_option=""
                            if [ $use_pp -le 0 ]; then
                                batch_size_option="--batch-size ${batch_size}"
                            fi
        
                            hf_model_name=$(basename "${HF_MODEL}")
                            LOG_FILE=output/${hf_model_name}_${input_size}_${output_size}_tp_${tp_size}_moe_ep_${moe_ep_size}_pp_${pp_size}_bs_${batch_size}.log
                            echo "Executing ${hf_model_name} with tp_size=${tp_size}, moe_ep_size=${moe_ep_size}, pp_size=${pp_size}, input_size=${input_size}, output_size=${output_size}, batch_size=${batch_size}, max_model_len=${max_model_len}, max_num_batched_tokens=${max_num_batched_tokens}"
                            dltrace_data_name="dltrace_data_${hf_model_name}_${tp_size}_${moe_ep_size}_${pp_size}_${input_size}_${output_size}_${batch_size}_${max_model_len}_${max_num_batched_tokens}"
                            rm dltrace_data -rf
                            rm cnperf_data_* -rf
                            CNPERF_VLOG_LEVEL=0-40 cnperf-cli record ${record_option} python3 ${BENCHMARK_CMD} \
                                --trust-remote-code \
                                --max-num-batched-tokens ${max_num_batched_tokens} \
                                --max-model-len ${max_model_len} \
                                --block-size 16 \
                                --model ${HF_MODEL} \
                                --tokenizer ${HF_MODEL} \
                                --dtype bfloat16 \
                                --input-len ${input_size} \
                                --output-len ${output_size} \
                                ${pp_option} ${tp_option} ${ep_option} \
                                --max-seq-len-to-capture ${max_seq_len_to_capture} \
                                --max-num-seqs ${max_num_seqs} \
                                ${batch_size_option} \
                                ${eager_option} ${ray_option} ${quantization_option} \
                                2>&1 | tee ${LOG_FILE}
                                # 检查日志文件中是否有 torch.OutOfMemoryError, Ceil of batch 或is larger than mlu blocks
                                if grep -E -q "torch\.OutOfMemoryError|Ceil of batch|is larger than mlu blocks" "$LOG_FILE"; then
                                    echo "Found one or more specified errors in the log file."
                                    break
                                else
                                    echo "No specified errors found."
                                fi
                            mv dltrace_data ${dltrace_data_name}
                            mv cnperf_data_* ${dltrace_data_name}/
                        done
                    done
                done
            done
        done
    done
done
