[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>
Signed-off-by: Sirra <asirra@amd.com>
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31805769600 |
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31807327119 |
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=31808848259 |
Merge amd/kimi-k3-agentic-perf-tuning (#2602) so the LMCache DRAM offload arm inherits its fixes rather than measuring the older stack: - aiter pybind11 internals fix, unblocking ROCM_AITER_FA MLA prefill - TritonMLA cudagraph support (FULL cudagraphs for DSpark) - KV block-pool clamp + --no-async-scheduling The last one matters most here. #2602 traces the mid-run engine crash to async scheduling: DSpark is exempt from vLLM's async-scheduling disable list, so max_concurrent_batches = 2, and with kv_role=kv_both (is_kv_consumer=True) the scheduler sets defer_block_free=True, leaving freed-but-still-referenced KV blocks on the free list. Our LMCache connector runs kv_role=kv_both, so it takes that same path -- and the crash we saw at c12 (block-pool assert in kv_cache_utils.popleft_n) is the same family as the block_pool.py ref_cnt assert #2602 describes, with the same concurrency scaling. Also bump the lmcache key's image to match the base key's new pin (nightly-ac7509e2b1). The key's premise is "same image, script and topology as kimik3-fp4-mi355x-vllm-agentic-mtp", so it has to track it. Conc list left at 4/8/10/12 rather than the base key's new 1/4/8/10/12/14/16/20 to keep the sweep bounded.
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31814935823 |
Merging #2602 brought its perf-changelog entry along, so its whole vllm-simple sweep (c1/c4/c8/c10/c12/c14/c16/c20) runs under this PR -- roughly doubling the GPU cost and landing its failures here. In the last run that arm went 2-for-9: - c8, c10, c14, c20-eval: illegal memory access, ALL on mia1-p01-g11 (a bad node -- it also took this PR's own c12 eval) - c1, c12, c16: RCCL all-gather watchdog timeout in _gather_logits during sampling, on three different nodes, at concurrencies as low as 1 Neither is caused by the LMCache arm, which went 3-for-3 on the same image and script (c4/c8/c12; c10 still running), with c12 clean at 98.5% GPU KV and none of the crash signatures that killed it twice before. Keep the merged code from #2602 -- the AITER pybind11 fix, TritonMLA cudagraph support and --no-async-scheduling are what made c12 survive -- but let #2602 sweep its own key on its own PR.
Two fixes in one: 1. Restore PR #2571's entry (dsv4-fp4-gb300-dynamo-vllm-agentic-mtp-agg / -disagg). My automated merge-conflict resolution dropped it while reconciling the append-only tail, which the changelog gate correctly rejected -- deletions are not permitted. 2. Drop #2602's entry, which arrived with the code merge and pulled its whole vllm-simple sweep (c1..c20) onto this PR. That is not a deletion relative to main, since the entry only exists on #2602's branch. Rebuilding as 'main verbatim + this PR's entry' makes both correct by construction: zero deletions vs main, one addition.
The previous merge left conflict markers in the file and my scripted fallback asserted before rewriting it, so the broken version was committed and pushed. Rebuild deterministically as 'upstream/main verbatim + this PR's entry': zero deletions vs main, valid YAML, 733 entries, #2602's entry absent so its sweep no longer runs here.
Signed-off-by: Sirra <asirra@amd.com>
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=31857937367 |
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31858613759 |
1 similar comment
|
see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31858613759 |
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.