Refresh MLX and add H3 MPP projection lab

Package.swift

@@ -510,7 +510,7 @@ if !isLinuxPackage {
 var packageDependencies: [Package.Dependency] = (useLinuxPrebuiltMLX ? [] : [
   .package(
     url: "https://github.com/sawfwair/mlx-swift",
-    revision: "3e6df6d8163a8f212061d15739eeeec12d5b89e3"
+    revision: "5bf3e46fecfb69cd3b559025fa99885ddd188731"
   )
 ]) + [
   .package(

Sources/MereRunCLI/Support/MLXBundleSupport.swift

@@ -27,8 +27,8 @@ enum MLXBundleSupport {
 
     static let expectedProvenance = MetallibProvenance(
       coreVersion: "0.32.1",
-      swiftRevision: "3e6df6d8163a8f212061d15739eeeec12d5b89e3",
-      kernelSourcesSHA256: "fb0c62d372d6aaa75edfbcb950d9dd797fce944a7df7bcde24dce2a672024be5"
+      swiftRevision: "5bf3e46fecfb69cd3b559025fa99885ddd188731",
+      kernelSourcesSHA256: "b791ce523bec5e6612766d9b00004fa66d3f3b1dbbbabd725b5d3c36cefbce41"
     )
 
     /// Relationship between a bundle's metallib version stamp

Sources/MereRunCore/MiniMaxH3/MiniMaxH3MPPProjection.swift

@@ -0,0 +1,158 @@
+import Foundation
+import MLX
+import MLXFast
+
+/// Experimental BF16 projection primitive for the H3 kernel lab.
+///
+/// The MPP shader structure and H3-specific tile choices are adapted from
+/// WeeTodd-Nodes commit e5b0e014db1abe4c86fedc195d12dfcd18562042 and
+/// translated to Swift/MLX. This primitive remains outside production model
+/// dispatch until the repository's exactness and clean-host benchmark gates
+/// qualify it on supported Apple GPUs.
+enum MiniMaxH3MPPProjection {
+    struct Tile: Equatable, Sendable {
+        let rows: Int
+        let columns: Int
+        let simdgroups: Int
+
+        init(rows: Int, columns: Int, simdgroups: Int) {
+            precondition(rows > 0)
+            precondition(columns > 0)
+            precondition(simdgroups > 0)
+            self.rows = rows
+            self.columns = columns
+            self.simdgroups = simdgroups
+        }
+    }
+
+    static let standardTile = Tile(rows: 32, columns: 64, simdgroups: 2)
+    static let feedForwardOutputTile = Tile(rows: 64, columns: 128, simdgroups: 8)
+
+    static func tile(inputDimension: Int, outputDimension: Int) -> Tile {
+        if inputDimension == 14_336, outputDimension == 5_376 {
+            return feedForwardOutputTile
+        }
+        return standardTile
+    }
+
+    static var isAvailable: Bool {
+        #if os(macOS)
+        let version = ProcessInfo.processInfo.operatingSystemVersion
+        return Device.defaultDevice().deviceType == .gpu
+            && version.majorVersion >= 26
+        #else
+        return false
+        #endif
+    }
+
+    /// Computes `source @ weight.T` for contiguous BF16 H3 projection tensors.
+    ///
+    /// Returning `nil` is the complete capability and shape fallback contract;
+    /// callers retain standard MLX matmul as the source of truth.
+    static func project(
+        source: MLXArray,
+        weight: MLXArray,
+        tile requestedTile: Tile? = nil
+    ) -> MLXArray? {
+        #if os(macOS)
+        guard isAvailable,
+              source.dtype == .bfloat16,
+              weight.dtype == .bfloat16,
+              source.ndim >= 2,
+              weight.ndim == 2,
+              source.dim(-1) == weight.dim(1) else {
+            return nil
+        }
+
+        let inputDimension = source.dim(-1)
+        let outputDimension = weight.dim(0)
+        let rows = source.size / inputDimension
+        guard rows > 0 else { return nil }
+
+        let tile = requestedTile ?? tile(
+            inputDimension: inputDimension,
+            outputDimension: outputDimension
+        )
+        let threadCount = 32 * tile.simdgroups
+        let outputShape = Array(source.shape.dropLast()) + [outputDimension]
+        return kernel(
+            [source, weight],
+            template: [
+                ("ROWS", rows),
+                ("OUTPUT_DIM", outputDimension),
+                ("INPUT_DIM", inputDimension),
+                ("TILE_M", tile.rows),
+                ("TILE_N", tile.columns),
+                ("SIMDGROUPS", tile.simdgroups),
+            ],
+            grid: (
+                divideRoundUp(outputDimension, by: tile.columns) * threadCount,
+                divideRoundUp(rows, by: tile.rows),
+                1
+            ),
+            threadGroup: (threadCount, 1, 1),
+            outputShapes: [outputShape],
+            outputDTypes: [.bfloat16]
+        )[0]
+        #else
+        return nil
+        #endif
+    }
+
+    private static func divideRoundUp(_ value: Int, by divisor: Int) -> Int {
+        (value + divisor - 1) / divisor
+    }
+
+    #if os(macOS)
+    private static let kernel = MLXFast.metalKernel(
+        name: "mere_h3_mpp_bf16_nt_matmul_v1",
+        inputNames: ["source", "weight"],
+        outputNames: ["output"],
+        source: """
+            auto matrix_a = tensor(
+                (device bfloat*)source,
+                dextents<int, 2>{INPUT_DIM, ROWS},
+                array<int, 2>{1, INPUT_DIM});
+            auto matrix_b = tensor(
+                (device bfloat*)weight,
+                dextents<int, 2>{INPUT_DIM, OUTPUT_DIM},
+                array<int, 2>{1, INPUT_DIM});
+            auto matrix_c = tensor(
+                (device bfloat*)output,
+                dextents<int, 2>{OUTPUT_DIM, ROWS},
+                array<int, 2>{1, OUTPUT_DIM});
+            constexpr auto descriptor = matmul2d_descriptor(
+                TILE_M,
+                TILE_N,
+                static_cast<int>(dynamic_extent),
+                false,
+                true,
+                false);
+            matmul2d<descriptor, execution_simdgroups<SIMDGROUPS>> operation;
+            auto tile_a = matrix_a.slice(
+                0,
+                threadgroup_position_in_grid.y * TILE_M);
+            auto tile_b = matrix_b.slice(
+                0,
+                threadgroup_position_in_grid.x * TILE_N);
+            auto tile_c = matrix_c.slice(
+                threadgroup_position_in_grid.x * TILE_N,
+                threadgroup_position_in_grid.y * TILE_M);
+            auto result = operation.template get_destination_cooperative_tensor<
+                decltype(tile_a), decltype(tile_b), bfloat>();
+            #pragma unroll
+            for (ushort index = 0; index < result.get_capacity(); ++index) {
+                result[index] = bfloat(0.0f);
+            }
+            operation.run(tile_a, tile_b, result);
+            result.store(tile_c);
+        """,
+        header: """
+            #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+            using namespace metal;
+            using namespace mpp::tensor_ops;
+        """,
+        ensureRowContiguous: true
+    )
+    #endif
+}

THIRD_PARTY_NOTICES.md

@@ -710,16 +710,41 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 ```
 
+### WeeTodd MiniMax-H3 MPP projection research
+
+- purpose: the lab-only BF16 Metal Performance Primitives projection shader
+  structure and measured H3 tile choices were adapted from
+  [`wee-todd/WeeTodd-Nodes`](https://github.com/wee-todd/WeeTodd-Nodes) at
+  commit `e5b0e014db1abe4c86fedc195d12dfcd18562042` and translated to Swift/MLX
+- distribution boundary: no WeeTodd model weights, runtime package, or Python
+  source files are vendored or linked; the adapted primitive remains outside
+  production dispatch until mere.run's exactness and benchmark gates qualify it
+- license: Apache License 2.0
+
+```text
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+```
+
 ### `vendor/mlx-swift_Cmlx.bundle`
 
 - purpose: bundled MLX Metal shader resources used by MLX-backed runtime paths
 - source project: [`sawfwair/mlx-swift`](https://github.com/sawfwair/mlx-swift),
   based on upstream [`ml-explore/mlx-swift`](https://github.com/ml-explore/mlx-swift)
   0.32.1
-- pinned package revision: `3e6df6d8163a8f212061d15739eeeec12d5b89e3`
-- embedded MLX revision: `b57bd7640f3f7c743b76a58478faaf1e8ee084f2`
+- pinned package revision: `5bf3e46fecfb69cd3b559025fa99885ddd188731`
+- embedded MLX revision: `31af89c4c21642236b8a2bc1358438512d9521e3`
 - generated-kernel source SHA-256:
-  `fb0c62d372d6aaa75edfbcb950d9dd797fce944a7df7bcde24dce2a672024be5`
+  `b791ce523bec5e6612766d9b00004fa66d3f3b1dbbbabd725b5d3c36cefbce41`
 - license: MIT
 
 ```

Tests/MereRunCoreTests/MiniMaxH3MPPProjectionTests.swift

@@ -0,0 +1,155 @@
+import Foundation
+import MLX
+import MLXRandom
+import XCTest
+@testable import MereRunCore
+
+final class MiniMaxH3MPPProjectionTests: MereRunCoreTestCase {
+    func testSelectsMeasuredFeedForwardOutputTile() {
+        XCTAssertEqual(
+            MiniMaxH3MPPProjection.tile(
+                inputDimension: 14_336,
+                outputDimension: 5_376
+            ),
+            .init(rows: 64, columns: 128, simdgroups: 8)
+        )
+        XCTAssertEqual(
+            MiniMaxH3MPPProjection.tile(
+                inputDimension: 5_376,
+                outputDimension: 21_504
+            ),
+            .init(rows: 32, columns: 64, simdgroups: 2)
+        )
+    }
+
+    #if os(macOS)
+    func testSmallProjectionMatchesMLXBitExactly() throws {
+        guard MiniMaxH3MPPProjection.isAvailable else {
+            throw XCTSkip(
+                "MPP projection parity requires macOS 26 and a Metal GPU."
+            )
+        }
+
+        MLXRandom.seed(2_026_081_015)
+        let source = MLXRandom.uniform(-0.5 ..< 0.5, [2, 37, 128])
+            .asType(.bfloat16)
+        let weight = MLXRandom.uniform(-0.5 ..< 0.5, [192, 128])
+            .asType(.bfloat16)
+        let reference = MLX.matmul(source, weight.T)
+        for tile in [
+            MiniMaxH3MPPProjection.standardTile,
+            MiniMaxH3MPPProjection.feedForwardOutputTile,
+        ] {
+            let candidate = try XCTUnwrap(
+                MiniMaxH3MPPProjection.project(
+                    source: source,
+                    weight: weight,
+                    tile: tile
+                )
+            )
+            MLX.eval(reference, candidate)
+
+            XCTAssertEqual(candidate.shape, [2, 37, 192])
+            XCTAssertEqual(candidate.dtype, .bfloat16)
+            XCTAssertTrue(MLX.arrayEqual(reference, candidate).item(Bool.self))
+        }
+    }
+
+    func testProductionShapeReleaseBenchmark() throws {
+        guard ProcessInfo.processInfo.environment["MERERUN_H3_MPP_BENCH"] == "1" else {
+            throw XCTSkip(
+                "Set MERERUN_H3_MPP_BENCH=1 to run the H3 MPP projection benchmark."
+            )
+        }
+        guard MiniMaxH3MPPProjection.isAvailable else {
+            throw XCTSkip("The H3 MPP projection benchmark requires macOS 26 and a Metal GPU.")
+        }
+
+        let rows = max(
+            1,
+            Int(ProcessInfo.processInfo.environment["MERERUN_H3_BENCH_ROWS"] ?? "")
+                ?? 14_958
+        )
+        let rounds = max(
+            2,
+            Int(ProcessInfo.processInfo.environment["MERERUN_H3_BENCH_ROUNDS"] ?? "")
+                ?? 4
+        )
+        for (name, inputDimension, outputDimension) in [
+            ("qkv", 5_376, 21_504),
+            ("attention-output", 7_168, 5_376),
+            ("feed-forward-input", 5_376, 28_672),
+            ("feed-forward-output", 14_336, 5_376),
+        ] {
+            try compareProductionShape(
+                name: name,
+                rows: rows,
+                inputDimension: inputDimension,
+                outputDimension: outputDimension,
+                rounds: rounds
+            )
+            MLX.Memory.clearCache()
+        }
+    }
+
+    private func compareProductionShape(
+        name: String,
+        rows: Int,
+        inputDimension: Int,
+        outputDimension: Int,
+        rounds: Int
+    ) throws {
+        let source = MLXRandom.uniform(
+            -0.25 ..< 0.25,
+            [1, rows, inputDimension]
+        ).asType(.bfloat16)
+        let weight = MLXRandom.uniform(
+            -0.25 ..< 0.25,
+            [outputDimension, inputDimension]
+        ).asType(.bfloat16)
+        let reference = MLX.matmul(source, weight.T)
+        let candidate = try XCTUnwrap(
+            MiniMaxH3MPPProjection.project(source: source, weight: weight)
+        )
+        MLX.eval(source, weight, reference, candidate)
+        XCTAssertTrue(MLX.arrayEqual(reference, candidate).item(Bool.self))
+
+        var bestMLX = Double.greatestFiniteMagnitude
+        var bestMPP = Double.greatestFiniteMagnitude
+        for round in 0..<rounds {
+            if round.isMultiple(of: 2) {
+                bestMLX = min(bestMLX, measure { MLX.matmul(source, weight.T) })
+                bestMPP = min(bestMPP, measureMPP(source: source, weight: weight))
+            } else {
+                bestMPP = min(bestMPP, measureMPP(source: source, weight: weight))
+                bestMLX = min(bestMLX, measure { MLX.matmul(source, weight.T) })
+            }
+        }
+
+        print(String(
+            format: "[h3-lab] mpp rows=%d projection=%@ %d->%d mlx_ms=%.3f "
+                + "mpp_ms=%.3f speedup=%.3fx exact=true",
+            rows,
+            name,
+            inputDimension,
+            outputDimension,
+            bestMLX * 1_000,
+            bestMPP * 1_000,
+            bestMLX / bestMPP
+        ))
+    }
+
+    private func measure(_ body: () -> MLXArray) -> Double {
+        let started = CFAbsoluteTimeGetCurrent()
+        MLX.eval(body())
+        return CFAbsoluteTimeGetCurrent() - started
+    }
+
+    private func measureMPP(source: MLXArray, weight: MLXArray) -> Double {
+        measure {
+            MiniMaxH3MPPProjection.project(source: source, weight: weight)
+                ?? MLX.matmul(source, weight.T)
+        }
+    }
+    #endif
+}