Skip to content

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 缓存 - #4307

Open
gbyu-amd wants to merge 15 commits into
mainfrom
amd/kimik3-atom-agentx-mono-statecache
Open

gbyu-amd wants to merge 15 commits into
mainfrom
amd/kimik3-atom-agentx-mono-statecache

Conversation

@gbyu-amd

@gbyu-amd gbyu-amd commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

kimik3-fp4-mi355x-atom-agentic-mtp moves to a new image and changes in two places:

  • Image rocm/atom-dev:nightly_202610131505 (ATOM 661ef7661) -> rocm/atom-dev:nightly_202610281457 (ATOM 06dbf6692, aiter with [HIP] [MoE] MXFP4 inline-sort: accept num_local_tokens, skip ids outside [0, E) ROCm/aiter#7020) in the master config. The recipe frontend is pinned to the matching digest sha256:96a9614276518c2619e88800dec5f276bc76a95691331b4abb2210184e1431af. The new image has Kimi-K3 mono decode and the one-token f_b_proj fix it depends on (fix(kimi-k3): give the bf16 f_b_proj a real copy of f_a at one token ROCm/ATOM#2626), and 16-row mono steps with K2 on the MLA layers ([Kimi-K3] mono: serve 16-row and drafter-tapped steps, optional MLA K2, scheduler cap ROCm/ATOM#2659).
  • Concurrency 1: ATOM mono decode (ATOM_MONO_ENABLE=1). Mono needs ATOM_GDN_SSM_DTYPE=fp32 and BF16 in_proj, f_b_proj, o_proj and shared-expert projections, so those patterns are added to the online_quant_config exclude list for this point. The draft-layer excludes are kept. ATOM now defaults mono on, so the base sets ATOM_MONO_ENABLE=0 for every other point. On the target image, mono refuses those points anyway (fp16 state, or a KV connector), so the explicit 0 changes nothing there.
  • Concurrency 48, 56 and 72: on-demand state checkpoints and a larger CPU tier, using the values from the ATOM PD recipe (feat(kimi-k3): add on-demand PD recipe with 128 GiB state cache ROCm/ATOM#2593). This sets ATOM_STATE_CHECKPOINT_DEMAND=1, raises the per-rank LMCACHE_MAX_LOCAL_CPU_SIZE from 128/192 to 224, and sets OFFLOAD_STATE_CPU_SIZE=128. The state size is a cap inside the shared LMCache pool, not a second pool. C48 moves out of the 0.343 block, and all three use dram-utilization: 0.6 (TOTAL_CPU_DRAM_GB 1799, about 60% of mi355x-amds DRAM).
  • Concurrency 1: one request at a time. ATOM_PREFILL_DECODE_INTERVAL=1000000 (PrefillDelayer holds new prefill while a request decodes), max-num-seqs: 2, cudagraph-capture-sizes: '[1,2]'. max-num-seqs 1 fails at startup because the DSpark drafter's warmup runs two sequences. Queued requests wait, so TTFT rises.
  • Every concurrency: AITER_FLYDSL_USE_HERD=1 and AITER_FLYDSL_HERD_MIN_M=64.
  • Apart from that, C4, C14 and C16 serve exactly as before. The published set stays [1, 4, 14, 16, 48, 56, 72].

Local MI355X AgentX results (semianalysis_cc_traces_weka_062126, 3600 s per point), all against the same-image baseline:

Change Point Output tok/s Other
mono C1 +10.9% ITL p50 6.27 -> 5.01 ms; ITL p90 -7% (inside C4's run-to-run spread)
DEMAND=1 + 224 / 128 C48 +3.2% TTFT p50 -2.5%
DEMAND=1 + 224 / 128 C56 +6.1% TTFT p50 -7.9%
DEMAND=1 + 224 / 128 C72 +11.8% TTFT p50 -28.7%
Supporting evidence
  • At C48/56/72, the state-tier cap goes from 1087 to 4351 entries per rank. Offload loads that failed and fell back to prefill drop from 10/11/29 to 0, and tokens served from the CPU tier rise 14-16%. The HBM KV pool is unchanged (within 0.7%).
  • These were single runs. C48 and C56 are inside the known run-to-run spread at those concurrencies; C72 is clearly outside it. The A/B moved the three knobs together, so their individual contributions are not separated.
  • At C14/C16 the CPU tier is almost unused (4 and 24 state loads per hour, about 0.5M tokens), so they keep their current sizing.
  • GSM8K 3-shot on the C1 mono config, with no forced acceptance and fix(kimi-k3): give the bf16 f_b_proj a real copy of f_a at one token ROCm/ATOM#2626 applied: flexible-extract 0.9530 / 0.9538 / 0.9553 at client concurrency 1 / 2 / 64 (ATOM CI baseline 0.95). Without fix(kimi-k3): give the bf16 f_b_proj a real copy of f_a at one token ROCm/ATOM#2626, the engine fails on the first one-token prefill chunk. With prompt-end anchors, that is any prompt of length 128k+1.
中文

kimik3-fp4-mi355x-atom-agentic-mtp 更换镜像,另有两处改动:

  • 镜像:主配置由 rocm/atom-dev:nightly_202610131505(ATOM 661ef7661)改为 rocm/atom-dev:nightly_202610281457(ATOM 06dbf6692,aiter 含 [HIP] [MoE] MXFP4 inline-sort: accept num_local_tokens, skip ids outside [0, E) ROCm/aiter#7020),配方 frontend 固定为对应的 digest sha256:96a9614276518c2619e88800dec5f276bc76a95691331b4abb2210184e1431af。新镜像包含 Kimi-K3 mono decode,以及 mono 所依赖的单 token f_b_proj 修复(fix(kimi-k3): give the bf16 f_b_proj a real copy of f_a at one token ROCm/ATOM#2626),以及支持 16 行、MLA 层也走 K2 的 mono decode([Kimi-K3] mono: serve 16-row and drafter-tapped steps, optional MLA K2, scheduler cap ROCm/ATOM#2659)。
  • 并发 1:ATOM mono decode(ATOM_MONO_ENABLE=1)。mono 要求 ATOM_GDN_SSM_DTYPE=fp32,并要求 in_proj、f_b_proj、o_proj 和共享专家投影层保持 BF16,因此该点的 online_quant_config 排除列表加上了这些规则,草稿层的排除规则保持不变。ATOM 现在默认开启 mono,所以 base 中对其他所有点设置 ATOM_MONO_ENABLE=0。在目标镜像上,这些点本来就会被 mono 拒绝(fp16 状态或存在 KV connector),因此显式写 0 不会改变它们的行为。
  • 并发 48、56、72:按需 state checkpoint,并扩大 CPU 缓存层,取值来自 ATOM PD 配方(feat(kimi-k3): add on-demand PD recipe with 128 GiB state cache ROCm/ATOM#2593)。具体是设置 ATOM_STATE_CHECKPOINT_DEMAND=1,把每个 rank 的 LMCACHE_MAX_LOCAL_CPU_SIZE 从 128/192 提高到 224,并设置 OFFLOAD_STATE_CPU_SIZE=128。state 大小是共享 LMCache 池内的上限,不是另一个独立的池。C48 移出 0.343 这一组,三个点都使用 dram-utilization: 0.6(TOTAL_CPU_DRAM_GB 1799,约占 mi355x-amds 内存的 60%)。
  • 并发 1:一次只服务一个请求。 ATOM_PREFILL_DECODE_INTERVAL=1000000(有请求在 decode 时,PrefillDelayer 暂缓新的 prefill),max-num-seqs: 2,cudagraph-capture-sizes: '[1,2]'。max-num-seqs 1 会因 DSpark 草稿模型预热需要两条序列而启动失败。排队的请求需要等待,因此 TTFT 会上升。
  • 所有并发点:设置 AITER_FLYDSL_USE_HERD=1 和 AITER_FLYDSL_HERD_MIN_M=64。
  • 除此之外,C4、C14、C16 的服务配置与之前完全相同。发布的并发集合保持 [1, 4, 14, 16, 48, 56, 72] 不变。

本地 MI355X AgentX 结果(semianalysis_cc_traces_weka_062126,每个点 3600 秒),均与同一镜像的基线对比:mono 使 C1 输出吞吐 +10.9%,ITL p50 从 6.27 降到 5.01 ms;DEMAND=1 + 224/128 使 C48/56/72 输出吞吐分别 +3.2% / +6.1% / +11.8%,C72 的 TTFT p50 降低 28.7%。三档的 offload 读回失败(退回 prefill)从 10/11/29 次降为 0,HBM KV 池不变。以上均为单次运行:C48/C56 的提升在已知波动范围内,C72 明显超出;三个参数是一起调整的,各自的贡献没有拆开。C14/C16 的 CPU 缓存层几乎没有被使用,因此保持现有配置。在应用 ROCm/ATOM#2626 的 C1 mono 配置上,GSM8K 3-shot(不强制接受率)在客户端并发 1/2/64 下 flexible-extract 分别为 0.9530 / 0.9538 / 0.9553(ATOM CI 基线 0.95)。

AI model disclosure

Prepared with Claude Code using claude-opus-5-5[1m] (local sweeps and analysis, recipe and master-config edits, changelog entry, PR description). No other model and no delegated agent contributed.

…ache at C48+

- Concurrency 1 and the new concurrency 2 run ATOM mono decode, with the
  fp32 SSM state and BF16 projection excludes it requires; every other
  point pins ATOM_MONO_ENABLE=0.
- Concurrency 48, 56 and 72 enable on-demand state checkpoints, a 224 GB
  per-rank LMCache pool and a 128 GB state cap inside it
  (dram-utilization 0.6 -> TOTAL_CPU_DRAM_GB 1799).

Image unchanged: mono at C1/C2 needs ROCm/ATOM#2626, which is not yet in
a nightly.

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

Copy link
Copy Markdown
Contributor

Thanks for the contribution!

  • Review: If this PR changes files owned by someone other than a repository admin or @SemiAnalysisAI/core, ask one eligible CODEOWNER to complete the latest PR_REVIEW_CHECKLIST.md before contacting a core maintainer on Slack. Follow the template exactly, including As a PR reviewer and CODEOWNER, I have reviewed this and have, so sign-off verification triggers.
  • PR verification: Sweeps run only when the PR appends an inferencex-e2e/perf-changelog.yaml entry and carries exactly one primary label: full-sweep-fail-fast (strongly recommended; canary plus per-matrix fail-fast), full-sweep-enabled (canary; matrix jobs continue after a failure), or non-canary-full-sweep-enabled (no canary or fail-fast). The modifiers all-evals, evals-only, and agentx-fast require a primary label. On fork PRs, a maintainer applies the label. See sweep labels and reuse.
  • After merging: PR authors must ensure all GitHub Actions jobs pass. Transient failures often pass on rerun; see how to rerun failed jobs.
中文

感谢你的贡献!

  • **审阅:**如果 PR 修改的文件归属于仓库管理员及 @SemiAnalysisAI/core 之外的 CODEOWNER,请先联系一位有资格的 CODEOWNER 填写最新的 PR_REVIEW_CHECKLIST.md,再通过 Slack 联系核心维护者。必须严格遵循模板,并保留 As a PR reviewer and CODEOWNER, I have reviewed this and have,才能触发签核验证。
  • **PR 验证:**只有当 PR 在 inferencex-e2e/perf-changelog.yaml 末尾追加了条目,并且恰好带有一个主标签时,才会运行扫描:full-sweep-fail-fast(强烈推荐;canary 加逐矩阵 fail-fast)、full-sweep-enabled(有 canary;矩阵任务在失败后继续运行)或 non-canary-full-sweep-enabled(无 canary,也无 fail-fast)。修饰标签 all-evals、evals-only 和 agentx-fast 必须与主标签一起使用。fork PR 的标签由维护者添加。参见扫描标签与复用。
  • **合并后:**PR 作者必须确保所有 GitHub Actions 任务通过。临时性失败通常可以通过重新运行恢复;参见重新运行失败任务的说明。

gbyu-amd and others added 4 commits October 19, 2026 04:38
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Only this PR's entry takes the #4307 link.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…x-mono-statecache

# Conflicts:
#	inferencex-e2e/perf-changelog.yaml
ATOM 53e4fb524 carries Kimi-K3 mono decode and the one-token f_b_proj
fix it depends on (ROCm/ATOM#2626). Recipe frontend pinned to the tag's
digest.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@gbyu-amd
gbyu-amd marked this pull request as ready for review October 19, 2026 15:47
@gbyu-amd
gbyu-amd requested a review from a team October 19, 2026 15:47
@gbyu-amd gbyu-amd added AMD agentx AgentX benchmarks, recipes, and infrastructure full-sweep-enabled Full sweep with canary gate; matrix jobs run to completion despite failures labels Oct 19, 2026
@claude

claude Bot commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Claude finished @gbyu-amd's task in 1m 16s —— View job


Review

LGTM - no blocking issues found.

What I checked:

  • Recipe selection for the new C2 point: override_tp8_c2 matches only on CONC=2. No other variant is bound to that concurrency, so select_recipe finds exactly one variant.
  • DRAM budget: The available DRAM is capped at MAX_AGENTIC_AVAILABLE_CPU_DRAM_MIB (2,861,022 MiB, about 3000 GB). With dram-utilization: 0.6 the matrix generates TOTAL_CPU_DRAM_GB=1799, which matches what the recipe pins for C48/C56/C72. The 0.343 block still generates 1028 for C14/C16, so the recipe/matrix equality check in validate_recipe passes. 1799 / TP8 gives 224 GB per rank, which matches LMCACHE_MAX_LOCAL_CPU_SIZE: '224'.
  • CUDA-graph ladder: The C2 ladder (1-32) follows the same 2×conc×(1+7) pattern as C1 (16) and C4 (64).
  • Images: The new master-config image rocm/atom-dev:nightly_202610191501 is on a public registry.
  • Changelog: The perf-changelog.yaml entry is appended at the end of the file and covers kimik3-fp4-mi355x-atom-agentic-mtp.

I could not run infx.matrix.changed because the uv run command needs approval in this environment. That check doesn't matter for this PR, since it already includes a changelog entry.

@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.

I reviewed this PR and didn't find any bugs. amd-master.yaml has designated CODEOWNERS, and the change retunes published benchmark points from single-run A/B results, so a look from one of them would still be worthwhile.

What was reviewed:

  • The DRAM budget for the 0.6 block. The runner cap of 2,861,022 MiB x 0.6 gives 1799 GB, so the recipe's TOTAL_CPU_DRAM_GB '1799' matches what the matrix produces. 1799/8 also covers the 224 GB/rank LMCache pool, and 0.343 still gives 1028 for C14/C16.
  • Recipe selection: the new override_tp8_c2 is the only variant with CONC '2', so select_recipe should find exactly one match. The &mono_env / &mono_quant aliases merge the same way the existing lmcache overrides already do.
  • Draft precision under mono: the &mono_quant exclude list keeps the unprefixed layers.* and context_proj draft exclusions and only adds more BF16 patterns. The perf-changelog entry is appended only at the end of the file and doesn't edit any earlier entry.
Extended reasoning...

This change updates one benchmark (kimik3-fp4-mi355x-atom-agentic-mtp) in three files. It bumps the image and its matching recipe digest, adds a C2 point, turns on mono decode for C1/C2 with a wider online-quant exclude list, retunes the LMCache/state settings for C48/56/72 under a new 0.6 DRAM block, and appends a perf-changelog entry. It touches no security-sensitive surface. The values I checked are consistent: the DRAM math, the unique recipe match per concurrency, and the draft-layer exclusions. I deferred instead of approving because amd-master.yaml has specific CODEOWNERS, and the new published points rest on single-run performance claims that the owners should confirm.

@github-actions

github-actions Bot commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

…x-mono-statecache

# Conflicts:
#	inferencex-e2e/perf-changelog.yaml
@gbyu-amd
gbyu-amd requested a review from zijiecode as a code owner October 20, 2026 13:55
@gbyu-amd

Copy link
Copy Markdown
Collaborator Author

@seungrokj @edwingao28 hi, could you kindly review this?

@gbyu-amd

Copy link
Copy Markdown
Collaborator Author

let's hold this pr first, will update the recipe later and re-sweep. Convert to draft now.

@gbyu-amd
gbyu-amd marked this pull request as draft October 21, 2026 02:29
gbyu-amd and others added 3 commits October 20, 2026 22:02
Concurrency 1 keeps mono decode; the published set returns to
[1, 4, 14, 16, 48, 56, 72].

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…x-mono-statecache

# Conflicts:
#	inferencex-e2e/perf-changelog.yaml
AITER_FLYDSL_USE_HERD=1 and AITER_FLYDSL_HERD_MIN_M=64 in the recipe
base env.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@seungrokj

Copy link
Copy Markdown
Collaborator

@seungrokj @edwingao28 hi, could you kindly review this?

@gbyu-amd can you plz check c1 interactivity ?

gbyu-amd and others added 2 commits October 26, 2026 10:07
PrefillDelayer (ATOM_PREFILL_DECODE_INTERVAL=1000000), max-num-seqs 2 and
CUDA-graph sizes [1, 2], per the Kimi-K3 recipe in ROCm/ATOM#2659.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…x-mono-statecache

# Conflicts:
#	inferencex-e2e/perf-changelog.yaml
gbyu-amd and others added 2 commits October 26, 2026 11:09
ATOM e893780f4 carries ROCm/ATOM#2659 (16-row Kimi-K3 mono, K2 on the MLA
layers). Recipe frontend pinned to the tag's digest.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@gbyu-amd
gbyu-amd marked this pull request as ready for review October 26, 2026 16:09
@claude

claude Bot commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Claude finished @gbyu-amd's task in 31s —— View job


Re-review

LGTM, no blocking issues found.

Since the last review, the C2 point was dropped, and C1 now uses max-num-seqs: 2 with cudagraph-capture-sizes: '[1,2]' and ATOM_PREFILL_DECODE_INTERVAL=1000000. HERD is now on at every point, the image moved to nightly_202610261558, and the base sets ATOM_MONO_ENABLE=0. What I checked:

  • Recipe selection: each concurrency in [1, 4, 14, 16, 48, 56, 72] matches exactly one override, and no CONC=2 variant is left behind.
  • DRAM: the 0.6 block yields 1799 GB, so C48, C56 and C72 match the recipe's TOTAL_CPU_DRAM_GB: '1799' (224 GB per rank at TP8). C14 and C16 stay at 1028.
  • C1 mono overrides: the env and quantization overrides apply only to C1. The exclude list keeps the draft-layer exclusions.
  • Changelog: the entry is appended at the end of perf-changelog.yaml.

Not checked: whether the recipe digest sha256:57477e7a… belongs to tag rocm/atom-dev:nightly_202610261558. The registry lookup needs approval in this environment, so please confirm it yourself. Also, the PR title still says "C1/C2", but C2 was removed.

@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.

I reviewed the current revision and found no bugs. amd-master.yaml has designated code owners, so one of them should still sign off on the retune.

What was reviewed:

  • The DRAM math. 0.343, 0.513 and 0.6 all imply the same host total of about 2998 GB, so 0.6 gives 1799 GB, and divided by TP8 that is the 224 GB/rank set in LMCACHE_MAX_LOCAL_CPU_SIZE. C48 also moved out of the 0.343 block and its TOTAL_CPU_DRAM_GB matches.
  • The C1 override. Its online_quant_config keeps every base exclude and adds the mono BF16 patterns. The fp32 SSM dtype and max-num-seqs 2 with capture sizes [1,2] are consistent with each other. The base ATOM_MONO_ENABLE=0 leaves the other points off.
  • The perf-changelog block. It is appended only at the tail, covers the one affected config key, and its block-style scenario-type matches existing entries.
Extended reasoning...

The change retunes one AMD agentic benchmark config (kimik3-fp4-mi355x-atom-agentic-mtp). It bumps the image and digest, enables mono decode at C1, enables on-demand state checkpoints and a larger LMCache tier at C48/56/72, and adds the matching perf-changelog entry. It touches no security-sensitive surface; it is benchmark YAML only. I deferred instead of approving because CODEOWNERS assigns inferencex-e2e/configs/amd-master.yaml to specific AMD owners. The timeline also has a comment from one of those owners whose content this review cannot see.

gbyu-amd and others added 2 commits October 28, 2026 11:22
…x-mono-statecache

# Conflicts:
#	inferencex-e2e/perf-changelog.yaml
Its aiter carries ROCm/aiter#7020, which keeps the tuned Kimi-K3 MXFP4
inline-sort MoE rows (C4/C14/C16 decode) under ATOM's real row count.
Recipe frontend pinned to the tag's digest.

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

This branch has not been deployed

No deployments
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 Full sweep with canary gate; matrix jobs run to completion despite failures

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants