From c3c316e1e6ec71ad52af05d4e33511235404139b Mon Sep 17 00:00:00 2001
From: Simba Zhang <solderzzc@gmail.com>
Date: Sat, 29 Aug 2026 06:40:33 -0700
Subject: [PATCH] fix: wire up Gemma4MTPBench target + repair regressed
 mlx-swift-lm/mlx-swift submodule pins
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

mlx-swift-lm was uncommitted-pinned backward past the commit that introduced
DualModelMTP/MTPTokenIterator/Gemma4AssistantModel, breaking
Gemma4MTPBench and forcing MTP speculative decoding to be commented out in
InferenceEngine.swift as a workaround. Repointing both submodules to their
current origin/main tips (mlx-swift-lm past the merged DSA stage-2 PR #61,
mlx-swift past the matching MLXFast.fromFp8 addition) restores those types
and lets the benchmark build again — no source workaround needed.

Also adds scripts/bootstrap_local_tests.sh, which mirrors CI's "Install MLX
Metal library" step (pip install mlx, copy its bundled metallib into every
built .xctest bundle) so `swift test` is runnable locally without the
manual cmake+make dance. Addresses the Tier 3 item in #128.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 Package.swift                    | 15 ++++++++++++++-
 mlx-swift                        |  2 +-
 scripts/bootstrap_local_tests.sh | 33 ++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100755 scripts/bootstrap_local_tests.sh

diff --git a/Package.swift b/Package.swift
index 70c4e01..8649fcd 100644
--- a/Package.swift
+++ b/Package.swift
@@ -9,7 +9,8 @@ let package = Package(
         .library(name: "DFlash", targets: ["DFlash"]),
         .executable(name: "SwiftLM", targets: ["SwiftLM"]),
         .executable(name: "SwiftBuddy", targets: ["SwiftBuddy"]),
-        .executable(name: "DFlashKernelBench", targets: ["DFlashKernelBench"])
+        .executable(name: "DFlashKernelBench", targets: ["DFlashKernelBench"]),
+        .executable(name: "Gemma4MTPBench", targets: ["Gemma4MTPBench"])
     ],
     dependencies: [
         // Local Apple MLX Swift fork for C++ extensions
@@ -53,6 +54,18 @@ let package = Package(
             ],
             path: "Sources/DFlashKernelBench"
         ),
+        // ── Gemma4 MTP Speculative Decoding Benchmark ───────────────
+        .executableTarget(
+            name: "Gemma4MTPBench",
+            dependencies: [
+                "MLXInferenceCore",
+                .product(name: "MLX", package: "mlx-swift"),
+                .product(name: "MLXLLM", package: "mlx-swift-lm"),
+                .product(name: "MLXLMCommon", package: "mlx-swift-lm"),
+                .product(name: "ArgumentParser", package: "swift-argument-parser"),
+            ],
+            path: "Sources/Gemma4MTPBench"
+        ),
         // ── STFT Audio Profiling Testing Script (macOS only) ───────────
         .executableTarget(
             name: "SwiftLMTestSTFT",
diff --git a/mlx-swift b/mlx-swift
index 133864c..5639a6d 160000
--- a/mlx-swift
+++ b/mlx-swift
@@ -1 +1 @@
-Subproject commit 133864c733c8d4178547f8fe92897da6a788368f
+Subproject commit 5639a6d9e6a7ab785e102d88d741879f529fce56
diff --git a/scripts/bootstrap_local_tests.sh b/scripts/bootstrap_local_tests.sh
new file mode 100755
index 0000000..f3770c6
--- /dev/null
+++ b/scripts/bootstrap_local_tests.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+# Makes `swift test` runnable locally without CI's help.
+#
+# A bare `swift test` aborts with "Failed to load the default metallib"
+# because Package.swift links MLX but nothing on a local machine ever builds
+# or installs mlx.metallib. CI works around this in .github/workflows/ci.yml
+# ("Install MLX Metal library" step) by pip-installing the `mlx` wheel and
+# copying its bundled metallib into every built .xctest bundle. This script
+# does the same thing locally.
+set -eo pipefail
+
+VENV_DIR="${MLX_METALLIB_VENV:-/tmp/swiftlm_mlx_venv}"
+
+echo "=> Building test harness (swift build --build-tests)..."
+swift build --build-tests
+
+echo "=> Installing MLX Metal library..."
+if [ ! -d "$VENV_DIR" ]; then
+    python3 -m venv "$VENV_DIR"
+fi
+"$VENV_DIR/bin/pip" install --quiet --upgrade mlx
+
+METALLIB=$(find "$VENV_DIR" -name "mlx.metallib" | head -1)
+if [ -z "$METALLIB" ]; then
+    echo "error: mlx.metallib not found after pip install mlx" >&2
+    exit 1
+fi
+
+cp "$METALLIB" .build/debug/ 2>/dev/null || true
+cp "$METALLIB" .build/release/ 2>/dev/null || true
+find .build -type d -name "MacOS" -exec cp "$METALLIB" {}/ \;
+
+echo "=> Done. Run tests with: swift test --skip-build"