Make MLX metallib generation automatic and stale libraries loudly detectable

Sources/MereRunCore/Support/MLXRuntimeVersion.swift

@@ -0,0 +1,26 @@
+#if canImport(Cmlx)
+import Cmlx
+#endif
+
+/// Version of the vendored MLX core this binary was compiled against.
+///
+/// Used to validate the AOT Metal kernel library (`default.metallib`) at
+/// startup: the metallib is built separately from `swift build` (see
+/// `scripts/build_mlx_metallib.sh`) and carries a version stamp; loading a
+/// library built from a different mlx core silently corrupts inference.
+public enum MLXRuntimeVersion {
+    /// The mlx core version string (e.g. "0.31.1"), or nil when the MLX C
+    /// runtime is not linked (Linux prebuilt MLX builds).
+    public static var coreVersion: String? {
+        #if canImport(Cmlx)
+        var str = mlx_string_new()
+        defer { mlx_string_free(str) }
+        guard mlx_version(&str) == 0, let data = mlx_string_data(str) else {
+            return nil
+        }
+        return String(cString: data)
+        #else
+        return nil
+        #endif
+    }
+}

Tests/MereRunCoreTests/MLXTestSupport.swift

@@ -75,7 +75,14 @@ enum MLXTestSupport {
         let binaryDir = executableURL.deletingLastPathComponent()
         guard fileManager.isWritableFile(atPath: binaryDir.path) else { return }
         let destination = binaryDir.appendingPathComponent("mlx.metallib")
-        guard !fileManager.fileExists(atPath: destination.path) else { return }
+        // Replace rather than skip: a leftover link/copy from an earlier run
+        // can point at a stale metallib, which silently corrupts kernels
+        // (see scripts/build_mlx_metallib.sh). A symlink already pointing at
+        // the current source is left alone so steady-state runs are no-ops.
+        if symlinkTarget(of: destination, fileManager: fileManager) == sourceMetallib.path {
+            return
+        }
+        try? fileManager.removeItem(at: destination)
 
         // Best-effort + race-safe: if concurrent tests call this at the same time,
         // one may succeed and the other will see "file exists" errors, which we ignore.
@@ -126,7 +133,15 @@ enum MLXTestSupport {
         }
 
         let destinationBundle = resourcesDir.appendingPathComponent(sourceBundle.lastPathComponent, isDirectory: true)
-        guard !fileManager.fileExists(atPath: destinationBundle.path) else { return }
+        // Replace rather than skip: a leftover symlink/copy from an earlier run
+        // can point at a stale bundle (this exact mechanism pinned tests to a
+        // pre-0.30 metallib and produced NaN attention past 1024 keys). A
+        // symlink already pointing at the current source is left alone so
+        // steady-state runs are no-ops.
+        if symlinkTarget(of: destinationBundle, fileManager: fileManager) == sourceBundle.path {
+            return
+        }
+        try? fileManager.removeItem(at: destinationBundle)
 
         // MLX's SwiftPM lookup expects `<someBundle>/Contents/Resources/<SWIFTPM_BUNDLE>.bundle/...`.
         // Symlink (or copy) the entire Cmlx bundle into the test bundle's Resources directory.
@@ -149,6 +164,16 @@ enum MLXTestSupport {
         }
     }
 
