From c8cd52467ee355d4c19d4af4c875cd352a440574 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Mon, 24 Aug 2026 13:55:54 +0800 Subject: [PATCH 1/5] feat(agentx): retune Kimi-K3 FP4 MI355X ATOM DSpark recipe on _0821 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refresh the MI355X Kimi-K3 ATOM AgentX submission onto image kimi_k3_agentic_0821 and replace concurrency [1, 4, 8, 10] with [1, 2, 4, 8, 12, 16, 32, 40, 56], retuning every point. Concurrency 1-4 is the latency floor: GPU-resident, no decode context parallelism, 7 draft tokens at golden AL 3.84, 8192-token prefill step. From concurrency 8 up decode is KV-bandwidth-bound over 100k+ token agentic contexts, so decode-context-parallel-size 8 shards the KV read across all 8 GPUs and the LMCache DRAM tier backs the paged KV. Concurrency 8 and 12 run 3 draft tokens at golden AL 3.00 with ReplaySSM and 96 in-GPU state checkpoint slots; from concurrency 16 up the KDA recurrent state moves to ATOM's CPU state-offload tier instead. Concurrency 32 and up serve without a draft model. Both acceptance lengths come from the committed golden curve in golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml (7 -> 3.84, 3 -> 3.00). Evaluations drop the flag and use real acceptance. LMCACHE_MAX_LOCAL_CPU_SIZE and OFFLOAD_STATE_CPU_SIZE are per rank, so the aggregate TOTAL_CPU_DRAM_GB is divided by TP as the agentic README requires: dram-utilization 0.534 gives 200 GB/rank of paged KV at concurrency 8 and 12, and 0.598 gives 224 GB/rank from concurrency 16 up, split 192 GB paged KV plus 32 GB for the state tier. 将 MI355X 上 Kimi-K3 的 ATOM AgentX 提交刷新到 kimi_k3_agentic_0821 镜像， 并把并发点从 [1, 4, 8, 10] 换成 [1, 2, 4, 8, 12, 16, 32, 40, 56]，逐点重调。 并发 1-4 为时延下界：全部驻留 GPU，不开 decode context parallel，草稿 7 token 对应 golden AL 3.84，prefill 步长 8192。并发 8 起，10 万 token 以上的 agentic 上下文使 decode 受 KV 带宽约束，因此用 decode-context-parallel-size 8 把 KV 读取 分摊到 8 张卡，并由 LMCache DRAM 层承接分页 KV。并发 8 与 12 使用草稿 3 token （golden AL 3.00）、开启 ReplaySSM 并保留 96 个 GPU 内 state checkpoint 槽位； 并发 16 起，KDA 循环状态改由 ATOM 的 CPU state-offload 层承载。并发 32 及以上 不加载草稿模型。 两个接受长度均取自仓库内已提交的 golden 曲线（7 -> 3.84，3 -> 3.00）；评测 （eval）不传该参数，使用真实接受率。 LMCACHE_MAX_LOCAL_CPU_SIZE 与 OFFLOAD_STATE_CPU_SIZE 均为每 rank 设置，因此按 agentic README 的要求将聚合预算 TOTAL_CPU_DRAM_GB 除以 TP：dram-utilization 0.534 在并发 8/12 下给到每 rank 200 GB 分页 KV；0.598 在并发 16 起给到每 rank 224 GB，拆为 192 GB 分页 KV + 32 GB state 层。 Co-Authored-By: Claude Opus 5 (1M context) --- .../agentic/kimik3_fp4_mi355x_atom_mtp.sh | 241 +++++++++++------- configs/amd-master.yaml | 37 +-- perf-changelog.yaml | 15 ++ 3 files changed, 190 insertions(+), 103 deletions(-) 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 922921f28..034bbb41f 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 67a060539..28895e2e6 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 a9d2492a0..aa143852c 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6387,3 +6387,18 @@ - "Replace the internal TileRT 0.1.5.post2+inferencex.1 queueing backport with the official 0.1.5.post3 PyPI wheel now that upstream has published it; remove the internal wheel builder and router patch" - "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 From d17b64fa198d0cd529a769fb9d669725bacb13c5 Mon Sep 17 00:00:00 2001 From: seungrokj <144636725+seungrokj@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:08:10 +0900 Subject: [PATCH 2/5] Update amd-master.yaml --- configs/amd-master.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 28895e2e6..073885f3a 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -671,7 +671,7 @@ kimik3-fp4-mi355x-atom-agentic-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 } + - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [32, 40, 56], spec-decoding: mtp } dsr1-fp4-mi355x-sglang-disagg: image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519 From cd4354b07530e7fdc788a6aa851d87cb8690eb02 Mon Sep 17 00:00:00 2001 From: seungrokj <144636725+seungrokj@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:30:44 +0900 Subject: [PATCH 3/5] Update amd-master.yaml --- configs/amd-master.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 073885f3a..28895e2e6 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -671,7 +671,7 @@ kimik3-fp4-mi355x-atom-agentic-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: [32, 40, 56], spec-decoding: mtp } + - { 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 From 1adf729df813d8fec3b38626be48ecb4978f3bb6 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 25 Aug 2026 13:59:06 +0800 Subject: [PATCH 4/5] fix(agentx): halve the Kimi-K3 ATOM paged-KV CPU pool to clear the NCCL barrier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every rank pins its own LMCache pool, and at the measured sizes (200 GB/rank at concurrency 8 and 12, 192 GB/rank above that) the eight ranks finished pinning more than 600 s apart on the CI nodes. 600 s is exactly PyTorch's hardcoded kProcessGroupNCCLDefaultTimeout, so the ranks that reached the barrier at the end of allocate_kv_cache() first timed out waiting for rank 0 to publish the ncclUniqueId, and their ModelRunner processes died before the server ever served a request. Every concurrency point with KV offload hung; the three GPU-resident points finished normally. Halve the paged-KV half of the budget: dram-utilization 0.534 -> 0.268 (100 GB/rank) and 0.598 -> 0.343 (128 GB/rank, split 96 GB paged KV plus the unchanged 32 GB state tier). Aggregate pinned memory drops from 1600 to 800 GB at concurrency 8 and 12, and from 1792 to 1024 GB above that. The smaller pool costs little: the CPU tier's measured hit rate is 0.0% at concurrency 8 and 12 and 0.3-2.4% above that, and 96 GB/rank still holds roughly 53M tokens, about 530 full 100k-token contexts. 每个 rank 各自 pin 自己的 LMCache 池。按实测尺寸（并发 8/12 每 rank 200 GB， 以上每 rank 192 GB），八个 rank 在 CI 节点上完成 pin 的时间相差超过 600 秒， 而 600 秒正是 PyTorch 写死的 kProcessGroupNCCLDefaultTimeout。先到达 allocate_kv_cache() 末尾那个 barrier 的 rank 等不到 rank 0 发布 ncclUniqueId 便超时，ModelRunner 进程随之死亡，服务从未处理过任何请求。所有开启 KV offload 的并发点都挂起，三个全 GPU 驻留的点则正常跑完。 将预算中分页 KV 的部分减半：dram-utilization 0.534 -> 0.268（每 rank 100 GB）， 0.598 -> 0.343（每 rank 128 GB，拆为 96 GB 分页 KV 加上保持不变的 32 GB state 层）。固定内存总量在并发 8/12 从 1600 GB 降到 800 GB，以上从 1792 GB 降到 1024 GB。 缩小池子的代价很小：CPU 层实测命中率在并发 8/12 为 0.0%，更高并发为 0.3%-2.4%；每 rank 96 GB 仍可容纳约 5300 万 token，约合 530 个 10 万 token 的完整上下文。 Co-Authored-By: Claude Opus 5 (1M context) --- configs/amd-master.yaml | 18 ++++++++++++++---- perf-changelog.yaml | 3 ++- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index e9bb87765..d09b2f13c 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -650,11 +650,21 @@ kimik3-fp4-mi355x-vllm-agentic-mtp: # 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 +# 0.268 -> 803 GB aggregate -> 100 GB/rank, all paged KV (conc 8, 12). +# 0.343 -> 1028 GB aggregate -> 128 GB/rank, split 96 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. +# The paged-KV half of that budget is deliberately smaller than the recipe was +# measured with (200 and 192 GB/rank). Every rank pins its own pool, and at the +# measured sizes the eight ranks finished pinning more than 600 s apart, which +# is exactly PyTorch's hardcoded NCCL process-group timeout: the ranks that +# arrived first at the barrier ending allocate_kv_cache() timed out waiting for +# rank 0 and the ModelRunner processes died before the server ever served a +# request. Halving the paged-KV pool halves the pinning work. It costs little: +# the CPU tier's measured hit rate was 0.0% at concurrency 8 and 12 and 0.3-2.4% +# above that, and 96 GB/rank still holds ~53M tokens, roughly 530 full 100k-token +# contexts. The state tier is left at 32 GB/rank. kimik3-fp4-mi355x-atom-agentic-mtp: image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0821 model: moonshotai/Kimi-K3 @@ -665,11 +675,11 @@ kimik3-fp4-mi355x-atom-agentic-mtp: multinode: false scenarios: agentic-coding: - - dram-utilization: 0.534 + - dram-utilization: 0.268 search-space: - { 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 + - dram-utilization: 0.343 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 } diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 875979feb..15a5f3d5f 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6397,7 +6397,8 @@ - "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." + - "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.268 lands 100 GB per rank of paged KV at concurrency 8 and 12; 0.343 lands 128 GB per rank from concurrency 16 up, split 96 GB paged KV plus 32 GB for the state tier." + - "Size the paged-KV tier at half of what the recipe was measured with (200 and 192 GB per rank). Every rank pins its own pool, and at the measured sizes the eight ranks finished pinning more than 600 seconds apart, which is exactly PyTorch's hardcoded NCCL process-group timeout: ranks reaching the barrier at the end of allocate_kv_cache first timed out waiting for rank 0 and the ModelRunner processes died before the server served a request. The measured CPU-tier hit rate is 0.0% at concurrency 8 and 12 and 0.3-2.4% above that, so the smaller pool costs little; 96 GB per rank still holds roughly 53M tokens. The state tier stays at 32 GB per rank." - "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." From 6af3214df23187dfba707ae50a787a3d03fb210a Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Wed, 26 Aug 2026 02:03:10 +0800 Subject: [PATCH 5/5] fix(agentx): read the Kimi-K3 checkpoint without mmap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Set ATOM_DISABLE_MMAP=true for every concurrency point. The mapped loader materializes each tensor out of the page cache, so a rank whose pages are reclaimed pays to fault them back in, and the eight ranks finish loading minutes apart. The barrier ending allocate_kv_cache() is a NCCL collective and PyTorch's process-group timeout is a compile-time 600 s that ATOM does not override, so the ranks that arrive first die waiting for the last one. Two jobs of the last sweep were lost this way, both on the same node, with per-rank drain times spread from 330 s to 2531 s while the read phase itself stayed even. Reading each shard explicitly costs more transient host memory but makes load time depend on the filesystem rather than on page-cache residency. 为所有并发点设置 ATOM_DISABLE_MMAP=true。映射方式的加载器是从 page cache 中 逐个张量取出实体，因此页面被回收的 rank 需要再次缺页读回，八个 rank 完成加载 的时刻相差数分钟。allocate_kv_cache() 末尾的 barrier 是一次 NCCL 集合通信， 而 PyTorch 的进程组超时是编译期写死的 600 秒且 ATOM 不做覆盖，因此先到达的 rank 会在等待最后一个时死亡。上一轮 sweep 有两个 job 因此丢失，且都在同一台 节点上，各 rank 的 drain 耗时从 330 秒散布到 2531 秒，而读取阶段本身是齐的。 显式读取每个分片会占用更多瞬时主机内存，但使加载耗时只取决于文件系统，而不再 取决于 page cache 的驻留情况。 Co-Authored-By: Claude Opus 5 (1M context) --- .../single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh | 10 ++++++++++ perf-changelog.yaml | 1 + 2 files changed, 11 insertions(+) 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 034bbb41f..b35c89d0b 100644 --- a/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh +++ b/benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh @@ -261,6 +261,16 @@ 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 +# Read the 1.56 TB checkpoint into anonymous buffers instead of mapping it. +# The mapped path materializes each tensor out of the page cache, so a rank +# that loses its pages to reclaim pays for them again later, and the eight +# ranks then finish loading minutes apart. That skew is fatal here: the +# barrier ending allocate_kv_cache() is a NCCL collective, and PyTorch's +# process-group timeout is a compile-time 600 s that ATOM does not override, +# so the ranks that arrive first die waiting for the last one. Reading +# explicitly costs more transient host memory per shard but makes the load +# time depend on the filesystem alone rather than on page-cache residency. +export ATOM_DISABLE_MMAP=true # 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. diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 15a5f3d5f..c08147e46 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6399,6 +6399,7 @@ - "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.268 lands 100 GB per rank of paged KV at concurrency 8 and 12; 0.343 lands 128 GB per rank from concurrency 16 up, split 96 GB paged KV plus 32 GB for the state tier." - "Size the paged-KV tier at half of what the recipe was measured with (200 and 192 GB per rank). Every rank pins its own pool, and at the measured sizes the eight ranks finished pinning more than 600 seconds apart, which is exactly PyTorch's hardcoded NCCL process-group timeout: ranks reaching the barrier at the end of allocate_kv_cache first timed out waiting for rank 0 and the ModelRunner processes died before the server served a request. The measured CPU-tier hit rate is 0.0% at concurrency 8 and 12 and 0.3-2.4% above that, so the smaller pool costs little; 96 GB per rank still holds roughly 53M tokens. The state tier stays at 32 GB per rank." + - "Set ATOM_DISABLE_MMAP=true at every concurrency so the 1.56 TB checkpoint is read into anonymous buffers instead of mapped. The mapped path materializes each tensor out of the page cache, so a rank that loses its pages to reclaim pays for them again and the eight ranks finish loading minutes apart; the barrier ending allocate_kv_cache is a NCCL collective whose process-group timeout is a compile-time 600 seconds that ATOM does not override, so the ranks arriving first die waiting for the last. Reading explicitly costs more transient host memory per shard but makes load time depend on the filesystem rather than on page-cache residency." - "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."