[AMD][AgentX] Kimi-K2.7-Code FP4 MI355X agentX vLLM #2126
+259
−0
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,235 @@ | ||
| #!/usr/bin/env bash | ||
| set -euo pipefail | ||
| set -x | ||
|
|
||
| # Agentic trace replay benchmark for Kimi-K2.7 FP4 on MI355X using vLLM. | ||
| # KV_OFFLOADING=none -> GPU KV only | ||
| # KV_OFFLOADING=dram KV_OFFLOAD_BACKEND=lmcache -> LMCache MP server + connector | ||
| # | ||
| # Required env vars: | ||
| # MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, EP_SIZE | ||
|
|
||
|
|
||
| source "$(dirname "$0")/../../benchmark_lib.sh" | ||
|
|
||
| check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE | ||
|
|
||
| if [[ -n "${SLURM_JOB_ID:-}" ]]; then | ||
| echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}" | ||
| fi | ||
|
|
||
| # ROCR/HIP visibility for vLLM 0.14+ | ||
| if [ -n "${ROCR_VISIBLE_DEVICES:-}" ]; then | ||
| export HIP_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" | ||
| fi | ||
|
|
||
| if [[ -n "${MODEL_PATH:-}" ]]; then | ||
| if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then | ||
| hf download "$MODEL" --local-dir "$MODEL_PATH" | ||
| fi | ||
| else | ||
| hf download "$MODEL" | ||
| export MODEL_PATH="$MODEL" | ||
| fi | ||
| rocm-smi || true | ||
| amd-smi || true | ||
|
|
||
| # ---- Resolve traces and install deps ---------------------------------------- | ||
| resolve_trace_source | ||
| install_agentic_deps | ||
|
|
||
| # Install amd-quark for MXFP4 (manual install due to ROCm vLLM bug) | ||
| pip install amd-quark | ||
|
|
||
| # Disable AITER RMSNorm for TP < 8 due to accuracy issues | ||
| if [ "${TP}" -lt 8 ]; then | ||
| export VLLM_ROCM_USE_AITER_RMSNORM=0 | ||
| fi | ||
| # Workaround for MEC FW <177 RCCL memory reclaim issue | ||
| version=$(rocm-smi --showfw 2>/dev/null | grep MEC | head -n 1 | awk '{print $NF}') | ||
| if [[ "$version" == "" || ${version:-0} -lt 177 ]]; then | ||
| export HSA_NO_SCRATCH_RECLAIM=1 | ||
| fi | ||
|
|
||
| export VLLM_ROCM_USE_AITER=1 | ||
| export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 | ||
| # Avoid intermittent symm_mem all-reduce rendezvous hang at engine init on | ||
| # MI35x nodes (see KIMIK27_CONC64_LMCACHE_RUNBOOK error #4). | ||
| export VLLM_ALLREDUCE_USE_SYMM_MEM="${VLLM_ALLREDUCE_USE_SYMM_MEM:-0}" | ||
|
|
||
| # ---- Server config ---------------------------------------------------------- | ||
| SERVER_LOG="$RESULT_DIR/server.log" | ||
| LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" | ||
| mkdir -p "$RESULT_DIR" | ||
|
|
||
| OFFLOAD_ARGS=() | ||
| PREFIX_CACHE_ARGS=() | ||
|
|
||
| # ---- LMCache config --------------------------------------------------------- | ||
| LMCACHE_PID="" | ||
|
|
||
| cleanup_lmcache_server() { | ||
| if [[ -n "$LMCACHE_PID" ]] && kill -0 "$LMCACHE_PID" 2>/dev/null; then | ||
| kill "$LMCACHE_PID" 2>/dev/null || true | ||
| wait "$LMCACHE_PID" 2>/dev/null || true | ||
| fi | ||
| } | ||
| trap cleanup_lmcache_server EXIT | ||
|
|
||
| wait_for_lmcache_ready() { | ||
| { set +x; } 2>/dev/null | ||
| local attempts="${LMCACHE_READY_ATTEMPTS:-120}" | ||
| local tail_pid="" | ||
|
|
||
| while [ ! -f "$LMCACHE_LOG" ]; do | ||
| if [[ -n "$LMCACHE_PID" ]] && ! kill -0 "$LMCACHE_PID" 2>/dev/null; then | ||
| echo "LMCache server died before creating log file. Exiting." >&2 | ||
| exit 1 | ||
| fi | ||
| sleep 1 | ||
| done | ||
|
|
||
| tail -f -n +1 "$LMCACHE_LOG" & | ||
| tail_pid=$! | ||
|
|
||
| for ((i = 1; i <= attempts; i++)); do | ||
| if curl --output /dev/null --silent --fail "http://127.0.0.1:${LMCACHE_HTTP_PORT}/healthcheck"; then | ||
| kill "$tail_pid" 2>/dev/null || true | ||
| wait "$tail_pid" 2>/dev/null || true | ||
| return 0 | ||
| fi | ||
| if [[ -n "$LMCACHE_PID" ]] && ! kill -0 "$LMCACHE_PID" 2>/dev/null; then | ||
| echo "LMCache server died before becoming healthy. Log follows:" >&2 | ||
| kill "$tail_pid" 2>/dev/null || true | ||
| wait "$tail_pid" 2>/dev/null || true | ||
| cat "$LMCACHE_LOG" >&2 || true | ||
| exit 1 | ||
| fi | ||
| sleep 1 | ||
| done | ||
|
|
||
| echo "Timed out waiting for LMCache server healthcheck. Log follows:" >&2 | ||
| kill "$tail_pid" 2>/dev/null || true | ||
| wait "$tail_pid" 2>/dev/null || true | ||
| cat "$LMCACHE_LOG" >&2 || true | ||
| exit 1 | ||
| } | ||
|
|
||
| if [[ "$KV_OFFLOADING" == "none" ]]; then | ||
| OFFLOAD_MODE="none" | ||
| else | ||
| OFFLOAD_MODE="${KV_OFFLOAD_BACKEND:?KV_OFFLOAD_BACKEND required when KV_OFFLOADING=dram}" | ||
| fi | ||
|
|
||
| case "$OFFLOAD_MODE" in | ||
| none) | ||
| # GPU-only KV baseline: keep on-GPU prefix caching ON (no DRAM offload) | ||
| # so it's apples-to-apples vs the lmcache cell. | ||
| PREFIX_CACHE_ARGS=(--enable-prefix-caching) | ||
| ;; | ||
| lmcache) | ||
| unset VLLM_USE_SIMPLE_KV_OFFLOAD | ||
|
|
||
| # Build LMCache against ROCm if the connector isn't importable. Clone to | ||
| # a container-local dir (NOT bind-mounted /workspace) so the next job's | ||
| # checkout `clean: true` won't trip over root-owned build artifacts. | ||
| if ! python3 -c "import lmcache.integration.vllm.lmcache_mp_connector" >/dev/null 2>&1; then | ||
| LMCACHE_SRC_DIR="${LMCACHE_SRC_DIR:-/opt/lmcache-src}" | ||
| LMCACHE_GIT_REF="${LMCACHE_GIT_REF:-aaf7c0d3}" | ||
| rm -rf "$LMCACHE_SRC_DIR" | ||
| git clone https://github.com/LMCache/LMCache.git "$LMCACHE_SRC_DIR" | ||
| ( cd "$LMCACHE_SRC_DIR" | ||
| git checkout "$LMCACHE_GIT_REF" | ||
| pip install -r requirements/build.txt | ||
| CXX=hipcc BUILD_WITH_HIP=1 pip install -e . --no-build-isolation ) | ||
| python3 -c "import lmcache.integration.vllm.lmcache_mp_connector" >/dev/null | ||
| fi | ||
|
|
||
| LMCACHE_HOST="${LMCACHE_HOST:-127.0.0.1}" | ||
| LMCACHE_PORT="${LMCACHE_PORT:-5555}" | ||
| LMCACHE_HTTP_PORT="${LMCACHE_HTTP_PORT:-8080}" | ||
| LMCACHE_CONNECT_HOST="${LMCACHE_CONNECT_HOST:-tcp://$LMCACHE_HOST}" | ||
| # LMCache L1 is SHM-backed: if L1 > /dev/shm free it silently disables | ||
| # SHM and falls back to the pickle path (crashes at load). Cap L1 to 90% | ||
| # of /dev/shm free so SHM stays enabled. | ||
| SHM_FREE_GB=$(df -BG --output=avail /dev/shm 2>/dev/null | tail -1 | tr -dc '0-9') | ||
| SHM_CAP_GB=$(( SHM_FREE_GB * 90 / 100 )) | ||
| LMCACHE_L1_SIZE_GB="${LMCACHE_L1_SIZE_GB:-$TOTAL_CPU_DRAM_GB}" | ||
| if [ -n "$SHM_CAP_GB" ] && [ "$SHM_CAP_GB" -gt 0 ] && [ "$LMCACHE_L1_SIZE_GB" -gt "$SHM_CAP_GB" ]; then | ||
| echo "Capping LMCACHE_L1_SIZE_GB $LMCACHE_L1_SIZE_GB -> $SHM_CAP_GB to fit /dev/shm (${SHM_FREE_GB}G free)" | ||
| LMCACHE_L1_SIZE_GB="$SHM_CAP_GB" | ||
| fi | ||
| LMCACHE_L1_INIT_SIZE_GB="${LMCACHE_L1_INIT_SIZE_GB:-20}" | ||
| LMCACHE_L1_READ_TTL_SECONDS="${LMCACHE_L1_READ_TTL_SECONDS:-7200}" | ||
| LMCACHE_CHUNK_SIZE="${LMCACHE_CHUNK_SIZE:-256}" | ||
| LMCACHE_MAX_WORKERS="${LMCACHE_MAX_WORKERS:-$((TP * 2))}" | ||
| export PYTHONHASHSEED="${PYTHONHASHSEED:-0}" | ||
| export LMCACHE_BLOCKING_TIMEOUT_SECS=60 | ||
|
|
||
| echo "Starting LMCache MP server..." | ||
| LMCACHE_CMD=( | ||
| lmcache server | ||
| --host "$LMCACHE_HOST" | ||
| --port "$LMCACHE_PORT" | ||
| --http-host "$LMCACHE_HOST" | ||
| --http-port "$LMCACHE_HTTP_PORT" | ||
| --l1-size-gb "$LMCACHE_L1_SIZE_GB" | ||
| --l1-init-size-gb "$LMCACHE_L1_INIT_SIZE_GB" | ||
| --l1-read-ttl-seconds "$LMCACHE_L1_READ_TTL_SECONDS" | ||
| --chunk-size "$LMCACHE_CHUNK_SIZE" | ||
| --max-workers "$LMCACHE_MAX_WORKERS" | ||
| --eviction-policy LRU | ||
| ) | ||
| printf '%q ' "${LMCACHE_CMD[@]}" > "$RESULT_DIR/lmcache_command.txt" | ||
| printf '\n' >> "$RESULT_DIR/lmcache_command.txt" | ||
| "${LMCACHE_CMD[@]}" > "$LMCACHE_LOG" 2>&1 & | ||
| LMCACHE_PID=$! | ||
| echo "LMCache server PID: $LMCACHE_PID" | ||
| wait_for_lmcache_ready | ||
|
|
||
| OFFLOAD_ARGS=( | ||
| --kv-transfer-config | ||
| "{\"kv_connector\":\"LMCacheMPConnector\",\"kv_connector_module_path\":\"lmcache.integration.vllm.lmcache_mp_connector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"lmcache.mp.host\":\"$LMCACHE_CONNECT_HOST\",\"lmcache.mp.port\":$LMCACHE_PORT}}" | ||
| ) | ||
| ;; | ||
| *) | ||
| echo "Error: unsupported KV_OFFLOAD_BACKEND '$OFFLOAD_MODE' (expected: lmcache)" >&2 | ||
| exit 1 | ||
| ;; | ||
| esac | ||
|
|
||
| EP_ARGS=() | ||
| if [ "$EP_SIZE" -gt 1 ]; then | ||
| EP_ARGS=(--enable-expert-parallel) | ||
| fi | ||
|
|
||
| echo "Starting vllm server..." | ||
| export PYTHONNOUSERSITE=1 | ||
|
|
||
| { set +x; } 2>/dev/null | ||
| VLLM_CMD=( | ||
| vllm serve "$MODEL_PATH" --served-model-name "$MODEL" | ||
| --host 0.0.0.0 | ||
| --port "$PORT" | ||
| --tensor-parallel-size="$TP" | ||
| "${EP_ARGS[@]}" | ||
| --gpu-memory-utilization 0.90 | ||
| --block-size=1 | ||
| --trust-remote-code | ||
| --max-num-seqs "$CONC" | ||
| --mm-encoder-tp-mode data | ||
| "${PREFIX_CACHE_ARGS[@]}" | ||
| "${OFFLOAD_ARGS[@]}" | ||
| ) | ||
| printf '%q ' "${VLLM_CMD[@]}" | tee "$RESULT_DIR/vllm_command.txt" | ||
| printf '\n' | tee -a "$RESULT_DIR/vllm_command.txt" | ||
| "${VLLM_CMD[@]}" > "$SERVER_LOG" 2>&1 & | ||
| SERVER_PID=$! | ||
| echo "Server PID: $SERVER_PID" | ||
|
|
||
| wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" | ||
|
|
||
| # ---- Run benchmark ---------------------------------------------------------- | ||
| build_replay_cmd "$RESULT_DIR" | ||
|
|
||
| run_agentic_replay_and_write_outputs "$RESULT_DIR" | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.