feat(agentx): Kimi-K3 ATOM mono decode at C1/C2 and on-demand state cache at C48+ / Kimi-K3 ATOM AgentX 在 C1/C2 启用 mono decode，C48 及以上启用按需 state 缓存

inferencex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml

@@ -6,7 +6,7 @@ base:
     gpus_per_node: 8
   frontend:
     type: atomesh
-    container_image: rocm/atom-dev@sha256:b78300a6caf6c1fec3b1f27b4f7333326c9d34911bfc7ae229d13d2ebbc43884
+    container_image: rocm/atom-dev@sha256:96a9614276518c2619e88800dec5f276bc76a95691331b4abb2210184e1431af
     enable_multiple_frontends: false
   observability:
     enabled: false
@@ -56,9 +56,13 @@ base:
         AITER_SITUV2_A4W4: '1'
         AITER_QUICK_REDUCE_QUANTIZATION: INT4
         AITER_FLYDSL_STAGE2_FP8: '1'
+        AITER_FLYDSL_USE_HERD: '1'
+        AITER_FLYDSL_HERD_MIN_M: '64'
         # Anchor-only state checkpointing.
         ATOM_STATE_CHECKPOINT_DEMAND: '0'
         ATOM_GDN_SSM_DTYPE: fp16
+        # ATOM defaults mono decode on; only concurrency 1 opts in.
+        ATOM_MONO_ENABLE: '0'
         # FlyDSL FP8 prefill attention; ATOM defaults it off.
         ATOM_USE_FLYDSL_FP8_PREFILL_ATTN: '1'
   benchmark:
@@ -78,7 +82,19 @@ override_tp8_c1:
   roles:
     agg:
       args:
-        cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]'
+        # max-num-seqs 1 fails the DSpark drafter's warmup, which runs two sequences.
+        max-num-seqs: 2
+        cudagraph-capture-sizes: '[1,2]'
+        # Mono loads in_proj, f_b_proj, o_proj and the two shared-expert
+        # projections as plain BF16 and refuses them quantized.
+        online_quant_config: '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*self_attn.[qkv]_conv1d*","*block_sparse_moe.experts*","*block_sparse_moe.routed_expert_*","*vision_tower*","*mm_projector*","layers.*.self_attn.fused_qkv_a_proj","layers.*.self_attn.q_b_proj","layers.*.self_attn.kv_b_proj","layers.*.self_attn.o_proj","layers.*.mlp.gate_up_proj","layers.*.mlp.down_proj","context_proj","*self_attn.in_proj*","*self_attn.f_b_proj*","*self_attn.o_proj*","*shared_experts.gate_up_proj*","*shared_experts.down_proj*"]}'
+      env:
+        ATOM_MONO_ENABLE: '1'
+        # The mono KDA kernel reads the SSM state pool as fp32 and refuses fp16.
+        ATOM_GDN_SSM_DTYPE: fp32
+        # Holds new prefill while any request decodes, so requests mostly decode
+        # one at a time.
+        ATOM_PREFILL_DECODE_INTERVAL: '1000000'
   benchmark:
     env:
       CONC: '1'
@@ -159,7 +175,11 @@ override_tp8_dcp8_c48_lmcache:
         extra-kv-connectors: *lmcache_offload
       env:
         <<: *lmcache_env
-        LMCACHE_MAX_LOCAL_CPU_SIZE: '128'
+        # On-demand state checkpoints at shared-prefix targets. The state cap
+        # is a limit inside the shared LMCache pool, not a second pool.
+        ATOM_STATE_CHECKPOINT_DEMAND: '1'
+        LMCACHE_MAX_LOCAL_CPU_SIZE: '224'
+        OFFLOAD_STATE_CPU_SIZE: '128'
         ATOM_ENABLE_REPLAYSSM: '1'
         ATOM_PREFILL_DECODE_INTERVAL: '4'
         ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000'
@@ -168,7 +188,7 @@ override_tp8_dcp8_c48_lmcache:
       CONC: '48'
       SPEC_DECODING: mtp
       KV_OFFLOADING: dram
-      TOTAL_CPU_DRAM_GB: '1028'
+      TOTAL_CPU_DRAM_GB: '1799'
 
 override_tp8_dcp8_c56_lmcache:
   roles:
@@ -181,7 +201,9 @@ override_tp8_dcp8_c56_lmcache:
         extra-kv-connectors: *lmcache_offload
       env:
         <<: *lmcache_env
-        LMCACHE_MAX_LOCAL_CPU_SIZE: '192'
+        ATOM_STATE_CHECKPOINT_DEMAND: '1'
+        LMCACHE_MAX_LOCAL_CPU_SIZE: '224'
+        OFFLOAD_STATE_CPU_SIZE: '128'
         ATOM_ENABLE_REPLAYSSM: '1'
         AITER_REUSE_IDENTICAL_COMM_GROUPS: '1'
         ATOM_PREFILL_DECODE_INTERVAL: '4'
@@ -191,7 +213,7 @@ override_tp8_dcp8_c56_lmcache:
       CONC: '56'
       SPEC_DECODING: mtp
       KV_OFFLOADING: dram
-      TOTAL_CPU_DRAM_GB: '1538'
+      TOTAL_CPU_DRAM_GB: '1799'
 
 override_tp8_dcp8_c72_lmcache:
   roles:
@@ -204,7 +226,9 @@ override_tp8_dcp8_c72_lmcache:
         extra-kv-connectors: *lmcache_offload
       env:
         <<: *lmcache_env
-        LMCACHE_MAX_LOCAL_CPU_SIZE: '192'
+        ATOM_STATE_CHECKPOINT_DEMAND: '1'
+        LMCACHE_MAX_LOCAL_CPU_SIZE: '224'
+        OFFLOAD_STATE_CPU_SIZE: '128'
         ATOM_ENABLE_REPLAYSSM: '1'
         AITER_REUSE_IDENTICAL_COMM_GROUPS: '1'
         ATOM_PREFILL_DECODE_INTERVAL: '4'
@@ -214,4 +238,4 @@ override_tp8_dcp8_c72_lmcache:
       CONC: '72'
       SPEC_DECODING: mtp
       KV_OFFLOADING: dram
-      TOTAL_CPU_DRAM_GB: '1538'
+      TOTAL_CPU_DRAM_GB: '1799'

inferencex-e2e/configs/amd-master.yaml

@@ -722,20 +722,21 @@ dsr1-fp8-mi355x-sglang-disagg-mtp:
 # 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, 4)        the latency floor, GPU-resident, dcp 1.
+#   interactive (1, 4)        the latency floor, GPU-resident, dcp 1; 1 runs
+#                             mono decode.
 #   mid         (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  (48, 56, 72)  same shape, with the draft model.
+#   throughput  (48, 56, 72)  same shape, with on-demand state checkpoints
+#                             and a 128 GB state cap inside the LMCache pool.
 # 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 14 through 48).
-#   0.513 -> 1538 GB aggregate -> 192 GB/rank (conc 56 and 72), where the
-#            wider in-flight window needs the deeper pool to keep the paged KV
-#            resident.
+#   0.343 -> 1028 GB aggregate -> 128 GB/rank (conc 14 and 16).
+#   0.6   -> 1799 GB aggregate -> 224 GB/rank (conc 48, 56 and 72), so the
+#            paged KV keeps ~96 GB/rank once the state tier reaches its cap.
 kimik3-fp4-mi355x-atom-agentic-mtp:
-  image: rocm/atom-dev:nightly_202610131505
+  image: rocm/atom-dev:nightly_202610281457
   model: moonshotai/Kimi-K3
   model-prefix: kimik3
   runner: cluster:mi355x-amds
@@ -748,10 +749,10 @@ kimik3-fp4-mi355x-atom-agentic-mtp:
     - dram-utilization: 0.343
       search-space:
       - { tp: 8, kv-offloading: none, conc-list: [1, 4], spec-decoding: mtp, srt-recipe: agentic.yaml }
-      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [14, 16, 48], spec-decoding: mtp, srt-recipe: agentic.yaml }
-    - dram-utilization: 0.513
+      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [14, 16], spec-decoding: mtp, srt-recipe: agentic.yaml }
+    - dram-utilization: 0.6
       search-space:
-      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [56, 72], spec-decoding: mtp, srt-recipe: agentic.yaml }
+      - { tp: 8, dcp-size: 8, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [48, 56, 72], spec-decoding: mtp, srt-recipe: agentic.yaml }
 
 # Kimi-K3 MXFP4 agentic-coding on MI355X via vLLM with DSpark speculative
 # decoding. Companion to kimik3-fp4-mi355x-atom-agentic-mtp: same checkpoint and

inferencex-e2e/perf-changelog.yaml

@@ -11499,3 +11499,15 @@
   pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4509
   scenario-type: [agentic-coding]
   append-only: true
+
+- config-keys:
+    - kimik3-fp4-mi355x-atom-agentic-mtp
+  scenario-type:
+    - agentic-coding
+  description:
+    - "Move the image from rocm/atom-dev:nightly_202610131505 (ATOM 661ef7661) to rocm/atom-dev:nightly_202610281457 (ATOM 06dbf6692), and pin the recipe frontend to that tag's digest sha256:96a9614276518c2619e88800dec5f276bc76a95691331b4abb2210184e1431af. The new image adds Kimi-K3 mono decode and the one-token f_b_proj fix it depends on (ROCm/ATOM#2626), and widens mono to 16-row steps with K2 on the MLA layers (ROCm/ATOM#2659). Its aiter keeps the tuned Kimi-K3 MXFP4 inline-sort MoE rows when ATOM passes the real row count (ROCm/aiter#7020)."
+    - "Serve concurrency 1 with ATOM mono decode (ATOM_MONO_ENABLE=1). Mono requires fp32 SSM state (ATOM_GDN_SSM_DTYPE=fp32) and BF16 in_proj, f_b_proj, o_proj and shared-expert projections, which are added to the online-quant exclude list. Every other point sets ATOM_MONO_ENABLE=0."
+    - "At concurrency 1, run one request at a time: the PrefillDelayer (ATOM_PREFILL_DECODE_INTERVAL=1000000) holds new prefill while a request decodes, with max-num-seqs 2 and CUDA-graph sizes [1, 2] (max-num-seqs 1 fails the DSpark drafter warmup)."
+    - "Set AITER_FLYDSL_USE_HERD=1 and AITER_FLYDSL_HERD_MIN_M=64 at every concurrency."
+    - "At concurrency 48, 56 and 72, enable on-demand state checkpoints (ATOM_STATE_CHECKPOINT_DEMAND=1), raise the per-rank LMCache pool to 224 GB, and cap retained state at 128 GB inside that pool (OFFLOAD_STATE_CPU_SIZE=128). dram-utilization 0.6 gives TOTAL_CPU_DRAM_GB 1799."
+  pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4307