From 14b77caa3f24021869db06edfcc994cc914fc4ab Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Sat, 15 Aug 2026 23:36:22 -0300 Subject: [PATCH 1/3] Refresh MLX runtime for H3 Metal --- Package.resolved | 2 +- Package.swift | 2 +- THIRD_PARTY_NOTICES.md | 6 +++--- .../Contents/Resources/default.metallib.version | 6 +++--- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/Package.resolved b/Package.resolved index a755cda5..a98fff84 100644 --- a/Package.resolved +++ b/Package.resolved @@ -33,7 +33,7 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/sawfwair/mlx-swift", "state" : { - "revision" : "3e6df6d8163a8f212061d15739eeeec12d5b89e3" + "revision" : "5bf3e46fecfb69cd3b559025fa99885ddd188731" } }, { diff --git a/Package.swift b/Package.swift index a9454cfa..2e95db00 100644 --- a/Package.swift +++ b/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( diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index ee7f946e..04d27bf7 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -716,10 +716,10 @@ SOFTWARE. - 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 ``` diff --git a/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version b/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version index e69afd1f..11bf4e7c 100644 --- a/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version +++ b/vendor/mlx-swift_Cmlx.bundle/Contents/Resources/default.metallib.version @@ -1,6 +1,6 @@ mlx-core-version: 0.32.1 mlx-swift-version: unknown -mlx-swift-revision: 3e6df6d8163a8f212061d15739eeeec12d5b89e3 -kernel-sources-sha256: fb0c62d372d6aaa75edfbcb950d9dd797fce944a7df7bcde24dce2a672024be5 -built-at: 2026-08-14T11:27:04Z +mlx-swift-revision: 5bf3e46fecfb69cd3b559025fa99885ddd188731 +kernel-sources-sha256: b791ce523bec5e6612766d9b00004fa66d3f3b1dbbbabd725b5d3c36cefbce41 +built-at: 2026-08-16T02:29:50Z metal-compiler: Apple metal version 32023.883 (metalfe-32023.883) From 34ede6e7402772693ab181629ef7d1e7dfc16231 Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Sat, 15 Aug 2026 23:46:31 -0300 Subject: [PATCH 2/3] Add H3 MPP projection lab candidate --- .../MiniMaxH3/MiniMaxH3MPPProjection.swift | 158 ++++++++++++++++++ THIRD_PARTY_NOTICES.md | 25 +++ .../MiniMaxH3MPPProjectionTests.swift | 155 +++++++++++++++++ docs/benchmarks/minimax-h3-h3c-transfer.md | 23 ++- scripts/h3-kernel-lab.sh | 10 +- 5 files changed, 368 insertions(+), 3 deletions(-) create mode 100644 Sources/MereRunCore/MiniMaxH3/MiniMaxH3MPPProjection.swift create mode 100644 Tests/MereRunCoreTests/MiniMaxH3MPPProjectionTests.swift diff --git a/Sources/MereRunCore/MiniMaxH3/MiniMaxH3MPPProjection.swift b/Sources/MereRunCore/MiniMaxH3/MiniMaxH3MPPProjection.swift new file mode 100644 index 00000000..10660a27 --- /dev/null +++ b/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{INPUT_DIM, ROWS}, + array{1, INPUT_DIM}); + auto matrix_b = tensor( + (device bfloat*)weight, + dextents{INPUT_DIM, OUTPUT_DIM}, + array{1, INPUT_DIM}); + auto matrix_c = tensor( + (device bfloat*)output, + dextents{OUTPUT_DIM, ROWS}, + array{1, OUTPUT_DIM}); + constexpr auto descriptor = matmul2d_descriptor( + TILE_M, + TILE_N, + static_cast(dynamic_extent), + false, + true, + false); + matmul2d> 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 + using namespace metal; + using namespace mpp::tensor_ops; + """, + ensureRowContiguous: true + ) + #endif +} diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 04d27bf7..03734bd6 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -710,6 +710,31 @@ 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 diff --git a/Tests/MereRunCoreTests/MiniMaxH3MPPProjectionTests.swift b/Tests/MereRunCoreTests/MiniMaxH3MPPProjectionTests.swift new file mode 100644 index 00000000..00e0a3ce --- /dev/null +++ b/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..%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 +} diff --git a/docs/benchmarks/minimax-h3-h3c-transfer.md b/docs/benchmarks/minimax-h3-h3c-transfer.md index 954bd207..86ae951a 100644 --- a/docs/benchmarks/minimax-h3-h3c-transfer.md +++ b/docs/benchmarks/minimax-h3-h3c-transfer.md @@ -18,6 +18,24 @@ materializations competitively requires an MLX/MLXFast primitive with a tiled quantized core and H3-specific epilogue, or an M5 TensorOps implementation. It is not achievable by further tuning the current standalone scalar kernels. +Resident BF16 has a separate M3/M4 opportunity. The refreshed MLX source has +a Metal 4 NAX GEMM implementation, but its runtime capability gate currently +requires Apple GPU generation 18 or newer, so it does not dispatch on the local +generation-16 M4 Max. The H3 lab now includes an independently gated MPP BF16 +projection candidate, adapted from WeeTodd's measured MiniMax-H3 tiles. Both +tile configurations pass the deterministic small-shape M4 Max bit-exact canary. +Production-shape timing remains unqualified because the first attempt correctly +stopped at the clean-host guard while unrelated builds and ML workloads were +active. The candidate is therefore not wired into model dispatch. Run the +isolated release arm only on a clean host: + +```bash +scripts/h3-kernel-lab.sh mpp-projections +``` + +This round deliberately targets the resident-BF16 M3/M4 path; M5-specific +TensorOps work is out of scope. + The quality-sensitive algorithm arms remain non-default. Reduced canvas, layer thinning, complete velocity reuse, and token reduction all produced material same-seed trajectory changes. The three-seed Ref2VA follow-up closed @@ -740,8 +758,9 @@ Algorithm gates: or memory from a production H3 block. 2. Add K2 as a BF16 head-major parity kernel before introducing INT8 projection arithmetic. This isolates layout correctness from quantization quality. -3. Prototype M5 INT8/TensorOps work in the pinned mlx-swift/MLXFast fork; retain - K1/K2 portable fallbacks in mere.run. +3. Qualify the resident-BF16 MPP projection candidate on M3/M4 at all four H3 + projection shapes and then through one exact 50-block forward. Keep it out + of model dispatch until those clean-host gates pass; M5 work is out of scope. 4. Qualify K3, K4, and K5 in that order because each consumes the preceding layout and activation contract. 5. Run A1-A4 only after the exact kernel baseline is stable, beginning with A1 diff --git a/scripts/h3-kernel-lab.sh b/scripts/h3-kernel-lab.sh index 5724e89c..63615139 100755 --- a/scripts/h3-kernel-lab.sh +++ b/scripts/h3-kernel-lab.sh @@ -131,6 +131,14 @@ case "$mode" in projections) run_release_test DiTShapeBenchTests/testMiniMaxH3QmmVsResidentBF16 ;; + mpp-projections) + export MERERUN_H3_MPP_BENCH=1 + export MERERUN_TEST_MLX_DEVICE=gpu + export MERERUN_H3_BENCH_ROWS="${MERERUN_H3_BENCH_ROWS:-14958}" + export MERERUN_H3_BENCH_ROUNDS="${MERERUN_H3_BENCH_ROUNDS:-4}" + run_release_test \ + MiniMaxH3MPPProjectionTests/testProductionShapeReleaseBenchmark + ;; modulation) export MERERUN_H3_BENCH_ROWS="${MERERUN_H3_BENCH_ROWS:-29018}" run_release_test DiTShapeBenchTests/testMiniMaxH3AdaLNRunModulation @@ -253,7 +261,7 @@ case "$mode" in run_release_test MiniMaxH3Tests/testInstalledAudioVAEDecodeMatchesReference ;; *) - print -u2 "usage: scripts/h3-kernel-lab.sh [quick|attention|attention-block|projections|modulation|gate-adaln|gate-adaln-int8|qkv-layout|qkv-direct|affine-oproj|affine-ffn|buffer-alias|exact-ref2va|block|post|dtype|turnover|boundary|gemm|gemm-block|vae|audio-parity]" + print -u2 "usage: scripts/h3-kernel-lab.sh [quick|attention|attention-block|projections|mpp-projections|modulation|gate-adaln|gate-adaln-int8|qkv-layout|qkv-direct|affine-oproj|affine-ffn|buffer-alias|exact-ref2va|block|post|dtype|turnover|boundary|gemm|gemm-block|vae|audio-parity]" exit 64 ;; esac From 5b7e0c04aeb775db1eff52cedced307473839eff Mon Sep 17 00:00:00 2001 From: Kyle McCullough Date: Sat, 15 Aug 2026 23:51:39 -0300 Subject: [PATCH 3/3] Align MLX metallib provenance --- Sources/MereRunCLI/Support/MLXBundleSupport.swift | 4 ++-- docs/mlx-swift-fork.md | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Sources/MereRunCLI/Support/MLXBundleSupport.swift b/Sources/MereRunCLI/Support/MLXBundleSupport.swift index ce4e5051..c02e8c9f 100644 --- a/Sources/MereRunCLI/Support/MLXBundleSupport.swift +++ b/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 diff --git a/docs/mlx-swift-fork.md b/docs/mlx-swift-fork.md index 1cdde403..73a78d0c 100644 --- a/docs/mlx-swift-fork.md +++ b/docs/mlx-swift-fork.md @@ -1,12 +1,12 @@ # mlx-swift fork policy and compiled-call overhead mere-run pins the public `sawfwair/mlx-swift` fork at -`3e6df6d8163a8f212061d15739eeeec12d5b89e3`. It is rebased onto upstream +`5bf3e46fecfb69cd3b559025fa99885ddd188731`. It is rebased onto upstream `mlx-swift` `da318704cc0e972b61dcca43c62cd15e545362ae`, including the upstream `MLXArray` finalizer fix and generated-source-list maintenance. The embedded `sawfwair/mlx` revision is -`b57bd7640f3f7c743b76a58478faaf1e8ee084f2`, based on upstream MLX -`bd5c3a2b170bb95340482e35b2a49fb08aea4de3` and retaining the 0.32.1 ABI. +`31af89c4c21642236b8a2bc1358438512d9521e3`, based on upstream MLX +`9ab977b5649154590d598ea5d545aa1b3c97f883` and retaining the 0.32.1 ABI. The owned patch stack carries the Linux/CUDA package bridge, executor-safe Swift streams, native affine 1-bit CUDA quantize/dequantize/QMV execution, the