Skip to content

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

Open
zejunchen-zejun wants to merge 7 commits into
mainfrom
amd/agentx_kimik3_atom_0821
Open

feat(agentx): retune Kimi-K3 FP4 MI355X ATOM DSpark recipe on _0821#2723
zejunchen-zejun wants to merge 7 commits into
mainfrom
amd/agentx_kimik3_atom_0821

Conversation

@zejunchen-zejun

@zejunchen-zejun zejunchen-zejun commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

Refresh the MI355X Kimi-K3 ATOM AgentX submission onto image kimi_k3_agentic_0821 and replace concurrency [1, 4, 8, 10] with [1, 2, 4, 8, 12, 16, 32, 40, 56], retuning every point.

Concurrency 1-4 is the latency floor: GPU-resident, no decode context parallelism, 7 draft tokens at golden AL 3.84, 8192-token prefill step. From concurrency 8 up decode is KV-bandwidth-bound over 100k+ token agentic contexts, so decode-context-parallel-size 8 shards the KV read across all 8 GPUs and the LMCache DRAM tier backs the paged KV. Concurrency 8 and 12 run 3 draft tokens at golden AL 3.00 with ReplaySSM and 96 in-GPU state checkpoint slots; from concurrency 16 up the KDA recurrent state moves to ATOM's CPU state-offload tier instead. Concurrency 32 and up serve without a draft model.

Both acceptance lengths come from the committed golden curve in golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml (7 -> 3.84, 3 -> 3.00). Evaluations drop the flag and use real acceptance.

LMCACHE_MAX_LOCAL_CPU_SIZE and OFFLOAD_STATE_CPU_SIZE are per rank, so the aggregate TOTAL_CPU_DRAM_GB is divided by TP as the agentic README requires: dram-utilization 0.534 gives 200 GB/rank of paged KV at concurrency 8 and 12, and 0.598 gives 224 GB/rank from concurrency 16 up, split 192 GB paged KV plus 32 GB for the state tier.


Note

Low Risk
Benchmark recipe, container image pin, and published concurrency matrix only—no application runtime or auth changes; risk is mis-tuned official numbers if the new knobs are wrong.

Overview
Refreshes the kimik3-fp4-mi355x-atom-agentic-mtp AgentX submission onto ATOM image kimi_k3_agentic_0821 and replaces concurrency [1, 4, 8, 10] with [1, 2, 4, 8, 12, 16, 32, 40, 56], with per-point retuning in kimik3_fp4_mi355x_atom_mtp.sh and matching matrix updates in configs/amd-master.yaml.

From concurrency 8 up, the recipe enables decode-context-parallel-size 8 (DCP_SIZE), LMCache DRAM for paged KV, and concurrency-specific batching, ReplaySSM, and state handling: low conc stays GPU-resident with 7 DSpark drafts at golden AL 3.84; 8–16 use 3 drafts at AL 3.00 (ReplaySSM + GPU checkpoints at 8/12, CPU state offload from 16); 32+ runs without speculative decoding. LMCache sizing now splits the per-rank CPU budget between paged KV and optional OFFLOAD_STATE_CPU_SIZE, uses LMCACHE_CHUNK_SIZE=1024 for DCP alignment, and halves the paged-KV pool vs the prior recipe to avoid multi-rank DRAM pinning exceeding the 600s NCCL barrier timeout. ATOM_DISABLE_MMAP=true avoids mmap-driven skewed checkpoint loads across ranks; ATOM_MLA_MAX_SPLIT_PER_BATCH is dropped.

perf-changelog.yaml documents the submission refresh for reviewers and downstream consumers.

Reviewed by Cursor Bugbot for commit 92166c9. Bugbot is set up for automated code reviews on this repo. Configure here.

zejunchen-zejun and others added 4 commits August 24, 2026 14:23
Refresh the MI355X Kimi-K3 ATOM AgentX submission onto image
kimi_k3_agentic_0821 and replace concurrency [1, 4, 8, 10] with
[1, 2, 4, 8, 12, 16, 32, 40, 56], retuning every point.

