diff --git a/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_sglang_mtp.sh b/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_sglang_mtp.sh index 11e40a70e..95ada0bb3 100644 --- a/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_sglang_mtp.sh +++ b/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_sglang_mtp.sh @@ -1,24 +1,71 @@ #!/usr/bin/env bash -set -eo pipefail +set -euo pipefail set -x - + +# Agentic trace replay benchmark for Kimi-K3 MXFP4 on MI355X / MI350X (gfx950) +# using vLLM. +# +# The server command is the AMD reference `vllm serve` for this model, i.e. the +# upstream vLLM recipe's amd block (vllm-project/recipes, +# https://recipes.vllm.ai/moonshotai/Kimi-K3) as run in practice: +# +# --trust-remote-code --moe-backend auto --tensor-parallel-size 8 +# --load-format auto --gpu-memory-utilization 0.95 --mm-encoder-tp-mode data +# --max-num-seqs 128 --max-num-batched-tokens 4096 --enable-auto-tool-choice +# --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 +# +# with env VLLM_ROCM_USE_AITER=1 SAFETENSORS_FAST_GPU=1 AITER_SITUV2_A8W4=1 +# AITER_BF16_FP8_MOE_BOUND=0 VLLM_USE_BREAKABLE_CUDAGRAPH=0. +# +# K3 is a 2.8T-parameter natively-multimodal MoE (896 routed experts, 16/token +# plus shared) on Kimi Delta Attention, gated MLA and Attention Residuals, with +# a 1M-token native context. +# +# TP=8 ONLY. The MXFP4 checkpoint is 1.561 TB decimal (1.420 TiB, 96 +# safetensors), ~195 GB/GPU across 8 GPUs of the 288 GB part; TP=4 would need +# ~390 GB/GPU and cannot load. Upstream strategy_min_gpus agrees (single_node_tp +# and multi_node_tep both 8, DEP 16+), which is why there is no DP-attention arm. +# +# Required env vars: +# MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, +# EP_SIZE +# +# Perf-search knobs. Each defaults to the reference command's value, so an +# otherwise-unset run reproduces the reference exactly: +# GPU_MEM_UTIL 0.95 (reference) +# MAX_NUM_BATCHED_TOKENS 8192 (default) +# AITER_A8W4 1 (reference; 0 = aiter a16w4 MoE path) +# LANGUAGE_MODEL_ONLY true +# KV_CACHE_DTYPE fp8 (default for every arm; =auto for a bf16 A/B) +# KV_BLOCK_SIZE unset (unset -> vLLM sizes the page; 128 under fp8) +# MAX_MODEL_LEN 1M +# SPEC_DECODE true (this is the _mtp DSpark recipe; =false for a no-spec A/B) +# SPEC_NUM_TOKENS 2 (DSpark draft length; validated by the _mtp config) + source "$(dirname "$0")/../../benchmark_lib.sh" - export EVAL_FRAMEWORK="lm-eval" - -check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION - -if [[ -n "$SLURM_JOB_ID" ]]; then - echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME" +wait_for_amd_gpu_clean + +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 under slurm cgroups. -if [ -n "$ROCR_VISIBLE_DEVICES" ]; then + +if [ "$TP" -ne 8 ]; then + echo "Error: Kimi-K3 MXFP4 is a 1.56 TB checkpoint and only fits at TP=8 on" >&2 + echo " 288 GB gfx950 parts (~195 GB/GPU). Got TP=$TP." >&2 + exit 1 +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 + +# `hf download` creates the target dir if missing and is itself idempotent. The +# 1.56 TB checkpoint is normally pre-staged, so these calls are a no-op there. +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 @@ -26,256 +73,283 @@ else hf download "$MODEL" export MODEL_PATH="$MODEL" fi + rocm-smi || true amd-smi || true - -# A server killed on this node minutes earlier (previous job, crashed run) -# can still be draining its ~1.4 TB of HBM: KFD reclaim takes minutes, and -# booting into a half-drained node fails RCCL init with HIP 'unhandled cuda -# error' / 'invalid argument' (observed as the mooncake-c64 CI failure). -# Wait for the GPUs to come back before launching. -# Per-GPU threshold: idle nodes hold a small driver/firmware VRAM baseline -# (observed up to ~4%/GPU, node-dependent), while a draining or occupied -# GPU sits at 50-90%. Require every GPU <= 10%. -GPU_CLEAN=false -for i in $(seq 1 90); do - VRAM_MAX=$(rocm-smi --showmemuse 2>/dev/null | grep -oE "GPU Memory Allocated \(VRAM%\): [0-9]+" | awk '{if ($NF > m) m = $NF} END {print m+0}') - if [ "${VRAM_MAX:-0}" -le 10 ]; then echo "GPUs clean (vram%max=$VRAM_MAX after $((i*10))s)"; GPU_CLEAN=true; break; fi - echo "waiting for prior-job GPU memory reclaim: vram%max=$VRAM_MAX"; sleep 10 -done -[ "$GPU_CLEAN" = "true" ] || { echo "Error: GPUs still draining prior job's memory after 15min" >&2; exit 1; } - + +# ---- Resolve traces and install deps ---------------------------------------- resolve_trace_source install_agentic_deps - + +# ---- Reference env block ---------------------------------------------------- +export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm +export VLLM_ROCM_USE_AITER=1 +export SAFETENSORS_FAST_GPU=1 +export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1 +export AITER_SITUV2_A8W4=1 +export AITER_BF16_FP8_MOE_BOUND=0 +export VLLM_USE_BREAKABLE_CUDAGRAPH=0 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 + +# Workaround for MEC FW <177 RCCL memory reclaim issue (shared with the other +# gfx950 recipes in this tree). +mec_version=$(rocm-smi --showfw 2>/dev/null | grep MEC | head -n 1 | awk '{print $NF}') +if [[ "$mec_version" == "" || ${mec_version:-0} -lt 177 ]]; then + export HSA_NO_SCRATCH_RECLAIM=1 +fi + +# 2.8T of weights off a shared/NFS mount takes far longer than the default. +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-7200}" + +# Long agentic turns against a 1M context: keep the client from timing out +# mid-request while the server is prefill-bound. +export AIPERF_HTTP_TCP_USER_TIMEOUT=900000 + +# ---- Server config ---------------------------------------------------------- SERVER_LOG="$RESULT_DIR/server.log" -ROUTER_LOG="$RESULT_DIR/router.log" mkdir -p "$RESULT_DIR" - -export PYTHONNOUSERSITE=1 -# Agentic warmup dispatches hundreds of large prompts at once; allow up to -# 15 minutes of TCP progress before AIPerf declares a connection dead. -export AIPERF_HTTP_TCP_USER_TIMEOUT=900000 -# AIPerf pins one pooled keep-alive connection per session (client-side -# keep-alive 300s) while uvicorn's default SGLANG_TIMEOUT_KEEP_ALIVE is 5s; -# inter-turn idle gaps can reuse a socket exactly as the server closes it. -# Outlast the client pool so the race cannot occur. -export SGLANG_TIMEOUT_KEEP_ALIVE=900 -# The DSA indexer's top-k v2 kernel (default since v0.5.14) is JIT-compiled -# from CUDA-only source (cooperative_groups.h) and cannot build for gfx950; -# v1 dispatches to the precompiled HIP op in sgl-kernel (upstream MI355X CI -# runs DSA models the same way). -export SGLANG_OPT_USE_TOPK_V2=false - -# HiCache L2 (host DRAM), optionally extended with Mooncake L3. -# KV_OFFLOADING=dram requires KV_OFFLOAD_BACKEND=hicache or mooncake. -# -# Per-arm L2 ratio (sizing rationale below) applies to both backends unless -# overridden via HICACHE_RATIO. TP arm (182.7 GB/rank device pool): the -# agentic-coding corpus saturates any fixed DRAM pool at conc ≥ 10; ratio 1.5 -# (~2.9 TB pinned) is the safe default for cluster:mi355x-amds nodes (~3.0 TB -# available DRAM per runners.yaml). ratio=2.5 (~4.8 TB) yields higher -# throughput at conc 10-12 but exceeds physical DRAM on these nodes and must -# be set via HICACHE_RATIO env-var override on nodes that can accommodate it. -# The DP-attention arm (159.4 GB/rank) only runs at conc >= 32, where the host -# tier just absorbs overflow - ratio 0.5 (~1.2 TB pinned, ~1.8 TB of load -# headroom) at negligible hit-rate cost (ratio 1.5 OOMs the host mid-storm at -# conc 48). -CACHE_ARGS=() -if agentic_kv_offload_enabled; then - if [ "$DP_ATTENTION" = "true" ]; then - HICACHE_RATIO="${HICACHE_RATIO:-0.5}" - else - # ratio=1.5 (~2.9 TB pinned): safe default within the ~3.0 TB DRAM - # available on cluster:mi355x-amds nodes. Set HICACHE_RATIO=2.5 via - # env-var override for maximum throughput on nodes with >4 TB DRAM. - HICACHE_RATIO="${HICACHE_RATIO:-1.5}" - fi - # write_through_selective skips DRAM writes for non-reusable KV blocks, - # reducing host-bus traffic without affecting the cache hit rate. - HICACHE_WRITE_POLICY="${HICACHE_WRITE_POLICY:-write_through_selective}" - HICACHE_IO_BACKEND="${HICACHE_IO_BACKEND:-direct}" - HICACHE_MEM_LAYOUT="${HICACHE_MEM_LAYOUT:-page_first_direct}" - case "$KV_OFFLOAD_BACKEND" in - hicache) - echo "HiCache (GPU+host DRAM only): ratio=$HICACHE_RATIO, write_policy=$HICACHE_WRITE_POLICY, io_backend=$HICACHE_IO_BACKEND, mem_layout=$HICACHE_MEM_LAYOUT" - CACHE_ARGS=( - --enable-hierarchical-cache - --hicache-ratio "$HICACHE_RATIO" - --hicache-write-policy "$HICACHE_WRITE_POLICY" - --hicache-io-backend "$HICACHE_IO_BACKEND" - --hicache-mem-layout "$HICACHE_MEM_LAYOUT" - ) - ;; - mooncake) - L3_PER_RANK_GB="${L3_PER_RANK_GB:-40}" - python3 -c "from mooncake.store import MooncakeDistributedStore" >/dev/null - MOONCAKE_MASTER_PORT=$((PORT + 12000)) - MOONCAKE_MASTER_LOG="$RESULT_DIR/mooncake_master.log" - MOONCAKE_CONFIG_PATH="$RESULT_DIR/mooncake_config.json" - cat > "$MOONCAKE_CONFIG_PATH" < "$MOONCAKE_MASTER_LOG" 2>&1 & - MOONCAKE_MASTER_PID=$! - sleep 2 - kill -0 "$MOONCAKE_MASTER_PID" - echo "HiCache+Mooncake: ratio=$HICACHE_RATIO, l3_per_rank=${L3_PER_RANK_GB} GB, dram_budget=${TOTAL_CPU_DRAM_GB} GB" - CACHE_ARGS=( - --enable-hierarchical-cache - --hicache-ratio "$HICACHE_RATIO" - --hicache-size 0 - --hicache-write-policy "$HICACHE_WRITE_POLICY" - --hicache-io-backend "$HICACHE_IO_BACKEND" - --hicache-mem-layout "$HICACHE_MEM_LAYOUT" - --hicache-storage-backend mooncake - --hicache-storage-prefetch-policy wait_complete - ) - ;; - *) - echo "Error: unsupported KV_OFFLOAD_BACKEND '$KV_OFFLOAD_BACKEND' (expected: hicache or mooncake)" >&2 - exit 1 - ;; - esac +trap cleanup_agentic_services EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +# ---- KV offload ------------------------------------------------------------- +# TOTAL_CPU_DRAM_GB is the aggregate host-DRAM budget the matrix generator +# derives from dram-utilization and the runner's available-cpu-dram-mib, capped +# at the 3,095,781 MiB (3 TB decimal) agentic limit. Per +# benchmarks/single_node/agentic/README.md it must be consumed as given and +# never replaced with a model-specific constant. +OFFLOAD_ARGS=() + +if agentic_kv_offload_enabled; then +case "${KV_OFFLOAD_BACKEND:-}" in + vllm-simple) + require_agentic_kv_offload_backend "$KV_OFFLOAD_BACKEND" + CPU_BYTES_PER_RANK=$(( TOTAL_CPU_DRAM_GB * 1000 * 1000 * 1000 / TP )) + # Identical prefixes must hash to identical block keys across ranks. + export PYTHONHASHSEED=42 + SIMPLE_LAZY_OFFLOAD="${SIMPLE_LAZY_OFFLOAD:-false}" + OFFLOAD_ARGS=( + --kv-transfer-config + "{\"kv_connector\":\"SimpleCPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"cpu_bytes_to_use_per_rank\":$CPU_BYTES_PER_RANK,\"lazy_offload\":$SIMPLE_LAZY_OFFLOAD}}" + ) + echo "SimpleCPUOffloadConnector: ${CPU_BYTES_PER_RANK} B/rank x ${TP} ranks, lazy_offload=$SIMPLE_LAZY_OFFLOAD" + ;; + lmcache) + require_agentic_kv_offload_backend "$KV_OFFLOAD_BACKEND" + + # Keep the image's tested torch/ROCm stack and install only LMCache's + # missing runtime dependencies, same as the MiniMax-M3 lmcache arm. + LMCACHE_VERSION="0.5.5rc3" + LMCACHE_ROCM_INDEX="https://github.com/LMCache/LMCache/releases/expanded_assets/nightly-rocm" + agentic_pip_install --quiet --no-cache-dir --no-deps \ + "sortedcontainers==2.4.0" \ + "opentelemetry-exporter-prometheus==0.61b0" \ + "cupy-rocm-7-0==14.1.1" \ + "lmcache==${LMCACHE_VERSION}" --find-links "$LMCACHE_ROCM_INDEX" + + # LMCache 0.5.5's transfer-channel layer eagerly imports the Mooncake + # backend (mooncake_te_impl.py -> `from mooncake.engine import + # TransferEngine`), whose native .so resolves all of its DT_NEEDED libs at + # import. The vLLM ROCm image ships none of them, so the import sanity + # check below (and the LMCache server) would otherwise fail with + # "ImportError: lib*.so: cannot open shared object file" (first libglog, + # then libjsoncpp, ...). Provision Mooncake's full runtime lib set from the + # distro before importing. apt-get install is idempotent, so run it + # whenever any of the libs is still missing rather than gating on one. + LMCACHE_NATIVE_LIBS=(libglog.so.0 libjsoncpp.so.25 libibverbs.so.1 librdmacm.so.1 libnuma.so.1) + for lib in "${LMCACHE_NATIVE_LIBS[@]}"; do + if ! ldconfig -p | grep -q "$lib"; then + apt-get update + apt-get install -y \ + libgoogle-glog0v5 libjsoncpp25 libibverbs1 librdmacm1 libnuma1 + break + fi + done + python3 -c \ + "import cupy; import lmcache.integration.vllm.lmcache_mp_connector; import opentelemetry.exporter.prometheus" \ + >/dev/null + + # One MP server for the node, per the Kimi-K3 recipe + # (docs.lmcache.ai/recipes/kimi_k3.html), with --chunk-size sized for + # THIS stack rather than the recipe's CUDA-path 768: the connector + # requires the chunk to be a multiple of every engine KV group's + # tokens_per_block, and the hybrid KDA/MLA layout here registers + # attention groups at 1536 ("Setting attention block size to 1536", + # run 31644990546) plus a KDA state group at 3072 (run 31645828378), + # so 3072 is the minimum valid chunk. The multi-group layout also + # requires one object group per sliding-window size: + # --separate-object-groups. + LMCACHE_PORT=6555 + LMCACHE_HTTP_PORT=8090 + LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" + + LMCACHE_L1_SIZE_GB="$TOTAL_CPU_DRAM_GB" + + LMCACHE_CMD=( + lmcache server + --host 127.0.0.1 + --port "$LMCACHE_PORT" + --http-host 127.0.0.1 + --http-port "$LMCACHE_HTTP_PORT" + --l1-size-gb "$LMCACHE_L1_SIZE_GB" + --l1-init-size-gb 10 + --chunk-size 12288 + --separate-object-groups + --enable-extra-logging + --extra-logging-interval 30 + --max-cpu-workers 8 + --max-gpu-workers 8 + --eviction-policy LRU + --supported-transfer-mode lmcache_driven + --shm-name "" + ) + append_command "$RESULT_DIR/lmcache_command.txt" "${LMCACHE_CMD[@]}" + "${LMCACHE_CMD[@]}" > "$LMCACHE_LOG" 2>&1 & + LMCACHE_PID=$! + wait_for_ready \ + --endpoint "http://127.0.0.1:${LMCACHE_HTTP_PORT}/healthcheck" \ + --log "$LMCACHE_LOG" \ + --pid "$LMCACHE_PID" \ + --sleep-interval 1 \ + --timeout 600 + + # 100k-330k-token agentic prefixes make single retrieves large; use the + # same MQ timeout headroom as the MiniMax-M3 arm. + 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.port\":$LMCACHE_PORT,\"lmcache.mp.mq_timeout\":6000.0}}" + ) + ;; + *) + echo "Error: unsupported KV_OFFLOAD_BACKEND='$KV_OFFLOAD_BACKEND' (expected vllm-simple or lmcache)" >&2 + exit 1 + ;; +esac fi - -# Arm selection. TP arm keeps the FP8 sibling's cookbook batch-shaping -# bands. -# -# NOTE: the DP-attention path below is currently DORMANT (no dp-attn arms -# in amd-master.yaml): DSA + dp-attention hangs a collective under -# long-context prefill on ROCm v0.5.14 (watchdog kills the scheduler with -# zero completions; reproduced with and without HiCache, with and without -# the DSv4 DP collective envs; short prompts are fine). Re-enable the -# config arm once upstream fixes the DSA DP prefill path. -# -# When active, the DP-attention (DEP) arm fronts the DP ranks with sglang-router -# using consistent hashing on the AIPerf correlation id so multi-turn -# sessions stay on the DP rank holding their radix/hicache prefix, and -# widens chunked-prefill (whole-engine, /dp ranks) like the B300 sibling. -USE_SGLANG_ROUTER=false -SGLANG_BACKEND_PORT="$PORT" -PARALLEL_ARGS=(--tp "$TP" --ep-size "$EP_SIZE") -MEM_FRACTION_STATIC=0.85 -if [ "$DP_ATTENTION" = "true" ]; then - USE_SGLANG_ROUTER=true - export AIPERF_HTTP_X_SMG_ROUTING_KEY_FROM_CORRELATION_ID=true - SGLANG_BACKEND_PORT=$((PORT + 1)) - SGLANG_ROUTER_METRICS_PORT=$((PORT + 10000)) - SGLANG_ROUTER_CMD=(python3 -m sglang_router.launch_router) - PARALLEL_ARGS+=(--dp "$TP" --enable-dp-attention) - CHUNKED_PREFILL_SIZE=32768 - export AGENTIC_WARMUP_GRACE_PERIOD=3600 - # Swap the DP gather collectives to gatherv/reduce-scatter on ROCm - # (dsv4_fp4_mi355x_sglang.sh precedent - the only green DP-attention - # config on this cluster/image): with the defaults the DSA DP path - # hangs a collective under long-context prefill load until the - # watchdog kills the scheduler (0/96 storm completions, twice). - export SGLANG_DP_USE_GATHERV=1 - export SGLANG_DP_USE_REDUCE_SCATTER=1 - export GPU_MAX_HW_QUEUES=5 -elif [ "$CONC" -le 16 ]; then - # Chunked prefill 32k: smaller chunks let the scheduler interleave decode - # steps between prefill chunks, reducing TPOT for concurrent sessions - # (improved interactivity vs the original 131072-token chunk). The reduced - # chunk size drops per-chunk activation headroom from ~7 GiB/rank to - # ~1.7 GiB/rank, so mem-fraction 0.85 is safe (0.85 OOMed at 131k: - # "Tried to allocate 6.86 GiB ... 5.15 GiB is free", run 29751563205). - CHUNKED_PREFILL_SIZE=32768 - MEM_FRACTION_STATIC=0.85 + +# ---- LLM server ------------------------------------------------------------ + +# ---- Parallelism ------------------------------------------------------------ +EP_ARGS=() +if [ "$EP_SIZE" -gt 1 ]; then + EP_ARGS=(--enable-expert-parallel) +fi + +# ---- Speculative / Util------------------------------------------------------ +case "$CONC" in + # No KV offload; the working set fits in HBM. + 1) + SYNTHETIC_ACCEPT_LEN=3.75 + SPEC_NUM_TOKENS=6 + GPU_MEM_UTIL=0.9 + MAX_NUM_BATCHED_TOKENS=16384 + ;; + 2|4|8|10|12|14) + SYNTHETIC_ACCEPT_LEN=3.00 + SPEC_NUM_TOKENS=3 + GPU_MEM_UTIL=0.9 + MAX_NUM_BATCHED_TOKENS=8192 + ;; + *) + SPEC_NUM_TOKENS=0 + GPU_MEM_UTIL=0.9 + MAX_NUM_BATCHED_TOKENS=8192 + ;; +esac + +SPEC_ARGS=() +if [ "$SPEC_NUM_TOKENS" -gt 0 ]; then +if [ "${EVAL_ONLY:-false}" = "true" ]; then + SPEC_ARGS=( + --speculative-config + "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"num_speculative_tokens\":$SPEC_NUM_TOKENS,\"method\":\"dspark\",\"attention_backend\":\"TRITON_MLA\",\"kv_cache_dtype\":\"fp8\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\": \"block\"}" + ) else - CHUNKED_PREFILL_SIZE=32768 - export AGENTIC_WARMUP_GRACE_PERIOD=3600 + SPEC_ARGS=( + --speculative-config + "{\"model\":\"Inferact/Kimi-K3-DSpark\",\"num_speculative_tokens\":$SPEC_NUM_TOKENS,\"method\":\"dspark\",\"attention_backend\":\"TRITON_MLA\",\"kv_cache_dtype\":\"fp8\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\": \"synthetic\", \"synthetic_acceptance_length\": $SYNTHETIC_ACCEPT_LEN}" + ) + fi +fi + +# ---- HIP graph ------------------------------------------------------------ +MAX_NUM_SEQS=$((2 * CONC)) +MAX_CUDAGRAPH_CAPTURE_SIZE=$((MAX_NUM_SEQS * (1 + SPEC_NUM_TOKENS))) +CUDAGRAPH_CAPTURE_SIZES="$(seq -s, 2 "$MAX_CUDAGRAPH_CAPTURE_SIZE")" +COMPILATION_CONFIG_ARGS=(--compilation-config "{\"mode\":3,\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"max_cudagraph_capture_size\":$MAX_CUDAGRAPH_CAPTURE_SIZE,\"custom_ops\":[\"+fused_rms_norm_gated\"],\"cudagraph_capture_sizes\":[$CUDAGRAPH_CAPTURE_SIZES]}") + +echo "Starting vllm server..." +export PYTHONNOUSERSITE=1 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS="${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-1200}" + + +# ---- DCP ------------------------------------------------------------ +# DCP shards decode KV across the TP ranks, so it must divide TP. +DCP_SIZE="${DCP_SIZE:-8}" +if [ $((TP % DCP_SIZE)) -ne 0 ]; then + echo "Error: TP='$TP' must be divisible by DCP_SIZE='$DCP_SIZE'" >&2 + exit 1 fi -# 2×CONC in-flight slots: MTP draft+verify transiently batches more tokens -# than CONC sessions; headroom prevents scheduler stalls under burst. -MAX_RUNNING_REQUESTS=$((2 * CONC)) -[ "$MAX_RUNNING_REQUESTS" -gt 256 ] && MAX_RUNNING_REQUESTS=256 -# SGLang interpolates a bs list [1..max_bs] automatically; cap at 64 to -# keep graph-capture memory bounded without giving up coverage. -CUDA_GRAPH_MAX_BS=$(( MAX_RUNNING_REQUESTS < 64 ? MAX_RUNNING_REQUESTS : 64 )) - -if [ "${EVAL_ONLY:-false}" != "true" ]; then - export SGLANG_SIMULATE_ACC_LEN=3.61 - export SGLANG_SIMULATE_ACC_METHOD=match-expected - export SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token +CP_ARGS=() +ATTN_BE_ARGS=() +if [ "$DCP_SIZE" -gt 1 ]; then + CP_ARGS+=(--decode-context-parallel-size "$DCP_SIZE" --dcp-comm-backend a2a) + ATTN_BE_ARGS+=(--attention-backend ROCM_AITER_MLA) fi - -SGLANG_CMD=( - python3 -m sglang.launch_server - --model-path "$MODEL_PATH" - --served-model-name "$MODEL" +export VLLM_USE_DIRECT_DCP_A2A=0 +export VLLM_USE_DIRECT_DCP_Q_GATHER=0 +export VLLM_USE_DIRECT_DCP_KV_GATHER=0 + +{ set +x; } 2>/dev/null +VLLM_CMD=( + vllm serve "$MODEL_PATH" --served-model-name "$MODEL" --host 0.0.0.0 - --port "$SGLANG_BACKEND_PORT" + --port "$PORT" --trust-remote-code - "${PARALLEL_ARGS[@]}" - --kv-cache-dtype fp8_e4m3 - --dsa-prefill-backend tilelang - --dsa-decode-backend tilelang - # GLM-5.2 emits the GLM-4.7-style tool-call format; glm47 is required for - # structured message.tool_calls (SWE-bench agentic evals die without it). - # The glm45 reasoning parser keeps hybrid thinking in reasoning_content. - --tool-call-parser glm47 - --reasoning-parser glm45 - --chunked-prefill-size "$CHUNKED_PREFILL_SIZE" - --mem-fraction-static "$MEM_FRACTION_STATIC" - --max-running-requests "$MAX_RUNNING_REQUESTS" - --cuda-graph-max-bs "$CUDA_GRAPH_MAX_BS" - --speculative-algorithm EAGLE - --speculative-num-steps 5 - --speculative-eagle-topk 1 - --speculative-num-draft-tokens 6 - "${CACHE_ARGS[@]}" - --watchdog-timeout 1800 - --enable-metrics + --moe-backend auto + --tensor-parallel-size "$TP" + "${EP_ARGS[@]}" + --load-format fastsafetensors + --gpu-memory-utilization "$GPU_MEM_UTIL" + --language-model-only + --max-num-seqs "$MAX_NUM_SEQS" + --enable-auto-tool-choice + --tool-call-parser kimi_k3 + --reasoning-parser kimi_k3 + --max-model-len 1048576 + --enable-prefix-caching + --kv-cache-dtype "fp8" + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" + --attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}' + "${ATTN_BE_ARGS[@]}" + "${COMPILATION_CONFIG_ARGS[@]}" + "${SPEC_ARGS[@]}" + "${OFFLOAD_ARGS[@]}" + "${CP_ARGS[@]}" ) - -printf '%q ' "${SGLANG_CMD[@]}" | tee "$RESULT_DIR/sglang_command.txt" -printf '\n' | tee -a "$RESULT_DIR/sglang_command.txt" - -echo "Starting SGLang server for MI355X..." -"${SGLANG_CMD[@]}" > "$SERVER_LOG" 2>&1 & +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 "$SGLANG_BACKEND_PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" - -if [ "$USE_SGLANG_ROUTER" = "true" ]; then - echo "Starting SGLang router on port $PORT for $TP DP ranks..." - "${SGLANG_ROUTER_CMD[@]}" \ - --worker-urls "http://localhost:$SGLANG_BACKEND_PORT" \ - --policy consistent_hashing \ - --request-id-headers x-correlation-id \ - --dp-aware \ - --host 0.0.0.0 \ - --port "$PORT" \ - --prometheus-host 127.0.0.1 \ - --prometheus-port "$SGLANG_ROUTER_METRICS_PORT" \ - --connect-timeout-secs 900 \ - --request-timeout-secs 14400 \ - --disable-health-check \ - --disable-retries > "$ROUTER_LOG" 2>&1 & - ROUTER_PID=$! - echo "Router PID: $ROUTER_PID" - wait_for_server_ready --port "$PORT" --server-log "$ROUTER_LOG" --server-pid "$ROUTER_PID" -fi - + +wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" + if [ "${EVAL_ONLY}" = "true" ]; then run_eval --port "$PORT" else build_replay_cmd "$RESULT_DIR" - REPLAY_CMD+=" --server-metrics http://localhost:$SGLANG_BACKEND_PORT/metrics" run_agentic_replay_and_write_outputs "$RESULT_DIR" -fi +fi \ No newline at end of file diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x.sh new file mode 100755 index 000000000..eff8ce940 --- /dev/null +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +set -euo pipefail + +export SPEC_DECODING=none +exec bash "$(dirname "$0")/kimik3_fp4_mi355x_mtp.sh" diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh index a8f2c42c1..81b3b2991 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh @@ -149,25 +149,18 @@ case "${KV_OFFLOAD_BACKEND:-}" in lmcache) require_agentic_kv_offload_backend "$KV_OFFLOAD_BACKEND" - # Keep the image's tested torch/ROCm stack and install only LMCache's - # missing runtime dependencies, same as the MiniMax-M3 lmcache arm. - LMCACHE_VERSION="0.5.5.dev60+rocm7.2" - LMCACHE_ROCM_INDEX="https://github.com/LMCache/LMCache/releases/expanded_assets/nightly-rocm" + # Keep the stock image's torch/ROCm stack and install the LMCache runtime + # dependencies used by the validated reference run. + LMCACHE_VERSION="0.5.5rc3+rocm7.2" + export KV_OFFLOAD_BACKEND_METADATA="{\"name\":\"lmcache\",\"version\":\"${LMCACHE_VERSION}\"}" + LMCACHE_RELEASE="v0.5.5rc3" + LMCACHE_ROCM_INDEX="https://github.com/LMCache/LMCache/releases/expanded_assets/${LMCACHE_RELEASE}-rocm" agentic_pip_install --quiet --no-cache-dir --no-deps \ "sortedcontainers==2.4.0" \ "opentelemetry-exporter-prometheus==0.61b0" \ "cupy-rocm-7-0==14.1.1" \ "lmcache==${LMCACHE_VERSION}" --find-links "$LMCACHE_ROCM_INDEX" - # LMCache 0.5.5's transfer-channel layer eagerly imports the Mooncake - # backend (mooncake_te_impl.py -> `from mooncake.engine import - # TransferEngine`), whose native .so resolves all of its DT_NEEDED libs at - # import. The vLLM ROCm image ships none of them, so the import sanity - # check below (and the LMCache server) would otherwise fail with - # "ImportError: lib*.so: cannot open shared object file" (first libglog, - # then libjsoncpp, ...). Provision Mooncake's full runtime lib set from the - # distro before importing. apt-get install is idempotent, so run it - # whenever any of the libs is still missing rather than gating on one. LMCACHE_NATIVE_LIBS=(libglog.so.0 libjsoncpp.so.25 libibverbs.so.1 librdmacm.so.1 libnuma.so.1) for lib in "${LMCACHE_NATIVE_LIBS[@]}"; do if ! ldconfig -p | grep -q "$lib"; then @@ -178,17 +171,16 @@ case "${KV_OFFLOAD_BACKEND:-}" in fi done python3 -c \ - "import cupy; import lmcache.integration.vllm.lmcache_mp_connector; import opentelemetry.exporter.prometheus" \ - >/dev/null + "import cupy; import opentelemetry.exporter.prometheus; from lmcache.v1.multiprocess.http_server import run_http_server" # One MP server for the node, per the Kimi-K3 recipe # (docs.lmcache.ai/recipes/kimi_k3.html), with --chunk-size sized for # THIS stack rather than the recipe's CUDA-path 768: the connector # requires the chunk to be a multiple of every engine KV group's - # tokens_per_block, and the hybrid KDA/MLA layout here registers - # attention groups at 1536 ("Setting attention block size to 1536", - # run 31644990546) plus a KDA state group at 3072 (run 31645828378), - # so 3072 is the minimum valid chunk. The multi-group layout also + # tokens_per_block. The hybrid KDA/MLA layout registers attention groups + # at 1536 tokens and a KDA state group at 3072. Use 12288 for every point + # so it is divisible by both group sizes and matches the tested LMCache + # scheduler geometry. The multi-group layout also # requires one object group per sliding-window size: # --separate-object-groups. LMCACHE_PORT=6555 @@ -205,7 +197,7 @@ case "${KV_OFFLOAD_BACKEND:-}" in --http-port "$LMCACHE_HTTP_PORT" --l1-size-gb "$LMCACHE_L1_SIZE_GB" --l1-init-size-gb 10 - --chunk-size 3072 + --chunk-size 12288 --separate-object-groups --enable-extra-logging --extra-logging-interval 30 @@ -248,20 +240,24 @@ if [ "$EP_SIZE" -gt 1 ]; then fi # ---- Speculative / Util------------------------------------------------------ -case "$CONC" in - # No KV offload; the working set fits in HBM. - 1) +case "${SPEC_DECODING:-mtp}:$CONC" in + mtp:1) SYNTHETIC_ACCEPT_LEN=3.75 SPEC_NUM_TOKENS=6 - GPU_MEM_UTIL=0.9 + GPU_MEM_UTIL=0.88 MAX_NUM_BATCHED_TOKENS=16384 ;; - 2|4|8|10|12|14) + mtp:2|mtp:4|mtp:8|mtp:10|mtp:12|mtp:14) SYNTHETIC_ACCEPT_LEN=3.00 SPEC_NUM_TOKENS=3 - GPU_MEM_UTIL=0.9 + GPU_MEM_UTIL=0.88 MAX_NUM_BATCHED_TOKENS=8192 ;; + none:40|none:44|none:48) + SPEC_NUM_TOKENS=0 + GPU_MEM_UTIL=0.88 + MAX_NUM_BATCHED_TOKENS=16384 + ;; *) SPEC_NUM_TOKENS=0 GPU_MEM_UTIL=0.85 @@ -285,10 +281,27 @@ else fi # ---- HIP graph ------------------------------------------------------------ -MAX_NUM_SEQS=$((2 * CONC)) -MAX_CUDAGRAPH_CAPTURE_SIZE=$((MAX_NUM_SEQS * (1 + SPEC_NUM_TOKENS))) -CUDAGRAPH_CAPTURE_SIZES="$(seq -s, 2 "$MAX_CUDAGRAPH_CAPTURE_SIZE")" -COMPILATION_CONFIG_ARGS=(--compilation-config "{\"mode\":3,\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"max_cudagraph_capture_size\":$MAX_CUDAGRAPH_CAPTURE_SIZE,\"custom_ops\":[\"+fused_rms_norm_gated\"],\"cudagraph_capture_sizes\":[$CUDAGRAPH_CAPTURE_SIZES]}") +SERVER_STREAM_ARGS=() +PREFIX_MATCH_ARGS=() +ATTENTION_CONFIG='{"mla_prefill_backend":"ROCM_AITER_FA"}' +COMPILATION_CUSTOM_OPS='["+fused_rms_norm_gated"]' +case "${SPEC_DECODING:-mtp}:$CONC" in + none:40|none:44|none:48) + MAX_NUM_SEQS=80 + MAX_CUDAGRAPH_CAPTURE_SIZE=4096 + CUDAGRAPH_CAPTURE_SIZES="$(seq -s, 1 "$MAX_NUM_SEQS"),128,256,512,1024,2048,4096" + SERVER_STREAM_ARGS=(--stream-interval 10) + PREFIX_MATCH_ARGS=(--prefix-match-unit 128) + ATTENTION_CONFIG='{"mla_prefill_backend":"ROCM_AITER_FA","use_prefill_query_quantization":true}' + COMPILATION_CUSTOM_OPS='["+fused_rms_norm_gated","+quant_fp8","+grouped_topk","+sparse_attn_indexer","none"]' + ;; + *) + MAX_NUM_SEQS=$((2 * CONC)) + MAX_CUDAGRAPH_CAPTURE_SIZE=$((MAX_NUM_SEQS * (1 + SPEC_NUM_TOKENS))) + CUDAGRAPH_CAPTURE_SIZES="$(seq -s, 2 "$MAX_CUDAGRAPH_CAPTURE_SIZE")" + ;; +esac +COMPILATION_CONFIG_ARGS=(--compilation-config "{\"mode\":3,\"cudagraph_mode\":\"FULL_AND_PIECEWISE\",\"max_cudagraph_capture_size\":$MAX_CUDAGRAPH_CAPTURE_SIZE,\"custom_ops\":$COMPILATION_CUSTOM_OPS,\"cudagraph_capture_sizes\":[$CUDAGRAPH_CAPTURE_SIZES]}") echo "Starting vllm server..." export PYTHONNOUSERSITE=1 @@ -305,8 +318,18 @@ fi CP_ARGS=() ATTN_BE_ARGS=() if [ "$DCP_SIZE" -gt 1 ]; then - CP_ARGS+=(--decode-context-parallel-size "$DCP_SIZE" --dcp-comm-backend a2a) - ATTN_BE_ARGS+=(--attention-backend TRITON_MLA) + CP_KV_CACHE_INTERLEAVE_SIZE=1 + if [ "${KV_OFFLOAD_BACKEND:-}" = "lmcache" ]; then + CP_KV_CACHE_INTERLEAVE_SIZE=1536 + fi + CP_ARGS+=( + --decode-context-parallel-size "$DCP_SIZE" + --dcp-comm-backend a2a + --cp-kv-cache-interleave-size "$CP_KV_CACHE_INTERLEAVE_SIZE" + ) + ATTN_BE_ARGS+=(--attention-backend ROCM_AITER_MLA) + export VLLM_ALLOW_DCP_FULL_CUDAGRAPH=1 + export PREFIX_CACHING_HASH_ALGO=sha256 fi export VLLM_USE_DIRECT_DCP_A2A=0 export VLLM_USE_DIRECT_DCP_Q_GATHER=0 @@ -329,10 +352,12 @@ VLLM_CMD=( --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 --max-model-len 1048576 + "${SERVER_STREAM_ARGS[@]}" --enable-prefix-caching + "${PREFIX_MATCH_ARGS[@]}" --kv-cache-dtype "fp8" --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" - --attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}' + --attention-config "$ATTENTION_CONFIG" "${ATTN_BE_ARGS[@]}" "${COMPILATION_CONFIG_ARGS[@]}" "${SPEC_ARGS[@]}" diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 0d380154a..6100a527c 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -625,7 +625,7 @@ dsr1-fp8-mi355x-sglang-disagg-mtp: - "DECODE_MTP_SIZE=2" kimik3-fp4-mi355x-vllm-agentic-mtp: - image: vllm/vllm-openai-rocm:nightly-1dc464d42681d22f38caf1fdc1eb632dc4421c45 + image: vllm/vllm-openai-rocm:nightly-7c5dc571cbd1064ecc8a9b1045637ff647aa22cb model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds @@ -636,8 +636,8 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: agentic-coding: - dram-utilization: 0.60 search-space: - - { tp: 8, kv-offloading: none, conc-list: [1] , spec-decoding: mtp} - - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.5.dev60+rocm7.2" }, conc-list: [4, 8, 10, 12, 14], spec-decoding: mtp } + - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.5rc3+rocm7.2" }, conc-list: [1, 8, 14], spec-decoding: mtp } + - { tp: 8, ep: 1, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.5.5rc3" }, conc-list: [44, 48, 52], spec-decoding: mtp } # Kimi-K3 MXFP4 agentic-coding benchmark on MI355X via ATOM with DSpark # speculative decoding. Acceptance is pinned to the committed golden curve in diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 4eb0a2caa..1e0f93567 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6821,3 +6821,22 @@ description: - "Refresh to collect TensorRT-LLM server metrics." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2774 + +- config-keys: + - kimik3-fp4-mi355x-vllm-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Refresh the Kimi-K3 MI355X AgentX recipe from the merged PR #2787 baseline to vLLM ROCm nightly 7c5dc571 and LMCache 0.5.5rc3+rocm7.2." + - "Run C1, C8, and C14 with TP8 DSpark and LMCache DRAM KV offload; run C40, C44, and C48 with TP8/DCP8 and no speculative decoding." + - "Use gpu-memory-utilization 0.88 for all six points, a 12288-token LMCache chunk, and a 1536-token DCP KV-cache interleave; C40/C44/C48 retain the 16K batched-token, max-num-seqs 80, and full CUDA-graph profile." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2804 + +- config-keys: + - kimik3-fp4-mi355x-vllm-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Switch the second MI355X K3 vLLM search-space arm from a no-spec TP8/DCP8 curve (conc 40/44/48, LMCache 0.5.5rc3+rocm7.2) to an MTP spec-decoding TP8/DCP8 curve (conc 44/48/52, LMCache 0.5.5rc3)." + - "Refresh the MI355X K3 vLLM agentic trace-replay benchmark script." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2809