diff --git a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh index 922921f284..034bbb41f5 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh @@ -5,8 +5,8 @@ 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. +# Companion to kimik3_fp4_mi355x_mtp.sh, which runs the same checkpoint 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 @@ -17,14 +17,14 @@ set -x # byte-for-byte and does not apply to this stack. # # Required env vars: -# MODEL, MODEL_PATH, TP, CONC, KV_OFFLOADING, KV_OFFLOAD_BACKEND, +# MODEL, MODEL_PATH, TP, DCP_SIZE, 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" +echo "MODEL=$MODEL TP=$TP DCP_SIZE=${DCP_SIZE:-1} 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" @@ -85,11 +85,117 @@ trap cleanup_agentic_services EXIT trap 'exit 130' INT trap 'exit 143' TERM +# ---- Per-concurrency knobs -------------------------------------------------- +# Concurrency 1-4 is the latency floor: everything GPU-resident, no decode +# context parallelism, the deepest draft the golden curve publishes, and an +# 8192-token prefill step. +# From concurrency 8 up, decode is KV-bandwidth-bound over 100k+ token agentic +# contexts, so DCP8 shards the KV read across all 8 GPUs and the LMCache DRAM +# tier backs the paged KV. +# +# The LMCache paged-KV tier itself is not switched here: it follows +# kv-offloading in configs/amd-master.yaml, which is `none` for concurrency 1-4 +# and `dram` from concurrency 8 up. What this block chooses is +# STATE_OFFLOAD_CPU_GIB, the per-rank slice of that CPU budget reserved for the +# Kimi Delta Attention recurrent state; 0 leaves the whole budget to the paged +# KV and keeps the state in GPU checkpoints instead. +case "$CONC" in + # No KV offload; the working set fits in HBM. + 1|2|4) + MAX_NUM_SEQS=32 + MAX_NUM_BATCHED_TOKENS=8192 + GPU_MEM_UTIL=0.88 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS="" + NUM_SPEC_TOKENS=7 + SPEC_DECODE_AL=3.84 + STATE_OFFLOAD_CPU_GIB=0 + ;; + # LMCache paged-KV tier on, whole per-rank budget; state stays in GPU + # checkpoints, which is what ReplaySSM replays from. + 8) + MAX_NUM_SEQS=32 + MAX_NUM_BATCHED_TOKENS=4096 + GPU_MEM_UTIL=0.88 + ATOM_ENABLE_REPLAYSSM=1 + STATE_CHECKPOINT_SLOTS=96 + NUM_SPEC_TOKENS=3 + SPEC_DECODE_AL=3.00 + STATE_OFFLOAD_CPU_GIB=0 + ;; + 12) + MAX_NUM_SEQS=24 + MAX_NUM_BATCHED_TOKENS=4096 + GPU_MEM_UTIL=0.88 + ATOM_ENABLE_REPLAYSSM=1 + STATE_CHECKPOINT_SLOTS=96 + NUM_SPEC_TOKENS=3 + SPEC_DECODE_AL=3.00 + STATE_OFFLOAD_CPU_GIB=0 + ;; + # LMCache paged-KV tier plus ATOM's CPU state tier, 32 GB/rank carved out + # for the KDA recurrent state. + 16) + MAX_NUM_SEQS=32 + MAX_NUM_BATCHED_TOKENS=8192 + GPU_MEM_UTIL=0.86 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS="" + NUM_SPEC_TOKENS=3 + SPEC_DECODE_AL=3.00 + STATE_OFFLOAD_CPU_GIB=32 + ;; + 32) + MAX_NUM_SEQS=64 + MAX_NUM_BATCHED_TOKENS=8192 + GPU_MEM_UTIL=0.86 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS="" + NUM_SPEC_TOKENS=0 + SPEC_DECODE_AL=0 + STATE_OFFLOAD_CPU_GIB=32 + ;; + 40) + MAX_NUM_SEQS=80 + MAX_NUM_BATCHED_TOKENS=8192 + GPU_MEM_UTIL=0.86 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS="" + NUM_SPEC_TOKENS=0 + SPEC_DECODE_AL=0 + STATE_OFFLOAD_CPU_GIB=32 + ;; + 56) + MAX_NUM_SEQS=72 + MAX_NUM_BATCHED_TOKENS=4096 + GPU_MEM_UTIL=0.88 + ATOM_ENABLE_REPLAYSSM=0 + STATE_CHECKPOINT_SLOTS="" + NUM_SPEC_TOKENS=0 + SPEC_DECODE_AL=0 + STATE_OFFLOAD_CPU_GIB=32 + ;; + *) + 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 + # ---- 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. +# alongside the paged KV. The paged KV rides this LMCache tier from +# concurrency 8 up; from concurrency 16 up the CPU state tier is switched on +# alongside it, 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 @@ -100,32 +206,39 @@ case "$KV_OFFLOAD_BACKEND" in 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. + # LMCACHE_MAX_LOCAL_CPU_SIZE and OFFLOAD_STATE_CPU_SIZE are 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. + PER_RANK_CPU_GB="$((TOTAL_CPU_DRAM_GB / TP))" + LMCACHE_CPU_GB="$((PER_RANK_CPU_GB - STATE_OFFLOAD_CPU_GIB))" + 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. + export LMCACHE_MAX_LOCAL_CPU_SIZE="$LMCACHE_CPU_GB" + # DCP-locked: the offload hash block is block-size(128) x dcp(8) = 1024, + # so the KV grid and the state-checkpoint grid coincide and the joint + # load aims both legs at one boundary. 512 or 2048 misaligns it. + export LMCACHE_CHUNK_SIZE=1024 + export OFFLOAD_KV_FOR_HYBRID=1 + # Statistics only -- per-step offload counters in the connector. Kept on + # because the submitted numbers were measured with it on. + export OFFLOAD_PROFILE=1 - # 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 + if [ "$STATE_OFFLOAD_CPU_GIB" -gt 0 ]; then + # CPU state-offload tier for the KDA recurrent state. + export OFFLOAD_STATE=1 + export OFFLOAD_STATE_CPU_SIZE="$STATE_OFFLOAD_CPU_GIB" + 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=32 + fi OFFLOAD_ARGS=( --kv-transfer-config @@ -148,76 +261,27 @@ 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 ------------------------------------------------------------ # https://github.com/SemiAnalysisAI/InferenceX/blob/main/golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml -# 6 draft tokens -> AL 3.75 -# 2 draft tokens -> AL 2.51 -# https://github.com/ROCm/ATOM/pull/1948 -if [ "$CONC" = 1 ]; then - SPEC_DECODE_AL=3.75 - NUM_SPEC_TOKENS=6 -else - SPEC_DECODE_AL=2.51 - NUM_SPEC_TOKENS=2 -fi - -if [ "${EVAL_ONLY}" = "true" ]; then +# 7 draft tokens -> AL 3.84 +# 3 draft tokens -> AL 3.00 +# Concurrency 32 and up serve without a draft model: past the throughput knee +# the draft forward no longer pays for itself against the resident batch. +SPEC_ARGS=() +if [ "$NUM_SPEC_TOKENS" -gt 0 ]; 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-length "$SPEC_DECODE_AL" - ) + if [ "${EVAL_ONLY}" != "true" ]; then + SPEC_ARGS+=(--spec-decode-acceptance-length "$SPEC_DECODE_AL") + fi fi echo "SPEC_DECODE_AL=$SPEC_DECODE_AL NUM_SPEC_TOKENS=$NUM_SPEC_TOKENS" @@ -229,6 +293,7 @@ ATOM_CMD=( --server-port "$PORT" --trust-remote-code --tensor-parallel-size "$TP" + --decode-context-parallel-size "${DCP_SIZE:-1}" --kv_cache_dtype fp8 --block-size 128 --max-num-seqs "$MAX_NUM_SEQS" diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 2195bad7e2..e9bb877656 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -638,21 +638,25 @@ kimik3-fp4-mi355x-vllm-agentic-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. +# speculative decoding. Acceptance is pinned to the committed golden curve in +# golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml +# at both draft lengths used here: 7 draft tokens -> AL 3.84 at concurrency 1-4, +# 3 draft tokens -> AL 3.00 from concurrency 8 up. +# Companion to kimik3-fp4-mi355x-vllm-agentic-mtp: same checkpoint, same runner. # 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). +# Concurrency 1-4 is the latency floor and stays GPU-resident. From concurrency +# 8 up, decode is KV-bandwidth-bound over 100k+ token contexts, so dcp-size 8 +# shards the KV read across all 8 GPUs and the LMCache DRAM tier backs the +# paged KV. +# Two utilization blocks because the per-rank CPU split differs: +# 0.534 -> 1601 GB aggregate -> 200 GB/rank, all paged KV (conc 8, 12). +# 0.598 -> 1793 GB aggregate -> 224 GB/rank, split 192 GB paged KV + 32 GB +# for ATOM's CPU state tier (conc 16 and up). K3 is a hybrid, so a +# resumed agentic turn needs the KDA recurrent state back and the +# paged KV tier alone cannot restore one. kimik3-fp4-mi355x-atom-agentic-mtp: - image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0820 + image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0821 model: moonshotai/Kimi-K3 model-prefix: kimik3 runner: cluster:mi355x-amds @@ -661,10 +665,13 @@ kimik3-fp4-mi355x-atom-agentic-mtp: multinode: false scenarios: agentic-coding: - - dram-utilization: 0.086 + - dram-utilization: 0.534 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 } + - { tp: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } + - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [8, 12], spec-decoding: mtp } + - dram-utilization: 0.598 + search-space: + - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [16, 32, 40, 56], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index d90a270383..875979feba 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6388,6 +6388,21 @@ - "Bump the shared TileRT dependency and router metadata to 0.1.5.post3 for both fixed-seq-length and AgentX configurations while retaining the AgentX 1800-second queue timeout" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2650 +- config-keys: + - kimik3-fp4-mi355x-atom-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Refresh the MI355X Kimi-K3 FP4 ATOM AgentX submission onto image kimi_k3_agentic_0821 and replace the published concurrency list [1, 4, 8, 10] with [1, 2, 4, 8, 12, 16, 32, 40, 56], retuning every point." + - "Shard the decode KV read across all 8 GPUs with decode-context-parallel-size 8 from concurrency 8 up. At 100k+ token agentic contexts decode is KV-bandwidth-bound; concurrency 1-4 leaves it off because a lightly loaded node gains nothing from the extra collectives." + - "Pin acceptance to the committed golden curve in golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml at both draft lengths: --spec-decode-acceptance-length 3.84 at 7 draft tokens for concurrency 1-4, and 3.00 at 3 draft tokens for concurrency 8 through 16. Evaluations drop the flag and use real acceptance. Concurrency 32 and up serve without a draft model." + - "Select max-num-seqs, max-num-batched-tokens, gpu-memory-utilization, ReplaySSM, and in-GPU state-checkpoint slots by concurrency: 32/8192/0.88 with ReplaySSM off at concurrency 1-4, 32 and 24 with a 4096-token prefill step and ReplaySSM on with 96 checkpoint slots at concurrency 8 and 12, and 32/64/80/72 with utilization 0.86-0.88 and ReplaySSM off from concurrency 16 up, where the Kimi Delta Attention recurrent state rides the CPU state-offload tier instead of in-GPU checkpoints." + - "Divide the aggregate TOTAL_CPU_DRAM_GB budget by TP for the per-rank LMCACHE_MAX_LOCAL_CPU_SIZE and OFFLOAD_STATE_CPU_SIZE. dram-utilization 0.534 lands 200 GB per rank of paged KV at concurrency 8 and 12; 0.598 lands 224 GB per rank from concurrency 16 up, split 192 GB paged KV plus 32 GB for the state tier." + - "Set LMCACHE_CHUNK_SIZE to 1024 rather than the 128-token page: with decode context parallelism the offload hash block is block-size times dcp, so any other chunk misaligns it." + - "Set OFFLOAD_GPU_STAGING_CHUNKS to 32 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." + - "Keep 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, a 128-token page, and prompt-end-anchor-only state checkpointing. Drop ATOM_MLA_MAX_SPLIT_PER_BATCH, which the retuned recipe does not set." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2716 + - config-keys: - dsv4-fp4-mi355x-sglang-agentic-mtp scenario-type: