[wip - contains patches that need to be removed] [amd agentx m3 atom]

benchmarks/single_node/agentic/apply_atom_pr2106_patch.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Apply ROCm/ATOM PR #2106 to the pinned image before serving.
+# The patch fixes the MiniMax-M3 EAGLE3 draft KV layout and block-table
+# freshness issues that can cause prefix-cache warmup OOMs or stale-KV reads.
+
+ATOM_ROOT="${ATOM_ROOT:-/app/ATOM}"
+# Pin the PR's current base/head range so a later PR update cannot silently
+# change an already reviewed benchmark run.
+PATCH_URL="https://github.com/ROCm/ATOM/compare/86476f716e9887a7cb2e423deb88712737f702f5...c63ab67e58d39f1c887d2b6af13ac0b9034a1a1a.diff"
+PATCH_FILE="$(mktemp /tmp/atom-pr2106.XXXXXX.patch)"
+trap 'rm -f "$PATCH_FILE"' EXIT
+PATCH_EXCLUDES=(
+    --exclude=atom/model_ops/attentions/backends.py
+    --exclude=atom/spec_decode/eagle_proposer.py
+)
+
+if [[ ! -d "$ATOM_ROOT" ]]; then
+    echo "ERROR: ATOM source tree not found at $ATOM_ROOT" >&2
+    exit 1
+fi
+
+curl -fsSL "$PATCH_URL" -o "$PATCH_FILE"
+
+if git -C "$ATOM_ROOT" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+    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"
+else
+    if [[ -f "$ATOM_ROOT/atom/spec_decode/draft_kv_layout.py" ]]; then
+        echo "ATOM PR #2106 already applied"
+        exit 0
+    fi
+    if patch --dry-run -p1 -d "$ATOM_ROOT" < "$PATCH_FILE" >/dev/null 2>&1; then
+        patch -p1 -d "$ATOM_ROOT" < "$PATCH_FILE"
+    else
+        echo "ERROR: ATOM PR #2106 does not apply cleanly to $ATOM_ROOT" >&2
+        exit 1
+    fi
+fi
+echo "Applied ROCm/ATOM PR #2106 to $ATOM_ROOT"

benchmarks/single_node/agentic/apply_atom_pr2147_patch.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Apply ROCm/ATOM PR #2147 on top of the pinned MiniMax image checkout.
+#
+# PR #2147 supersedes PR #2106 for the MiniMax-M3 EAGLE3 draft-KV layout work:
+# both fix the draft/target KV pool disagreement that can OOM during prefix-cache
+# warmup. The pinned image predates #2147's merge base, so we fast-forward the
+# in-image tree to the current PR head rather than applying the PR diff alone.
+#
+# A small MiniMax EAGLE3 follow-up fix is applied after the forward diff because
+# the upstream check compared kernel block sizes (16 vs 128) instead of the
+# scheduler block size both sides actually index (128).
+
+ATOM_ROOT="${ATOM_ROOT:-/app/ATOM}"
+PINNED_ATOM_SHA="00760297ef69af7ab5d345af9c8fc6da00f5314d"
+PR2147_HEAD_SHA="71ad8c69c5377738657907e24c5aff8672c0f004"
+PATCH_URL="https://github.com/ROCm/ATOM/compare/${PINNED_ATOM_SHA}...${PR2147_HEAD_SHA}.diff"
+PATCH_FILE="$(mktemp /tmp/atom-pr2147.XXXXXX.patch)"
+MARKER_FILE="$ATOM_ROOT/.inferencex-pr2147-applied"
+DRAFT_KV_FILE="$ATOM_ROOT/atom/spec_decode/draft_kv.py"
+
+if [[ ! -d "$ATOM_ROOT" ]]; then
+    echo "ERROR: ATOM source tree not found at $ATOM_ROOT" >&2
+    exit 1
+fi
+
+trap 'rm -f "$PATCH_FILE"' EXIT
+
+if [[ -f "$MARKER_FILE" ]]; then
+    echo "ATOM PR #2147 already applied ($(cat "$MARKER_FILE"))"
+    exit 0
+fi
+
+curl -fsSL "$PATCH_URL" -o "$PATCH_FILE"
+git -C "$ATOM_ROOT" apply --check "$PATCH_FILE"
+git -C "$ATOM_ROOT" apply "$PATCH_FILE"
+
+python3 - "$DRAFT_KV_FILE" <<'PY'
+from pathlib import Path
+import sys
+
+path = Path(sys.argv[1])
+text = path.read_text()
+old = """        if pool.block_size != target.block_size:
+            raise ValueError(
+                f"the draft's blocks are {pool.block_size} tokens and the "
+                f"target's {target.block_size}; a draft is indexed by the "
+                "target's block tables, so it cannot yet block at anything else"
+            )"""
+new = """        pool_scheduler_block = pool.block_size * (
+            runner.block_size // pool.block_size
+        )
+        target_scheduler_block = target.block_size * target.block_ratio
+        if pool_scheduler_block != target_scheduler_block:
+            raise ValueError(
+                f"the draft indexes {pool_scheduler_block}-token scheduler "
+                f"blocks but the target indexes {target_scheduler_block}-token "
+                "scheduler blocks; a draft is indexed by the target's block "
+                "tables, so the two must agree at scheduler granularity"
+            )"""
+if old not in text:
+    if "pool_scheduler_block" in text:
+        print("MiniMax EAGLE3 scheduler-block fix already present")
+    else:
+        raise SystemExit(
+            "ERROR: expected draft_kv.py block-size guard not found after PR #2147 apply"
+        )
+else:
+    path.write_text(text.replace(old, new, 1))
+    print("Applied MiniMax EAGLE3 scheduler-block alignment fix to draft_kv.py")
+PY
+
+printf '%s\n' "$PR2147_HEAD_SHA" > "$MARKER_FILE"
+echo "Applied ROCm/ATOM PR #2147 (${PINNED_ATOM_SHA} -> ${PR2147_HEAD_SHA}) to $ATOM_ROOT"