diff --git a/mlx/backend/common/gemma4_expert_qmm.h b/mlx/backend/common/gemma4_expert_qmm.h index 155260ad79..1a1ce37542 100644 --- a/mlx/backend/common/gemma4_expert_qmm.h +++ b/mlx/backend/common/gemma4_expert_qmm.h @@ -119,11 +119,11 @@ inline Gemma4ExpertQMMRoute classify_gemma4_expert_qmm( 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) { + !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; } const bool gemma4 = input.expert_count == 128; @@ -183,9 +183,8 @@ struct Gemma4ExpertQMMCounterSnapshot { bool armed{false}; uint64_t attempts() const { - return hits + fallback_nax + fallback_outer_route + - fallback_quantization + fallback_topology + - fallback_assignment_count + fallback_geometry + + return hits + fallback_nax + fallback_outer_route + fallback_quantization + + fallback_topology + fallback_assignment_count + fallback_geometry + fallback_metallib_unavailable + fallback_sortedness_retracted; } }; diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp index 54f51c5536..c2b9d66b3d 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -7,8 +7,8 @@ #include #include -#include #include +#include #include #include #include @@ -70,13 +70,13 @@ MetalAllocator::MetalAllocator(Device& d) // crash. The value may only LOWER the ceiling (it is clamped to the OS limit) // — raising it above what the hardware/OS reports would invite the very crash // this guards against. Strictly validated: a plain unsigned decimal that - // consumes the whole string, is non-zero, and does not overflow; anything else - // (empty, sign, junk, range error) is ignored and the OS limit stands. + // consumes the whole string, is non-zero, and does not overflow; anything + // else (empty, sign, junk, range error) is ignored and the OS limit stands. if (const char* rl = std::getenv("MLX_RESOURCE_LIMIT")) { while (*rl == ' ' || *rl == '\t') { ++rl; } - if (*rl >= '0' && *rl <= '9') { // unsigned decimal only (reject sign/junk) + if (*rl >= '0' && *rl <= '9') { // unsigned decimal only (reject sign/junk) errno = 0; char* end = nullptr; unsigned long long v = std::strtoull(rl, &end, 10); @@ -158,11 +158,13 @@ Buffer MetalAllocator::malloc(size_t size) { if (!buf) { size_t mem_required = get_active_memory() + get_cache_memory() + size; + auto pool = metal::new_scoped_memory_pool(); + // If we have a lot of memory pressure try to reclaim memory from the cache. - // NOTE: release_cached_buffers takes a BYTES-to-free target; when the buffers - // are tiny this frees only a few entries even though the COUNT is the binding - // constraint, so the byte path alone cannot bound num_resources_ (see the - // count-aware reclaim below). + // NOTE: release_cached_buffers takes a BYTES-to-free target; when the + // buffers are tiny this frees only a few entries even though the COUNT is + // the binding constraint, so the byte path alone cannot bound + // num_resources_ (see the count-aware reclaim below). if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) { num_resources_ -= buffer_cache_.release_cached_buffers(mem_required - gc_limit_); @@ -171,18 +173,18 @@ Buffer MetalAllocator::malloc(size_t size) { // Count-aware reclaim (Darkbloom): the Metal resource COUNT limit // (resource_limit_, ~iogpu.rsrc_limit/499000) is independent of byte usage. // Under churn with many distinct buffer shapes (varied prompt lengths, - // growing KV caches, multiple co-resident models) freed buffers are recycled - // into the size-keyed cache and never reused at that exact size, so the cache - // ENTRY COUNT creeps toward the limit while byte usage stays modest — the - // byte-driven trim above never fires (its threshold is ~physical RAM). Once - // the count crosses a high-water mark, proactively clear the cache (pure - // reuse pool — clearing only costs re-allocation, never correctness) so the - // count drops back to the live working set. This makes the count limit - // unreachable by any request mix / batching method, while the existing byte - // limits keep total memory below physical RAM. + // growing KV caches, multiple co-resident models) freed buffers are + // recycled into the size-keyed cache and never reused at that exact size, + // so the cache ENTRY COUNT creeps toward the limit while byte usage stays + // modest — the byte-driven trim above never fires (its threshold is + // ~physical RAM). Once the count crosses a high-water mark, proactively + // clear the cache (pure reuse pool — clearing only costs re-allocation, + // never correctness) so the count drops back to the live working set. This + // makes the count limit unreachable by any request mix / batching method, + // while the existing byte limits keep total memory below physical RAM. if (resource_limit_ > 0 && num_resources_ >= (resource_limit_ * resource_high_water_num_) / - resource_high_water_den_) { + resource_high_water_den_) { num_resources_ -= buffer_cache_.clear(); } diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 3282f45f0e..c4d7de0586 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -589,8 +589,7 @@ Device::Device() : device_(load_device()), residency_sets_(device_.get()) { auto pool = new_scoped_memory_pool(); default_library_ = NS::TransferPtr(load_default_library(device_.get())); - std::string expert_qmm_env = - env::get_var("MLX_GATHER_QMM_EXPERT_SLICES", ""); + std::string expert_qmm_env = env::get_var("MLX_GATHER_QMM_EXPERT_SLICES", ""); std::transform( expert_qmm_env.begin(), expert_qmm_env.end(), @@ -615,16 +614,14 @@ Device::Device() : device_(load_device()), residency_sets_(device_.get()) { "alN_true_bm_32_bn_32_bk_32"; auto has_default_function = [this](const char* name) { auto ns_name = NS::String::string(name, NS::ASCIIStringEncoding); - auto function = - NS::TransferPtr(default_library_->newFunction(ns_name)); + auto function = NS::TransferPtr(default_library_->newFunction(ns_name)); return function.get() != nullptr; }; // All expert-tile symbols ship from one source-matched metallib // (scripts/fetch-metallib.sh completeness contract), so availability is // all-or-nothing: a metallib missing any of them predates this revision // and must fail the whole route closed. - gemma4_expert_qmm_aot_available_ = - has_default_function(descriptor_kernel) && + gemma4_expert_qmm_aot_available_ = has_default_function(descriptor_kernel) && has_default_function(descriptor_kernel_e256) && has_default_function(tile_kernel); if (gemma4_expert_qmm_requested_ && gemma4_expert_qmm_aot_available_) { @@ -1043,8 +1040,7 @@ void gemma4_expert_qmm_diagnostics_snapshot( diagnostics->fallback_outer_route = counters.fallback_outer_route; diagnostics->fallback_quantization = counters.fallback_quantization; diagnostics->fallback_topology = counters.fallback_topology; - diagnostics->fallback_assignment_count = - counters.fallback_assignment_count; + diagnostics->fallback_assignment_count = counters.fallback_assignment_count; diagnostics->fallback_geometry = counters.fallback_geometry; diagnostics->fallback_metallib_unavailable = counters.fallback_metallib_unavailable; @@ -1080,8 +1076,7 @@ extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_clear_and_arm(void) { } } -extern "C" void -mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm( +extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot_and_disarm( mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics) { gemma4_expert_qmm_diagnostics_snapshot(diagnostics, true); } diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index 99b7caa7c7..83faa6ff33 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -10,6 +10,7 @@ #include #include "mlx/array.h" +#include "mlx/backend/common/gemma4_expert_qmm.h" #include "mlx/backend/common/metal_kernel.h" #include "mlx/backend/common/gemma4_expert_qmm.h" #include "mlx/backend/metal/resident.h" @@ -196,10 +197,6 @@ class MLX_API Device { const MTLFCList& func_consts = {}, const std::vector& linked_functions = {}); - ResidencySets& residency_sets() { - return residency_sets_; - } - bool gemma4_expert_qmm_requested() const { return gemma4_expert_qmm_requested_; } @@ -240,6 +237,10 @@ class MLX_API Device { gemma4_expert_qmm_counters_.clear_and_arm(); } + ResidencySets& residency_sets() { + return residency_sets_; + } + private: NS::SharedPtr build_library_( const std::string& source_string, diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6831cfd294..759b66d598 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -2615,8 +2615,7 @@ template 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); + const uint violation_vote = simd_or((boundary_ok && adjacent_ok) ? 0u : 1u); if (simd_lid == 0) { violation_votes[simd_gid] = violation_vote; } @@ -2650,8 +2649,7 @@ template threadgroup_barrier(mem_flags::mem_threadgroup); } - const uint descriptor_count = - inclusive_tile_offsets[expert_count - 1]; + 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: @@ -2666,8 +2664,7 @@ template uint expert_lower = 0; uint expert_upper = expert_count; while (expert_lower < expert_upper) { - const uint midpoint = - expert_lower + (expert_upper - expert_lower) / 2; + const uint midpoint = expert_lower + (expert_upper - expert_lower) / 2; if (inclusive_tile_offsets[midpoint] <= slot) { expert_lower = midpoint + 1; } else { @@ -2677,10 +2674,8 @@ template 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); + 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); } } @@ -2737,8 +2732,7 @@ template < 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; + reinterpret_cast(w) + expert * expert_w_stride; scales += expert * expert_sb_stride; biases += expert * expert_sb_stride; diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index ce5b45229c..c0696593bc 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -15,9 +15,9 @@ namespace mlx::core { +using metal::classify_gemma4_expert_qmm; using metal::Gemma4ExpertQMMRoute; using metal::Gemma4ExpertQMMRouteInput; -using metal::classify_gemma4_expert_qmm; namespace { @@ -1691,8 +1691,7 @@ Gemma4ExpertQMMRoute try_gemma4_expert_qmm( compute_encoder.set_bytes(N, c++); compute_encoder.dispatch_threadgroups( - MTL::Size((N + bn - 1) / bn, max_tile_count, 1), - MTL::Size(32, wn, wm)); + MTL::Size((N + bn - 1) / bn, max_tile_count, 1), MTL::Size(32, wn, wm)); return Gemma4ExpertQMMRoute::hit; } @@ -1721,8 +1720,7 @@ void gather_qmm_rhs( route_input.requested = true; route_input.outer_route = true; route_input.nax_available = true; - d.record_armed_gemma4_expert_qmm( - classify_gemma4_expert_qmm(route_input)); + d.record_armed_gemma4_expert_qmm(classify_gemma4_expert_qmm(route_input)); } return gather_qmm_rhs_nax( /* const array& x_ = */ x_, @@ -1798,8 +1796,7 @@ void gather_qmm_rhs( route_input.biases_contiguous = biases_ && biases_->flags().row_contiguous; route_input.group_size = group_size; route_input.bits = bits; - route_input.expert_count = - w.size() / w.shape(-1) / w.shape(-2); + route_input.expert_count = w.size() / w.shape(-1) / w.shape(-2); route_input.assignments = M; route_input.index_count = indices.size(); route_input.k = K; @@ -2078,11 +2075,8 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { // record below and the dispatch decision evaluate this one predicate so a // future tuning change cannot desynchronize them. // TODO: Tune 16 and 4 here a bit better. -static constexpr bool takes_sorted_rhs_route( - int M, - int B, - int E, - bool right_sorted) { +static constexpr bool +takes_sorted_rhs_route(int M, int B, int E, bool right_sorted) { return M == 1 && B >= 16 && right_sorted && B / E >= 4; } @@ -2116,8 +2110,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { Gemma4ExpertQMMRouteInput route_input; route_input.requested = true; route_input.outer_route = false; - d.record_armed_gemma4_expert_qmm( - classify_gemma4_expert_qmm(route_input)); + d.record_armed_gemma4_expert_qmm(classify_gemma4_expert_qmm(route_input)); } // We are walking x in order and w is also in order so we can batch up the diff --git a/tests/array_tests.cpp b/tests/array_tests.cpp index 68a4bed3a8..b33d328eeb 100644 --- a/tests/array_tests.cpp +++ b/tests/array_tests.cpp @@ -601,6 +601,7 @@ TEST_CASE("test array shared buffer") { array b = array(buf_b, shape, float32, deleter); eval(a + b); + synchronize(); } TEST_CASE("test make empty array") { diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index c0c3071a56..18018a4632 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -714,11 +714,10 @@ TEST_CASE("test layer norm vjp bias grad race") { CHECK(worst <= 1e-5); } - TEST_CASE("test Gemma 4 expert QMM pure route table") { + using metal::classify_gemma4_expert_qmm; using metal::Gemma4ExpertQMMRoute; using metal::Gemma4ExpertQMMRouteInput; - using metal::classify_gemma4_expert_qmm; auto gate_up = [](int assignments) { Gemma4ExpertQMMRouteInput input; @@ -787,8 +786,7 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { } auto exact = gate_up(4096); - auto check_miss = [&exact]( - auto mutate, Gemma4ExpertQMMRoute expected) { + auto check_miss = [&exact](auto mutate, Gemma4ExpertQMMRoute expected) { auto input = exact; mutate(input); CHECK(classify_gemma4_expert_qmm(input) == expected); @@ -815,8 +813,7 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { [](auto& x) { x.group_size = 32; }, Gemma4ExpertQMMRoute::fallback_quantization); check_miss( - [](auto& x) { x.bits = 8; }, - Gemma4ExpertQMMRoute::fallback_quantization); + [](auto& x) { x.bits = 8; }, Gemma4ExpertQMMRoute::fallback_quantization); check_miss( [](auto& x) { x.indices_uint32 = false; }, Gemma4ExpertQMMRoute::fallback_quantization); @@ -851,11 +848,9 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { [](auto& x) { x.expert_count = 127; }, Gemma4ExpertQMMRoute::fallback_topology); check_miss( - [](auto& x) { x.x_rank = 4; }, - Gemma4ExpertQMMRoute::fallback_topology); + [](auto& x) { x.x_rank = 4; }, Gemma4ExpertQMMRoute::fallback_topology); check_miss( - [](auto& x) { x.w_rank = 2; }, - Gemma4ExpertQMMRoute::fallback_topology); + [](auto& x) { x.w_rank = 2; }, Gemma4ExpertQMMRoute::fallback_topology); check_miss( [](auto& x) { x.scales_rank = 2; }, Gemma4ExpertQMMRoute::fallback_topology); @@ -875,11 +870,9 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { Gemma4ExpertQMMRoute::fallback_assignment_count); } check_miss( - [](auto& x) { x.w_dim2 = 176; }, - Gemma4ExpertQMMRoute::fallback_geometry); + [](auto& x) { x.w_dim2 = 176; }, Gemma4ExpertQMMRoute::fallback_geometry); check_miss( - [](auto& x) { x.w_dim1 += 1; }, - Gemma4ExpertQMMRoute::fallback_geometry); + [](auto& x) { x.w_dim1 += 1; }, Gemma4ExpertQMMRoute::fallback_geometry); check_miss( [](auto& x) { x.k += 32; @@ -887,8 +880,7 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { }, Gemma4ExpertQMMRoute::fallback_geometry); check_miss( - [](auto& x) { x.n -= 32; }, - Gemma4ExpertQMMRoute::fallback_geometry); + [](auto& x) { x.n -= 32; }, Gemma4ExpertQMMRoute::fallback_geometry); check_miss( [](auto& x) { x.aot_available = false; }, Gemma4ExpertQMMRoute::fallback_metallib_unavailable); @@ -902,9 +894,9 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { } TEST_CASE("test Qwen 3.6 expert QMM pure route table") { + using metal::classify_gemma4_expert_qmm; using metal::Gemma4ExpertQMMRoute; using metal::Gemma4ExpertQMMRouteInput; - using metal::classify_gemma4_expert_qmm; // Base input: Qwen 3.5/3.6 35B-A3B expert projection at W4/g64, // parametrized by whole-projection [E=256, n, k]. @@ -967,8 +959,7 @@ TEST_CASE("test Qwen 3.6 expert QMM pure route table") { } auto exact = qwen(4096, 2048, 1024); - auto check_miss = [&exact]( - auto mutate, Gemma4ExpertQMMRoute expected) { + auto check_miss = [&exact](auto mutate, Gemma4ExpertQMMRoute expected) { auto input = exact; mutate(input); CHECK(classify_gemma4_expert_qmm(input) == expected); @@ -998,11 +989,9 @@ TEST_CASE("test Qwen 3.6 expert QMM pure route table") { }, Gemma4ExpertQMMRoute::fallback_geometry); check_miss( - [](auto& x) { x.w_dim2 = 128; }, - Gemma4ExpertQMMRoute::fallback_geometry); + [](auto& x) { x.w_dim2 = 128; }, Gemma4ExpertQMMRoute::fallback_geometry); check_miss( - [](auto& x) { x.n -= 32; }, - Gemma4ExpertQMMRoute::fallback_geometry); + [](auto& x) { x.n -= 32; }, Gemma4ExpertQMMRoute::fallback_geometry); // T=128 chunks (1024 assignments) intentionally stay on the legacy path. for (int assignments : {8, 1024, 4095, 4097}) { check_miss( @@ -1014,8 +1003,7 @@ TEST_CASE("test Qwen 3.6 expert QMM pure route table") { Gemma4ExpertQMMRoute::fallback_assignment_count); } check_miss( - [](auto& x) { x.bits = 8; }, - Gemma4ExpertQMMRoute::fallback_quantization); + [](auto& x) { x.bits = 8; }, Gemma4ExpertQMMRoute::fallback_quantization); check_miss( [](auto& x) { x.aot_available = false; }, Gemma4ExpertQMMRoute::fallback_metallib_unavailable); @@ -1060,9 +1048,10 @@ TEST_CASE("test Gemma 4 expert QMM counter invariant") { counters.reset(); snapshot = counters.snapshot(); CHECK(snapshot.attempts() == 0); - CHECK(snapshot.attempts() == snapshot.hits + snapshot.fallback_nax + - snapshot.fallback_outer_route + snapshot.fallback_quantization + - snapshot.fallback_topology + + CHECK( + snapshot.attempts() == + snapshot.hits + snapshot.fallback_nax + snapshot.fallback_outer_route + + snapshot.fallback_quantization + snapshot.fallback_topology + snapshot.fallback_assignment_count + snapshot.fallback_geometry + snapshot.fallback_metallib_unavailable + snapshot.fallback_sortedness_retracted); @@ -1097,10 +1086,12 @@ TEST_CASE("test Gemma 4 expert QMM arm disarm cycle") { CHECK(interval.hits == 1); CHECK(interval.fallback_sortedness_retracted == 1); CHECK(interval.fallback_metallib_unavailable == 1); - CHECK(interval.attempts() == interval.hits + interval.fallback_nax + - interval.fallback_outer_route + interval.fallback_quantization + - interval.fallback_topology + interval.fallback_assignment_count + - interval.fallback_geometry + interval.fallback_metallib_unavailable + + CHECK( + interval.attempts() == + interval.hits + interval.fallback_nax + interval.fallback_outer_route + + interval.fallback_quantization + interval.fallback_topology + + interval.fallback_assignment_count + interval.fallback_geometry + + interval.fallback_metallib_unavailable + interval.fallback_sortedness_retracted); // The snapshot stays readable while disarmed.