+    private static func symlinkTarget(of url: URL, fileManager: FileManager) -> String? {
+        guard let target = try? fileManager.destinationOfSymbolicLink(atPath: url.path) else {
+            return nil
+        }
+        if target.hasPrefix("/") {
+            return target
+        }
+        return url.deletingLastPathComponent().appendingPathComponent(target).standardizedFileURL.path
+    }
+
     private static func resolveXCTestBundleURL(executableURL: URL, arguments: [String]) -> URL? {
         // When invoked directly from the test bundle executable:
         //   .../<name>.xctest/Contents/MacOS/<exe>

Tests/MereRunCoreTests/SDPAVectorKernelTests.swift

@@ -0,0 +1,98 @@
+import Foundation
+import MLX
+import MLXFast
+import MLXRandom
+import XCTest
+
+/// Regression tests for the fused single-token-decode attention kernels.
+///
+/// On 2026-07-03 a stale AOT metallib (built from pre-0.30 mlx sources but
+/// loaded by the mlx 0.31.x host dispatch) made `MLXFast.scaledDotProductAttention`
+/// return garbage, nondeterministically, whenever the key length crossed 1024 —
+/// the point where the Metal backend switches from the 1-pass to the 2-pass
+/// sdpa_vector kernel on 'd'/'s'-class Apple GPUs. Every MLX text runtime
+/// (Q35, Gemma4, LFM2, Psi) produced incoherent output past ~1024 tokens of
+/// context while short prompts stayed clean, so nothing short caught it.
+///
+/// These tests pin the exact failing shape (Q35 decode: q [1,16,1,256] bf16,
+/// GQA 16:2, cache-style strided k/v views) on both sides of the 1024
+/// boundary, comparing against an unfused fp32 reference and requiring
+/// bit-identical results across repeated runs.
+final class SDPAVectorKernelTests: MereRunCoreTestCase {
+    private func skipUnlessGPU() throws {
+        guard Device.defaultDevice().deviceType == .gpu else {
+            throw XCTSkip("SDPA vector kernels are Metal-only; set MERERUN_TEST_MLX_DEVICE=gpu to run them.")
+        }
+    }
+
+    /// Unfused fp32 attention with GQA expansion.
+    private func referenceAttention(
+        queries: MLXArray,
+        keys: MLXArray,
+        values: MLXArray,
+        scale: Float
+    ) -> MLXArray {
+        let repeats = queries.dim(1) / keys.dim(1)
+        let expandedKeys = MLX.repeated(keys.asType(.float32), count: repeats, axis: 1)
+        let expandedValues = MLX.repeated(values.asType(.float32), count: repeats, axis: 1)
+        let scores = MLX.matmul(queries.asType(.float32), expandedKeys.transposed(0, 1, 3, 2)) * scale
+        return MLX.matmul(MLX.softmax(scores, axis: -1), expandedValues)
+    }
+
+    func testSingleTokenDecodeAcross1024Boundary() throws {
+        try skipUnlessGPU()
+
+        let headDim = 1 << 8
+        let scale = 1.0 / Float(Double(headDim).squareRoot())
+
+        // 1023 stays on the 1-pass kernel (control); 1100 and 3300 engage the
+        // 2-pass kernel that the stale metallib corrupted.
+        for keyLength in [1023, 1100, 3300] {
+            MLXRandom.seed(0)
+            let queries = (MLXRandom.normal([1, 16, 1, headDim]) * 0.5).asType(.bfloat16)
+
+            // Mirror KVCacheSimple: a step-256 preallocated buffer sliced back
+            // to the live length, so k/v reach the kernel as strided views.
+            let capacity = ((keyLength + 255) / 256) * 256
+            let keyBuffer = (MLXRandom.normal([1, 2, capacity, headDim]) * 0.5).asType(.bfloat16)
+            let valueBuffer = (MLXRandom.normal([1, 2, capacity, headDim]) * 0.5).asType(.bfloat16)
+            let keys = keyBuffer[0..., 0..., ..<keyLength, 0...]
+            let values = valueBuffer[0..., 0..., ..<keyLength, 0...]
+
+            var outputs: [MLXArray] = []
+            for _ in 0..<3 {
+                let out = MLXFast.scaledDotProductAttention(
+                    queries: queries,
+                    keys: keys,
+                    values: values,
+                    scale: scale,
+                    mask: .none
+                )
+                MLX.eval(out)
+                outputs.append(out.asType(.float32))
+            }
+
+            let reference = referenceAttention(queries: queries, keys: keys, values: values, scale: scale)
+            MLX.eval(reference)
+
+            let referenceError = MLX.abs(outputs[0] - reference).max().item(Float.self)
+            XCTAssertFalse(
+                referenceError.isNaN,
+                "SDPA produced NaN at keyLength=\(keyLength) — broken/stale metallib kernel?"
+            )
+            // bf16 rounding lands around 2e-4; kernel corruption is O(1) or NaN.
+            XCTAssertLessThanOrEqual(
+                referenceError, 0.02,
+                "SDPA diverged from unfused reference at keyLength=\(keyLength)"
+            )
+
+            for run in 1..<outputs.count {
+                let drift = MLX.abs(outputs[run] - outputs[0]).max().item(Float.self)
+                XCTAssertEqual(
+                    drift, 0,
+                    "SDPA nondeterministic at keyLength=\(keyLength) (run \(run) differs) — broken/stale metallib kernel?"
+                )
+            }
+        }
+    }
+}

scripts/build_mere_run_app.sh

@@ -41,6 +41,12 @@ fi
 swift "${swift_app_args[@]}"
 swift "${swift_cli_args[@]}"
 
+# Regenerate + stamp the MLX Metal kernel library from the current checkout.
+# `swift build` never compiles the .metal sources, and shipping a stale
+# leftover metallib silently corrupts inference (gibberish, nondeterministic
+# generation past ~1024 tokens of context). See scripts/build_mlx_metallib.sh.
+"${repo_root}/scripts/build_mlx_metallib.sh" --configuration "$configuration"
+
 build_dir="$(swift "${swift_bin_path_args[@]}")"
 executable="${build_dir}/mere.run.app"
 cli_executable="${build_dir}/mere.run"
@@ -76,11 +82,14 @@ if [[ -d "${repo_root}/skills/use-mere-run" ]]; then
   cp -R "${repo_root}/skills/use-mere-run" "${resources}/skills/use-mere-run"
 fi
 
-# Frameworks/bundles co-located beside the CLI so its @executable_path rpath resolves.
+# Frameworks co-located beside the CLI so its @executable_path rpath resolves.
+# MLX's resource-only .bundle is not embedded here because stricter codesign
+# treats unsigned nested .bundle directories under Helpers as invalid code.
+# The stamped flat Resources/default.metallib layout is enough for runtime
+# lookup and is verified below.
 for asset in \
   "${build_dir}/llama.framework" \
   "${build_dir}/magentart.framework" \
-  "${build_dir}/mlx-swift_Cmlx.bundle" \
   "${build_dir}/Resources"
 do
   if [[ -e "$asset" ]]; then
@@ -95,6 +104,10 @@ if [[ -d "${repo_root}/vendor/ds4" ]]; then
   cp -R "${repo_root}/vendor/ds4" "${cli_payload}/vendor/ds4"
 fi
 
+# Refuse to ship a metallib whose stamp doesn't match the checkout it was
+# supposedly built from.
+"${repo_root}/scripts/build_mlx_metallib.sh" --verify-only "${cli_payload}/Resources"
+
 plutil -create xml1 "${contents}/Info.plist"
 plutil -insert CFBundleExecutable -string "mere.run.app" "${contents}/Info.plist"
 plutil -insert CFBundleIdentifier -string "run.mere.MereRunApp" "${contents}/Info.plist"
@@ -139,8 +152,15 @@ sign() {
   codesign "${args[@]}" "$@"
 }
 
-# 1. Co-located frameworks/bundles — no entitlements.
+# 1. Co-located executable frameworks/bundles — no entitlements. SwiftPM also
+#    places resource-only .bundle directories here; those have no executable
+#    to sign and are sealed as bundle resources in step 3.
 while IFS= read -r -d '' asset; do
+  if [[ "$asset" == *.bundle ]]; then
+    executable_name="$(/usr/libexec/PlistBuddy -c 'Print :CFBundleExecutable' \
+      "${asset}/Contents/Info.plist" 2>/dev/null || true)"
+    [[ -n "$executable_name" && -e "${asset}/Contents/MacOS/${executable_name}" ]] || continue
+  fi
   sign "" "$asset"
 done < <(find "$helpers" \( -name '*.framework' -o -name '*.bundle' \) -prune -print0)