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..7b624ddeb --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-backend-common-gemma4_expert_qmm.h @@ -0,0 +1,287 @@ +// Copyright © 2023-2024 Apple Inc. + +#pragma once + +#include + +#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" { +#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; + uint64_t fallback_sortedness_retracted; +} 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, + fallback_sortedness_retracted, +}; + +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}; + 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_sortedness_retracted; + } +}; + +class Gemma4ExpertQMMCounters { + public: + bool armed() const { + return armed_.load(std::memory_order_relaxed); + } + + // 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) { + 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; + case Gemma4ExpertQMMRoute::fallback_sortedness_retracted: + counter = &fallback_sortedness_retracted_; + 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), + fallback_sortedness_retracted_.load(std::memory_order_relaxed), + armed_.load(std::memory_order_relaxed), + }; + } + + Gemma4ExpertQMMCounterSnapshot snapshot_and_disarm() { + 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; + } + + 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); + fallback_sortedness_retracted_.store(0, std::memory_order_relaxed); + } + + void clear_and_arm() { + reset(); + armed_.store(true, std::memory_order_relaxed); + } + + private: + std::atomic 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}; + std::atomic fallback_sortedness_retracted_{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..5f59ef14c --- /dev/null +++ b/Source/Cmlx/include/mlx/gemma4_expert_qmm.h @@ -0,0 +1,73 @@ +// Copyright © 2026 Apple Inc. + +#ifndef MLX_GEMMA4_EXPERT_QMM_H +#define MLX_GEMMA4_EXPERT_QMM_H + +#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; + uint64_t fallback_sortedness_retracted; +} 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 + +// 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) == 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"); +_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"); +_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 + +#endif diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx index d5a240408..9b0d1b4cb 160000 --- a/Source/Cmlx/mlx +++ b/Source/Cmlx/mlx @@ -1 +1 @@ -Subproject commit d5a240408508f2be37f1a4893da0b415e8c0db55 +Subproject commit 9b0d1b4cbb9924b5075098d2fa71a25891c89e8f diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 12b5c85c6..6e7b233fd 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,232 @@ 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]], + 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. + 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); + } + + // 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. 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. + 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; + } + 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; + } + 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); + + // 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) { + // 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. + // 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..b8dcdebe8 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,232 @@ 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]], + 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. + 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); + } + + // 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. 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. + 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; + } + 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; + } + 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); + + // 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) { + // 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. + // 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..bd9f4ca30 100644 --- a/Source/MLX/GPU+Metal.swift +++ b/Source/MLX/GPU+Metal.swift @@ -245,4 +245,82 @@ 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 + public let fallbackSortednessRetracted: UInt64 + + /// Sum of all mutually exclusive fallback reasons. + public var fallbacks: UInt64 { + fallbackNAX + fallbackOuterRoute + fallbackQuantization + + fallbackTopology + fallbackAssignmentCount + fallbackGeometry + + fallbackMetallibUnavailable + fallbackSortednessRetracted + } + + 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 + fallbackSortednessRetracted = value.fallback_sortedness_retracted + } + } + + /// 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..f9681e36f --- /dev/null +++ b/Tests/MLXTests/SortedGatherQuantizedMMTests.swift @@ -0,0 +1,487 @@ +// 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 +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 + ) { + // 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 { + 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..d25cc3bec 100755 --- a/tools/update-mlx-xcodeproj.sh +++ b/tools/update-mlx-xcodeproj.sh @@ -32,6 +32,42 @@ 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 +# 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 +# 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 + # 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",