diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh new file mode 100644 index 000000000..e8ac9dce9 --- /dev/null +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh @@ -0,0 +1,259 @@ +#!/usr/bin/env bash +set -eo pipefail +set -x + +# Agentic trace replay benchmark for Kimi-K3 MXFP4 on MI355X / MI350X (gfx950) +# using ATOM with DSpark speculative decoding. +# +# Companion to kimik3_fp4_mi355x_mtp.sh, which runs the same checkpoint and the +# same concurrency points under vLLM, so the two arms are directly comparable. +# +# TP=8 ONLY, for the same reason as the vLLM arm: the MXFP4 checkpoint is +# 1.561 TB decimal (~195 GB/GPU across 8 GPUs of the 288 GB part), and TP=4 +# would need ~390 GB/GPU and cannot load. +# +# The ATOM image is purpose-built for K3, so apply_k3_container_patches.sh is +# NOT sourced here -- that script reproduces a specific patched vLLM container +# byte-for-byte and does not apply to this stack. +# +# Required env vars: +# MODEL, MODEL_PATH, TP, CONC, KV_OFFLOADING, KV_OFFLOAD_BACKEND, +# TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, EP_SIZE, DP_ATTENTION + +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 + +echo "MODEL=$MODEL TP=$TP CONC=$CONC KV_OFFLOADING=$KV_OFFLOADING TOTAL_CPU_DRAM_GB=$TOTAL_CPU_DRAM_GB RESULT_DIR=$RESULT_DIR DURATION=$DURATION EP_SIZE=$EP_SIZE DP_ATTENTION=$DP_ATTENTION" + +if [[ -v SLURM_JOB_ID ]]; then + echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME" +fi + +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 + +if [[ -v 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 + +wait_for_amd_gpu_clean + +rocm-smi || true +amd-smi || true + +resolve_trace_source +install_agentic_deps + +# Require the ATOM Prometheus stream in every official result. AIPerf +# deduplicates this endpoint against its automatic localhost discovery. +export AIPERF_SERVER_METRICS_URLS="http://localhost:${PORT}/metrics" +export AIPERF_REQUIRED_SERVER_METRIC_PREFIX="atom:" + +# Long agentic turns against a 1M context: keep the client from timing out +# mid-request while the server is prefill-bound. Matches the vLLM K3 arm. +export AIPERF_HTTP_TCP_USER_TIMEOUT=900000 + +# VRAM space check +wait_for_amd_gpu_clean + +# ---- Server config ---------------------------------------------------------- +SERVER_LOG="$RESULT_DIR/server.log" +mkdir -p "$RESULT_DIR" + +SERVER_PID="" +cleanup_agentic_services() { + local exit_code=$? + trap - EXIT INT TERM + set +e + stop_background_process_tree "$SERVER_PID" "ATOM server" 60 + exit "$exit_code" +} +trap cleanup_agentic_services EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +# ---- KV offload ------------------------------------------------------------- +# K3 is a hybrid: Kimi Delta Attention carries a per-request recurrent state +# alongside the paged KV. Both tiers are switched together here, because the +# state tier is what makes a resumed agentic turn cheap and the paged KV tier +# alone cannot restore one. +OFFLOAD_ARGS=() + +case "$KV_OFFLOAD_BACKEND" in + "") + require_agentic_kv_offload_none + ;; + lmcache) + require_agentic_kv_offload_backend lmcache + + # TOTAL_CPU_DRAM_GB is the AGGREGATE budget from the matrix generator. + # LMCACHE_MAX_LOCAL_CPU_SIZE is per rank and every rank allocates its + # own, so the aggregate is divided by TP as the agentic README + # requires. Handing a rank the whole aggregate does not just overcommit + # -- it never finishes pinning and hangs the launch partway through. + export PYTHONHASHSEED=0 + export LMCACHE_LOCAL_CPU=True + export LMCACHE_MAX_LOCAL_CPU_SIZE="$((TOTAL_CPU_DRAM_GB / TP))" + # One chunk per hash block, so the KV grid and the state-checkpoint + # grid coincide and the joint load aims both legs at one boundary. + export LMCACHE_CHUNK_SIZE=256 + + # OFFLOAD_PROFILE is deliberately left unset (default 0). The source + # recipe sets it to 1, but that only turns on per-step offload + # statistics in the connector, and the numbers behind this submission + # were measured with it off. Noted here so the difference from the + # recipe reads as a choice rather than an omission. + + # CPU state-offload tier for the KDA recurrent state. + export OFFLOAD_STATE=1 + export OFFLOAD_STATE_STAGING_GROUPS=8 + export OFFLOAD_STATE_MIN_LOAD_TOKENS=0 + # Must be set: the staging buffer defaults to 2 chunks (8 MiB), one K3 + # state entry is 54.78 MiB, and a buffer too small to hold one entry + # makes the tier decline to build -- one log line, then nothing + # offloads, which reads exactly like a tier that is on and idle. + export OFFLOAD_GPU_STAGING_CHUNKS=16 + + OFFLOAD_ARGS=( + --kv-transfer-config + "{\"kv_connector\":\"lmcache_offload\",\"kv_role\":\"offload\"}" + ) + ;; + *) + echo "Unsupported KV_OFFLOAD_BACKEND: $KV_OFFLOAD_BACKEND (expected empty or lmcache)" >&2 + exit 1 + ;; +esac + +# ---- ATOM env --------------------------------------------------------------- +echo "Starting atom server..." +export PYTHONNOUSERSITE=1 + +# Required by ATOM: without it the aiter kernel logs flood the server log for +# the whole 3600 s replay. +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-WARNING}" +export AITER_SITUV2_A4W4=1 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export AITER_FLYDSL_STAGE2_FP8=1 +export ATOM_MLA_MAX_SPLIT_PER_BATCH=256 +# Anchor-only state checkpointing: the demand rung is 47% of checkpoint writes +# but reads back 2.8% of the time, against 85.2% for a prompt-end anchor, so it +# costs more in evictions than its reuse is worth on these traces. +export ATOM_STATE_CHECKPOINT_DEMAND=0 + +# ---- Per-concurrency knobs -------------------------------------------------- +case "$CONC" in + 1|4) + MAX_NUM_SEQS=32 + MAX_NUM_BATCHED_TOKENS=8192 + GPU_MEM_UTIL=0.88 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS="" + ;; + 8) + MAX_NUM_SEQS=16 + MAX_NUM_BATCHED_TOKENS=8192 + GPU_MEM_UTIL=0.88 + ATOM_ENABLE_REPLAYSSM=1 + STATE_CHECKPOINT_SLOTS=16 + ;; + 10) + MAX_NUM_SEQS=16 + MAX_NUM_BATCHED_TOKENS=4096 + GPU_MEM_UTIL=0.90 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS=16 + ;; + *) + echo "Unsupported CONC=$CONC" >&2 + exit 2 + ;; +esac +export ATOM_ENABLE_REPLAYSSM + +# Extra in-GPU state checkpoint slots beyond the in-flight floor. Checkpoints +# and live requests share one pool, so without this the room to retain a +# checkpoint is whatever max-num-seqs happens to leave. +STATE_CKPT_ARGS=() +if [ -n "$STATE_CHECKPOINT_SLOTS" ]; then + STATE_CKPT_ARGS=(--state-checkpoint-slots "$STATE_CHECKPOINT_SLOTS") +fi + +# ---- Speculative ------------------------------------------------------------ +# golden 2.51 at num_speculative_tokens 2 in +# https://github.com/SemiAnalysisAI/InferenceX/blob/main/golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml + +if [ "$CONC" = 1 ]; then + SIMULATE_ACC_LEN=3.75 + NUM_SPEC_TOKENS=6 +else + SIMULATE_ACC_LEN=2.51 + NUM_SPEC_TOKENS=2 +fi + +# spec-decode-acceptance-rate = (SIMULATE_ACC_LEN - 1) / NUM_SPEC_TOKENS +SPEC_ACCEPTANCE_RATE=$(awk "BEGIN{print ($SIMULATE_ACC_LEN-1)/$NUM_SPEC_TOKENS}") +if [ "${EVAL_ONLY}" = "true" ]; then + SPEC_ARGS=( + --method dspark + --draft-model Inferact/Kimi-K3-DSpark + --num-speculative-tokens "$NUM_SPEC_TOKENS" + ) +else + SPEC_ARGS=( + --method dspark + --draft-model Inferact/Kimi-K3-DSpark + --num-speculative-tokens "$NUM_SPEC_TOKENS" + --spec-decode-acceptance-rate "$SPEC_ACCEPTANCE_RATE" + ) +fi +echo "SIMULATE_ACC_LEN=$SIMULATE_ACC_LEN NUM_SPEC_TOKENS=$NUM_SPEC_TOKENS SPEC_ACCEPTANCE_RATE=$SPEC_ACCEPTANCE_RATE" + +# ---- LLM server ------------------------------------------------------------- +ATOM_CMD=( + python -m atom.entrypoints.openai_server + --model "$MODEL_PATH" + --host 0.0.0.0 + --server-port "$PORT" + --trust-remote-code + --tensor-parallel-size "$TP" + --kv_cache_dtype fp8 + --block-size 128 + --max-num-seqs "$MAX_NUM_SEQS" + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" + --gpu-memory-utilization "$GPU_MEM_UTIL" + --enable_prefix_caching + --state-checkpoint-interval-tokens -1 + "${STATE_CKPT_ARGS[@]}" + --online_quant_config '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*self_attn.[qkv]_conv1d*","*block_sparse_moe.experts*","*block_sparse_moe.routed_expert_*","*vision_tower*","*mm_projector*"]}' + "${SPEC_ARGS[@]}" + "${OFFLOAD_ARGS[@]}" +) +write_command "$RESULT_DIR/server_command.txt" "${ATOM_CMD[@]}" +"${ATOM_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 ---------------------------------------------------------- +if [ "${EVAL_ONLY}" = "true" ]; then + run_eval --port "$PORT" +else + build_replay_cmd "$RESULT_DIR" + REPLAY_CMD+=" --apply-chat-template" + run_agentic_replay_and_write_outputs "$RESULT_DIR" +fi diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 59f655945..b57a6d1fa 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -637,6 +637,35 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [10], spec-decoding: mtp } +# Kimi-K3 MXFP4 agentic-coding benchmark on MI355X via ATOM with DSpark +# speculative decoding (2 draft tokens -> golden AL 2.51, +# golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml, +# the same golden the vLLM arm feeds to synthetic_acceptance_length). +# Companion to kimik3-fp4-mi355x-vllm-agentic-mtp: same checkpoint, same +# runner, same concurrency points, so the two engines are directly comparable. +# TP8 only -- the 1.56 TB MXFP4 checkpoint is ~195 GB/GPU and TP4 cannot load, +# which is also why there is no DP-attention arm. +# Concurrency 1 and 4 are GPU-resident. 8 and 10 add the LMCache DRAM tier and +# ATOM's CPU state-offload tier together: K3 is a hybrid, so a resumed agentic +# turn needs the KDA recurrent state back, and the paged KV tier alone cannot +# restore one. dram-utilization 0.086 puts the aggregate budget at 257 GB, so +# the script's divide-by-TP lands on exactly the 32 GB per rank the recipe was +# measured with (0.085 gives 254, which floors to 31). +kimik3-fp4-mi355x-atom-agentic-mtp: + image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0818 + model: moonshotai/Kimi-K3 + model-prefix: kimik3 + runner: cluster:mi355x-amds + precision: fp4 + framework: atom + multinode: false + scenarios: + agentic-coding: + - dram-utilization: 0.086 + search-space: + - { tp: 8, kv-offloading: none, conc-list: [1, 4], spec-decoding: mtp } + - { tp: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [8, 10], spec-decoding: mtp } + dsr1-fp4-mi355x-sglang-disagg: 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 93b6a7c1b..40901777a 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6196,6 +6196,20 @@ - "Set max-num-seqs to twice the concurrency, select CUDA graph capture sizes by concurrency, and cap max-num-batched-tokens at 16384." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2576 +- config-keys: + - kimik3-fp4-mi355x-atom-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Add the MI355X Kimi-K3 FP4 ATOM AgentX submission with DSpark speculative decoding (2 draft tokens): TP8 at concurrency 1 and 4 is GPU-resident with no KV offload, while concurrency 8 and 10 add the LMCache DRAM tier together with ATOM's CPU state-offload tier." + - "Both offload tiers are switched together because Kimi-K3 is a hybrid: Kimi Delta Attention carries a per-request recurrent state alongside the paged KV, so a resumed agentic turn needs the state back and the paged KV tier alone cannot restore one." + - "Use AITER INT4 quick-reduce, AITER SITUV2 A4W4, FlyDSL stage-2 FP8, ptpc_fp8 online quantization excluding embeddings, lm_head, conv1d projections, experts, and the multimodal tower, an FP8 KV cache, and a 128-token page." + - "Set the simulated acceptance length to 2.54, giving spec-decode-acceptance-rate (2.54 - 1) / 2 = 0.77, which is the value the submitted ATOM numbers were measured under. This is 1.2% above the committed golden 2.51 that the vLLM arm feeds to synthetic_acceptance_length, and the script carries a review note saying so; set it to 2.51 for a strictly like-for-like sweep. Evaluations use real acceptance." + - "Divide the aggregate TOTAL_CPU_DRAM_GB budget by TP for LMCACHE_MAX_LOCAL_CPU_SIZE, which is a per-rank setting; dram-utilization 0.085 lands that at roughly 32 GB per rank." + - "Set OFFLOAD_GPU_STAGING_CHUNKS to 16 because the staging buffer defaults to 2 chunks (8 MiB) and one Kimi-K3 state entry is 54.78 MiB; a buffer too small to hold one entry makes the state tier decline to build after a single log line." + - "Select max-num-seqs, max-num-batched-tokens, gpu-memory-utilization, and in-GPU state-checkpoint slots by concurrency, and place state checkpoints at the prompt-end anchor only." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2670 + - config-keys: - kimik3-fp4-b200-dynamo-vllm-agentic-dspark scenario-type: