[wip - contains patches that need to be removed] [amd agentx m3 atom] - #2816
[wip - contains patches that need to be removed] [amd agentx m3 atom]#2816charxwu wants to merge 8 commits into
Conversation
新增 MiniMax-M3 ATOM AgentX recipe，固定可复现镜像并覆盖 MI355X TP4 concurrency sweep。 Align the ATOM serving entrypoint, prefix-cache settings, and generated matrix with the validated MI355X smoke path. 对齐 ATOM 服务入口、prefix cache 参数以及已验证的 MI355X smoke 路径矩阵。 Co-authored-by: Cursor <cursoragent@cursor.com>
修复 MiniMax ATOM prefix cache 默认关闭的问题。 Enable prefix caching by default so generated AgentX runs match the validated smoke configuration. 默认开启 prefix caching，使生成的 AgentX 运行与已验证 smoke 配置一致。 Co-authored-by: Cursor <cursoragent@cursor.com>
按 AgentX 并发度动态设置 ATOM max-num-seqs。 Scale ATOM scheduler capacity to two sequence slots per configured concurrency lane. 根据配置的并发 lane 数量为每个 lane 分配两个 ATOM scheduler sequence slot。 Co-authored-by: Cursor <cursoragent@cursor.com>
为 MiniMax AgentX 适配 ATOM PR #2106 补丁，并修复 served model 名称不匹配导致的 warmup 失败。 Co-authored-by: Cursor <cursoragent@cursor.com>
为 MiniMax ATOM AgentX MTP recipe 增加无 offloading 的 TP2 并发搜索空间，并更新矩阵快照校验。 Co-authored-by: Cursor <cursoragent@cursor.com>
让 MiniMax ATOM AgentX launcher 接受 TP2 search-space 配置。 Co-authored-by: Cursor <cursoragent@cursor.com>
将 MiniMax ATOM AgentX launcher 的 GPU memory utilization 提升到 0.9。 Co-authored-by: Cursor <cursoragent@cursor.com>
charxwu
changed the title
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 3 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit d08ed52. Configure here.
| STATE_CKPT_ARGS=() | ||
| if [ -n "$STATE_CHECKPOINT_SLOTS" ]; then | ||
| STATE_CKPT_ARGS=(--state-checkpoint-slots "$STATE_CHECKPOINT_SLOTS") | ||
| fi |
There was a problem hiding this comment.
Checkpoint slots never forwarded
Medium Severity
For concurrencies 8–28 the recipe enables ReplaySSM and builds STATE_CKPT_ARGS with 96 extra --state-checkpoint-slots, but those flags are never added to ATOM_CMD. Checkpoints then share only leftover max-num-seqs capacity, so the intended retained state pool is never reserved.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit d08ed52. Configure here.
| else | ||
| echo "ERROR: ATOM PR #2106 does not apply cleanly to $ATOM_ROOT" >&2 | ||
| exit 1 | ||
| fi |
There was a problem hiding this comment.
Inconsistent patch apply paths
Medium Severity
PATCH_EXCLUDES is used only on the git apply. The “already applied” reverse-check runs against the full diff, and the non-git patch path applies every file. After a successful partial apply, a rerun fails the reverse-check and then fails apply because the kept hunks are already present. A non-git tree also tries the two files the waiver says do not apply.
Reviewed by Cursor Bugbot for commit d08ed52. Configure here.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Beyond the inline findings, I also checked whether the per-concurrency case "$CONC" block's MAX_NUM_SEQS/MAX_NUM_BATCHED_TOKENS/GPU_MEM_UTIL settings are dead code now that they're overwritten right after the case statement — that overwrite is intentional (the adjacent comment states the fixed values are the official validated MiniMax-M3 launch settings), so it isn't a separate bug beyond what's already flagged inline.
Extended reasoning...
Confirmed findings already posted inline cover real, non-trivial bugs: the patch idempotency check in apply_atom_pr2106_patch.sh can never validate correctly against the excluded-file patch; SPEC_ARGS unconditionally enables EAGLE3 with 3 draft tokens even when the CONC=40/56 branches intend to disable it (NUM_SPEC_TOKENS=0); STATE_CKPT_ARGS is computed but never appended to ATOM_CMD, silently dropping --state-checkpoint-slots; and the waiver doc is misnamed relative to docs/PR_REVIEW_CHECKLIST.md's requirement. These are exactly the kind of functional discrepancies between comments/intent and actual behavior that warrant a human look, and the bug-hunt exit reason (max_bugs) means the search was cut off rather than exhausted, so no additional approval signal should be inferred. I verified the case block in minimaxm3_fp4_mi355x_atom_mtp.sh directly and confirmed the described overwrite pattern (MAX_NUM_SEQS=$((2*CONC)), fixed MAX_NUM_BATCHED_TOKENS=32768, GPU_MEM_UTIL=0.9 right after the case) is real, but distinguished it from a bug since the surrounding comment documents it as intentional — while the state-checkpoint-slots and spec-token gating issues in that same block are genuine bugs already captured inline. Given the number and severity of confirmed findings, this does not meet the bar for approval, and no prior review body exists on this PR to avoid duplicating, so a short informational defer note is appropriate.
Findings marked 🟡 are optional suggestions and need no follow-up push.
| if git -C "$ATOM_ROOT" apply --reverse --check "$PATCH_FILE" >/dev/null 2>&1; then | ||
| echo "ATOM PR #2106 already applied" | ||
| exit 0 | ||
| fi | ||
|
|
||
| # These files were structurally rewritten after the pinned image's ATOM | ||
| # commit. Their freshness guard is orthogonal to the draft-KV OOM fix and | ||
| # is intentionally omitted until it can be ported against that revision. | ||
| git -C "$ATOM_ROOT" apply --check "${PATCH_EXCLUDES[@]}" "$PATCH_FILE" | ||
| git -C "$ATOM_ROOT" apply "${PATCH_EXCLUDES[@]}" "$PATCH_FILE" |
There was a problem hiding this comment.
🔴 The 'already applied' idempotency check at line 27 runs git apply --reverse --check on the FULL patch, but the patch is only ever forward-applied with PATCH_EXCLUDES (backends.py, eagle_proposer.py excluded at lines 14-17). Since those two files never get patched, the whole-patch reverse-check can never succeed once already applied, so every re-run against an already-patched ATOM_ROOT falls through to the excludes-scoped apply --check at line 35, which now fails (hunks already applied) and aborts under set -e. Fix: scope the idempotency check to the same exclude set used to apply the patch (or detect a marker unaffected by the excludes), so an already-patched tree is recognized on every invocation, not just the first.
Extended reasoning...
First run: ATOM_ROOT is unpatched, reverse --check fails as expected, so the script applies the patch with PATCH_EXCLUDES (lines 35-36), leaving backends.py and eagle_proposer.py unmodified by design. Second run against the same ATOM_ROOT (e.g. a persistent/shared AMD runner container reused across the matrix's per-concurrency jobs, or a retried/rerun job) — line 27's reverse --check on the unfiltered patch requires ALL hunks including the excluded files' to reverse cleanly; those files are still in pre-patch state so reversal fails, so the 'already applied' fast path is never taken. The script proceeds to line 35 git apply --check (with excludes) against files that are already forward-patched, which fails because the expected pre-patch context no longer matches. Under set -euo pipefail this aborts the entire script with a raw git error instead of the friendly already-applied message, killing the whole benchmark launch for that concurrency point.
Verification: normal (conditional on re-run against a persistent ATOM_ROOT — the very scenario the idempotency block exists for). Line 27 git apply --reverse --check "$PATCH_FILE" checks the UNFILTERED patch (no PATCH_EXCLUDES), but the patch is only ever forward-applied with atom/model_ops/attentions/backends.py and atom/spec_decode/eagle_proposer.py excluded (lines 14-17, 35-36). Those two files are…
| STATE_CKPT_ARGS=() | ||
| if [ -n "$STATE_CHECKPOINT_SLOTS" ]; then | ||
| STATE_CKPT_ARGS=(--state-checkpoint-slots "$STATE_CHECKPOINT_SLOTS") | ||
| fi | ||
|
|
||
| # ---- KV offload ------------------------------------------------------------- | ||
| # K3 is a hybrid: MiniMax-M3 attention carries a per-request recurrent state | ||
| # alongside the paged KV. The paged KV rides this LMCache tier from | ||
| # concurrency 8 up; from concurrency 16 up the CPU state tier is switched on | ||
| # alongside it, because the state tier is what makes a resumed agentic turn | ||
| # cheap and the paged KV tier alone cannot restore one. | ||
| OFFLOAD_ARGS=() | ||
|
|
||
| case "$KV_OFFLOAD_BACKEND" in | ||
| "") | ||
| require_agentic_kv_offload_none | ||
| ;; | ||
| lmcache) | ||
| require_agentic_kv_offload_backend lmcache | ||
|
|
||
| # TOTAL_CPU_DRAM_GB is the AGGREGATE budget from the matrix generator. | ||
| # LMCACHE_MAX_LOCAL_CPU_SIZE and OFFLOAD_STATE_CPU_SIZE are per rank and | ||
| # every rank allocates its own, so the aggregate is divided by TP as the | ||
| # agentic README requires. Handing a rank the whole aggregate does not | ||
| # just overcommit -- it never finishes pinning and hangs the launch | ||
| # partway through. | ||
| PER_RANK_CPU_GB="$((TOTAL_CPU_DRAM_GB / TP))" | ||
| LMCACHE_CPU_GB="$((PER_RANK_CPU_GB - STATE_OFFLOAD_CPU_GIB))" | ||
|
|
||
| export PYTHONHASHSEED=0 | ||
| export LMCACHE_LOCAL_CPU=True | ||
| export LMCACHE_MAX_LOCAL_CPU_SIZE="$LMCACHE_CPU_GB" | ||
| # DCP-locked: the offload hash block is block-size(128) x dcp(8) = 1024, | ||
| # so the KV grid and the state-checkpoint grid coincide and the joint | ||
| # load aims both legs at one boundary. 512 or 2048 misaligns it. | ||
| export LMCACHE_CHUNK_SIZE=1024 | ||
| export OFFLOAD_KV_FOR_HYBRID=1 | ||
| # Statistics only -- per-step offload counters in the connector. Kept on | ||
| # because the submitted numbers were measured with it on. | ||
| export OFFLOAD_PROFILE=1 | ||
|
|
||
| if [ "$STATE_OFFLOAD_CPU_GIB" -gt 0 ]; then | ||
| # CPU state-offload tier for the attention state. | ||
| export OFFLOAD_STATE=1 | ||
| export OFFLOAD_STATE_CPU_SIZE="$STATE_OFFLOAD_CPU_GIB" | ||
| export OFFLOAD_STATE_STAGING_GROUPS=8 | ||
| export OFFLOAD_STATE_MIN_LOAD_TOKENS=0 | ||
| # Must be set: the staging buffer defaults to 2 chunks (8 MiB), one | ||
| # K3 state entry is 54.78 MiB, and a buffer too small to hold one | ||
| # entry makes the tier decline to build -- one log line, then | ||
| # nothing offloads, which reads exactly like a tier that is on and | ||
| # idle. | ||
| export OFFLOAD_GPU_STAGING_CHUNKS=32 | ||
| fi | ||
|
|
||
| OFFLOAD_ARGS=( | ||
| --kv-transfer-config | ||
| "{\"kv_connector\":\"lmcache_offload\",\"kv_role\":\"offload\"}" | ||
| ) | ||
| ;; | ||
| *) | ||
| echo "Unsupported KV_OFFLOAD_BACKEND: $KV_OFFLOAD_BACKEND (expected empty or lmcache)" >&2 | ||
| exit 1 | ||
| ;; | ||
| esac | ||
|
|
||
| # ---- ATOM env --------------------------------------------------------------- | ||
| echo "Starting atom server..." | ||
| export PYTHONNOUSERSITE=1 | ||
|
|
||
| # Required by ATOM: without it the aiter kernel logs flood the server log for | ||
| # the whole 3600 s replay. | ||
| export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-WARNING}" | ||
| export AITER_SITUV2_A4W4=1 | ||
| export AITER_QUICK_REDUCE_QUANTIZATION=INT4 | ||
| export AITER_FLYDSL_STAGE2_FP8=1 | ||
| export ATOM_FORCE_ATTN_TRITON=1 | ||
| # Anchor-only state checkpointing: the demand rung is 47% of checkpoint writes | ||
| # but reads back 2.8% of the time, against 85.2% for a prompt-end anchor, so it | ||
| # costs more in evictions than its reuse is worth on these traces. | ||
|
|
||
| # ---- Speculative ------------------------------------------------------------ | ||
| # golden_al_distribution/minimaxm3_eagle3_gqa.yaml | ||
| # 3 draft tokens -> AL 2.78 | ||
| # 3 draft tokens -> AL 2.78 | ||
| # Concurrency 32 and up serve without a draft model: past the throughput knee | ||
| # the draft forward no longer pays for itself against the resident batch. | ||
| # Use the official MiniMax-M3 EAGLE3 draft model and three speculative tokens. | ||
| SPEC_ARGS=( | ||
| --method eagle3 | ||
| --draft-model Inferact/MiniMax-M3-EAGLE3 | ||
| --num-speculative-tokens 3 | ||
| ) | ||
| echo "SPEC_DECODE_AL=$SPEC_DECODE_AL NUM_SPEC_TOKENS=$NUM_SPEC_TOKENS" | ||
|
|
||
| # ---- LLM server ------------------------------------------------------------- | ||
| ATOM_CMD=( | ||
| python -m atom.entrypoints.openai_server | ||
| --model "$MODEL_PATH" | ||
| --served-model-name "$MODEL" | ||
| --host 0.0.0.0 | ||
| --server-port "$PORT" | ||
| --trust-remote-code | ||
| --tensor-parallel-size "$TP" | ||
| --kv_cache_dtype fp8 | ||
| --block-size 128 | ||
| --max-num-seqs "$MAX_NUM_SEQS" | ||
| --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" | ||
| --gpu-memory-utilization "$GPU_MEM_UTIL" | ||
| --index-cache-dtype fp8 | ||
| --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' | ||
| --online_quant_config '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","vision_tower","multi_modal_projector","patch_merge_mlp","*block_sparse_moe"]}' | ||
| "${SPEC_ARGS[@]}" |
There was a problem hiding this comment.
🔴 STATE_CKPT_ARGS (built from STATE_CHECKPOINT_SLOTS=96 for CONC 8-28 to reserve extra in-GPU state-checkpoint slots per the script's own comment) is computed but never appended to ATOM_CMD, so --state-checkpoint-slots is never passed to the server and ATOM_ENABLE_REPLAYSSM runs with no reserved slot room. Fix: add "${STATE_CKPT_ARGS[@]}" to the ATOM_CMD array alongside SPEC_ARGS/OFFLOAD_ARGS.
Extended reasoning...
For CONC in {8,10,12,14,15,20,24,28} the case block sets STATE_CHECKPOINT_SLOTS=96 and builds STATE_CKPT_ARGS=(--state-checkpoint-slots 96) at line 186-188. ATOM_CMD is assembled at line 282-296 and only appends "${SPEC_ARGS[@]}" and "${OFFLOAD_ARGS[@]}" (line 298-299); STATE_CKPT_ARGS is never referenced again in the file. The server therefore starts with ATOM_ENABLE_REPLAYSSM=1 but no extra reserved slots, so retained checkpoints silently compete with in-flight requests for whatever max-num-seqs leaves — exactly the failure mode the comment above the case block warns about.
Verification: normal. STATE_CKPT_ARGS is genuinely dead code — never reaches the server. At lines 186-188 the script builds STATE_CKPT_ARGS=(--state-checkpoint-slots "$STATE_CHECKPOINT_SLOTS") when STATE_CHECKPOINT_SLOTS is non-empty. grep confirms STATE_CKPT_ARGS is referenced ONLY at 186-188 and nowhere else in the file. The ATOM_CMD array (lines 282-300) appends only "${SPEC_ARGS[@]}" (298) and…
| SPEC_ARGS=( | ||
| --method eagle3 | ||
| --draft-model Inferact/MiniMax-M3-EAGLE3 | ||
| --num-speculative-tokens 3 | ||
| ) |
There was a problem hiding this comment.
🟡 (optional) SPEC_ARGS unconditionally enables eagle3 with 3 speculative tokens regardless of NUM_SPEC_TOKENS, so the CONC=40/56 case branches that set NUM_SPEC_TOKENS=0 (meant to disable the draft model past the throughput knee, per the comment above SPEC_ARGS) have no effect. Fix: gate SPEC_ARGS on NUM_SPEC_TOKENS -gt 0, omitting --method/--draft-model/--num-speculative-tokens when it is 0.
Extended reasoning...
Lines 157/167 set NUM_SPEC_TOKENS=0, SPEC_DECODE_AL=0 for CONC 40 and 56. Line 279 only echoes these values; SPEC_ARGS at lines 274-278 is a fixed array always containing --method eagle3 --draft-model Inferact/MiniMax-M3-EAGLE3 --num-speculative-tokens 3, unconditionally appended to ATOM_CMD. Any manual/TP8 run at CONC=40 or 56 (the script's header says TP8 is accepted for larger-memory variants/manual smoke tests) launches with the draft model enabled contrary to the documented design, wasting the draft forward pass past the intended throughput knee.
Verification: nit. Real dead-variable/logic defect, but it does not degrade any shipped configuration. SPEC_ARGS (lines 274-278) is a fixed array hardcoding --method eagle3 --draft-model Inferact/MiniMax-M3-EAGLE3 --num-speculative-tokens 3, appended unconditionally at line 298. Grep confirms NUM_SPEC_TOKENS is referenced only in the echo at line 279 — never used to gate/build SPEC_ARGS. Therefore the…
| @@ -0,0 +1,47 @@ | |||
| # Inference-engine patch waiver — PR #2106 | |||
There was a problem hiding this comment.
🟡 (optional) Waiver file is named after the upstream ROCm/ATOM PR (2106) rather than the InferenceX PR that introduces the patch, contradicting docs/PR_REVIEW_CHECKLIST.md item 10 which requires docs/waiver/<PR_NUMBER>.md to be named after "the PR that introduces the patch and filed in that same PR". The codeowner-signoff-verify workflow independently checks this exact rule as a merge-gating status, so a sign-off could be flagged/blocked or an automated mapping of waivers to introducing PRs would resolve to the wrong PR. Fix: rename the waiver file to match this InferenceX PR's own number once assigned (not the upstream PR being ported).
Extended reasoning...
.github/workflows/codeowner-signoff-verify.yml documents (item 10) that the sole exception to "no patches to the inference engine/serving stack" is a filled-out docs/waiver/<PR_NUMBER>.md, named after the PR that introduces the patch and filed in that same PR. Here the patch-application script (apply_atom_pr2106_patch.sh) and its invocation are introduced by this InferenceX PR, but the waiver doc is named 2106.md — the number of the upstream ROCm/ATOM PR being ported, not this InferenceX PR's number. Any tooling or reviewer following the checklist literally would look for docs/waiver/.md and not find it, since the file present maps to the wrong PR.
Verification: nit. The naming-convention violation is real. docs/PR_REVIEW_CHECKLIST.md item 10 states the sole exception to "does not patch the inference engine or serving stack" 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] ... the upstream PR/issue link" — i.e. the filename must be the…
functionstackx
changed the title
将 MiniMax ATOM AgentX launcher 与 MI355X vLLM 配置对齐：EAGLE3-GQA draft、synthetic acceptance length 2.78，并启用 thinking_mode。 Co-authored-by: Cursor <cursoragent@cursor.com>
2 participants
Summary / 摘要
This PR adds the MiniMax-M3 MXFP4 ATOM AgentX MTP recipe and applies a pinned patch derived from ROCm/ATOM PR #2106. The upstream PR is still waiting to be merged, so the patch is required temporarily until the fix is included in the ATOM image.
本 PR 增加 MiniMax-M3 MXFP4 ATOM AgentX MTP recipe，并应用源自 ROCm/ATOM PR #2106 的固定版本补丁。上游 PR 目前仍在等待合并，因此在修复进入 ATOM 镜像前需要临时应用该补丁。
Problem and fix / 问题与修复
The MiniMax-M3 EAGLE3 draft KV cache warmup could run out of memory because the draft KV pool was converted to a contiguous layout during prefill.
The patch allocates the draft KV cache in flash layout directly, avoiding the temporary full-pool
permute().contiguous()allocation that causes warmup OOM.The launcher also sets
--served-model-nameto match AIPerf requests, enables prefix caching, and uses FP8 KV/index cache.MiniMax-M3 EAGLE3 draft KV cache 在 warmup 期间可能因预填充阶段将 draft KV pool 转换为 contiguous layout 而发生显存不足。
补丁直接以 flash layout 分配 draft KV cache，避免临时执行整个 pool 的
permute().contiguous()所产生的额外内存，从而解决 warmup OOM。launcher 同时设置
--served-model-name以匹配 AIPerf 请求，启用 prefix caching，并使用 FP8 KV/index cache。Validation and performance / 验证与性能
Actions run 33707153934 completed successfully for the configured concurrency sweep.
GSM8K 5-shot accuracy passed:
flexible-extract 0.9454 ± 0.0063;strict-match 0.9462 ± 0.0062. The run usednum_concurrent=32, batch size65, andmax_gen_toks=16384.The configured search space completed successfully with no issues: TP4 conc
[1, 2, 4, 5, 8, 10, 12, 15, 20, 24, 28, 32]; TP2 conc[1, 2, 5]; no KV offloading.TP4 conc=8 warmup: 15/15 requests succeeded, 0 errors, no OOM.
Matrix tests:
127 passed.Actions run 33707153934 已成功完成配置的并发扫描。
GSM8K 5-shot 准确率验证通过：
flexible-extract 0.9454 ± 0.0063；strict-match 0.9462 ± 0.0062。运行参数为num_concurrent=32、batch size65、max_gen_toks=16384。配置中的 search space 已全部成功运行，没有问题：TP4 并发
[1, 2, 4, 5, 8, 10, 12, 15, 20, 24, 28, 32]；TP2 并发[1, 2, 5]；不开启 KV offloading。TP4 conc=8 warmup：15/15 个请求成功，0 个错误，未发生 OOM。
matrix tests：
127 passed。Follow-up / 后续工作
Remove the temporary patch application once ROCm/ATOM PR #2106 is merged and the corresponding ATOM image is available.
待 ROCm/ATOM PR #2106 合并且对应 ATOM 镜像发布后，移除临时补丁应用逻辑。
Note
Medium Risk
Runtime patching of the serving stack omits part of upstream #2106 (block-table freshness guards), and benchmark results depend on that patch until a newer ATOM image ships.
Overview
Adds
minimaxm3-fp4-mi355x-atom-agentic-mtpto the AMD master matrix: pinnedrocm/atom-devimage, agentic-coding sweeps at TP4 (conc 1–32) and TP2 (conc 1, 2, 5), all no KV offload, with MTP spec decoding.Introduces
minimaxm3_fp4_mi355x_atom_mtp.shas the vLLM-comparable ATOM arm: official MiniMax FP8 KV/index cache, EAGLE3-GQA draft with synthetic AL 2.78, prefix caching,thinking_mode, AIPerfatom:metrics, and optional LMCache hybrid offload hooks for future matrix points. Before launch it appliesapply_atom_pr2106_patch.sh, which idempotently installs a pinned ROCm/ATOM #2106 diff (draft KV flash layout / warmup OOM fix), excluding two freshness-guard files that do not apply cleanly to the image’s ATOM revision—documented indocs/waiver/2106.md.perf-changelog.yamland the sweep manifest test expectations are updated for the new config (65 manifest rows).Reviewed by Cursor Bugbot for commit 4153ec0. Bugbot is set up for automated code reviews on this repo. Configure here.