Skip to content

feat(agentx): widen the Kimi-K3 ATOM AgentX DSpark draft to 4 tokens and bump the image / 将 Kimi-K3 ATOM AgentX 的 DSpark 草稿扩到 4 个 token 并升级镜像 #4085

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

Merged
edwingao28 merged 5 commits into from
Oct 15, 2026
Merged

feat(agentx): widen the Kimi-K3 ATOM AgentX DSpark draft to 4 tokens and bump the image / 将 Kimi-K3 ATOM AgentX 的 DSpark 草稿扩到 4 个 token 并升级镜像 #4085

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
bc2444e
feat(agentx): bump Kimi-K3 FP4 MI355X ATOM image and widen the DSpark…
gbyu-amd Oct 15, 2026
6224baf
chore(changelog): point the Kimi-K3 ATOM AgentX entry at its PR
gbyu-amd Oct 15, 2026
e713918
Merge remote-tracking branch 'origin/main' into amd/kimik3-atom-agent…
gbyu-amd Oct 15, 2026
4d7b7b8
perf(amd): Kimi-K3 MI355X ATOM AgentX on nightly_202610131505 with a …
seungrokj Oct 15, 2026
4506661
Merge branch 'main' into amd/kimik3-atom-agentx-mtp4
seungrokj Oct 15, 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
23 changes: 9 additions & 14 deletions ...ncex-e2e/benchmarks/single_node/srt-slurm-recipes/kimik3/atom/mi355x-fp4-mtp/agentic.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ base:
gpus_per_node: 8
frontend:
type: atomesh
container_image: rocm/atom-dev@sha256:f00a9dc588a380038fdad54cd57a956043cb316ddda54ebb1b62a24f8caa343b
container_image: rocm/atom-dev@sha256:b78300a6caf6c1fec3b1f27b4f7333326c9d34911bfc7ae229d13d2ebbc43884
enable_multiple_frontends: false
observability:
enabled: false
Expand Down Expand Up @@ -74,8 +74,6 @@ base:
AIPERF_REQUIRED_SERVER_METRIC_PREFIX: 'atom:'
AIPERF_APPLY_CHAT_TEMPLATE: 'true'

