[AMD][AgentX] Add Kimi-K3 MXFP4 MI355X vLLM agentic MTP recipe

[seungrokj]

Summary

• Adds kimik3-fp4-mi355x-vllm-agentic-mtp: single-node TP8 agentic-coding recipe for Kimi-K3 (MXFP4, 2.8T MoE, KDA/MLA hybrid, 1M native context) on MI355X (gfx950) with vLLM + DSpark MTP speculative decoding.
• Container patches (apply_k3_container_patches.sh) delivered as an idempotent git apply wrapper (KDA state_indices coercion, triton_mla cudagraph support, rocm_aiter_mla small-head verify, mla_gluon int64 KV stride) plus Triton 3.7.0.
• Piecewise cudagraph capture: compilation-config mode:3 with an explicit even-step cudagraph_capture_sizes list up to 2*CONC*(1+SPEC_NUM_TOKENS) (6*CONC at spec=2); --max-num-seqs = 1*CONC.
• Spec sampling uses synthetic rejection with synthetic_acceptance_length 2.45; a8w4 MoE via VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1.
• configs/amd-master.yaml: adds the recipe entry with conc-list: [4, 8, 10, 12, 14, 16].

Test plan

• Full-sweep validation green on MI355X (gfx950) across the conc ladder
• apply_k3_container_patches.sh applies cleanly in-container and is a no-op on re-run
• CUDA graph capture completes (no KDA state_indices / mla_gluon faults)

🤖 Generated with Claude Code

[github-actions]

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]

🔴 Line 770 hardcodes KV_CACHE_DTYPE_ARGS=(--kv-cache-dtype "auto") as a literal instead of (--kv-cache-dtype "$KV_CACHE_DTYPE"), so the served vLLM instance always launches with bf16 KV regardless of the $KV_CACHE_DTYPE variable computed earlier in the script. This silently disables the recipe's documented primary optimization (fp8 KV) and, when an LMCache/fp8 offload arm is enabled via KV_CACHE_DTYPE=fp8, causes an N=1536-vs-768 block-size mismatch between the LMCache sizing logic and the actual bf16 server allocation.

Extended reasoning...

The bug: benchmarks/single_node/agentic/kimik3_fp4_mi355x_mtp.sh:770 sets:

KV_CACHE_DTYPE_ARGS=(--kv-cache-dtype "auto")

as a hardcoded string literal. This array is the only place the vllm serve invocation gets its --kv-cache-dtype flag (line 818, "${KV_CACHE_DTYPE_ARGS[@]}" inside VLLM_CMD). Nowhere in the script is KV_CACHE_DTYPE_ARGS reassigned from the $KV_CACHE_DTYPE shell variable.

Where the computed value goes dead: $KV_CACHE_DTYPE is computed at line 206 (KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-auto}"), force-set to fp8 for the vllm-simple-fp8/vllm-simple-fp8-lazy offload backends (line 689), and used to branch the LMCache unified-block-size derivation (LMCACHE_UNIFIED_BLOCK = 1536 for fp8 vs 768 otherwise, around line 272). All of that computation is legitimate and consistent — it's just never wired into the actual served flag.

Why the surrounding code doesn't catch this: The script has extensive comments asserting fp8 KV is "the DEFAULT for every arm on this model" and "the single largest lever available" (headers around lines 47 and 189-203), and the mla_gluon patch is explicitly said to "gate on KV_CACHE_DTYPE=fp8 and must see it already set." None of that intent is enforced at the point where the vLLM command is actually assembled — the literal at line 770 silently overrides it.

Concrete proof-by-example:

1. A user sets KV_OFFLOAD_BACKEND=vllm-simple-fp8 (or any path that forces KV_CACHE_DTYPE=fp8 at line 689).
2. The LMCache/offload sizing logic reads $KV_CACHE_DTYPE == fp8 and derives LMCACHE_UNIFIED_BLOCK=1536, then computes MAX_NUM_BATCHED_TOKENS and LMCACHE_CHUNK_SIZE under the assumption N=1536 (an fp8-sized page).
3. At line 770, KV_CACHE_DTYPE_ARGS is still the literal (--kv-cache-dtype "auto").
4. vllm serve launches with --kv-cache-dtype auto, i.e., bf16 KV pages, whose actual unified block size is 768, not 1536.
5. The server's real N=768 now violates the N <= max_num_batched_tokens < 2*N constraint the earlier code derived for N=1536 — this is exactly the ValueError: Mamba-hybrid models with LMCache require block_size <= max_num_batched_tokens < 2*block_size failure the script's own comments warn about a few hundred lines earlier.
6. Even outside the LMCache arms, any attempt at an fp8-vs-bf16 A/B via KV_CACHE_DTYPE=fp8 env override silently does nothing — the server always runs bf16, doubling KV memory per token (~217 vs ~108 KiB/token) on a 1M-context recipe that is explicitly memory-tight.

