[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=31273221156
see unofficial run visualizer at https://inferencex.semianalysis.com/evaluation?unofficialRun=31273221156

[github-actions]

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

[github-actions]

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

[github-actions]

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

[github-actions]

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