# TP×DP matrix experiment for DeepSeek-V4-Flash on Ascend 910C (8 NPUs / 16 dies)
# using vLLM-Ascend.
# Tests vLLM with three parallel configurations:
#   TP=2, DP=4  -> 2 dies per replica, 4 replicas
#   TP=4, DP=2  -> 4 dies per replica, 2 replicas
#   TP=8, DP=1  -> 8 dies, no data parallelism
#
# Platform: ascend_910c (see platforms/ascend_910c.env).
# Host: 910c.1 / NPU-NODE61, openEuler 22.03 SP4 aarch64, driver 25.5.2, CANN 9.0.0.
#
# DeepSeek-V4-Flash w8a8-mtp weights are present at /mnt/models/DeepSeek-V4-Flash-w8a8-mtp
# (downloaded from ModelScope Eco-Tech/DeepSeek-V4-Flash-w8a8-mtp, ~280 GiB, 70 shards).

EXPERIMENT="dsv4_910c_vllm_tp_dp_matrix"
MODEL_NAME="DeepSeek-V4-Flash"
# Real DSV4-Flash w8a8-mtp directory (verified present & chown'd to shishi).
MODEL_PATH="${MODEL_PATH:-/mnt/models/DeepSeek-V4-Flash-w8a8-mtp}"
SERVED_MODEL_NAME="dsv4"

VLLM_PORT="${VLLM_PORT:-30052}"

# Dedicated container name so this experiment never touches other 910c runs.
CONTAINER_NAME="vllm-ascend-dsv4-910c"

# Python interpreter for the benchmark client inside the vllm-ascend container.
CONTAINER_PYTHON="/usr/local/python3.12.13/bin/python3"

# vllm-ascend image. Use the general v0.23 A3 image for DSV4 (the GLM5.2-tuned
# variant carries GLM-specific patches and is NOT compatible with DSV4).
# Verified present locally: quay.io/ascend/vllm-ascend:v0.23.0rc1-a3-openeuler
USE_DOCKER="${USE_DOCKER:-1}"
DOCKER_IMAGE="${DOCKER_IMAGE:-local/vllm-ascend:0.23-a3-dsv4-sglang}"

# Benchmark client Docker image. The vllm-ascend images built on this host
# (local/vllm-ascend:0.23-a3-*-sglang, see envs/ASCEND_910C_ENV_SETUP.md §8.1)
# carry sglang 0.5.2 bench_serving, so the client runs as a standalone container
# (TORCH_DEVICE_BACKEND_AUTOLOAD=0 injected by the runner, no NPU needed).
DOCKER_CLIENT_IMAGE="${DOCKER_CLIENT_IMAGE:-local/vllm-ascend:0.23-a3-dsv4-sglang}"
USE_DOCKER_CLIENT="${USE_DOCKER_CLIENT:-1}"

# Device selection. We mount all 8 cards (16 dies) via --device /dev/davinci0..15
# in start_vllm_docker.sh. ASCEND_VISIBLE_DEVICES is kept for parity with the
# shared library but the explicit --device flags are the authoritative path on
# this host (Ascend Docker Runtime injection was unreliable here).
export ASCEND_VISIBLE_DEVICES="${ASCEND_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"

# Runtime working directory for logs, pid files, and tmp.
RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"

# Parallel configurations to test. Format: "TP DP"
# A3 910C has 16 dies (8 cards x 2 dies/card). TP/DP address dies, so TP*DP=16
# uses all dies. DSV4-Flash w8a8 weights ~280 GiB total, sharded across TP dies:
#   TP=4  DP=4 -> 4 dies/replica x 4 replicas (39 GiB/die, tight KV cache)
#   TP=8  DP=2 -> 8 dies/replica x 2 replicas (35 GiB/die, balanced)
#   TP=16 DP=1 -> 16 dies/replica x 1 replica (17.5 GiB/die, max KV cache)
# Note: A3 TP=4 == H20 TP=4 per-card-equivalent (A3 has 2 dies/card).
#   TP=2, DP=4 -> 4 dies/replica x 4 replicas (smallest TP, most replicas)
#   TP=4, DP=2 -> 4 dies/replica x 2 replicas
#   TP=8, DP=1 -> 8 dies/replica x 1 replica  (largest TP, max weight sharding)
if [[ -n "${PARALLEL_CONFIGS_STR:-}" ]]; then
  declare -a PARALLEL_CONFIGS=()
  for pair in $PARALLEL_CONFIGS_STR; do
    PARALLEL_CONFIGS+=("${pair//,/ }")
  done
else
  declare -a PARALLEL_CONFIGS=(
    "4 4"
    "8 2"
    "16 1"
  )
fi

# vLLM-Ascend server settings for DSV4-Flash (w8a8-mtp).
# Notes:
#   - KV cache dtype fp8 is supported on 910C; fall back to fp16 if the image rejects it.
#   - block-size 128 matches Ascend page semantics (910C favors 128).
#   - MAX_MODEL_LEN: DSV4-Flash supports up to 1M context; cap at 131072 for the
#     matrix sweep (extend to 1M via matrix.json once TP=8 is verified).
#   - gpu-memory-utilization maps to NPU HBM fraction on vllm-ascend (0.9).
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.9}"
KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-}"
BLOCK_SIZE="${BLOCK_SIZE:-128}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-1048576}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-64}"
# Max tokens per batch (official A3 tutorial: 10240).
MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-10240}"
# API server count (official: 1).
API_SERVER_COUNT="${API_SERVER_COUNT:-1}"