Why it doesn't crash the currently-shipped sweep: The one search-space arm actually configured in configs/amd-master.yaml for this recipe uses kv-offloading: none, which skips the offload code path entirely, so KV_CACHE_DTYPE stays at its own default (auto) — coincidentally matching the hardcoded literal. The shipped matrix will not crash, but its results will silently be bf16-KV numbers being reported under a recipe whose whole premise (per its own comments) is fp8 KV as the primary lever, and enabling the fully-implemented LMCache/fp8 arms in the future would immediately hit the engine-init failure described above.

The fix is a one-line change: KV_CACHE_DTYPE_ARGS=(--kv-cache-dtype "$KV_CACHE_DTYPE").

[claude]

🟡 Inside the lmcache* KV-offload branch, MAX_NUM_SEQS and GPU_MEM_UTIL are deliberately lowered (via ${VAR:-default}) to leave GPU headroom for LMCache's staging buffers, per a detailed comment documenting a real OOM this fixes. But the unconditional plain assignments MAX_NUM_SEQS=$((1 * CONC)) and GPU_MEM_UTIL="0.9" later in the script clobber those values before vllm serve runs, so any LMCache offload arm would silently reintroduce the exact OOM the branch exists to prevent. Fix by making those assignments respect an already-set value, e.g. MAX_NUM_SEQS=${MAX_NUM_SEQS:-$((1 * CONC))} and GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.9}".

Extended reasoning...

The LMCache offload branch (case "${KV_OFFLOAD_BACKEND:-}" in lmcache|lmcache-k27|lmcache-budget|lmcache-k3)) sets MAX_NUM_SEQS="${MAX_NUM_SEQS:-32}" and GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.8}". The extensive comment directly above these lines documents a real production incident (run 30559676068): the LMCache connector's GPU-side staging buffers halve the effective OOM threshold compared to the no-offload case (num_computed_tokens=184,320 at only 5.2% KV usage vs. 360,960+ at 10-21% without the connector), and explains the two specific levers (max_num_seqs 128→32, gpu_memory_utilization 0.88→0.85) chosen to compensate.

The problem is that further down the script, after the offload case block has already run, two more assignments execute unconditionally for every code path:

MAX_NUM_SEQS=$((1 * CONC))
...
GPU_MEM_UTIL="0.9"

These are plain assignments, not ${VAR:-default} forms, so they always overwrite whatever the offload branch computed. Since they run after the case block in the script's control flow and before VLLM_CMD is built (which reads $MAX_NUM_SEQS and $GPU_MEM_UTIL for --max-num-seqs and --gpu-memory-utilization), the offload-branch's carefully-derived low-memory-headroom values never reach the actual vllm serve invocation.

Concretely: set KV_OFFLOAD_BACKEND=lmcache (or lmcache-k27/lmcache-budget/lmcache-k3) and trigger agentic_kv_offload_enabled. The case block runs and sets MAX_NUM_SEQS=32, GPU_MEM_UTIL=0.8 (assuming the env vars are otherwise unset). Then, unconditionally, MAX_NUM_SEQS=$((1 * CONC)) runs — at conc=16 this becomes MAX_NUM_SEQS=16, and even at higher conc this is decoupled from the intended 32 value entirely — and GPU_MEM_UTIL="0.9" overwrites 0.8. VLLM_CMD is then built with --gpu-memory-utilization 0.9 --max-num-seqs $MAX_NUM_SEQS, i.e. exactly the higher-memory-pressure configuration the comment block says caused HSA_STATUS_ERROR_OUT_OF_RESOURCES. The LMCache arm would OOM again, and the mitigation the author wrote (and documented at length) never actually executes.

Nothing in the existing code guards against this: the two blocks are independent plain-assignment statements with no ordering dependency check, and there's no assertion or comment near lines 779/790 acknowledging that they must not run when the offload branch already set these values.

The fix is straightforward and matches the pattern already used for MAX_NUM_SEQS/GPU_MEM_UTIL at the top of the offload branch: change both to ${VAR:-default} form, e.g. MAX_NUM_SEQS="${MAX_NUM_SEQS:-$((1 * CONC))}" and GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.9}". That way the no-offload path (which never sets these vars beforehand) is unaffected, while the LMCache arms' pre-set values survive through to vllm serve.

