Skip to content

feat(agentx): retune Kimi-K3 FP4 MI355X ATOM DSpark recipe on _0821 (mirror of #2723) #2725

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
seungrokj wants to merge 8 commits into
base: main
Choose a base branch
from
+201 −103
Open

feat(agentx): retune Kimi-K3 FP4 MI355X ATOM DSpark recipe on _0821 (mirror of #2723) #2725

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c8cd524
feat(agentx): retune Kimi-K3 FP4 MI355X ATOM DSpark recipe on _0821
zejunchen-zejun Aug 24, 2026
d17b64f
Update amd-master.yaml
seungrokj Aug 24, 2026
cd4354b
Update amd-master.yaml
seungrokj Aug 25, 2026
b0134e1
Merge branch 'main' into amd/agentx_kimik3_atom_0821
seungrokj Aug 25, 2026
1adf729
fix(agentx): halve the Kimi-K3 ATOM paged-KV CPU pool to clear the NC…
zejunchen-zejun Aug 25, 2026
d11b64f
Move conc 4 into DCP arm; drop conc 16 from Kimi-K3 ATOM mtp sweep
seungrokj Aug 25, 2026
65ddbd7
Update amd-master.yaml
seungrokj Aug 25, 2026
8fcfc62
Merge branch 'main' into amd/agentx_kimik3_atom_0821_zejun
seungrokj Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
241 changes: 153 additions & 88 deletions benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand Down Expand Up @@ -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

@cquil11 cquil11 Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why dynamic SPED_DECODE_AL ?

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
Expand All @@ -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
Expand All @@ -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"

Expand All @@ -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"
Expand Down
47 changes: 32 additions & 15 deletions configs/amd-master.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -638,21 +638,35 @@ 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.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_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
Expand All @@ -661,10 +675,13 @@ kimik3-fp4-mi355x-atom-agentic-mtp:
multinode: false
scenarios:
agentic-coding:
- dram-utilization: 0.086
- dram-utilization: 0.268
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.343
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 }

@cursor cursor Bot Aug 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing concurrency 16 point

High Severity

Concurrency 16 was removed from the 0.343 DRAM search space. The PR targets the published list including 16, the config comments still describe state offload for concurrency 16 and up, and the launcher already has a dedicated CONC=16 recipe, so that retuned point never runs.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit d11b64f. Configure here.


dsr1-fp4-mi355x-sglang-disagg:
image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519
Expand Down
Loading
Loading