[AMD][WIP][AGENTX]: Kimi-K3 DSpark on MI355X #2496
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
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,209 @@ | ||
| #!/usr/bin/env bash | ||
| set -euo pipefail | ||
| set -x | ||
|
|
||
| # Agentic trace replay benchmark for Kimi-K3 MXFP4 on MI355X using vLLM, | ||
| # DSpark speculative decoding, and LMCache MP host-DRAM KV offload. | ||
| # | ||
| # Required env vars: | ||
| # MODEL, TP, CONC, KV_OFFLOADING, KV_OFFLOAD_BACKEND, | ||
| # TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, EP_SIZE, PORT | ||
|
|
||
| source "$(dirname "$0")/../../benchmark_lib.sh" | ||
|
|
||
| check_env_vars \ | ||
| MODEL TP CONC KV_OFFLOADING KV_OFFLOAD_BACKEND TOTAL_CPU_DRAM_GB \ | ||
| RESULT_DIR DURATION EP_SIZE PORT | ||
|
|
||
| if [ "$TP" -ne 8 ]; then | ||
| echo "Error: Kimi-K3 MXFP4 requires TP=8 on MI355X; got TP=$TP." >&2 | ||
| exit 1 | ||
| fi | ||
|
|
||
| if [ "$EP_SIZE" -gt 1 ]; then | ||
| echo "Error: this Kimi-K3 DSpark recipe supports pure TP8 only; got EP_SIZE=$EP_SIZE." >&2 | ||
| exit 1 | ||
| fi | ||
|
|
||
| require_agentic_kv_offload_backend lmcache | ||
|
|
||
| if [[ -n "${SLURM_JOB_ID:-}" ]]; then | ||
| echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}" | ||
| fi | ||
|
|
||
| if [[ -n "${ROCR_VISIBLE_DEVICES:-}" ]]; then | ||
| export HIP_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" | ||
| fi | ||
|
|
||
| # Do not start the large Kimi-K3 load while a previous Slurm job is still | ||
| # releasing VRAM on the allocated MI355X GPUs. | ||
| wait_for_amd_gpu_clean | ||
|
|
||
| # The cluster launcher mounts its persistent Hugging Face cache here. MODEL_PATH | ||
| # may instead point at a pre-staged snapshot. | ||
| export HF_HUB_CACHE="${HF_HUB_CACHE:-/models/huggingface_hub}" | ||
| 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_trace_source | ||
| install_agentic_deps | ||
|
|
||
| # Environment used by the validated Kimi-K3 + DSpark + LMCache image. | ||
| export PYTHONNOUSERSITE=1 | ||
| export PYTHONHASHSEED=42 | ||
| export VLLM_ROCM_USE_AITER=1 | ||
| export VLLM_ROCM_USE_AITER_MLA=1 | ||
| export VLLM_ROCM_USE_SKINNY_GEMM=0 | ||
| export AITER_SITUV2_A8W4=1 | ||
| export AITER_BF16_FP8_MOE_BOUND=0 | ||
| export VLLM_ENABLE_K3_LATENT_MOE_TAIL_FUSION=1 | ||
| export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS="${VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS:-1200}" | ||
| export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-7200}" | ||
| export HSA_NO_SCRATCH_RECLAIM=1 | ||
| export SAFETENSORS_FAST_GPU=1 | ||
|
|
||
| SERVER_LOG="$RESULT_DIR/server.log" | ||
| LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" | ||
| mkdir -p "$RESULT_DIR" | ||
|
|
||
| SERVER_PID="" | ||
| LMCACHE_PID="" | ||
|
|
||
| cleanup_services() { | ||
| local exit_code=$? | ||
| trap - EXIT INT TERM | ||
| set +e | ||
| stop_background_process_tree "$SERVER_PID" "vLLM server" 60 | ||
| stop_background_process_tree "$LMCACHE_PID" "LMCache server" 30 | ||
| exit "$exit_code" | ||
| } | ||
| trap cleanup_services EXIT | ||
| trap 'exit 130' INT | ||
| trap 'exit 143' TERM | ||
|
|
||
| wait_for_lmcache_ready() { | ||
| { set +x; } 2>/dev/null | ||
| local attempts="${LMCACHE_READY_ATTEMPTS:-1800}" | ||
| local paths=(/healthcheck /health /v1/health /status /) | ||
| local i path | ||
|
|
||
| for ((i = 1; i <= attempts; i++)); do | ||
| for path in "${paths[@]}"; do | ||
| if curl --output /dev/null --silent --fail \ | ||
| "http://127.0.0.1:${LMCACHE_HTTP_PORT}${path}"; then | ||
| echo "LMCache server healthy after ${i}s (endpoint ${path})" | ||
| set -x | ||
| return 0 | ||
| fi | ||
| done | ||
|
|
||
| if ! kill -0 "$LMCACHE_PID" 2>/dev/null; then | ||
| echo "LMCache server exited before becoming healthy. Log follows:" >&2 | ||
| tail -200 "$LMCACHE_LOG" >&2 || true | ||
| exit 1 | ||
| fi | ||
|
|
||
| if ((i % 60 == 0)); then | ||
| echo "Still waiting for LMCache (${i}s/${attempts}s, L1=${LMCACHE_L1_SIZE_GB}GB)" | ||
| fi | ||
| sleep 1 | ||
| done | ||
|
|
||
| echo "Timed out waiting for LMCache on port $LMCACHE_HTTP_PORT. Log follows:" >&2 | ||
| tail -200 "$LMCACHE_LOG" >&2 || true | ||
| exit 1 | ||
| } | ||
|
|
||
| # This is the pool size validated on the MI355X host. Keep it overridable for | ||
| # hosts with a different /dev/shm allocation, but never exceed the matrix's DRAM | ||
| # budget. | ||
| LMCACHE_L1_SIZE_GB="${LMCACHE_L1_SIZE_GB:-906}" | ||
| if [ "$LMCACHE_L1_SIZE_GB" -gt "$TOTAL_CPU_DRAM_GB" ]; then | ||
| echo "Error: LMCache L1=${LMCACHE_L1_SIZE_GB}GB exceeds the generated DRAM budget ${TOTAL_CPU_DRAM_GB}GB." >&2 | ||
| exit 1 | ||
| fi | ||
|
|
||
| LMCACHE_HTTP_PORT="${LMCACHE_HTTP_PORT:-8080}" | ||
| LMCACHE_MP_PORT="${LMCACHE_MP_PORT:-6000}" | ||
| LMCACHE_CMD=( | ||
| lmcache server | ||
| --host 127.0.0.1 | ||
| --port "$LMCACHE_MP_PORT" | ||
| --http-host 127.0.0.1 | ||
| --http-port "$LMCACHE_HTTP_PORT" | ||
| --l1-size-gb "$LMCACHE_L1_SIZE_GB" | ||
| --l1-init-size-gb 20 | ||
| --l1-read-ttl-seconds 7200 | ||
| --chunk-size 1536 | ||
| --max-workers 8 | ||
| --eviction-trigger-watermark 0.85 | ||
| --eviction-ratio 0.10 | ||
| --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 | ||
|
|
||
| # AgentX concurrency counts live session trees, not individual requests. | ||
| # Subagent fan-out can push instantaneous request concurrency above CONC, so | ||
| # leave 2x headroom rather than clipping those bursts at the scheduler. | ||
| MAX_NUM_SEQS="${MAX_NUM_SEQS:-$((2 * CONC))}" | ||
| SPEC_NUM_TOKENS="${SPEC_NUM_TOKENS:-3}" | ||
| SPEC_DRAFT_MODEL="${SPEC_DRAFT_MODEL:-Inferact/Kimi-K3-DSpark}" | ||
| SPEC_DRAFT_MODEL_PATH="${SPEC_DRAFT_MODEL_PATH:-$SPEC_DRAFT_MODEL}" | ||
| if [[ "$SPEC_DRAFT_MODEL_PATH" == "$SPEC_DRAFT_MODEL" ]]; then | ||
| hf download "$SPEC_DRAFT_MODEL" | ||
| fi | ||
| SPEC_CONFIG="{\"model\":\"${SPEC_DRAFT_MODEL_PATH}\",\"num_speculative_tokens\":${SPEC_NUM_TOKENS},\"method\":\"dspark\",\"attention_backend\":\"TRITON_MLA\",\"draft_sample_method\":\"probabilistic\",\"rejection_sample_method\":\"block\"}" | ||
|
Collaborator
There was a problem hiding this comment. InferenceX/benchmarks/single_node/agentic/kimik3_fp4_b300_vllm_mtp.sh Lines 172 to 191 in 09e6cd0 |
||
| 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\":\"tcp://127.0.0.1\",\"lmcache.mp.port\":${LMCACHE_MP_PORT}}}" | ||
|
|
||
| VLLM_CMD=( | ||
| vllm serve "$MODEL_PATH" | ||
| --served-model-name "$MODEL" | ||
| --host 0.0.0.0 | ||
| --port "$PORT" | ||
| --trust-remote-code | ||
| --moe-backend auto | ||
| --tensor-parallel-size "$TP" | ||
| --load-format auto | ||
| --gpu-memory-utilization 0.85 | ||
| --mm-encoder-tp-mode data | ||
| --max-num-seqs "$MAX_NUM_SEQS" | ||
| --max-num-batched-tokens 3000 | ||
| --enable-auto-tool-choice | ||
| --tool-call-parser kimi_k3 | ||
| --reasoning-parser kimi_k3 | ||
| --enable-prefix-caching | ||
| --mamba-cache-mode align | ||
| --kv-cache-dtype fp8 | ||
| --attention-backend TRITON_MLA | ||
| --compilation-config '{"cudagraph_mode":"PIECEWISE","custom_ops":["+fused_rms_norm_gated"]}' | ||
| --speculative-config "$SPEC_CONFIG" | ||
| --kv-transfer-config "$KV_TRANSFER_CONFIG" | ||
| ) | ||
| 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 "vLLM 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 | ||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| # Serving-stack patch waiver for PR #2496 | ||
|
|
||
| ## What is patched | ||
|
|
||
| The Kimi-K3 MI355X AgentX configuration uses | ||
| `haihub007/hc-gpu:kimik3-dspark-lmcache-pr50619-20260805`, which contains the | ||
| production-file portion of | ||
| [vLLM PR #50619](https://github.com/vllm-project/vllm/pull/50619) rebased onto | ||
| the image's vLLM revision `124154a88`. It fixes Kimi-K3 TP8 DSpark FP8 MLA | ||
| verification by preserving native causal multi-token Gluon queries and the | ||
| target model's CUDA-graph policy. The published image digest is | ||
| `sha256:3b5b0131f89be08d34d431f2c4e45e32606f9dff0e3aace79272489e788e6eda`. | ||
|
|
||
| ## Why the unmodified image cannot run this benchmark | ||
|
|
||
| The image's current vLLM expands each DSpark verification token into an | ||
| independent Gluon MLA row. During engine warmup, 16 requests with the | ||
| four-token draft configuration produced 144 rows and exceeded the kernel's | ||
| maximum batch size of 128, preventing EngineCore from starting. Reducing the | ||
| draft to three tokens and temporarily enabling eager execution lowers risk, | ||
| but neither is a complete serving-stack correction for multi-token FP8 MLA | ||
| verification. | ||
|
|
||
| The failure is recorded in | ||
| [run 30985536172](https://github.com/SemiAnalysisAI/InferenceX/actions/runs/30985536172/job/92239885715). | ||
| The upstream Kimi-K3 ROCm roadmap is tracked in | ||
| [vLLM issue #50682](https://github.com/vllm-project/vllm/issues/50682). | ||
|
|
||
| ## Removal plan | ||
|
|
||
| Remove this waiver after PR #50619, or its upstream replacement, is merged and | ||
| included in an upstream `vllm/vllm-openai-rocm` image validated for Kimi-K3 | ||
| TP8 DSpark with FP8 KV cache and LMCache. Replace the personal image with that | ||
| upstream image in the same update. |
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
Submodule aiperf
updated
16 files
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🔴 This new agentic recipe drives the same Kimi-K3 trace-replay path as
kimik3_fp4_b300_vllm.sh/kimik3_fp4_b300_vllm_mtp.shbut never exportsAIPERF_HTTP_TCP_USER_TIMEOUT(set by every other vLLM agentic recipe, incl. the same-SKUdsv4_fp4_mi355x_vllm.sh) orVLLM_HTTP_TIMEOUT_KEEP_ALIVE(set by the identical-model B300 siblings, which document a real production incident where omitting it aborted a whole benchmark job). Without these, a long warmup request or an inter-turn idle gap can trip AIPerf into treating a healthy server as dead, aborting the entire agentic sweep.Extended reasoning...
The bug:
kimik3_fp4_mi355x_agentic_dspark.shsets up an extensive environment block (lines 65-92) for the validated Kimi-K3 + DSpark + LMCache image, but it never exportsAIPERF_HTTP_TCP_USER_TIMEOUTorVLLM_HTTP_TIMEOUT_KEEP_ALIVE. Every other vLLM agentic recipe in the repo exportsAIPERF_HTTP_TCP_USER_TIMEOUT=900000(e.g.dsv4_fp4_mi355x_vllm.sh:61, the same MI355X/vLLM SKU+framework), and the two same-model B300 siblings —kimik3_fp4_b300_vllm.shandkimik3_fp4_b300_vllm_mtp.sh— additionally setVLLM_HTTP_TIMEOUT_KEEP_ALIVE=900with an explicit comment explaining why it's necessary.Why this matters — the documented failure mode:
kimik3_fp4_b300_vllm.sh:77-86describes a real production incident: AIPerf pins one pooled keep-alive HTTP connection per agentic session and reuses it across conversation turns. If the inter-turn idle gap exceeds the server's keep-alive timeout (5s by default), the client can grab the socket for reuse in the exact instant the server closes it, producing anaiohttp.ServerDisconnectedError. AIPerf treats this as a terminal warmup failure — not a retryable blip — and aborts the entire benchmark job. The comment explicitly citesrun 30324907690, where this "killed the dram c4 arm ~15 min into run" against a server that "kept serving after the client gave up." Separately,AIPERF_HTTP_TCP_USER_TIMEOUT=900000exists because agentic warmup dispatches large prompts at once, and the client needs up to 15 minutes of TCP progress before AIPerf gives up on the connection as dead — this is the universal convention across every vLLM/sglang agentic recipe in the repo.The code path that triggers it: This script drives the identical agentic mechanism as the B300 siblings —
resolve_trace_source,build_replay_cmd, andrun_agentic_replay_and_write_outputs(lines 92, 189-192) — for the same Kimi-K3 model with the same DSpark draft, on a 1.5+ TB checkpoint whose warmup requests are long and whose per-turn idle gaps are exactly the pattern that triggers the keep-alive race. Nothing inbenchmark_lib.shsets a default for either variable — each recipe is individually responsible for exporting them — and grepping the new file confirms neither is present.Why nothing else prevents it: The script does set generous timeouts elsewhere (
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200,VLLM_ENGINE_READY_TIMEOUT_S=7200), showing the author is aware that this workload needs extended timeouts, but those don't cover the AIPerf-side TCP dead-connection detection or the server-side HTTP keep-alive window that caused the documented B300 incident.VLLM_USE_RUST_FRONTENDisn't set here, so the timeout technically applies to uvicorn'stimeout_keep_aliveon the standard FastAPI frontend (default 5s) rather than the Rust frontend — but the default is the same 5s, so the same race is structurally possible; the incident report just happens to have been diagnosed on the Rust frontend.Step-by-step proof of the failure:
vllm servewith a large model and starts an agentic trace-replay session viarun_agentic_replay_and_write_outputs.VLLM_HTTP_TIMEOUT_KEEP_ALIVEis unset), the server closes the idle socket.aiohttp.ServerDisconnectedError.Fix: Add the same two exports used by the B300 siblings to this script's environment block:
Impact if unfixed: Since the PR's own test plan leaves the "MI355X End-to-End AgentX smoke validation" and "canonical full sweep" checkboxes unfixed, this failure mode has not yet been exercised against this new recipe, so the risk is currently unvalidated. The consequence when it does trigger is not degraded numbers but a spurious, terminal abort of the whole agentic sweep (wasted GPU time on the 8xMI355X node, a false-negative CI run), matching the exact incident the identical-model B300 recipes were patched to avoid.