# Per-TP parameter overrides (applied in start_vllm_docker.sh via case $TP).
# DSV4-Flash w8a8 weight per die = ~280GiB / TP. Each die has 64GB HBM.
# Defaults below are the verified-OK values (from the 2026-07-29..30 adaptive
# runs' server_cmd.txt); 1048576-context OOMs on TP4/TP8 (KV cache budget).
#   TP=4:  ~39 GiB/die weights -> 25 GiB for KV cache; cap context to avoid OOM
#   TP=8:  ~35 GiB/die weights -> 29 GiB for KV cache; balanced
#   TP=16: ~17.5 GiB/die weights -> 46 GiB for KV cache; full context
TP4_GPU_MEMORY_UTILIZATION="${TP4_GPU_MEMORY_UTILIZATION:-0.9}"
TP4_MAX_MODEL_LEN="${TP4_MAX_MODEL_LEN:-32768}"
TP4_MAX_NUM_SEQS="${TP4_MAX_NUM_SEQS:-128}"
TP8_GPU_MEMORY_UTILIZATION="${TP8_GPU_MEMORY_UTILIZATION:-0.9}"
TP8_MAX_MODEL_LEN="${TP8_MAX_MODEL_LEN:-65536}"
TP8_MAX_NUM_SEQS="${TP8_MAX_NUM_SEQS:-256}"
TP16_GPU_MEMORY_UTILIZATION="${TP16_GPU_MEMORY_UTILIZATION:-0.9}"
TP16_MAX_MODEL_LEN="${TP16_MAX_MODEL_LEN:-131072}"
TP16_MAX_NUM_SEQS="${TP16_MAX_NUM_SEQS:-256}"

# vLLM-Ascend-specific launch flags injected by start_vllm_docker.sh.
# attention backend for 910C: use the fused/atb attention path. Adjust per image.
VLLM_ASCEND_ATTENTION_BACKEND="${VLLM_ASCEND_ATTENTION_BACKEND:-atb}"

# DSV4-Flash-specific server args (passed through to vllm serve in
# start_vllm_docker.sh). These are REQUIRED for DSV4 -- GLM-5.2 does not need them.
#   --tokenizer-mode deepseek_v4        : DSV4 custom tokenizer
#   --tool-call-parser deepseek_v4      : DSV4 tool-call parser
#   --enable-auto-tool-choice           : enable automatic tool choice
#   --reasoning-parser deepseek_v4      : DSV4 reasoning parser
#   --enable-expert-parallel            : shard 256 experts across EP ranks
#   --quantization ascend               : use modelslim w8a8 quantization path
#   --safetensors-load-strategy prefetch: prefetch shards for faster load
#   --async-scheduling                  : overlap CPU scheduling with NPU compute
#   MTP speculative decoding via SPEC_CONFIG
DSV4_TOKENIZER_MODE="${DSV4_TOKENIZER_MODE:-deepseek_v4}"
DSV4_TOOL_CALL_PARSER="${DSV4_TOOL_CALL_PARSER:-deepseek_v4}"
DSV4_REASONING_PARSER="${DSV4_REASONING_PARSER:-deepseek_v4}"
DSV4_QUANTIZATION="${DSV4_QUANTIZATION:-ascend}"
DSV4_SAFETENSORS_LOAD_STRATEGY="${DSV4_SAFETENSORS_LOAD_STRATEGY:-prefetch}"
# MTP speculative config (1 speculative token). JSON string, kept single-quoted
# in the launcher to avoid shell mangling.
DSV4_SPEC_CONFIG="${DSV4_SPEC_CONFIG:-{\"num_speculative_tokens\": 1, \"method\": \"mtp\", \"enforce_eager\": true}}"
DSV4_COMPILATION_CONFIG="${DSV4_COMPILATION_CONFIG:-{\"cudagraph_mode\": \"FULL_DECODE_ONLY\"}}"
DSV4_ADDITIONAL_CONFIG="${DSV4_ADDITIONAL_CONFIG:-{\"ascend_compilation_config\":{\"enable_npugraph_ex\":true,\"enable_static_kernel\":false},\"enable_cpu_binding\": true,\"multistream_overlap_shared_expert\":true}}"
DSV4_ENABLE_EXPERT_PARALLEL="${DSV4_ENABLE_EXPERT_PARALLEL:-1}"
DSV4_ENABLE_ASYNC_SCHEDULING="${DSV4_ENABLE_ASYNC_SCHEDULING:-1}"
# MTP speculative decoding (disabled for fair comparison with H20 which has no MTP).
DSV4_ENABLE_MTP="${DSV4_ENABLE_MTP:-0}"

# Model-loader extra config (multithread load, 128 threads) to speed up 280GiB load.
DSV4_MODEL_LOADER_EXTRA_CONFIG="${DSV4_MODEL_LOADER_EXTRA_CONFIG:-{\"enable_multithread_load\": \"true\", \"num_threads\": 128}}"

# Dataset used by sglang.bench_serving --dataset-name random.
DATASET_PATH="${DATASET_PATH:-${ROOT_DIR}/datasets/ShareGPT_V3_unfiltered_cleaned_split.json}"

# Matrix and concurrency rules are defined in matrix.json by default.
MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR:-.}/matrix.json}"
MATRIX_MODE="${MATRIX_MODE:-Y}"

export CONCURRENCY_SAMPLES="${CONCURRENCY_SAMPLES:-2}"

SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}"
GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}"

DRY_RUN="${DRY_RUN:-0}"
GRID_LIMIT="${GRID_LIMIT:-0}"