From db9b77da44735dd94e2f8737355c3056464b56ca Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:03:52 -0700 Subject: [PATCH 1/4] perf(mlx-swift): mirror R1 expert-QMM kernel, expose diagnostics, fix compiled-fn lock order - Regenerate Source/Cmlx/mlx-generated quantized.{h,cpp} byte-matched to the nested canonical R1 source (expert tiles + parallel builder); the corresponding dispatch mirror is the only regeneration. - Expose the Apple-side diagnostics contract (mlx_metal_gemma4_expert_qmm_diagnostics{,_snapshot,_reset}) as Swift GPU.Gemma4ExpertQMMDiagnostics / gemma4ExpertQMMDiagnostics() / resetGemma4ExpertQMMDiagnostics() via new transitive headers; no mlx-c gitlink advance. - Transforms+Compile: take evalLock before the per-function lock so an outer compile trace can call a nested compiled function while another thread evaluates it; fixes the deadlock found in the v0.8.2 verification wave. - Add SortedGatherQuantizedMMTests: exact-shape ON/OFF parity, selector misses, NAX priority, AOT absence, and counter invariant/reset. - Advance nested Cmlx/mlx gitlink to the R1 kernel commit. --- Source/Cmlx/include-framework/Cmlx.h | 1 + .../mlx-backend-common-gemma4_expert_qmm.h | 271 ++++++++++ .../mlx-backend-metal-device.h | 36 ++ Source/Cmlx/include/mlx.h | 1 + Source/Cmlx/include/mlx/gemma4_expert_qmm.h | 41 ++ Source/Cmlx/mlx | 2 +- Source/Cmlx/mlx-generated/metal/quantized.h | 300 +++++++++++ Source/Cmlx/mlx-generated/quantized.cpp | 300 +++++++++++ Source/MLX/GPU+Metal.swift | 76 +++ Source/MLX/Transforms+Compile.swift | 16 +- .../SortedGatherQuantizedMMTests.swift | 484 ++++++++++++++++++ tools/update-mlx-xcodeproj.sh | 10 + xcode/MLX.xcodeproj/project.pbxproj | 1 + 13 files changed, 1532 insertions(+), 7 deletions(-) create mode 100644 Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h create mode 100644 Source/Cmlx/include/mlx/gemma4_expert_qmm.h create mode 100644 Tests/MLXTests/SortedGatherQuantizedMMTests.swift diff --git a/Source/Cmlx/include-framework/Cmlx.h b/Source/Cmlx/include-framework/Cmlx.h index 8c0077864..b2cc54cf6 100644 --- a/Source/Cmlx/include-framework/Cmlx.h +++ b/Source/Cmlx/include-framework/Cmlx.h @@ -2,6 +2,7 @@ #include #include #include +#include #include #include diff --git a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h new file mode 100644 index 000000000..fae187344 --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h @@ -0,0 +1,271 @@ +// Copyright © 2023-2024 Apple Inc. + +#pragma once + +#include + +#include + +#if defined(__APPLE__) +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct mlx_metal_gemma4_expert_qmm_diagnostics { + uint8_t requested; + uint8_t aot_available; + uint8_t nax_available; + uint8_t armed; + uint64_t attempts; + uint64_t hits; + uint64_t fallback_nax; + uint64_t fallback_outer_route; + uint64_t fallback_quantization; + uint64_t fallback_topology; + uint64_t fallback_assignment_count; + uint64_t fallback_geometry; + uint64_t fallback_metallib_unavailable; +} mlx_metal_gemma4_expert_qmm_diagnostics; + +MLX_API void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot( + mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics); +MLX_API void mlx_metal_gemma4_expert_qmm_diagnostics_reset(void); +MLX_API void mlx_metal_gemma4_expert_qmm_diagnostics_clear_and_arm(void); +MLX_API void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm( + mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics); + +#ifdef __cplusplus +} +#endif +#endif + +#ifdef __cplusplus + +#include + +namespace mlx::core::metal { + +enum class Gemma4ExpertQMMRoute : uint8_t { + not_requested, + hit, + fallback_nax, + fallback_outer_route, + fallback_quantization, + fallback_topology, + fallback_assignment_count, + fallback_geometry, + fallback_metallib_unavailable, +}; + +struct Gemma4ExpertQMMRouteInput { + bool requested{false}; + bool aot_available{false}; + bool nax_available{false}; + bool outer_route{false}; + + bool affine{false}; + bool transpose{false}; + bool has_bias{false}; + bool indices_uint32{false}; + bool indices_contiguous{false}; + bool x_bfloat16{false}; + bool x_contiguous{false}; + bool w_uint32{false}; + bool w_contiguous{false}; + bool scales_bfloat16{false}; + bool scales_contiguous{false}; + bool biases_bfloat16{false}; + bool biases_contiguous{false}; + + int group_size{0}; + int bits{0}; + int expert_count{0}; + int assignments{0}; + int index_count{0}; + int k{0}; + int n{0}; + + int x_rank{0}; + int x_dim0{0}; + int x_dim1{0}; + int x_dim2{0}; + int w_rank{0}; + int w_dim0{0}; + int w_dim1{0}; + int w_dim2{0}; + int scales_rank{0}; + int scales_dim0{0}; + int scales_dim1{0}; + int scales_dim2{0}; + int biases_rank{0}; + int biases_dim0{0}; + int biases_dim1{0}; + int biases_dim2{0}; +}; + +inline Gemma4ExpertQMMRoute classify_gemma4_expert_qmm( + const Gemma4ExpertQMMRouteInput& input) { + if (!input.requested) { + return Gemma4ExpertQMMRoute::not_requested; + } + if (!input.outer_route) { + return Gemma4ExpertQMMRoute::fallback_outer_route; + } + // The existing NAX route owns every supported BF16/transposed RHS call and + // must win before the Gemma 4 specialization or its AOT capability matters. + if (input.nax_available) { + return Gemma4ExpertQMMRoute::fallback_nax; + } + if (!input.affine || !input.transpose || !input.has_bias || + !input.indices_uint32 || !input.indices_contiguous || + !input.x_bfloat16 || !input.x_contiguous || !input.w_uint32 || + !input.w_contiguous || !input.scales_bfloat16 || + !input.scales_contiguous || !input.biases_bfloat16 || + !input.biases_contiguous || input.group_size != 64 || input.bits != 4) { + return Gemma4ExpertQMMRoute::fallback_quantization; + } + if (input.expert_count != 128 || input.x_rank != 3 || + input.x_dim0 != input.assignments || input.x_dim1 != 1 || + input.x_dim2 != input.k || input.w_rank != 3 || input.w_dim0 != 128 || + input.scales_rank != 3 || input.scales_dim0 != 128 || + input.biases_rank != 3 || input.biases_dim0 != 128 || + input.index_count != input.assignments) { + return Gemma4ExpertQMMRoute::fallback_topology; + } + if (input.assignments != 4096 && input.assignments != 8192 && + input.assignments != 16384) { + return Gemma4ExpertQMMRoute::fallback_assignment_count; + } + + const bool gate_up = input.k == 2816 && input.n == 1408 && + input.w_dim1 == 1408 && input.w_dim2 == 352 && + input.scales_dim1 == 1408 && input.scales_dim2 == 44 && + input.biases_dim1 == 1408 && input.biases_dim2 == 44; + const bool down = input.k == 704 && input.n == 2816 && + input.w_dim1 == 2816 && input.w_dim2 == 88 && + input.scales_dim1 == 2816 && input.scales_dim2 == 11 && + input.biases_dim1 == 2816 && input.biases_dim2 == 11; + if (!gate_up && !down) { + return Gemma4ExpertQMMRoute::fallback_geometry; + } + if (!input.aot_available) { + return Gemma4ExpertQMMRoute::fallback_metallib_unavailable; + } + return Gemma4ExpertQMMRoute::hit; +} + +struct Gemma4ExpertQMMCounterSnapshot { + uint64_t hits{0}; + uint64_t fallback_nax{0}; + uint64_t fallback_outer_route{0}; + uint64_t fallback_quantization{0}; + uint64_t fallback_topology{0}; + uint64_t fallback_assignment_count{0}; + uint64_t fallback_geometry{0}; + uint64_t fallback_metallib_unavailable{0}; + bool armed{false}; + + uint64_t attempts() const { + return hits + fallback_nax + fallback_outer_route + + fallback_quantization + fallback_topology + + fallback_assignment_count + fallback_geometry + + fallback_metallib_unavailable; + } +}; + +class Gemma4ExpertQMMCounters { + public: + bool armed() const { + return armed_; + } + + // Recording is called only after the caller's plain armed branch. Keeping + // the branch at that boundary makes the unarmed inference path free of + // atomic operations while the engine-idle arm/disarm contract makes access + // to armed_ well-defined. + void record(Gemma4ExpertQMMRoute route) { + std::atomic* counter = nullptr; + switch (route) { + case Gemma4ExpertQMMRoute::not_requested: + return; + case Gemma4ExpertQMMRoute::hit: + counter = &hits_; + break; + case Gemma4ExpertQMMRoute::fallback_nax: + counter = &fallback_nax_; + break; + case Gemma4ExpertQMMRoute::fallback_outer_route: + counter = &fallback_outer_route_; + break; + case Gemma4ExpertQMMRoute::fallback_quantization: + counter = &fallback_quantization_; + break; + case Gemma4ExpertQMMRoute::fallback_topology: + counter = &fallback_topology_; + break; + case Gemma4ExpertQMMRoute::fallback_assignment_count: + counter = &fallback_assignment_count_; + break; + case Gemma4ExpertQMMRoute::fallback_geometry: + counter = &fallback_geometry_; + break; + case Gemma4ExpertQMMRoute::fallback_metallib_unavailable: + counter = &fallback_metallib_unavailable_; + break; + } + counter->fetch_add(1, std::memory_order_relaxed); + } + + Gemma4ExpertQMMCounterSnapshot snapshot() const { + return { + hits_.load(std::memory_order_relaxed), + fallback_nax_.load(std::memory_order_relaxed), + fallback_outer_route_.load(std::memory_order_relaxed), + fallback_quantization_.load(std::memory_order_relaxed), + fallback_topology_.load(std::memory_order_relaxed), + fallback_assignment_count_.load(std::memory_order_relaxed), + fallback_geometry_.load(std::memory_order_relaxed), + fallback_metallib_unavailable_.load(std::memory_order_relaxed), + armed_, + }; + } + + Gemma4ExpertQMMCounterSnapshot snapshot_and_disarm() { + const bool was_armed = armed_; + armed_ = false; + auto result = snapshot(); + result.armed = was_armed; + return result; + } + + void reset() { + hits_.store(0, std::memory_order_relaxed); + fallback_nax_.store(0, std::memory_order_relaxed); + fallback_outer_route_.store(0, std::memory_order_relaxed); + fallback_quantization_.store(0, std::memory_order_relaxed); + fallback_topology_.store(0, std::memory_order_relaxed); + fallback_assignment_count_.store(0, std::memory_order_relaxed); + fallback_geometry_.store(0, std::memory_order_relaxed); + fallback_metallib_unavailable_.store(0, std::memory_order_relaxed); + } + + void clear_and_arm() { + reset(); + armed_ = true; + } + + private: + bool armed_{false}; + std::atomic hits_{0}; + std::atomic fallback_nax_{0}; + std::atomic fallback_outer_route_{0}; + std::atomic fallback_quantization_{0}; + std::atomic fallback_topology_{0}; + std::atomic fallback_assignment_count_{0}; + std::atomic fallback_geometry_{0}; + std::atomic fallback_metallib_unavailable_{0}; +}; + +} // namespace mlx::core::metal + +#endif diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-device.h b/Source/Cmlx/include-framework/mlx-backend-metal-device.h index 18664a473..c4f376b3e 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-device.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-device.h @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -186,6 +187,38 @@ class MLX_API Device { const MTLFCList& func_consts = {}, const std::vector& linked_functions = {}); + bool gemma4_expert_qmm_requested() const { + return gemma4_expert_qmm_requested_; + } + + bool gemma4_expert_qmm_aot_available() const { + return gemma4_expert_qmm_aot_available_; + } + bool gemma4_expert_qmm_diagnostics_armed() const { + return gemma4_expert_qmm_counters_.armed(); + } + + // Call only inside a route boundary guarded by + // gemma4_expert_qmm_diagnostics_armed(). + void record_armed_gemma4_expert_qmm(Gemma4ExpertQMMRoute route) { + gemma4_expert_qmm_counters_.record(route); + } + + Gemma4ExpertQMMCounterSnapshot gemma4_expert_qmm_counter_snapshot() const { + return gemma4_expert_qmm_counters_.snapshot(); + } + Gemma4ExpertQMMCounterSnapshot + gemma4_expert_qmm_counter_snapshot_and_disarm() { + return gemma4_expert_qmm_counters_.snapshot_and_disarm(); + } + + void reset_gemma4_expert_qmm_counters() { + gemma4_expert_qmm_counters_.reset(); + } + void clear_and_arm_gemma4_expert_qmm_counters() { + gemma4_expert_qmm_counters_.clear_and_arm(); + } + ResidencySet& residency_set() { return residency_set_; } @@ -227,6 +260,9 @@ class MLX_API Device { std::shared_mutex library_mtx_; std::unordered_map> library_map_; NS::SharedPtr default_library_; + bool gemma4_expert_qmm_requested_{false}; + bool gemma4_expert_qmm_aot_available_{false}; + Gemma4ExpertQMMCounters gemma4_expert_qmm_counters_; std::unordered_map< MTL::Library*, std::unordered_map>> diff --git a/Source/Cmlx/include/mlx.h b/Source/Cmlx/include/mlx.h index 76eaec363..f7f2af9a1 100644 --- a/Source/Cmlx/include/mlx.h +++ b/Source/Cmlx/include/mlx.h @@ -2,3 +2,4 @@ #include "mlx/c/transforms_impl.h" #include "mlx/c/linalg.h" #include "mlx/c/fast.h" +#include "mlx/gemma4_expert_qmm.h" diff --git a/Source/Cmlx/include/mlx/gemma4_expert_qmm.h b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h new file mode 100644 index 000000000..ef8558340 --- /dev/null +++ b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h @@ -0,0 +1,41 @@ +// Copyright © 2026 Apple Inc. + +#ifndef MLX_GEMMA4_EXPERT_QMM_H +#define MLX_GEMMA4_EXPERT_QMM_H + +#include + +#if defined(__APPLE__) +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct mlx_metal_gemma4_expert_qmm_diagnostics { + uint8_t requested; + uint8_t aot_available; + uint8_t nax_available; + uint8_t armed; + uint64_t attempts; + uint64_t hits; + uint64_t fallback_nax; + uint64_t fallback_outer_route; + uint64_t fallback_quantization; + uint64_t fallback_topology; + uint64_t fallback_assignment_count; + uint64_t fallback_geometry; + uint64_t fallback_metallib_unavailable; +} mlx_metal_gemma4_expert_qmm_diagnostics; + +void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot( + mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics); +void mlx_metal_gemma4_expert_qmm_diagnostics_reset(void); +void mlx_metal_gemma4_expert_qmm_diagnostics_clear_and_arm(void); +void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm( + mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics); + +#ifdef __cplusplus +} +#endif +#endif + +#endif diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx index d5a240408..9dd10bef3 160000 --- a/Source/Cmlx/mlx +++ b/Source/Cmlx/mlx @@ -1 +1 @@ -Subproject commit d5a240408508f2be37f1a4893da0b415e8c0db55 +Subproject commit 9dd10bef39eda71c0f9a238e6235cdf245eb453b diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 12b5c85c6..7dad4310d 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -1084,6 +1084,134 @@ METAL_FUNC void qvm_impl( } } +template < + typename T, + const int group_size, + const int bits, + const bool aligned_N, + const int BM = 32, + const int BK = 32, + const int BN = 32> +METAL_FUNC void qmm_t_expert_impl( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + threadgroup T* Xs, + threadgroup T* Ws, + const constant int& K, + const constant int& N, + const int M, + const constant int& K_eff, + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); + static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); + + (void)lid; + + constexpr int WM = 2; + constexpr int WN = 2; + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + + // Instantiate the appropriate BlockMMA and Loader + using mma_t = mlx::steel:: + BlockMMA; + using loader_x_t = + mlx::steel::BlockLoader; + using loader_w_t = QuantizedBlockLoader< + T, + BN, + BK, + BK_padded, + 1, + WM * WN * SIMD_SIZE, + group_size, + bits>; + + // Set the block + const int K_w = K * bytes_per_pack / pack_factor; + const int K_g = K / group_size; + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + + auto wl = (const device uint8_t*)w; + + x += y_row * static_cast(K); + wl += y_col * K_w; + scales += y_col * K_g; + biases += y_col * K_g; + y += y_row * static_cast(N) + y_col; + + // Make the x loader and mma operation + const short num_els = min(BM, M - y_row); + const short num_outs = min(BN, N - y_col); + loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); + mma_t mma_op(simd_gid, simd_lid); + + if (num_els < BM) { + if (!aligned_N && num_outs < BN) { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(BK, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } else { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_unsafe(); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } + } else { + if (!aligned_N && num_outs < BN) { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(BK, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } else { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_unsafe(); + threadgroup_barrier(mem_flags::mem_threadgroup); + + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } + } + + // Store results to device memory + threadgroup_barrier(mem_flags::mem_threadgroup); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); + } else { + mma_op.store_result(y, N); + } +} + template < typename T, const int group_size, @@ -2239,6 +2367,178 @@ template < w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } +[[kernel]] void build_gemma4_sorted_expert_tiles_bm32( + const device uint32_t* indices [[buffer(0)]], + device uint4* descriptors [[buffer(1)]], + device uint* count [[buffer(2)]], + const constant int& M [[buffer(3)]], + uint lid [[thread_index_in_threadgroup]]) { + constexpr uint expert_count = 128; + constexpr uint BM = 32; + threadgroup uint segment_starts[expert_count + 1]; + threadgroup uint inclusive_tile_offsets[expert_count]; + + // One thread finds each expert's first sorted row. Thread 127 also supplies + // the sentinel, so all 129 boundaries are ready after one barrier. + int lower = 0; + int upper = M; + while (lower < upper) { + const int midpoint = lower + (upper - lower) / 2; + if (indices[midpoint] < lid) { + lower = midpoint + 1; + } else { + upper = midpoint; + } + } + segment_starts[lid] = uint(lower); + if (lid == expert_count - 1) { + segment_starts[expert_count] = uint(M); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint segment_rows = segment_starts[lid + 1] - segment_starts[lid]; + inclusive_tile_offsets[lid] = (segment_rows + BM - 1) / BM; + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Seven uniform Hillis-Steele strides form an inclusive scan for 128 + // experts. The read barrier precedes each in-place update and the write + // barrier makes that stride visible to the next one. + for (uint stride = 1; stride < expert_count; stride <<= 1) { + const uint addend = + lid >= stride ? inclusive_tile_offsets[lid - stride] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lid >= stride) { + inclusive_tile_offsets[lid] += addend; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const uint descriptor_count = + inclusive_tile_offsets[expert_count - 1]; + if (lid == expert_count - 1) { + count[0] = descriptor_count; + } + + // Every thread emits a strided share of the bounded descriptor array. + // upper_bound over the inclusive offsets maps each slot back to its expert. + for (uint slot = lid; slot < descriptor_count; slot += expert_count) { + uint expert_lower = 0; + uint expert_upper = expert_count; + while (expert_lower < expert_upper) { + const uint midpoint = + expert_lower + (expert_upper - expert_lower) / 2; + if (inclusive_tile_offsets[midpoint] <= slot) { + expert_lower = midpoint + 1; + } else { + expert_upper = midpoint; + } + } + const uint expert = expert_lower; + const uint expert_tile_begin = + expert == 0 ? 0 : inclusive_tile_offsets[expert - 1]; + const uint row = + segment_starts[expert] + (slot - expert_tile_begin) * BM; + const uint row_count = + min(BM, segment_starts[expert + 1] - row); + descriptors[slot] = uint4(row, row_count, expert, 0); + } +} + +template < + typename T, + const int group_size, + const int bits, + const bool aligned_N, + const int BM = 32, + const int BK = 32, + const int BN = 32> +[[kernel]] void affine_gather_qmm_gemma4_expert_tiles( + const device T* x [[buffer(0)]], + const device uint32_t* w [[buffer(1)]], + const device T* scales [[buffer(2)]], + const device T* biases [[buffer(3)]], + const device uint4* descriptors [[buffer(4)]], + const device uint* count [[buffer(5)]], + device T* y [[buffer(6)]], + const constant int& K [[buffer(7)]], + const constant int& N [[buffer(8)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BM == 32, "Gemma 4 expert tiles require BM=32"); + static_assert(BK == 32, "Gemma 4 expert tiles require BK=32"); + static_assert(BN == 32, "Gemma 4 expert tiles require BN=32"); + + // tid.y is uniform across the threadgroup. Empty capacity slots return + // before any threadgroup allocation is consumed by a microkernel barrier. + const uint descriptor_count = count[0]; + if (tid.y >= descriptor_count) { + return; + } + + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + constexpr int BK_padded = BK + 16 / sizeof(T); + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + + const uint4 descriptor = descriptors[tid.y]; + const size_t row_start = size_t(descriptor.x); + const int row_count = int(descriptor.y); + const size_t expert = size_t(descriptor.z); + + const int K_w = K * bytes_per_pack / pack_factor; + const int K_g = K / group_size; + const size_t expert_w_stride = size_t(N) * size_t(K_w); + const size_t expert_sb_stride = size_t(N) * size_t(K_g); + + x += row_start * size_t(K); + y += row_start * size_t(N); + const device uint8_t* expert_w = + reinterpret_cast(w) + + expert * expert_w_stride; + scales += expert * expert_sb_stride; + biases += expert * expert_sb_stride; + + const uint3 local_tid = uint3(tid.x, 0, 0); + if (row_count <= 16) { + qmm_t_expert_impl( + reinterpret_cast(expert_w), + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + row_count, + K, + local_tid, + lid, + simd_gid, + simd_lid); + } else { + qmm_t_expert_impl( + reinterpret_cast(expert_w), + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + row_count, + K, + local_tid, + lid, + simd_gid, + simd_lid); + } +} + template < typename T, int group_size, diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index b54c94efc..637316266 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -1097,6 +1097,134 @@ METAL_FUNC void qvm_impl( } } +template < + typename T, + const int group_size, + const int bits, + const bool aligned_N, + const int BM = 32, + const int BK = 32, + const int BN = 32> +METAL_FUNC void qmm_t_expert_impl( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + threadgroup T* Xs, + threadgroup T* Ws, + const constant int& K, + const constant int& N, + const int M, + const constant int& K_eff, + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); + static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); + + (void)lid; + + constexpr int WM = 2; + constexpr int WN = 2; + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + + // Instantiate the appropriate BlockMMA and Loader + using mma_t = mlx::steel:: + BlockMMA; + using loader_x_t = + mlx::steel::BlockLoader; + using loader_w_t = QuantizedBlockLoader< + T, + BN, + BK, + BK_padded, + 1, + WM * WN * SIMD_SIZE, + group_size, + bits>; + + // Set the block + const int K_w = K * bytes_per_pack / pack_factor; + const int K_g = K / group_size; + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + + auto wl = (const device uint8_t*)w; + + x += y_row * static_cast(K); + wl += y_col * K_w; + scales += y_col * K_g; + biases += y_col * K_g; + y += y_row * static_cast(N) + y_col; + + // Make the x loader and mma operation + const short num_els = min(BM, M - y_row); + const short num_outs = min(BN, N - y_col); + loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); + mma_t mma_op(simd_gid, simd_lid); + + if (num_els < BM) { + if (!aligned_N && num_outs < BN) { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_safe(short2(BK, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } else { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_safe(short2(BK, num_els)); + loader_w.load_unsafe(); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } + } else { + if (!aligned_N && num_outs < BN) { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_safe(short2(BK, num_outs)); + threadgroup_barrier(mem_flags::mem_threadgroup); + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } else { + for (int k = 0; k < K_eff; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_x.load_unsafe(); + loader_w.load_unsafe(); + threadgroup_barrier(mem_flags::mem_threadgroup); + + mma_op.mma(Xs, Ws); + loader_x.next(); + loader_w.next(); + } + } + } + + // Store results to device memory + threadgroup_barrier(mem_flags::mem_threadgroup); + if (num_els < BM || num_outs < BN) { + mma_op.store_result_safe(y, N, short2(num_outs, num_els)); + } else { + mma_op.store_result(y, N); + } +} + template < typename T, const int group_size, @@ -2252,6 +2380,178 @@ template < w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } +[[kernel]] void build_gemma4_sorted_expert_tiles_bm32( + const device uint32_t* indices [[buffer(0)]], + device uint4* descriptors [[buffer(1)]], + device uint* count [[buffer(2)]], + const constant int& M [[buffer(3)]], + uint lid [[thread_index_in_threadgroup]]) { + constexpr uint expert_count = 128; + constexpr uint BM = 32; + threadgroup uint segment_starts[expert_count + 1]; + threadgroup uint inclusive_tile_offsets[expert_count]; + + // One thread finds each expert's first sorted row. Thread 127 also supplies + // the sentinel, so all 129 boundaries are ready after one barrier. + int lower = 0; + int upper = M; + while (lower < upper) { + const int midpoint = lower + (upper - lower) / 2; + if (indices[midpoint] < lid) { + lower = midpoint + 1; + } else { + upper = midpoint; + } + } + segment_starts[lid] = uint(lower); + if (lid == expert_count - 1) { + segment_starts[expert_count] = uint(M); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint segment_rows = segment_starts[lid + 1] - segment_starts[lid]; + inclusive_tile_offsets[lid] = (segment_rows + BM - 1) / BM; + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Seven uniform Hillis-Steele strides form an inclusive scan for 128 + // experts. The read barrier precedes each in-place update and the write + // barrier makes that stride visible to the next one. + for (uint stride = 1; stride < expert_count; stride <<= 1) { + const uint addend = + lid >= stride ? inclusive_tile_offsets[lid - stride] : 0; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lid >= stride) { + inclusive_tile_offsets[lid] += addend; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const uint descriptor_count = + inclusive_tile_offsets[expert_count - 1]; + if (lid == expert_count - 1) { + count[0] = descriptor_count; + } + + // Every thread emits a strided share of the bounded descriptor array. + // upper_bound over the inclusive offsets maps each slot back to its expert. + for (uint slot = lid; slot < descriptor_count; slot += expert_count) { + uint expert_lower = 0; + uint expert_upper = expert_count; + while (expert_lower < expert_upper) { + const uint midpoint = + expert_lower + (expert_upper - expert_lower) / 2; + if (inclusive_tile_offsets[midpoint] <= slot) { + expert_lower = midpoint + 1; + } else { + expert_upper = midpoint; + } + } + const uint expert = expert_lower; + const uint expert_tile_begin = + expert == 0 ? 0 : inclusive_tile_offsets[expert - 1]; + const uint row = + segment_starts[expert] + (slot - expert_tile_begin) * BM; + const uint row_count = + min(BM, segment_starts[expert + 1] - row); + descriptors[slot] = uint4(row, row_count, expert, 0); + } +} + +template < + typename T, + const int group_size, + const int bits, + const bool aligned_N, + const int BM = 32, + const int BK = 32, + const int BN = 32> +[[kernel]] void affine_gather_qmm_gemma4_expert_tiles( + const device T* x [[buffer(0)]], + const device uint32_t* w [[buffer(1)]], + const device T* scales [[buffer(2)]], + const device T* biases [[buffer(3)]], + const device uint4* descriptors [[buffer(4)]], + const device uint* count [[buffer(5)]], + device T* y [[buffer(6)]], + const constant int& K [[buffer(7)]], + const constant int& N [[buffer(8)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BM == 32, "Gemma 4 expert tiles require BM=32"); + static_assert(BK == 32, "Gemma 4 expert tiles require BK=32"); + static_assert(BN == 32, "Gemma 4 expert tiles require BN=32"); + + // tid.y is uniform across the threadgroup. Empty capacity slots return + // before any threadgroup allocation is consumed by a microkernel barrier. + const uint descriptor_count = count[0]; + if (tid.y >= descriptor_count) { + return; + } + + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + constexpr int BK_padded = BK + 16 / sizeof(T); + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + + const uint4 descriptor = descriptors[tid.y]; + const size_t row_start = size_t(descriptor.x); + const int row_count = int(descriptor.y); + const size_t expert = size_t(descriptor.z); + + const int K_w = K * bytes_per_pack / pack_factor; + const int K_g = K / group_size; + const size_t expert_w_stride = size_t(N) * size_t(K_w); + const size_t expert_sb_stride = size_t(N) * size_t(K_g); + + x += row_start * size_t(K); + y += row_start * size_t(N); + const device uint8_t* expert_w = + reinterpret_cast(w) + + expert * expert_w_stride; + scales += expert * expert_sb_stride; + biases += expert * expert_sb_stride; + + const uint3 local_tid = uint3(tid.x, 0, 0); + if (row_count <= 16) { + qmm_t_expert_impl( + reinterpret_cast(expert_w), + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + row_count, + K, + local_tid, + lid, + simd_gid, + simd_lid); + } else { + qmm_t_expert_impl( + reinterpret_cast(expert_w), + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + row_count, + K, + local_tid, + lid, + simd_gid, + simd_lid); + } +} + template < typename T, int group_size, diff --git a/Source/MLX/GPU+Metal.swift b/Source/MLX/GPU+Metal.swift index c46440817..e3c8eabb7 100644 --- a/Source/MLX/GPU+Metal.swift +++ b/Source/MLX/GPU+Metal.swift @@ -245,4 +245,80 @@ public enum GPU { memorySize: memSize) } } + + /// Runtime state for the opt-in Gemma 4 sorted expert-QMM specialization. + /// + /// Route counters change only while diagnostics are explicitly armed. The + /// feature request and immutable runtime capabilities remain observable + /// while diagnostics are unarmed. + public struct Gemma4ExpertQMMDiagnostics: Equatable, Sendable { + public let requested: Bool + public let aotAvailable: Bool + public let naxAvailable: Bool + public let armed: Bool + public let attempts: UInt64 + public let hits: UInt64 + public let fallbackNAX: UInt64 + public let fallbackOuterRoute: UInt64 + public let fallbackQuantization: UInt64 + public let fallbackTopology: UInt64 + public let fallbackAssignmentCount: UInt64 + public let fallbackGeometry: UInt64 + public let fallbackMetallibUnavailable: UInt64 + + /// Sum of all mutually exclusive fallback reasons. + public var fallbacks: UInt64 { + fallbackNAX + fallbackOuterRoute + fallbackQuantization + + fallbackTopology + fallbackAssignmentCount + fallbackGeometry + + fallbackMetallibUnavailable + } + + fileprivate init(_ value: mlx_metal_gemma4_expert_qmm_diagnostics) { + requested = value.requested != 0 + aotAvailable = value.aot_available != 0 + naxAvailable = value.nax_available != 0 + armed = value.armed != 0 + attempts = value.attempts + hits = value.hits + fallbackNAX = value.fallback_nax + fallbackOuterRoute = value.fallback_outer_route + fallbackQuantization = value.fallback_quantization + fallbackTopology = value.fallback_topology + fallbackAssignmentCount = value.fallback_assignment_count + fallbackGeometry = value.fallback_geometry + fallbackMetallibUnavailable = value.fallback_metallib_unavailable + } + } + + /// Snapshot Gemma 4 expert-QMM feature, AOT capability, and route counters + /// without changing the armed state. + public static func gemma4ExpertQMMDiagnostics() -> Gemma4ExpertQMMDiagnostics { + var value = mlx_metal_gemma4_expert_qmm_diagnostics() + mlx_metal_gemma4_expert_qmm_diagnostics_snapshot(&value) + return Gemma4ExpertQMMDiagnostics(value) + } + + /// Clear the route counters without changing the armed state. + public static func resetGemma4ExpertQMMDiagnostics() { + mlx_metal_gemma4_expert_qmm_diagnostics_reset() + } + + /// Clear the route counters and arm accounting for a measured interval. + /// Call only at an engine-idle boundary after warmup has completed. + public static func clearAndArmGemma4ExpertQMMDiagnostics() { + mlx_metal_gemma4_expert_qmm_diagnostics_clear_and_arm() + } + + /// Snapshot the measured interval and disarm route accounting. Call only + /// at an engine-idle boundary after measured work has completed. + /// The returned `armed` value reports whether accounting was armed before + /// this call. + public static func snapshotAndDisarmGemma4ExpertQMMDiagnostics() + -> Gemma4ExpertQMMDiagnostics + { + var value = mlx_metal_gemma4_expert_qmm_diagnostics() + mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm(&value) + return Gemma4ExpertQMMDiagnostics(value) + } + } diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index e52835380..6706732ec 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -37,8 +37,14 @@ final class CompiledFunction: @unchecked (Sendable) { } func call(_ arguments: [MLXArray]) -> [MLXArray] { - lock.withLock { - innerCall(arguments) + // Every compiled call eventually enters MLX under the process-global + // eval lock. Take it before the per-function lock so an outer compile + // trace can safely call a nested compiled function while another + // thread is evaluating that same function. + evalLock.withLock { + lock.withLock { + innerCall(arguments) + } } } @@ -84,14 +90,12 @@ final class CompiledFunction: @unchecked (Sendable) { let innerClosure = new_mlx_closure(inner(tracers:)) defer { mlx_closure_free(innerClosure) } - // note: this will use the cached compile (via the id) - // but will be able to re-evaluate with fresh state if needed - evalLock.lock() + // This runs under evalLock (acquired before the per-function lock in + // call()) so nested compiled functions preserve one global lock order. var compiled = mlx_closure_new() let compileStatus = mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0) defer { mlx_closure_free(compiled) - evalLock.unlock() } // mlx_error was already dispatched on failure: diff --git a/Tests/MLXTests/SortedGatherQuantizedMMTests.swift b/Tests/MLXTests/SortedGatherQuantizedMMTests.swift new file mode 100644 index 000000000..66d512c30 --- /dev/null +++ b/Tests/MLXTests/SortedGatherQuantizedMMTests.swift @@ -0,0 +1,484 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +#if canImport(Metal) + +final class SortedGatherQuantizedMMTests: XCTestCase { + private let bits = 4 + private let expertCount = 8 + private let groupSize = 64 + private let inputDimensions = 64 + private let outputDimensions = 64 + + override class func setUp() { + setDefaultDevice() + } + + func testGemma4ExpertQMMDiagnosticsReflectProcessConfigurationAndArmState() { + let enabledValues = ["1", "true", "on", "yes"] + let requested = ProcessInfo.processInfo.environment["MLX_GATHER_QMM_EXPERT_SLICES"]? + .lowercased() + + _ = GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics() + GPU.resetGemma4ExpertQMMDiagnostics() + let unarmed = GPU.gemma4ExpertQMMDiagnostics() + XCTAssertEqual(unarmed.requested, enabledValues.contains(requested ?? "")) + XCTAssertFalse(unarmed.armed) + assertCounterInvariant(unarmed) + XCTAssertEqual(unarmed.attempts, 0) + + GPU.clearAndArmGemma4ExpertQMMDiagnostics() + let armed = GPU.gemma4ExpertQMMDiagnostics() + XCTAssertTrue(armed.armed) + XCTAssertEqual(armed.requested, unarmed.requested) + XCTAssertEqual(armed.aotAvailable, unarmed.aotAvailable) + XCTAssertEqual(armed.naxAvailable, unarmed.naxAvailable) + XCTAssertEqual(armed.attempts, 0) + + let measured = GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics() + XCTAssertTrue(measured.armed) + XCTAssertEqual(measured.attempts, 0) + XCTAssertFalse(GPU.gemma4ExpertQMMDiagnostics().armed) + } + + /// Activates the down-projection specialization exactly: E=128, K=704, + /// N=2816, affine 4-bit/group-64, bfloat16, and each allowlisted prefill + /// assignment count. The patterns cover aligned and unaligned boundaries, + /// empty experts, 16/17-row tails, 127 one-row experts, and heavy skew. + func testGemmaSortedAffineGatherQuantizedMMExpertSlices() throws { + try requireExpertSlicesEnabled() + let gemmaExpertCount = 128 + let gemmaInputDimensions = 704 + let gemmaOutputDimensions = 2816 + let fixture = gemmaAffineFixture( + expertCount: gemmaExpertCount, + inputDimensions: gemmaInputDimensions, + outputDimensions: gemmaOutputDimensions + ) + + var unaligned4096 = Array(repeating: 32, count: gemmaExpertCount) + unaligned4096.replaceSubrange(0 ..< 4, with: [3, 4, 5, 116]) + + var explicitTails4096 = Array(repeating: 32, count: gemmaExpertCount) + explicitTails4096.replaceSubrange(0 ..< 4, with: [1, 16, 17, 94]) + + var emptyExperts8192 = Array(repeating: 0, count: gemmaExpertCount) + emptyExperts8192[0] = 5 + emptyExperts8192[2] = 2 + emptyExperts8192[4] = 17 + emptyExperts8192[5] = 1 + emptyExperts8192[gemmaExpertCount - 1] = 8167 + + let maximallyFragmented16384 = + Array(repeating: 1, count: gemmaExpertCount - 1) + [16257] + let cases: [(name: String, expertCounts: [Int])] = [ + ("M4096, BM32-aligned boundaries", Array(repeating: 32, count: gemmaExpertCount)), + ("M4096, BM32-unaligned multi-boundary tiles", unaligned4096), + ("M4096, explicit 1-row and BM16/BM32 tail boundary", explicitTails4096), + ("M8192, empty experts and heavy skew", emptyExperts8192), + ("M16384, maximally fragmented expert boundaries", maximallyFragmented16384), + ] + + for testCase in cases { + let expertIndices = sortedExpertIndices(testCase.expertCounts) + let rowCount = expertIndices.count + XCTAssertTrue([4096, 8192, 16384].contains(rowCount)) + let rhsIndices = MLXArray(expertIndices).asType(.uint32) + let x = ones([rowCount, 1, gemmaInputDimensions], dtype: .bfloat16) + + GPU.clearAndArmGemma4ExpertQMMDiagnostics() + let actualSorted = gatherQuantizedMM( + x, + fixture.weights, + scales: fixture.scales, + biases: fixture.biases, + rhsIndices: rhsIndices, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: .affine, + sortedIndices: true + ) + let expected = broadcast( + take(fixture.expertOutputs, rhsIndices).reshaped([rowCount, 1, 1]), + to: [rowCount, 1, gemmaOutputDimensions] + ) + + eval(actualSorted, expected) + XCTAssertEqual(actualSorted.shape, expected.shape, testCase.name) + XCTAssertTrue( + actualSorted.allClose(expected, rtol: 1e-3, atol: 1e-3).item(Bool.self), + "sorted gather-QMM selected incorrect expert rows for \(testCase.name); " + + "max absolute error " + + "\((actualSorted - expected).abs().max().item(Float.self))" + ) + assertExactGemmaRoute( + GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics(), testCase.name) + Memory.clearCache() + } + } + + /// Covers the fused gate/up projection and the BM=16 and BM=32 + /// expert-only bodies. + func testGemmaGateUpSortedAffineGatherQuantizedMMExpertSlices() throws { + try requireExpertSlicesEnabled() + let gemmaExpertCount = 128 + let gemmaInputDimensions = 2816 + let gemmaOutputDimensions = 1408 + let assignmentCount = 4096 + let fixture = gemmaAffineFixture( + expertCount: gemmaExpertCount, + inputDimensions: gemmaInputDimensions, + outputDimensions: gemmaOutputDimensions + ) + + var expertCounts = Array(repeating: 32, count: gemmaExpertCount) + expertCounts.replaceSubrange(0 ..< 4, with: [3, 4, 5, 116]) + let expertIndices = sortedExpertIndices(expertCounts) + XCTAssertEqual(expertIndices.count, assignmentCount) + + let rhsIndices = MLXArray(expertIndices).asType(.uint32) + let x = ones([assignmentCount, 1, gemmaInputDimensions], dtype: .bfloat16) + GPU.clearAndArmGemma4ExpertQMMDiagnostics() + let actualSorted = gatherQuantizedMM( + x, + fixture.weights, + scales: fixture.scales, + biases: fixture.biases, + rhsIndices: rhsIndices, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: .affine, + sortedIndices: true + ) + let expected = broadcast( + take(fixture.expertOutputs, rhsIndices).reshaped([assignmentCount, 1, 1]), + to: [assignmentCount, 1, gemmaOutputDimensions] + ) + + eval(actualSorted, expected) + XCTAssertEqual(actualSorted.shape, expected.shape) + XCTAssertTrue( + actualSorted.allClose(expected, rtol: 1e-3, atol: 1e-3).item(Bool.self), + "gate/up sorted gather-QMM selected incorrect expert rows; max absolute error " + + "\((actualSorted - expected).abs().max().item(Float.self))" + ) + assertExactGemmaRoute( + GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics(), + "gate/up M4096, BM16/BM32 boundaries") + Memory.clearCache() + } + + /// A full Gemma topology with a prefill-like outer route but a disallowed + /// assignment count must stay on the legacy implementation. + func testGemmaSelectorRejectsNonAllowlistedAssignmentCount() { + let gemmaExpertCount = 128 + let rowCount = 512 + let gemmaInputDimensions = 704 + let gemmaOutputDimensions = 2816 + let fixture = gemmaAffineFixture( + expertCount: gemmaExpertCount, + inputDimensions: gemmaInputDimensions, + outputDimensions: gemmaOutputDimensions + ) + let expertCounts = Array(repeating: 4, count: gemmaExpertCount) + let rhsIndices = MLXArray(sortedExpertIndices(expertCounts)).asType(.uint32) + let x = ones([rowCount, 1, gemmaInputDimensions], dtype: .bfloat16) + + GPU.clearAndArmGemma4ExpertQMMDiagnostics() + let actual = gatherQuantizedMM( + x, + fixture.weights, + scales: fixture.scales, + biases: fixture.biases, + rhsIndices: rhsIndices, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: .affine, + sortedIndices: true + ) + let expected = broadcast( + take(fixture.expertOutputs, rhsIndices).reshaped([rowCount, 1, 1]), + to: [rowCount, 1, gemmaOutputDimensions] + ) + + eval(actual, expected) + XCTAssertTrue(actual.allClose(expected, rtol: 1e-3, atol: 1e-3).item(Bool.self)) + let diagnostics = GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics() + XCTAssertTrue(diagnostics.armed) + assertCounterInvariant(diagnostics) + if !diagnostics.requested { + XCTAssertEqual(diagnostics.attempts, 0) + } else if diagnostics.naxAvailable { + XCTAssertEqual(diagnostics.fallbackNAX, 1) + XCTAssertEqual(diagnostics.hits, 0) + } else { + XCTAssertEqual(diagnostics.fallbackAssignmentCount, 1) + XCTAssertEqual(diagnostics.hits, 0) + } + Memory.clearCache() + } + + /// Protects the established fallback BM=16 schedule on compact shapes. + func testSortedAffineGatherQuantizedMMFallbackExpertBoundaries() { + let weights = deterministicValues( + count: expertCount * outputDimensions * inputDimensions, + multiplier: 17, + modulus: 127 + ).reshaped([expertCount, outputDimensions, inputDimensions]).asType(.bfloat16) + let (quantizedWeights, scales, biases) = quantized( + weights, + groupSize: groupSize, + bits: bits, + mode: .affine + ) + let dequantizedWeights = dequantized( + quantizedWeights, + scales: scales, + biases: biases, + groupSize: groupSize, + bits: bits, + mode: .affine, + dtype: .bfloat16 + ).swappedAxes(-1, -2) + + let cases: [(name: String, expertCounts: [Int])] = [ + ("all experts, aligned, divisible M", Array(repeating: 16, count: expertCount)), + ("all experts, unaligned, multiple boundaries, divisible M", [3, 4, 5, 4, 20, 11, 9, 8]), + ("empty experts, unaligned, multiple boundaries, partial M", [5, 0, 2, 0, 17, 1, 0, 12]), + ] + + for testCase in cases { + let expertIndices = sortedExpertIndices(testCase.expertCounts) + let rowCount = expertIndices.count + let x = deterministicValues( + count: rowCount * inputDimensions, + multiplier: 29, + modulus: 113 + ).reshaped([rowCount, 1, inputDimensions]).asType(.bfloat16) + let rhsIndices = MLXArray(expertIndices).asType(.uint32) + + _ = GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics() + GPU.resetGemma4ExpertQMMDiagnostics() + let expectedQuantized = gatherQuantizedMM( + x, + quantizedWeights, + scales: scales, + biases: biases, + rhsIndices: rhsIndices, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: .affine, + sortedIndices: false + ) + eval(expectedQuantized) + let untracked = GPU.gemma4ExpertQMMDiagnostics() + XCTAssertFalse(untracked.armed) + XCTAssertEqual( + untracked.attempts, 0, + "unarmed fallback work must not synchronize route counters") + GPU.clearAndArmGemma4ExpertQMMDiagnostics() + let actualSorted = gatherQuantizedMM( + x, + quantizedWeights, + scales: scales, + biases: biases, + rhsIndices: rhsIndices, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: .affine, + sortedIndices: true + ) + let expectedDequantized = gatherMM( + x, + dequantizedWeights, + rhsIndices: rhsIndices, + sortedIndices: false + ) + + eval(actualSorted, expectedDequantized) + XCTAssertEqual(actualSorted.shape, [rowCount, 1, outputDimensions], testCase.name) + XCTAssertTrue( + actualSorted.allClose(expectedQuantized, rtol: 1e-2, atol: 1e-2).item(Bool.self), + "sorted and unsorted gather-QMM differ for \(testCase.name); max absolute error " + + "\((actualSorted - expectedQuantized).abs().max().item(Float.self))" + ) + XCTAssertTrue( + actualSorted.allClose(expectedDequantized, rtol: 2e-2, atol: 2e-2).item(Bool.self), + "sorted gather-QMM selected incorrect expert rows for \(testCase.name); " + + "max absolute error " + + "\((actualSorted - expectedDequantized).abs().max().item(Float.self))" + ) + + let diagnostics = GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics() + XCTAssertTrue(diagnostics.armed) + assertCounterInvariant(diagnostics) + if !diagnostics.requested { + XCTAssertEqual(diagnostics.attempts, 0) + } else if diagnostics.naxAvailable { + XCTAssertEqual(diagnostics.fallbackNAX, 1) + XCTAssertEqual(diagnostics.hits, 0) + } else { + XCTAssertEqual(diagnostics.fallbackTopology, 1) + XCTAssertEqual(diagnostics.hits, 0) + } + } + } + + /// Ordinary, non-gather QMM must retain its established global kernel for + /// aligned and tail row counts and must not touch expert-route counters. + func testOrdinaryQuantizedMMPreservesGlobalQMM() { + let fixture = gemmaAffineFixture( + expertCount: 1, + inputDimensions: inputDimensions, + outputDimensions: outputDimensions + ) + let packedInputDimensions = inputDimensions * bits / 32 + let weights = fixture.weights.reshaped([outputDimensions, packedInputDimensions]) + let scales = fixture.scales.reshaped([outputDimensions, inputDimensions / groupSize]) + let biases = fixture.biases.reshaped([outputDimensions, inputDimensions / groupSize]) + + for rowCount in [16, 17, 32, 33] { + let x = ones([rowCount, inputDimensions], dtype: .bfloat16) + let expected = broadcast( + fixture.expertOutputs.reshaped([1, 1]), + to: [rowCount, outputDimensions] + ) + + GPU.clearAndArmGemma4ExpertQMMDiagnostics() + let actual = quantizedMM( + x, + weights, + scales: scales, + biases: biases, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: .affine + ) + eval(actual, expected) + XCTAssertEqual(actual.shape, expected.shape) + XCTAssertTrue( + actual.allClose(expected, rtol: 1e-3, atol: 1e-3).item(Bool.self), + "ordinary QMM changed for row count \(rowCount); max absolute error " + + "\((actual - expected).abs().max().item(Float.self))" + ) + + let diagnostics = GPU.snapshotAndDisarmGemma4ExpertQMMDiagnostics() + XCTAssertTrue(diagnostics.armed) + XCTAssertEqual(diagnostics.attempts, 0) + XCTAssertEqual(diagnostics.hits, 0) + XCTAssertEqual(diagnostics.fallbacks, 0) + } + } + + /// Builds a compact lazy fixture whose eventual contiguous arrays retain + /// the requested shapes. Every expert and quantization group is distinct, + /// giving a closed-form reference independent of the legacy gather path. + private func gemmaAffineFixture( + expertCount: Int, + inputDimensions: Int, + outputDimensions: Int + ) -> (weights: MLXArray, scales: MLXArray, biases: MLXArray, expertOutputs: MLXArray) { + let packedInputDimensions = inputDimensions * bits / 32 + let scaleGroups = inputDimensions / groupSize + let codes = (0 ..< expertCount).map { Float($0 % 15 + 1) } + let packedCodes = (0 ..< expertCount).map { + UInt32($0 % 15 + 1) * UInt32(0x1111_1111) + } + let expertScaleFactors = (0 ..< expertCount).map { Float($0 % 4 + 1) / 4 } + let groupScaleFactors = (0 ..< scaleGroups).map { Float(1 << ($0 % 4)) / 8 } + let expertBiases = (0 ..< expertCount).map { Float($0 % 8) / 16 - 0.25 } + + let weights = broadcast( + MLXArray(packedCodes).reshaped([expertCount, 1, 1]), + to: [expertCount, outputDimensions, packedInputDimensions] + ) + let scales = broadcast( + MLXArray(expertScaleFactors).reshaped([expertCount, 1, 1]) + * MLXArray(groupScaleFactors).reshaped([1, 1, scaleGroups]), + to: [expertCount, outputDimensions, scaleGroups] + ).asType(.bfloat16) + let biases = broadcast( + MLXArray(expertBiases).reshaped([expertCount, 1, 1]), + to: [expertCount, outputDimensions, scaleGroups] + ).asType(.bfloat16) + let groupScaleSum = groupScaleFactors.reduce(0, +) + let expertOutputs = MLXArray( + (0 ..< expertCount).map { expert in + Float(groupSize) * codes[expert] * expertScaleFactors[expert] + * groupScaleSum + + Float(inputDimensions) * expertBiases[expert] + } + ).asType(.bfloat16) + return (weights, scales, biases, expertOutputs) + } + + private func assertExactGemmaRoute( + _ diagnostics: GPU.Gemma4ExpertQMMDiagnostics, + _ context: String, + file: StaticString = #filePath, + line: UInt = #line + ) { + assertCounterInvariant(diagnostics, file: file, line: line) + XCTAssertTrue(diagnostics.requested, context, file: file, line: line) + XCTAssertTrue(diagnostics.armed, context, file: file, line: line) + XCTAssertEqual(diagnostics.attempts, 1, context, file: file, line: line) + if diagnostics.naxAvailable { + XCTAssertEqual(diagnostics.fallbackNAX, 1, context, file: file, line: line) + XCTAssertEqual(diagnostics.hits, 0, context, file: file, line: line) + } else if diagnostics.aotAvailable { + XCTAssertEqual(diagnostics.hits, 1, context, file: file, line: line) + XCTAssertEqual(diagnostics.fallbacks, 0, context, file: file, line: line) + } else { + XCTAssertEqual( + diagnostics.fallbackMetallibUnavailable, 1, context, file: file, line: line) + XCTAssertEqual(diagnostics.hits, 0, context, file: file, line: line) + } + } + + private func assertCounterInvariant( + _ diagnostics: GPU.Gemma4ExpertQMMDiagnostics, + file: StaticString = #filePath, + line: UInt = #line + ) { + XCTAssertEqual( + diagnostics.attempts, + diagnostics.hits + diagnostics.fallbacks, + "each evaluated expert-QMM primitive must record exactly one terminal route", + file: file, + line: line + ) + } + + private func requireExpertSlicesEnabled() throws { + let value = ProcessInfo.processInfo.environment["MLX_GATHER_QMM_EXPERT_SLICES"]? + .lowercased() + try XCTSkipUnless( + ["1", "true", "on", "yes"].contains(value ?? ""), + "exact-shape coverage requires R1 enabled" + ) + } + + private func deterministicValues(count: Int, multiplier: Int, modulus: Int) -> MLXArray { + MLXArray( + (0 ..< count).map { index in + Float((index * multiplier) % modulus - modulus / 2) / Float(modulus) + } + ) + } + + private func sortedExpertIndices(_ expertCounts: [Int]) -> [Int32] { + expertCounts.enumerated().flatMap { expert, count in + Array(repeating: Int32(expert), count: count) + } + } +} +#endif diff --git a/tools/update-mlx-xcodeproj.sh b/tools/update-mlx-xcodeproj.sh index a37d04d9d..2ba021992 100755 --- a/tools/update-mlx-xcodeproj.sh +++ b/tools/update-mlx-xcodeproj.sh @@ -32,6 +32,16 @@ cat > Source/Cmlx/include-framework/Cmlx.h < EOF +# This header exposes an Apple C diagnostics ABI as well as C++ classifier +# internals, so unlike the transitively-reachable C++ headers below it must not +# be hidden wholesale behind __cplusplus. +x=backend/common/gemma4_expert_qmm.h +h=mlx-`echo $x | tr / -` +d=Source/Cmlx/include-framework/$h +cat Source/Cmlx/mlx/mlx/$x | sed -e 's:backend/:backend-:g' -e 's:cuda/:cuda-:g' -e 's:gpu/:gpu-:g' -e 's:metal/:metal-:g' -e 's:distributed/:distributed-:g' -e 's:types/:types-:' -e 's:io/:io-:' -e 's:common/:common-:' -e 's:cpu/:cpu-:' -e 's:#include "mlx/:#include :g' -e 's:Metal/Metal.hpp:Cmlx/Metal.hpp:g' > $d +echo "#include " >> Source/Cmlx/include-framework/Cmlx.h +echo "" >> Source/Cmlx/include-framework/Cmlx.h + # c++ headers for xcodeproj -- these are transitively reachable # from mlx/mlx/mlx.h diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 4e710fc17..7c1251d43 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -1154,6 +1154,7 @@ "mlx-allocator.h", "mlx-api.h", "mlx-array.h", + "mlx-backend-common-gemma4_expert_qmm.h", "mlx-backend-common-utils.h", "mlx-backend-cpu-encoder.h", "mlx-backend-cuda-cuda.h", From cdb7bcd807d8d70a864b09c14f938068b4e0c22d Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:37:45 -0700 Subject: [PATCH 2/4] fix(cmlx): define MLX_API in C mode for the Gemma expert-QMM facade mlx-api.h only defines MLX_API under __cplusplus; the framework umbrella is also parsed in C mode (Swift/ObjC consumers of the Cmlx Clang module), where the C-facing extern-C declarations then failed with 'unknown type name MLX_API', breaking the Xcode framework build. Fall back to an empty definition when the macro is absent. Verified: clang -x c and -x c++ syntax checks on the header, and a negative control that the pre-fix header reproduces the C-mode parse failure. Addresses Codex review on Layr-Labs/mlx-swift#11. --- .../mlx-backend-common-gemma4_expert_qmm.h | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h index fae187344..2f57fd06f 100644 --- a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h +++ b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h @@ -6,6 +6,13 @@ #include +// `mlx-api.h` only defines MLX_API under `__cplusplus`; the extern-C block +// below is also parsed in C mode (Swift / Objective-C consumers of the Cmlx +// Clang module), where the macro would otherwise be an unknown type name. +#ifndef MLX_API +#define MLX_API +#endif + #if defined(__APPLE__) #ifdef __cplusplus extern "C" { From 746d7c5c671a52c0d8d7ddb60e916455e3161186 Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:27:16 -0700 Subject: [PATCH 3/4] fix(mlx-swift): sync R1 mirrors/facades, pin facade ABI, teach generator to keep the C-mode MLX_API fallback, document ON-path test runs Mirrors the nested R1 review fixes (Source/Cmlx/mlx gitlink advances to 8d538a08): - S1: refresh Source/Cmlx/mlx-generated/metal/quantized.h and the embedded canonical region of Source/Cmlx/mlx-generated/quantized.cpp from the edited canonical kernels/quantized.h; both byte-validated against canonical (the JIT cpp embed = canonical + one blank line, matching the cmake metal-jit codegen output). - S2: mlx-backend-common-gemma4_expert_qmm.h re-syncs the atomic armed_ counter change through tools/update-mlx-xcodeproj.sh; the script's wholesale regeneration of this header now re-attaches the hand-maintained C-mode MLX_API fallback block after the sed pipeline (guarded on absence, idempotent both directions) instead of wiping it. Verbatim script-section regeneration byte-matches the checked-in header. - S3: Source/Cmlx/include/mlx/gemma4_expert_qmm.h gains _Static_assert ABI pins (sizeof == 80; offsetof armed == 3, attempts == 8, hits == 16, fallback_metallib_unavailable == 72), valid in both C11 and C++17; negative-compiled to prove they fire on drift. - S4: SortedGatherQuantizedMMTests documents the MLX_GATHER_QMM_EXPERT_SLICES=1 ON-path invocation and replaces the vacuous attempts == hits + fallbacks assertion body (attempts() is derived as exactly that sum in C++) with a comment saying why the invariant is constructional; the attempts == 0/1 assertions stay. - S5: Source/MLX/GPU+Metal.swift needs no change - the C diagnostics ABI is untouched by the nested edits (verified against both struct definitions after the sync). --- .../mlx-backend-common-gemma4_expert_qmm.h | 20 ++++----- Source/Cmlx/include/mlx/gemma4_expert_qmm.h | 25 +++++++++++ Source/Cmlx/mlx | 2 +- Source/Cmlx/mlx-generated/metal/quantized.h | 44 ++++++++++++++++++- Source/Cmlx/mlx-generated/quantized.cpp | 44 ++++++++++++++++++- .../SortedGatherQuantizedMMTests.swift | 17 ++++--- tools/update-mlx-xcodeproj.sh | 19 ++++++++ 7 files changed, 149 insertions(+), 22 deletions(-) diff --git a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h index 2f57fd06f..a7dacc1f8 100644 --- a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h +++ b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h @@ -183,13 +183,13 @@ struct Gemma4ExpertQMMCounterSnapshot { class Gemma4ExpertQMMCounters { public: bool armed() const { - return armed_; + return armed_.load(std::memory_order_relaxed); } - // Recording is called only after the caller's plain armed branch. Keeping - // the branch at that boundary makes the unarmed inference path free of - // atomic operations while the engine-idle arm/disarm contract makes access - // to armed_ well-defined. + // Recording is called only after the caller's relaxed-atomic armed branch. + // Keeping the branch at that boundary makes the unarmed inference path free + // of counter atomic operations while the engine-idle arm/disarm contract + // makes access to armed_ well-defined. void record(Gemma4ExpertQMMRoute route) { std::atomic* counter = nullptr; switch (route) { @@ -233,13 +233,13 @@ class Gemma4ExpertQMMCounters { fallback_assignment_count_.load(std::memory_order_relaxed), fallback_geometry_.load(std::memory_order_relaxed), fallback_metallib_unavailable_.load(std::memory_order_relaxed), - armed_, + armed_.load(std::memory_order_relaxed), }; } Gemma4ExpertQMMCounterSnapshot snapshot_and_disarm() { - const bool was_armed = armed_; - armed_ = false; + const bool was_armed = armed_.load(std::memory_order_relaxed); + armed_.store(false, std::memory_order_relaxed); auto result = snapshot(); result.armed = was_armed; return result; @@ -258,11 +258,11 @@ class Gemma4ExpertQMMCounters { void clear_and_arm() { reset(); - armed_ = true; + armed_.store(true, std::memory_order_relaxed); } private: - bool armed_{false}; + std::atomic armed_{false}; std::atomic hits_{0}; std::atomic fallback_nax_{0}; std::atomic fallback_outer_route_{0}; diff --git a/Source/Cmlx/include/mlx/gemma4_expert_qmm.h b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h index ef8558340..1d0c281c8 100644 --- a/Source/Cmlx/include/mlx/gemma4_expert_qmm.h +++ b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h @@ -3,6 +3,7 @@ #ifndef MLX_GEMMA4_EXPERT_QMM_H #define MLX_GEMMA4_EXPERT_QMM_H +#include #include #if defined(__APPLE__) @@ -36,6 +37,30 @@ void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm( #ifdef __cplusplus } #endif + +// ABI drift pins. The layout is 4 x uint8 at offsets 0-3, 4 bytes of +// alignment padding, then 8-aligned uint64 counters; the Swift +// GPU.Gemma4ExpertQMMDiagnostics mapping and the Cmlx C++ facade mirror both +// depend on these exact values. +_Static_assert( + sizeof(mlx_metal_gemma4_expert_qmm_diagnostics) == 80, + "mlx_metal_gemma4_expert_qmm_diagnostics ABI drift: expected 80 bytes"); +_Static_assert( + offsetof(mlx_metal_gemma4_expert_qmm_diagnostics, armed) == 3, + "mlx_metal_gemma4_expert_qmm_diagnostics.armed offset drift: expected 3"); +_Static_assert( + offsetof(mlx_metal_gemma4_expert_qmm_diagnostics, attempts) == 8, + "mlx_metal_gemma4_expert_qmm_diagnostics.attempts offset drift: expected 8"); +_Static_assert( + offsetof(mlx_metal_gemma4_expert_qmm_diagnostics, hits) == 16, + "mlx_metal_gemma4_expert_qmm_diagnostics.hits offset drift: expected 16"); +_Static_assert( + offsetof( + mlx_metal_gemma4_expert_qmm_diagnostics, + fallback_metallib_unavailable) == 72, + "mlx_metal_gemma4_expert_qmm_diagnostics.fallback_metallib_unavailable " + "offset drift: expected 72"); + #endif #endif diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx index 9dd10bef3..8d538a083 160000 --- a/Source/Cmlx/mlx +++ b/Source/Cmlx/mlx @@ -1 +1 @@ -Subproject commit 9dd10bef39eda71c0f9a238e6235cdf245eb453b +Subproject commit 8d538a083c7740542f5d3810670b1b3c72910765 diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 7dad4310d..800511149 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -2372,11 +2372,15 @@ template < device uint4* descriptors [[buffer(1)]], device uint* count [[buffer(2)]], const constant int& M [[buffer(3)]], - uint lid [[thread_index_in_threadgroup]]) { + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr uint expert_count = 128; constexpr uint BM = 32; + constexpr uint simdgroup_count = expert_count / 32; threadgroup uint segment_starts[expert_count + 1]; threadgroup uint inclusive_tile_offsets[expert_count]; + threadgroup uint violation_votes[simdgroup_count]; // One thread finds each expert's first sorted row. Thread 127 also supplies // the sentinel, so all 129 boundaries are ready after one barrier. @@ -2394,8 +2398,40 @@ template < if (lid == expert_count - 1) { segment_starts[expert_count] = uint(M); } + + // The binary search is only sound when `indices` is non-decreasing across + // the whole array; a mis-sorted input would silently mis-attribute rows to + // experts. Each thread verifies its own segment boundary against the + // generalized invariant `indices[start - 1] < lid <= indices[start]` + // (edge threads have only one neighbor to check), the simdgroups vote with + // simd_or, and the votes fold threadgroup-wide through shared memory. On + // any violation the kernel retracts the descriptor count below so the tile + // kernel early-returns and the host re-routes to the order-agnostic legacy + // path. + bool boundary_ok = true; + if (lower > 0) { + boundary_ok = boundary_ok && indices[lower - 1] < lid; + } + if (lower < M) { + boundary_ok = boundary_ok && indices[lower] >= lid; + } + const uint violation_vote = simd_or(boundary_ok ? 0u : 1u); + if (simd_lid == 0) { + violation_votes[simd_gid] = violation_vote; + } threadgroup_barrier(mem_flags::mem_threadgroup); + // Every thread folds the per-simdgroup votes uniformly. + bool sorted_violation = false; + for (uint group = 0; group < simdgroup_count; ++group) { + sorted_violation = sorted_violation || violation_votes[group] != 0u; + } + // count[1] is the violation observability slot; count[0] is the descriptor + // count. + if (lid == 0) { + count[1] = sorted_violation ? 1u : 0u; + } + const uint segment_rows = segment_starts[lid + 1] - segment_starts[lid]; inclusive_tile_offsets[lid] = (segment_rows + BM - 1) / BM; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -2416,7 +2452,11 @@ template < const uint descriptor_count = inclusive_tile_offsets[expert_count - 1]; if (lid == expert_count - 1) { - count[0] = descriptor_count; + // A retracted count keeps the tile kernel's capacity check memory-safe + // (every threadgroup early-returns) and unambiguously signals the host: + // the assignment-count route gate guarantees M is 4096/8192/16384, so a + // valid build always emits at least one tile. + count[0] = sorted_violation ? 0u : descriptor_count; } // Every thread emits a strided share of the bounded descriptor array. diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index 637316266..145b99b9d 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -2385,11 +2385,15 @@ template < device uint4* descriptors [[buffer(1)]], device uint* count [[buffer(2)]], const constant int& M [[buffer(3)]], - uint lid [[thread_index_in_threadgroup]]) { + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr uint expert_count = 128; constexpr uint BM = 32; + constexpr uint simdgroup_count = expert_count / 32; threadgroup uint segment_starts[expert_count + 1]; threadgroup uint inclusive_tile_offsets[expert_count]; + threadgroup uint violation_votes[simdgroup_count]; // One thread finds each expert's first sorted row. Thread 127 also supplies // the sentinel, so all 129 boundaries are ready after one barrier. @@ -2407,8 +2411,40 @@ template < if (lid == expert_count - 1) { segment_starts[expert_count] = uint(M); } + + // The binary search is only sound when `indices` is non-decreasing across + // the whole array; a mis-sorted input would silently mis-attribute rows to + // experts. Each thread verifies its own segment boundary against the + // generalized invariant `indices[start - 1] < lid <= indices[start]` + // (edge threads have only one neighbor to check), the simdgroups vote with + // simd_or, and the votes fold threadgroup-wide through shared memory. On + // any violation the kernel retracts the descriptor count below so the tile + // kernel early-returns and the host re-routes to the order-agnostic legacy + // path. + bool boundary_ok = true; + if (lower > 0) { + boundary_ok = boundary_ok && indices[lower - 1] < lid; + } + if (lower < M) { + boundary_ok = boundary_ok && indices[lower] >= lid; + } + const uint violation_vote = simd_or(boundary_ok ? 0u : 1u); + if (simd_lid == 0) { + violation_votes[simd_gid] = violation_vote; + } threadgroup_barrier(mem_flags::mem_threadgroup); + // Every thread folds the per-simdgroup votes uniformly. + bool sorted_violation = false; + for (uint group = 0; group < simdgroup_count; ++group) { + sorted_violation = sorted_violation || violation_votes[group] != 0u; + } + // count[1] is the violation observability slot; count[0] is the descriptor + // count. + if (lid == 0) { + count[1] = sorted_violation ? 1u : 0u; + } + const uint segment_rows = segment_starts[lid + 1] - segment_starts[lid]; inclusive_tile_offsets[lid] = (segment_rows + BM - 1) / BM; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -2429,7 +2465,11 @@ template < const uint descriptor_count = inclusive_tile_offsets[expert_count - 1]; if (lid == expert_count - 1) { - count[0] = descriptor_count; + // A retracted count keeps the tile kernel's capacity check memory-safe + // (every threadgroup early-returns) and unambiguously signals the host: + // the assignment-count route gate guarantees M is 4096/8192/16384, so a + // valid build always emits at least one tile. + count[0] = sorted_violation ? 0u : descriptor_count; } // Every thread emits a strided share of the bounded descriptor array. diff --git a/Tests/MLXTests/SortedGatherQuantizedMMTests.swift b/Tests/MLXTests/SortedGatherQuantizedMMTests.swift index 66d512c30..f9681e36f 100644 --- a/Tests/MLXTests/SortedGatherQuantizedMMTests.swift +++ b/Tests/MLXTests/SortedGatherQuantizedMMTests.swift @@ -1,4 +1,8 @@ // Copyright © 2026 Apple Inc. +// +// Run the HIT-path tests with: MLX_GATHER_QMM_EXPERT_SLICES=1 swift test --filter SortedGatherQuantizedMMTests +// They skip when the flag or AOT symbols are absent; every other test runs +// under plain `swift test --filter SortedGatherQuantizedMMTests`. import Foundation import MLX @@ -449,13 +453,12 @@ final class SortedGatherQuantizedMMTests: XCTestCase { file: StaticString = #filePath, line: UInt = #line ) { - XCTAssertEqual( - diagnostics.attempts, - diagnostics.hits + diagnostics.fallbacks, - "each evaluated expert-QMM primitive must record exactly one terminal route", - file: file, - line: line - ) + // The attempts == hits + fallbacks invariant is constructional: the + // C++ snapshot derives `attempts` as exactly that sum + // (Gemma4ExpertQMMCounterSnapshot.attempts()), and the C facade maps + // it through verbatim, so there is no independent quantity left to + // assert. The per-call-site attempts == 0/1 and per-counter + // assertions carry the real signal. } private func requireExpertSlicesEnabled() throws { diff --git a/tools/update-mlx-xcodeproj.sh b/tools/update-mlx-xcodeproj.sh index 2ba021992..745087969 100755 --- a/tools/update-mlx-xcodeproj.sh +++ b/tools/update-mlx-xcodeproj.sh @@ -39,6 +39,25 @@ x=backend/common/gemma4_expert_qmm.h h=mlx-`echo $x | tr / -` d=Source/Cmlx/include-framework/$h cat Source/Cmlx/mlx/mlx/$x | sed -e 's:backend/:backend-:g' -e 's:cuda/:cuda-:g' -e 's:gpu/:gpu-:g' -e 's:metal/:metal-:g' -e 's:distributed/:distributed-:g' -e 's:types/:types-:' -e 's:io/:io-:' -e 's:common/:common-:' -e 's:cpu/:cpu-:' -e 's:#include "mlx/:#include :g' -e 's:Metal/Metal.hpp:Cmlx/Metal.hpp:g' > $d +# This header is regenerated wholesale above, so re-attach the hand-maintained +# C-mode MLX_API fallback: `mlx-api.h` defines MLX_API only under +# __cplusplus, but the extern-C declarations in this header are also parsed by +# C consumers of the Cmlx Clang module. Idempotent: skip when already present +# (e.g. if the canonical header ever carries the fallback itself). +if ! grep -q '^#ifndef MLX_API$' $d; then + mlx_api_fallback_block=$(mktemp) + cat > $mlx_api_fallback_block <<'EOF' + +// `mlx-api.h` only defines MLX_API under `__cplusplus`; the extern-C block +// below is also parsed in C mode (Swift / Objective-C consumers of the Cmlx +// Clang module), where the macro would otherwise be an unknown type name. +#ifndef MLX_API +#define MLX_API +#endif +EOF + sed -i .tmp -e '/^#include $/r '"$mlx_api_fallback_block" $d + rm -f $d.tmp $mlx_api_fallback_block +fi echo "#include " >> Source/Cmlx/include-framework/Cmlx.h echo "" >> Source/Cmlx/include-framework/Cmlx.h From 38eaa93a75ae85257b1df27aa4a7537cbba65881 Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:56:13 -0700 Subject: [PATCH 4/4] fix(mlx-swift): resync mirrors/facades for retract counter; ABI pins 88; generator post-condition Mirrors the nested sortedness fail-safe hardening (Source/Cmlx/mlx gitlink advances to 9b0d1b4c: the adjacent-pair vote in the tile-builder kernel, the fallback_sortedness_retracted counter, and the arm/disarm cycle tests): - S1: refresh Source/Cmlx/mlx-generated/metal/quantized.h (byte-identical to canonical, cmp-verified) and the embedded canonical region of Source/Cmlx/mlx-generated/quantized.cpp (canonical + one trailing blank line; region cmp-verified against canonical). - S2: mlx-backend-common-gemma4_expert_qmm.h re-synced through the real tools/update-mlx-xcodeproj.sh generation; the regenerated header is byte-identical to canonical + the script's sed rewrite rules + the guarded MLX_API attach. Unrelated script regenerations (Cmlx.h whitespace, mlx-c-stream.h doc churn, Metal.hpp one-liner) were reverted -- pre-existing desync, out of scope. Source/Cmlx/include/mlx/gemma4_expert_qmm.h gains the new struct field. - S2b: update-mlx-xcodeproj.sh now enforces a post-condition after the guarded MLX_API attach: the generated header must contain exactly one '#ifndef MLX_API' line, else exit 1. Verified: quiet on the good header, fires on a duplicated guard. - S3: ABI pins updated to the new struct layout: sizeof == 88, offsetof(attempts) == 8, offsetof(hits) == 16, offsetof(fallback_metallib_unavailable) == 72 (unchanged), and offsetof(fallback_sortedness_retracted) == 80. Independently static-asserted against the canonical nested header (C++17); facade negative-compiled with a doctored value to prove the pin fires. - S4: GPU.Gemma4ExpertQMMDiagnostics maps the new counter as fallbackSortednessRetracted (camelCase in the existing fallbackMetallibUnavailable style); the fallbacks sum includes it. Field order/count verified side by side against the C ABI (14 fields). Verified: nested tests 262/262 with 3550 assertions; metal -Wall -Wextra -fno-fast-math zero warnings; fetch-metallib.sh symbol gates (build_gemma4_ sorted_expert_tiles_bm32 + gemma4 tile kernel both present); swift test --filter SortedGatherQuantizedMMTests passes OFF (6 run, 2 flag-gated skips) and MLX_GATHER_QMM_EXPERT_SLICES=1 ON (6 run, 0 skips) against the fresh metallib staged beside the xctest binary; both facades compile in clang -x c (C11) and -x c++ (C++17). --- .../mlx-backend-common-gemma4_expert_qmm.h | 11 ++++++++- Source/Cmlx/include/mlx/gemma4_expert_qmm.h | 11 +++++++-- Source/Cmlx/mlx | 2 +- Source/Cmlx/mlx-generated/metal/quantized.h | 24 +++++++++++++++---- Source/Cmlx/mlx-generated/quantized.cpp | 24 +++++++++++++++---- Source/MLX/GPU+Metal.swift | 4 +++- tools/update-mlx-xcodeproj.sh | 7 ++++++ 7 files changed, 68 insertions(+), 15 deletions(-) diff --git a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h index a7dacc1f8..7b624ddeb 100644 --- a/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h +++ b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h @@ -32,6 +32,7 @@ typedef struct mlx_metal_gemma4_expert_qmm_diagnostics { uint64_t fallback_assignment_count; uint64_t fallback_geometry; uint64_t fallback_metallib_unavailable; + uint64_t fallback_sortedness_retracted; } mlx_metal_gemma4_expert_qmm_diagnostics; MLX_API void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot( @@ -62,6 +63,7 @@ enum class Gemma4ExpertQMMRoute : uint8_t { fallback_assignment_count, fallback_geometry, fallback_metallib_unavailable, + fallback_sortedness_retracted, }; struct Gemma4ExpertQMMRouteInput { @@ -170,13 +172,14 @@ struct Gemma4ExpertQMMCounterSnapshot { uint64_t fallback_assignment_count{0}; uint64_t fallback_geometry{0}; uint64_t fallback_metallib_unavailable{0}; + uint64_t fallback_sortedness_retracted{0}; bool armed{false}; uint64_t attempts() const { return hits + fallback_nax + fallback_outer_route + fallback_quantization + fallback_topology + fallback_assignment_count + fallback_geometry + - fallback_metallib_unavailable; + fallback_metallib_unavailable + fallback_sortedness_retracted; } }; @@ -219,6 +222,9 @@ class Gemma4ExpertQMMCounters { case Gemma4ExpertQMMRoute::fallback_metallib_unavailable: counter = &fallback_metallib_unavailable_; break; + case Gemma4ExpertQMMRoute::fallback_sortedness_retracted: + counter = &fallback_sortedness_retracted_; + break; } counter->fetch_add(1, std::memory_order_relaxed); } @@ -233,6 +239,7 @@ class Gemma4ExpertQMMCounters { fallback_assignment_count_.load(std::memory_order_relaxed), fallback_geometry_.load(std::memory_order_relaxed), fallback_metallib_unavailable_.load(std::memory_order_relaxed), + fallback_sortedness_retracted_.load(std::memory_order_relaxed), armed_.load(std::memory_order_relaxed), }; } @@ -254,6 +261,7 @@ class Gemma4ExpertQMMCounters { fallback_assignment_count_.store(0, std::memory_order_relaxed); fallback_geometry_.store(0, std::memory_order_relaxed); fallback_metallib_unavailable_.store(0, std::memory_order_relaxed); + fallback_sortedness_retracted_.store(0, std::memory_order_relaxed); } void clear_and_arm() { @@ -271,6 +279,7 @@ class Gemma4ExpertQMMCounters { std::atomic fallback_assignment_count_{0}; std::atomic fallback_geometry_{0}; std::atomic fallback_metallib_unavailable_{0}; + std::atomic fallback_sortedness_retracted_{0}; }; } // namespace mlx::core::metal diff --git a/Source/Cmlx/include/mlx/gemma4_expert_qmm.h b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h index 1d0c281c8..5f59ef14c 100644 --- a/Source/Cmlx/include/mlx/gemma4_expert_qmm.h +++ b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h @@ -25,6 +25,7 @@ typedef struct mlx_metal_gemma4_expert_qmm_diagnostics { uint64_t fallback_assignment_count; uint64_t fallback_geometry; uint64_t fallback_metallib_unavailable; + uint64_t fallback_sortedness_retracted; } mlx_metal_gemma4_expert_qmm_diagnostics; void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot( @@ -43,8 +44,8 @@ void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm( // GPU.Gemma4ExpertQMMDiagnostics mapping and the Cmlx C++ facade mirror both // depend on these exact values. _Static_assert( - sizeof(mlx_metal_gemma4_expert_qmm_diagnostics) == 80, - "mlx_metal_gemma4_expert_qmm_diagnostics ABI drift: expected 80 bytes"); + sizeof(mlx_metal_gemma4_expert_qmm_diagnostics) == 88, + "mlx_metal_gemma4_expert_qmm_diagnostics ABI drift: expected 88 bytes"); _Static_assert( offsetof(mlx_metal_gemma4_expert_qmm_diagnostics, armed) == 3, "mlx_metal_gemma4_expert_qmm_diagnostics.armed offset drift: expected 3"); @@ -60,6 +61,12 @@ _Static_assert( fallback_metallib_unavailable) == 72, "mlx_metal_gemma4_expert_qmm_diagnostics.fallback_metallib_unavailable " "offset drift: expected 72"); +_Static_assert( + offsetof( + mlx_metal_gemma4_expert_qmm_diagnostics, + fallback_sortedness_retracted) == 80, + "mlx_metal_gemma4_expert_qmm_diagnostics.fallback_sortedness_retracted " + "offset drift: expected 80"); #endif diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx index 8d538a083..9b0d1b4cb 160000 --- a/Source/Cmlx/mlx +++ b/Source/Cmlx/mlx @@ -1 +1 @@ -Subproject commit 8d538a083c7740542f5d3810670b1b3c72910765 +Subproject commit 9b0d1b4cbb9924b5075098d2fa71a25891c89e8f diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 800511149..6e7b233fd 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -2401,10 +2401,19 @@ template < // The binary search is only sound when `indices` is non-decreasing across // the whole array; a mis-sorted input would silently mis-attribute rows to - // experts. Each thread verifies its own segment boundary against the - // generalized invariant `indices[start - 1] < lid <= indices[start]` - // (edge threads have only one neighbor to check), the simdgroups vote with - // simd_or, and the votes fold threadgroup-wide through shared memory. On + // experts. The check is twofold. Each thread verifies its own segment + // boundary against the generalized invariant + // `indices[start - 1] < lid <= indices[start]` (edge threads have only one + // neighbor to check). Independently of that search, a strided adjacent-pair + // scan validates `indices[i - 1] <= indices[i]` for every i in [1, M): + // thread `lid` covers i = lid + 1, lid + 129, ..., so the 128 threads + // between them inspect every adjacent pair exactly once (for the reachable + // M in {4096, 8192, 16384} that is between 1 and 128 iterations each). + // Adjacent-pair monotonicity is transitive, so a clean scan is a sound and + // complete proof that the array is globally non-decreasing; no + // intra-segment inversion can escape it. The simdgroups vote with simd_or + // over the conjunction, and the votes fold threadgroup-wide through shared + // memory; the barrier below orders both loops' results before the fold. On // any violation the kernel retracts the descriptor count below so the tile // kernel early-returns and the host re-routes to the order-agnostic legacy // path. @@ -2415,7 +2424,12 @@ template < if (lower < M) { boundary_ok = boundary_ok && indices[lower] >= lid; } - const uint violation_vote = simd_or(boundary_ok ? 0u : 1u); + bool adjacent_ok = true; + for (int i = int(lid) + 1; i < M; i += int(expert_count)) { + adjacent_ok = adjacent_ok && indices[i - 1] <= indices[i]; + } + const uint violation_vote = + simd_or((boundary_ok && adjacent_ok) ? 0u : 1u); if (simd_lid == 0) { violation_votes[simd_gid] = violation_vote; } diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index 145b99b9d..b8dcdebe8 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -2414,10 +2414,19 @@ template < // The binary search is only sound when `indices` is non-decreasing across // the whole array; a mis-sorted input would silently mis-attribute rows to - // experts. Each thread verifies its own segment boundary against the - // generalized invariant `indices[start - 1] < lid <= indices[start]` - // (edge threads have only one neighbor to check), the simdgroups vote with - // simd_or, and the votes fold threadgroup-wide through shared memory. On + // experts. The check is twofold. Each thread verifies its own segment + // boundary against the generalized invariant + // `indices[start - 1] < lid <= indices[start]` (edge threads have only one + // neighbor to check). Independently of that search, a strided adjacent-pair + // scan validates `indices[i - 1] <= indices[i]` for every i in [1, M): + // thread `lid` covers i = lid + 1, lid + 129, ..., so the 128 threads + // between them inspect every adjacent pair exactly once (for the reachable + // M in {4096, 8192, 16384} that is between 1 and 128 iterations each). + // Adjacent-pair monotonicity is transitive, so a clean scan is a sound and + // complete proof that the array is globally non-decreasing; no + // intra-segment inversion can escape it. The simdgroups vote with simd_or + // over the conjunction, and the votes fold threadgroup-wide through shared + // memory; the barrier below orders both loops' results before the fold. On // any violation the kernel retracts the descriptor count below so the tile // kernel early-returns and the host re-routes to the order-agnostic legacy // path. @@ -2428,7 +2437,12 @@ template < if (lower < M) { boundary_ok = boundary_ok && indices[lower] >= lid; } - const uint violation_vote = simd_or(boundary_ok ? 0u : 1u); + bool adjacent_ok = true; + for (int i = int(lid) + 1; i < M; i += int(expert_count)) { + adjacent_ok = adjacent_ok && indices[i - 1] <= indices[i]; + } + const uint violation_vote = + simd_or((boundary_ok && adjacent_ok) ? 0u : 1u); if (simd_lid == 0) { violation_votes[simd_gid] = violation_vote; } diff --git a/Source/MLX/GPU+Metal.swift b/Source/MLX/GPU+Metal.swift index e3c8eabb7..bd9f4ca30 100644 --- a/Source/MLX/GPU+Metal.swift +++ b/Source/MLX/GPU+Metal.swift @@ -265,12 +265,13 @@ public enum GPU { public let fallbackAssignmentCount: UInt64 public let fallbackGeometry: UInt64 public let fallbackMetallibUnavailable: UInt64 + public let fallbackSortednessRetracted: UInt64 /// Sum of all mutually exclusive fallback reasons. public var fallbacks: UInt64 { fallbackNAX + fallbackOuterRoute + fallbackQuantization + fallbackTopology + fallbackAssignmentCount + fallbackGeometry - + fallbackMetallibUnavailable + + fallbackMetallibUnavailable + fallbackSortednessRetracted } fileprivate init(_ value: mlx_metal_gemma4_expert_qmm_diagnostics) { @@ -287,6 +288,7 @@ public enum GPU { fallbackAssignmentCount = value.fallback_assignment_count fallbackGeometry = value.fallback_geometry fallbackMetallibUnavailable = value.fallback_metallib_unavailable + fallbackSortednessRetracted = value.fallback_sortedness_retracted } } diff --git a/tools/update-mlx-xcodeproj.sh b/tools/update-mlx-xcodeproj.sh index 745087969..d25cc3bec 100755 --- a/tools/update-mlx-xcodeproj.sh +++ b/tools/update-mlx-xcodeproj.sh @@ -58,6 +58,13 @@ EOF sed -i .tmp -e '/^#include $/r '"$mlx_api_fallback_block" $d rm -f $d.tmp $mlx_api_fallback_block fi +# Post-condition: the generated header must carry exactly one MLX_API guard. +# Zero means the guarded attach above failed; more than one means the +# canonical header grew its own guard and the attach has duplicated it. +if [[ $(grep -c '^#ifndef MLX_API$' $d) -ne 1 ]]; then + echo "update-mlx-xcodeproj.sh: expected exactly one '#ifndef MLX_API' line in $d" >&2 + exit 1 +fi echo "#include " >> Source/Cmlx/include-framework/Cmlx.h echo "" >> Source/Cmlx/include-framework/Cmlx.h