Concurrency 1-4 is the latency floor: GPU-resident, no decode context
parallelism, 7 draft tokens at golden AL 3.84, 8192-token prefill step.
From concurrency 8 up decode is KV-bandwidth-bound over 100k+ token agentic
contexts, so decode-context-parallel-size 8 shards the KV read across all
8 GPUs and the LMCache DRAM tier backs the paged KV. Concurrency 8 and 12
run 3 draft tokens at golden AL 3.00 with ReplaySSM and 96 in-GPU state
checkpoint slots; from concurrency 16 up the KDA recurrent state moves to
ATOM's CPU state-offload tier instead. Concurrency 32 and up serve without
a draft model.

Both acceptance lengths come from the committed golden curve in
golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml
(7 -> 3.84, 3 -> 3.00). Evaluations drop the flag and use real acceptance.

LMCACHE_MAX_LOCAL_CPU_SIZE and OFFLOAD_STATE_CPU_SIZE are per rank, so the
aggregate TOTAL_CPU_DRAM_GB is divided by TP as the agentic README requires:
dram-utilization 0.534 gives 200 GB/rank of paged KV at concurrency 8 and 12,
and 0.598 gives 224 GB/rank from concurrency 16 up, split 192 GB paged KV
plus 32 GB for the state tier.

将 MI355X 上 Kimi-K3 的 ATOM AgentX 提交刷新到 kimi_k3_agentic_0821 镜像,
并把并发点从 [1, 4, 8, 10] 换成 [1, 2, 4, 8, 12, 16, 32, 40, 56],逐点重调。

并发 1-4 为时延下界:全部驻留 GPU,不开 decode context parallel,草稿 7 token
对应 golden AL 3.84,prefill 步长 8192。并发 8 起,10 万 token 以上的 agentic
上下文使 decode 受 KV 带宽约束,因此用 decode-context-parallel-size 8 把 KV 读取
分摊到 8 张卡,并由 LMCache DRAM 层承接分页 KV。并发 8 与 12 使用草稿 3 token
(golden AL 3.00)、开启 ReplaySSM 并保留 96 个 GPU 内 state checkpoint 槽位;
并发 16 起,KDA 循环状态改由 ATOM 的 CPU state-offload 层承载。并发 32 及以上
不加载草稿模型。

两个接受长度均取自仓库内已提交的 golden 曲线(7 -> 3.84,3 -> 3.00);评测
(eval)不传该参数,使用真实接受率。

LMCACHE_MAX_LOCAL_CPU_SIZE 与 OFFLOAD_STATE_CPU_SIZE 均为每 rank 设置,因此按
agentic README 的要求将聚合预算 TOTAL_CPU_DRAM_GB 除以 TP:dram-utilization
0.534 在并发 8/12 下给到每 rank 200 GB 分页 KV;0.598 在并发 16 起给到每 rank
224 GB,拆为 192 GB 分页 KV + 32 GB state 层。

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

Thanks for the contribution! Please reach out to respective companies' CODEOWNER to fill in the latest PR_REVIEW_CHECKLIST.md before pinging core maintainer on Slack for review. In order for the signoff PR check bot to trigger, you must follow the PR_REVIEW_CHECKLIST.md template correctly, including the phrase As a PR reviewer and CODEOWNER, I have reviewed this and have.

For PR verification, add the full-sweep-fail-fast label (strongly recommended) to this PR — the benchmark sweep only runs on labeled PRs. Use full-sweep-enabled only if you need matrix jobs to keep running past a failure.

PR authors are responsible for ensuring that after merging, all GitHub Action jobs fully pass. A lot of the time, failures are just flakes and simply re-running the failed jobs will fix it. See GitHub's docs on re-running failed jobs


