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

benchmarks/single_node/agentic/kimik3_fp4_mi355x_atom_mtp.sh

@@ -8,9 +8,9 @@ set -x
 # apply_k3_container_patches.sh is not sourced: it targets a patched vLLM container.
 #
 # Serving bands, one fresh server per concurrency point:
-#   interactive (1, 2, 4)         DCP1, DSpark 7, no LMCache
-#   mid         (8, 12, 14, 16)   DCP8, DSpark 3, LMCache 128 GB/rank
-#   throughput  (32, 40, 56, 64)  DCP8, no draft model, LMCache 128 or 192 GB/rank
+#   interactive (1, 2, 4)                 DCP1, DSpark 7, no LMCache
+#   mid         (8, 16)                   DCP8, DSpark 3, LMCache 128 GB/rank
+#   throughput  (32, 40, 48, 56, 64, 72)  DCP8, no draft, LMCache 128 or 192 GB/rank
 #
 # Required env vars:
 #   MODEL, MODEL_PATH, TP, DCP_SIZE, CONC, KV_OFFLOADING, KV_OFFLOAD_BACKEND,
@@ -95,90 +95,31 @@ trap 'exit 143' TERM
 # so DCP8 (from dcp-size in configs/amd-master.yaml), LMCache DRAM tier, and
 # ReplaySSM rebuilding KDA state from the checkpoint ring. Throughput band:
 # past the knee the draft forward no longer pays for itself, so no draft.
-# CUDAGRAPH_MAX_NUM_SEQS defaults to 2 * CONC; only concurrency 14 pins it
-# because it runs the concurrency 16 server verbatim.
+MAX_NUM_BATCHED_TOKENS=8192
+GPU_MEM_UTIL=0.90
 AITER_REUSE_IDENTICAL_COMM_GROUPS=0
-CUDAGRAPH_MAX_NUM_SEQS=""
 case "$CONC" in
     1|2|4)
         MAX_NUM_SEQS=32
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.88
         ATOM_ENABLE_REPLAYSSM=0
         NUM_SPEC_TOKENS=7
         SPEC_DECODE_AL=3.84
         ;;
-    8)
+    8|16)
         MAX_NUM_SEQS=32
-        MAX_NUM_BATCHED_TOKENS=4096
-        GPU_MEM_UTIL=0.88
         ATOM_ENABLE_REPLAYSSM=1
         NUM_SPEC_TOKENS=3
         SPEC_DECODE_AL=3.00
         ;;
-    12)
-        MAX_NUM_SEQS=24
-        MAX_NUM_BATCHED_TOKENS=4096
-        GPU_MEM_UTIL=0.88
-        ATOM_ENABLE_REPLAYSSM=1
-        NUM_SPEC_TOKENS=3
-        SPEC_DECODE_AL=3.00
-        ;;
-    14)
-        # The concurrency 16 server verbatim; only the client concurrency is 14.
-        # Deriving the graph width from 2 * CONC (28) dies during graph warmup.
-        MAX_NUM_SEQS=32
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.86
-        ATOM_ENABLE_REPLAYSSM=1
-        NUM_SPEC_TOKENS=3
-        SPEC_DECODE_AL=3.00
-        CUDAGRAPH_MAX_NUM_SEQS=32
-        ;;
-    16)
-        MAX_NUM_SEQS=32
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.86
-        ATOM_ENABLE_REPLAYSSM=1
-        NUM_SPEC_TOKENS=3
-        SPEC_DECODE_AL=3.00
-        ;;
-    32)
-        MAX_NUM_SEQS=64
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.86
+    32|40|48|56|64|72)
+        MAX_NUM_SEQS=$((2 * CONC))
         ATOM_ENABLE_REPLAYSSM=0
         NUM_SPEC_TOKENS=0
         SPEC_DECODE_AL=0
-        ;;
-    40)
-        MAX_NUM_SEQS=80
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.86
-        ATOM_ENABLE_REPLAYSSM=0
-        NUM_SPEC_TOKENS=0
-        SPEC_DECODE_AL=0
-        ;;
-    # The two widest points are the only ones that reuse identical AITER
-    # communicator groups, and the only ones given the 192 GB/rank LMCache
-    # budget rather than 128.
-    56)
-        MAX_NUM_SEQS=112
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.86
-        ATOM_ENABLE_REPLAYSSM=0
-        NUM_SPEC_TOKENS=0
-        SPEC_DECODE_AL=0
-        AITER_REUSE_IDENTICAL_COMM_GROUPS=1
-        ;;
-    64)
-        MAX_NUM_SEQS=128
-        MAX_NUM_BATCHED_TOKENS=8192
-        GPU_MEM_UTIL=0.86
-        ATOM_ENABLE_REPLAYSSM=0
-        NUM_SPEC_TOKENS=0
-        SPEC_DECODE_AL=0
-        AITER_REUSE_IDENTICAL_COMM_GROUPS=1
+        # Only the three widest points reuse identical AITER communicator groups.
+        if [ "$CONC" -ge 56 ]; then
+            AITER_REUSE_IDENTICAL_COMM_GROUPS=1
+        fi
         ;;
     *)
         echo "Unsupported CONC=$CONC" >&2
