|
| 1 | +#!/usr/bin/env bash |
| 2 | +set -euo pipefail |
| 3 | +set -x |
| 4 | + |
| 5 | +source "$(dirname "$0")/../../benchmark_lib.sh" |
| 6 | + |
| 7 | +check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION |
| 8 | + |
| 9 | + |
| 10 | +if [[ -n "${SLURM_JOB_ID:-}" ]]; then |
| 11 | + echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}" |
| 12 | +fi |
| 13 | + |
| 14 | +DRAFT_MODEL="lightseekorg/kimi-k2.6-eagle3.1-mla" |
| 15 | + |
| 16 | +if [[ -n "${MODEL_PATH:-}" ]]; then |
| 17 | + if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then |
| 18 | + hf download "$MODEL" --local-dir "$MODEL_PATH" |
| 19 | + fi |
| 20 | + DRAFT_MODEL_PATH="/data/models/${DRAFT_MODEL##*/}" |
| 21 | + if [[ ! -d "$DRAFT_MODEL_PATH" || -z "$(ls -A "$DRAFT_MODEL_PATH" 2>/dev/null)" ]]; then |
| 22 | + hf download "$DRAFT_MODEL" --local-dir "$DRAFT_MODEL_PATH" |
| 23 | + fi |
| 24 | +else |
| 25 | + hf download "$MODEL" |
| 26 | + export MODEL_PATH="$MODEL" |
| 27 | + hf download "$DRAFT_MODEL" |
| 28 | + DRAFT_MODEL_PATH="$DRAFT_MODEL" |
| 29 | +fi |
| 30 | +nvidia-smi |
| 31 | + |
| 32 | +resolve_trace_source |
| 33 | +install_agentic_deps |
| 34 | + |
| 35 | +SERVER_LOG="$RESULT_DIR/server.log" |
| 36 | +mkdir -p "$RESULT_DIR" |
| 37 | + |
| 38 | +SERVER_PID="" |
| 39 | + |
| 40 | +cleanup_agentic_services() { |
| 41 | + local exit_code=$? |
| 42 | + trap - EXIT INT TERM |
| 43 | + set +e |
| 44 | + stop_background_process_tree "$SERVER_PID" "vLLM server" 60 |
| 45 | + exit "$exit_code" |
| 46 | +} |
| 47 | +trap cleanup_agentic_services EXIT |
| 48 | +trap 'exit 130' INT |
| 49 | +trap 'exit 143' TERM |
| 50 | + |
| 51 | +DCP_SIZE="${DCP_SIZE:-1}" |
| 52 | +DCP_ARGS=() |
| 53 | +if [[ "$DCP_SIZE" -gt 1 ]]; then |
| 54 | + DCP_ARGS+=(--decode-context-parallel-size "$DCP_SIZE" --dcp-comm-backend a2a) |
| 55 | + NUM_SPEC_TOKENS=3 |
| 56 | + SYNTHETIC_ACCEPT_LEN=2.88 |
| 57 | + SPEC_ARGS=(--speculative-config "{\"method\":\"eagle3\",\"model\":\"$DRAFT_MODEL_PATH\",\"num_speculative_tokens\":$NUM_SPEC_TOKENS,\"rejection_sample_method\":\"synthetic\",\"synthetic_acceptance_length\":$SYNTHETIC_ACCEPT_LEN,\"attention_backend\":\"TOKENSPEED_MLA\"}") |
| 58 | + ATTN_CONFIG='{"mla_prefill_backend":"TOKENSPEED_MLA"}' |
| 59 | + COMPILATION_CONFIG='{"pass_config":{"fuse_allreduce_rms":false}}' |
| 60 | +else |
| 61 | + NUM_SPEC_TOKENS=4 |
| 62 | + SYNTHETIC_ACCEPT_LEN=3.24 |
| 63 | + SPEC_ARGS=(--speculative-config "{\"method\":\"eagle3\",\"model\":\"$DRAFT_MODEL_PATH\",\"num_speculative_tokens\":$NUM_SPEC_TOKENS,\"rejection_sample_method\":\"synthetic\",\"synthetic_acceptance_length\":$SYNTHETIC_ACCEPT_LEN}") |
| 64 | + ATTN_CONFIG='{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":true}' |
| 65 | + COMPILATION_CONFIG='{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}' |
| 66 | +fi |
| 67 | +ATTN_BACKEND_ARGS=(--attention-backend TOKENSPEED_MLA) |
| 68 | + |
| 69 | +OFFLOAD_ARGS=() |
| 70 | + |
| 71 | +if agentic_kv_offload_enabled; then |
| 72 | + case "$KV_OFFLOAD_BACKEND" in |
| 73 | + native) |
| 74 | + export VLLM_USE_SIMPLE_KV_OFFLOAD=1 |
| 75 | + CPU_OFFLOAD_BYTES=$((TOTAL_CPU_DRAM_GB * 1024 * 1024 * 1024)) |
| 76 | + OFFLOAD_ARGS=( |
| 77 | + --disable-hybrid-kv-cache-manager |
| 78 | + --kv-transfer-config |
| 79 | + "{\"kv_connector\":\"SimpleCPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"cpu_bytes_to_use\":$CPU_OFFLOAD_BYTES,\"lazy_offload\":false}}" |
| 80 | + ) |
| 81 | + ;; |
| 82 | + *) echo "Error: unsupported KV_OFFLOAD_BACKEND value '$KV_OFFLOAD_BACKEND' with EAGLE3 (expected: native)" >&2; exit 1 ;; |
| 83 | + esac |
| 84 | +fi |
| 85 | + |
| 86 | + |
| 87 | +GMU=0.90 |
| 88 | +if [[ "$DCP_SIZE" -gt 1 && "$KV_OFFLOADING" == "none" ]]; then |
| 89 | + GMU=0.85 |
| 90 | +fi |
| 91 | + |
| 92 | +echo "Starting vllm server..." |
| 93 | +export PYTHONNOUSERSITE=1 |
| 94 | + |
| 95 | +export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm |
| 96 | + |
| 97 | +{ set +x; } 2>/dev/null |
| 98 | +VLLM_CMD=( |
| 99 | + vllm serve "$MODEL_PATH" --served-model-name "$MODEL" |
| 100 | + --host 0.0.0.0 |
| 101 | + --port "$PORT" |
| 102 | + --kv-cache-dtype fp8 |
| 103 | + --trust-remote-code |
| 104 | + --block-size 64 |
| 105 | + --language-model-only |
| 106 | + --gpu-memory-utilization "$GMU" |
| 107 | + --max-num-seqs "$CONC" |
| 108 | + "${ATTN_BACKEND_ARGS[@]}" |
| 109 | + --attention-config "$ATTN_CONFIG" |
| 110 | + --compilation-config "$COMPILATION_CONFIG" |
| 111 | + --max-cudagraph-capture-size 2048 |
| 112 | + --max-num-batched-tokens 16384 |
| 113 | + --stream-interval 10 |
| 114 | + --enable-prefix-caching |
| 115 | + --tensor-parallel-size "$TP" |
| 116 | + "${SPEC_ARGS[@]}" |
| 117 | + "${DCP_ARGS[@]}" |
| 118 | + "${OFFLOAD_ARGS[@]}" |
| 119 | +) |
| 120 | +printf '%q ' "${VLLM_CMD[@]}" | tee "$RESULT_DIR/vllm_command.txt" |
| 121 | +printf '\n' | tee -a "$RESULT_DIR/vllm_command.txt" |
| 122 | +"${VLLM_CMD[@]}" > "$SERVER_LOG" 2>&1 & |
| 123 | +SERVER_PID=$! |
| 124 | +echo "Server PID: $SERVER_PID" |
| 125 | + |
| 126 | +wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" |
| 127 | + |
| 128 | +build_replay_cmd "$RESULT_DIR" |
| 129 | + |
| 130 | +run_agentic_replay_and_write_outputs "$RESULT_DIR" |
0 commit comments