# One variant per point. Full graphs cover every verify batch up to 2x CONC
# times the verify window (1 + draft tokens).
override_tp8_c1:
roles:
agg:
Expand All @@ -99,7 +97,7 @@ override_tp8_dcp8_c14_lmcache:
agg:
args:
decode-context-parallel-size: 8
num-speculative-tokens: 3
num-speculative-tokens: 4
cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112]'
extra-kv-connectors: &lmcache_offload
- kv_connector: lmcache_offload
Expand Down Expand Up @@ -135,7 +133,7 @@ override_tp8_dcp8_c16_lmcache:
agg:
args:
decode-context-parallel-size: 8
num-speculative-tokens: 3
num-speculative-tokens: 4
cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128]'
extra-kv-connectors: *lmcache_offload
env:
Expand All @@ -155,15 +153,14 @@ override_tp8_dcp8_c48_lmcache:
agg:
args:
decode-context-parallel-size: 8
method: null
draft-model: null
num-speculative-tokens: null
num-speculative-tokens: 4
max-num-seqs: 96
cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96]'
Comment thread
gbyu-amd marked this conversation as resolved.
Resolved
extra-kv-connectors: *lmcache_offload
env:
<<: *lmcache_env
LMCACHE_MAX_LOCAL_CPU_SIZE: '128'
ATOM_ENABLE_REPLAYSSM: '1'
ATOM_PREFILL_DECODE_INTERVAL: '4'
ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000'
benchmark:
Expand All @@ -178,15 +175,14 @@ override_tp8_dcp8_c56_lmcache:
agg:
args:
decode-context-parallel-size: 8
method: null
draft-model: null
num-speculative-tokens: null
num-speculative-tokens: 4
max-num-seqs: 112
cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112]'
extra-kv-connectors: *lmcache_offload
env:
<<: *lmcache_env
LMCACHE_MAX_LOCAL_CPU_SIZE: '192'
ATOM_ENABLE_REPLAYSSM: '1'
AITER_REUSE_IDENTICAL_COMM_GROUPS: '1'
ATOM_PREFILL_DECODE_INTERVAL: '4'
ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000'
Expand All @@ -202,15 +198,14 @@ override_tp8_dcp8_c72_lmcache:
agg:
args:
decode-context-parallel-size: 8
method: null
draft-model: null
num-speculative-tokens: null
num-speculative-tokens: 4
max-num-seqs: 144
cudagraph-capture-sizes: '[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144]'
extra-kv-connectors: *lmcache_offload
env:
<<: *lmcache_env
LMCACHE_MAX_LOCAL_CPU_SIZE: '192'
ATOM_ENABLE_REPLAYSSM: '1'
AITER_REUSE_IDENTICAL_COMM_GROUPS: '1'
ATOM_PREFILL_DECODE_INTERVAL: '4'
ATOM_PREFILL_DELAYER_MAX_QUEUE_MS: '5000'
Expand Down
6 changes: 3 additions & 3 deletions inferencex-e2e/configs/amd-master.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,7 +696,7 @@ dsr1-fp8-mi355x-sglang-disagg-mtp:
# speculative decoding. Acceptance is pinned to the committed golden curve in
# infx/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 and 4,
# 3 draft tokens -> AL 3.00 at concurrency 14 and 16.
# 4 draft tokens -> AL 3.36 from concurrency 14 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.
Expand All @@ -706,15 +706,15 @@ dsr1-fp8-mi355x-sglang-disagg-mtp:
# 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, no draft model.
# throughput (48, 56, 72) same shape, with the 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 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.
kimik3-fp4-mi355x-atom-agentic-mtp:
image: rocm/atom-dev:nightly_202609251613
image: rocm/atom-dev:nightly_202610131505
model: moonshotai/Kimi-K3
model-prefix: kimik3
runner: cluster:mi355x-amds
Expand Down
13 changes: 13 additions & 0 deletions inferencex-e2e/perf-changelog.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9866,3 +9866,16 @@
description:
- "Upgrade agentx-harness from 1.0.9 to 1.0.10: pin SPAWN subagents to one client worker to prevent duplicate session state accumulating across workers. Metric collection, calculations, and exports are unchanged."
pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4119

- config-keys:
- kimik3-fp4-mi355x-atom-agentic-mtp
scenario-type:
- agentic-coding
description:
- "Move the MI355X Kimi-K3 FP4 ATOM AgentX submission from image rocm/atom-dev:nightly_202609251613 (ATOM 0.1.7.dev36+g68e0df5e7) to rocm/atom-dev:nightly_202610131505 (ATOM 0.1.8.dev44+g661ef7661), and pin the recipe frontend to that tag's digest sha256:b78300a6caf6c1fec3b1f27b4f7333326c9d34911bfc7ae229d13d2ebbc43884."
- "Keep the published concurrency set [1, 4, 14, 16, 48, 56, 72]; no points are added or removed."
- "Widen the DSpark draft from 3 to 4 tokens at concurrency 14 and 16, which moves the golden acceptance length from 3.00 to 3.36."
- "Serve concurrency 48, 56 and 72 with the same 4-token DSpark draft and ReplaySSM (ATOM_ENABLE_REPLAYSSM=1) instead of no draft model. Concurrency 1 and 4 keep the 7-token draft at AL 3.84. max-num-seqs, the CUDA-graph ladders, dcp-size, LMCache sizing and the prefill-delayer settings are unchanged."
- "The new image carries the Kimi-K3 prefill gated RMSNorm pipeline (ROCm/ATOM#2458) and moves the LMCache wheel that the in-process lmcache_offload connector imports from 0.5.6.dev98 to 0.5.6.dev139 (ROCm/ATOM#2438)."
- "No change to the Inferact/Kimi-K3-DSpark draft's precision: the checkpoint is BF16 throughout, and at ATOM 661ef7661 online_quant_config still excludes every draft linear (layers.*, context_proj), so its weights and activations stay BF16 at every drafted point, including the newly drafted concurrency 48, 56 and 72, and it shares the target's FP8 KV cache."
pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4085
Loading