@@ -299,6 +240,19 @@ ATOM_CMD=(
     --host 0.0.0.0
     --server-port "$PORT"
     --trust-remote-code
+    # ATOM reads the tool-call wire format off the chat template at startup
+    # ("auto", the default), and this image already resolves it correctly --
+    # the server log says `Tool-call format: kimi_k3 (from the chat template)`.
+    # Pinning it is for the day the template stops rendering a tools payload:
+    # auto then falls back to delivering tool calls as plain text, which the
+    # vendor verifier reads as a model that cannot call tools. An explicit name
+    # fails closed instead, and the log line changes to `(from
+    # --tool-call-parser)` so which path ran is visible.
+    #
+    # This is ATOM's own flag. It has no --enable-auto-tool-choice, and
+    # --reasoning-parser belongs to the mesh router, not this server; argparse
+    # here is strict, so either one would fail the launch.
+    --tool-call-parser kimi_k3
     --tensor-parallel-size "$TP"
     --decode-context-parallel-size "${DCP_SIZE}"
     --kv_cache_dtype fp8

configs/amd-master.yaml

@@ -596,27 +596,25 @@ kimik3-fp4-mi355x-vllm-agentic-mtp:
 # 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 at concurrency 8 through 16.
+# 3 draft tokens -> AL 3.00 at concurrency 8 and 16.
 # 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.
 # Three serving bands, tracking the ATOM recipe:
-#   interactive (1, 2, 4)        the latency floor, GPU-resident, dcp 1.
-#   mid         (8, 12, 14, 16)  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.
-#   throughput  (32, 40, 56, 64) same shape, no draft model.
+#   interactive (1, 2, 4)                 the latency floor, GPU-resident, dcp 1.
+#   mid         (8, 16)                   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.
+#   throughput  (32, 40, 48, 56, 64, 72)  same shape, no draft model.
 # Two utilization blocks because the per-rank LMCache pool differs, and the
 # launcher divides the aggregate budget by TP:
-#   0.343 -> 1028 GB aggregate -> 128 GB/rank (conc 8 through 40).
-#   0.513 -> 1538 GB aggregate -> 192 GB/rank (conc 56 and 64), where the wider
-#            in-flight window needs the deeper pool to keep the paged KV
+#   0.343 -> 1028 GB aggregate -> 128 GB/rank (conc 8 through 48).
+#   0.513 -> 1538 GB aggregate -> 192 GB/rank (conc 56 through 72), where the
+#            wider in-flight window needs the deeper pool to keep the paged KV
 #            resident.
-# Concurrency 14 is deliberate: it runs the concurrency 16 server verbatim,
-# including the pinned CUDA-graph width, and only moves the client concurrency.
 kimik3-fp4-mi355x-atom-agentic-mtp:
-  image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0907
+  image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0911
   model: moonshotai/Kimi-K3
   model-prefix: kimik3
   runner: cluster:mi355x-amds
@@ -628,10 +626,10 @@ kimik3-fp4-mi355x-atom-agentic-mtp:
     - dram-utilization: 0.343
       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, 14, 16, 32, 40], spec-decoding: mtp }
+      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [8, 16, 32, 40, 48], spec-decoding: mtp }
     - dram-utilization: 0.513
       search-space:
-      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [56, 64], spec-decoding: mtp }
+      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [56, 64, 72], spec-decoding: mtp }
 
 minimaxm3-fp4-mi355x-atom-agentic-mtp:
   image: rocm/atom-dev:nightly_202609140645-lirzhang-triton-build

perf-changelog.yaml

@@ -7924,6 +7924,21 @@
     - "Switch HiCache defaults to --hicache-io-backend kernel and --hicache-mem-layout page_first (from direct / page_first_direct). Ratio 1.5 and write_through are unchanged."
   pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/3118
 
+- config-keys:
+    - kimik3-fp4-mi355x-atom-agentic-mtp
+  scenario-type:
+    - agentic-coding
+  description:
+    - "Move the MI355X Kimi-K3 FP4 ATOM AgentX submission to image rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_kimi_k3_agentic_0911, tracking recipes/Agentic-Kimi-K3.md at ROCm/ATOM#2131 as retuned by ROCm/ATOM#2208."
+    - "Publish the concurrency set [1, 2, 4, 8, 16, 32, 40, 48, 56, 64, 72], replacing [1, 2, 4, 8, 12, 14, 16, 32, 40, 56, 64]."
+    - "Set gpu-memory-utilization to 0.90 and max-num-batched-tokens to 8192 on every point."
+    - "Set max-num-seqs to 32 through concurrency 16 and to 2 x concurrency from concurrency 32 up, giving 64, 80, 96, 112, 128, and 144 at concurrency 32, 40, 48, 56, 64, and 72."
+    - "Set AITER_REUSE_IDENTICAL_COMM_GROUPS=1 at concurrency 56, 64, and 72, and 0 from concurrency 1 through 48."
+    - "Drop the concurrency-14 CUDA-graph pin with the point. It set CUDAGRAPH_MAX_NUM_SEQS=32 in the case block, which the capture-width derivation overwrote with 2 x concurrency two statements later, so it never reached the server. Capture behaviour is unchanged: --level 3 with --cudagraph-mode FULL over [2 .. 2 x concurrency x (1 + draft tokens)]."
+    - "Carry the two paged-KV dram-utilization blocks onto the new concurrency set: 0.343 lands 128 GB per rank for concurrency 8 through 48, and 0.513 lands 192 GB per rank for concurrency 56 through 72."
+    - "Draft depth and synthetic acceptance (7 draft tokens at AL 3.84 for concurrency 1 through 4, 3 at AL 3.00 for concurrency 8 and 16, no draft from 32 up), ReplaySSM placement, dcp-size 8 from concurrency 8 up, the FP8 KV cache, the 128-token page, --state-checkpoint-interval-tokens -1 with anchor-only state checkpointing, ptpc_fp8 online quantization, LMCACHE_CHUNK_SIZE 1024, the NUMA pinning, the draft pre-staging retry, --served-model-name $MODEL, and the AgentX client are unchanged from #2852."
+  pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/3207
+
 - config-keys:
     - dsv4-fp4-mi355x-sglang-agentic-mtp
   scenario-type: