feat(agentx): retune Kimi-K3 FP4 MI355X ATOM DSpark recipe on _0821

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,37 @@ 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
+# 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.
 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 +303,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"

configs/amd-master.yaml

@@ -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
@@ -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: [16, 32, 40, 56], spec-decoding: mtp }
 
 dsr1-fp4-mi355x-sglang-disagg:
   image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-20260519