diff --git a/Sources/MereRunCLI/Support/MLXBundleSupport.swift b/Sources/MereRunCLI/Support/MLXBundleSupport.swift index b9d94a0e..8ce5b1c3 100644 --- a/Sources/MereRunCLI/Support/MLXBundleSupport.swift +++ b/Sources/MereRunCLI/Support/MLXBundleSupport.swift @@ -3,51 +3,72 @@ import ArgumentParser import Darwin #endif import Foundation +import MereRunCore enum MLXBundleSupport { #if os(Linux) static func ensureAvailable(quiet _: Bool) throws {} #else + private static let bundleName = "mlx-swift_Cmlx.bundle" + private static let stampName = "default.metallib.version" + + /// Relationship between a bundle's metallib version stamp + /// (`default.metallib.version`, written by scripts/build_mlx_metallib.sh) + /// and the mlx core version this binary was compiled against. `swift build` + /// never regenerates the metallib, so a leftover library from an older + /// mlx-swift silently corrupts inference — the stamp is the only guard. + enum MetallibStamp { + case matched + case unstamped + case mismatched(stamped: String, expected: String) + /// The runtime mlx version is unavailable (e.g. prebuilt MLX); no + /// judgement is possible. + case unvalidatable + } + static func ensureAvailable(quiet: Bool) throws { - let bundleName = "mlx-swift_Cmlx.bundle" let fm = FileManager.default let execDir = executableDirectory() let destBundleURL = execDir.appendingPathComponent(bundleName, isDirectory: true) - if fm.fileExists(atPath: destBundleURL.path) { - try installCompatibilityMetallibs(bundleURL: destBundleURL, executableDir: execDir) + if hasMetallib(destBundleURL) { + try validateAndFinish( + bundleURL: destBundleURL, + executableDir: execDir, + quiet: quiet, + justInstalled: false + ) return } - let installCandidates = bundleInstallCandidates(executableDir: execDir) - for candidateDir in installCandidates { - let sourceBundleURL = candidateDir.appendingPathComponent(bundleName, isDirectory: true) - guard fm.fileExists(atPath: sourceBundleURL.path) else { - continue - } + let flatResourcesURL = execDir.appendingPathComponent("Resources", isDirectory: true) + if hasMetallib(resourcesURL: flatResourcesURL) { + try validateFlatResources( + resourcesURL: flatResourcesURL, + executableDir: execDir, + quiet: quiet + ) + return + } + // No usable bundle next to the executable (missing entirely, or a + // husk without a metallib inside): copy one in, preferring candidates + // whose stamp matches this binary. + if let source = locateCandidateBundle(executableDir: execDir) { try? fm.removeItem(at: destBundleURL) - try fm.copyItem(at: sourceBundleURL, to: destBundleURL) - try installCompatibilityMetallibs(bundleURL: destBundleURL, executableDir: execDir) + try fm.copyItem(at: source, to: destBundleURL) if !quiet { - CLIStderr.write("[mererun] Installed \(bundleName) for mlx-swift Metal shaders.\n") + CLIStderr.write("[mererun] Installed \(bundleName) for mlx-swift Metal shaders (from \(source.path)).\n") } + try validateAndFinish( + bundleURL: destBundleURL, + executableDir: execDir, + quiet: quiet, + justInstalled: true + ) return } - let rootCandidates = workspaceRootCandidates(executableDir: execDir) - for root in rootCandidates { - if let sourceBundleURL = findMlxSwiftBundle(workspaceRoot: root) { - try? fm.removeItem(at: destBundleURL) - try fm.copyItem(at: sourceBundleURL, to: destBundleURL) - try installCompatibilityMetallibs(bundleURL: destBundleURL, executableDir: execDir) - if !quiet { - CLIStderr.write("[mererun] Installed \(bundleName) for mlx-swift Metal shaders.\n") - } - return - } - } - if ProcessInfo.processInfo.environment["DYLD_FRAMEWORK_PATH"] != nil { if !quiet { CLIStderr.write("[mererun] warning: \(bundleName) not found near executable; relying on DYLD_FRAMEWORK_PATH.\n") @@ -57,22 +78,248 @@ enum MLXBundleSupport { throw ValidationError( """ - Missing mlx-swift Metal shaders (\(bundleName)). + Missing mlx-swift Metal shaders (\(bundleName) or Resources/default.metallib). + + `swift build` does not generate the Metal kernel library. Build a + stamped one from the current mlx-swift checkout, then rerun: + + scripts/build_mlx_metallib.sh - Build the package once to generate the metallib bundle, then rerun: - swift build + (Run `swift package resolve` first if .build/checkouts/mlx-swift + is missing.) """ ) } + private static func validateFlatResources( + resourcesURL: URL, + executableDir: URL, + quiet _: Bool + ) throws { + switch stampStatus(resourcesURL: resourcesURL) { + case .matched, .unvalidatable: + return + + case .unstamped: + CLIStderr.write( + """ + [mererun] WARNING: the MLX Metal shader library has no version stamp: + [mererun] \(resourcesURL.path) + [mererun] Its provenance is unknown, so it may not match this binary\(expectedVersionSuffix()). + [mererun] A stale metallib silently corrupts output (gibberish, nondeterministic + [mererun] generation past ~1024 tokens of context). Rebuild and stamp it with: + [mererun] scripts/build_mlx_metallib.sh + + """ + ) + + case .mismatched(let stamped, let expected): + if let replacement = locateCandidateBundle(executableDir: executableDir, matchedOnly: true) { + try installCompatibilityMetallibs(bundleURL: replacement, executableDir: executableDir) + CLIStderr.write("[mererun] Replaced stale MLX metallib (built for mlx \(stamped), this binary needs \(expected)) with matching copy from \(replacement.path).\n") + return + } + + if mismatchOverrideEnabled { + CLIStderr.write("[mererun] WARNING: running with a mismatched MLX metallib (built for mlx \(stamped), this binary needs \(expected)) because MERERUN_ALLOW_METALLIB_MISMATCH=1. Expect corrupted output.\n") + return + } + + throw ValidationError( + """ + Stale MLX Metal shader library. + + metallib built for mlx core: \(stamped) + this binary requires: \(expected) + resources: \(resourcesURL.path) + + A mismatched metallib produces silently corrupted inference + (gibberish, nondeterministic generation past ~1024 tokens of + context). Rebuild it from the current checkout: + + scripts/build_mlx_metallib.sh + + Emergency override (unsafe): MERERUN_ALLOW_METALLIB_MISMATCH=1 + """ + ) + } + } + + private static func validateAndFinish( + bundleURL: URL, + executableDir: URL, + quiet: Bool, + justInstalled: Bool + ) throws { + switch stampStatus(bundleURL: bundleURL) { + case .matched, .unvalidatable: + try installCompatibilityMetallibs(bundleURL: bundleURL, executableDir: executableDir) + + case .unstamped: + CLIStderr.write( + """ + [mererun] WARNING: the MLX Metal shader library has no version stamp: + [mererun] \(bundleURL.path) + [mererun] Its provenance is unknown, so it may not match this binary\(expectedVersionSuffix()). + [mererun] A stale metallib silently corrupts output (gibberish, nondeterministic + [mererun] generation past ~1024 tokens of context). Rebuild and stamp it with: + [mererun] scripts/build_mlx_metallib.sh + + """ + ) + try installCompatibilityMetallibs(bundleURL: bundleURL, executableDir: executableDir) + + case .mismatched(let stamped, let expected): + if !justInstalled, + let replacement = locateCandidateBundle(executableDir: executableDir, matchedOnly: true) { + let fm = FileManager.default + try? fm.removeItem(at: bundleURL) + try fm.copyItem(at: replacement, to: bundleURL) + CLIStderr.write("[mererun] Replaced stale MLX metallib (built for mlx \(stamped), this binary needs \(expected)) with matching copy from \(replacement.path).\n") + try installCompatibilityMetallibs(bundleURL: bundleURL, executableDir: executableDir) + return + } + + if mismatchOverrideEnabled { + CLIStderr.write("[mererun] WARNING: running with a mismatched MLX metallib (built for mlx \(stamped), this binary needs \(expected)) because MERERUN_ALLOW_METALLIB_MISMATCH=1. Expect corrupted output.\n") + try installCompatibilityMetallibs(bundleURL: bundleURL, executableDir: executableDir) + return + } + + throw ValidationError( + """ + Stale MLX Metal shader library. + + metallib built for mlx core: \(stamped) + this binary requires: \(expected) + bundle: \(bundleURL.path) + + A mismatched metallib produces silently corrupted inference + (gibberish, nondeterministic generation past ~1024 tokens of + context). Rebuild it from the current checkout: + + scripts/build_mlx_metallib.sh + + Emergency override (unsafe): MERERUN_ALLOW_METALLIB_MISMATCH=1 + """ + ) + } + } + + private static var mismatchOverrideEnabled: Bool { + let raw = (ProcessInfo.processInfo.environment["MERERUN_ALLOW_METALLIB_MISMATCH"] ?? "").lowercased() + return raw == "1" || raw == "true" || raw == "yes" + } + + private static func expectedVersionSuffix() -> String { + guard let expected = MLXRuntimeVersion.coreVersion.flatMap(normalizedVersion) else { + return "" + } + return " (mlx core \(expected))" + } + + static func stampStatus(bundleURL: URL) -> MetallibStamp { + let resourcesURL = bundleURL + .appendingPathComponent("Contents", isDirectory: true) + .appendingPathComponent("Resources", isDirectory: true) + return stampStatus(resourcesURL: resourcesURL) + } + + private static func stampStatus(resourcesURL: URL) -> MetallibStamp { + guard let expected = MLXRuntimeVersion.coreVersion.flatMap(normalizedVersion) else { + return .unvalidatable + } + let stampURL = resourcesURL.appendingPathComponent(stampName, isDirectory: false) + guard let contents = try? String(contentsOf: stampURL, encoding: .utf8), + let stamped = stampField("mlx-core-version", in: contents).flatMap(normalizedVersion) else { + return .unstamped + } + return stamped == expected ? .matched : .mismatched(stamped: stamped, expected: expected) + } + + private static func stampField(_ key: String, in contents: String) -> String? { + for line in contents.split(separator: "\n") { + let parts = line.split(separator: ":", maxSplits: 1) + guard parts.count == 2, + parts[0].trimmingCharacters(in: .whitespaces) == key else { + continue + } + return parts[1].trimmingCharacters(in: .whitespaces) + } + return nil + } + + /// Reduces a version string to its MAJOR.MINOR.PATCH numeric prefix, so + /// dev-build suffixes ("0.31.1.dev20260115+abc") still compare equal. + private static func normalizedVersion(_ raw: String) -> String? { + let numeric = raw.trimmingCharacters(in: .whitespaces).prefix { $0.isNumber || $0 == "." } + let parts = numeric.split(separator: ".").prefix(3) + guard !parts.isEmpty else { + return nil + } + return parts.joined(separator: ".") + } + + private static func hasMetallib(_ bundleURL: URL) -> Bool { + let resourcesURL = bundleURL + .appendingPathComponent("Contents", isDirectory: true) + .appendingPathComponent("Resources", isDirectory: true) + return hasMetallib(resourcesURL: resourcesURL) + } + + private static func hasMetallib(resourcesURL: URL) -> Bool { + let metallibURL = resourcesURL.appendingPathComponent("default.metallib", isDirectory: false) + return FileManager.default.fileExists(atPath: metallibURL.path) + } + + /// Finds a bundle to copy next to the executable. Preference order: + /// stamp-matched candidates, then unstamped ones, then mismatched ones + /// (`matchedOnly` restricts to the first tier, for self-healing). + private static func locateCandidateBundle(executableDir: URL, matchedOnly: Bool = false) -> URL? { + var firstUnstamped: URL? + var firstMismatched: URL? + + for candidate in candidateBundleURLs(executableDir: executableDir) { + guard hasMetallib(candidate) else { + continue + } + switch stampStatus(bundleURL: candidate) { + case .matched: + return candidate + case .unvalidatable: + // No runtime version to compare against; without a better + // signal the first available bundle wins (legacy behaviour). + return matchedOnly ? nil : candidate + case .unstamped: + if firstUnstamped == nil { firstUnstamped = candidate } + case .mismatched: + if firstMismatched == nil { firstMismatched = candidate } + } + } + + return matchedOnly ? nil : (firstUnstamped ?? firstMismatched) + } + + private static func candidateBundleURLs(executableDir: URL) -> [URL] { + var urls: [URL] = [] + for candidateDir in bundleInstallCandidates(executableDir: executableDir) { + urls.append(candidateDir.appendingPathComponent(bundleName, isDirectory: true)) + } + for root in workspaceRootCandidates(executableDir: executableDir) { + urls.append(contentsOf: findMlxSwiftBundles(workspaceRoot: root)) + } + var seen = Set() + return urls.filter { seen.insert($0.standardizedFileURL.path).inserted } + } + private static func installCompatibilityMetallibs(bundleURL: URL, executableDir: URL) throws { let fm = FileManager.default - let defaultMetallibURL = bundleURL + let bundleResources = bundleURL .appendingPathComponent("Contents", isDirectory: true) .appendingPathComponent("Resources", isDirectory: true) - .appendingPathComponent("default.metallib", isDirectory: false) + let sourceMetallib = bundleResources.appendingPathComponent("default.metallib", isDirectory: false) - guard fm.fileExists(atPath: defaultMetallibURL.path) else { + guard fm.fileExists(atPath: sourceMetallib.path) else { return } @@ -84,23 +331,57 @@ enum MLXBundleSupport { resourcesDir.appendingPathComponent("mlx.metallib", isDirectory: false), resourcesDir.appendingPathComponent("default.metallib", isDirectory: false), ] - for destination in destinations { + try refreshCopy(from: sourceMetallib, to: destination) + } + + let sourceStamp = bundleResources.appendingPathComponent(stampName, isDirectory: false) + if fm.fileExists(atPath: sourceStamp.path) { + try? refreshCopy(from: sourceStamp, to: resourcesDir.appendingPathComponent(stampName, isDirectory: false)) + } + } + + /// Copies `source` over `destination` unless the destination already has + /// the same size. (The previous skip-if-exists behaviour let stale flat + /// copies outlive bundle updates forever.) The copy goes through a temp + /// file + rename so concurrent CLI startups never observe a torn file. + private static func refreshCopy(from source: URL, to destination: URL) throws { + let fm = FileManager.default + if let sourceSize = fileSize(source), + let destinationSize = fileSize(destination), + sourceSize == destinationSize { + return + } + + let tmp = destination + .deletingLastPathComponent() + .appendingPathComponent(".\(destination.lastPathComponent).tmp-\(getpid())", isDirectory: false) + try? fm.removeItem(at: tmp) + do { + try fm.copyItem(at: source, to: tmp) if fm.fileExists(atPath: destination.path) { - continue + _ = try fm.replaceItemAt(destination, withItemAt: tmp) + } else { + try fm.moveItem(at: tmp, to: destination) } - try? fm.removeItem(at: destination) - do { - try fm.copyItem(at: defaultMetallibURL, to: destination) - } catch { - if fm.fileExists(atPath: destination.path) { - continue - } - throw error + } catch { + try? fm.removeItem(at: tmp) + // Tolerate races with a concurrent startup that produced the file. + if fm.fileExists(atPath: destination.path) { + return } + throw error } } + private static func fileSize(_ url: URL) -> Int64? { + guard let attributes = try? FileManager.default.attributesOfItem(atPath: url.path), + let size = attributes[.size] as? Int64 else { + return nil + } + return size + } + private static func executableDirectory() -> URL { // `Bundle.main.executableURL` / `argv[0]` can reflect the *invocation* path (e.g. a symlink in // `/usr/local/bin`). `proc_pidpath` gives us the actual on-disk executable path. @@ -210,12 +491,13 @@ enum MLXBundleSupport { return candidates } - private static func findMlxSwiftBundle(workspaceRoot: URL) -> URL? { + private static func findMlxSwiftBundles(workspaceRoot: URL) -> [URL] { let fm = FileManager.default + var found: [URL] = [] let vendoredBundle = workspaceRoot.appendingPathComponent("vendor/mlx-swift_Cmlx.bundle", isDirectory: true) if fm.fileExists(atPath: vendoredBundle.path) { - return vendoredBundle + found.append(vendoredBundle) } let explicitCandidates: [String] = [ @@ -230,14 +512,14 @@ enum MLXBundleSupport { for relative in explicitCandidates { let url = workspaceRoot.appendingPathComponent(relative, isDirectory: true) if fm.fileExists(atPath: url.path) { - return url + found.append(url) } } let buildRoot = workspaceRoot.appendingPathComponent(".build", isDirectory: true) var isDir: ObjCBool = false guard fm.fileExists(atPath: buildRoot.path, isDirectory: &isDir), isDir.boolValue else { - return nil + return found } let derivedDataDirs: [URL] @@ -258,7 +540,7 @@ enum MLXBundleSupport { dir.appendingPathComponent("Build/Products/Release/mlx-swift_Cmlx.bundle", isDirectory: true), ] for candidate in candidates where fm.fileExists(atPath: candidate.path) { - return candidate + found.append(candidate) } } @@ -275,12 +557,12 @@ enum MLXBundleSupport { includingPropertiesForKeys: [.isDirectoryKey], options: [.skipsHiddenFiles, .skipsPackageDescendants] ) else { - return nil + return found } for case let url as URL in enumerator { if url.lastPathComponent == "mlx-swift_Cmlx.bundle" { - return url + found.append(url) } if skipped.contains(url.lastPathComponent), (try? url.resourceValues(forKeys: [.isDirectoryKey]).isDirectory) == true { @@ -288,7 +570,7 @@ enum MLXBundleSupport { } } - return nil + return found } #endif } diff --git a/Sources/MereRunCore/Support/MLXRuntimeVersion.swift b/Sources/MereRunCore/Support/MLXRuntimeVersion.swift new file mode 100644 index 00000000..3901907b --- /dev/null +++ b/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 + } +} diff --git a/Tests/MereRunCoreTests/MLXTestSupport.swift b/Tests/MereRunCoreTests/MLXTestSupport.swift index 84734550..4a2898e8 100644 --- a/Tests/MereRunCoreTests/MLXTestSupport.swift +++ b/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 `/Contents/Resources/.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: // .../.xctest/Contents/MacOS/ diff --git a/Tests/MereRunCoreTests/SDPAVectorKernelTests.swift b/Tests/MereRunCoreTests/SDPAVectorKernelTests.swift new file mode 100644 index 00000000..d38aa2e1 --- /dev/null +++ b/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..., ../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) diff --git a/scripts/build_mlx_metallib.sh b/scripts/build_mlx_metallib.sh new file mode 100755 index 00000000..c6e5ccc8 --- /dev/null +++ b/scripts/build_mlx_metallib.sh @@ -0,0 +1,282 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Build mlx-swift's AOT Metal kernel library (default.metallib) from the +# CURRENT dependency checkout and stamp it with the versions it was built from. +# +# Why this exists: plain `swift build` never compiles the .metal sources — +# there is no metallib rule in the SwiftPM manifest — yet the MLX runtime +# hard-requires the library and loads whatever file it finds with no version +# validation. A metallib left over from an older mlx-swift silently corrupts +# inference: on 2026-07-03 a pre-0.30 metallib paired with the mlx 0.31.1 +# host dispatch produced gibberish, nondeterministic decode for every MLX +# text model once the KV length crossed 1024 (2-pass SDPA kernel ABI change). +# +# The produced library is accompanied by a `default.metallib.version` sidecar +# recording the mlx core version, the mlx-swift pin, and a hash of the kernel +# sources. `mere.run` validates the sidecar at startup (MLXBundleSupport) and +# refuses to run against a mismatched library. +# +# Usage: +# scripts/build_mlx_metallib.sh Build and install into +# .build (debug+release). +# scripts/build_mlx_metallib.sh --configuration release +# Install into one config. +# scripts/build_mlx_metallib.sh --output DIR Build and write the +# metallib + sidecar pair +# into DIR only. +# scripts/build_mlx_metallib.sh --verify-only PATH Compare an existing +# bundle/sidecar against +# the current checkout; +# exit 1 if stale. +# +# Environment: +# MERERUN_MLX_SWIFT_CHECKOUT Override the mlx-swift checkout location +# (default: .build/checkouts/mlx-swift). + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +checkout="${MERERUN_MLX_SWIFT_CHECKOUT:-$repo_root/.build/checkouts/mlx-swift}" +gen_dir="$checkout/Source/Cmlx/mlx-generated/metal" +version_header="$checkout/Source/Cmlx/mlx/mlx/version.h" +resolved="$repo_root/Package.resolved" +stamp_name="default.metallib.version" + +usage() { + sed -n '3,36p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//' +} + +mode="install" +output_dir="" +verify_target="" +configurations=(debug release) + +while [[ $# -gt 0 ]]; do + case "$1" in + --output) + [[ $# -ge 2 ]] || { echo "error: --output needs a directory" >&2; exit 64; } + mode="output"; output_dir="$2"; shift 2 ;; + --configuration) + [[ $# -ge 2 ]] || { echo "error: --configuration needs debug|release|all" >&2; exit 64; } + case "$2" in + debug|release) configurations=("$2") ;; + all) configurations=(debug release) ;; + *) echo "error: unknown configuration: $2" >&2; exit 64 ;; + esac + shift 2 ;; + --verify-only) + [[ $# -ge 2 ]] || { echo "error: --verify-only needs a bundle/sidecar path" >&2; exit 64; } + mode="verify"; verify_target="$2"; shift 2 ;; + -h|--help) usage; exit 0 ;; + *) echo "error: unknown argument: $1" >&2; usage >&2; exit 64 ;; + esac +done + +if [[ ! -d "$gen_dir" ]]; then + cat >&2 <&2 + exit 1 +fi + +read -r swift_pin_version swift_pin_revision < <(python3 - "$resolved" <<'EOF' +import json, sys +with open(sys.argv[1]) as f: + doc = json.load(f) +pins = doc.get("pins") or doc.get("object", {}).get("pins", []) +for pin in pins: + identity = (pin.get("identity") or pin.get("package", "")).lower() + if identity == "mlx-swift": + state = pin.get("state", {}) + print(state.get("version", "unknown"), state.get("revision", "unknown")) + break +else: + print("unknown", "unknown") +EOF +) + +sources_hash="$(cd "$gen_dir" && find . -type f \( -name '*.metal' -o -name '*.h' \) -print0 \ + | sort -z | xargs -0 shasum -a 256 | shasum -a 256 | cut -d' ' -f1)" + +# --- Verify mode ------------------------------------------------------------ + +find_sidecar() { + local target="$1" + local candidates=( + "$target" + "$target/$stamp_name" + "$target/Contents/Resources/$stamp_name" + "$target/Resources/$stamp_name" + ) + for c in "${candidates[@]}"; do + if [[ -f "$c" && "$(basename "$c")" == "$stamp_name" ]]; then + echo "$c" + return 0 + fi + done + return 1 +} + +stamp_field() { + awk -F': ' -v key="$1" '$1 == key { print $2; exit }' "$2" +} + +if [[ "$mode" == "verify" ]]; then + if ! sidecar="$(find_sidecar "$verify_target")"; then + echo "STALE: no $stamp_name found under $verify_target" >&2 + echo " (unstamped metallib — provenance unknown; rebuild with scripts/build_mlx_metallib.sh)" >&2 + exit 1 + fi + stamped_core="$(stamp_field "mlx-core-version" "$sidecar")" + stamped_rev="$(stamp_field "mlx-swift-revision" "$sidecar")" + stamped_hash="$(stamp_field "kernel-sources-sha256" "$sidecar")" + status=0 + [[ "$stamped_core" == "$core_version" ]] || { echo "STALE: mlx core version $stamped_core != checkout $core_version" >&2; status=1; } + [[ "$stamped_rev" == "$swift_pin_revision" ]] || { echo "STALE: mlx-swift revision $stamped_rev != pinned $swift_pin_revision" >&2; status=1; } + [[ "$stamped_hash" == "$sources_hash" ]] || { echo "STALE: kernel source hash mismatch" >&2; status=1; } + if [[ $status -eq 0 ]]; then + echo "OK: $sidecar matches mlx-swift $swift_pin_version (mlx core $core_version)" + fi + exit $status +fi + +# --- Build ------------------------------------------------------------------ + +workdir="$(mktemp -d -t mlx-metallib)" +trap 'rm -rf "$workdir"' EXIT + +metal_sources=("$gen_dir"/*.metal) +if [[ ${#metal_sources[@]} -eq 0 ]]; then + echo "error: no .metal sources in $gen_dir" >&2 + exit 1 +fi + +echo "[metallib] compiling ${#metal_sources[@]} kernels from mlx-swift $swift_pin_version (mlx core $core_version)" +# -fno-fast-math is mandatory: Metal defaults to fast math, and mlx's kernels +# require IEEE semantics (mlx's own CMake passes the same flag). +for src in "${metal_sources[@]}"; do + base="$(basename "$src" .metal)" + xcrun -sdk macosx metal \ + -x metal -Wall -Wextra -fno-fast-math -Wno-c++17-extensions -Wno-c++20-extensions \ + -c "$src" -I "$gen_dir" -o "$workdir/$base.air" \ + 2> "$workdir/$base.err" & +done +wait + +failed=0 +for src in "${metal_sources[@]}"; do + base="$(basename "$src" .metal)" + if [[ ! -f "$workdir/$base.air" ]]; then + echo "error: metal compile failed for $base.metal:" >&2 + cat "$workdir/$base.err" >&2 + failed=1 + fi +done +[[ $failed -eq 0 ]] || exit 1 + +xcrun -sdk macosx metallib "$workdir"/*.air -o "$workdir/default.metallib" + +metal_compiler="$(xcrun -sdk macosx metal --version 2>/dev/null | head -1 || echo unknown)" +cat > "$workdir/$stamp_name" < "$plist" <<'EOF' + + + + + CFBundleDevelopmentRegion + en + CFBundleIdentifier + mlx-swift.Cmlx.resources + CFBundleInfoDictionaryVersion + 6.0 + CFBundleName + mlx-swift_Cmlx + CFBundlePackageType + BNDL + + +EOF +} + +if [[ "$mode" == "output" ]]; then + install_pair "$output_dir" + echo "[metallib] wrote $output_dir/default.metallib (+ $stamp_name)" + exit 0 +fi + +# Install into the SwiftPM build tree: the canonical bundle location, plus the +# flat compatibility locations the MLX loader (and MLXTestSupport) probe. +for config in "${configurations[@]}"; do + build_root="$repo_root/.build/arm64-apple-macosx/$config" + install_pair "$build_root/mlx-swift_Cmlx.bundle/Contents/Resources" + ensure_bundle_info_plist "$build_root/mlx-swift_Cmlx.bundle" + install_pair "$build_root/Resources" + cp -f "$workdir/default.metallib" "$build_root/Resources/mlx.metallib" + cp -f "$workdir/default.metallib" "$build_root/mlx.metallib" + echo "[metallib] installed into .build/arm64-apple-macosx/$config" +done + +# The repo vendors a copy (tracked in git) so plain `swift build` users get +# working Metal shaders without the full Xcode metal toolchain. It is the +# FIRST location the bundle/test resolvers consult, so it must be regenerated +# on every mlx-swift bump — commit the refreshed pair when it changes. +vendor_bundle="$repo_root/vendor/mlx-swift_Cmlx.bundle" +if [[ -d "$vendor_bundle" ]]; then + # The vendor pair is tracked in git; skip the refresh when the existing + # stamp already matches this exact build (same sources, same pin) so + # routine script runs don't churn the tracked sidecar's built-at line. + vendor_stamp="$vendor_bundle/Contents/Resources/$stamp_name" + if [[ -f "$vendor_stamp" ]] \ + && [[ "$(stamp_field "kernel-sources-sha256" "$vendor_stamp")" == "$sources_hash" ]] \ + && [[ "$(stamp_field "mlx-swift-revision" "$vendor_stamp")" == "$swift_pin_revision" ]] \ + && [[ "$(stamp_field "mlx-core-version" "$vendor_stamp")" == "$core_version" ]]; then + echo "[metallib] vendored $vendor_bundle already current; left untouched" + else + install_pair "$vendor_bundle/Contents/Resources" + ensure_bundle_info_plist "$vendor_bundle" + echo "[metallib] refreshed vendored $vendor_bundle (tracked in git — commit if changed)" + fi +fi + +echo "[metallib] stamp: mlx core $core_version, mlx-swift $swift_pin_version ($swift_pin_revision)" diff --git a/scripts/install.sh b/scripts/install.sh index 1050315c..44341fd1 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -185,19 +185,34 @@ fi # Install MLX Metal shader resources alongside the binary. # mlx-swift looks for metallib files in a Resources/ directory next to the executable. MLX_BUNDLE="$BIN_DEST_DIR/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib" +MLX_STAMP="$BIN_DEST_DIR/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version" if [[ "$install_platform" == "Darwin" && -f "$MLX_BUNDLE" ]]; then echo "[mere.run] installing MLX Metal shaders..." + if [[ ! -f "$MLX_STAMP" ]]; then + # The CLI validates the stamp at startup; an unstamped library cannot be + # checked against the binary and warns on every run. Packages built via + # scripts/build_mere_run_app.sh always carry the stamp. + echo "[mere.run] WARNING: packaged MLX Metal library has no version stamp;" >&2 + echo "[mere.run] the CLI cannot verify it matches this binary. Repackage" >&2 + echo "[mere.run] with scripts/build_mlx_metallib.sh to add the stamp." >&2 + fi RESOURCES_DIR="$BIN_DEST_DIR/Resources" if [[ "$can_install_without_sudo" == true ]]; then mkdir -p "$RESOURCES_DIR" cp -f "$MLX_BUNDLE" "$RESOURCES_DIR/default.metallib" cp -f "$MLX_BUNDLE" "$RESOURCES_DIR/mlx.metallib" cp -f "$MLX_BUNDLE" "$BIN_DEST_DIR/mlx.metallib" + if [[ -f "$MLX_STAMP" ]]; then + cp -f "$MLX_STAMP" "$RESOURCES_DIR/default.metallib.version" + fi else sudo mkdir -p "$RESOURCES_DIR" sudo cp -f "$MLX_BUNDLE" "$RESOURCES_DIR/default.metallib" sudo cp -f "$MLX_BUNDLE" "$RESOURCES_DIR/mlx.metallib" sudo cp -f "$MLX_BUNDLE" "$BIN_DEST_DIR/mlx.metallib" + if [[ -f "$MLX_STAMP" ]]; then + sudo cp -f "$MLX_STAMP" "$RESOURCES_DIR/default.metallib.version" + fi fi fi diff --git a/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib b/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib index a92e56cc..0285a441 100644 Binary files a/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib and b/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib differ diff --git a/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version b/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version new file mode 100644 index 00000000..c80cccdd --- /dev/null +++ b/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version @@ -0,0 +1,6 @@ +mlx-core-version: 0.31.1 +mlx-swift-version: 0.31.4 +mlx-swift-revision: dc43e62d7055353c7f99fa071a4e71d29dfddc44 +kernel-sources-sha256: 8cc1c1ac4c7679a81ff0be1c8f67cd487386c36b164bf7fe6cb13960025f28bf +built-at: 2026-07-03T20:03:03Z +metal-compiler: Apple metal version 32023.883 (metalfe-32023.883)