One reviewer flagged this as a duplicate of a broader "bug_002" umbrella finding — that's a synthesis/tracking concern about how findings were split across reports, not a reason to doubt the underlying defect, which all four independent verifiers reproduced against the current source and confirmed as real. Severity is nit rather than normal because the only search-space arm actually shipped in configs/amd-master.yaml for this recipe uses kv-offloading: none, so the offload branch is never entered by the configured sweep and lines 779/790 currently hold exactly the intended no-offload values — this PR's CI will not be affected. But the LMCache offload path is fully implemented, documented as an intended future arm ("both should be walked back toward the no-offload values once this arm produces a number"), and reachable via env vars today, so anyone activating it (in this PR or a future config change) will hit this immediately.

[github-actions]

see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31090250586
see unofficial run visualizer at https://inferencex.semianalysis.com/evaluation?unofficialRun=31090250586

[github-actions]

see unofficial run visualizer at https://inferencex.semianalysis.com/inference?unofficialRun=31554596012
see unofficial run visualizer at https://inferencex.semianalysis.com/evaluation?unofficialRun=31554596012

[seungrokj]

/reuse-sweep-run

[seungrokj]

will work on waiver.md & vllm recipe update

[seungrokj]

recipe: vllm-project/recipes#788

once all patches are merged to the vllm, will update the final recipe (some variables will work only when all patches are meged)

[billishyahao]

As a PR reviewer and CODEOWNER, I have reviewed this and have:

• Verified that as of the moment of typing this, this is the latest version of PR_REVIEW_CHECKLIST.md
• Verified that the general code quality meets the InferenceX standard and does not make the code quality any worse.
• Verified that this PR has passed PR validation. Please link to GitHub Action workflow that shows this. https://github.com/semiAnalysisAI/InferenceX/actions/runs/31554596012
• Verified that this PR passes evals. Please link to GitHub Action workflow that shows this. https://github.com/semiAnalysisAI/InferenceX/actions/runs/31554596012
• Verified that speculative decoding PRs uses chat templates to align the AL distribution to real world
• For agentic workloads: verified that speculative-decoding configs (EAGLE / MTP / draft models) run with simulated synthetic acceptance, with the acceptance-length value taken from the committed golden AL curve in golden_al_distribution/ for that model, thinking mode, and draft length. A submission may choose any supported draft length, but it may not substitute a different acceptance target.
• Verified against the current MODELS.md that this PR does not submit a deprecated model, scenario, or model-scenario combination.
• Verified that the model architecture isn't changed with benchmark hacks like using --hf-overrides to skipping indexer for every x layers on models that don't natively support this. As a general rule, we won't accept optimizations that reduces the number of model architecture FLOPs. Anything that makes that same computation run faster is fair game; FLOPs at lower precisions is fine, given that the config passes private evals. As an general north star princple, we should only use optimizations which is used in production by customers that care about accuracy
• If an company claims that they support vLLM/SGLang as first class LLM inference engines on their hardware, I have verified that the respective vLLM submission made using upstream https://hub.docker.com/u/vllm docker repo, upstream SGLang https://hub.docker.com/u/lmsysorg docker repo. The only exceptions are for new hardware, such as MI455X UALoE72, Vera Rubin NVL72, Rubin NVL8, etc., and for new model architectures where there is an actual reason why vLLM/SGLang does not fundamentally support them yet as supported by vLLM/SGLang community maintainers
• If an company claims that they support vLLM/SGLang as first class upstream in-tree LLM inference engines on their hardware, I have have verified that the respective vLLM/SGLang submission has been made before additional frameworks (TRT-LLM, ATOM, etc.). The only exceptions are for new hardware, such as MI455X UALoE72, Vera Rubin NVL72, Rubin NVL8, etc., and for new model architectures where there is an actual reason why vLLM/SGLang does not fundamentally support them yet.
• Verified that every single-node vLLM/SGLang recipe in this PR is documented in the official vLLM recipes and/or the SGLang cookbook: vllm-project/recipes#788
  • I linked the corresponding upstream PR in the vLLM recipe repo or SGLang repo and verified that it is MERGED before this InferenceX PR merges. An opened, draft, or closed-without-merge upstream PR does not satisfy this requirement. If the matching recipe was already published, I linked the published recipe/cookbook page in the additional detail section below. https://docs.sglang.io/cookbook/autoregressive/DeepSeek/DeepSeek-V4#hw=mi355x&variant=pro&quant=fp4&strategy=high-throughput&nodes=single
