[AMD] [AGENTX] Kimi-K3 Performance Tuning - #2602
Conversation
Signed-off-by: Sirra <asirra@amd.com>
Signed-off-by: Sirra <asirra@amd.com>
Signed-off-by: Sirra <asirra@amd.com>
Signed-off-by: Sirra <asirra@amd.com>
|
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 For PR verification, add the 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 模板，包括保留英文语句 如需进行 PR 验证，请为此 PR 添加 PR 作者有责任确保合并后所有 GitHub Action 任务完全通过。 很多时候失败只是偶发抖动（flake），重新运行失败的任务即可解决。参见 GitHub 关于重新运行失败任务的文档 |
ajith-sirra-amd
added
AMD
agentx
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| agentic-coding: | ||
| - dram-utilization: 0.50 | ||
| search-space: | ||
| - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} | ||
| # - { tp: 8, kv-offloading: none, conc-list: [1, 4, 8] , spec-decoding: mtp} | ||
| - { tp: 8, ep: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [10], spec-decoding: mtp } |
There was a problem hiding this comment.
🔴 AGENTS.md:21 requires a perf-changelog.yaml entry for every change that can affect benchmark performance or any recipe addition/modification, but this PR is missing one despite bumping the kimik3-fp4-mi355x-vllm-agentic-mtp image, disabling the tp:8/conc-list [1,4,8] search-space arm in configs/amd-master.yaml, and flipping the default MLA_PREFILL_BACKEND from FLASH_ATTN to ROCM_AITER_FA in kimik3_fp4_mi355x_mtp.sh. Please append an entry documenting these changes and the measured throughput deltas (+4.4%/+20.1% prefill, 5.52x DSpark) before merging.
Extended reasoning...
AGENTS.md:21 states a non-negotiable invariant: "Every change that can affect benchmark performance and every recipe addition or modification requires a new perf-changelog.yaml entry. The file is append-only and byte-sensitive. Preserve all existing bytes and separator whitespace, and append only at the tail." This is an explicit repository-mandated check, not a stylistic suggestion, and it applies squarely here.
This PR makes several performance-affecting modifications to the kimik3-fp4-mi355x-vllm-agentic-mtp recipe: it bumps the pinned image from nightly-cb8104839c141609d99f1254459ef3a4f1bd4263 to nightly-3ee2df30337a301164c46ae444b76ee67e71c106 in configs/amd-master.yaml, comments out the tp:8/kv-offloading:none/conc-list [1,4,8] search-space arm (removing it from the sweep entirely), and in kimik3_fp4_mi355x_mtp.sh flips the default MLA_PREFILL_BACKEND from FLASH_ATTN to ROCM_AITER_FA plus wires in two new patch scripts (apply_aiter_pybind11_fix.sh and apply_triton_mla_cudagraph_fix.sh) that the PR itself claims produce measured throughput deltas of +4.4%/+20.1% on prefill and a 5.52x speedup for DSpark spec-decode. Each of these is independently sufficient to trigger the AGENTS.md rule — the image bump and search-space edit are recipe modifications to configs/amd-master.yaml, and the backend flip plus patch scripts directly change benchmark performance/behavior.
Despite this, perf-changelog.yaml is not in the PRs changed-files list at all — only apply_aiter_pybind11_fix.sh, apply_triton_mla_cudagraph_fix.sh, kimik3_fp4_mi355x_mtp.sh, and configs/amd-master.yaml were touched. I confirmed the file exists at the repo root (368KB, append-only) and its existing entries follow a config-keys / description / pr-link format per PR (see the tail entries for PRs #2550, #2554, #2562, #2570) — none reference this PR or the kimik3-fp4-mi355x-vllm-agentic-mtp recipe.
Nothing in the existing code prevents this omission — there is no CI check enforced here that fails a PR for a missing changelog entry visible in this diff, so the requirement relies entirely on author/reviewer discipline per AGENTS.md. Concretely, a reviewer or downstream consumer relying on perf-changelog.yaml as the source of truth for why a recipes performance profile changed (e.g. someone investigating a throughput regression/improvement report for kimik3-fp4-mi355x-vllm-agentic-mtp after this merges) would find no record explaining the image bump, the removed search-space arm, or the backend default flip and its claimed gains.
Step-by-step proof:
- Search
perf-changelog.yamlfor anyconfig-keysentry containingkimik3-fp4-mi355x-vllm-agentic-mtp— none exists in the current file (confirmed via the tail of the file, and the PR does not add one since the file is absent from the diff). - Diff
configs/amd-master.yaml: thekimik3-fp4-mi355x-vllm-agentic-mtpimage tag changes and thetp:8, conc-list:[1,4,8]line is commented out — both are recipe modifications per AGENTS.md:21. - Diff
kimik3_fp4_mi355x_mtp.sh:MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-ROCM_AITER_FA}"replaces the priorFLASH_ATTNdefault, and the PR description quantifies this as a measurable performance change. - Per AGENTS.md:21, each of steps 2–3 alone "requires a new perf-changelog.yaml entry" — the PR should append one covering the image bump, the removed search-space arm, and the backend default change with a pointer to this PR, before merge.
Fix: append a new entry at the tail of perf-changelog.yaml (preserving all existing bytes/whitespace per the byte-sensitive requirement) with config-keys: [kimik3-fp4-mi355x-vllm-agentic-mtp], a description summarizing the AITER pybind11 fix, the TritonMLA cudagraph fix, the MLA_PREFILL_BACKEND default flip, the removed search-space arm, and the measured throughput deltas, plus a pr-link to this PR.
| MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-ROCM_AITER_FA}" | ||
| MLA_PREFILL_ARGS=() | ||
| if [ -n "$MLA_PREFILL_BACKEND" ]; then | ||
| MLA_PREFILL_ARGS=( | ||
| --attention-config | ||
| "{\"mla_prefill_backend\":\"$MLA_PREFILL_BACKEND\"}" | ||
| ) | ||
| fi |
There was a problem hiding this comment.
🟡 Minor dead-code nit: on line 214, MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-ROCM_AITER_FA}" already guarantees the variable is non-empty (bash's :- fires on both unset and empty), so the if [ -n "$MLA_PREFILL_BACKEND" ] on line 216 is always true and the implicit empty-MLA_PREFILL_ARGS else-branch is unreachable. Consider dropping the guard and building MLA_PREFILL_ARGS unconditionally, matching how COMPILATION_CONFIG_ARGS/SPEC_ARGS are built elsewhere in this file.
Extended reasoning...
What the bug is: Line 214 sets MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-ROCM_AITER_FA}". Bash's ${var:-default} form (with the colon) substitutes default whenever var is either unset or set to the empty string — this is distinct from ${var-default} (no colon), which only substitutes on unset. Because of the colon, after line 214 executes, $MLA_PREFILL_BACKEND is guaranteed to be non-empty: it is either the caller's non-empty override or the literal ROCM_AITER_FA fallback.
The code path: Line 216 then guards the array construction with if [ -n "$MLA_PREFILL_BACKEND" ]; then MLA_PREFILL_ARGS=(--attention-config ...); fi. Since the variable can never be empty at this point, the condition is always true, and the implicit else (which would leave MLA_PREFILL_ARGS=() empty) is dead code that can never execute — not even via MLA_PREFILL_BACKEND="" as an explicit env override, since that empty string is itself replaced by the default on line 214 before the check ever runs.
Why nothing catches this today: There's no test or lint that exercises unreachable branches in this shell script, and the guard reads as plausible defensive code at a glance (mirroring the real, reachable guards used for EP_ARGS at line ~199 and OFFLOAD_ARGS's case statement), so it's easy to miss that this particular guard's precondition was already foreclosed on the prior line.
Impact: None functionally — behavior is identical whether or not the if is present, since the branch it guards against can never be taken. This is purely a readability/maintainability nit.
Proof by walkthrough:
- Caller leaves
MLA_PREFILL_BACKENDunset → line 214::-fires →MLA_PREFILL_BACKEND=ROCM_AITER_FA. - Caller sets
MLA_PREFILL_BACKEND=""→ line 214::-fires (colon form applies to empty too) →MLA_PREFILL_BACKEND=ROCM_AITER_FA. - Caller sets
MLA_PREFILL_BACKEND=FLASH_ATTN→ line 214::-does not fire →MLA_PREFILL_BACKEND=FLASH_ATTN. - In all three cases, line 216's
[ -n "$MLA_PREFILL_BACKEND" ]evaluates to true, soMLA_PREFILL_ARGSis always built with the--attention-configflag — the else-path (empty array) is unreachable in every case.
Suggested fix: Drop the if/fi and build MLA_PREFILL_ARGS unconditionally, e.g.:
MLA_PREFILL_BACKEND="${MLA_PREFILL_BACKEND:-ROCM_AITER_FA}"
MLA_PREFILL_ARGS=(
--attention-config
"{\"mla_prefill_backend\":\"$MLA_PREFILL_BACKEND\"}"
)This matches the pattern already used for COMPILATION_CONFIG_ARGS and SPEC_ARGS earlier in the same file, both of which are built unconditionally without a reachable empty-array path.
Signed-off-by: Sirra <asirra@amd.com>
Signed-off-by: Sirra <asirra@amd.com>
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31764459116 |
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31765562914 |
Signed-off-by: Sirra <asirra@amd.com>
Signed-off-by: Sirra <asirra@amd.com>
…SemiAnalysisAI/InferenceX into amd/kimi-k3-agentic-perf-tuning
Signed-off-by: Sirra <asirra@amd.com>
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31783189021 |
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31783353670 |
Signed-off-by: Sirra <asirra@amd.com>
Signed-off-by: Sirra <asirra@amd.com>
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31791138854 |
1 similar comment
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31791138854 |
Signed-off-by: Sirra <asirra@amd.com>
2 participants
Summary: [AMD] [AGENTX] KIMI-K3 Perf Tuning
Commit
c0119e6be3— 3 files changed, 212 insertions(+), 5 deletions(-)What changed
1. New:
apply_aiter_pybind11_fix.shFixes a pybind11 version mismatch between AITER's prebuilt
.sofiles (built against torch's bundled pybind11, internals v11) and its JIT builder (which injects standalone pybind11 v12 via-I, outranking torch's-isystempath). Since pybind11 keeps a separate type registry per internals version, JIT modules couldn't see types registered by the prebuilt core, crashing the server during warmup:The fix patches
aiter/jit/utils/cpp_extension.pyto prefer torch's bundled pybind11, and purges stale JIT-built.soartifacts compiled against the wrong version. It's idempotent and verifies the mismatch exists before touching anything.Impact: unblocks
ROCM_AITER_FAfor MLA prefill (previously had to fall back toFLASH_ATTN). On 8x MI355X / Kimi-K3 MXFP4 TP8:2. New:
apply_triton_mla_cudagraph_fix.shFixes DSpark spec-decode falling back to fully eager execution on ROCm.
TritonMLAMetadataBuilderdeclared_cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE, which silently downgradedFULL_AND_PIECEWISE→PIECEWISEcudagraph mode whenever spec-decode was enabled, giving the DSpark drafterCUDAGraphMode.NONEwith no warning logged.TRITON_MLAcan't simply be swapped forROCM_AITER_MLAsince it's the only ROCm MLA backend supporting non-causal multi-token decode (required by the DSpark draft).The fix patches the backend to declare
AttentionCGSupport.UNIFORM_BATCHinstead, restoring full cudagraph capture.Impact: on 8x MI355X, Kimi-K3 MXFP4 TP8, DSpark, single stream, 600-token generations:
3. Modified:
kimik3_fp4_mi355x_mtp.shinstall_agentic_deps.MLA_PREFILL_BACKENDfromFLASH_ATTNtoROCM_AITER_FA, now that the pybind11 fix resolves the underlying AITER crash.MLA_PREFILL_BACKEND=FLASH_ATTNremains available as a fallback knob.Net effect
Both fixes together unlock AITER's faster MLA prefill path and restore full CUDA-graph capture for spec-decode, compounding into a substantial throughput improvement for Kimi-K3 MXFP4 on MI355X.