感谢你的贡献!请联系相应公司的 CODEOWNER 填写最新的 PR_REVIEW_CHECKLIST.md,然后再在 Slack 上联系核心维护者进行审阅。为了触发 signoff PR 检查机器人,你必须正确遵循 PR_REVIEW_CHECKLIST.md 模板,包括保留英文语句 As a PR reviewer and CODEOWNER, I have reviewed this and have

如需进行 PR 验证,请为此 PR 添加 full-sweep-fail-fast 标签(强烈推荐)— 基准测试 sweep 仅在带有标签的 PR 上运行。仅当需要矩阵任务在失败后继续运行时才使用 full-sweep-enabled

PR 作者有责任确保合并后所有 GitHub Action 任务完全通过。 很多时候失败只是偶发抖动(flake),重新运行失败的任务即可解决。参见 GitHub 关于重新运行失败任务的文档

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Beyond the inline finding, I also checked the new DCP_SIZE env var: it's documented in the "Required env vars" docstring but not added to the check_env_vars call, so it's optional in practice (defaults to 1 via ${DCP_SIZE:-1}) rather than actually enforced — not a correctness bug, just a docstring/enforcement mismatch worth noting alongside the confirmed spec-decoding label issue.

Extended reasoning...

Verified the confirmed inline finding directly against the diff: for CONC=32/40/56 the script sets NUM_SPEC_TOKENS=0 and SPEC_ARGS stays empty (no draft model), while configs/amd-master.yaml's conc-list [16, 32, 40, 56] search-space entry sets spec-decoding: mtp uniformly for the whole block, so results at 32/40/56 will be mislabeled. Also independently confirmed the ruled-out DCP_SIZE issue: line 20's docstring lists DCP_SIZE as required, but line 25's check_env_vars call omits it, so it is never enforced (it silently defaults via ${DCP_SIZE:-1} at the echo and ATOM_CMD construction) — a minor doc/enforcement gap, not a functional bug, so it doesn't change the disposition beyond the already-reported finding. Also spot-checked that perf-changelog.yaml's new PR #2716 entry was appended after the prior tail entry without disturbing existing bytes, consistent with the append-only convention.

Comment thread configs/amd-master.yaml
- { 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.598
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 }

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 The conc-list [16, 32, 40, 56] search-space entry sets spec-decoding: mtp for the whole block, but the script (kimik3_fp4_mi355x_atom_mtp.sh) sets NUM_SPEC_TOKENS=0 for CONC=32/40/56, meaning no draft model is loaded and SPEC_ARGS stays empty at those points -- only CONC=16 actually runs mtp spec decoding.

Extended reasoning...

SPEC_DECODING=mtp is forwarded as an env var and written verbatim into result metadata (benchmark_lib.sh:1222 -> process_agentic_result.py:224, spec_decoding field). Runs at concurrency 32, 40 and 56 will be tagged spec_decoding="mtp" in the published results even though the server ran dense decode with no draft model, mislabeling those data points in aggregation/dashboards as speculative-decoding throughput when they are not.

Verification: nit. The factual chain the candidate describes is real and reachable. configs/amd-master.yaml:674 puts the whole block under one arm: { tp: 8, dcp-size: 8, kv-offloading: dram, ..., conc-list: [16, 32, 40, 56], spec-decoding: mtp }. In kimik3_fp4_mi355x_atom_mtp.sh the CONC 32/40/56 cases set NUM_SPEC_TOKENS=0, and the new speculative block only populates SPEC_ARGS when `[ "$NUM_SPEC_TOKENS"

@github-actions

Copy link
Copy Markdown
Contributor

@github-actions

Copy link
Copy Markdown
Contributor

…CL barrier

Every rank pins its own LMCache pool, and at the measured sizes (200 GB/rank at
concurrency 8 and 12, 192 GB/rank above that) the eight ranks finished pinning
more than 600 s apart on the CI nodes. 600 s is exactly PyTorch's hardcoded
kProcessGroupNCCLDefaultTimeout, so the ranks that reached the barrier at the
end of allocate_kv_cache() first timed out waiting for rank 0 to publish the
ncclUniqueId, and their ModelRunner processes died before the server ever
served a request. Every concurrency point with KV offload hung; the three
GPU-resident points finished normally.

Halve the paged-KV half of the budget: dram-utilization 0.534 -> 0.268 (100
GB/rank) and 0.598 -> 0.343 (128 GB/rank, split 96 GB paged KV plus the
unchanged 32 GB state tier). Aggregate pinned memory drops from 1600 to 800 GB
at concurrency 8 and 12, and from 1792 to 1024 GB above that.

The smaller pool costs little: the CPU tier's measured hit rate is 0.0% at
concurrency 8 and 12 and 0.3-2.4% above that, and 96 GB/rank still holds
roughly 53M tokens, about 530 full 100k-token contexts.

每个 rank 各自 pin 自己的 LMCache 池。按实测尺寸(并发 8/12 每 rank 200 GB,
以上每 rank 192 GB),八个 rank 在 CI 节点上完成 pin 的时间相差超过 600 秒,
而 600 秒正是 PyTorch 写死的 kProcessGroupNCCLDefaultTimeout。先到达
allocate_kv_cache() 末尾那个 barrier 的 rank 等不到 rank 0 发布 ncclUniqueId
便超时,ModelRunner 进程随之死亡,服务从未处理过任何请求。所有开启 KV
offload 的并发点都挂起,三个全 GPU 驻留的点则正常跑完。

将预算中分页 KV 的部分减半:dram-utilization 0.534 -> 0.268(每 rank 100 GB),
0.598 -> 0.343(每 rank 128 GB,拆为 96 GB 分页 KV 加上保持不变的 32 GB state
层)。固定内存总量在并发 8/12 从 1600 GB 降到 800 GB,以上从 1792 GB 降到
1024 GB。

缩小池子的代价很小:CPU 层实测命中率在并发 8/12 为 0.0%,更高并发为
0.3%-2.4%;每 rank 96 GB 仍可容纳约 5300 万 token,约合 530 个 10 万 token
的完整上下文。

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

@github-actions

Copy link
Copy Markdown
Contributor

2 similar comments
@github-actions

Copy link
Copy Markdown
Contributor

@github-actions

Copy link
Copy Markdown
Contributor

Set ATOM_DISABLE_MMAP=true for every concurrency point. The mapped loader
materializes each tensor out of the page cache, so a rank whose pages are
reclaimed pays to fault them back in, and the eight ranks finish loading
minutes apart. 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. Two
jobs of the last sweep were lost this way, both on the same node, with
per-rank drain times spread from 330 s to 2531 s while the read phase itself
stayed even.

Reading each shard explicitly costs more transient host memory but makes load
time depend on the filesystem rather than on page-cache residency.

为所有并发点设置 ATOM_DISABLE_MMAP=true。映射方式的加载器是从 page cache 中
逐个张量取出实体,因此页面被回收的 rank 需要再次缺页读回,八个 rank 完成加载
的时刻相差数分钟。allocate_kv_cache() 末尾的 barrier 是一次 NCCL 集合通信,
而 PyTorch 的进程组超时是编译期写死的 600 秒且 ATOM 不做覆盖,因此先到达的
rank 会在等待最后一个时死亡。上一轮 sweep 有两个 job 因此丢失,且都在同一台
节点上,各 rank 的 drain 耗时从 330 秒散布到 2531 秒,而读取阶段本身是齐的。

显式读取每个分片会占用更多瞬时主机内存,但使加载耗时只取决于文件系统,而不再
取决于 page cache 的驻留情况。

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown
Contributor

@github-actions

Copy link
Copy Markdown
Contributor

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

agentx AgentX benchmarks, recipes, and infrastructure AMD full-sweep-enabled

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants