diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_vllm_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_vllm_mtp.sh new file mode 100755 index 000000000..acdca8ae4 --- /dev/null +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_vllm_mtp.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +# Agentic trace replay benchmark for Kimi-K3 (MXFP4) on MI355X (gfx950) using +# vLLM with DSpark speculative decoding at level 2 (num_speculative_tokens=2). +# AMD sister of benchmarks/single_node/agentic/kimik3_fp4_b300_vllm_mtp.sh; the +# harness scaffolding follows the AMD convention in dsv4_fp4_mi355x_vllm_mtp.sh +# (env-var contract, HF-id model resolution, ROCR->HIP mapping, server metrics, +# replay invocation). The target-server deltas are all ROCm/AITER: +# - AITER MLA on the asm-padded persistent route +# (VLLM_ROCM_AITER_MLA_ASM_PADDING=asm) instead of FLASHINFER_MLA, so the +# speculative-config carries no attention_backend override. +# - AITER MoE (VLLM_ROCM_USE_AITER_MOE=1, --moe-backend aiter). +# +# Validated bring-up on rocm/pytorch-private:hy-kk-08092026 (vLLM 0.26.1rc1.dev, +# 8x gfx950, TP8) at native 1M context: asm AITER MLA + AITER MoE, fp8 KV cache, +# --gpu-memory-utilization 0.95, --max-model-len 1048576, --max-num-batched-tokens +# 4096. The asm-MLA 1M warmup only converges at that util/chunk, and concurrent +# long-context prefills are bounded with --max-num-seqs (<=8) so the MoE/linear +# activations stay within the ~15 GB/GPU headroom left after fp8 KV. GPU-resident +# KV only -- no DRAM offload arm is wired for this recipe. +# +# Model resolution is by HF id: MODEL (moonshotai/Kimi-K3) and DRAFT_MODEL +# (Inferact/Kimi-K3-DSpark) are passed straight through and the server resolves +# them (the launcher leaves MODEL_PATH unset on single-node runs). +# +# Required env vars: +# MODEL, TP, CONC, KV_OFFLOADING, TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, +# EP_SIZE, DP_ATTENTION +# +# TP8 is the only single-node layout: the MXFP4 checkpoint does not fit below 8 +# GPUs. This recipe ships the pure-TP8 profile (EP_SIZE=1, DP_ATTENTION=false). + +source "$(dirname "$0")/../../benchmark_lib.sh" + +check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION + +if [ "$TP" -ne 8 ]; then + echo "Error: Kimi-K3 on MI355X requires TP=8 (the MXFP4 checkpoint does not fit at TP<8), got TP='$TP'" >&2 + exit 1 +fi + +if [ "$EP_SIZE" -gt 1 ]; then + echo "Error: this recipe ships the pure-TP8 profile; EP_SIZE='$EP_SIZE' is not wired yet" >&2 + exit 1 +fi + +if [ "$DP_ATTENTION" = "true" ]; then + echo "Error: this recipe ships the pure-TP8 profile; DP_ATTENTION=true is not wired yet" >&2 + exit 1 +fi + +if [[ -n "${SLURM_JOB_ID:-}" ]]; then + echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}" +fi + +DRAFT_MODEL="${DRAFT_MODEL:-Inferact/Kimi-K3-DSpark}" + +# Resolve by HF id: pass the names straight to the server, which downloads / +# resolves the checkpoints. The mi355x launcher leaves MODEL_PATH unset on the +# single-node path, so this falls through to the HF-name branch. +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 + DRAFT_MODEL_PATH="${WRITABLE_MODELS_DIR:-/data/models}/${DRAFT_MODEL##*/}" + if [[ ! -d "$DRAFT_MODEL_PATH" || -z "$(ls -A "$DRAFT_MODEL_PATH" 2>/dev/null)" ]]; then + hf download "$DRAFT_MODEL" --local-dir "$DRAFT_MODEL_PATH" + fi +else + hf download "$MODEL" + export MODEL_PATH="$MODEL" + hf download "$DRAFT_MODEL" + DRAFT_MODEL_PATH="$DRAFT_MODEL" +fi + +if [ -n "${ROCR_VISIBLE_DEVICES:-}" ]; then + export HIP_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" +fi +rocm-smi || true + +# ---- Resolve traces and install deps ---------------------------------------- +resolve_trace_source +install_agentic_deps +# Nightly ROCm image may be missing runtime deps; ensure they are present. +agentic_pip_install --quiet Pillow fastapi uvicorn + +export AIPERF_HTTP_TCP_USER_TIMEOUT=900000 +# AIPerf scrapes the vLLM engine's own /metrics for the server-side throughput +# columns (pure TP: engine == public endpoint). +export AIPERF_SERVER_METRICS_URLS="http://localhost:${PORT}/metrics" +export AIPERF_REQUIRED_SERVER_METRIC_PREFIX="vllm:" +# Loading the MXFP4 shards past the default readiness window. +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-7200}" +# AIPerf pins one pooled keep-alive connection per agentic session and reuses it +# across turns; outlast the client pool so an inter-turn idle gap cannot race the +# server closing the socket (aiohttp ServerDisconnectedError -> warmup failure). +export VLLM_HTTP_TIMEOUT_KEEP_ALIVE="${VLLM_HTTP_TIMEOUT_KEEP_ALIVE:-900}" + +# ---- Server config ---------------------------------------------------------- +SERVER_LOG="$RESULT_DIR/server.log" +mkdir -p "$RESULT_DIR" + +# GPU-resident KV only. fp8 KV at 0.95 util already fills HBM after weights; +# no DRAM-offload arm is validated for this recipe. +case "${KV_OFFLOAD_BACKEND:-}" in + "") + require_agentic_kv_offload_none + ;; + *) + echo "Error: kimik3 MI355X ships GPU-resident KV only; KV_OFFLOAD_BACKEND='$KV_OFFLOAD_BACKEND' is not wired" >&2 + exit 1 + ;; +esac + +# ---- DSpark speculative decoding -------------------------------------------- +# DSpark level 2 (num_speculative_tokens=2, the value adopted for this recipe) +# on the Inferact/Kimi-K3-DSpark draft head, probabilistic drafting with +# synthetic acceptance pinned to the committed golden AL, per the AgentX policy +# in golden_al_distribution/README.md (a submission chooses the draft length, +# not the acceptance target). Mirrors the B300 sibling; the only delta is the +# ROCm image serves MLA via AITER, so no attention_backend override is set. +# EVAL_ONLY switches to real block verification (synthetic acceptance commits +# drafts regardless of target logits and would zero the SWE-bench score). +NUM_SPEC_TOKENS=2 +# Committed golden AL at K=2 on the probabilistic/block curve +# (golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml: +# thinking_on 2 -> 2.51). +SYNTHETIC_ACCEPT_LEN=2.51 +if [ "${EVAL_ONLY:-false}" = "true" ]; then + SPEC_CONFIG="{\"method\": \"dspark\", \"model\": \"$DRAFT_MODEL_PATH\", \"num_speculative_tokens\": $NUM_SPEC_TOKENS, \"draft_sample_method\": \"probabilistic\", \"rejection_sample_method\": \"block\"}" +else + SPEC_CONFIG="{\"method\": \"dspark\", \"model\": \"$DRAFT_MODEL_PATH\", \"num_speculative_tokens\": $NUM_SPEC_TOKENS, \"draft_sample_method\": \"probabilistic\", \"rejection_sample_method\": \"synthetic\", \"synthetic_acceptance_length\": $SYNTHETIC_ACCEPT_LEN}" +fi + +# AgentX concurrency counts session trees, not requests; subagent fan-out can +# push instantaneous request concurrency above CONC, so leave 2x scheduler +# headroom -- but cap at 8. The asm-MLA 1M warmup and long-context prefill +# activations only stay within the per-GPU headroom at max-num-seqs <= 8 (higher +# OOMs the MoE/linear activations on 200k+ ISL bursts). +MAX_NUM_SEQS=$((2 * CONC)) +if [ "$MAX_NUM_SEQS" -gt 8 ]; then + MAX_NUM_SEQS=8 +fi + +PARALLEL_ARGS=(--tensor-parallel-size "$TP" --data-parallel-size 1) + +echo "Starting vllm server..." +set -x +# AITER MLA + MoE, asm-padded persistent MLA route (the config validated at 1M). +export VLLM_ROCM_USE_AITER=1 +export VLLM_ROCM_USE_AITER_MOE=1 +export VLLM_ROCM_USE_AITER_MLA=1 +export VLLM_ROCM_AITER_MLA_ASM_PADDING=asm +export VLLM_USE_BREAKABLE_CUDAGRAPH=0 +export VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1 +export SAFETENSORS_FAST_GPU=1 + +{ set +x; } 2>/dev/null +VLLM_CMD=( + vllm serve "$MODEL_PATH" --served-model-name "$MODEL" + --host 0.0.0.0 + --port "$PORT" + --trust-remote-code + --distributed-executor-backend mp + --kv-cache-dtype fp8 + "${PARALLEL_ARGS[@]}" + --gpu-memory-utilization 0.95 + --max-model-len 1048576 + --max-num-batched-tokens 4096 + --max-num-seqs "$MAX_NUM_SEQS" + --moe-backend aiter + --enable-prefix-caching + --speculative-config "$SPEC_CONFIG" + --reasoning-parser kimi_k3 + --tool-call-parser kimi_k3 + --enable-auto-tool-choice + --disable-uvicorn-access-log +) +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" + +if [ "${EVAL_ONLY:-false}" = "true" ]; then + run_eval --port "$PORT" +else + build_replay_cmd "$RESULT_DIR" + run_agentic_replay_and_write_outputs "$RESULT_DIR" +fi diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index c6580af43..135341cd3 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -1250,6 +1250,41 @@ dsv4-fp4-mi355x-vllm-agentic-mtp: # while MTP creates two. Restore these points after the upstream hybrid # KV recovery fix lands: https://github.com/vllm-project/vllm/pull/45497 +kimik3-fp4-mi355x-vllm-agentic-dspark: + # AMD MI355X (gfx950) sister of kimik3-fp4-b300-vllm-agentic-dspark: DSpark + # speculative decoding on the Inferact/Kimi-K3-DSpark draft head at level 2 + # (num_speculative_tokens 2), probabilistic drafting with synthetic acceptance + # pinned to the committed golden AL 2.51 (golden_al_distribution/ + # kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml), + # per the AgentX policy: draft length is the submission's choice, the + # acceptance target is not. + # + # Validated serve config on rocm/pytorch-private:hy-kk-08092026 (vLLM + # 0.26.1rc1.dev, TP8), agentic-coding at native 1M context: + # VLLM_ROCM_USE_AITER=1 / _MOE=1 / _MLA=1, asm-padded MLA + # (VLLM_ROCM_AITER_MLA_ASM_PADDING=asm) + # --kv-cache-dtype fp8 # halves KV; required for a stable 1M warmup + # --max-model-len 1048576 # model native max_position_embeddings + # --gpu-memory-utilization 0.95 # asm-MLA 1M warmup only converges at 0.95 + # --max-num-batched-tokens 4096 # chunked prefill + # --max-num-seqs 8 # bounds concurrent long-context prefills so + # # MoE/linear activations stay within the + # # ~15 GB/GPU headroom left after fp8 KV + # GPU-resident KV only (no DRAM offload): fp8 KV at 0.95 util holds ~4.3M + # tokens on 8x gfx950. TP8-only -- the MXFP4 checkpoint does not fit below 8 + # GPUs. conc 1/2/4/8 verification sweep. + image: rocm/pytorch-private:hy-kk-08092026 + model: moonshotai/Kimi-K3 + model-prefix: kimik3 + runner: cluster:mi355x-amds + precision: fp4 + framework: vllm + multinode: false + scenarios: + agentic-coding: + - search-space: + - { tp: 8, ep: 1, dp-attn: false, kv-offloading: none, spec-decoding: mtp, conc-list: [1, 2, 4, 8] } + dsr1-fp4-mi355x-sglang-disagg-mtp: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 model: amd/DeepSeek-R1-0528-MXFP4-v2 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 448540716..eb02c3c25 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5736,3 +5736,15 @@ - "Upgrade SGLang from nightly-dev-cu13-20260711-7de33ce8 to nightly-dev-cu13-20260719-99f5a6f4, the latest nightly before the scheduler WAR-barrier regression in SGLang #31687" - "Restore the pre-watchdog AgentX concurrency grid and use AIPerf's default 300-second per-trajectory idle-gap cap" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2319 + +- config-keys: + - kimik3-fp4-mi355x-vllm-agentic-dspark + scenario-type: + - agentic-coding + description: + - "Add Kimi-K3 (MXFP4) agentic-coding recipe on MI355X (gfx950), TP8, concurrency sweep [1, 2, 4, 8]" + - "Enable DSpark speculative decoding at level 2 (num_speculative_tokens=2) on the Inferact/Kimi-K3-DSpark draft head, probabilistic drafting with synthetic acceptance pinned to golden AL 2.51" + - "Enable AITER MLA on the asm-padded persistent route (VLLM_ROCM_AITER_MLA_ASM_PADDING=asm) plus AITER MoE" + - "Enable fp8 KV cache (--kv-cache-dtype fp8) at native 1M context (--max-model-len 1048576)" + - "Set --gpu-memory-utilization 0.95, --max-num-batched-tokens 4096, and cap --max-num-seqs at 8 to bound long-context prefill activations" + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2551