• [] Verified that this PR does not patch the inference engine or serving stack — the pinned image must run as shipped. This covers .patch files / git apply / patch, inline patches embedded in benchmark scripts (e.g. a python3/sed heredoc that rewrites installed engine sources before serving), in-place edits of site-packages, monkey-patching, overwriting container files, and installing forked/rebuilt engine wheels on top of the pinned image. The only exception is a patch covered by a filled-out waiver at docs/waiver/<PR_NUMBER>.md — named after the PR that introduces the patch and filed in that same PR, stating what is patched, why the unmodified upstream image cannot run this benchmark, the upstream PR/issue link, and the removal plan — which I have linked below in the additional detail section. @seungrokj will work on the waiver.md
• If any of the above criteria cannot reasonably be satisfied, I have provided additional reasoning below.

Additional detail section:

• insert any additional info here

Signed: billishyahao

[Klaud-Cold]

❌❌❌ REJECTED ❌❌❌

@billishyahao — blocking: this PR patches the inference engine inside the pinned image with no waiver (docs/waiver/2508.md does not exist and docs/waiver/ is empty), the sign-off itself leaves that checklist item unchecked, and the benchmark's --kv-cache-dtype fp8 is a MAJOR arg not in the published upstream recipe's AMD config.

✅ Check 0 (CODEOWNER): PASS — billishyahao is a listed owner of configs/amd-master.yaml; the other changed paths are catch-all-only.
✅ Check 1 (sweep on in-PR commit): PASS — run 31554596012 ran on head 0baf642 with all executed agentic / (c1/c4/c8, dram c10) and agentic eval / check-runs success (fixed-seq-len rows correctly skipped for this agentic-only config).
✅ Check 2 (evals pass): PASS — GSM8K em_strict 0.9666 ≥ 0.90 bar (kimik3 uses the default threshold), run used this PR's image vllm/vllm-openai-rocm:nightly-cb8104839c....
❌ Check 3 (recipe): FAIL — --kv-cache-dtype fp8 (a MAJOR arg) is missing from the published Kimi-K3 recipe's AMD configuration: recipes PR #788 is MERGED and the published recipe sets fp8 KV only in its Blackwell block; the AMD block leaves KV dtype at default (bf16). The patch stack itself carries "vllm #51011 fix fp8 KV cache decode on the AITER MLA backend", confirming the shipped image doesn't serve this config. Model/TP8/MXFP4/DSpark(TRITON_MLA, probabilistic) all match; draft length 2 vs 7, fastsafetensors, cudagraph mode, and synthetic acceptance are InferenceX-specific (informational). Note: the sign-off's published-page link points at DeepSeek-V4's SGLang cookbook, not this recipe.
✅ Check 4 (reuse command): PASS — /reuse-sweep-run posted by seungrokj (COLLABORATOR).
❌ Check 5 (latest checklist): FAIL — the engine-patch item is left unchecked with no explanation in the additional detail section, which is the untouched template placeholder ("insert any additional info here"); the inline "@seungrokj will work on the waiver.md" concedes the requirement is unmet. All other current-template items are present and checked.
✅ Check 6 (upstream image + ordering): PASS — vllm/vllm-openai-rocm:nightly-cb8104839c... is from the upstream vLLM org; framework is vllm, so engine-first ordering is satisfied.
✅ Check 7 (deprecated models): PASS — Kimi-K3 agentic-coding with DSpark is the active combination in MODELS.md as of 2026-08-12.
✅ Check 8 (architecture hacks): PASS — no --hf-overrides/layer-trimming; --language-model-only is the recipe's documented text_only opt-in and removes no FLOPs from the text workload.
✅ Check 9 (chat templates): PASS — replay drives /v1/chat/completions with --endpoint-type chat (benchmarks/benchmark_lib.sh:1930); evals use local-chat-completions --apply_chat_template.
❌ Check 10 (engine patches): FAIL — no waiver exists. apply_k3_container_patches.sh (invoked by kimik3_fp4_mi355x_mtp.sh before vllm serve) git-applies/rewrites ~10 installed vLLM+aiter source files in the pinned image, downloads binary kernels and rebuilds aiter module_mla_asm (ROCm/aiter#4521), and pip-installs Triton 3.7.0 over the image. This PR does not add docs/waiver/2508.md, docs/waiver/ is empty on main, and the sign-off links no waiver.
✅ Check 11 (golden AL): PASS — benchmark arm pins rejection_sample_method: synthetic with synthetic_acceptance_length: 2.51, matching the committed golden AL for kimi-k3 thinking_on at num_speculative_tokens=2 with probabilistic draft + block rejection (golden_al_distribution/kimik3_dspark_probabilistic_sample_method_block_rejection_sample_method.yaml); eval arm runs real block rejection. (PR body/changelog text says 2.45 — stale; the script pins 2.51.)

[seungrokj]

@functionstackx @cquil11
added docs/waiver/2508.md
recipe --kv-cache-dtype fp8 update will be after full patch merge to vllm.