From d52134113ce48d6395d0d566773b1fe7ac4fdd80 Mon Sep 17 00:00:00 2001 From: anupsv <6407789+anupsv@users.noreply.github.com> Date: Mon, 15 Jun 2026 12:41:29 -0700 Subject: [PATCH 01/84] Bound Metal buffer COUNT, not just bytes, in MetalAllocator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Metal allocator throws `[metal::malloc] Resource limit (N) exceeded` when num_resources_ (the live+cached Metal buffer COUNT) reaches resource_limit_ (the iogpu.rsrc_limit sysctl, default ~499000). Freed buffers are recycled into a size-keyed cache whose only trim is by BYTES (release_cached_buffers takes a bytes-to-free target, max_pool_size_ ~= physical RAM). Under churn with many distinct buffer shapes (varied prompt lengths, growing KV caches, multiple co-resident models) the cache fills with entries never reused at that exact size, so the COUNT climbs to the limit while byte usage stays modest and the byte trim never fires — the process crashes mid-inference on a machine with most of its RAM free. malloc() now also reclaims by count: when num_resources_ crosses a 90% high-water mark of resource_limit_, it clears the (pure-reuse) buffer cache so the count drops back to the live working set. Clearing the cache only costs re-allocation, never correctness, so the count limit becomes unreachable by any request mix or batching method while the existing byte limits keep total memory bounded. Adds get_num_resources()/get_resource_limit() to the public memory API (metal + no_gpu + cuda backends) so the count and its ceiling are observable from callers. Adds an MLX_RESOURCE_LIMIT env override that can only LOWER the ceiling (clamped to the OS limit, strictly validated) to exercise the trim deterministically and as an operator safety valve. --- mlx/backend/cuda/allocator.cpp | 7 ++++ mlx/backend/metal/allocator.cpp | 57 +++++++++++++++++++++++++++++++- mlx/backend/metal/allocator.h | 19 +++++++++++ mlx/backend/no_gpu/allocator.cpp | 6 ++++ mlx/memory.h | 15 +++++++++ 5 files changed, 103 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index 718ae33e9c..7185a77e57 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -436,6 +436,13 @@ size_t get_memory_limit() { size_t get_cache_memory() { return cu::allocator().get_cache_memory(); } +size_t get_num_resources() { + // CUDA allocator does not track a Metal-style resource count; report 0. + return 0; +} +size_t get_resource_limit() { + return 0; +} size_t set_cache_limit(size_t limit) { return cu::allocator().set_cache_limit(limit); } diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp index c15bf3bdf7..ed033d7b5a 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -7,7 +7,10 @@ #include #include +#include +#include #include +#include namespace mlx::core { @@ -49,6 +52,30 @@ MetalAllocator::MetalAllocator() auto max_rec_size = std::get(info.at("max_recommended_working_set_size")); resource_limit_ = std::get(info.at("resource_limit")); + // Optional override (Darkbloom): MLX_RESOURCE_LIMIT lets an operator/test pin + // the Metal resource-COUNT ceiling below the OS default (~499000). Used to + // deterministically exercise the count-aware high-water trim, and as a safety + // valve to force earlier cache reclamation on a box seeing the resource-limit + // 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. + 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) + errno = 0; + char* end = nullptr; + unsigned long long v = std::strtoull(rl, &end, 10); + bool consumed_all = end != rl && *end == '\0'; + if (consumed_all && errno != ERANGE && v > 0 && + v <= std::numeric_limits::max()) { + resource_limit_ = std::min(static_cast(v), resource_limit_); + } + } + } block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize); gc_limit_ = std::min(static_cast(0.95 * max_rec_size), block_limit_); max_pool_size_ = block_limit_; @@ -130,12 +157,34 @@ Buffer MetalAllocator::malloc(size_t size) { auto pool = metal::new_scoped_memory_pool(); - // If we have a lot of memory pressure try to reclaim memory from the cache + // 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). if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) { num_resources_ -= buffer_cache_.release_cached_buffers(mem_required - gc_limit_); } + // 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. + if (resource_limit_ > 0 && + num_resources_ >= (resource_limit_ * resource_high_water_num_) / + resource_high_water_den_) { + num_resources_ -= buffer_cache_.clear(); + } + // Allocate new buffer if needed if (num_resources_ >= resource_limit_) { std::ostringstream msg; @@ -272,6 +321,12 @@ void reset_peak_memory() { size_t get_cache_memory() { return metal::allocator().get_cache_memory(); } +size_t get_num_resources() { + return metal::allocator().get_num_resources(); +} +size_t get_resource_limit() { + return metal::allocator().get_resource_limit(); +} void clear_cache() { return metal::allocator().clear_cache(); } diff --git a/mlx/backend/metal/allocator.h b/mlx/backend/metal/allocator.h index 5e177b3d3e..b8930a0469 100644 --- a/mlx/backend/metal/allocator.h +++ b/mlx/backend/metal/allocator.h @@ -37,6 +37,17 @@ class MetalAllocator : public allocator::Allocator { size_t get_cache_memory() { return buffer_cache_.cache_size(); }; + // Live Metal resource (buffer) COUNT and its hard ceiling. The count limit + // (default iogpu.rsrc_limit, ~499000) is independent of the byte limits and + // is what malloc() throws on when reached; exposed so callers can observe and + // bound it (it can be far higher than byte usage implies when many tiny + // buffers accumulate in the cache). + size_t get_num_resources() { + return num_resources_; + }; + size_t get_resource_limit() { + return resource_limit_; + }; size_t set_cache_limit(size_t limit); size_t set_memory_limit(size_t limit); size_t get_memory_limit(); @@ -71,6 +82,14 @@ class MetalAllocator : public allocator::Allocator { size_t num_resources_{0}; size_t resource_limit_{0}; + // Count-aware cache-reclaim high-water mark (Darkbloom): when num_resources_ + // reaches resource_high_water_num_/den_ of resource_limit_ (90%), malloc() + // proactively clears the (pure-reuse) buffer cache so the resource COUNT can + // never reach resource_limit_ and throw, regardless of buffer-byte sizes. + // Integer fraction to avoid float work in the allocation hot path. + static constexpr size_t resource_high_water_num_ = 9; + static constexpr size_t resource_high_water_den_ = 10; + std::mutex mutex_; }; diff --git a/mlx/backend/no_gpu/allocator.cpp b/mlx/backend/no_gpu/allocator.cpp index abb83e50e4..d970a4448d 100644 --- a/mlx/backend/no_gpu/allocator.cpp +++ b/mlx/backend/no_gpu/allocator.cpp @@ -123,6 +123,12 @@ size_t get_memory_limit() { size_t get_cache_memory() { return 0; } +size_t get_num_resources() { + return 0; +} +size_t get_resource_limit() { + return 0; +} size_t set_cache_limit(size_t) { return 0; } diff --git a/mlx/memory.h b/mlx/memory.h index f4eabc9976..4e82512ba9 100644 --- a/mlx/memory.h +++ b/mlx/memory.h @@ -33,6 +33,21 @@ MLX_API void reset_peak_memory(); * */ MLX_API size_t get_cache_memory(); +/* Get the number of live Metal resources (buffers). + * + * This is a COUNT, independent of byte usage. The Metal backend throws when it + * reaches the resource limit (see get_resource_limit). Many small cached + * buffers can push this count high while byte usage stays low. + * */ +MLX_API size_t get_num_resources(); + +/* Get the hard ceiling on the number of live Metal resources (buffers). + * + * Defaults to the iogpu.rsrc_limit sysctl (~499000 when unset). Allocation + * throws once get_num_resources() reaches this value. + * */ +MLX_API size_t get_resource_limit(); + /* Set the memory limit. * The memory limit is a guideline for the maximum amount of memory to use * during graph evaluation. If the memory limit is exceeded and there is no From aa480bd8ee48eccc3185267a14d4159e5b380a0c Mon Sep 17 00:00:00 2001 From: anupsv <6407789+anupsv@users.noreply.github.com> Date: Mon, 15 Jun 2026 12:41:29 -0700 Subject: [PATCH 02/84] Bound Metal buffer COUNT, not just bytes, in MetalAllocator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Metal allocator throws `[metal::malloc] Resource limit (N) exceeded` when num_resources_ (the live+cached Metal buffer COUNT) reaches resource_limit_ (the iogpu.rsrc_limit sysctl, default ~499000). Freed buffers are recycled into a size-keyed cache whose only trim is by BYTES (release_cached_buffers takes a bytes-to-free target, max_pool_size_ ~= physical RAM). Under churn with many distinct buffer shapes (varied prompt lengths, growing KV caches, multiple co-resident models) the cache fills with entries never reused at that exact size, so the COUNT climbs to the limit while byte usage stays modest and the byte trim never fires — the process crashes mid-inference on a machine with most of its RAM free. malloc() now also reclaims by count: when num_resources_ crosses a 90% high-water mark of resource_limit_, it clears the (pure-reuse) buffer cache so the count drops back to the live working set. Clearing the cache only costs re-allocation, never correctness, so the count limit becomes unreachable by any request mix or batching method while the existing byte limits keep total memory bounded. Adds get_num_resources()/get_resource_limit() to the public memory API (metal + no_gpu + cuda backends) so the count and its ceiling are observable from callers. Adds an MLX_RESOURCE_LIMIT env override that can only LOWER the ceiling (clamped to the OS limit, strictly validated) to exercise the trim deterministically and as an operator safety valve. --- mlx/backend/cuda/allocator.cpp | 7 ++++ mlx/backend/metal/allocator.cpp | 59 +++++++++++++++++++++++++++++++- mlx/backend/metal/allocator.h | 19 ++++++++++ mlx/backend/no_gpu/allocator.cpp | 6 ++++ mlx/memory.h | 15 ++++++++ 5 files changed, 105 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index b04d72da0e..36eb494e2d 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -422,6 +422,13 @@ size_t get_memory_limit() { size_t get_cache_memory() { return cu::allocator().get_cache_memory(); } +size_t get_num_resources() { + // CUDA allocator does not track a Metal-style resource count; report 0. + return 0; +} +size_t get_resource_limit() { + return 0; +} size_t set_cache_limit(size_t limit) { return cu::allocator().set_cache_limit(limit); } diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp index 60459c67c2..fd8b546995 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -7,8 +7,11 @@ #include #include +#include #include +#include #include +#include namespace mlx::core { @@ -60,6 +63,30 @@ MetalAllocator::MetalAllocator(Device& d) auto max_rec_size = std::get(info.at("max_recommended_working_set_size")); resource_limit_ = std::get(info.at("resource_limit")); + // Optional override (Darkbloom): MLX_RESOURCE_LIMIT lets an operator/test pin + // the Metal resource-COUNT ceiling below the OS default (~499000). Used to + // deterministically exercise the count-aware high-water trim, and as a safety + // valve to force earlier cache reclamation on a box seeing the resource-limit + // 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. + 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) + errno = 0; + char* end = nullptr; + unsigned long long v = std::strtoull(rl, &end, 10); + bool consumed_all = end != rl && *end == '\0'; + if (consumed_all && errno != ERANGE && v > 0 && + v <= std::numeric_limits::max()) { + resource_limit_ = std::min(static_cast(v), resource_limit_); + } + } + } block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize); gc_limit_ = std::min(static_cast(0.95 * max_rec_size), block_limit_); max_pool_size_ = block_limit_; @@ -131,12 +158,36 @@ Buffer MetalAllocator::malloc(size_t size) { if (!buf) { size_t mem_required = get_active_memory() + get_cache_memory() + size; - // If we have a lot of memory pressure try to reclaim memory from the cache + 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). if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) { num_resources_ -= buffer_cache_.release_cached_buffers(mem_required - gc_limit_); } + // 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. + if (resource_limit_ > 0 && + num_resources_ >= (resource_limit_ * resource_high_water_num_) / + resource_high_water_den_) { + num_resources_ -= buffer_cache_.clear(); + } + // Allocate new buffer if needed if (num_resources_ >= resource_limit_) { std::ostringstream msg; @@ -272,6 +323,12 @@ void reset_peak_memory() { size_t get_cache_memory() { return metal::allocator().get_cache_memory(); } +size_t get_num_resources() { + return metal::allocator().get_num_resources(); +} +size_t get_resource_limit() { + return metal::allocator().get_resource_limit(); +} void clear_cache() { return metal::allocator().clear_cache(); } diff --git a/mlx/backend/metal/allocator.h b/mlx/backend/metal/allocator.h index 4cbbfb0adc..9bf205ed06 100644 --- a/mlx/backend/metal/allocator.h +++ b/mlx/backend/metal/allocator.h @@ -36,6 +36,17 @@ class MetalAllocator : public allocator::Allocator { size_t get_cache_memory() { return buffer_cache_.cache_size(); }; + // Live Metal resource (buffer) COUNT and its hard ceiling. The count limit + // (default iogpu.rsrc_limit, ~499000) is independent of the byte limits and + // is what malloc() throws on when reached; exposed so callers can observe and + // bound it (it can be far higher than byte usage implies when many tiny + // buffers accumulate in the cache). + size_t get_num_resources() { + return num_resources_; + }; + size_t get_resource_limit() { + return resource_limit_; + }; size_t set_cache_limit(size_t limit); size_t set_memory_limit(size_t limit); size_t get_memory_limit(); @@ -72,6 +83,14 @@ class MetalAllocator : public allocator::Allocator { size_t num_resources_{0}; size_t resource_limit_{0}; + // Count-aware cache-reclaim high-water mark (Darkbloom): when num_resources_ + // reaches resource_high_water_num_/den_ of resource_limit_ (90%), malloc() + // proactively clears the (pure-reuse) buffer cache so the resource COUNT can + // never reach resource_limit_ and throw, regardless of buffer-byte sizes. + // Integer fraction to avoid float work in the allocation hot path. + static constexpr size_t resource_high_water_num_ = 9; + static constexpr size_t resource_high_water_den_ = 10; + std::mutex mutex_; }; diff --git a/mlx/backend/no_gpu/allocator.cpp b/mlx/backend/no_gpu/allocator.cpp index a800e381a7..3d2515c53d 100644 --- a/mlx/backend/no_gpu/allocator.cpp +++ b/mlx/backend/no_gpu/allocator.cpp @@ -203,6 +203,12 @@ size_t get_memory_limit() { size_t get_cache_memory() { return allocator::common_allocator().get_cache_memory(); } +size_t get_num_resources() { + return 0; +} +size_t get_resource_limit() { + return 0; +} size_t set_cache_limit(size_t limit) { return allocator::common_allocator().set_cache_limit(limit); } diff --git a/mlx/memory.h b/mlx/memory.h index f4eabc9976..4e82512ba9 100644 --- a/mlx/memory.h +++ b/mlx/memory.h @@ -33,6 +33,21 @@ MLX_API void reset_peak_memory(); * */ MLX_API size_t get_cache_memory(); +/* Get the number of live Metal resources (buffers). + * + * This is a COUNT, independent of byte usage. The Metal backend throws when it + * reaches the resource limit (see get_resource_limit). Many small cached + * buffers can push this count high while byte usage stays low. + * */ +MLX_API size_t get_num_resources(); + +/* Get the hard ceiling on the number of live Metal resources (buffers). + * + * Defaults to the iogpu.rsrc_limit sysctl (~499000 when unset). Allocation + * throws once get_num_resources() reaches this value. + * */ +MLX_API size_t get_resource_limit(); + /* Set the memory limit. * The memory limit is a guideline for the maximum amount of memory to use * during graph evaluation. If the memory limit is exceeded and there is no From a4b2b4c2a42444f4984893ae8ae901dd72f4ecaa Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:52:43 -0700 Subject: [PATCH 03/84] perf(mlx): opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder (#4) * perf(mlx): add opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder Adds a distinctly-named expert QMM implementation for the Gemma 4 26B-A4B MoE production shapes, gated by MLX_GATHER_QMM_EXPERT_SLICES: - qmm_t_expert_impl: BM32 expert tile body (BM16 fallback rows) taking a private/by-value row count; the shared qmm_t_impl constant-address ABI and all ordinary gathered/batched/dense QMM routes are unchanged. - build_gemma4_sorted_expert_tiles_bm32: one 128-thread threadgroup replaces the reference design's single-GPU-thread serial builder; parallel expert-range binary search, Hillis-Steele scan, and strided upper-bound descriptor emission. - Selector runs after the NAX-first route and requires affine BF16 transposed inputs, 4-bit gs=64 weights, 128 experts, assignment counts of exactly 4096/8192/16384, and the exact gate/up or down rank-3 shapes; every miss keeps the legacy route. NAX engagement is non-engagement, never bypassed. - device.{h,cpp}: one-shot request resolution, nonthrowing dual-symbol AOT probe/prewarm, relaxed-atomic diagnostics (requested, aotAvailable, naxAvailable, hits, per-class fallbacks). - gpu_tests: exact-shape arithmetic parity, fallback, and counter invariant probes. Retention standing (2026-08-09 production matrix): opt-in experiment. Standalone profile dropped (prefill -10.2% vs bracket); paired weighted-unsort+R1 profile retained-final (prefill +1.8%, TTFT -7.5%, decode +3.3%, arrival E2E +12.0%). NOTE: this source post-dates the benchmarked binaries/metallib (post-measurement kernel-body edit); rebuild and re-verify before any performance claim. * fix(mlx): fail-safe sortedness check in gemma expert tile builder; counter/atomic hygiene Review-wave fixes for the R1 expert-QMM path: - N1 (sortedness trust): build_gemma4_sorted_expert_tiles_bm32 now verifies each thread's post-binary-search segment boundary against the generalized invariant indices[start - 1] < lid <= indices[start] (edge threads check their single neighbor), votes per simdgroup via simd_or, folds the votes through threadgroup memory, and on any violation retracts count[0] to 0 (tile kernel then early-returns) and records the violation in count[1]; the buffer ABI is unchanged (count index 1 was previously unused). try_gemma4_expert_qmm allocates the second count element, drains the encoder after the builder, and re-routes a retracted call to the order-agnostic legacy path instead of dispatching the tile kernel (zero count is unambiguous: the selector's assignment gate guarantees M is 4096/8192/16384). - N2 (route-condition duplication): the sorted-RHS gate literal that appeared (negated) in the diagnostics record and in the dispatch decision is now the shared static constexpr predicate takes_sorted_rhs_route, so future tuning of the 16/4 thresholds cannot desynchronize counter vs route. - N3 (per-call bias normalization): gather_qmm_rhs no longer spends ensure_row_contiguous on biases before classification reads the raw tensor's fields; normalization runs only inside the winning-route branch (hit semantics unchanged; the legacy block keeps its own normalization point and ordering). - N4 (armed_ data race): Gemma4ExpertQMMCounters::armed_ is now std::atomic with relaxed loads/stores in armed(), snapshot(), snapshot_and_disarm() (read-then-write order preserved) and clear_and_arm(); the class remains non-copyable, now enforced. * fix(mlx): make the R1 sortedness fail-safe sound; proper retract attribution F1: the per-expert boundary vote was a partial detector -- an inversion inside a segment used by no other expert's boundary could escape, so "re-route on any violation" overclaimed. build_gemma4_sorted_expert_tiles_bm32 now also runs a strided adjacent-pair scan: thread lid checks indices[i-1] <= indices[i] for i = lid+1; i < M; i += 128, covering every adjacent pair in [1, M) exactly once (1..128 iterations at the reachable M in {4096,8192,16384}). Adjacent-pair monotonicity is transitive, so a clean scan is a sound and complete sortedness oracle; it folds into the same simd_or/threadgroup vote and the same retract (count[0]=0, count[1]=1). The boundary checks stay as cheap, precise diagnostics. F2: retracts were write-only in count[1] and surfaced as fallback_metallib_unavailable -- misattribution in the only observable surface. A dedicated fallback_sortedness_retracted counter now rides the GemmA4 route counters and the C diagnostics ABI (sizeof 80 -> 88, new uint64 at offset 80; existing offsets unchanged). try_gemma4_expert_qmm returns the route class: count[0]==0 with count[1]==1 records fallback_sortedness_retracted, any other unusable build keeps fallback_metallib_unavailable, then re-routes to the legacy path as before. F4: new doctest drives the full armed() -> clear_and_arm() -> snapshot_and_disarm() cycle and the attempts == hits + fallbacks invariant including the new class; the route-table and counter-invariant tests now cover fallback_sortedness_retracted. Verified: cmake tests 262/262 + 3550 assertions pass; metal -Wall -Wextra -fno-fast-math compile of kernels/quantized.metal is warning-free. --- mlx/backend/common/gemma4_expert_qmm.h | 280 +++++++++++++++++ mlx/backend/metal/device.cpp | 106 +++++++ mlx/backend/metal/device.h | 36 +++ mlx/backend/metal/kernels/quantized.h | 354 ++++++++++++++++++++++ mlx/backend/metal/kernels/quantized.metal | 15 +- mlx/backend/metal/quantized.cpp | 207 ++++++++++++- tests/gpu_tests.cpp | 273 +++++++++++++++++ 7 files changed, 1265 insertions(+), 6 deletions(-) create mode 100644 mlx/backend/common/gemma4_expert_qmm.h diff --git a/mlx/backend/common/gemma4_expert_qmm.h b/mlx/backend/common/gemma4_expert_qmm.h new file mode 100644 index 0000000000..21a0de634d --- /dev/null +++ b/mlx/backend/common/gemma4_expert_qmm.h @@ -0,0 +1,280 @@ +// Copyright © 2023-2024 Apple Inc. + +#pragma once + +#include + +#include "mlx/api.h" + +#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/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 2a8e15afd7..ccef4a4b0f 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -1,5 +1,7 @@ // Copyright © 2023-2024 Apple Inc. +#include +#include #include #include @@ -557,6 +559,42 @@ MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() { Device::Device() : device_(load_device()), residency_set_(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::transform( + expert_qmm_env.begin(), + expert_qmm_env.end(), + expert_qmm_env.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + gemma4_expert_qmm_requested_ = expert_qmm_env == "1" || + expert_qmm_env == "true" || expert_qmm_env == "on" || + expert_qmm_env == "yes"; + + constexpr const char* descriptor_kernel = + "build_gemma4_sorted_expert_tiles_bm32"; + constexpr const char* tile_kernel = + "affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_" + "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)); + return function.get() != nullptr; + }; + gemma4_expert_qmm_aot_available_ = + has_default_function(descriptor_kernel) && + has_default_function(tile_kernel); + if (gemma4_expert_qmm_requested_ && gemma4_expert_qmm_aot_available_) { + try { + // Resolve both pipelines once so missing or incompatible packaged AOT + // assets fail closed before an inference command encoder is touched. + get_kernel(descriptor_kernel); + get_kernel(tile_kernel); + } catch (...) { + gemma4_expert_qmm_aot_available_ = false; + } + } arch_ = env::metal_gpu_arch(); if (arch_.empty()) { arch_ = std::string(device_->architecture()->name()->utf8String()); @@ -932,3 +970,71 @@ bool is_nax_available() { } } // namespace mlx::core::metal + +#if defined(__APPLE__) +namespace { +void gemma4_expert_qmm_diagnostics_snapshot( + mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics, + bool disarm) { + if (diagnostics == nullptr) { + return; + } + *diagnostics = {}; + try { + auto& d = mlx::core::metal::device(mlx::core::Device::gpu); + const auto counters = disarm + ? d.gemma4_expert_qmm_counter_snapshot_and_disarm() + : d.gemma4_expert_qmm_counter_snapshot(); + diagnostics->requested = d.gemma4_expert_qmm_requested(); + diagnostics->aot_available = d.gemma4_expert_qmm_aot_available(); + diagnostics->nax_available = mlx::core::metal::is_nax_available(); + diagnostics->armed = counters.armed; + diagnostics->attempts = counters.attempts(); + diagnostics->hits = counters.hits; + diagnostics->fallback_nax = counters.fallback_nax; + 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_geometry = counters.fallback_geometry; + diagnostics->fallback_metallib_unavailable = + counters.fallback_metallib_unavailable; + diagnostics->fallback_sortedness_retracted = + counters.fallback_sortedness_retracted; + } catch (...) { + // Diagnostics are optional. A missing Metal device must remain observable + // as an all-zero snapshot rather than escaping an exception through C ABI. + } +} +} // namespace + +extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_snapshot( + mlx_metal_gemma4_expert_qmm_diagnostics* diagnostics) { + gemma4_expert_qmm_diagnostics_snapshot(diagnostics, false); +} + +extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_reset(void) { + try { + mlx::core::metal::device(mlx::core::Device::gpu) + .reset_gemma4_expert_qmm_counters(); + } catch (...) { + // Reset is best-effort on hosts without an accessible Metal device. + } +} + +extern "C" void mlx_metal_gemma4_expert_qmm_diagnostics_clear_and_arm(void) { + try { + mlx::core::metal::device(mlx::core::Device::gpu) + .clear_and_arm_gemma4_expert_qmm_counters(); + } catch (...) { + // Arming is best-effort on hosts without an accessible Metal device. + } +} + +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); +} +#endif diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index bed0cd636e..e45b257f21 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -11,6 +11,7 @@ #include #include "mlx/array.h" +#include "mlx/backend/common/gemma4_expert_qmm.h" #include "mlx/backend/metal/resident.h" #include "mlx/device.h" @@ -185,6 +186,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_; } @@ -226,6 +259,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/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 12b5c85c6a..6e7b233fdc 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index d632cfbda5..7c3fd51e1a 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -155,4 +155,17 @@ instantiate_quantized_groups(6) \ instantiate_quantized_groups(8) -instantiate_quantized_all() // clang-format on +instantiate_quantized_all() + +instantiate_kernel( + "affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_alN_true_bm_32_bn_32_bk_32", + affine_gather_qmm_gemma4_expert_tiles, + bfloat16_t, + 64, + 4, + true, + 32, + 32, + 32) + + // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index cfba4c0f8d..fec36db7eb 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -14,6 +14,10 @@ namespace mlx::core { +using metal::Gemma4ExpertQMMRoute; +using metal::Gemma4ExpertQMMRouteInput; +using metal::classify_gemma4_expert_qmm; + namespace { template @@ -1214,6 +1218,92 @@ void gather_qmm_rhs_nax( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +Gemma4ExpertQMMRoute try_gemma4_expert_qmm( + const array& x, + const array& w, + const array& scales, + const array& biases, + const array& indices, + array& out, + int M, + int N, + int K, + metal::Device& d, + const Stream& s) { + constexpr const char* descriptor_kernel_name = + "build_gemma4_sorted_expert_tiles_bm32"; + constexpr const char* tile_kernel_name = + "affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_" + "alN_true_bm_32_bn_32_bk_32"; + + MTL::ComputePipelineState* descriptor_kernel = nullptr; + MTL::ComputePipelineState* tile_kernel = nullptr; + try { + descriptor_kernel = d.get_kernel(descriptor_kernel_name); + tile_kernel = d.get_kernel(tile_kernel_name); + } catch (...) { + return Gemma4ExpertQMMRoute::fallback_metallib_unavailable; + } + + constexpr int bm = 32; + constexpr int bn = 32; + constexpr int wm = 2; + constexpr int wn = 2; + constexpr int expert_count = 128; + const int max_tile_count = (M + bm - 1) / bm + expert_count - 1; + + array descriptors({max_tile_count, 4}, uint32, nullptr, {}); + descriptors.set_data(allocator::malloc(descriptors.nbytes())); + // count[0] is the descriptor count; count[1] is the builder's sortedness + // violation observability slot. + array tile_count({2}, uint32, nullptr, {}); + tile_count.set_data(allocator::malloc(tile_count.nbytes())); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.add_temporary(descriptors); + compute_encoder.add_temporary(tile_count); + compute_encoder.set_compute_pipeline_state(descriptor_kernel); + compute_encoder.set_input_array(indices, 0); + compute_encoder.set_output_array(descriptors, 1); + compute_encoder.set_output_array(tile_count, 2); + compute_encoder.set_bytes(M, 3); + compute_encoder.dispatch_threads( + MTL::Size(expert_count, 1, 1), MTL::Size(expert_count, 1, 1)); + + // The descriptor builder fail-safes to count[0] == 0 when the purportedly + // sorted indices violate the non-decreasing invariant its binary search + // relies on. A zero count is unambiguous here: the selector's assignment + // gate guarantees M is one of 4096/8192/16384, so a valid build always + // emits at least one tile. Drain the stream and re-route retracted calls + // to the order-agnostic legacy path rather than running the tile kernel. + compute_encoder.synchronize(); + const uint32_t* counts = tile_count.data(); + if (counts[0] == 0) { + // count[1] flags a detected sortedness violation; attribute the retract + // to its own bucket. Any other unusable build keeps the metallib bucket. + return counts[1] == 1u + ? Gemma4ExpertQMMRoute::fallback_sortedness_retracted + : Gemma4ExpertQMMRoute::fallback_metallib_unavailable; + } + + compute_encoder.set_compute_pipeline_state(tile_kernel); + int c = 0; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(biases, c++); + compute_encoder.set_input_array(descriptors, c++); + compute_encoder.set_input_array(tile_count, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + 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)); + return Gemma4ExpertQMMRoute::hit; +} + void gather_qmm_rhs( const array& x_, const array& w_, @@ -1230,8 +1320,18 @@ void gather_qmm_rhs( metal::Device& d, const Stream& s, const std::string mode) { - if (metal::is_nax_available() && transpose && - (env::enable_tf32() || x_.dtype() != float32)) { + const bool nax_takes_call = metal::is_nax_available() && transpose && + (env::enable_tf32() || x_.dtype() != float32); + if (nax_takes_call) { + if (d.gemma4_expert_qmm_diagnostics_armed() && + d.gemma4_expert_qmm_requested()) { + Gemma4ExpertQMMRouteInput route_input; + 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)); + } return gather_qmm_rhs_nax( /* const array& x_ = */ x_, /* const array& w_ = */ w_, @@ -1275,6 +1375,83 @@ void gather_qmm_rhs( array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + if (d.gemma4_expert_qmm_requested()) { + auto shape_dim = [](const array& value, int axis) { + return value.ndim() > axis ? value.shape(axis) : 0; + }; + Gemma4ExpertQMMRouteInput route_input; + route_input.requested = true; + route_input.aot_available = d.gemma4_expert_qmm_aot_available(); + route_input.nax_available = false; + route_input.outer_route = true; + route_input.affine = mode == "affine"; + route_input.transpose = transpose; + route_input.has_bias = biases_.has_value(); + route_input.indices_uint32 = indices.dtype() == uint32; + route_input.indices_contiguous = indices.flags().row_contiguous; + route_input.x_bfloat16 = x.dtype() == bfloat16; + route_input.x_contiguous = x.flags().row_contiguous; + route_input.w_uint32 = w.dtype() == uint32; + route_input.w_contiguous = w.flags().row_contiguous; + route_input.scales_bfloat16 = scales.dtype() == bfloat16; + route_input.scales_contiguous = scales.flags().row_contiguous; + // Classification reads the raw bias tensor; normalization is spent only + // in the winning-route branch below. The legacy block at the end of this + // function retains its original normalization point and ordering. + route_input.biases_bfloat16 = biases_ && biases_->dtype() == bfloat16; + 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.assignments = M; + route_input.index_count = indices.size(); + route_input.k = K; + route_input.n = N; + route_input.x_rank = x.ndim(); + route_input.x_dim0 = shape_dim(x, 0); + route_input.x_dim1 = shape_dim(x, 1); + route_input.x_dim2 = shape_dim(x, 2); + route_input.w_rank = w.ndim(); + route_input.w_dim0 = shape_dim(w, 0); + route_input.w_dim1 = shape_dim(w, 1); + route_input.w_dim2 = shape_dim(w, 2); + route_input.scales_rank = scales.ndim(); + route_input.scales_dim0 = shape_dim(scales, 0); + route_input.scales_dim1 = shape_dim(scales, 1); + route_input.scales_dim2 = shape_dim(scales, 2); + if (biases_) { + route_input.biases_rank = biases_->ndim(); + route_input.biases_dim0 = shape_dim(*biases_, 0); + route_input.biases_dim1 = shape_dim(*biases_, 1); + route_input.biases_dim2 = shape_dim(*biases_, 2); + } + + auto route = classify_gemma4_expert_qmm(route_input); + if (route == Gemma4ExpertQMMRoute::hit) { + // A hit requires has_bias, so dereferencing biases_ is safe. + array expert_biases = ensure_row_contiguous(*biases_, d, s); + route = try_gemma4_expert_qmm( + x, w, scales, expert_biases, indices, out, M, N, K, d, s); + if (route == Gemma4ExpertQMMRoute::hit) { + if (d.gemma4_expert_qmm_diagnostics_armed()) { + d.record_armed_gemma4_expert_qmm(route); + } + return; + } + // A retracted build attributes to its own counter bucket + // (fallback_sortedness_retracted); missing AOT kernels keep the + // metallib bucket. In both cases the legacy route below produces the + // correct result. + } + if (d.gemma4_expert_qmm_diagnostics_armed()) { + d.record_armed_gemma4_expert_qmm(route); + } + } + + // Legacy gather path. Its normalization and dispatch order intentionally + // remain the global behavior for every non-exact call. + // TODO: Tune the block sizes int bm = 16, bn = 32, bk = 32; int wm = 1, wn = 2; @@ -1455,6 +1632,18 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { return; } +// Single source for the sorted-RHS expert route gate. Both the diagnostics +// 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) { + return M == 1 && B >= 16 && right_sorted && B / E >= 4; +} + void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); @@ -1479,11 +1668,19 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; auto mode = quantization_mode_to_string(mode_); + if (d.gemma4_expert_qmm_diagnostics_armed() && + d.gemma4_expert_qmm_requested() && + !takes_sorted_rhs_route(M, B, E, right_sorted_)) { + Gemma4ExpertQMMRouteInput route_input; + route_input.requested = true; + route_input.outer_route = false; + 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 // matmuls and reuse reading x and w. - // - // TODO: Tune 16 and 4 here a bit better. - if (M == 1 && B >= 16 && right_sorted_ == true && B / E >= 4) { + if (takes_sorted_rhs_route(M, B, E, right_sorted_)) { gather_qmm_rhs( x, w, diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index d920da49d9..b531d76c5f 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -6,6 +6,7 @@ #include #include "doctest/doctest.h" +#include "mlx/backend/common/gemma4_expert_qmm.h" #include "mlx/mlx.h" using namespace mlx::core; @@ -686,3 +687,275 @@ 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::Gemma4ExpertQMMRoute; + using metal::Gemma4ExpertQMMRouteInput; + using metal::classify_gemma4_expert_qmm; + + auto gate_up = [](int assignments) { + Gemma4ExpertQMMRouteInput input; + input.requested = true; + input.aot_available = true; + input.outer_route = true; + input.affine = true; + input.transpose = true; + input.has_bias = true; + input.indices_uint32 = true; + input.indices_contiguous = true; + input.x_bfloat16 = true; + input.x_contiguous = true; + input.w_uint32 = true; + input.w_contiguous = true; + input.scales_bfloat16 = true; + input.scales_contiguous = true; + input.biases_bfloat16 = true; + input.biases_contiguous = true; + input.group_size = 64; + input.bits = 4; + input.expert_count = 128; + input.assignments = assignments; + input.index_count = assignments; + input.k = 2816; + input.n = 1408; + input.x_rank = 3; + input.x_dim0 = assignments; + input.x_dim1 = 1; + input.x_dim2 = 2816; + input.w_rank = 3; + input.w_dim0 = 128; + input.w_dim1 = 1408; + input.w_dim2 = 352; + input.scales_rank = 3; + input.scales_dim0 = 128; + input.scales_dim1 = 1408; + input.scales_dim2 = 44; + input.biases_rank = 3; + input.biases_dim0 = 128; + input.biases_dim1 = 1408; + input.biases_dim2 = 44; + return input; + }; + auto down = [&gate_up](int assignments) { + auto input = gate_up(assignments); + input.k = 704; + input.n = 2816; + input.x_dim2 = 704; + input.w_dim1 = 2816; + input.w_dim2 = 88; + input.scales_dim1 = 2816; + input.scales_dim2 = 11; + input.biases_dim1 = 2816; + input.biases_dim2 = 11; + return input; + }; + + for (int assignments : {4096, 8192, 16384}) { + CHECK( + classify_gemma4_expert_qmm(gate_up(assignments)) == + Gemma4ExpertQMMRoute::hit); + CHECK( + classify_gemma4_expert_qmm(down(assignments)) == + Gemma4ExpertQMMRoute::hit); + } + + auto exact = gate_up(4096); + auto check_miss = [&exact]( + auto mutate, Gemma4ExpertQMMRoute expected) { + auto input = exact; + mutate(input); + CHECK(classify_gemma4_expert_qmm(input) == expected); + }; + check_miss( + [](auto& x) { x.requested = false; }, + Gemma4ExpertQMMRoute::not_requested); + check_miss( + [](auto& x) { x.nax_available = true; }, + Gemma4ExpertQMMRoute::fallback_nax); + check_miss( + [](auto& x) { x.outer_route = false; }, + Gemma4ExpertQMMRoute::fallback_outer_route); + check_miss( + [](auto& x) { x.affine = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.transpose = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.has_bias = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.group_size = 32; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.bits = 8; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.indices_uint32 = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.indices_contiguous = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.x_bfloat16 = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.x_contiguous = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.w_uint32 = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.w_contiguous = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.scales_bfloat16 = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.scales_contiguous = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.biases_bfloat16 = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.biases_contiguous = false; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.expert_count = 127; }, + Gemma4ExpertQMMRoute::fallback_topology); + check_miss( + [](auto& x) { x.x_rank = 4; }, + Gemma4ExpertQMMRoute::fallback_topology); + check_miss( + [](auto& x) { x.w_rank = 2; }, + Gemma4ExpertQMMRoute::fallback_topology); + check_miss( + [](auto& x) { x.scales_rank = 2; }, + Gemma4ExpertQMMRoute::fallback_topology); + check_miss( + [](auto& x) { x.biases_rank = 2; }, + Gemma4ExpertQMMRoute::fallback_topology); + check_miss( + [](auto& x) { x.index_count -= 1; }, + Gemma4ExpertQMMRoute::fallback_topology); + for (int assignments : {8, 16, 32, 4095, 4097}) { + check_miss( + [assignments](auto& x) { + x.assignments = assignments; + x.index_count = assignments; + x.x_dim0 = assignments; + }, + Gemma4ExpertQMMRoute::fallback_assignment_count); + } + check_miss( + [](auto& x) { x.w_dim2 = 176; }, + Gemma4ExpertQMMRoute::fallback_geometry); + check_miss( + [](auto& x) { x.w_dim1 += 1; }, + Gemma4ExpertQMMRoute::fallback_geometry); + check_miss( + [](auto& x) { + x.k += 32; + x.x_dim2 = x.k; + }, + Gemma4ExpertQMMRoute::fallback_geometry); + check_miss( + [](auto& x) { x.n -= 32; }, + Gemma4ExpertQMMRoute::fallback_geometry); + check_miss( + [](auto& x) { x.aot_available = false; }, + Gemma4ExpertQMMRoute::fallback_metallib_unavailable); + + auto nax_without_aot = exact; + nax_without_aot.nax_available = true; + nax_without_aot.aot_available = false; + CHECK( + classify_gemma4_expert_qmm(nax_without_aot) == + Gemma4ExpertQMMRoute::fallback_nax); +} + +TEST_CASE("test Gemma 4 expert QMM counter invariant") { + metal::Gemma4ExpertQMMCounters counters; + using Route = metal::Gemma4ExpertQMMRoute; + counters.record(Route::not_requested); + counters.record(Route::hit); + counters.record(Route::fallback_nax); + counters.record(Route::fallback_outer_route); + counters.record(Route::fallback_quantization); + counters.record(Route::fallback_topology); + counters.record(Route::fallback_assignment_count); + counters.record(Route::fallback_geometry); + counters.record(Route::fallback_metallib_unavailable); + counters.record(Route::fallback_sortedness_retracted); + + auto snapshot = counters.snapshot(); + CHECK(snapshot.hits == 1); + CHECK(snapshot.fallback_nax == 1); + CHECK(snapshot.fallback_outer_route == 1); + CHECK(snapshot.fallback_quantization == 1); + CHECK(snapshot.fallback_topology == 1); + CHECK(snapshot.fallback_assignment_count == 1); + CHECK(snapshot.fallback_geometry == 1); + CHECK(snapshot.fallback_metallib_unavailable == 1); + CHECK(snapshot.fallback_sortedness_retracted == 1); + CHECK(snapshot.attempts() == 9); + + 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 + + snapshot.fallback_assignment_count + snapshot.fallback_geometry + + snapshot.fallback_metallib_unavailable + + snapshot.fallback_sortedness_retracted); +} + +TEST_CASE("test Gemma 4 expert QMM arm disarm cycle") { + metal::Gemma4ExpertQMMCounters counters; + using Route = metal::Gemma4ExpertQMMRoute; + + // Counters start disarmed with an empty interval. + CHECK(!counters.armed()); + + // Arm: the interval opens with zeroed counters. + counters.clear_and_arm(); + CHECK(counters.armed()); + CHECK(counters.snapshot().attempts() == 0); + + // Record across the measured interval, including the retract class the + // sortedness fail-safe attributes mis-sorted indices to. + counters.record(Route::hit); + counters.record(Route::fallback_sortedness_retracted); + counters.record(Route::fallback_metallib_unavailable); + + // Disarm snapshots the interval and reports the previous armed state. + auto interval = counters.snapshot_and_disarm(); + CHECK(interval.armed); + CHECK(!counters.armed()); + + // The attempts == hits + sum(fallback classes) invariant holds across the + // cycle, with the sortedness-retract class included in the sum. + CHECK(interval.attempts() == 3); + 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 + + interval.fallback_sortedness_retracted); + + // The snapshot stays readable while disarmed. + CHECK(!counters.snapshot().armed); + CHECK(counters.snapshot().attempts() == 3); + + // Re-arming clears the interval again, and disarming it reports armed. + counters.clear_and_arm(); + CHECK(counters.armed()); + auto reopened = counters.snapshot_and_disarm(); + CHECK(reopened.armed); + CHECK(reopened.attempts() == 0); + CHECK(!counters.armed()); +} \ No newline at end of file From d3c82db012162f26206caf3864dae8e01274830c Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Thu, 13 Aug 2026 15:45:09 -0700 Subject: [PATCH 04/84] =?UTF-8?q?perf(metal):=20E=3D256=20expert-tile=20ro?= =?UTF-8?q?ute=20+=20trust=20+=20gpu::eval=20UAF=20fix=20=E2=80=94=20darkb?= =?UTF-8?q?loom-base=20mirror=20(#7)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(metal): instantiate E=256 expert-tile route for Qwen 3.5/3.6 MoE prefill (mirror of Cmlx/mlx 58fab469) * fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval (mirror) * perf(metal): trust mode skips retract readback (mirror) --- mlx/backend/common/gemma4_expert_qmm.h | 40 ++++--- mlx/backend/metal/device.cpp | 18 ++- mlx/backend/metal/device.h | 9 ++ mlx/backend/metal/eval.cpp | 12 +- mlx/backend/metal/kernels/quantized.h | 18 +-- mlx/backend/metal/kernels/quantized.metal | 14 +++ mlx/backend/metal/quantized.cpp | 45 ++++++-- tests/gpu_tests.cpp | 130 ++++++++++++++++++++++ 8 files changed, 250 insertions(+), 36 deletions(-) diff --git a/mlx/backend/common/gemma4_expert_qmm.h b/mlx/backend/common/gemma4_expert_qmm.h index 21a0de634d..155260ad79 100644 --- a/mlx/backend/common/gemma4_expert_qmm.h +++ b/mlx/backend/common/gemma4_expert_qmm.h @@ -126,11 +126,14 @@ inline Gemma4ExpertQMMRoute classify_gemma4_expert_qmm( !input.biases_contiguous || input.group_size != 64 || input.bits != 4) { return Gemma4ExpertQMMRoute::fallback_quantization; } - if (input.expert_count != 128 || input.x_rank != 3 || + const bool gemma4 = input.expert_count == 128; + const bool qwen36 = input.expert_count == 256; + if ((!gemma4 && !qwen36) || 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.x_dim2 != input.k || input.w_rank != 3 || + input.w_dim0 != input.expert_count || input.scales_rank != 3 || + input.scales_dim0 != input.expert_count || input.biases_rank != 3 || + input.biases_dim0 != input.expert_count || input.index_count != input.assignments) { return Gemma4ExpertQMMRoute::fallback_topology; } @@ -139,15 +142,26 @@ inline Gemma4ExpertQMMRoute classify_gemma4_expert_qmm( 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) { + // Whole-projection geometry for one expert matrix [E, n, k] at W4/g64: + // packed weight columns k/8 (eight 4-bit values per uint32) and one + // scale/bias per 64-wide group, k/64 columns. The quantization gate above + // guarantees bits==4 and group_size==64, so the divisions are exact. + auto projection = [&input](int k, int n) { + return input.k == k && input.n == n && input.w_dim1 == n && + input.w_dim2 == k / 8 && input.scales_dim1 == n && + input.scales_dim2 == k / 64 && input.biases_dim1 == n && + input.biases_dim2 == k / 64; + }; + // Gemma 4 26B-A4B (E=128): gate/up [128,1408,2816] and down [128,2816,704]. + // Qwen 3.5/3.6 35B-A3B (E=256): fused gate_up [256,1024,2048], split + // gate/up [256,512,2048], and down [256,2048,512]. The tile kernel itself + // is expert-count agnostic; only the descriptor builder instantiation + // differs (one thread per expert). + const bool hit_geometry = gemma4 + ? (projection(2816, 1408) || projection(704, 2816)) + : (projection(2048, 1024) || projection(2048, 512) || + projection(512, 2048)); + if (!hit_geometry) { return Gemma4ExpertQMMRoute::fallback_geometry; } if (!input.aot_available) { diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index ccef4a4b0f..030cfad2ac 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -569,10 +569,18 @@ Device::Device() : device_(load_device()), residency_set_(device_.get()) { [](unsigned char c) { return static_cast(std::tolower(c)); }); gemma4_expert_qmm_requested_ = expert_qmm_env == "1" || expert_qmm_env == "true" || expert_qmm_env == "on" || - expert_qmm_env == "yes"; + expert_qmm_env == "yes" || expert_qmm_env == "trust"; + // "trust" additionally skips the descriptor-retract readback: the caller + // asserts its sorted-indices contract is machine-guaranteed (the Swift + // SwitchGLU path sorts on-device), so the host never drains the stream to + // observe a retracted build. Under trust, a genuinely mis-sorted input + // produces undefined tile output instead of the legacy fallback. + gemma4_expert_qmm_trust_sorted_ = expert_qmm_env == "trust"; constexpr const char* descriptor_kernel = "build_gemma4_sorted_expert_tiles_bm32"; + constexpr const char* descriptor_kernel_e256 = + "build_sorted_expert_tiles_bm32_e256"; constexpr const char* tile_kernel = "affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_" "alN_true_bm_32_bn_32_bk_32"; @@ -582,14 +590,20 @@ Device::Device() : device_(load_device()), residency_set_(device_.get()) { 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) && + has_default_function(descriptor_kernel_e256) && has_default_function(tile_kernel); if (gemma4_expert_qmm_requested_ && gemma4_expert_qmm_aot_available_) { try { - // Resolve both pipelines once so missing or incompatible packaged AOT + // Resolve the pipelines once so missing or incompatible packaged AOT // assets fail closed before an inference command encoder is touched. get_kernel(descriptor_kernel); + get_kernel(descriptor_kernel_e256); get_kernel(tile_kernel); } catch (...) { gemma4_expert_qmm_aot_available_ = false; diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index e45b257f21..697b3ade07 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -190,6 +190,14 @@ class MLX_API Device { return gemma4_expert_qmm_requested_; } + // MLX_GATHER_QMM_EXPERT_SLICES=trust: skip the descriptor-retract + // readback in the expert-tile route (no mid-eval stream drain). The + // caller asserts sorted indices are machine-guaranteed; a violation + // yields undefined tile output instead of the legacy fallback. + bool gemma4_expert_qmm_trust_sorted() const { + return gemma4_expert_qmm_trust_sorted_; + } + bool gemma4_expert_qmm_aot_available() const { return gemma4_expert_qmm_aot_available_; } @@ -260,6 +268,7 @@ class MLX_API Device { std::unordered_map> library_map_; NS::SharedPtr default_library_; bool gemma4_expert_qmm_requested_{false}; + bool gemma4_expert_qmm_trust_sorted_{false}; bool gemma4_expert_qmm_aot_available_{false}; Gemma4ExpertQMMCounters gemma4_expert_qmm_counters_; std::unordered_map< diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp index e6826253de..fcc572b476 100644 --- a/mlx/backend/metal/eval.cpp +++ b/mlx/backend/metal/eval.cpp @@ -30,7 +30,6 @@ void eval(array& arr) { auto pool = metal::new_scoped_memory_pool(); auto s = arr.primitive().stream(); auto& encoder = metal::get_command_encoder(s); - auto* command_buffer = encoder.get_command_buffer(); auto outputs = arr.outputs(); { @@ -41,7 +40,8 @@ void eval(array& arr) { inputs = arr.inputs(); } - debug_set_primitive_buffer_label(command_buffer, arr.primitive()); + debug_set_primitive_buffer_label( + encoder.get_command_buffer(), arr.primitive()); arr.primitive().eval_gpu(arr.inputs(), outputs); } std::unordered_set> buffers; @@ -63,7 +63,13 @@ void eval(array& arr) { scheduler::notify_task_completion(s); }); } else { - command_buffer->addCompletedHandler( + // Fetch the command buffer AFTER eval_gpu: primitives that synchronize + // mid-eval (e.g. the expert-tile route's descriptor retract check) + // commit and REPLACE the encoder's buffer, so a pointer captured before + // eval_gpu would be stale here — attaching the buffer-liveness handler + // to it is a use-after-free. The current buffer holds the tail of this + // primitive's work, which is exactly what the inputs must outlive. + encoder.get_command_buffer()->addCompletedHandler( [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {}); } } diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6e7b233fdc..2d0f60b980 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -2367,7 +2367,11 @@ 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( +// Descriptor builder for the sorted expert-tile route. One thread per expert +// (threadgroup size == NE); instantiated for NE=128 (Gemma 4) and NE=256 +// (Qwen 3.5/3.6 MoE) in quantized.metal. +template +[[kernel]] void build_sorted_expert_tiles_bm32( const device uint32_t* indices [[buffer(0)]], device uint4* descriptors [[buffer(1)]], device uint* count [[buffer(2)]], @@ -2375,15 +2379,15 @@ template < 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 expert_count = uint(NE); 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. + // One thread finds each expert's first sorted row. The last thread also + // supplies the sentinel, so all NE+1 boundaries are ready after one barrier. int lower = 0; int upper = M; while (lower < upper) { @@ -2406,9 +2410,9 @@ template < // `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 + // thread `lid` covers i = lid + 1, lid + NE + 1, ..., so the NE 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). + // M in {4096, 8192, 16384} that is a bounded number of 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 @@ -2450,7 +2454,7 @@ template < 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 + // log2(NE) uniform Hillis-Steele strides form an inclusive scan over the // 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) { diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 7c3fd51e1a..2c7a59dff3 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -168,4 +168,18 @@ instantiate_kernel( 32, 32) +// Sorted expert-tile descriptor builders. The E=128 instantiation keeps the +// historical Gemma 4 host name; E=256 serves Qwen 3.5/3.6 MoE. The tile +// kernel instantiation above is expert-count agnostic (K/N are runtime +// arguments) and is shared by both routes. +instantiate_kernel( + "build_gemma4_sorted_expert_tiles_bm32", + build_sorted_expert_tiles_bm32, + 128) + +instantiate_kernel( + "build_sorted_expert_tiles_bm32_e256", + build_sorted_expert_tiles_bm32, + 256) + // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index fec36db7eb..3adf093311 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1230,8 +1230,14 @@ Gemma4ExpertQMMRoute try_gemma4_expert_qmm( int K, metal::Device& d, const Stream& s) { - constexpr const char* descriptor_kernel_name = - "build_gemma4_sorted_expert_tiles_bm32"; + // The classifier admits exactly E=128 (Gemma 4) and E=256 (Qwen 3.5/3.6 + // MoE); w is rank-3 [E, N_w, K_w] by the same gate. The tile kernel is + // expert-count agnostic — only the descriptor builder (one thread per + // expert) is instantiated per expert count. + const int expert_count = w.shape(0); + const char* descriptor_kernel_name = expert_count == 256 + ? "build_sorted_expert_tiles_bm32_e256" + : "build_gemma4_sorted_expert_tiles_bm32"; constexpr const char* tile_kernel_name = "affine_gather_qmm_gemma4_expert_tiles_bfloat16_t_gs_64_b_4_" "alN_true_bm_32_bn_32_bk_32"; @@ -1249,7 +1255,9 @@ Gemma4ExpertQMMRoute try_gemma4_expert_qmm( constexpr int bn = 32; constexpr int wm = 2; constexpr int wn = 2; - constexpr int expert_count = 128; + // Upper bound on descriptors: sum over experts of ceil(rows_e / bm) with + // sum(rows_e) == M. With k experts carrying a nonzero remainder the total + // is at most (M - k) / bm + k <= M / bm + E - 1 for the reachable M/E. const int max_tile_count = (M + bm - 1) / bm + expert_count - 1; array descriptors({max_tile_count, 4}, uint32, nullptr, {}); @@ -1276,14 +1284,29 @@ Gemma4ExpertQMMRoute try_gemma4_expert_qmm( // gate guarantees M is one of 4096/8192/16384, so a valid build always // emits at least one tile. Drain the stream and re-route retracted calls // to the order-agnostic legacy path rather than running the tile kernel. - compute_encoder.synchronize(); - const uint32_t* counts = tile_count.data(); - if (counts[0] == 0) { - // count[1] flags a detected sortedness violation; attribute the retract - // to its own bucket. Any other unusable build keeps the metallib bucket. - return counts[1] == 1u - ? Gemma4ExpertQMMRoute::fallback_sortedness_retracted - : Gemma4ExpertQMMRoute::fallback_metallib_unavailable; + // + // MLX_GATHER_QMM_EXPERT_SLICES=trust skips this drain entirely: the tile + // grid below is already over-dispatched (max_tile_count threadgroups; + // the kernel early-returns slots >= count[0]), so the readback exists + // ONLY to observe a retracted build. Under trust the caller asserts the + // sorted contract is machine-guaranteed (the Swift SwitchGLU prefill + // path sorts on-device just before this call); a genuine violation then + // yields undefined output for this matmul instead of the legacy result. + // Measured cost of the drain: ~120 stream drains per 512-token prefill + // chunk (3 gathers x 40 MoE layers) — it cancels the tile kernel's + // 12-23%/unit win end-to-end. A device-side legacy fallback (dispatch- + // diet item 1.3) would make trust the only behavior. + if (!d.gemma4_expert_qmm_trust_sorted()) { + compute_encoder.synchronize(); + const uint32_t* counts = tile_count.data(); + if (counts[0] == 0) { + // count[1] flags a detected sortedness violation; attribute the + // retract to its own bucket. Any other unusable build keeps the + // metallib bucket. + return counts[1] == 1u + ? Gemma4ExpertQMMRoute::fallback_sortedness_retracted + : Gemma4ExpertQMMRoute::fallback_metallib_unavailable; + } } compute_encoder.set_compute_pipeline_state(tile_kernel); diff --git a/tests/gpu_tests.cpp b/tests/gpu_tests.cpp index b531d76c5f..3cd4eca6a5 100644 --- a/tests/gpu_tests.cpp +++ b/tests/gpu_tests.cpp @@ -875,6 +875,136 @@ TEST_CASE("test Gemma 4 expert QMM pure route table") { Gemma4ExpertQMMRoute::fallback_nax); } +TEST_CASE("test Qwen 3.6 expert QMM pure route table") { + 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]. + auto qwen = [](int assignments, int k, int n) { + Gemma4ExpertQMMRouteInput input; + input.requested = true; + input.aot_available = true; + input.outer_route = true; + input.affine = true; + input.transpose = true; + input.has_bias = true; + input.indices_uint32 = true; + input.indices_contiguous = true; + input.x_bfloat16 = true; + input.x_contiguous = true; + input.w_uint32 = true; + input.w_contiguous = true; + input.scales_bfloat16 = true; + input.scales_contiguous = true; + input.biases_bfloat16 = true; + input.biases_contiguous = true; + input.group_size = 64; + input.bits = 4; + input.expert_count = 256; + input.assignments = assignments; + input.index_count = assignments; + input.k = k; + input.n = n; + input.x_rank = 3; + input.x_dim0 = assignments; + input.x_dim1 = 1; + input.x_dim2 = k; + input.w_rank = 3; + input.w_dim0 = 256; + input.w_dim1 = n; + input.w_dim2 = k / 8; + input.scales_rank = 3; + input.scales_dim0 = 256; + input.scales_dim1 = n; + input.scales_dim2 = k / 64; + input.biases_rank = 3; + input.biases_dim0 = 256; + input.biases_dim1 = n; + input.biases_dim2 = k / 64; + return input; + }; + + // Fused gate_up, split gate/up, and down projections hit at the chunked + // prefill assignment counts (T x top-8 for T in {512, 1024, 2048}). + for (int assignments : {4096, 8192, 16384}) { + CHECK( + classify_gemma4_expert_qmm(qwen(assignments, 2048, 1024)) == + Gemma4ExpertQMMRoute::hit); + CHECK( + classify_gemma4_expert_qmm(qwen(assignments, 2048, 512)) == + Gemma4ExpertQMMRoute::hit); + CHECK( + classify_gemma4_expert_qmm(qwen(assignments, 512, 2048)) == + Gemma4ExpertQMMRoute::hit); + } + + auto exact = qwen(4096, 2048, 1024); + auto check_miss = [&exact]( + auto mutate, Gemma4ExpertQMMRoute expected) { + auto input = exact; + mutate(input); + CHECK(classify_gemma4_expert_qmm(input) == expected); + }; + // Expert counts other than the two instantiated builders miss on topology. + check_miss( + [](auto& x) { + x.expert_count = 255; + x.w_dim0 = 255; + x.scales_dim0 = 255; + x.biases_dim0 = 255; + }, + Gemma4ExpertQMMRoute::fallback_topology); + // E=256 with Gemma geometry (and vice versa) must miss on geometry: the + // shape table is tied to the expert count, never mixed. + check_miss( + [](auto& x) { + x.k = 2816; + x.n = 1408; + x.x_dim2 = 2816; + x.w_dim1 = 1408; + x.w_dim2 = 352; + x.scales_dim1 = 1408; + x.scales_dim2 = 44; + x.biases_dim1 = 1408; + x.biases_dim2 = 44; + }, + Gemma4ExpertQMMRoute::fallback_geometry); + check_miss( + [](auto& x) { x.w_dim2 = 128; }, + Gemma4ExpertQMMRoute::fallback_geometry); + check_miss( + [](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( + [assignments](auto& x) { + x.assignments = assignments; + x.index_count = assignments; + x.x_dim0 = assignments; + }, + Gemma4ExpertQMMRoute::fallback_assignment_count); + } + check_miss( + [](auto& x) { x.bits = 8; }, + Gemma4ExpertQMMRoute::fallback_quantization); + check_miss( + [](auto& x) { x.aot_available = false; }, + Gemma4ExpertQMMRoute::fallback_metallib_unavailable); + + // The Gemma table must also reject Qwen geometry under E=128. + auto gemma_with_qwen_geometry = exact; + gemma_with_qwen_geometry.expert_count = 128; + gemma_with_qwen_geometry.w_dim0 = 128; + gemma_with_qwen_geometry.scales_dim0 = 128; + gemma_with_qwen_geometry.biases_dim0 = 128; + CHECK( + classify_gemma4_expert_qmm(gemma_with_qwen_geometry) == + Gemma4ExpertQMMRoute::fallback_geometry); +} + TEST_CASE("test Gemma 4 expert QMM counter invariant") { metal::Gemma4ExpertQMMCounters counters; using Route = metal::Gemma4ExpertQMMRoute; From 9b6575c35434efa588fff7e408ad55bb7617c62d Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:15:06 +0800 Subject: [PATCH 05/84] Return tuple in meshgrid (#4229) --- python/src/ops.cpp | 6 +++--- python/tests/test_ops.py | 5 +++++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index d8892a45d2..bdcf3dde4b 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3386,14 +3386,14 @@ void init_ops(nb::module_& m) { mx::StreamOrDevice s) { std::vector arrays = nb::cast>(arrays_); - return mx::meshgrid(arrays, sparse, indexing, s); + return nb::tuple(nb::cast(mx::meshgrid(arrays, sparse, indexing, s))); }, "arrays"_a, "sparse"_a = false, "indexing"_a = "xy", "stream"_a = nb::none(), nb::sig( - "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> array"), + "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> tuple[array, ...]"), R"pbdoc( Generate multidimensional coordinate grids from 1-D coordinate arrays @@ -3406,7 +3406,7 @@ void init_ops(nb::module_& m) { Defaults to ``'xy'``. Returns: - list(array): The output arrays. + tuple(array): The output arrays. )pbdoc"); m.def( "repeat", diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 1b46237ce5..83f95ea5fc 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2132,6 +2132,11 @@ def test_meshgrid(self): x = mx.array([1, 2, 3], dtype=mx.int32) y = np.array([1, 2, 3], dtype=np.int32) + # Test return type is a tuple + self.assertIsInstance(mx.meshgrid(x), tuple) + self.assertIsInstance(mx.meshgrid(x, x), tuple) + self.assertIsInstance(mx.meshgrid(x, x, x, sparse=True), tuple) + # Test single input a_mlx = mx.meshgrid(x) a_np = np.meshgrid(y) From 1d717bd3c562a45e6c0f6e195d413ea16d22898f Mon Sep 17 00:00:00 2001 From: AK <144495202+AKnassa@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:38:31 -0400 Subject: [PATCH 06/84] Add endpoint parameter to linspace (#4184) Co-authored-by: Cheng --- mlx/ops.cpp | 6 ++++- mlx/ops.h | 16 +++++++++-- python/src/ops.cpp | 8 +++++- python/tests/test_double.py | 8 +++++- python/tests/test_ops.py | 53 ++++++++++++++++++++++++++++++++++++- tests/ops_tests.cpp | 18 ++++++++++++- 6 files changed, 102 insertions(+), 7 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 9c8db3a26d..ef229fe87f 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -271,6 +271,7 @@ array linspace( double start, double stop, int num /* = 50 */, + bool endpoint /* = true */, Dtype dtype /* = float32 */, StreamOrDevice s /* = {} */) { if (num < 0) { @@ -282,8 +283,11 @@ array linspace( return astype(array({start}), dtype, s); } auto inner_type = dtype == float64 ? float64 : float32; + // Without the endpoint the samples are spaced so that `stop` would be the + // next one after the last, i.e. the step is (stop - start) / num. + auto denominator = endpoint ? num - 1 : num; array t = - divide(arange(0, num, inner_type, s), array(num - 1, inner_type), s); + divide(arange(0, num, inner_type, s), array(denominator, inner_type), s); array t_bar = subtract(array(1, inner_type), t, s); return astype( add(multiply(t_bar, array(start, inner_type), s), diff --git a/mlx/ops.h b/mlx/ops.h index 01e0a99286..f597753b1e 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -38,13 +38,25 @@ MLX_API array arange(int start, int stop, int step, StreamOrDevice s = {}); MLX_API array arange(int start, int stop, StreamOrDevice s = {}); MLX_API array arange(int stop, StreamOrDevice s = {}); -/** A 1D array of `num` evenly spaced numbers in the range `[start, stop]` */ +/** + * A 1D array of `num` evenly spaced numbers in the range `[start, stop]`, or + * in the half-open range `[start, stop)` when `endpoint` is false. + */ MLX_API array linspace( double start, double stop, - int num = 50, + int num, + bool endpoint, Dtype dtype = float32, StreamOrDevice s = {}); +inline array linspace( + double start, + double stop, + int num = 50, + Dtype dtype = float32, + StreamOrDevice s = {}) { + return linspace(start, stop, num, true, dtype, s); +} /** Convert an array to the given data type. */ MLX_API array astype(array a, Dtype dtype, StreamOrDevice s = {}); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index bdcf3dde4b..4dc5114bf7 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1644,22 +1644,25 @@ void init_ops(nb::module_& m) { [](Scalar start, Scalar stop, int num, + bool endpoint, std::optional dtype, mx::StreamOrDevice s) { return mx::linspace( scalar_to_double(start), scalar_to_double(stop), num, + endpoint, dtype.value_or(mx::float32), s); }, "start"_a, "stop"_a, "num"_a = 50, + "endpoint"_a = true, "dtype"_a.none() = mx::float32, "stream"_a = nb::none(), nb::sig( - "def linspace(start: scalar, stop: scalar, num: int | None = 50, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"), + "def linspace(start: scalar, stop: scalar, num: int | None = 50, endpoint: bool = True, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate ``num`` evenly spaced numbers over interval ``[start, stop]``. @@ -1667,6 +1670,9 @@ void init_ops(nb::module_& m) { start (scalar): Starting value. stop (scalar): Stopping value. num (int, optional): Number of samples, defaults to ``50``. + endpoint (bool, optional): If ``True``, ``stop`` is the last + sample. Otherwise it is not included and the samples are spaced + over the half-open interval ``[start, stop)``. Default: ``True``. dtype (Dtype, optional): Specifies the data type of the output, default to ``float32``. diff --git a/python/tests/test_double.py b/python/tests/test_double.py index 65603cd937..3186e7e573 100644 --- a/python/tests/test_double.py +++ b/python/tests/test_double.py @@ -336,9 +336,15 @@ def test_python_float_keeps_double_precision(self): def test_linspace(self): with mx.stream(mx.cpu): - vals = mx.linspace(0, math.pi, 2, mx.float64) + vals = mx.linspace(0, math.pi, 2, dtype=mx.float64) self.assertEqual(vals.tolist()[1], math.pi) + vals = mx.linspace(0, math.pi, 4, endpoint=False, dtype=mx.float64) + self.assertEqual(vals.dtype, mx.float64) + self.assertTrue( + np.allclose(vals.tolist(), np.linspace(0, math.pi, 4, endpoint=False)) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 83f95ea5fc..cdf0bc1b22 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2950,7 +2950,7 @@ def test_linspace(self): self.assertEqualArray(a, expected) # Test int64 dtype - b = mx.linspace(0, 10, 5, mx.int64) + b = mx.linspace(0, 10, 5, dtype=mx.int64) expected = mx.array(np.linspace(0, 10, 5, dtype=int)) self.assertEqualArray(b, expected) @@ -2977,6 +2977,57 @@ def test_linspace(self): self.assertEqual(d[0], a) self.assertEqual(d[-1], b) + def test_linspace_endpoint(self): + # endpoint=True is the default and matches the old behaviour + a = mx.linspace(0, 1, 5, endpoint=True) + self.assertEqualArray(a, mx.array(np.linspace(0, 1, 5, endpoint=True))) + self.assertEqualArray(a, mx.linspace(0, 1, 5)) + + # endpoint=False drops the stop value and uses a step of + # (stop - start) / num instead of (stop - start) / (num - 1) + for num in [0, 1, 2, 5, 50]: + b = mx.linspace(0, 10, num, endpoint=False) + expected = mx.array(np.linspace(0, 10, num, endpoint=False)) + self.assertEqualArray(b, expected) + + c = mx.linspace(-2.7, -0.7, 7, endpoint=False) + self.assertEqualArray(c, mx.array(np.linspace(-2.7, -0.7, 7, endpoint=False))) + + # endpoint is the fourth positional argument, before dtype, as in numpy + self.assertEqualArray( + mx.linspace(0, 10, 5, False), mx.array(np.linspace(0, 10, 5, False)) + ) + + # dtype still applies + d = mx.linspace(0, 10, 5, False, mx.int64) + self.assertEqual(d.dtype, mx.int64) + self.assertEqualArray( + d, mx.array(np.linspace(0, 10, 5, endpoint=False, dtype=int)) + ) + + # the start is kept and the stop is excluded + e = mx.linspace(3.0, 4.0, 4, endpoint=False).tolist() + self.assertEqual(e[0], 3.0) + self.assertNotIn(4.0, e) + + # decreasing ranges drop the stop value too + f = mx.linspace(10, 0, 5, endpoint=False) + self.assertEqualArray(f, mx.array(np.linspace(10, 0, 5, endpoint=False))) + + # start == stop keeps every sample at that value + g = mx.linspace(5, 5, 4, endpoint=False) + self.assertEqualArray(g, mx.array(np.linspace(5, 5, 4, endpoint=False))) + + # integer dtype truncates fractional steps, as in numpy + h = mx.linspace(0, 10, 3, endpoint=False, dtype=mx.int32) + self.assertEqualArray( + h, mx.array(np.linspace(0, 10, 3, endpoint=False, dtype=np.int32)) + ) + + # num must still be non-negative + with self.assertRaises(ValueError): + mx.linspace(0, 1, -1, endpoint=False) + def test_repeat(self): # Setup data for the tests data = mx.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]]) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index f7a2b8ab92..3da0a2950b 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3348,13 +3348,29 @@ TEST_CASE("test linspace") { auto expected = array({0.0f, 2.5f, 5.0f, 7.5f, 10.0f}, {5}); CHECK(array_equal(x, expected).item()); - x = linspace(0, 10, 5, int32); + x = linspace(0, 10, 5, true, int32); expected = array({0, 2, 5, 7, 10}, {5}); CHECK(array_equal(x, expected).item()); x = linspace(0, 1, 0); expected = array(std::initializer_list{}, {0}); CHECK(array_equal(x, expected).item()); + + x = linspace(0, 10, 5, false); + expected = array({0.0f, 2.0f, 4.0f, 6.0f, 8.0f}, {5}); + CHECK(array_equal(x, expected).item()); + + x = linspace(0, 10, 5, false, int32); + expected = array({0, 2, 4, 6, 8}, {5}); + CHECK(array_equal(x, expected).item()); + + x = linspace(1, 10, 1, false); + expected = array({1.0f}, {1}); + CHECK(array_equal(x, expected).item()); + + x = linspace(0, 1, 0, false); + expected = array(std::initializer_list{}, {0}); + CHECK(array_equal(x, expected).item()); } TEST_CASE("test quantize dequantize") { From 306bdcd18dce9b6734ea76b2b842abf3aa4af1f6 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Thu, 13 Aug 2026 22:39:17 -0700 Subject: [PATCH 07/84] Fix vmap of partition/argpartition dropping the kth argument (#4116) --- mlx/primitives.cpp | 4 +- python/tests/test_vmap.py | 79 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 3bafd19407..3c3d4fc604 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -616,7 +616,7 @@ std::pair, std::vector> ArgPartition::vmap( assert(axes.size() == 1); int axis_left = axes[0] >= 0 && axes[0] <= axis_; - return {{argpartition(inputs[0], axis_ + axis_left, stream())}, axes}; + return {{argpartition(inputs[0], kth_, axis_ + axis_left, stream())}, axes}; } std::vector ArgPartition::vjp( @@ -3421,7 +3421,7 @@ std::pair, std::vector> Partition::vmap( assert(axes.size() == 1); int axis_left = axes[0] >= 0 && axes[0] <= axis_; - return {{partition(inputs[0], axis_ + axis_left, stream())}, axes}; + return {{partition(inputs[0], kth_, axis_ + axis_left, stream())}, axes}; } bool Partition::is_equivalent(const Primitive& other) const { diff --git a/python/tests/test_vmap.py b/python/tests/test_vmap.py index 99c30a2dc2..ec97ed6fdf 100644 --- a/python/tests/test_vmap.py +++ b/python/tests/test_vmap.py @@ -252,6 +252,85 @@ def test_vmap_argreduce(self): expected = mx.array([2, 1]) self.assertTrue(mx.array_equal(out, expected)) + def _unstack(self, x, axis): + return [s.squeeze(axis) for s in mx.split(x, x.shape[axis], axis=axis)] + + def test_vmap_partition(self): + # Distinct values so each lane has a single valid kth element + a = mx.random.permutation(2 * 3 * 4).reshape(2, 3, 4).astype(mx.float32) + + for in_axis in (0, 1, 2): + slices = self._unstack(a, in_axis) + # Axis of the batched output that the inner axis maps onto + out_axes_map = [d for d in range(a.ndim) if d != in_axis] + for axis in (0, 1, -1): + oaxis = out_axes_map[axis if axis >= 0 else axis + 2] + for kth in range(slices[0].shape[axis]): + expected = mx.stack( + [mx.partition(x, kth, axis=axis) for x in slices], + axis=in_axis, + ) + pivot = mx.take(expected, mx.array([kth]), axis=oaxis) + + out = mx.vmap( + lambda x: mx.partition(x, kth, axis=axis), + in_axes=in_axis, + out_axes=in_axis, + )(a) + self.assertEqual(out.shape, expected.shape) + # partition only pins the kth element; the two sides are + # an arbitrary permutation, so compare against the sorted + # input rather than element-wise. + self.assertTrue( + mx.array_equal(mx.sort(out, axis=oaxis), mx.sort(a, axis=oaxis)) + ) + self.assertTrue( + mx.array_equal(mx.take(out, mx.array([kth]), axis=oaxis), pivot) + ) + + idx = mx.vmap( + lambda x: mx.argpartition(x, kth, axis=axis), + in_axes=in_axis, + out_axes=in_axis, + )(a) + self.assertEqual(idx.shape, expected.shape) + gathered = mx.take_along_axis(a, idx, axis=oaxis) + self.assertTrue( + mx.array_equal( + mx.sort(gathered, axis=oaxis), mx.sort(a, axis=oaxis) + ) + ) + self.assertTrue( + mx.array_equal( + mx.take(gathered, mx.array([kth]), axis=oaxis), pivot + ) + ) + + def test_vmap_topk(self): + a = mx.random.permutation(2 * 3 * 4).reshape(2, 3, 4).astype(mx.float32) + + for in_axis in (0, 1, 2): + slices = self._unstack(a, in_axis) + out_axes_map = [d for d in range(a.ndim) if d != in_axis] + for axis in (0, 1, -1): + oaxis = out_axes_map[axis if axis >= 0 else axis + 2] + for k in range(1, slices[0].shape[axis] + 1): + out = mx.vmap( + lambda x: mx.topk(x, k, axis=axis), + in_axes=in_axis, + out_axes=in_axis, + )(a) + expected = mx.stack( + [mx.topk(x, k, axis=axis) for x in slices], axis=in_axis + ) + self.assertEqual(out.shape, expected.shape) + # topk does not promise an order within the k elements + self.assertTrue( + mx.array_equal( + mx.sort(out, axis=oaxis), mx.sort(expected, axis=oaxis) + ) + ) + def test_vmap_mean(self): a = mx.arange(8).reshape(2, 4) out = mx.vmap(mx.mean)(a) From d9ad465542b4c33f46acd89c028053a10b652577 Mon Sep 17 00:00:00 2001 From: anchor Date: Fri, 14 Aug 2026 16:40:16 +0800 Subject: [PATCH 08/84] Fix nan_to_num replacing inf with 0 for float16 and bfloat16 (#4222) Co-authored-by: codeAnqiang-ma <273298913+codeAnqiang-ma@users.noreply.github.com> Co-authored-by: Cheng --- mlx/ops.cpp | 7 ++++--- python/tests/test_ops.py | 12 +++++++++++- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index ef229fe87f..7a9afcfcc7 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -16,6 +16,7 @@ #include "mlx/primitives.h" #include "mlx/transforms.h" #include "mlx/transforms_impl.h" +#include "mlx/types/limits.h" #include "mlx/utils.h" namespace mlx::core { @@ -2111,11 +2112,11 @@ array nan_to_num( auto type_to_max = [](const auto& dtype) -> float { if (dtype == float32) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else if (dtype == bfloat16) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else if (dtype == float16) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else { std::ostringstream msg; msg << "[nan_to_num] Does not yet support given type: " << dtype << "."; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index cdf0bc1b22..2b6f9324f5 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2267,7 +2267,7 @@ def test_nan_to_num(self): self.assertTrue(np.allclose(out_mx, out_np)) for t in [mx.float32, mx.float16]: - a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]) + a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]).astype(t) out_mx = mx.nan_to_num(a) out_np = np.nan_to_num(a) self.assertTrue(np.allclose(out_mx, out_np)) @@ -2277,6 +2277,16 @@ def test_nan_to_num(self): out_mx = mx.nan_to_num(a, nan=0.0, posinf=1000, neginf=-1000) self.assertTrue(np.allclose(out_mx, out_np)) + # bfloat16 has no numpy analogue; infinities should clamp to the + # dtype's largest finite value, not 0 + a = mx.array([float("inf"), 6.9, float("nan"), float("-inf")]).astype( + mx.bfloat16 + ) + out_mx = mx.nan_to_num(a) + bf_max = mx.finfo(mx.bfloat16).max + expected = mx.array([bf_max, 6.9, 0.0, -bf_max]).astype(mx.bfloat16) + self.assertTrue(mx.array_equal(out_mx, expected)) + def test_pad_reflect_symmetric(self): # mx.pad reflect/symmetric must match numpy.pad exactly. Covers # in-bounds, multi-reflect (pad larger than the axis, exercising the From bd5c3a2b170bb95340482e35b2a49fb08aea4de3 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:42:06 -0700 Subject: [PATCH 09/84] Fix einsum not broadcasting batch dimensions in batched tensordot (#4125) Co-authored-by: Cheng --- mlx/einsum.cpp | 7 ++++++- python/tests/test_einsum.py | 35 ++++++++++++++++++++++++++++++++++- 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/mlx/einsum.cpp b/mlx/einsum.cpp index b683733d0b..705a5fc2cb 100644 --- a/mlx/einsum.cpp +++ b/mlx/einsum.cpp @@ -356,7 +356,7 @@ array batch_tensordot( std::vector b_batch, std::vector b_concat, StreamOrDevice s) { - // Broadcast contracting dimensions + // Broadcast contracting and batch dimensions. { auto a_shape = a.shape(); auto b_shape = b.shape(); @@ -365,6 +365,11 @@ array batch_tensordot( a_shape[a_contract[i]] = d; b_shape[b_contract[i]] = d; } + for (int i = 0; i < a_batch.size(); ++i) { + auto d = std::max(a.shape(a_batch[i]), b.shape(b_batch[i])); + a_shape[a_batch[i]] = d; + b_shape[b_batch[i]] = d; + } a = broadcast_to(a, a_shape, s); b = broadcast_to(b, b_shape, s); } diff --git a/python/tests/test_einsum.py b/python/tests/test_einsum.py index a73ea38187..c87a6dc45f 100644 --- a/python/tests/test_einsum.py +++ b/python/tests/test_einsum.py @@ -188,7 +188,6 @@ def test_broadcasting(self): a = mx.full((5, 1), 1.0) b = mx.full((8, 2), 1.0) a_mx = mx.einsum("ab,bc->c", a, b) - return a_np = np.einsum("ab,bc->c", a, b) self.assertTrue(np.array_equal(a_mx, a_np)) @@ -358,6 +357,40 @@ def inputs_for_case(test_case): with self.assertRaises(ValueError): mx.einsum(test_case[1], *inputs) + def test_ellipses_broadcast(self): + # Size 1 batch dimensions covered by an ellipsis have to broadcast + # against the other operands, including when the smaller operand + # comes first. + shape_pairs = [ + ((1, 3, 4), (2, 4, 5)), + ((2, 3, 4), (1, 4, 5)), + ((1, 1, 3, 4), (5, 2, 4, 5)), + ((5, 1, 3, 4), (1, 2, 4, 5)), + ] + for sa, sb in shape_pairs: + a = mx.random.uniform(shape=sa) + b = mx.random.uniform(shape=sb) + mx_out = mx.einsum("...ij,...jk->...ik", a, b) + np_out = np.einsum("...ij,...jk->...ik", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + for sa, sb in [((1, 4), (5, 4)), ((5, 4), (1, 4))]: + a = mx.random.uniform(shape=sa) + b = mx.random.uniform(shape=sb) + mx_out = mx.einsum("...i,...i->...", a, b) + np_out = np.einsum("...i,...i->...", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + # Same thing with explicit labels rather than an ellipsis + a = mx.random.uniform(shape=(1, 3, 4)) + b = mx.random.uniform(shape=(2, 4, 5)) + mx_out = mx.einsum("bij,bjk->bik", a, b) + np_out = np.einsum("bij,bjk->bik", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 56e026d8a340e1b00a651d385722e11c1fbaa9f1 Mon Sep 17 00:00:00 2001 From: Rohan Gautam Date: Fri, 14 Aug 2026 03:31:17 -0700 Subject: [PATCH 10/84] Dequantize in float32 (#4241) --- mlx/backend/metal/kernels/fp_quantized.h | 7 +- mlx/backend/metal/kernels/fp_quantized_nax.h | 7 +- mlx/backend/metal/kernels/quantized.h | 73 +++++++++++--------- mlx/backend/metal/kernels/quantized_nax.h | 73 +++++++++++--------- 4 files changed, 88 insertions(+), 72 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 6e77569f56..8c963030f2 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -137,11 +137,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { template inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + const float s = float(scale); if constexpr (bits == 4) { - w_local[0] = scale * Dequantize<4, U>{}(w); - w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = scale * Dequantize<8, U>{}(w); + w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); } } diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index cf64ff7f46..57712d9bf2 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -61,11 +61,12 @@ struct Dequantize { template inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + const float s = float(scale); if constexpr (bits == 4) { - w_local[0] = scale * Dequantize<4, U>{}(w); - w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = scale * Dequantize<8, U>{}(w); + w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); } } diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6d87dc770f..f628c05612 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -489,17 +489,16 @@ inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { bits == 8, "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + const float s = float(scale); + const float b = float(bias); + if (bits == 2) { - U s[4] = { - scale, - scale / static_cast(4.0f), - scale / static_cast(16.0f), - scale / static_cast(64.0f)}; + float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { - w_local[4 * i] = s[0] * (w[i] & 0x03) + bias; - w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias; - w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias; - w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias; + w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); + w_local[4 * i + 1] = static_cast(sc[1] * (w[i] & 0x0c) + b); + w_local[4 * i + 2] = static_cast(sc[2] * (w[i] & 0x30) + b); + w_local[4 * i + 3] = static_cast(sc[3] * (w[i] & 0xc0) + b); } } @@ -508,22 +507,24 @@ inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { w_local += 8 * i; w += 3 * i; - w_local[0] = (w[0] & 0x7) * scale + bias; - w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias; - w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias; - w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias; - w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias; - w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias; - w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias; - w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias; + w_local[0] = static_cast((w[0] & 0x7) * s + b); + w_local[1] = static_cast(((w[0] & 0x38) >> 3) * s + b); + w_local[2] = + static_cast((((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * s + b); + w_local[3] = static_cast(((w[1] & 0xe) >> 1) * s + b); + w_local[4] = static_cast(((w[1] & 0x70) >> 4) * s + b); + w_local[5] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * s + b); + w_local[6] = static_cast(((w[2] & 0x1c) >> 2) * s + b); + w_local[7] = static_cast(((w[2] & 0xe0) >> 5) * s + b); } } else if (bits == 4) { - U s[2] = {scale, scale / static_cast(16.0f)}; + float sc[2] = {s, s / 16.0f}; for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias; - w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias; + w_local[2 * i] = static_cast(sc[0] * (w[i] & 0x0f) + b); + w_local[2 * i + 1] = static_cast(sc[1] * (w[i] & 0xf0) + b); } } @@ -532,14 +533,18 @@ inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { w_local += 8 * i; w += 5 * i; - w_local[0] = (w[0] & 0x1f) * scale + bias; - w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias; - w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias; - w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias; - w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias; - w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias; - w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias; - w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias; + w_local[0] = static_cast((w[0] & 0x1f) * s + b); + w_local[1] = + static_cast((((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * s + b); + w_local[2] = static_cast(((w[1] & 0x7c) >> 2) * s + b); + w_local[3] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * s + b); + w_local[4] = + static_cast((((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * s + b); + w_local[5] = static_cast(((w[3] & 0x3e) >> 1) * s + b); + w_local[6] = + static_cast((((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * s + b); + w_local[7] = static_cast(((w[4] & 0xf8) >> 3) * s + b); } } @@ -547,16 +552,18 @@ inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { for (int i = 0; i < (N / 4); i++) { w_local += 4 * i; w += 3 * i; - w_local[0] = (w[0] & 0x3f) * scale + bias; - w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias; - w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias; - w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias; + w_local[0] = static_cast((w[0] & 0x3f) * s + b); + w_local[1] = + static_cast((((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * s + b); + w_local[2] = + static_cast((((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * s + b); + w_local[3] = static_cast(((w[2] >> 2) & 0x3f) * s + b); } } else if (bits == 8) { for (int i = 0; i < N; i++) { - w_local[i] = scale * w[i] + bias; + w_local[i] = static_cast(s * w[i] + b); } } } diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 31e51a5b7e..db20c64390 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -491,17 +491,16 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { bits == 8, "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + const float s = float(scale); + const float b = float(bias); + if (bits == 2) { - U s[4] = { - scale, - scale / static_cast(4.0f), - scale / static_cast(16.0f), - scale / static_cast(64.0f)}; + float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { - w_local[4 * i] = s[0] * (w[i] & 0x03) + bias; - w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias; - w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias; - w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias; + w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); + w_local[4 * i + 1] = static_cast(sc[1] * (w[i] & 0x0c) + b); + w_local[4 * i + 2] = static_cast(sc[2] * (w[i] & 0x30) + b); + w_local[4 * i + 3] = static_cast(sc[3] * (w[i] & 0xc0) + b); } } @@ -510,22 +509,24 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 3 * i; - w_local[0] = (w[0] & 0x7) * scale + bias; - w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias; - w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias; - w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias; - w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias; - w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias; - w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias; - w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias; + w_local[0] = static_cast((w[0] & 0x7) * s + b); + w_local[1] = static_cast(((w[0] & 0x38) >> 3) * s + b); + w_local[2] = + static_cast((((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * s + b); + w_local[3] = static_cast(((w[1] & 0xe) >> 1) * s + b); + w_local[4] = static_cast(((w[1] & 0x70) >> 4) * s + b); + w_local[5] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * s + b); + w_local[6] = static_cast(((w[2] & 0x1c) >> 2) * s + b); + w_local[7] = static_cast(((w[2] & 0xe0) >> 5) * s + b); } } else if (bits == 4) { - U s[2] = {scale, scale / static_cast(16.0f)}; + float sc[2] = {s, s / 16.0f}; for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias; - w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias; + w_local[2 * i] = static_cast(sc[0] * (w[i] & 0x0f) + b); + w_local[2 * i + 1] = static_cast(sc[1] * (w[i] & 0xf0) + b); } } @@ -534,14 +535,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 5 * i; - w_local[0] = (w[0] & 0x1f) * scale + bias; - w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias; - w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias; - w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias; - w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias; - w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias; - w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias; - w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias; + w_local[0] = static_cast((w[0] & 0x1f) * s + b); + w_local[1] = + static_cast((((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * s + b); + w_local[2] = static_cast(((w[1] & 0x7c) >> 2) * s + b); + w_local[3] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * s + b); + w_local[4] = + static_cast((((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * s + b); + w_local[5] = static_cast(((w[3] & 0x3e) >> 1) * s + b); + w_local[6] = + static_cast((((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * s + b); + w_local[7] = static_cast(((w[4] & 0xf8) >> 3) * s + b); } } @@ -549,16 +554,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { for (int i = 0; i < (N / 4); i++) { w_local += 4 * i; w += 3 * i; - w_local[0] = (w[0] & 0x3f) * scale + bias; - w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias; - w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias; - w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias; + w_local[0] = static_cast((w[0] & 0x3f) * s + b); + w_local[1] = + static_cast((((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * s + b); + w_local[2] = + static_cast((((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * s + b); + w_local[3] = static_cast(((w[2] >> 2) & 0x3f) * s + b); } } else if (bits == 8) { for (int i = 0; i < N; i++) { - w_local[i] = scale * w[i] + bias; + w_local[i] = static_cast(s * w[i] + b); } } } From 6ba2d44b315207c53ab985dee36d979f97787a61 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 14 Aug 2026 03:32:18 -0700 Subject: [PATCH 11/84] chore: Reject complex in erf and erfinv (#4243) --- mlx/ops.cpp | 6 ++++++ python/tests/test_ops.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 7a9afcfcc7..f1e6b32027 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3433,6 +3433,9 @@ array sigmoid(const array& a, StreamOrDevice s /* = {} */) { } array erf(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[erf] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); return array( a.shape(), @@ -3442,6 +3445,9 @@ array erf(const array& a, StreamOrDevice s /* = {} */) { } array erfinv(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[erfinv] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); return array( a.shape(), diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 2b6f9324f5..18556e7d2c 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1161,6 +1161,13 @@ def test_erf(self): expected = np.array([math.erf(i) for i in inputs]) self.assertTrue(np.allclose(mx.erf(x), expected)) + # Complex is not supported and has to say so rather than abort + z = mx.array([1 + 2j], mx.complex64) + with self.assertRaises(ValueError): + mx.erf(z) + with self.assertRaises(ValueError): + mx.erfinv(z) + def test_erfinv(self): inputs = [-5.0, -1.0, 0.5, 0.0, 0.5, 1.0, 5.0] x = mx.array(inputs) From 3d23f7d8792a243a7f730a05300c1135aaf8298f Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 14 Aug 2026 03:35:10 -0700 Subject: [PATCH 12/84] Fix cpu compilation failure of abs with uint (#4240) Co-authored-by: Cheng --- mlx/backend/cpu/simd/base_simd.h | 10 +++++++++- python/tests/test_compile.py | 7 +++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cpu/simd/base_simd.h b/mlx/backend/cpu/simd/base_simd.h index d69e69ecf3..1ae883f7e5 100644 --- a/mlx/backend/cpu/simd/base_simd.h +++ b/mlx/backend/cpu/simd/base_simd.h @@ -84,7 +84,6 @@ Simd recip(Simd in) { DEFAULT_UNARY(operator-, std::negate{}) DEFAULT_UNARY(operator!, std::logical_not{}) -DEFAULT_UNARY(abs, std::abs) DEFAULT_UNARY(acos, std::acos) DEFAULT_UNARY(acosh, std::acosh) DEFAULT_UNARY(asin, std::asin) @@ -103,6 +102,15 @@ DEFAULT_UNARY(sqrt, std::sqrt) DEFAULT_UNARY(tan, std::tan) DEFAULT_UNARY(tanh, std::tanh) +template +Simd abs(Simd in) { + if constexpr (std::is_unsigned_v) { + return in; + } else { + return std::abs(in.value); + } +} + template Simd log1p(Simd in) { if constexpr (is_complex) { diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 7a2c6b9d0d..76e8916538 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1569,6 +1569,13 @@ def p(x): expected = w[::-1, :, ::-1, :] + 1.0 self.assertTrue(mx.array_equal(p(w[::-1, :, ::-1, :]), expected)) + def test_compile_abs_unsigned(self): + # abs has to compile for the wider unsigned types too + fun = lambda x: mx.abs(x) + 1 + for dtype in [mx.uint8, mx.uint16, mx.uint32, mx.uint64]: + x = mx.array([1, 2, 3], dtype) + self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 4947e3b615964bf52d868540b12861f685cbbad5 Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Sat, 15 Aug 2026 07:47:42 +0800 Subject: [PATCH 13/84] Fix quantize matrix multiplication floor issue (#4251) --- mlx/backend/metal/quantized.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index f659e16c93..5db3893226 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -649,7 +649,7 @@ void qvm_split_k( constexpr int bk = 32; int bn = std::min(group_size, 32) * num_simdgroups; MTL::Size group_dims = MTL::Size(bk, num_simdgroups, 1); - MTL::Size grid_dims = MTL::Size(M, N / bn, B); + MTL::Size grid_dims = MTL::Size(M, (N + bn - 1) / bn, B); auto x_shape = x.shape(); auto x_strides = x.strides(); From adf21deabd32c4fe3726ab2d373ff2165c3f27d6 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Fri, 14 Aug 2026 19:58:16 -0400 Subject: [PATCH 14/84] Only use MPI backend for world size > 1 (#4210) --- mlx/distributed/mpi/mpi.cpp | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/mlx/distributed/mpi/mpi.cpp b/mlx/distributed/mpi/mpi.cpp index 3b176e6e67..ea3960edc4 100644 --- a/mlx/distributed/mpi/mpi.cpp +++ b/mlx/distributed/mpi/mpi.cpp @@ -166,6 +166,11 @@ struct MPIWrapper { if (!is_available()) { return false; } + // MPI_Init is an error to call twice, and init() can run more than once + // when it returns without a group. + if (initialized_) { + return true; + } bool success = init(nullptr, nullptr) == MPI_SUCCESS; // Initialize custom types and ops @@ -495,6 +500,19 @@ std::shared_ptr init(bool strict /* = false */) { return nullptr; } + // Open MPI initializes a world of size 1 for a program that was not started + // with mpirun, which is not a distributed group. + int size = 1; + mpi().size(mpi().world(), &size); + if (size <= 1) { + if (strict) { + throw std::runtime_error( + "[mpi] The world has a single process. Launch with mpirun to " + "initialize the mpi backend."); + } + return nullptr; + } + return std::make_shared(mpi().world(), true); } From 140faa8ae7f1a146bbe03c6ff0991b35856af2c0 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Fri, 14 Aug 2026 20:59:33 -0700 Subject: [PATCH 15/84] chore: Reject complex in expm1, sigmoid and arctan2 (#4257) --- mlx/ops.cpp | 9 +++++++++ python/tests/test_ops.py | 10 ++++++++++ 2 files changed, 19 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index f1e6b32027..d154bd3d19 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3273,6 +3273,9 @@ array exp(const array& a, StreamOrDevice s /* = {} */) { } array expm1(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[expm1] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); auto input = astype(a, dtype, s); return array( @@ -3319,6 +3322,9 @@ array arctan(const array& a, StreamOrDevice s /* = {} */) { } array arctan2(const array& a, const array& b, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64 || b.dtype() == complex64) { + throw std::invalid_argument("[arctan2] Not supported for complex64."); + } auto dtype = at_least_float(promote_types(a.dtype(), b.dtype())); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); auto shape = inputs[0].shape(); @@ -3426,6 +3432,9 @@ array logaddexp(const array& a, const array& b, StreamOrDevice s /* = {} */) { } array sigmoid(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[sigmoid] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); auto input = astype(a, dtype, s); return array( diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 18556e7d2c..dd0a93001a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1155,6 +1155,16 @@ def test_expm1(self): np.seterr(over=errs["over"]) self.assertTrue(np.allclose(result, expected, rtol=1e-3, atol=1e-4)) + # Complex is not supported and has to say so rather than quietly + # computing on the real part + z = mx.array([1 + 2j], mx.complex64) + with self.assertRaises(ValueError): + mx.expm1(z) + with self.assertRaises(ValueError): + mx.sigmoid(z) + with self.assertRaises(ValueError): + mx.arctan2(z, z) + def test_erf(self): inputs = [-5, 0.0, 0.5, 1.0, 2.0, 10.0] x = mx.array(inputs) From 9ab977b5649154590d598ea5d545aa1b3c97f883 Mon Sep 17 00:00:00 2001 From: katlun-lgtm Date: Sat, 15 Aug 2026 06:07:04 -0400 Subject: [PATCH 16/84] Decompose small kernel-depth 3D convs into 2D convs (#3785) Co-authored-by: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com> Co-authored-by: Cheng --- mlx/backend/metal/conv.cpp | 104 +++++++++++++++++++++++++++++++++++++ python/tests/test_conv.py | 32 ++++++++++-- 2 files changed, 131 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 926f31f05a..5c32b3d963 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -5,6 +5,7 @@ #include "mlx/backend/gpu/copy.h" #include "mlx/backend/gpu/slicing.h" +#include "mlx/backend/metal/binary.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/kernels/defines.h" @@ -743,6 +744,100 @@ void pad_and_slice_conv_3D_gpu( intermediate, intermediate.strides(), {0}, intermediate.data_size()); } +void conv_2D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const std::vector& padding, + const std::vector& wt_strides, + const std::vector& wt_dilation, + const std::vector& in_dilation, + const int groups, + bool flip, + std::vector& copies); + +void small_kd_conv_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<3>& conv_params, + std::vector& copies) { + const int H = conv_params.iS[1]; + const int W = conv_params.iS[2]; + const int C = conv_params.C; + const int O = conv_params.O; + const int KD = conv_params.wS[0]; + const int KH = conv_params.wS[1]; + const int KW = conv_params.wS[2]; + const int OD = conv_params.oS[0]; + const int OH = conv_params.oS[1]; + const int OW = conv_params.oS[2]; + + array acc({OD, OH, OW, O}, out.dtype(), nullptr, {}); + for (int kd = 0; kd < KD; ++kd) { + array in_2d({OD, H, W, C}, in.dtype(), nullptr, {}); + in_2d.copy_shared_buffer( + in, + {static_cast(H) * W * C, + static_cast(W) * C, + static_cast(C), + 1}, + {true, true, false}, + static_cast(OD) * H * W * C, + static_cast(kd) * H * W * C); + + array wt_2d({O, KH, KW, C}, wt.dtype(), nullptr, {}); + wt_2d.copy_shared_buffer( + wt, + {static_cast(KD) * KH * KW * C, + static_cast(KW) * C, + static_cast(C), + 1}, + {false, false, false}, + static_cast(O - 1) * KD * KH * KW * C + + static_cast(KH) * KW * C, + static_cast(kd) * KH * KW * C); + + array conv_out({OD, OH, OW, O}, out.dtype(), nullptr, {}); + conv_2D_gpu( + s, + d, + in_2d, + wt_2d, + conv_out, + {conv_params.pad[1], conv_params.pad[2]}, + {conv_params.str[1], conv_params.str[2]}, + {conv_params.kdil[1], conv_params.kdil[2]}, + {conv_params.idil[1], conv_params.idil[2]}, + /* groups = */ 1, + conv_params.flip, + copies); + + if (kd == 0) { + acc = conv_out; + } else { + binary_op_gpu_inplace({acc, conv_out}, acc, "Add", s); + copies.push_back(conv_out); + } + } + + // Output shape is [1, OD, OH, OW, O]. + out.copy_shared_buffer( + acc, + {static_cast(OD) * OH * OW * O, + static_cast(OH) * OW * O, + static_cast(OW) * O, + static_cast(O), + 1}, + {true, true, false}, + static_cast(OD) * OH * OW * O, + 0); +} + void dispatch_conv_3D_gpu( const Stream& s, metal::Device& d, @@ -773,6 +868,15 @@ void dispatch_conv_3D_gpu( auto in = ensure_row_contiguous(in_pre, d, s); auto wt = ensure_row_contiguous(wt_pre, d, s); + // Decompose 3D conv to per-frame 2D convs + constexpr int kSmallKdLimit3D = 7; + if (is_idil_one && mod16_channels && conv_params.groups == 1 && + conv_params.N == 1 && conv_params.wS[0] <= kSmallKdLimit3D && + conv_params.str[0] == 1 && conv_params.kdil[0] == 1 && + conv_params.pad[0] == 0) { + return small_kd_conv_3D_gpu(s, d, in, wt, out, conv_params, copies); + } + // Perform the implicit gemm if (is_idil_one && mod16_channels) { return implicit_gemm_conv_3D_gpu(s, d, in, wt, out, conv_params); diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index c5f9a2c1b2..331d9fc2af 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -14,7 +14,7 @@ import torch.nn.functional as F has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -309,9 +309,11 @@ def run_conv2D( lambda x: mx.array(x).astype(mx_dtype), (in_np, wt_np) ) in_pt, wt_pt = map( - lambda x: torch.from_numpy(x.transpose(0, 3, 1, 2)) - .to("cpu") - .to(torch_dtype), + lambda x: ( + torch.from_numpy(x.transpose(0, 3, 1, 2)) + .to("cpu") + .to(torch_dtype) + ), (in_np, wt_np), ) @@ -1069,7 +1071,6 @@ def test_repeated_conv(self): @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_depthwise(self): - # fmt: off shapes = ( # N, H, W, C kH, kW, O, strides, padding, groups @@ -1221,6 +1222,27 @@ def test_conv2d_large_filter_small_channels(self): y_hat = mx.conv2d(x, w, (1, 1), (1, 1)) self.assertTrue(mx.allclose(y, y_hat, rtol=1e-3, atol=1e-3)) + def test_conv_3D_small_kd_decomposition(self): + # Exercises the small kernel-depth 3D -> KD x 2D decomposition (#3625): + # N=1, small KD, depth stride/dilation 1, no depth padding, mod16 channels. + # Validated against the CPU reference, which uses a different code path. + for T, H, W, Cin, Cout, kd, kh, kw in [ + (5, 16, 16, 32, 32, 3, 3, 3), # canonical 3x3x3 (2D hits Winograd) + (4, 12, 10, 16, 48, 3, 3, 3), # Cout != Cin + (6, 14, 14, 32, 32, 1, 3, 3), # KD = 1 + (5, 12, 12, 16, 16, 5, 1, 1), # larger KD, 1x1 spatial + (4, 10, 10, 32, 16, 2, 3, 3), # KD = 2 + ]: + x = mx.random.normal((1, T, H, W, Cin)) + w = mx.random.normal((Cout, kd, kh, kw, Cin)) + y_gpu = mx.conv_general(x, w, stride=(1, 1, 1)) + y_cpu = mx.conv_general(x, w, stride=(1, 1, 1), stream=mx.cpu) + mx.eval(y_gpu, y_cpu) + self.assertTrue( + mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4), + f"3D small-kd mismatch T{T} H{H} W{W} C{Cin}->{Cout} k{kd}{kh}{kw}", + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 052d42818ba56a18aaa914edc5bcac24c7a2b473 Mon Sep 17 00:00:00 2001 From: robertomeroni <150194833+robertomeroni@users.noreply.github.com> Date: Sun, 16 Aug 2026 06:27:33 +0200 Subject: [PATCH 17/84] Fix Metal sort of a view with a negative stride (#4252) --- mlx/backend/metal/kernels/sort.h | 10 +++++++--- python/tests/test_ops.py | 21 +++++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/kernels/sort.h b/mlx/backend/metal/kernels/sort.h index 068d43d126..ea2640bace 100644 --- a/mlx/backend/metal/kernels/sort.h +++ b/mlx/backend/metal/kernels/sort.h @@ -388,8 +388,11 @@ template < using ValT = typename sort_kernel::ValT; using IdxT = typename sort_kernel::IdxT; - auto in_block_idx = elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); - auto out_block_idx = elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); + // Signed offsets: a non-sorted axis may have a negative stride. + auto in_block_idx = + elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); + auto out_block_idx = + elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); inp += in_block_idx; out += out_block_idx; @@ -532,7 +535,8 @@ template < BLOCK_THREADS, N_PER_THREAD>; - auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); + // Signed offset: a non-sorted axis may have a negative stride. + auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); inp += block_idx; out_vals += tid.y * size_sorted_axis; out_idxs += tid.y * size_sorted_axis; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index dd0a93001a..edc5c28eb9 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2718,6 +2718,27 @@ def test_sort(self): y_np = np.sort(np.array(a), axis=-1) self.assertTrue(np.array_equal(y_np, y_mx)) + # Negative stride on an axis that is not sorted, single and multi block + np.random.seed(0) + for dtype in ("int32", "float32"): + for size in (4, 32769): + with self.subTest(dtype=dtype, size=size): + a_np = np.random.uniform(0, 100, size=(3, size)) + a_np = a_np.astype(getattr(np, dtype)) + a_mx = mx.array(a_np)[::-1, :] + a_np = a_np[::-1, :] + + b_np = np.sort(a_np, axis=-1) + self.assertTrue(np.array_equal(b_np, mx.sort(a_mx, axis=-1))) + + idx = mx.argsort(a_mx, axis=-1) + self.assertTrue( + np.array_equal(b_np, mx.take_along_axis(a_mx, idx, axis=-1)) + ) + + b_mx = mx.partition(a_mx, 1, axis=-1) + self.assertTrue(np.array_equal(b_np[:, 1], np.array(b_mx)[:, 1])) + def test_partition(self): shape = (3, 4, 5) for dtype in ("int32", "float32"): From b34545d7d96fef69fbcb0d9ad29ac2fd62fff4a3 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 16 Aug 2026 00:20:33 -0700 Subject: [PATCH 18/84] Mirror the depth axis in the decomposed 3D conv when flipped (#4277) --- mlx/backend/metal/conv.cpp | 6 +++++- python/tests/test_conv.py | 19 ++++++++++++------- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 5c32b3d963..217f70edd6 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -790,6 +790,10 @@ void small_kd_conv_3D_gpu( static_cast(OD) * H * W * C, static_cast(kd) * H * W * C); + // The 2D conv only flips the last two kernel axes, so mirror the depth + // axis here when the convolution is flipped. + const int kd_wt = conv_params.flip ? KD - 1 - kd : kd; + array wt_2d({O, KH, KW, C}, wt.dtype(), nullptr, {}); wt_2d.copy_shared_buffer( wt, @@ -800,7 +804,7 @@ void small_kd_conv_3D_gpu( {false, false, false}, static_cast(O - 1) * KD * KH * KW * C + static_cast(KH) * KW * C, - static_cast(kd) * KH * KW * C); + static_cast(kd_wt) * KH * KW * C); array conv_out({OD, OH, OW, O}, out.dtype(), nullptr, {}); conv_2D_gpu( diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index 331d9fc2af..d8bf7f8f66 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -1235,13 +1235,18 @@ def test_conv_3D_small_kd_decomposition(self): ]: x = mx.random.normal((1, T, H, W, Cin)) w = mx.random.normal((Cout, kd, kh, kw, Cin)) - y_gpu = mx.conv_general(x, w, stride=(1, 1, 1)) - y_cpu = mx.conv_general(x, w, stride=(1, 1, 1), stream=mx.cpu) - mx.eval(y_gpu, y_cpu) - self.assertTrue( - mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4), - f"3D small-kd mismatch T{T} H{H} W{W} C{Cin}->{Cout} k{kd}{kh}{kw}", - ) + # flip mirrors every kernel axis, including the decomposed depth + for flip in (False, True): + y_gpu = mx.conv_general(x, w, stride=(1, 1, 1), flip=flip) + y_cpu = mx.conv_general( + x, w, stride=(1, 1, 1), flip=flip, stream=mx.cpu + ) + mx.eval(y_gpu, y_cpu) + self.assertTrue( + mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4), + f"3D small-kd mismatch T{T} H{H} W{W} " + f"C{Cin}->{Cout} k{kd}{kh}{kw} flip={flip}", + ) if __name__ == "__main__": From a1e0e0b56904cb8eef3a4671b747833c94b0b9fb Mon Sep 17 00:00:00 2001 From: Fu Xiaonan Date: Sun, 16 Aug 2026 15:22:32 +0800 Subject: [PATCH 19/84] Fix Metal row reductions on negative-stride views (#4267) Co-authored-by: Fu Xiaonan <214359569+FU-max-boop@users.noreply.github.com> --- mlx/backend/metal/kernels/reduction/reduce_row.h | 3 ++- python/tests/test_reduce.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/kernels/reduction/reduce_row.h b/mlx/backend/metal/kernels/reduction/reduce_row.h index 936d75bb52..b55c83f315 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_row.h +++ b/mlx/backend/metal/kernels/reduction/reduce_row.h @@ -337,7 +337,8 @@ template < // lid.x * N_READS breaks the per_thread_row_reduce interface a bit. Maybe it // needs a small refactor. - in += elem_to_loc(out_idx, shape, strides, ndim) + lid.x * N_READS; + in += + elem_to_loc(out_idx, shape, strides, ndim) + IdxT(lid.x) * N_READS; LoopedElemToLoc 2)> loop(reduce_ndim); const device T* row; diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 6ac8fc1504..164e2dd803 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -1,6 +1,5 @@ # Copyright © 2023 Apple Inc. -import unittest from itertools import combinations, permutations import mlx.core as mx @@ -47,6 +46,16 @@ def test_expand_sums(self): np.allclose(z_npy, np.array(z_mlx), atol=1e-4) ) + def test_row_reduce_negative_stride(self): + x_npy = np.arange(1, 131).reshape(2, 65)[::-1] + x_mlx = mx.arange(1, 131).reshape(2, 65)[::-1] + + for op in ["sum", "max", "min", "mean", "var"]: + with self.subTest(op=op): + expected = getattr(np, op)(x_npy, axis=-1) + actual = getattr(mx, op)(x_mlx, axis=-1) + self.assertTrue(np.allclose(expected, actual)) + def test_dtypes(self): int_dtypes = [ "int8", From 9f5f7931c01bfe52866ff48b10275577048de4e7 Mon Sep 17 00:00:00 2001 From: Hao Xu Date: Sun, 16 Aug 2026 03:35:44 -0700 Subject: [PATCH 20/84] [CUDA] Fix custom kernel cache collision for same name, different source (#4273) Co-authored-by: Cheng --- mlx/backend/cuda/custom_kernel.cpp | 4 +++- python/tests/test_fast.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cuda/custom_kernel.cpp b/mlx/backend/cuda/custom_kernel.cpp index 9b5bd38b7f..c230656d80 100644 --- a/mlx/backend/cuda/custom_kernel.cpp +++ b/mlx/backend/cuda/custom_kernel.cpp @@ -310,9 +310,11 @@ void CustomKernel::eval_gpu( // Compile the custom kernel std::string kernel_name = (is_precompiled_) ? name_ : "mlx::core::cu::" + name_; + std::string module_name = + fmt::format("{}_{:x}", name_, std::hash{}(source_)); cu::JitModule& mod = cu::get_jit_module( encoder.device(), - name_, + module_name, [&]() { return std::make_tuple( is_precompiled_, source_, std::vector{kernel_name}); diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 5dacaa605c..ba5b8f3138 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1058,6 +1058,35 @@ def call_kernel(a, source): self.assertTrue(mx.array_equal(out_a, a * 2.0)) self.assertTrue(mx.array_equal(out_b, a + 100.0)) + @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") + def test_cuda_kernel_same_name_different_source(self): + # The CUDA module cache was keyed on the kernel name alone, so the + # second kernel here silently ran the first one's code. Metal had the + # same bug, fixed in #3833. + def call_kernel(a, source): + kernel = mx.fast.cuda_kernel( + name="dup_name", + input_names=["inp"], + output_names=["out"], + source=source, + ) + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.arange(32, dtype=mx.float32) + elem = "auto e = cooperative_groups::this_grid().thread_rank();" + out_a = call_kernel(a, f"{elem} out[e] = inp[e] * 2.0f;") + out_b = call_kernel(a, f"{elem} out[e] = inp[e] + 100.0f;") + mx.eval(out_a, out_b) + self.assertTrue(mx.array_equal(out_a, a * 2.0)) + self.assertTrue(mx.array_equal(out_b, a + 100.0)) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") def test_custom_metal_kernel_math_mode(self): with self.assertRaises(ValueError): From 2542e08b97203ad2ac3997af1708ea65ebc4a3ef Mon Sep 17 00:00:00 2001 From: Feli <89400571+FeliGame@users.noreply.github.com> Date: Sun, 16 Aug 2026 18:36:15 +0800 Subject: [PATCH 21/84] Fix ops rejecting integers larger than INT32_MAX (#4255) Co-authored-by: Feli Co-authored-by: Cheng --- python/src/ops.cpp | 6 +++--- python/tests/test_ops.py | 8 +++++++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 4dc5114bf7..611e281dbf 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -24,10 +24,10 @@ namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; -using Scalar = std::variant; +using Scalar = std::variant; mx::Dtype scalar_to_dtype(Scalar s) { - if (std::holds_alternative(s)) { + if (std::holds_alternative(s)) { return mx::int32; } else if (std::holds_alternative(s)) { return mx::float32; @@ -37,7 +37,7 @@ mx::Dtype scalar_to_dtype(Scalar s) { } double scalar_to_double(Scalar s) { - if (auto pv = std::get_if(&s); pv) { + if (auto pv = std::get_if(&s); pv) { return static_cast(*pv); } else if (auto pv = std::get_if(&s); pv) { return *pv; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index edc5c28eb9..1237c80e04 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1623,7 +1623,7 @@ def test_arange_overload_dispatch(self): a = mx.arange(float("inf"), 1, float("inf")) with self.assertRaises(ValueError): a = mx.arange(float("inf"), 1, 5) - with self.assertRaises(TypeError): + with self.assertRaises(ValueError): INT_MAX = 2147483647 a = mx.arange(0, INT_MAX + 1, 1) @@ -1743,6 +1743,12 @@ def test_arange_corner_cases_cast(self): expected = [0] self.assertListEqual(a.tolist(), expected) + n = mx.iinfo(mx.int32).max + result = mx.arange(n - 1, n + 3) + self.assertEqual(result.shape, (4,)) + self.assertEqual(result.dtype, mx.int32) + self.assertEqual(result.tolist(), [n - 1, n, -2147483648, -2147483647]) + def test_hanning_general(self): a = mx.hanning(10) expected = np.hanning(10) From 4b342c00be759e2e33f6b36bf4d18abb14e1f975 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 16 Aug 2026 03:37:23 -0700 Subject: [PATCH 22/84] Fix var/std for complex numbers (#4260) --- mlx/ops.cpp | 12 +++++++++++- python/tests/test_ops.py | 10 ++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index d154bd3d19..9378040292 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2455,7 +2455,17 @@ array var( StreamOrDevice s /* = {}*/) { auto dtype = at_least_float(a.dtype()); auto mu = mean(a, axes, /* keepdims= */ true, s); - auto v = sum(square(subtract(a, mu, s), s), axes, keepdims, s); + auto d = subtract(a, mu, s); + // The variance of complex values is the mean squared magnitude. Squaring the + // deviations directly gives a complex result which can even be negative, so + // multiply by the conjugate instead. + auto sq = issubdtype(dtype, complexfloating) + ? real(multiply(d, conjugate(d, s), s), s) + : square(d, s); + if (issubdtype(dtype, complexfloating)) { + dtype = float32; + } + auto v = sum(sq, axes, keepdims, s); if (ddof != 0) { auto normalizer = maximum( diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 1237c80e04..8ee3b45ab4 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -985,11 +985,21 @@ def test_var(self): out = mx.var(x, ddof=3) self.assertEqual(out.item(), float("inf")) + x = mx.array([1 + 2j, -3 - 4j, 0.5 - 0.25j]) + x_np = np.array(x) + self.assertEqual(mx.var(x).dtype, mx.float32) + self.assertAlmostEqual(mx.var(x).item(), x_np.var().item(), places=5) + def test_std(self): x = mx.random.uniform(shape=(5, 5)) x_np = np.array(x) self.assertAlmostEqual(mx.std(x).item(), x_np.std().item(), places=6) + x = mx.array([1 + 2j, -3 - 4j, 0.5 - 0.25j]) + x_np = np.array(x) + self.assertEqual(mx.std(x).dtype, mx.float32) + self.assertAlmostEqual(mx.std(x).item(), x_np.std().item(), places=5) + def test_abs(self): a = mx.array([-1.0, 1.0, -2.0, 3.0]) result = mx.abs(a) From 3973edd3ff89e36d8541c46766280e6412b4db75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Sun, 16 Aug 2026 13:37:40 +0300 Subject: [PATCH 23/84] Fix int32 overflow in conv padded input and pad shapes (#4258) Co-authored-by: Cheng --- mlx/backend/cpu/conv.cpp | 9 +++++++-- mlx/backend/metal/conv.cpp | 15 ++++++++++----- mlx/ops.cpp | 5 ++++- tests/ops_tests.cpp | 8 ++++++++ 4 files changed, 29 insertions(+), 8 deletions(-) diff --git a/mlx/backend/cpu/conv.cpp b/mlx/backend/cpu/conv.cpp index 70b5f270f0..17bc6cb9ff 100644 --- a/mlx/backend/cpu/conv.cpp +++ b/mlx/backend/cpu/conv.cpp @@ -814,7 +814,11 @@ void explicit_gemm_conv_1D_cpu( auto& encoder = cpu::get_command_encoder(stream); // Pad input - Shape padded_shape = {N, iH + padding_lo[0] + padding_hi[0], C}; + Shape padded_shape = { + N, + safe_cast( + static_cast(iH) + padding_lo[0] + padding_hi[0], "conv"), + C}; array in_padded(padded_shape, conv_dtype, nullptr, {}); // Fill with zeros @@ -961,7 +965,8 @@ void explicit_gemm_conv_ND_cpu( Shape padded_shape(in.shape().size()); padded_shape.front() = N; for (size_t i = 0; i < iDim.size(); i++) { - padded_shape[i + 1] = iDim[i] + padding_lo[i] + padding_hi[i]; + padded_shape[i + 1] = safe_cast( + static_cast(iDim[i]) + padding_lo[i] + padding_hi[i], "conv"); } padded_shape.back() = C; array in_padded(padded_shape, conv_dtype, nullptr, {}); diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 217f70edd6..85e137bf4f 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -902,15 +902,20 @@ void winograd_conv_2D_gpu( array& out, const MLXConvParams<2>& conv_params, std::vector& copies_w) { + // Round the padded spatial dims up to the Winograd tile in int64 so the + // rounding cannot overflow int32 just below the limit. + int64_t pad_h = static_cast(conv_params.iS[0]) + + 2 * static_cast(conv_params.pad[0]); + int64_t pad_w = static_cast(conv_params.iS[1]) + + 2 * static_cast(conv_params.pad[1]); + pad_h = 6 * ((pad_h - 2 + 5) / 6) + 2; + pad_w = 6 * ((pad_w - 2 + 5) / 6) + 2; Shape padded_shape = { conv_params.N, - conv_params.iS[0] + 2 * conv_params.pad[0], - conv_params.iS[1] + 2 * conv_params.pad[1], + safe_cast(pad_h, "conv"), + safe_cast(pad_w, "conv"), conv_params.C}; - padded_shape[1] = 6 * ((padded_shape[1] - 2 + 5) / 6) + 2; - padded_shape[2] = 6 * ((padded_shape[2] - 2 + 5) / 6) + 2; - array in_padded(std::move(padded_shape), in.dtype(), nullptr, {}); // Fill with zeros diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 9378040292..e447daa513 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1633,7 +1633,10 @@ array pad( } auto ax = axes[i] < 0 ? a.ndim() + axes[i] : axes[i]; - out_shape[ax] += low_pad_size[i] + high_pad_size[i]; + out_shape[ax] = safe_cast( + static_cast(out_shape[ax]) + low_pad_size[i] + + high_pad_size[i], + "pad"); } if (mode == "constant") { diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 3da0a2950b..09236f1da1 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -4491,6 +4491,14 @@ TEST_CASE("test conv shape overflow") { Shape{1, 8, 8, 1}); } +TEST_CASE("test pad shape overflow") { + // A padding sum that overflows int32 is rejected, not wrapped. + // https://github.com/ml-explore/mlx/issues/3611 + const int imax = 2147483647; + CHECK_THROWS_AS( + pad(zeros({8}), {0}, Shape{imax}, Shape{imax}), std::overflow_error); +} + TEST_CASE("test fp8 conversion") { for (auto t : {float32, float16, bfloat16}) { array in({-1.125, -1.0, 0.0, 1.0, 1.125, 4.5, 448.0}, t); From bac07b3249a74ab443abee89ad59a6e55a024e23 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 16 Aug 2026 03:39:53 -0700 Subject: [PATCH 24/84] chore: Reject complex in remainder (#4270) --- mlx/ops.cpp | 3 +++ python/tests/test_ops.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index e447daa513..edcf6670c2 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -3196,6 +3196,9 @@ array floor_divide( array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); + if (issubdtype(dtype, complexfloating)) { + throw std::invalid_argument("[remainder] Complex type not supported."); + } auto inputs = broadcast_arrays( {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s); auto shape = inputs[0].shape(); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 8ee3b45ab4..a09b96669e 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -344,6 +344,14 @@ def test_divide(self): self.assertEqual(z.item(), 2) def test_remainder(self): + # Complex is not supported and has to say so rather than quietly + # computing a componentwise remainder, which no other library defines + z = mx.array([7 + 3j], mx.complex64) + with self.assertRaises(ValueError): + mx.remainder(z, z) + with self.assertRaises(ValueError): + z % z + for dt in [mx.int32, mx.float32, mx.float16, mx.bfloat16]: x = mx.array(2, dtype=dt) y = mx.array(4, dtype=dt) From c2bcf47ee312a26cdb6d1dee76aec43e03121303 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Sun, 16 Aug 2026 06:52:06 -0400 Subject: [PATCH 25/84] chore: Compare the macOS SDK version as a version when gating JACCL (#4286) --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9d7b5baa30..3c8057c383 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -363,7 +363,7 @@ target_include_directories( if(MLX_BUILD_CPU AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin" AND DEFINED MACOS_SDK_VERSION - AND MACOS_SDK_VERSION GREATER_EQUAL 26.2) + AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2) add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib ${CMAKE_BINARY_DIR}/jaccl) endif() From a44fc8c06671107e93eff21c7b3159c505db4cca Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Sun, 16 Aug 2026 19:44:45 -0400 Subject: [PATCH 26/84] Clamp ring socket transfers so a payload of 2 GiB or more can be sent (#4281) Co-authored-by: Cheng --- mlx/distributed/ring/ring.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index 3e0c2a3221..9a81010e34 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -36,6 +37,10 @@ constexpr const size_t ALL_SUM_BUFFERS = 2; constexpr const int CONN_ATTEMPTS = 5; constexpr const int CONN_WAIT = 1000; constexpr const char* RING_TAG = "[ring]"; +// send(2) and recv(2) reject a length above INT_MAX with EINVAL, so a single +// transfer of 2 GiB or more fails outright rather than being carried in +// pieces. +constexpr const size_t MAX_IO_BYTES = 1024 * 1024 * 1024; using GroupImpl = mlx::core::distributed::detail::GroupImpl; using json = nlohmann::json; @@ -174,7 +179,8 @@ class SocketThread { if (!recvs_.empty()) { auto& task = recvs_.front(); - ssize_t r = ::recv(fd_, task.buffer, task.size, 0); + ssize_t r = + ::recv(fd_, task.buffer, std::min(task.size, MAX_IO_BYTES), 0); if (r > 0) { task.buffer = static_cast(task.buffer) + r; task.size -= r; @@ -191,7 +197,8 @@ class SocketThread { } if (!sends_.empty()) { auto& task = sends_.front(); - ssize_t r = ::send(fd_, task.buffer, task.size, 0); + ssize_t r = + ::send(fd_, task.buffer, std::min(task.size, MAX_IO_BYTES), 0); if (r > 0) { task.buffer = static_cast(task.buffer) + r; task.size -= r; From bbebc8f293c0e43f8e2c15148387fde89f52f900 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Sun, 16 Aug 2026 16:45:27 -0700 Subject: [PATCH 27/84] chore: Use normalize_axis_index in split/unstack/partition/topk (#4288) --- mlx/ops.cpp | 54 ++++++++--------------------------------------------- 1 file changed, 8 insertions(+), 46 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index edcf6670c2..1c09e728d0 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1143,13 +1143,7 @@ std::vector split( const Shape& indices, int axis, StreamOrDevice s /* = {} */) { - auto ax = axis < 0 ? axis + a.ndim() : axis; - if (ax < 0 || ax >= a.ndim()) { - std::ostringstream msg; - msg << "Invalid axis (" << axis << ") passed to split" - << " for array with shape " << a.shape() << "."; - throw std::invalid_argument(msg.str()); - } + auto ax = normalize_axis_index(axis, a.ndim(), "[split] "); if (indices.empty()) { return {a}; @@ -1191,20 +1185,14 @@ split(const array& a, const Shape& indices, StreamOrDevice s /* = {} */) { std::vector split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) { - auto ax = axis < 0 ? axis + a.ndim() : axis; - if (ax < 0 || ax >= a.ndim()) { - std::ostringstream msg; - msg << "Invalid axis " << axis << " passed to split" - << " for array with shape " << a.shape() << "."; - throw std::invalid_argument(msg.str()); - } + auto ax = normalize_axis_index(axis, a.ndim(), "[split] "); if (num_splits <= 0) { std::ostringstream msg; msg << "[split] num_splits must be positive and non-zero but got " << num_splits << "."; throw std::invalid_argument(msg.str()); } - auto q_and_r = std::ldiv(a.shape(axis), num_splits); + auto q_and_r = std::ldiv(a.shape(ax), num_splits); if (q_and_r.rem) { std::ostringstream msg; msg << "Array split does not result in sub arrays with equal size:" @@ -1217,7 +1205,7 @@ split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) { for (int i = 0; i < indices.size(); ++i) { indices[i] = (i + 1) * split_size; } - return split(a, indices, axis, s); + return split(a, indices, ax, s); } std::vector @@ -1228,13 +1216,7 @@ split(const array& a, int num_splits, StreamOrDevice s /* = {} */) { std::vector unstack(const array& a, int axis, StreamOrDevice s /* = {} */) { auto ndim = static_cast(a.ndim()); - auto ax = axis < 0 ? axis + ndim : axis; - if (ax < 0 || ax >= ndim) { - std::ostringstream msg; - msg << "[unstack] Invalid axis " << axis << " for array with " << ndim - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + auto ax = normalize_axis_index(axis, ndim, "[unstack] "); auto n = a.shape(ax); std::vector res; res.reserve(n); @@ -2851,14 +2833,7 @@ array partition( int axis, StreamOrDevice s /* = {} */) { // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[partition] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - int axis_ = axis < 0 ? axis + a.ndim() : axis; + int axis_ = normalize_axis_index(axis, a.ndim(), "[partition] "); int kth_ = kth < 0 ? kth + a.shape(axis) : kth; if (kth_ < 0 || kth_ >= a.shape(axis_)) { std::ostringstream msg; @@ -2892,14 +2867,7 @@ array argpartition( int axis, StreamOrDevice s /* = {} */) { // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[argpartition] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - int axis_ = axis < 0 ? axis + a.ndim() : axis; + int axis_ = normalize_axis_index(axis, a.ndim(), "[argpartition] "); int kth_ = kth < 0 ? kth + a.shape(axis) : kth; if (kth_ < 0 || kth_ >= a.shape(axis_)) { std::ostringstream msg; @@ -2954,13 +2922,7 @@ array topk(const array& a, int k, StreamOrDevice s /* = {}*/) { /** Returns topk elements of the array along a given axis. */ array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) { // Check for valid axis - int axis_ = axis < 0 ? axis + a.ndim() : axis; - if (axis_ < 0 || axis_ >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[topk] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } + int axis_ = normalize_axis_index(axis, a.ndim(), "[topk] "); if (k < 0 || k > a.shape(axis_)) { std::ostringstream msg; msg << "[topk] Received invalid k=" << k << " along axis " << axis From 61a4867be62594aa2a243910b401ac0bf3e02525 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 17 Aug 2026 09:08:28 +0900 Subject: [PATCH 28/84] Remove grouped output in CI (#4195) --- .github/actions/build-macos/action.yml | 12 ++-------- .github/actions/build-wheel/action.yml | 10 -------- .github/actions/build/action.yml | 4 ---- .github/actions/setup/action.yml | 18 -------------- .github/actions/test-linux/action.yml | 31 ++++--------------------- .github/actions/test-macos/action.yml | 7 ++---- .github/actions/test-wheel/action.yml | 2 -- .github/actions/test-windows/action.yml | 10 ++------ .github/workflows/build_and_test.yml | 1 + 9 files changed, 12 insertions(+), 83 deletions(-) diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml index 84055009e9..b7a69aaca0 100644 --- a/.github/actions/build-macos/action.yml +++ b/.github/actions/build-macos/action.yml @@ -15,10 +15,7 @@ runs: steps: - name: Install dependencies shell: bash - run: | - echo "::group::Install dependencies" - uv pip install 'build<=1.4.2' setuptools - echo "::endgroup::" + run: uv pip install 'build<=1.4.2' setuptools - name: Build wheel shell: bash @@ -26,17 +23,13 @@ runs: DEBUG: 1 CMAKE_ARGS: ${{ inputs.cmake-args }} MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - echo "::group::Build wheel" - python -m build -w - echo "::endgroup::" + run: python -m build -w - name: Build CPP only shell: bash env: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | - echo "::group::Build CPP only" if ${{ contains(inputs.cmake-args, 'CMAKE_BUILD_TYPE') }} ; then cmake . -B build ${{ inputs.cmake-args }} else @@ -44,4 +37,3 @@ runs: -DCMAKE_BUILD_TYPE=Debug fi cmake --build build -j $(sysctl -n hw.physicalcpu) - echo "::endgroup::" diff --git a/.github/actions/build-wheel/action.yml b/.github/actions/build-wheel/action.yml index 96cca9eadc..6177bf1849 100644 --- a/.github/actions/build-wheel/action.yml +++ b/.github/actions/build-wheel/action.yml @@ -34,13 +34,11 @@ runs: - name: Install dependencies shell: bash run: | - echo "::group::Install dependencies" uv pip install 'build<=1.4.2' setuptools if ${{ runner.os == 'Linux' }} ; then uv pip install auditwheel patchelf fi mkdir -p wheelhouse - echo "::endgroup::" - name: Build frontend package if: inputs.build-frontend == 'true' @@ -49,16 +47,13 @@ runs: CMAKE_ARGS: ${{ inputs.cmake-args }} MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | - echo "::group::Build frontend package" python setup.py clean --all MLX_BUILD_STAGE=1 python -m build -w - echo "::endgroup::" - name: Post-process frontend package if: inputs.build-frontend == 'true' shell: bash run: | - echo "::group::Post-process frontend package" if ${{ runner.os == 'Linux' }} ; then auditwheel repair dist/mlx-*.whl \ --plat manylinux_2_35_${{ inputs.arch-tag }} \ @@ -67,7 +62,6 @@ runs: else mv dist/mlx-*.whl wheelhouse/ fi - echo "::endgroup::" - name: Build backend package if: inputs.build-backend == 'true' @@ -76,16 +70,13 @@ runs: CMAKE_ARGS: ${{ inputs.cmake-args }} MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | - echo "::group::Build backend package" python setup.py clean --all MLX_BUILD_STAGE=2 python -m build -w - echo "::endgroup::" - name: Post-process backend package if: inputs.build-backend == 'true' shell: bash run: | - echo "::group::Post-process backend package" if ${{ runner.os == 'Linux' }} ; then if [ -f dist/mlx_cpu*.whl ]; then auditwheel repair dist/mlx_cpu*.whl \ @@ -112,4 +103,3 @@ runs: mv dist/mlx_metal*.whl wheelhouse/ fi fi - echo "::endgroup::" diff --git a/.github/actions/build/action.yml b/.github/actions/build/action.yml index 5756f6e146..037e0ea6bc 100644 --- a/.github/actions/build/action.yml +++ b/.github/actions/build/action.yml @@ -18,11 +18,9 @@ runs: DEBUG: ${{ inputs.debug == 'true' && 1 || 0 }} CMAKE_ARGS: ${{ inputs.cmake-args }} run: | - echo "::group::Install Python package" # Install cpu-only torch to save space. uv pip install torch --torch-backend=cpu uv pip install --no-build-isolation -e ".[dev]" -v - echo "::endgroup::" - name: Build CPP only shell: bash @@ -35,10 +33,8 @@ runs: CCACHE_BASEDIR: ${{ github.workspace }}/build/cpp/mlx CCACHE_NOHASHDIR: true run: | - echo "::group::Build CPP only" cmake . -B build/cpp/mlx ${{ inputs.cmake-args }} \ -DBUILD_SHARED_LIBS=ON \ -DCMAKE_BUILD_TYPE=${{ inputs.debug == 'true' && 'Debug' || 'Release' }} cmake --build build/cpp/mlx \ -j ${{ runner.os == 'Windows' && '$NUMBER_OF_PROCESSORS' || '$(nproc)' }} - echo "::endgroup::" diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 4a1a27cf9c..3c71b13d72 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -39,30 +39,25 @@ runs: if: runner.os == 'Linux' shell: bash run: | - echo "::group::Install common dependencies" sudo apt-get update sudo apt-get install -y --no-install-recommends \ gdb g++ ninja-build zip \ libblas-dev liblapack-dev liblapacke-dev \ openmpi-bin openmpi-common libopenmpi-dev - echo "::endgroup::" - name: Install macOS dependencies if: runner.os == 'macOS' shell: bash run: | - echo "::group::Install macOS dependencies" brew update brew install openmpi xcodebuild -showComponent MetalToolchain sysctl -a | grep machdep.cpu - echo "::endgroup::" - name: Setup Windows environment if: runner.os == 'Windows' shell: cmd run: | - echo "::group::Setup environment" :: Find out path to Visual Studio. pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\" for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x @@ -78,7 +73,6 @@ runs: set CCACHE_SLOPPINESS=include_file_ctime,include_file_mtime :: Export to all steps. >>%GITHUB_ENV% set - echo "::endgroup::" - uses: astral-sh/setup-uv@v8.2.0 with: @@ -120,7 +114,6 @@ runs: if: runner.os != 'Windows' shell: bash run: | - echo "::group::Setup Python venv" uv venv --python ${{ inputs.python-version }} --managed-python # Make sure all builds use the same cmake binary. uv pip install cmake @@ -132,18 +125,15 @@ runs: if ${{ startsWith(inputs.toolkit, 'cuda') }} ; then echo MLX_PTX_CACHE_DIR=/tmp/mlx-ptx-cache >> $GITHUB_ENV fi - echo "::endgroup::" - name: Setup Python venv (Windows) if: runner.os == 'Windows' shell: cmd run: | - echo "::group::Setup Python venv" uv venv --python ${{ inputs.python-version }}${{ runner.arch == 'arm64' && '-arm64' || ''}} || exit /b uv pip install cmake call ".venv/Scripts/activate.bat" >>%GITHUB_ENV% set - echo "::endgroup::" - name: Install CUDA toolkit (Linux) if: runner.os == 'Linux' && startsWith(inputs.toolkit, 'cuda') @@ -156,7 +146,6 @@ runs: "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0" } run: | - echo "::group::Install CUDA toolkit" # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is # Jetson specific. SBSA means Arm Server Base System Architecture. ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }} @@ -167,7 +156,6 @@ runs: libnccl2 libnccl-dev \ ${{ fromJson(env.PACKAGES)[inputs.toolkit] }} echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH - echo "::endgroup::" - name: Install CUDA Toolkit (Windows) if: runner.os == 'Windows' && startsWith(inputs.toolkit, 'cuda') @@ -186,7 +174,6 @@ runs: "cuda-13.0": ["cudart_13.0", "nvcc_13.0", "cublas_13.0", "cublas_dev_13.0", "cufft_13.0", "cufft_dev_13.0", "nvrtc_13.0", "nvrtc_dev_13.0", "crt_13.0", "nvvm_13.0", "nvptxcompiler_13.0"], } run: | - echo "::group::Install CUDA toolkit" $ErrorActionPreference = "Stop" $cudaUrl = "${{ fromJson(env.INSTALLERS)[inputs.toolkit] }}" $cudaInstaller = "./install.exe" @@ -201,7 +188,6 @@ runs: Start-Process -FilePath $cudaInstaller -ArgumentList "$args" -NoNewWindow -Wait $cudaPath = (Resolve-Path "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\*").path echo "$cudaPath\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append - echo "::endgroup::" - name: Install cuDNN (Windows) if: runner.os == 'Windows' && startsWith(inputs.toolkit, 'cuda') @@ -215,7 +201,6 @@ runs: "cuda-13.0": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda13-archive.zip" } run: | - echo "::group::Install cuDNN" $ErrorActionPreference = "Stop" $cudnnUrl = "${{ fromJson(env.ARCHIVES)[inputs.toolkit] }}" $cudnnZip = "cudnn.zip" @@ -229,13 +214,11 @@ runs: Expand-Archive -Path $cudnnZip -DestinationPath cudnn-extracted $cudnnDir = (Get-ChildItem -Path cudnn-extracted -Directory)[0].FullName echo "cudnnDir=$($cudnnDir -replace '\\', '/')" | Out-File -FilePath $env:GITHUB_OUTPUT - echo "::endgroup::" - name: Generate CMake args id: cmake-args shell: bash run: | - echo "::group::Generate CMake args" cmakeArgs=( "-G Ninja" ) @@ -292,4 +275,3 @@ runs: IFS=" " echo ${cmakeArgs[*]} echo "cmakeArgs=${cmakeArgs[*]}" >> $GITHUB_OUTPUT - echo "::endgroup::" diff --git a/.github/actions/test-linux/action.yml b/.github/actions/test-linux/action.yml index 9d64e416de..ea6ac2af88 100644 --- a/.github/actions/test-linux/action.yml +++ b/.github/actions/test-linux/action.yml @@ -8,83 +8,62 @@ runs: id: gpu-check shell: bash run: | - echo "::group::Check GPU support" if __nvcc_device_query ; then echo "good=true" >> $GITHUB_OUTPUT else echo "good=false" >> $GITHUB_OUTPUT fi echo - echo "::endgroup::" - name: Run MPI tests if: steps.gpu-check.outputs.good == 'false' shell: bash - run: | - echo "::group::MPI tests" - mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py - echo "::endgroup::" + run: mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py - name: Run distributed tests if: steps.gpu-check.outputs.good == 'false' shell: bash run: | - echo "::group::Distributed tests" mlx.launch --verbose -n 8 python python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) if grep -Fq '[WARN]' stderr.log ; then grep -F '[WARN]' stderr.log echo "Distributed ring test failed"; exit 1; fi - echo "::endgroup::" - name: Run Python tests - CPU if: steps.gpu-check.outputs.good == 'false' shell: bash env: DEVICE: cpu - run: | - echo "::group::Python tests - CPU" - python -m unittest discover python/tests -v - echo "::endgroup::" + run: python -m unittest discover python/tests -v - name: Run Python tests - GPU if: steps.gpu-check.outputs.good == 'true' shell: bash env: DEVICE: gpu - run: | - echo "::group::Python tests - GPU" - python -m tests discover python/tests -v - echo "::endgroup::" + run: python -m tests discover python/tests -v - name: Run CPP tests - CPU shell: bash env: DEVICE: cpu - run: | - echo "::group::CPP tests - CPU" - ./build/cpp/mlx/tests/tests - echo "::endgroup::" + run: ./build/cpp/mlx/tests/tests - name: Run CPP tests - GPU if: steps.gpu-check.outputs.good == 'true' shell: bash env: DEVICE: gpu - run: | - echo "::group::CPP tests - GPU" - ./build/cpp/mlx/tests/tests -sfe="*linalg_tests.cpp" - echo "::endgroup::" + run: ./build/cpp/mlx/tests/tests -sfe="*linalg_tests.cpp" - name: Show stack trace on crash if: failure() shell: bash run: | - echo "::group::Show stack trace on crash" set +e sleep 10 if coredumpctl list; then coredumpctl debug --debugger-arguments="-batch -ex 'thread apply all bt'" fi - echo "::endgroup::" diff --git a/.github/actions/test-macos/action.yml b/.github/actions/test-macos/action.yml index f03e77c8df..aa043fb66c 100644 --- a/.github/actions/test-macos/action.yml +++ b/.github/actions/test-macos/action.yml @@ -12,12 +12,9 @@ runs: steps: - name: Install tests dependencies shell: bash - run: | - echo "::group::Install tests dependencies" - uv pip install tensorflow - echo "::endgroup::" + run: uv pip install tensorflow - - name: Run Python tests + - name: Run tests shell: bash env: METAL_DEBUG_ERROR_MODE: 0 diff --git a/.github/actions/test-wheel/action.yml b/.github/actions/test-wheel/action.yml index db68579bbb..b745a4d2f9 100644 --- a/.github/actions/test-wheel/action.yml +++ b/.github/actions/test-wheel/action.yml @@ -30,7 +30,6 @@ runs: - name: Test local packages shell: bash run: | - echo "::group::Test local packages" # Fall back to PyPI if PyTorch index has issues. uv pip install torch --torch-backend=cpu uv pip install numpy @@ -48,4 +47,3 @@ runs: exit 1 fi python -m unittest discover -v python/tests - echo "::endgroup::" diff --git a/.github/actions/test-windows/action.yml b/.github/actions/test-windows/action.yml index ac812d8af9..a7510f57e2 100644 --- a/.github/actions/test-windows/action.yml +++ b/.github/actions/test-windows/action.yml @@ -8,16 +8,10 @@ runs: shell: bash env: DEVICE: cpu - run: | - echo "::group::Python tests - CPU" - python -m unittest discover python/tests -v - echo "::endgroup::" + run: python -m unittest discover python/tests -v - name: Run CPP tests - CPU shell: bash env: DEVICE: cpu - run: | - echo "::group::CPP tests - CPU" - ./build/cpp/mlx/tests.exe -tce="*gguf*" - echo "::endgroup::" + run: ./build/cpp/mlx/tests.exe -tce="*gguf*" diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 6de4218b2c..d5dbc329f0 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -77,6 +77,7 @@ jobs: name: macOS (${{ matrix.macos-target }}, ${{ matrix.toolkit }}) if: github.repository == 'ml-explore/mlx' strategy: + fail-fast: false matrix: macos-target: ['14.0', '15.0', '26.2'] toolkit: ['cpu', 'metal', 'jit'] From 0bb81212cad4e1bb9fb31018cc84b7d8dcabdbe5 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 17 Aug 2026 09:16:19 +0900 Subject: [PATCH 29/84] [CUDA] Fix finding cuda 13 headers in JIT compilation (#3995) --- mlx/backend/cuda/jit_module.cpp | 16 +++++++++++++++- setup.py | 7 ++++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp index 3de1ddb018..0d107a8e18 100644 --- a/mlx/backend/cuda/jit_module.cpp +++ b/mlx/backend/cuda/jit_module.cpp @@ -49,6 +49,20 @@ const std::filesystem::path& default_cuda_toolkit_path() { return cached_path; } +// Get the dirname of nvidia python package that contains CUDA headers. +inline const char* cudart_dirname() { +#if CUDART_VERSION < 13000 + return "cuda_runtime"; +#elif CUDART_VERSION < 14000 + return "cu13"; +#else + static_assert( + false, + "Please find out the newest dirname under site-packages/nvidia " + "and add it in this function."); +#endif +} + // Return the --include-path args used for invoking NVRTC. const std::vector& include_path_args() { static std::vector cached_args = []() { @@ -72,7 +86,7 @@ const std::vector& include_path_args() { } // Add path to CUDA runtime headers, try local-installed python package // first and then system-installed headers. - path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include"; + path = root_dir.parent_path() / "nvidia" / cudart_dirname() / "include"; if (!std::filesystem::exists(path)) { const char* home = std::getenv("CUDA_HOME"); if (!home) { diff --git a/setup.py b/setup.py index 3c2f138048..835375599e 100644 --- a/setup.py +++ b/setup.py @@ -320,9 +320,10 @@ def get_tag(self) -> tuple[str, str, str]: ] elif toolkit == 13: install_requires += [ - "nvidia-cublas", - "nvidia-cufft", - "nvidia-cuda-nvrtc", + "nvidia-cublas==13.*", + "nvidia-cufft==12.*", + "nvidia-cuda-nvrtc==13.*", + "nvidia-cuda-runtime==13.*", ] else: raise ValueError(f"Unknown toolkit {toolkit}") From d9e2b0d40c9595e6ebc02b09699fd72294b78641 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 17 Aug 2026 11:17:12 +0900 Subject: [PATCH 30/84] Refactor wheel building script (#3818) --- .github/actions/build-macos/action.yml | 2 +- .github/actions/build-wheel/action.yml | 6 +- .github/actions/setup/action.yml | 3 +- .github/workflows/release.yml | 4 ++ setup.py | 92 ++++++++++++++++---------- 5 files changed, 67 insertions(+), 40 deletions(-) diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml index b7a69aaca0..7bcd4e0d09 100644 --- a/.github/actions/build-macos/action.yml +++ b/.github/actions/build-macos/action.yml @@ -15,7 +15,7 @@ runs: steps: - name: Install dependencies shell: bash - run: uv pip install 'build<=1.4.2' setuptools + run: uv pip install build setuptools - name: Build wheel shell: bash diff --git a/.github/actions/build-wheel/action.yml b/.github/actions/build-wheel/action.yml index 6177bf1849..2e15efd93f 100644 --- a/.github/actions/build-wheel/action.yml +++ b/.github/actions/build-wheel/action.yml @@ -34,7 +34,7 @@ runs: - name: Install dependencies shell: bash run: | - uv pip install 'build<=1.4.2' setuptools + uv pip install build setuptools if ${{ runner.os == 'Linux' }} ; then uv pip install auditwheel patchelf fi @@ -48,7 +48,7 @@ runs: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w + MLX_BUILD_FRONTEND_PACKAGE=1 python -m build -w - name: Post-process frontend package if: inputs.build-frontend == 'true' @@ -71,7 +71,7 @@ runs: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w + MLX_BUILD_BACKEND_PACKAGE=1 python -m build -w - name: Post-process backend package if: inputs.build-backend == 'true' diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 3c71b13d72..947c0ff1f0 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -41,7 +41,7 @@ runs: run: | sudo apt-get update sudo apt-get install -y --no-install-recommends \ - gdb g++ ninja-build zip \ + gdb g++ ninja-build unzip \ libblas-dev liblapack-dev liblapacke-dev \ openmpi-bin openmpi-common libopenmpi-dev @@ -50,6 +50,7 @@ runs: shell: bash run: | brew update + brew trust aws/tap # suppress warning in github actions brew install openmpi xcodebuild -showComponent MetalToolchain sysctl -a | grep machdep.cpu diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bd94af697b..14b463e3a7 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -86,6 +86,7 @@ jobs: with: cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: false + - run: unzip -l wheelhouse/*.whl - uses: actions/upload-artifact@v7 with: name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} @@ -127,6 +128,7 @@ jobs: with: cmake-args: ${{ steps.setup.outputs.cmake-args }} build-frontend: false + - run: unzip -l wheelhouse/*.whl - uses: actions/upload-artifact@v7 with: name: backend-${{ matrix.toolkit }}-${{ runner.os }}-${{ runner.arch }} @@ -168,6 +170,8 @@ jobs: macos-target: '26.2' cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: ${{ matrix.python-version == '3.10' }} + - name: Display content of the wheels + run: for WHEEL in wheelhouse/*.whl; do unzip -l $WHEEL; echo; done - name: Upload frontend packages uses: actions/upload-artifact@v7 with: diff --git a/setup.py b/setup.py index 835375599e..52cc4f7492 100644 --- a/setup.py +++ b/setup.py @@ -53,7 +53,22 @@ def get_version(): return version -build_stage = int(os.environ.get("MLX_BUILD_STAGE", 0)) +# Release builds for PyPi are separated into 2 packages: +# +# Frontend package: +# - Triggered with `MLX_BUILD_FRONTEND_PACKAGE=1` +# - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) +# - Wheel has Python ABI and platform tags +# - Wheel should be built for the cross-product of python version and platforms +# - Package name is "mlx" and it depends on backend packages (e.g. mlx-metal, mlx-cuda) +# Backend package: +# - Triggered with `MLX_BUILD_BACKEND_PACKAGE=1` +# - Include headers and backend binaries. +# - Wheel has only platform tags +# - Wheel should be built only for different platforms +# - Package name is back-end specific, e.g mlx-metal, mlx-cuda +build_frontend = int(os.environ.get("MLX_BUILD_FRONTEND_PACKAGE", 0)) +build_backend = int(os.environ.get("MLX_BUILD_BACKEND_PACKAGE", 0)) build_macos = platform.system() == "Darwin" build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "") @@ -77,9 +92,7 @@ def finalize_options(self) -> None: self.build_temp = os.path.dirname(self.build_temp) def build_extension(self, ext: CMakeExtension) -> None: - # Must be in this form due to bug in .resolve() only fixed in Python 3.10+ - ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) # type: ignore[no-untyped-call] - extdir = ext_fullpath.parent.resolve() + extdir = self._get_ext_dir(ext) debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug cfg = "Debug" if debug else "Release" @@ -88,17 +101,9 @@ def build_extension(self, ext: CMakeExtension) -> None: if not build_temp.exists(): build_temp.mkdir(parents=True) - install_prefix = extdir - pybind_out_dir = extdir - if build_stage == 1: - # Don't include MLX libraries in the wheel - install_prefix = build_temp - elif build_stage == 2: - # Don't include Python bindings in the wheel - pybind_out_dir = build_temp cmake_args = [ - f"-DCMAKE_INSTALL_PREFIX={install_prefix}", - f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={pybind_out_dir}", + f"-DCMAKE_INSTALL_PREFIX={extdir}", + f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={extdir}", f"-DCMAKE_BUILD_TYPE={cfg}", f"-DPython_EXECUTABLE={sys.executable}", "-DMLX_BUILD_PYTHON_BINDINGS=ON", @@ -113,8 +118,7 @@ def build_extension(self, ext: CMakeExtension) -> None: if "CMAKE_ARGS" in os.environ: cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item] - # For release wheel force building for all supported arches. - if build_stage == 2 and build_cuda: + if build_backend and build_cuda: # Last arch is always real and virtual for forward-compatibility cuda_archs = [ "75-real", @@ -186,16 +190,50 @@ def run(self): ["cmake", "--install", build_temp, "--component", "core_stub"], check=True, ) + # Copy the type stubs to extdir so they are included in wheels. + stubs_dir = Path("python/mlx/core") + if stubs_dir.exists(): + extdir = self._get_ext_dir(ext) + self.copy_tree(stubs_dir, extdir / "core") + + def _get_ext_dir(self, ext): + # Must be in this form due to bug in .resolve() only fixed in Python 3.10+ + ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) # type: ignore[no-untyped-call] + return ext_fullpath.parent.resolve() class MLXBdistWheel(bdist_wheel): def get_tag(self) -> tuple[str, str, str]: impl, abi, plat_name = super().get_tag() - if build_stage == 2: + if build_backend: impl = self.python_tag abi = "none" return (impl, abi, plat_name) + def write_wheelfile(self, *args, **kwargs) -> None: + super().write_wheelfile(*args, **kwargs) + + mlx_dir = Path(self.bdist_dir, "mlx") + + def is_backend_file(file): + if file.is_relative_to(Path(mlx_dir, "lib")): + return True + if file.is_relative_to(Path(mlx_dir, "include")): + return True + if file.is_relative_to(Path(mlx_dir, "share")): + return True + if file.suffix == ".dll": + return True + return False + + if build_frontend or build_backend: + for file in Path(self.bdist_dir).rglob("*"): + if not file.is_relative_to(mlx_dir) or not file.is_file(): + continue + bf = is_backend_file(file) + if (build_frontend and bf) or (build_backend and not bf): + file.unlink() + # Read the content of README.md with open(Path(__file__).parent / "README.md", encoding="utf-8") as f: @@ -260,24 +298,8 @@ def get_tag(self) -> tuple[str, str, str]: } install_requires = [] - # Release builds for PyPi are in two stages. - # Each stage should be run from a clean build: - # python setup.py clean --all - # - # Stage 1: - # - Triggered with `MLX_BUILD_STAGE=1` - # - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) - # - Wheel has Python ABI and platform tags - # - Wheel should be built for the cross-product of python version and platforms - # - Package name is mlx and it depends on subpackage in stage 2 (e.g. mlx-metal) - # Stage 2: - # - Triggered with `MLX_BUILD_STAGE=2` - # - Includes only backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) - # - Wheel has only platform tags - # - Wheel should be built only for different platforms - # - Package name is back-end specific, e.g mlx-metal - if build_stage != 2: - if build_stage == 1: + if not build_backend: + if build_frontend: install_requires.append( f'mlx-metal=={version}; platform_system == "Darwin"' ) From 8e00a2d9dd143908a4b16a7f0802d5799afd557b Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 17 Aug 2026 11:19:41 +0900 Subject: [PATCH 31/84] Make mx.compile cache erasing thread safe (#4248) Co-authored-by: yentur --- mlx/compile.cpp | 91 +++++++++++++++++++++++------------- mlx/compile_impl.h | 16 ++++--- mlx/scheduler.cpp | 2 + python/src/random.cpp | 4 ++ python/src/random.h | 3 ++ python/src/stream.cpp | 6 ++- python/src/transforms.cpp | 37 ++++----------- python/tests/test_compile.py | 47 +++++++++++++++++++ 8 files changed, 138 insertions(+), 68 deletions(-) diff --git a/mlx/compile.cpp b/mlx/compile.cpp index 12d7397be4..bb17c44962 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -298,7 +299,7 @@ std::uintptr_t get_function_address(const std::function& fun) { return reinterpret_cast(*fun_ptr); } -class CompilerCache { +class CompileCache { public: struct CacheEntry { CacheEntry(Stream stream, bool shapeless) @@ -313,15 +314,38 @@ class CompilerCache { std::shared_ptr extra; }; - // Returns a reference to a CacheEntry which can be updated - // by the caller to avoid copying large tapes / inputs / outputs - CacheEntry& find( + CompileCache() { + // Make sure the allocator is fully initialized before the compiler cache. + allocator::allocator(); + } + + // Returns a reference to a CacheEntry which can be updated by the caller to + // avoid copying large tapes / inputs / outputs, with the shared_ptr of + // entries to avoid getting erased during compilation. + std::tuple>> find( std::uintptr_t fun_id, const std::vector& inputs, bool shapeless, const std::vector& constants) { - // Find the cache entries for |fun_id|. - std::vector& entries = cache_[fun_id]; + // Find the cache entries for |fun_id| in a thread-safe way. + auto entries_ptr = [&]() { + // Lookup with shared lock. + { + std::shared_lock lock(mutex_); + auto it = cache_.find(fun_id); + if (it != cache_.end()) { + return it->second; + } + } + // Insertion with exclusive lock. + std::unique_lock lock(mutex_); + auto& ptr = cache_[fun_id]; + if (!ptr) { + ptr = std::make_shared>(); + } + return ptr; + }(); + auto& entries = *entries_ptr; // Compare if 2 arrays have same shape and dtype. auto has_same_shape_and_dtype = [shapeless]( @@ -359,40 +383,37 @@ class CompilerCache { // Check the inputs match and return if so if (has_same_shape_and_dtype(inputs, entry.inputs) && constants == entry.constants) { - return entry; + return {entry, std::move(entries_ptr)}; } } // Otherwise append a new cache entry entries.push_back(CacheEntry{stream, shapeless}); - return entries.back(); + return {entries.back(), std::move(entries_ptr)}; } void erase(std::uintptr_t fun_id) { + std::unique_lock lock(mutex_); cache_.erase(fun_id); } void clear() { + std::unique_lock lock(mutex_); cache_.clear(); } - bool empty() { - return cache_.empty(); - } - private: - CompilerCache() { - // Make sure the allocator is fully - // initialized before the compiler cache - allocator::allocator(); - } - - friend CompilerCache& compiler_cache(); - std::unordered_map> cache_; + // The cache may get its key erased from a separate thread, but its value is + // only added and modified in the thread of creation. + // Put value in a shared_ptr to avoid race condition when erasing happened + // during compilation for the same function. + std::unordered_map>> + cache_; + std::shared_mutex mutex_; }; -CompilerCache& compiler_cache() { - static thread_local CompilerCache compiler_cache_; - return compiler_cache_; +std::shared_ptr& compile_cache_unsafe() { + static thread_local auto cache = std::make_shared(); + return cache; } std::tuple, std::vector, std::shared_ptr> @@ -1120,7 +1141,9 @@ ArrayFnWithExtra compile( } // Find a cache entry with the correct inputs - auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants); + auto [entry, entries_ptr] = + compile_cache_unsafe()->find(fun_id, inputs, shapeless, constants); + static_assert(std::is_reference_v); // No matching cache entry existed, so compile if (entry.empty) { @@ -1192,16 +1215,20 @@ std::function(const std::vector&)> compile( }; } -void compile_erase(std::uintptr_t fun_id) { - detail::compiler_cache().erase(fun_id); +CompileCacheWeakPtr compile_cache() { + return compile_cache_unsafe(); } -void compile_clear_cache() { - detail::compiler_cache().clear(); +void compile_erase(const CompileCacheWeakPtr& cache, std::uintptr_t fun_id) { + if (auto p = cache.lock()) { + p->erase(fun_id); + } } -bool compile_cache_empty() { - return detail::compiler_cache().empty(); +void compile_clear_cache(const CompileCacheWeakPtr& cache) { + if (auto p = cache.lock()) { + p->clear(); + } } } // namespace detail @@ -1221,8 +1248,8 @@ std::function(const std::vector&)> compile( auto pfun = std::shared_ptr< std::function(const std::vector&)>>( new std::function(const std::vector&)>{fun}, - [](auto* p) { - detail::compile_erase(reinterpret_cast(p)); + [cache = detail::compile_cache()](auto* p) { + detail::compile_erase(cache, reinterpret_cast(p)); delete p; }); fun_id = reinterpret_cast(pfun.get()); diff --git a/mlx/compile_impl.h b/mlx/compile_impl.h index cd3313be2c..1afe07b931 100644 --- a/mlx/compile_impl.h +++ b/mlx/compile_impl.h @@ -27,15 +27,19 @@ MLX_API ArrayFnWithExtra compile( bool shapeless, std::vector constants); -// Erase cached compile functions -MLX_API void compile_erase(std::uintptr_t fun_id); +// Get the compiler cache of current thread. +class CompileCache; +using CompileCacheWeakPtr = std::weak_ptr; +MLX_API CompileCacheWeakPtr compile_cache(); + +// Erase cached compile function. +MLX_API void compile_erase( + const CompileCacheWeakPtr& cache, + std::uintptr_t fun_id); // Clear the compiler cache causing a recompilation of all compiled functions // when called again. -MLX_API void compile_clear_cache(); - -// Return true if the cache is empty. -MLX_API bool compile_cache_empty(); +MLX_API void compile_clear_cache(const CompileCacheWeakPtr& cache); bool compile_available_for_device(const Device& device); diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp index 7507917f5b..6a0fdcf942 100644 --- a/mlx/scheduler.cpp +++ b/mlx/scheduler.cpp @@ -3,6 +3,7 @@ #include "mlx/scheduler.h" #include "mlx/backend/cpu/eval.h" #include "mlx/backend/gpu/eval.h" +#include "mlx/compile_impl.h" #include "mlx/utils.h" namespace mlx::core { @@ -27,6 +28,7 @@ void synchronize() { } void clear_streams() { + detail::compile_clear_cache(detail::compile_cache()); cpu::clear_streams(); gpu::clear_streams(); } diff --git a/python/src/random.cpp b/python/src/random.cpp index 8485faea41..10b82b8921 100644 --- a/python/src/random.cpp +++ b/python/src/random.cpp @@ -64,6 +64,10 @@ PyKeySequence& default_key() { return ks; } +void reset_random_state() { + default_key().reset(); +} + // A process-global sentinel for `mx.random.state`. Since it is the same object // on every thread, capturing it (e.g. with `mx.compile`) is thread-independent; // the pytree traversal in trees.cpp resolves it to the calling thread's key. diff --git a/python/src/random.h b/python/src/random.h index 2baf9d92f1..02d81d4c74 100644 --- a/python/src/random.h +++ b/python/src/random.h @@ -9,6 +9,9 @@ namespace mx = mlx::core; namespace nb = nanobind; +// Clear the `mx.random.state` python object in current thread. +void reset_random_state(); + // The process-global `mx.random.state` sentinel. nb::object random_state_sentinel(); diff --git a/python/src/stream.cpp b/python/src/stream.cpp index 004301a45c..467518e991 100644 --- a/python/src/stream.cpp +++ b/python/src/stream.cpp @@ -9,6 +9,7 @@ #include "mlx/stream.h" #include "mlx/utils.h" +#include "python/src/random.h" namespace mx = mlx::core; namespace nb = nanobind; @@ -137,7 +138,10 @@ void init_stream(nb::module_& m) { R"pbdoc(Make a new stream that will be unique per thread.)pbdoc"); m.def( "clear_streams", - &mx::clear_streams, + []() { + reset_random_state(); + mx::clear_streams(); + }, R"pbdoc(Destroy all streams created in current thread.)pbdoc"); nb::class_(m, "StreamContext", R"pbdoc( diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index 1d7aa8b9b1..1ec20a1375 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -406,29 +406,13 @@ auto py_vmap( }; } -void ensure_compile_cache_cleanup() { - // Make sure each thread using mx.compile would clear its compile cache - // before python interpreter exits. - struct ThreadCleanup { - ~ThreadCleanup() { - if (!mx::detail::compile_cache_empty()) { - nb::gil_scoped_acquire gil; - mx::detail::compile_clear_cache(); - } - } - }; - static thread_local auto clear_cache = []() { - mx::detail::compile_clear_cache(); - return ThreadCleanup{}; - }(); -} - struct PyCompiledFun { nb::callable fun; std::uintptr_t fun_id; nb::object captured_inputs; nb::object captured_outputs; bool shapeless; + mx::detail::CompileCacheWeakPtr cache; // Data to attach to the compiled function that contains the python output // structure and the number of arrays in said structure. @@ -456,15 +440,16 @@ struct PyCompiledFun { PyCompiledFun& operator=(PyCompiledFun&& other) = delete; PyCompiledFun(PyCompiledFun&& other) : fun(std::move(other.fun)), - fun_id(reinterpret_cast(fun.ptr())) { + fun_id(reinterpret_cast(fun.ptr())), + captured_inputs(std::move(other.captured_inputs)), + captured_outputs(std::move(other.captured_outputs)), + shapeless(other.shapeless), + cache(other.cache) { other.fun_id = 0; - captured_inputs = std::move(other.captured_inputs); - captured_outputs = std::move(other.captured_outputs); - shapeless = other.shapeless; }; nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) { - ensure_compile_cache_cleanup(); + cache = mx::detail::compile_cache(); // Flat array inputs std::vector inputs; @@ -599,7 +584,7 @@ struct PyCompiledFun { ~PyCompiledFun() { nb::gil_scoped_acquire gil; - mx::detail::compile_erase(fun_id); + mx::detail::compile_erase(cache, fun_id); fun.reset(); captured_inputs.reset(); captured_outputs.reset(); @@ -1553,10 +1538,4 @@ void init_transforms(nb::module_& m) { A callable that recomputes intermediate states during gradient computation. )pbdoc"); - - // Ensure the main thread cleanup will happen before the interpreter goes - // away. As a result if the other threads join the main thread we should have - // a clean tear-down. - auto atexit = nb::module_::import_("atexit"); - atexit.attr("register")(nb::cpp_function(&mx::detail::compile_clear_cache)); } diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 76e8916538..1e1b20b05d 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -87,6 +87,7 @@ def worker(): results.append((y.item(), z.item())) except Exception as e: errors.append(e) + mx.clear_streams() for _ in range(3): thread = threading.Thread(target=worker) @@ -98,6 +99,50 @@ def worker(): raise errors[0] self.assertEqual(results, [(2.0, 2.0)] * 3) + def test_compile_release_on_another_thread(self): + # A function traced on one thread but released on another must still + # drop its cache entry, otherwise a later compile of the same id gets + # handed the dead function's tape instead of being traced again. + traces = [] + + def fun(x): + traces.append(1) + return x + 1 + + holder = {} + traced = threading.Event() + released = threading.Event() + errors = [] + + def worker(): + try: + holder["fn"] = mx.compile(fun) + mx.eval(holder["fn"](mx.array([1.0]))) + traced.set() + self.assertTrue(released.wait(10)) + # The same callable, so the same id. + fn = mx.compile(fun) + mx.eval(fn(mx.array([1.0]))) + except Exception as e: + errors.append(e) + finally: + traced.set() + mx.clear_streams() + + # The tracing thread has to outlive the release, on exit it would tear + # down its cache anyway. + thread = threading.Thread(target=worker) + thread.start() + self.assertTrue(traced.wait(10)) + holder.clear() + gc.collect() + released.set() + thread.join() + + if errors: + raise errors[0] + self.assertEqual(len(traces), 2) + def test_compile_grad(self): def loss_fn(x): return mx.exp(x).sum() @@ -449,6 +494,7 @@ def test_compile_rng_across_threads(self): def grab(): state_from_thread["s"] = mx.random.state + mx.clear_streams() t = threading.Thread(target=grab) t.start() @@ -482,6 +528,7 @@ def worker(): results["seed_changes"] = not bool( mx.allclose(c, e, 1e-2, 1e-2).item() ) + mx.clear_streams() t = threading.Thread(target=worker) t.start() From 9a795735ad9a42664e08f42361b405ed570bcf1a Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 17 Aug 2026 13:09:50 +0900 Subject: [PATCH 32/84] Add builds for free-threaded python (#3812) --- .github/workflows/release.yml | 4 ++-- python/src/CMakeLists.txt | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 14b463e3a7..d10493efff 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -58,7 +58,7 @@ jobs: matrix: os: ['Linux', 'Windows'] arch: ['x86_64', 'aarch64'] - python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] + python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t'] # There is no cp310 binary for Windows on arm. exclude: - os: 'Windows' @@ -140,7 +140,7 @@ jobs: if: github.repository == 'ml-explore/mlx' strategy: matrix: - python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] + python-version: *pyver runs-on: 'macos-26' env: *build-env steps: diff --git a/python/src/CMakeLists.txt b/python/src/CMakeLists.txt index 447271500b..0798add410 100644 --- a/python/src/CMakeLists.txt +++ b/python/src/CMakeLists.txt @@ -2,6 +2,7 @@ nanobind_add_module( core NB_STATIC STABLE_ABI + FREE_THREADED LTO NOMINSIZE NB_DOMAIN From 98a188cb8e836a527a1a26b3ba5557bdfdb18a91 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 17 Aug 2026 02:46:39 -0700 Subject: [PATCH 33/84] Fix int32 overflow in concatenate/repeat/kron (#4303) --- mlx/ops.cpp | 16 +++++++++++----- mlx/primitives.cpp | 4 +++- python/tests/test_ops.py | 21 +++++++++++++++++++++ 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 1c09e728d0..5bc69eb2f5 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1312,7 +1312,9 @@ array concatenate( }; auto shape = arrays[0].shape(); - shape[ax] = 0; + // Accumulate the concatenation axis in 64 bits so a total that does not fit + // in a shape dimension is reported rather than silently wrapping. + int64_t concat_size = 0; // Make the output shape and validate that all arrays have the same shape // except for the concatenation axis. for (auto& a : arrays) { @@ -1331,8 +1333,9 @@ array concatenate( throw_invalid_shapes(); } } - shape[ax] += a.shape(ax); + concat_size += a.shape(ax); } + shape[ax] = safe_cast(concat_size, "concatenate"); // Promote all the arrays to the same type auto dtype = result_type(arrays); @@ -1408,7 +1411,8 @@ array repeat(const array& arr, int repeats, int axis, StreamOrDevice s) { // Reshape back into a contiguous array where S_axis is now S_axis * repeats shape.erase(shape.begin() + axis + 1); - shape[axis] *= repeats; + shape[axis] = + safe_cast(static_cast(shape[axis]) * repeats, "repeat"); out = reshape(out, shape, s); return out; @@ -3646,11 +3650,13 @@ array kron(const array& a, const array& b, StreamOrDevice s /* = {} */) { for (int i = ndim - 1, j = a.ndim() - 1; j >= 0; j--, i--) { a_shape[2 * i] = a.shape(j); - out_shape[i] *= a.shape(j); + out_shape[i] = + safe_cast(static_cast(out_shape[i]) * a.shape(j), "kron"); } for (int i = ndim - 1, j = b.ndim() - 1; j >= 0; j--, i--) { b_shape[2 * i + 1] = b.shape(j); - out_shape[i] *= b.shape(j); + out_shape[i] = + safe_cast(static_cast(out_shape[i]) * b.shape(j), "kron"); } return reshape( diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 3c3d4fc604..d4975e87f7 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -1187,9 +1187,11 @@ bool Concatenate::is_equivalent(const Primitive& other) const { std::vector Concatenate::output_shapes( const std::vector& inputs) { auto shape = inputs[0].shape(); + int64_t concat_size = shape[axis_]; for (int i = 1; i < inputs.size(); ++i) { - shape[axis_] += inputs[i].shape(axis_); + concat_size += inputs[i].shape(axis_); } + shape[axis_] = safe_cast(concat_size, "concatenate"); return {std::move(shape)}; } diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a09b96669e..86f0226856 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -126,6 +126,27 @@ def test_shape_overflow_error(self): mx.broadcast_to(a, [too_big, 1]) self.assertIn(str(too_big), str(cm.exception)) + # A concatenation axis that does not fit is computed rather than given, + # so it has to be reported instead of wrapping into a bogus dimension. + # These stay lazy, so nothing near this size is allocated. + big = mx.zeros(2**30) + for parts in (3, 4, 5): + with self.assertRaises(OverflowError) as cm: + mx.concatenate([big] * parts) + self.assertIn(str(2**30 * parts), str(cm.exception)) + + # repeat and kron multiply a dimension, and used to wrap into a + # negative or zero one that only surfaced later as a confusing reshape + # error naming a shape the caller never asked for. + for parts in (2, 3, 4): + with self.assertRaises(OverflowError) as cm: + mx.repeat(big, parts) + self.assertIn(str(2**30 * parts), str(cm.exception)) + + with self.assertRaises(OverflowError) as cm: + mx.kron(mx.zeros(2**16), mx.zeros(2**16)) + self.assertIn(str(2**32), str(cm.exception)) + # Negative overflow (< int32 min) is caught too. too_negative = -(2**31) - 1 with self.assertRaises(OverflowError) as cm: From e0fe5403e0000dcd8aa31a966cb22cb4990060e2 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Mon, 17 Aug 2026 02:47:07 -0700 Subject: [PATCH 34/84] python: Widen list elements that do not fit in int32 to int64 (#4305) --- python/src/convert.cpp | 31 ++++++++++++++++++++++++------- python/tests/test_array.py | 27 +++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 7 deletions(-) diff --git a/python/src/convert.cpp b/python/src/convert.cpp index 9941358e84..a3da76ef53 100644 --- a/python/src/convert.cpp +++ b/python/src/convert.cpp @@ -505,7 +505,8 @@ PyScalarT validate_shape( T list, const mx::Shape& shape, int idx, - bool& all_python_primitive_elements) { + bool& all_python_primitive_elements, + bool& has_wide_int) { if (idx >= shape.size()) { throw std::invalid_argument("Initialization encountered extra dimension."); } @@ -524,13 +525,18 @@ PyScalarT validate_shape( PyScalarT t; if (nb::isinstance(l)) { t = validate_shape( - nb::cast(l), shape, idx + 1, all_python_primitive_elements); + nb::cast(l), + shape, + idx + 1, + all_python_primitive_elements, + has_wide_int); } else if (nb::isinstance(*list.begin())) { t = validate_shape( nb::cast(l), shape, idx + 1, - all_python_primitive_elements); + all_python_primitive_elements, + has_wide_int); } else if (nb::isinstance(l)) { all_python_primitive_elements = false; auto arr = nb::cast(l); @@ -549,6 +555,13 @@ PyScalarT validate_shape( t = pybool; } else if (nb::isinstance(l)) { t = pyint; + // Match the scalar path, which widens to int64 rather than failing + // when a python int does not fit in int32. + auto val = nb::cast(l); + if (val > std::numeric_limits::max() || + val < std::numeric_limits::min()) { + has_wide_int = true; + } } else if (nb::isinstance(l)) { t = pyfloat; } else if (PyComplex_Check(l.ptr())) { @@ -594,7 +607,8 @@ mx::array array_from_list_impl( T pl, const PyScalarT& inferred_type, std::optional specified_type, - const mx::Shape& shape) { + const mx::Shape& shape, + bool has_wide_int) { // Make the array switch (inferred_type) { case pybool: { @@ -603,7 +617,8 @@ mx::array array_from_list_impl( return mx::array(vals.begin(), shape, specified_type.value_or(mx::bool_)); } case pyint: { - auto dtype = specified_type.value_or(mx::int32); + auto dtype = + specified_type.value_or(has_wide_int ? mx::int64 : mx::int32); if (dtype == mx::int64) { std::vector vals; fill_vector(pl, vals); @@ -663,11 +678,13 @@ mx::array array_from_list_impl(T pl, std::optional dtype) { // Validate the shape and type bool all_python_primitive_elements = true; - auto type = validate_shape(pl, shape, 0, all_python_primitive_elements); + bool has_wide_int = false; + auto type = + validate_shape(pl, shape, 0, all_python_primitive_elements, has_wide_int); if (all_python_primitive_elements) { // `pl` does not contain mlx arrays - return array_from_list_impl(pl, type, dtype, shape); + return array_from_list_impl(pl, type, dtype, shape, has_wide_int); } // `pl` contains mlx arrays diff --git a/python/tests/test_array.py b/python/tests/test_array.py index aee2ab8872..a920b39b38 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -523,6 +523,33 @@ def test_double_keeps_precision(self): out = mx.array([x], dtype=mx.float64).item() self.assertEqual(out, x) + def test_construction_from_lists_wide_ints(self): + # A python int that does not fit in int32 widens to int64, the same + # rule the scalar path already uses. It used to raise std::bad_cast. + for value in (2**31, 2**40, -(2**31) - 1, -(2**40)): + for make in ( + lambda v: [v], + lambda v: (v,), + lambda v: [[v]], + lambda v: [v, 1], + ): + x = mx.array(make(value)) + self.assertEqual(x.dtype, mx.int64, msg=f"{value} {make(value)}") + self.assertEqual(x.flatten()[0].item(), value) + self.assertEqual(mx.array(value).dtype, mx.int64) + + # Values that still fit keep int32, including both boundaries. + for value in (0, 1, 2**31 - 1, -(2**31)): + x = mx.array([value]) + self.assertEqual(x.dtype, mx.int32, msg=str(value)) + self.assertEqual(x[0].item(), value) + + # An explicit dtype still wins. + self.assertEqual(mx.array([2**40], mx.int64).dtype, mx.int64) + self.assertEqual(mx.array([1, 2], mx.int64).dtype, mx.int64) + # A float in the list still makes it float, not int64. + self.assertEqual(mx.array([2**40, 1.5]).dtype, mx.float32) + def test_construction_from_lists_of_mlx_arrays(self): dtypes = [ mx.bool_, From 06f154bcf55f5a6be304f1aad2e85ed5b7a39316 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 17 Aug 2026 19:39:04 +0900 Subject: [PATCH 35/84] Propagate CPU errors to events (#3742) Co-authored-by: Alessio Pollero --- mlx/array.h | 1 + mlx/backend/common/load.cpp | 2 +- mlx/backend/cpu/encoder.h | 9 +- mlx/backend/cuda/event.cu | 190 +++++++++++++++++------------------ mlx/backend/cuda/event.h | 20 +++- mlx/backend/cuda/fence.cpp | 20 ++-- mlx/backend/metal/device.cpp | 43 ++++---- mlx/backend/metal/device.h | 12 +-- mlx/backend/metal/event.cpp | 37 +++---- mlx/backend/metal/event.h | 10 +- mlx/backend/metal/fence.cpp | 3 +- mlx/backend/no_gpu/event.cpp | 48 +++++---- mlx/backend/no_gpu/fence.cpp | 50 +++------ mlx/error.h | 49 +++++++++ mlx/event.h | 32 ++++++ mlx/fence.h | 7 +- mlx/scheduler.cpp | 116 +++++++++++++++++++-- mlx/scheduler.h | 70 ++++--------- python/tests/test_load.py | 35 +++++++ 19 files changed, 459 insertions(+), 295 deletions(-) create mode 100644 mlx/error.h diff --git a/mlx/array.h b/mlx/array.h index 8e14ca4726..3f45e9cb9d 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -426,6 +426,7 @@ class MLX_API array { } void detach_event() const { + array_desc_->event.check_error(); array_desc_->event = Event{}; } diff --git a/mlx/backend/common/load.cpp b/mlx/backend/common/load.cpp index ce41963de7..b53c92483c 100644 --- a/mlx/backend/common/load.cpp +++ b/mlx/backend/common/load.cpp @@ -51,7 +51,7 @@ void Load::eval_cpu(const std::vector& inputs, array& out) { } }; auto fut = io::thread_pool().enqueue(std::move(read_task)).share(); - scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.wait(); }); + scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.get(); }); } } // namespace mlx::core diff --git a/mlx/backend/cpu/encoder.h b/mlx/backend/cpu/encoder.h index cd015623f6..eb45d64ca0 100644 --- a/mlx/backend/cpu/encoder.h +++ b/mlx/backend/cpu/encoder.h @@ -46,11 +46,10 @@ struct MLX_API CommandEncoder { auto task = std::bind(std::forward(f), std::forward(args)...); if (num_ops_ == 0) { scheduler::notify_new_task(stream_); - auto task_wrap = [s = stream_, task = std::move(task)]() mutable { - task(); - scheduler::notify_task_completion(s); - }; - scheduler::enqueue(stream_, std::move(task_wrap)); + scheduler::enqueue(stream_, std::move(task)); + // Notify completion separately as |task| may throw exception. + scheduler::enqueue( + stream_, [s = stream_] { scheduler::notify_task_completion(s); }); } else { scheduler::enqueue(stream_, std::move(task)); } diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index b73937ec38..d3b6f97f5d 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -113,10 +113,7 @@ void CudaEvent::init_pool() { cuda_event_pool(); } -// Wraps CudaEvent with a few features: -// 1. The class can be copied. -// 2. Make wait/record work with CPU streams. -// 3. Add checks for waiting on un-recorded event. +// Wraps CudaEvent so it can be copied. class CopyableCudaEvent { public: explicit CopyableCudaEvent(Device& d) @@ -126,32 +123,24 @@ class CopyableCudaEvent { cudaEventDisableTiming | cudaEventBlockingSync)) {} void wait() { + check_recorded(); event_->wait(); } void wait(Stream s) { - if (s.device == mlx::core::Device::cpu) { - scheduler::enqueue(s, [*this]() mutable { - check_recorded(); - event_->wait(); - }); - } else { - check_recorded(); - auto& encoder = cu::get_command_encoder(s); - encoder.commit(); - event_->wait(encoder.stream()); - } + assert(s.device == mlx::core::Device::gpu); + check_recorded(); + auto& encoder = cu::get_command_encoder(s); + encoder.commit(); + event_->wait(encoder.stream()); } void record(Stream s) { - if (s.device == mlx::core::Device::cpu) { - throw std::runtime_error("CudaEvent can not wait on CPU stream."); - } else { - auto& encoder = cu::get_command_encoder(s); - encoder.commit(); - event_->record(encoder.stream()); - recorded_ = true; - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = cu::get_command_encoder(s); + encoder.commit(); + event_->record(encoder.stream()); + recorded_ = true; } bool is_signaled() const { @@ -213,6 +202,11 @@ auto check_gpu_coherency() { return coherency; } +const CudaStream& signal_stream() { + static CudaStream stream(device(0)); + return stream; +} + AtomicEvent::AtomicEvent(Device& d) { void* buf; cudaError_t (*cuda_free)(void*); @@ -264,14 +258,11 @@ void AtomicEvent::wait(cudaStream_t stream, uint32_t value) { void AtomicEvent::wait(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::wait(s)"); - if (s.device == mlx::core::Device::cpu) { - scheduler::enqueue(s, [*this, value]() mutable { wait(value); }); - } else { - auto& encoder = get_command_encoder(s); - encoder.commit(); - wait(encoder.stream(), value); - encoder.add_completed_handler([buf = buf_]() {}); - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = get_command_encoder(s); + encoder.commit(); + wait(encoder.stream(), value); + encoder.add_completed_handler([buf = buf_]() {}); } void AtomicEvent::signal(uint32_t value) { @@ -289,17 +280,11 @@ void AtomicEvent::signal(cudaStream_t stream, uint32_t value) { void AtomicEvent::signal(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::signal(s)"); - if (s.device == mlx::core::Device::cpu) { - // Signal through a GPU stream so the atomic is updated in GPU - updating - // the atomic in CPU sometimes does not get GPU notified. - scheduler::enqueue( - s, [*this, value]() mutable { signal(signal_stream(), value); }); - } else { - auto& encoder = get_command_encoder(s); - encoder.commit(); - signal(encoder.stream(), value); - encoder.add_completed_handler([buf = buf_]() {}); - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = get_command_encoder(s); + encoder.commit(); + signal(encoder.stream(), value); + encoder.add_completed_handler([buf = buf_]() {}); } bool AtomicEvent::is_signaled(uint32_t val) const { @@ -319,9 +304,21 @@ uint32_t AtomicEvent::value() const { } } -const CudaStream& AtomicEvent::signal_stream() { - static CudaStream stream(device(0)); - return stream; +/////////////////////////////////////////////////////////////////////////////// +// EventImpl implementations +/////////////////////////////////////////////////////////////////////////////// + +void EventImpl::ensure_created(Stream s, uint64_t signal_value) { + if (is_created()) { + return; + } + auto& d = cu::device(s.device); + if (s.device == mlx::core::Device::cpu || signal_value > 1) { + nvtx3::mark("Using slow AtomicEvent"); + atomic = std::make_unique(d); + } else { + cuda = std::make_unique(d); + } } } // namespace cu @@ -330,86 +327,85 @@ const CudaStream& AtomicEvent::signal_stream() { // Event implementations /////////////////////////////////////////////////////////////////////////////// -namespace { - -struct EventImpl { - // CudaEvent is preferred when possible because it is fast, however we have - // to fallback to AtomicEvent in following cases: - // 1. the event is used to wait/signal a cpu stream; - // 2. signal value other than 1 has been specified. - std::unique_ptr cuda; - std::unique_ptr atomic; - - bool is_created() const { - return cuda || atomic; - } - - void ensure_created(Stream s, uint64_t signal_value) { - if (is_created()) { - return; - } - auto& d = cu::device(s.device); - if (s.device == mlx::core::Device::cpu || signal_value > 1) { - nvtx3::mark("Using slow AtomicEvent"); - atomic = std::make_unique(d); - } else { - cuda = std::make_unique(d); - } - } -}; - -} // namespace - Event::Event(Stream s) : stream_(s) { - event_ = std::shared_ptr( - new EventImpl(), [](void* ptr) { delete static_cast(ptr); }); + event_ = std::make_shared(); } void Event::wait() { - auto* event = static_cast(event_.get()); - assert(event->is_created()); - if (event->cuda) { + check_error(); + auto& event = cast(); + assert(event.is_created()); + if (event.cuda) { assert(value() == 1); - event->cuda->wait(); + event.cuda->wait(); } else { - event->atomic->wait(value()); + event.atomic->wait(value()); } CHECK_CUDA_ERROR(cudaPeekAtLastError()); + check_error(); } void Event::wait(Stream s) { - auto* event = static_cast(event_.get()); - assert(event->is_created()); - if (event->cuda) { + auto& event = cast(); + assert(event.is_created()); + if (event.cuda) { assert(value() == 1); - event->cuda->wait(s); + if (s.device == mlx::core::Device::cpu) { + scheduler::wait_event(s, *this, [value = value()](Event& self) { + self.cast().cuda->wait(); + }); + } else { + event.cuda->wait(s); + } } else { - event->atomic->wait(s, value()); + if (s.device == mlx::core::Device::cpu) { + scheduler::wait_event(s, *this, [value = value()](Event& self) { + self.cast().atomic->wait(value); + }); + } else { + event.atomic->wait(s, value()); + } } } void Event::signal(Stream s) { - auto* event = static_cast(event_.get()); - event->ensure_created(s, value()); - if (event->cuda) { + auto& event = cast(); + event.ensure_created(s, value()); + if (event.cuda) { assert(value() == 1); - event->cuda->record(s); + if (s.device == mlx::core::Device::cpu) { + throw std::runtime_error("CudaEvent can not wait on CPU stream."); + } else { + event.cuda->record(s); + } } else { - event->atomic->signal(s, value()); + if (s.device == mlx::core::Device::cpu) { + // Signal through a GPU stream so the atomic is updated in GPU - updating + // the atomic in CPU sometimes does not get GPU notified. + scheduler::signal_event(s, *this, [value = value()](Event& self) { + self.cast().atomic->signal(cu::signal_stream(), value); + }); + } else { + event.atomic->signal(s, value()); + } } } bool Event::is_signaled() const { - auto* event = static_cast(event_.get()); - if (!event->is_created()) { + auto& event = cast(); + if (!event.is_created()) { return false; } - if (event->cuda) { + if (event.cuda) { assert(value() == 1); - return event->cuda->is_signaled(); + return event.cuda->is_signaled(); } else { - return event->atomic->is_signaled(value()); + return event.atomic->is_signaled(value()); } } +std::atomic& Event::error() { + return cast().error; +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/event.h b/mlx/backend/cuda/event.h index 53afeb0117..fdeb6a0e78 100644 --- a/mlx/backend/cuda/event.h +++ b/mlx/backend/cuda/event.h @@ -13,6 +13,7 @@ namespace mlx::core::cu { +class CopyableCudaEvent; class Device; // RAII-managed move-only wrapper of cudaEvent_t. @@ -66,8 +67,6 @@ class AtomicEvent { uint32_t value() const; private: - const CudaStream& signal_stream(); - uint32_t* ptr() const { return static_cast(buf_.get()); } @@ -76,4 +75,21 @@ class AtomicEvent { std::shared_ptr buf_; }; +struct EventImpl { + std::atomic error; + + // CudaEvent is preferred when possible because it is fast, however we have + // to fallback to AtomicEvent in following cases: + // 1. the event is used to wait/signal a cpu stream; + // 2. signal value other than 1 has been specified. + std::unique_ptr cuda; + std::unique_ptr atomic; + + bool is_created() const { + return cuda || atomic; + } + + void ensure_created(Stream s, uint64_t signal_value); +}; + } // namespace mlx::core::cu diff --git a/mlx/backend/cuda/fence.cpp b/mlx/backend/cuda/fence.cpp index c6a41f0e60..3a3acdba09 100644 --- a/mlx/backend/cuda/fence.cpp +++ b/mlx/backend/cuda/fence.cpp @@ -9,22 +9,23 @@ namespace mlx::core { struct FenceImpl { uint32_t count; - cu::AtomicEvent event; + Event event; + + FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} }; Fence::Fence(Stream s) { - fence_ = std::shared_ptr( - new FenceImpl{0, cu::device(s.device)}, - [](void* ptr) { delete static_cast(ptr); }); + fence_ = std::make_shared(0, s); + // Ensure that we use AtomicEvent. + cast().event.cast().ensure_created(s, 2); } void Fence::wait(Stream s, const array&) { - auto* fence = static_cast(fence_.get()); - fence->event.wait(fence->count); + cast().event.wait(); } void Fence::update(Stream s, const array& a, bool cross_device) { - auto* fence = static_cast(fence_.get()); + auto& f = cast(); if (cross_device) { // Move to managed memory if there is a device switch auto& cbuf = @@ -35,8 +36,9 @@ void Fence::update(Stream s, const array& a, bool cross_device) { cu::allocator().move_to_unified_memory(cbuf, encoder.stream()); } } - fence->count++; - fence->event.signal(s, fence->count); + f.count++; + f.event.set_value(f.count); + f.event.signal(s); } } // namespace mlx::core diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 65df5c108c..2f25f894e4 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -496,19 +496,15 @@ void CommandEncoder::end_encoding() { all_inputs_.clear(); } -void CommandEncoder::signal_event( - std::shared_ptr event, - uint64_t value) { +void CommandEncoder::signal_event(Event event, uint64_t value) { end_encoding(); - buffer_->encodeSignalEvent(event->mtl_event(), value); + buffer_->encodeSignalEvent(event.cast().mtl_event(), value); signal_events_.push_back({std::move(event), value}); } -void CommandEncoder::wait_event( - std::shared_ptr event, - uint64_t value) { +void CommandEncoder::wait_event(Event event, uint64_t value) { end_encoding(); - buffer_->encodeWait(event->mtl_event(), value); + buffer_->encodeWait(event.cast().mtl_event(), value); wait_events_.push_back(std::move(event)); } @@ -525,35 +521,37 @@ void CommandEncoder::commit(std::function completion) { [&error_ = error_, wait_events = std::move(wait_events_), signal_events = std::move(signal_events_), - completion = std::move(completion)](MTL::CommandBuffer* cbuf) { + completion = std::move(completion)](MTL::CommandBuffer* cbuf) mutable { if (completion) { completion(); } // If any of the waited event has error in it, poison the encoder. for (auto& event : wait_events) { - if (event->error()) { - error_ = event->error(); + if (error_.store_if_valid(event.load_error())) { break; } } // Set error only when no error happended before, to preserve the // earliest error. - if (!error_ && cbuf->status() == MTL::CommandBufferStatusError) { - error_ = std::make_shared(fmt::format( - "[METAL] Command buffer execution failed: {}.", - cbuf->error()->localizedDescription()->utf8String())); + bool has_error = error_.valid(); + if (!has_error && cbuf->status() == MTL::CommandBufferStatusError) { + error_.set_message( + std::make_shared(fmt::format( + "[METAL] Command buffer execution failed: {}.", + cbuf->error()->localizedDescription()->utf8String()))); + has_error = true; } // Poison all the signaled events when error happened. - if (error_) { + if (has_error) { for (auto& [event, value] : signal_events) { - event->set_error(error_); + event.set_error(error_); } } // Metal won't signal the events for us on error, manually signal them // to avoid infinite waiting. if (cbuf->status() == MTL::CommandBufferStatusError) { for (auto& [event, value] : signal_events) { - event->signal(value); + event.cast().signal(value); } } }); @@ -570,20 +568,17 @@ void CommandEncoder::synchronize() { commit(); cbuf->waitUntilCompleted(); - if (error_ && !exiting_) { - auto error = std::move(error_); - throw std::runtime_error(*error); + if (!exiting_) { + error_.check(); } } MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() { if (!encoder_) { + error_.check(); encoder_ = NS::RetainPtr( buffer_->computeCommandEncoder(MTL::DispatchTypeConcurrent)); fence_ = NS::TransferPtr(device_.mtl_device()->newFence()); - // Reset error when user starts to encode new commands, they are supposed to - // have handled the error in synchronize() or Event::wait(). - error_.reset(); } return encoder_.get(); } diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index 3bb1e9e3b3..2d22283351 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -6,7 +6,6 @@ #include #include #include -#include #include #include @@ -21,7 +20,6 @@ using MTLFCList = std::vector>; class Device; -class EventImpl; class MLX_API CommandEncoder { public: @@ -92,8 +90,8 @@ class MLX_API CommandEncoder { void barrier(); void end_encoding(); - void wait_event(std::shared_ptr event, uint64_t value); - void signal_event(std::shared_ptr event, uint64_t value); + void wait_event(Event event, uint64_t value); + void signal_event(Event event, uint64_t value); bool needs_commit() const; void commit(std::function completion = nullptr); void synchronize(); @@ -119,11 +117,11 @@ class MLX_API CommandEncoder { uint64_t sets_attached_{0}; // The events hooked to current command buffer. - std::vector> wait_events_; - std::vector, uint64_t>> signal_events_; + std::vector wait_events_; + std::vector> signal_events_; // Error from previous commited command buffer. - std::shared_ptr error_; + Error error_; // Encoder for issuing GPU commands. // The members are used within a single ComputeCommandEncoder and will be diff --git a/mlx/backend/metal/event.cpp b/mlx/backend/metal/event.cpp index 77f48f0838..38a387c9c0 100644 --- a/mlx/backend/metal/event.cpp +++ b/mlx/backend/metal/event.cpp @@ -26,26 +26,13 @@ EventImpl::~EventImpl() { } void EventImpl::wait(uint64_t value) { - check_error(); mtl_event_->waitUntilSignaledValue(value, -1); // never times out - check_error(); } void EventImpl::signal(uint64_t value) { mtl_event_->setSignaledValue(value); } -void EventImpl::set_error(std::shared_ptr error) { - std::atomic_store(&error_, std::move(error)); -} - -void EventImpl::check_error() { - auto error = std::atomic_exchange(&error_, {}); - if (error) { - throw std::runtime_error(*error); - } -} - } // namespace metal /////////////////////////////////////////////////////////////////////////////// @@ -57,36 +44,40 @@ Event::Event(Stream stream) : stream_(stream) { } void Event::wait() { - static_cast(event_.get())->wait(value()); + check_error(); + cast().wait(value()); + check_error(); } void Event::wait(Stream stream) { - auto impl = std::static_pointer_cast(event_); if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [impl = std::move(impl), value = value()]() { - impl->wait(value); + scheduler::wait_event(stream, *this, [value = value()](Event& self) { + self.cast().wait(value); }); } else { auto& encoder = metal::get_command_encoder(stream); - encoder.wait_event(std::move(impl), value()); + encoder.wait_event(*this, value()); } } void Event::signal(Stream stream) { - auto impl = std::static_pointer_cast(event_); if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [impl = std::move(impl), value = value()]() { - impl->signal(value); + scheduler::signal_event(stream, *this, [value = value()](Event& self) { + self.cast().signal(value); }); } else { auto& encoder = metal::get_command_encoder(stream); - encoder.signal_event(std::move(impl), value()); + encoder.signal_event(*this, value()); } } bool Event::is_signaled() const { - auto* mtl_event = static_cast(event_.get())->mtl_event(); + auto* mtl_event = cast().mtl_event(); return mtl_event->signaledValue() >= value(); } +std::atomic& Event::error() { + return cast().error(); +} + } // namespace mlx::core diff --git a/mlx/backend/metal/event.h b/mlx/backend/metal/event.h index c5c82a7cd3..d1e43fa02f 100644 --- a/mlx/backend/metal/event.h +++ b/mlx/backend/metal/event.h @@ -12,20 +12,18 @@ class EventImpl { void wait(uint64_t value); void signal(uint64_t value); - void set_error(std::shared_ptr error); - void check_error(); - const auto& error() const { + auto& error() { return error_; } - auto* mtl_event() { + auto* mtl_event() const { return mtl_event_.get(); } private: - // TODO: Use std::atomic when it gets supported in Xcode. - std::shared_ptr error_; + // All streams outlive events so pointers would be always valid. + std::atomic error_; NS::SharedPtr mtl_event_; }; diff --git a/mlx/backend/metal/fence.cpp b/mlx/backend/metal/fence.cpp index 6fdd57a5f6..70dd0e33bd 100644 --- a/mlx/backend/metal/fence.cpp +++ b/mlx/backend/metal/fence.cpp @@ -41,8 +41,7 @@ struct FenceImpl { }; Fence::Fence(Stream stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - fence_ = std::shared_ptr(new FenceImpl(stream), dtor); + fence_ = std::make_shared(stream); } void Fence::wait(Stream stream, const array& x) { diff --git a/mlx/backend/no_gpu/event.cpp b/mlx/backend/no_gpu/event.cpp index 6dde047ab4..8966b77613 100644 --- a/mlx/backend/no_gpu/event.cpp +++ b/mlx/backend/no_gpu/event.cpp @@ -12,42 +12,52 @@ struct EventCounter { uint64_t value{0}; std::mutex mtx; std::condition_variable cv; + std::atomic error; + + void wait(uint64_t val) { + std::unique_lock lk(mtx); + if (value >= val) { + return; + } + cv.wait(lk, [this, val] { return value >= val; }); + } }; Event::Event(Stream stream) : stream_(stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - event_ = std::shared_ptr(new EventCounter{}, dtor); + event_ = std::make_shared(); } void Event::wait() { - auto ec = static_cast(event_.get()); - std::unique_lock lk(ec->mtx); - if (ec->value >= value()) { - return; - } - ec->cv.wait(lk, [value = value(), ec] { return ec->value >= value; }); + check_error(); + cast().wait(value()); + check_error(); } void Event::wait(Stream stream) { - scheduler::enqueue(stream, [*this]() mutable { wait(); }); + scheduler::wait_event(stream, *this, [value = value()](Event& self) { + self.cast().wait(value); + }); } void Event::signal(Stream stream) { - scheduler::enqueue(stream, [*this]() mutable { - auto ec = static_cast(event_.get()); + scheduler::signal_event(stream, *this, [value = value()](Event& self) { + auto& ec = self.cast(); { - std::lock_guard lk(ec->mtx); - ec->value = value(); + std::lock_guard lk(ec.mtx); + ec.value = value; } - ec->cv.notify_all(); + ec.cv.notify_all(); }); } bool Event::is_signaled() const { - auto ec = static_cast(event_.get()); - { - std::lock_guard lk(ec->mtx); - return (ec->value >= value()); - } + auto& ec = cast(); + std::lock_guard lk(ec.mtx); + return ec.value >= value(); +} + +std::atomic& Event::error() { + return cast().error; } + } // namespace mlx::core diff --git a/mlx/backend/no_gpu/fence.cpp b/mlx/backend/no_gpu/fence.cpp index cd66d23cfe..05852c860b 100644 --- a/mlx/backend/no_gpu/fence.cpp +++ b/mlx/backend/no_gpu/fence.cpp @@ -1,54 +1,30 @@ // Copyright © 2024 Apple Inc. -#include -#include - #include "mlx/fence.h" -#include "mlx/scheduler.h" +#include "mlx/event.h" namespace mlx::core { struct FenceImpl { - uint32_t count{0}; - uint32_t value{0}; - std::mutex mtx; - std::condition_variable cv; + uint32_t count; + Event event; + + FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} }; -Fence::Fence(Stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - fence_ = std::shared_ptr(new FenceImpl{}, dtor); +Fence::Fence(Stream s) { + fence_ = std::make_shared(0, s); } -void Fence::wait(Stream stream, const array&) { - auto& f = *static_cast(fence_.get()); - if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable { - auto& f = *static_cast(fence_.get()); - std::unique_lock lk(f.mtx); - if (f.value >= count) { - return; - } - f.cv.wait(lk, [&f, count] { return f.value >= count; }); - }); - } else { - throw std::runtime_error("[Fence::wait] Invalid stream."); - } +void Fence::wait(Stream s, const array&) { + cast().event.wait(s); } -void Fence::update(Stream stream, const array&, bool) { - auto& f = *static_cast(fence_.get()); +void Fence::update(Stream s, const array&, bool) { + auto& f = cast(); f.count++; - if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable { - auto& f = *static_cast(fence_.get()); - std::unique_lock lk(f.mtx); - f.value = count; - f.cv.notify_all(); - }); - } else { - throw std::runtime_error("[Fence::update] Invalid stream."); - } + f.event.set_value(f.count); + f.event.signal(s); } } // namespace mlx::core diff --git a/mlx/error.h b/mlx/error.h new file mode 100644 index 0000000000..ba1164f192 --- /dev/null +++ b/mlx/error.h @@ -0,0 +1,49 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include +#include + +namespace mlx::core { + +class Error { + public: + // TODO: Use std::atomic when it gets supported in Xcode. + using Message = std::shared_ptr; + + void set_message(Message msg) { + std::atomic_store(&message_, std::move(msg)); + } + + bool valid() const { + auto msg = std::atomic_load(&message_); + return msg.get(); + } + + // If |ptr| is a valid event, copy and return true. + bool store_if_valid(const Error* ptr) { + if (ptr && this != ptr) { + Message msg = std::atomic_load(&ptr->message_); + if (msg) { + set_message(std::move(msg)); + return true; + } + } + return false; + } + + // If current error is valid, throw and clear. + void check() { + auto msg = std::atomic_exchange(&message_, {}); + if (msg) { + throw std::runtime_error(*msg); + } + } + + private: + Message message_; +}; + +} // namespace mlx::core diff --git a/mlx/event.h b/mlx/event.h index 66a6a75df5..cf2d5cc7d6 100644 --- a/mlx/event.h +++ b/mlx/event.h @@ -5,6 +5,7 @@ #include #include +#include "mlx/error.h" #include "mlx/stream.h" namespace mlx::core { @@ -26,6 +27,26 @@ class Event { // Check if the event has been signaled at its current value bool is_signaled() const; + // Associate an error to the event + void set_error(Error& err) { + error().store(&err); + } + + // Get the error associated with the event + Error* load_error() const { + if (!valid()) { + return nullptr; + } + return error().load(); + } + + // Throw and clear the associated error + void check_error() { + if (auto* p = load_error(); p) { + p->check(); + } + } + // Check if the event is valid bool valid() const { return event_ != nullptr; @@ -47,7 +68,18 @@ class Event { return stream_; } + template + auto& cast() const { + return *static_cast(event_.get()); + } + private: + std::atomic& error(); + + const std::atomic& error() const { + return const_cast(this)->error(); + } + // Default constructed stream should never be used // since the event is not yet valid Stream stream_{0, Device::cpu}; diff --git a/mlx/fence.h b/mlx/fence.h index 0ececdb6d7..3fd5da333b 100644 --- a/mlx/fence.h +++ b/mlx/fence.h @@ -32,8 +32,13 @@ class Fence { void update(Stream stream, const array& x, bool cross_device); void wait(Stream stream, const array& x); + template + auto& cast() const { + return *static_cast(fence_.get()); + } + private: - std::shared_ptr fence_{nullptr}; + std::shared_ptr fence_; }; } // namespace mlx::core diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp index 6a0fdcf942..ec86091bfc 100644 --- a/mlx/scheduler.cpp +++ b/mlx/scheduler.cpp @@ -1,9 +1,12 @@ // Copyright © 2023-2026 Apple Inc. -#include "mlx/scheduler.h" +#include +#include + #include "mlx/backend/cpu/eval.h" #include "mlx/backend/gpu/eval.h" #include "mlx/compile_impl.h" +#include "mlx/scheduler.h" #include "mlx/utils.h" namespace mlx::core { @@ -35,6 +38,58 @@ void clear_streams() { namespace scheduler { +struct StreamThread { + std::mutex mtx; + std::queue> q; + std::condition_variable cond; + bool stop; + std::thread thread; + Error error; + + StreamThread() : stop(false), thread(&StreamThread::thread_fn, this) {} + + ~StreamThread() { + { + std::lock_guard lk(mtx); + stop = true; + } + cond.notify_one(); + thread.join(); + } + + void thread_fn() { + while (true) { + std::function task; + { + std::unique_lock lk(mtx); + cond.wait(lk, [this] { return !this->q.empty() || this->stop; }); + if (q.empty() && stop) { + return; + } + task = std::move(q.front()); + q.pop(); + } + + task(); + } + } + + void enqueue(std::function f) { + if (is_main_thread()) { + error.check(); + } + { + std::lock_guard lk(mtx); + if (stop) { + throw std::runtime_error( + "Cannot enqueue work after stream is stopped."); + } + q.emplace(std::move(f)); + } + cond.notify_one(); + } +}; + Scheduler::Scheduler() { is_main_thread(); gpu::init(); @@ -43,23 +98,62 @@ Scheduler::Scheduler() { Scheduler::~Scheduler() = default; void Scheduler::enqueue(Stream s, std::function task) { - StreamThread* st = nullptr; + auto& st = get_thread(s); + st.enqueue([&st, task = std::move(task)]() mutable { + try { + task(); + } catch (const std::exception& error) { + // Set error to stream only when no error happended before, to preserve + // the earliest error. + if (!st.error.valid()) { + st.error.set_message(std::make_shared(error.what())); + } + } + }); +} + +void Scheduler::wait_event( + Stream s, + Event event, + std::function task) { + assert(s.device == Device::cpu); + auto& st = get_thread(s); + st.enqueue([&st, event = std::move(event), task = std::move(task)]() mutable { + task(event); + // Poison current stream if the waited event has error. + st.error.store_if_valid(event.load_error()); + }); +} + +void Scheduler::signal_event( + Stream s, + Event event, + std::function task) { + assert(s.device == Device::cpu); + auto& st = get_thread(s); + st.enqueue([&st, event = std::move(event), task = std::move(task)]() mutable { + // Poison the signal event if current stream has error. + if (st.error.valid()) { + event.set_error(st.error); + } + task(event); + }); +} + +StreamThread& Scheduler::get_thread(Stream s) { { std::shared_lock lock(threads_mtx_); auto it = threads_.find(s.index); if (it != threads_.end()) { - st = it->second.get(); + return *it->second.get(); } } - if (!st) { - std::unique_lock lock(threads_mtx_); - auto it = threads_.find(s.index); - if (it == threads_.end()) { - it = threads_.emplace(s.index, std::make_unique()).first; - } - st = it->second.get(); + std::unique_lock lock(threads_mtx_); + auto it = threads_.find(s.index); + if (it == threads_.end()) { + it = threads_.emplace(s.index, std::make_unique()).first; } - st->enqueue(std::move(task)); + return *it->second.get(); } // Leak the scheduler singleton on all platforms. During static destruction, diff --git a/mlx/scheduler.h b/mlx/scheduler.h index c84ab62855..7c05b83689 100644 --- a/mlx/scheduler.h +++ b/mlx/scheduler.h @@ -3,66 +3,19 @@ #pragma once #include -#include #include #include -#include #include #include "mlx/api.h" #include "mlx/backend/gpu/eval.h" #include "mlx/device.h" #include "mlx/stream.h" +#include "mlx/utils.h" namespace mlx::core::scheduler { -struct StreamThread { - std::mutex mtx; - std::queue> q; - std::condition_variable cond; - bool stop; - std::thread thread; - - StreamThread() : stop(false), thread(&StreamThread::thread_fn, this) {} - - ~StreamThread() { - { - std::lock_guard lk(mtx); - stop = true; - } - cond.notify_one(); - thread.join(); - } - - void thread_fn() { - while (true) { - std::function task; - { - std::unique_lock lk(mtx); - cond.wait(lk, [this] { return !this->q.empty() || this->stop; }); - if (q.empty() && stop) { - return; - } - task = std::move(q.front()); - q.pop(); - } - - task(); - } - } - - void enqueue(std::function f) { - { - std::lock_guard lk(mtx); - if (stop) { - throw std::runtime_error( - "Cannot enqueue work after stream is stopped."); - } - q.emplace(std::move(f)); - } - cond.notify_one(); - } -}; +class StreamThread; class MLX_API Scheduler { public: @@ -76,6 +29,8 @@ class MLX_API Scheduler { Scheduler& operator=(Scheduler&&) = delete; void enqueue(Stream s, std::function task); + void wait_event(Stream s, Event event, std::function task); + void signal_event(Stream s, Event event, std::function task); void notify_new_task(const Stream& stream) { { @@ -110,6 +65,8 @@ class MLX_API Scheduler { private: friend Stream mlx::core::new_stream(Device d); + StreamThread& get_thread(Stream s); + int n_active_tasks_{0}; std::unordered_map> threads_; std::shared_mutex threads_mtx_; @@ -120,8 +77,19 @@ class MLX_API Scheduler { MLX_API Scheduler& scheduler(); template -void enqueue(const Stream& stream, F&& f) { - scheduler().enqueue(stream, std::forward(f)); +inline void enqueue(Stream s, F&& f) { + scheduler().enqueue(s, std::forward(f)); +} + +// Like enqueue but the task is used for processing the passed event. +template +inline void wait_event(Stream s, Event event, F&& f) { + scheduler().wait_event(s, std::move(event), std::forward(f)); +} + +template +inline void signal_event(Stream s, Event event, F&& f) { + scheduler().signal_event(s, std::move(event), std::forward(f)); } inline int n_active_tasks() { diff --git a/python/tests/test_load.py b/python/tests/test_load.py index 1c52f333a6..9dd9bff838 100644 --- a/python/tests/test_load.py +++ b/python/tests/test_load.py @@ -88,6 +88,41 @@ def test_load_npy_dtype(self): with self.assertRaises(Exception): out = mx.load(save_file, stream=mx.cpu) + def test_load_npy_read_error(self): + save_file = os.path.join(self.test_dir, "truncated.npy") + expected = np.arange(16, dtype=np.float32) + np.save(save_file, expected) + with open(save_file, "r+b") as f: + f.truncate(os.path.getsize(save_file) - expected.nbytes) + + out = mx.load(save_file, stream=mx.cpu) + with self.assertRaises(RuntimeError): + mx.eval(out) + + def test_async_load_npy_read_error_across_streams(self): + save_file = os.path.join(self.test_dir, "truncated_async.npy") + expected = np.arange(16, dtype=np.float32) + np.save(save_file, expected) + with open(save_file, "r+b") as f: + f.truncate(os.path.getsize(save_file) - expected.nbytes) + + producer_stream = mx.new_stream(mx.cpu) + consumer_stream = mx.new_stream(mx.cpu) + out = mx.add( + mx.load(save_file, stream=producer_stream), + 1.0, + stream=consumer_stream, + ) + with self.assertRaises(RuntimeError): + mx.eval(out) + # The error should propagate on both streams, but the Event impl of + # CUDA backend signals via gpu stream which adds a Fence wait which + # does a synchronous wait, so error surfaced early in producer_stream + # before poisoning the producer_stream. + if not mx.cuda.is_available(): + with self.assertRaises(RuntimeError): + mx.synchronize(producer_stream) + def test_save_and_load_safetensors(self): test_file = os.path.join(self.test_dir, "test.safetensors") with self.assertRaises(Exception): From d331598a77833af1ce661f16f5bf5ed9a371c2a1 Mon Sep 17 00:00:00 2001 From: Zhiqi Zhang Date: Tue, 18 Aug 2026 08:34:09 +0800 Subject: [PATCH 36/84] Fix mx.arange dtype inference overflow regression (#4324) --- python/src/ops.cpp | 10 +++++-- python/tests/test_ops.py | 61 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 611e281dbf..b8390b036d 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. +#include #include #include #include @@ -27,8 +28,11 @@ using namespace nb::literals; using Scalar = std::variant; mx::Dtype scalar_to_dtype(Scalar s) { - if (std::holds_alternative(s)) { - return mx::int32; + if (auto pv = std::get_if(&s); pv) { + return (*pv > std::numeric_limits::max() || + *pv < std::numeric_limits::min()) + ? mx::int64 + : mx::int32; } else if (std::holds_alternative(s)) { return mx::float32; } else { @@ -1513,7 +1517,7 @@ void init_ops(nb::module_& m) { start (float or int, optional): Starting value which defaults to ``0``. stop (float or int, optional): Stopping value. step (float or int, optional): Increment which defaults to ``1``. - dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``. + dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``, or ``int64`` if any of ``start``, ``stop``, or ``step`` does not fit in ``int32``. Returns: array: The range of values. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 86f0226856..054a1ca6b3 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1665,6 +1665,8 @@ def test_arange_overload_dispatch(self): with self.assertRaises(ValueError): INT_MAX = 2147483647 a = mx.arange(0, INT_MAX + 1, 1) + with self.assertRaises(ValueError): + a = mx.arange(0, 2**40) a = mx.arange(5) expected = [0, 1, 2, 3, 4] @@ -1728,6 +1730,57 @@ def test_arange_inferred_dtype(self): a = mx.arange(1.0, 3.0, 0.2, dtype=mx.int32) self.assertEqual(a.dtype, mx.int32) + # Integers that do not fit in int32 widen the inferred dtype to int64, + # matching the scalar inference of mx.array and numpy. + a = mx.arange(2**40, 2**40 + 3) + self.assertEqual(a.dtype, mx.int64) + self.assertListEqual(a.tolist(), [2**40, 2**40 + 1, 2**40 + 2]) + + a = mx.arange(-(2**40), -(2**40) + 3) + self.assertEqual(a.dtype, mx.int64) + self.assertListEqual(a.tolist(), [-(2**40), -(2**40) + 1, -(2**40) + 2]) + + # int32 boundaries themselves still infer int32. + a = mx.arange(2**31 - 3, 2**31 - 1) + self.assertEqual(a.dtype, mx.int32) + self.assertListEqual(a.tolist(), [2**31 - 3, 2**31 - 2]) + + a = mx.arange(-(2**31), -(2**31) + 2) + self.assertEqual(a.dtype, mx.int32) + self.assertListEqual(a.tolist(), [-(2**31), -(2**31) + 1]) + + # The first values that no longer fit widen as well. + a = mx.arange(-(2**31) - 1, -(2**31) + 1) + self.assertEqual(a.dtype, mx.int64) + self.assertListEqual(a.tolist(), [-(2**31) - 1, -(2**31)]) + + # A large step also widens the inferred dtype. + a = mx.arange(2**40, 2**40 + 3, 2**40) + self.assertEqual(a.dtype, mx.int64) + self.assertListEqual(a.tolist(), [2**40]) + + a = mx.arange(stop=2, step=2**40) + self.assertEqual(a.dtype, mx.int64) + self.assertListEqual(a.tolist(), [0]) + + # A negative step with widened values. + a = mx.arange(2**40 + 3, 2**40, -1) + self.assertEqual(a.dtype, mx.int64) + self.assertListEqual(a.tolist(), [2**40 + 3, 2**40 + 2, 2**40 + 1]) + + # The stop-only overload widens too, even for an empty result. + a = mx.arange(stop=2**40, step=-1) + self.assertEqual(a.dtype, mx.int64) + self.assertEqual(a.shape, (0,)) + + # An explicit dtype takes precedence over the widened inference. + a = mx.arange(2**40, 2**40 + 3, dtype=mx.int32) + self.assertEqual(a.dtype, mx.int32) + + # A float in the mix still infers float32. + a = mx.arange(0.5, 2**40, 2**39) + self.assertEqual(a.dtype, mx.float32) + def test_arange_corner_cases_cast(self): a = mx.arange(0, 3, 0.2, dtype=mx.int32) expected = [0] * 15 @@ -1782,9 +1835,17 @@ def test_arange_corner_cases_cast(self): expected = [0] self.assertListEqual(a.tolist(), expected) + # The range crossing the int32 limit widens the dtype to int64 instead + # of saturating or wrapping. n = mx.iinfo(mx.int32).max result = mx.arange(n - 1, n + 3) self.assertEqual(result.shape, (4,)) + self.assertEqual(result.dtype, mx.int64) + self.assertEqual(result.tolist(), [n - 1, n, n + 1, n + 2]) + + # An explicit dtype keeps the previous wrapping behaviour. + result = mx.arange(n - 1, n + 3, dtype=mx.int32) + self.assertEqual(result.shape, (4,)) self.assertEqual(result.dtype, mx.int32) self.assertEqual(result.tolist(), [n - 1, n, -2147483648, -2147483647]) From cee24b30c50ef4e1022795bef0f847c91aff502c Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 18 Aug 2026 09:34:26 +0900 Subject: [PATCH 37/84] Add workflow to update pull request limit bypass list (#4320) --- .github/scripts/find-eligible-contributors.js | 220 ++++++++++++++++++ .github/scripts/update-bypass-list.js | 161 +++++++++++++ .github/workflows/update_bypass_list.yml | 21 ++ .pre-commit-config.yaml | 1 + python/tests/test_load.py | 13 +- 5 files changed, 409 insertions(+), 7 deletions(-) create mode 100644 .github/scripts/find-eligible-contributors.js create mode 100644 .github/scripts/update-bypass-list.js create mode 100644 .github/workflows/update_bypass_list.yml diff --git a/.github/scripts/find-eligible-contributors.js b/.github/scripts/find-eligible-contributors.js new file mode 100644 index 0000000000..f5a05396ce --- /dev/null +++ b/.github/scripts/find-eligible-contributors.js @@ -0,0 +1,220 @@ +#!/usr/bin/env node + +/** + * find-eligible-contributors.js + * + * Finds users who have opened pull requests on a GitHub repo, + * are NOT collaborators on that repo, and match: + * - fewer than 5 currently open PRs + * - more than 2 merged PRs + * - merge rate (merged / closed) > 50% + * + * Usage: + * GITHUB_TOKEN=xxxx node find-eligible-contributors.js + */ + +const GITHUB_API = 'https://api.github.com'; + +function authHeaders(token) { + return { + Authorization: `Bearer ${token}`, + Accept: 'application/vnd.github+json', + 'X-GitHub-Api-Version': '2022-11-28', + 'User-Agent': 'find-eligible-contributors-script', + }; +} + +async function fetchAllPages(url, token, { stopWhen } = {}) { + const results = []; + let page = 1; + const perPage = 100; + const concurrentRequests = 5; + + while (true) { + const responses = await Promise.all( + Array.from({ length: concurrentRequests }, async (_, index) => { + const pageUrl = new URL(url); + pageUrl.searchParams.set('per_page', String(perPage)); + pageUrl.searchParams.set('page', String(page + index)); + + const res = await fetch(pageUrl, { headers: authHeaders(token) }); + if (!res.ok) { + const body = await res.text(); + throw new Error(`GitHub API error ${res.status} for ${pageUrl}: ${body}`); + } + return res.json(); + }) + ); + + let end = false; + for (const data of responses) { + if (!Array.isArray(data) || data.length === 0) { + end = true; + continue; + } + if (data.length < perPage) { + end = true; + } + + for (const item of data) { + if (stopWhen && stopWhen(item)) break; + results.push(item); + } + } + + if (end) { + break; + } + + page += concurrentRequests; + } + + return results; +} + +async function fetchPullRequests(owner, repo, token) { + const url = `${GITHUB_API}/repos/${owner}/${repo}/pulls?state=all&sort=created&direction=desc`; + + const oneYearAgo = new Date(); + oneYearAgo.setFullYear(oneYearAgo.getFullYear() - 1); + + return fetchAllPages(url, token, { + stopWhen: (pr) => new Date(pr.created_at) < oneYearAgo, + }); +} + +async function fetchCollaborators(owner, repo, token) { + const url = `${GITHUB_API}/repos/${owner}/${repo}/collaborators`; + const collaborators = await fetchAllPages(url, token); + return new Set(collaborators.map((c) => c.login)); +} + +/** + * Aggregate per-user PR stats from the list of PRs. + * Returns Map + */ +function aggregateUserStats(pullRequests) { + const stats = new Map(); + + for (const pr of pullRequests) { + const login = pr.user?.login; + if (!login) continue; + + if (!stats.has(login)) { + stats.set(login, { + open: 0, + closed: 0, + merged: 0, + mostRecentCreatedAt: null, + }); + } + const s = stats.get(login); + + if (pr.state === 'open') { + s.open += 1; + } else if (pr.state === 'closed') { + s.closed += 1; + if (pr.merged_at) { + s.merged += 1; + } + } + + const createdAt = new Date(pr.created_at); + if (!s.mostRecentCreatedAt || createdAt > s.mostRecentCreatedAt) { + s.mostRecentCreatedAt = createdAt; + } + } + + return stats; +} + +/** + * Apply eligibility filters: + * - fewer than 5 currently open PRs + * - more than 2 merged PRs + * - merge rate (merged / closed) > 50% + * - not a collaborator + * - has an open PR or at least one PR within the past 6 weeks + */ +function filterEligibleUsers(stats, collaboratorLogins) { + const SIX_WEEKS_MS = 6 * 7 * 24 * 60 * 60 * 1000; + const cutoff = new Date((new Date).getTime() - SIX_WEEKS_MS); + const eligible = []; + + for (const [login, s] of stats.entries()) { + if (s.open >= 5) continue; + if (s.merged <= 2) continue; + if (collaboratorLogins.has(login)) continue; + if (s.open == 0 && s.mostRecentCreatedAt < cutoff) continue; + if ((s.merged / s.closed) < 0.5) continue; + + eligible.push(login); + } + return eligible; +} + +/** + * Given a repo (owner/repo) and a GitHub token, returns the login names of + * eligible external contributors matching all filters. + * + * @param {Object} params + * @param {string} params.owner - Repo owner (user or org). + * @param {string} params.repo - Repo name. + * @param {string} params.token - GitHub token with repo read (and ideally push) access. + * @returns {Promise} + */ +async function findEligibleContributors({ owner, repo, token } = {}) { + if (!owner || !repo) { + throw new Error('findEligibleContributors requires both "owner" and "repo".'); + } + if (!token) { + throw new Error('findEligibleContributors requires a "token".'); + } + + const [pullRequests, collaboratorLogins] = await Promise.all([ + fetchPullRequests(owner, repo, token), + fetchCollaborators(owner, repo, token), + ]); + + const stats = aggregateUserStats(pullRequests); + return filterEligibleUsers(stats, collaboratorLogins); +} + +// ---- CLI entry point ---- + +function isRunAsCLI() { + return import.meta.url === `file://${process.argv[1]}`; +} + +async function runCLI() { + const [, , owner, repo] = process.argv; + if (!owner || !repo) { + console.error('Usage: node find-eligible-contributors.js '); + process.exit(1); + } + const token = process.env.GITHUB_TOKEN; + if (!token) { + console.error('Error: set GITHUB_TOKEN environment variable with a valid GitHub token.'); + process.exit(1); + } + + console.error(`Fetching pull requests and collaborators for ${owner}/${repo} in parallel...`); + const eligibleLogins = await findEligibleContributors({ owner, repo, token }); + + console.log(JSON.stringify(eligibleLogins, null, 2)); + console.error(`\n${eligibleLogins.length} user(s) match the criteria.`); +} + +if (isRunAsCLI()) { + runCLI().catch((err) => { + console.error('Fatal error:', err); + process.exit(1); + }); +} + +export { findEligibleContributors }; diff --git a/.github/scripts/update-bypass-list.js b/.github/scripts/update-bypass-list.js new file mode 100644 index 0000000000..0f0f0817a8 --- /dev/null +++ b/.github/scripts/update-bypass-list.js @@ -0,0 +1,161 @@ +#!/usr/bin/env node + +/** + * update-bypass-list.js + * + * Syncs a repo's PR interaction-limits bypass list with the current + * set of eligible external contributors. + * + * Usage: + * GITHUB_TOKEN=xxxx node update-bypass-list.js [--dry-run] + */ + +import { findEligibleContributors } from './find-eligible-contributors.js'; + +const GITHUB_API = 'https://api.github.com'; + +function authHeaders(token) { + return { + Authorization: `Bearer ${token}`, + Accept: 'application/vnd.github+json', + 'X-GitHub-Api-Version': '2026-03-10', + 'User-Agent': 'update-bypass-list-script', + 'Content-Type': 'application/json', + }; +} + +function bypassListUrl(owner, repo) { + return `${GITHUB_API}/repos/${owner}/${repo}/interaction-limits/pulls/bypass-list`; +} + +async function fetchBypassList(owner, repo, token) { + const url = bypassListUrl(owner, repo); + const res = await fetch(url, { headers: authHeaders(token) }); + + if (res.status === 404) { + return new Set(); + } + + if (!res.ok) { + const body = await res.text(); + throw new Error(`Failed to fetch bypass list (${res.status}): ${body}`); + } + + const data = await res.json(); + const logins = data.map((user) => user?.login).filter(Boolean); + + return new Set(logins); +} + +async function addUsersToBypassList(owner, repo, token, usernames) { + if (usernames.length === 0) return; + + const res = await fetch(bypassListUrl(owner, repo), { + method: 'PUT', + headers: authHeaders(token), + body: JSON.stringify({ users: usernames }), + }); + + if (!res.ok) { + const body = await res.text(); + throw new Error(`Failed to add users [${usernames.join(', ')}] (${res.status}): ${body}`); + } +} + +async function removeUsersFromBypassList(owner, repo, token, usernames) { + if (usernames.length === 0) return; + + const res = await fetch(bypassListUrl(owner, repo), { + method: 'DELETE', + headers: authHeaders(token), + body: JSON.stringify({ users: usernames }), + }); + + if (!res.ok) { + const body = await res.text(); + throw new Error(`Failed to remove users [${usernames.join(', ')}] (${res.status}): ${body}`); + } +} + +function diffLists(eligibleLogins, currentBypassLogins) { + const toAdd = [...eligibleLogins].filter((login) => !currentBypassLogins.has(login)); + const toRemove = [...currentBypassLogins].filter((login) => !eligibleLogins.has(login)); + return { toAdd, toRemove }; +} + +/** + * Update the bypass list to eligible contributors. + * + * @param {Object} params + * @param {string} params.owner + * @param {string} params.repo + * @param {string} params.token + * @param {boolean} [params.dryRun] - If true, only computes the diff, makes no API writes. + */ +async function updateBypassList({ owner, repo, token, dryRun = false }) { + console.error(`Getting eligible external contributors for ${owner}/${repo}...`); + const eligibleLogins = new Set(await findEligibleContributors({ owner, repo, token })); + console.error(`Found ${eligibleLogins.size} eligible user(s).`); + + console.error(`Fetching current bypass list...`); + const currentBypassLogins = await fetchBypassList(owner, repo, token); + console.error(`Current bypass list has ${currentBypassLogins.size} user(s).`); + + const { toAdd, toRemove } = diffLists(eligibleLogins, currentBypassLogins); + + console.error(`\nUsers to add (${toAdd.length}): ${toAdd.join(', ') || '(none)'}`); + console.error(`Users to remove (${toRemove.length}): ${toRemove.join(', ') || '(none)'}`); + + if (dryRun) { + console.error('\nDry run mode: no changes will be made.'); + return; + } + + try { + await addUsersToBypassList(owner, repo, token, toAdd); + if (toAdd.length > 0) console.error(`Added ${toAdd.length} user(s) to bypass list.`); + } catch (err) { + console.error(`Error adding users: ${err.message}`); + } + + try { + await removeUsersFromBypassList(owner, repo, token, toRemove); + if (toRemove.length > 0) console.error(`Removed ${toRemove.length} user(s) from bypass list.`); + } catch (err) { + console.error(`Error removing users: ${err.message}`); + } +} + +// ---- CLI entry point ---- + +function isRunAsCLI() { + return import.meta.url === `file://${process.argv[1]}`; +} + +async function runCLI() { + const args = process.argv.slice(2); + const dryRun = args.includes('--dry-run'); + const [owner, repo] = args.filter((a) => !a.startsWith('--')); + + if (!owner || !repo) { + console.error('Usage: node update-bypass-list.js [--dry-run]'); + process.exit(1); + } + + const token = process.env.GITHUB_TOKEN; + if (!token) { + console.error('Error: set GITHUB_TOKEN environment variable with a valid GitHub token.'); + process.exit(1); + } + + await updateBypassList({ owner, repo, token, dryRun }); +} + +if (isRunAsCLI()) { + runCLI().catch((err) => { + console.error('Fatal error:', err); + process.exit(1); + }); +} + +export { updateBypassList }; diff --git a/.github/workflows/update_bypass_list.yml b/.github/workflows/update_bypass_list.yml new file mode 100644 index 0000000000..72f3e2fd29 --- /dev/null +++ b/.github/workflows/update_bypass_list.yml @@ -0,0 +1,21 @@ +name: 'Update bypass list' +description: 'Update interaction-limits bypass list with eligible contributors' + +on: + workflow_dispatch: + schedule: + - cron: 0 * * * * + +permissions: + contents: write + +jobs: + update_bypass_list: + name: 'Update bypass list' + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v7 + - env: + # Can only use PAT with permissions: admin:org_hook, public_repo + GITHUB_TOKEN: ${{ secrets.ZCBENZ_TOKEN_UPDATE_BYPASS_LIST }} + run: node .github/scripts/update-bypass-list.js ml-explore mlx diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0345848e97..859d0f660e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,6 +9,7 @@ repos: rev: v21.1.8 hooks: - id: clang-format + files: \.(h|cpp)$ # Using this mirror lets us use mypyc-compiled black, which is about 2x faster - repo: https://github.com/psf/black-pre-commit-mirror rev: 26.1.0 diff --git a/python/tests/test_load.py b/python/tests/test_load.py index 9dd9bff838..f5947de471 100644 --- a/python/tests/test_load.py +++ b/python/tests/test_load.py @@ -115,13 +115,12 @@ def test_async_load_npy_read_error_across_streams(self): ) with self.assertRaises(RuntimeError): mx.eval(out) - # The error should propagate on both streams, but the Event impl of - # CUDA backend signals via gpu stream which adds a Fence wait which - # does a synchronous wait, so error surfaced early in producer_stream - # before poisoning the producer_stream. - if not mx.cuda.is_available(): - with self.assertRaises(RuntimeError): - mx.synchronize(producer_stream) + # Depending on backend the error might be caught early before poisoning + # the producer_stream, but still sync to clear the errors. + try: + mx.synchronize(producer_stream) + except Exception: + pass def test_save_and_load_safetensors(self): test_file = os.path.join(self.test_dir, "test.safetensors") From 3a6219917e4535575ce5bce2fc2ba27a483a709b Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Mon, 17 Aug 2026 18:44:13 -0700 Subject: [PATCH 38/84] Support head dimension 72 in Metal full attention (#4330) --- benchmarks/python/sdpa_bench.py | 11 +++++++- .../steel/attn/kernels/steel_attention.metal | 1 + .../metal/scaled_dot_product_attention.cpp | 8 +++--- python/tests/test_fast_sdpa.py | 27 +++++++++++++++++++ 4 files changed, 43 insertions(+), 4 deletions(-) diff --git a/benchmarks/python/sdpa_bench.py b/benchmarks/python/sdpa_bench.py index bd279f0ead..4130b05760 100644 --- a/benchmarks/python/sdpa_bench.py +++ b/benchmarks/python/sdpa_bench.py @@ -180,6 +180,15 @@ def get_gflop_count(B, M, N, K): ( 1, 2048, 32121, 64, 32, 8), ) + shapes_72 = ( + # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 1024, 1024, 72, 32, 8), + ( 1, 2048, 2048, 72, 32, 8), + ( 1, 4096, 4096, 72, 32, 8), + ( 1, 4096, 5000, 72, 32, 8), + ( 1, 2048, 32121, 72, 32, 8), + ) + shapes_80 = ( # ( B, qsl, ksl, head_dim, n_qh, n_kvh) ( 1, 1024, 1024, 80, 32, 8), @@ -208,7 +217,7 @@ def get_gflop_count(B, M, N, K): ) # fmt: on - shapes = shapes_64 + shapes_80 + shapes_96 + shapes_128 + shapes = shapes_64 + shapes_72 + shapes_80 + shapes_96 + shapes_128 masks = [None, "bool", "causal"] diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index 7bddfcb054..4bb9ff5873 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal @@ -15,6 +15,7 @@ instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) #define instantiate_attn_mask_helper(iname, itype) \ diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index cb4af8523e..bb8ad808b8 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -174,7 +174,9 @@ void sdpa_full_self_attention_metal( bool do_causal_, const std::optional& mask, const std::optional& sinks) { - if (metal::is_nax_available() && q.shape(3) != 80 && + // NAX tiles the head dim in units of kU=16 and steps TD by 2, so it needs + // a multiple of 32; 72 and 80 take the classic steel kernel. + if (metal::is_nax_available() && q.shape(3) != 80 && q.shape(3) != 72 && (env::enable_tf32() || q.dtype() != float32)) { return sdpa_full_self_attention_nax( /* const Stream& s = */ s, @@ -627,8 +629,8 @@ bool ScaledDotProductAttention::use_fallback( query_head_dim == 256)) || (query_head_dim == 192 && value_head_dim == 128); const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 96 || - query_head_dim == 128); + (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 || + query_head_dim == 96 || query_head_dim == 128); const bool sdpa_full_supported_mask = !has_mask || has_arr_mask || (query_sequence_length <= key_sequence_length && do_causal); diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 2418997bc8..a5fdb0fbb7 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -117,6 +117,33 @@ def mlx_primitives_sdpa(q, k, v, scale, mask=None): class TestFastSDPA(mlx_tests.MLXTestCase): + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_head_dim_72(self): + B, D, qH, kH = (1, 72, 8, 2) + for qL, kL, dtype, mask_str in product( + (64, 65), + (128, 127), + (mx.float16, mx.bfloat16, mx.float32), + (None, "additive", "bool", "causal"), + ): + with self.subTest(qL=qL, kL=kL, dtype=dtype, mask=mask_str): + q, k, v, scale, mask = prepare_inputs( + B, qL, kL, D, qH, kH, mask_str, False, dtype + ) + ref = mlx_ref_attn(q, k, v, scale, mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + + if dtype == mx.float32: + atol = 1e-5 + elif dtype == mx.bfloat16: + atol = 5e-3 + else: + atol = 3e-4 + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") def test_sdpa_head_dim_96(self): B, D, qH, kH = (1, 96, 8, 2) From 29006584fee8deb0407510fefd3a7e6301124d3c Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 18 Aug 2026 18:47:21 +0900 Subject: [PATCH 39/84] Patch bump to 0.32.2 (#4333) --- .github/workflows/update_bypass_list.yml | 6 ++++-- mlx/version.h | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/workflows/update_bypass_list.yml b/.github/workflows/update_bypass_list.yml index 72f3e2fd29..a129c51ce2 100644 --- a/.github/workflows/update_bypass_list.yml +++ b/.github/workflows/update_bypass_list.yml @@ -3,8 +3,10 @@ description: 'Update interaction-limits bypass list with eligible contributors' on: workflow_dispatch: - schedule: - - cron: 0 * * * * + pull_request_target: + types: + - opened + - closed permissions: contents: write diff --git a/mlx/version.h b/mlx/version.h index dc8b3a8630..9de9f2e558 100644 --- a/mlx/version.h +++ b/mlx/version.h @@ -6,7 +6,7 @@ #define MLX_VERSION_MAJOR 0 #define MLX_VERSION_MINOR 32 -#define MLX_VERSION_PATCH 1 +#define MLX_VERSION_PATCH 2 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) From f889b463ee82c088c65a546b23cc0b57e719b1b4 Mon Sep 17 00:00:00 2001 From: Tanish Jain Date: Tue, 18 Aug 2026 15:37:46 +0530 Subject: [PATCH 40/84] Preserve subnormal float values when casting to bool (#4224) --- mlx/backend/metal/kernels/copy.h | 46 +++++++++---------- .../metal/kernels/reduction/reduce_all.h | 4 +- .../metal/kernels/reduction/reduce_col.h | 14 +++--- .../metal/kernels/reduction/reduce_row.h | 6 +-- mlx/backend/metal/kernels/utils.h | 24 ++++++++++ python/tests/test_ops.py | 15 ++++++ 6 files changed, 74 insertions(+), 35 deletions(-) diff --git a/mlx/backend/metal/kernels/copy.h b/mlx/backend/metal/kernels/copy.h index cf22347ee5..95ed69b760 100644 --- a/mlx/backend/metal/kernels/copy.h +++ b/mlx/backend/metal/kernels/copy.h @@ -9,11 +9,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } } @@ -27,11 +27,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } } @@ -46,11 +46,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } } @@ -65,11 +65,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } } @@ -81,7 +81,7 @@ template constant const int64_t& src_stride [[buffer(3)]], uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); - dst[index] = static_cast(src[src_idx]); + dst[index] = cast_to(src[src_idx]); } template @@ -93,7 +93,7 @@ template uint2 grid_dim [[threads_per_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * index.y; - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -106,7 +106,7 @@ template auto src_idx = elem_to_loc_3(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -123,14 +123,14 @@ template if (N == 1) { IdxT dst_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); return; } auto xshape = src_shape[ndim - 1]; IdxT dst_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z); auto src_xstride = src_strides[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[dst_idx + i] = static_cast(src[src_idx]); + dst[dst_idx + i] = cast_to(src[src_idx]); src_idx += src_xstride; } } @@ -144,7 +144,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -156,7 +156,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -168,7 +168,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -187,14 +187,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } @@ -211,7 +211,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -225,7 +225,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -239,7 +239,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -262,14 +262,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } diff --git a/mlx/backend/metal/kernels/reduction/reduce_all.h b/mlx/backend/metal/kernels/reduction/reduce_all.h index e0d08392c0..47ad63fbbd 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_all.h +++ b/mlx/backend/metal/kernels/reduction/reduce_all.h @@ -37,13 +37,13 @@ template < for (IdxT b = 0; b < blocks; b++) { for (int i = 0; i < N_READS; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } in += lsize.x * N_READS; } if (extra > 0) { for (int i = 0; i < extra; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } } diff --git a/mlx/backend/metal/kernels/reduction/reduce_col.h b/mlx/backend/metal/kernels/reduction/reduce_col.h index c109faf0bc..b1546adb55 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_col.h +++ b/mlx/backend/metal/kernels/reduction/reduce_col.h @@ -43,13 +43,13 @@ template row = in + loop.location(); if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -125,7 +125,7 @@ template for (IdxT r = gid.z * lsize.y + lid.y; r < total_rows; r += lsize.y * gsize.z) { row = in + loop.location(); - total = op(static_cast(*row), total); + total = op(cast_to(*row), total); loop.next(lsize.y * gsize.z, reduce_shape, reduce_strides); } @@ -207,13 +207,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -352,13 +352,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); diff --git a/mlx/backend/metal/kernels/reduction/reduce_row.h b/mlx/backend/metal/kernels/reduction/reduce_row.h index b55c83f315..09d1d88348 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_row.h +++ b/mlx/backend/metal/kernels/reduction/reduce_row.h @@ -34,7 +34,7 @@ METAL_FUNC void per_thread_row_reduce( for (int i = 0; i < blocks; i++) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } inputs[j] += lsize_x * N_READS; @@ -46,13 +46,13 @@ METAL_FUNC void per_thread_row_reduce( if (index + N_READS <= extra) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } else { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; index + i < extra; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } diff --git a/mlx/backend/metal/kernels/utils.h b/mlx/backend/metal/kernels/utils.h index 266f27e91c..f15928282f 100644 --- a/mlx/backend/metal/kernels/utils.h +++ b/mlx/backend/metal/kernels/utils.h @@ -446,3 +446,27 @@ template struct ConditionalType { using type = T; }; + +/////////////////////////////////////////////////////////////////////////////// +// Type casting utils +/////////////////////////////////////////////////////////////////////////////// + +template +inline U cast_to(T val) { + return static_cast(val); +} + +template <> +inline bool cast_to(float val) { + return (as_type(val) & 0x7FFFFFFF) != 0; +} + +template <> +inline bool cast_to(bfloat16_t val) { + return (as_type(val) & 0x7FFF) != 0; +} + +template <> +inline bool cast_to(complex64_t val) { + return cast_to(val.real) || cast_to(val.imag); +} diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 054a1ca6b3..40d3f5644a 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1355,6 +1355,21 @@ def test_any(self): self.assertEqual(mx.any(a, axis=0).tolist(), [True, False]) self.assertEqual(mx.any(a, axis=1).tolist(), [True, False]) + def test_subnormal_bool_cast(self): + f32_sub = mx.array(np.array([0x00000001], dtype=np.uint32)).view(mx.float32) + f16_sub = mx.array(np.array([0x0001], dtype=np.uint16)).view(mx.float16) + bf16_sub = mx.array(np.array([0x0001], dtype=np.uint16)).view(mx.bfloat16) + + self.assertTrue(f32_sub.astype(mx.bool_).item()) + self.assertTrue(f16_sub.astype(mx.bool_).item()) + self.assertTrue(bf16_sub.astype(mx.bool_).item()) + self.assertTrue(mx.any(f32_sub).item()) + self.assertTrue(mx.any(f16_sub).item()) + self.assertTrue(mx.any(bf16_sub).item()) + self.assertTrue(mx.all(f32_sub).item()) + self.assertTrue(mx.all(f16_sub).item()) + self.assertTrue(mx.all(bf16_sub).item()) + def test_stop_gradient(self): def func(x): return mx.sum(2 * x + mx.stop_gradient(3 * x)) From 6172852cdd2bdf1df01f343f4628a6f30ef834d1 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:09:20 -0700 Subject: [PATCH 41/84] python: Support assigning through a bare Ellipsis index (#4314) --- python/src/indexing.cpp | 2 ++ python/tests/test_array.py | 22 ++++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp index 3df4c96882..1dce537880 100644 --- a/python/src/indexing.cpp +++ b/python/src/indexing.cpp @@ -778,6 +778,8 @@ mlx_compute_scatter_args( return mlx_scatter_args_int(src, obj, vals); } else if (nb::isinstance(obj)) { return mlx_scatter_args_nd(src, nb::cast(obj), vals); + } else if (nb::isinstance(obj)) { + return {{}, broadcast_to(vals, src.shape()), {}}; } else if (obj.is_none()) { return {{}, broadcast_to(vals, src.shape()), {}}; } else if (nb::isinstance(obj)) { diff --git a/python/tests/test_array.py b/python/tests/test_array.py index a920b39b38..2c55a3ba20 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -1279,6 +1279,28 @@ def test_setitem(self): a[0:2] = 3 self.assertEqual(a.tolist(), [3, 3, 1]) + # Assigning through a bare Ellipsis, like a[:] and a[None] + e = mx.zeros((2, 3), mx.int32) + e[...] = 5 + self.assertEqual(e.tolist(), [[5, 5, 5], [5, 5, 5]]) + + # Broadcasting an array update through Ellipsis + e[...] = mx.array([1, 2, 3]) + self.assertEqual(e.tolist(), [[1, 2, 3], [1, 2, 3]]) + + e[...] = mx.zeros((2, 3), mx.int32) + self.assertEqual(e.tolist(), [[0, 0, 0], [0, 0, 0]]) + + # Scalar array + e = mx.array(0) + e[...] = 7 + self.assertEqual(e.item(), 7) + + # Shapes that cannot broadcast are still rejected + e = mx.zeros((2, 3), mx.int32) + with self.assertRaises(ValueError): + e[...] = mx.array([1, 2]) + a[0:3] = 4 self.assertEqual(a.tolist(), [4, 4, 4]) From d5841be95f68eba13bce5ab6abd673260bf12f74 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 18 Aug 2026 03:14:03 -0700 Subject: [PATCH 42/84] Fix divmod truncating the quotient for floats (#4108) Co-authored-by: Cheng --- mlx/backend/cpu/binary.cpp | 16 ++++++++++++++-- mlx/backend/cuda/device/binary_ops.cuh | 11 ++++++++++- mlx/backend/metal/kernels/binary_ops.h | 25 ++++++++++++++++--------- python/tests/test_ops.py | 16 ++++++++++++++++ 4 files changed, 56 insertions(+), 12 deletions(-) diff --git a/mlx/backend/cpu/binary.cpp b/mlx/backend/cpu/binary.cpp index 9cca16d869..90b0378f6a 100644 --- a/mlx/backend/cpu/binary.cpp +++ b/mlx/backend/cpu/binary.cpp @@ -47,10 +47,22 @@ void DivMod::eval_cpu( out_b = array::unsafe_weak_copy(out_b), bopt]() mutable { auto integral_op = [](auto x, auto y) { - return std::make_pair(x / y, x % y); + auto q = x / y; + auto r = x % y; + if constexpr (std::is_signed_v) { + if (r != 0 && (r < 0) != (y < 0)) { + q -= 1; + r += y; + } + } + return std::make_pair(q, r); }; auto float_op = [](auto x, auto y) { - return std::make_pair(std::trunc(x / y), std::fmod(x, y)); + auto r = std::fmod(x, y); + if (r != 0 && (r < 0) != (y < 0)) { + r += y; + } + return std::make_pair(std::floor(x / y), r); }; dispatch_all_types(out_a.dtype(), [&](auto type_tag) { diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh index b0b7962807..4368864465 100644 --- a/mlx/backend/cuda/device/binary_ops.cuh +++ b/mlx/backend/cuda/device/binary_ops.cuh @@ -17,9 +17,18 @@ struct FloorDivide { template __device__ T operator()(T x, T y) { if constexpr (cuda::std::is_integral_v) { + auto q = x / y; + if constexpr (cuda::std::is_signed_v) { + if (x % y != 0 && (x < 0) != (y < 0)) { + q -= 1; + } + } + return q; + } else if constexpr (is_complex_v) { + // Complex is not supported, simply make compiler happy. return x / y; } else { - return cuda::std::trunc(x / y); + return cuda::std::floor(x / y); } } }; diff --git a/mlx/backend/metal/kernels/binary_ops.h b/mlx/backend/metal/kernels/binary_ops.h index 863d6369e2..37650c7d93 100644 --- a/mlx/backend/metal/kernels/binary_ops.h +++ b/mlx/backend/metal/kernels/binary_ops.h @@ -16,20 +16,27 @@ struct Add { struct FloorDivide { template - T operator()(T x, T y) thread { + metal::enable_if_t & !metal::is_signed_v, T> + operator()(T x, T y) thread { return x / y; } - template <> - float operator()(float x, float y) thread { - return trunc(x / y); + template + metal::enable_if_t & metal::is_signed_v, T> + operator()(T x, T y) thread { + auto q = x / y; + if (x % y != 0 && (x < 0) != (y < 0)) { + q -= 1; + } + return q; } - template <> - half operator()(half x, half y) thread { - return trunc(x / y); + template + metal::enable_if_t, T> operator()(T x, T y) thread { + return floor(x / y); } template <> - bfloat16_t operator()(bfloat16_t x, bfloat16_t y) thread { - return trunc(x / y); + complex64_t operator()(complex64_t x, complex64_t y) thread { + // Complex is not supported, simply make compiler happy. + return x / y; } }; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 40d3f5644a..e5761c78b2 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3322,6 +3322,22 @@ def test_divmod(self): np.allclose(np_out[0], mx_out[0]), msg=f"Shapes {s1} {s2}, Type {t}" ) + # Mixed signs floor, matching python's divmod and numpy, so + # q * b + r == a holds + av = [-7, 7, -7, 7, -1, 1, -5, 5, 6, -6] + bv = [2, 2, -2, -2, 3, -3, 3, -3, 3, 3] + a, b = mx.array(av), mx.array(bv) + q, r = mx.divmod(a, b) + self.assertEqual(q.tolist(), [x // y for x, y in zip(av, bv)]) + self.assertEqual(r.tolist(), [x % y for x, y in zip(av, bv)]) + self.assertTrue(mx.array_equal(q * b + r, a)) + + af = mx.array([-7.0, 7.0, -7.5, 7.5]) + bf = mx.array([2.0, -2.0, 2.0, -2.0]) + q, r = mx.divmod(af, bf) + self.assertTrue(mx.array_equal(q, mx.array([-4.0, -4.0, -4.0, -4.0]))) + self.assertTrue(mx.array_equal(q * bf + r, af)) + def test_tile(self): self.assertCmpNumpy([(2,), [2]], mx.tile, np.tile) self.assertCmpNumpy([(2, 3, 4), [2]], mx.tile, np.tile) From 772bc8a00100d12a91eca7365ff676fc4ac724cd Mon Sep 17 00:00:00 2001 From: hojin12312 Date: Tue, 18 Aug 2026 21:14:10 +0900 Subject: [PATCH 43/84] Add force_fused option to scaled_dot_product_attention (#4185) --- .../cuda/scaled_dot_product_attention.cpp | 45 +++++- .../scaled_dot_product_attention.metal | 2 + .../steel/attn/kernels/steel_attention.h | 4 +- .../steel/attn/kernels/steel_attention.metal | 2 + .../metal/scaled_dot_product_attention.cpp | 136 +++++++++++++----- mlx/backend/no_gpu/primitives.cpp | 14 +- mlx/fast.cpp | 15 +- mlx/fast.h | 1 + mlx/fast_primitives.h | 15 +- python/src/fast.cpp | 31 +++- python/tests/test_fast_sdpa.py | 94 ++++++++++++ 11 files changed, 306 insertions(+), 53 deletions(-) diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index ca411e91c6..286d500e2c 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -549,6 +549,33 @@ void sdpa_vector( namespace fast { +namespace { + +std::tuple has_fused_kernel( + const array& q, + const array& k, + const array& v, + bool has_arr_mask, + bool do_causal, + bool output_logsumexp, + Stream s) { + if (s.device != Device::gpu) { + return {false, "the fused kernels require a GPU stream."}; + } + if (!supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) && + !supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp)) { + std::ostringstream msg; + msg << "neither the cuDNN attention nor the vector attention kernel " + << "supports this configuration; got query shape " << q.shape() + << ", key shape " << k.shape() << ", value shape " << v.shape() + << " with dtype " << q.dtype() << "."; + return {false, msg.str()}; + } + return {true, ""}; +} + +} // namespace + bool ScaledDotProductAttention::use_fallback( const array& q, const array& k, @@ -558,13 +585,21 @@ bool ScaledDotProductAttention::use_fallback( bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s) { - if (s.device == Device::cpu) { - return true; + auto [has_fused, reason] = + has_fused_kernel(q, k, v, has_arr_mask, do_causal, output_logsumexp, s); + if (force_fused) { + if (!has_fused) { + std::ostringstream msg; + msg << "[scaled_dot_product_attention] force_fused=True but no fused " + "kernel is available: " + << reason; + throw std::invalid_argument(msg.str()); + } + return false; } - - return !supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) && - !supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp); + return !has_fused; } bool ScaledDotProductAttention::supports_bool_mask() { diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 44ccb834e9..84486d62c3 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -33,10 +33,12 @@ using namespace metal; instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ instantiate_sdpa_vector(type, 192, 128) \ + instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ + instantiate_sdpa_vector_aggregation(type, 192) \ instantiate_sdpa_vector_aggregation(type, 256) instantiate_sdpa_vector_heads(float) diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h index 0d9628e834..29fa7ba396 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h @@ -428,7 +428,7 @@ template < for (short id = 0; id < TD; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } @@ -438,7 +438,7 @@ template < Vtile.template load( &Vs[Vs_offset + kk * LDV_tgp + dd]); - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index 4bb9ff5873..fbd84004f0 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal @@ -12,6 +12,8 @@ attention, dtype, bq, bk, bd, wm, wn, mtype, float) #define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ + instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index bb8ad808b8..acff685790 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -160,6 +160,7 @@ void sdpa_full_self_attention_nax( MTL::Size grid_dims = MTL::Size(NQ, H, B); MTL::Size group_dims = MTL::Size(32, wm, wn); + check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -174,9 +175,8 @@ void sdpa_full_self_attention_metal( bool do_causal_, const std::optional& mask, const std::optional& sinks) { - // NAX tiles the head dim in units of kU=16 and steps TD by 2, so it needs - // a multiple of 32; 72 and 80 take the classic steel kernel. - if (metal::is_nax_available() && q.shape(3) != 80 && q.shape(3) != 72 && + if (metal::is_nax_available() && + (q.shape(3) == 64 || q.shape(3) == 96 || q.shape(3) == 128) && (env::enable_tf32() || q.dtype() != float32)) { return sdpa_full_self_attention_nax( /* const Stream& s = */ s, @@ -325,6 +325,7 @@ void sdpa_full_self_attention_metal( MTL::Size grid_dims = MTL::Size(NQ, H, B); MTL::Size group_dims = MTL::Size(32, wm, wn); + check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -591,28 +592,23 @@ void sdpa_vector_2pass( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } -} // namespace - -bool ScaledDotProductAttention::use_fallback( +std::tuple has_fused_kernel( const array& q, const array& k, const array& v, bool has_mask, bool has_arr_mask, bool do_causal, - bool is_training, bool output_logsumexp, Stream s) { - if (is_training) { - // It's faster for training on Metal to use the unfused SDPA for both - // forward and backward. - return true; + if (s.device != Device::gpu) { + return {false, "the fused kernels require a GPU (Metal) stream."}; } if (output_logsumexp) { - return true; - } - if (s.device == Device::cpu) { - return true; + return { + false, + "the fused forward does not produce the logsumexp required for " + "the fused VJP; use default routing when training."}; } const int value_head_dim = v.shape(-1); @@ -623,27 +619,103 @@ bool ScaledDotProductAttention::use_fallback( const int num_kv_heads = k.shape(1); const int gqa_factor = num_query_heads / num_kv_heads; - const bool sdpa_vector_supported_head_dim = - (query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || - query_head_dim == 256)) || - (query_head_dim == 192 && value_head_dim == 128); - const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 || - query_head_dim == 96 || query_head_dim == 128); + std::ostringstream msg; + if (query_sequence_length > 8) { + const bool supported_head_dim = query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 || + query_head_dim == 96 || query_head_dim == 128 || + query_head_dim == 192 || query_head_dim == 256); + if (!supported_head_dim) { + msg << "the full attention kernel supports head dims " + << "{64, 72, 80, 96, 128, 192, 256} with matching query/value head " + << "dims; got query head dim " << query_head_dim + << " and value head dim " << value_head_dim << "."; + return {false, msg.str()}; + } + if (has_mask && !has_arr_mask && + !(query_sequence_length <= key_sequence_length && do_causal)) { + msg << "the full attention kernel with a causal mask requires the " + << "query sequence to be no longer than the key sequence; got " + << "query length " << query_sequence_length << " and key length " + << key_sequence_length << "."; + return {false, msg.str()}; + } + } else { + const bool supported_head_dim = + (query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 96 || + query_head_dim == 128 || query_head_dim == 192 || + query_head_dim == 256)) || + (query_head_dim == 192 && value_head_dim == 128); + if (!supported_head_dim) { + msg << "the vector attention kernel supports head dims " + << "{64, 96, 128, 192, 256} with matching query/value head dims, " + << "or query head dim 192 with value head dim 128; got query head " + << "dim " << query_head_dim << " and value head dim " + << value_head_dim << "."; + return {false, msg.str()}; + } + if (query_sequence_length > key_sequence_length) { + msg << "the vector attention kernel requires the query sequence to be " + << "no longer than the key sequence; got query length " + << query_sequence_length << " and key length " << key_sequence_length + << "."; + return {false, msg.str()}; + } + if (query_sequence_length * gqa_factor > 32) { + msg << "the vector attention kernel requires the query length times " + << "the GQA factor to be at most 32; got query length " + << query_sequence_length << " and GQA factor " << gqa_factor << "."; + return {false, msg.str()}; + } + } + return {true, ""}; +} - const bool sdpa_full_supported_mask = !has_mask || has_arr_mask || - (query_sequence_length <= key_sequence_length && do_causal); +} // namespace - const bool supports_sdpa_full = query_sequence_length > 8 && - sdpa_full_supported_mask && sdpa_full_supported_head_dim; +bool ScaledDotProductAttention::use_fallback( + const array& q, + const array& k, + const array& v, + bool has_mask, + bool has_arr_mask, + bool do_causal, + bool is_training, + bool output_logsumexp, + bool force_fused, + Stream s) { + auto [has_fused, reason] = has_fused_kernel( + q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp, s); + if (force_fused) { + if (!has_fused) { + std::ostringstream msg; + msg << "[scaled_dot_product_attention] force_fused=True but no fused " + "kernel is available: " + << reason; + throw std::invalid_argument(msg.str()); + } + return false; + } - const bool supports_sdpa_vector = (query_sequence_length <= 8) && - (query_sequence_length <= key_sequence_length) && - sdpa_vector_supported_head_dim && - (query_sequence_length * gqa_factor) <= 32; + if (is_training) { + // It's faster for training on Metal to use the unfused SDPA for both + // forward and backward. + return true; + } + if (!has_fused) { + return true; + } - return !(supports_sdpa_full || supports_sdpa_vector); + // Unfused path is faster for following shapes. + const int query_sequence_length = q.shape(2); + const int query_head_dim = q.shape(-1); + const int value_head_dim = v.shape(-1); + if (query_sequence_length > 8) { + return query_head_dim == 192 || query_head_dim == 256; + } else { + return query_head_dim == value_head_dim && query_head_dim == 192; + } } bool ScaledDotProductAttention::supports_bool_mask() { diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index b7d7a19467..7f60d0d83a 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -32,20 +32,26 @@ bool fast::ScaledDotProductAttention::use_fallback( bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s) { + if (force_fused) { + throw std::invalid_argument( + "[scaled_dot_product_attention] force_fused=True but no fused " + "kernel is available in CPU backend."); + } return true; } -bool fast::ScaledDotProductAttention::supports_bool_mask() { - return false; -} - bool fast::ScaledDotProductAttentionVJP::use_fallback( const array& q, Stream s) { return true; } +bool fast::ScaledDotProductAttention::supports_bool_mask() { + return false; +} + NO_GPU(Abs) NO_GPU(Add) NO_GPU(AddMM) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a668fe9abd..df2beebd84 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -618,7 +618,8 @@ array scaled_dot_product_attention( const std::string& mask_mode /* = "" */, std::optional mask_arr /* = {} */, const std::optional& sinks /* = {} */, - StreamOrDevice s /* = {}*/) { + bool force_fused /* = false */, + StreamOrDevice s /* = {} */) { for (const auto& tensor : {queries, keys, values}) { if (tensor.ndim() != 4) { std::ostringstream msg; @@ -834,6 +835,7 @@ array scaled_dot_product_attention( do_causal, is_training, output_logsumexp, + force_fused, stream)) { if (has_bool_mask && !ScaledDotProductAttention::supports_bool_mask()) { // Convert bool mask to additive mask. @@ -846,7 +848,13 @@ array scaled_dot_product_attention( } Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)}; auto primitive = std::make_shared( - stream, fallback, scale, do_causal, has_sinks, output_logsumexp); + stream, + fallback, + scale, + do_causal, + has_sinks, + output_logsumexp, + force_fused); if (output_logsumexp) { return array::make_arrays( {std::move(out_shape), Shape{q.shape(0), q.shape(1), q.shape(2), 1}}, @@ -912,7 +920,8 @@ bool ScaledDotProductAttention::is_equivalent(const Primitive& other) const { static_cast(other); return scale_ == a_other.scale_ && do_causal_ == a_other.do_causal_ && has_sinks_ == a_other.has_sinks_ && - output_logsumexp_ == a_other.output_logsumexp_; + output_logsumexp_ == a_other.output_logsumexp_ && + force_fused_ == a_other.force_fused_; } bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { diff --git a/mlx/fast.h b/mlx/fast.h index 934fadc2b7..c5f664df79 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -53,6 +53,7 @@ MLX_API array scaled_dot_product_attention( const std::string& mask_mode = "", std::optional mask_arr = {}, const std::optional& sinks = {}, + bool force_fused = false, StreamOrDevice s = {}); using TemplateArg = std::variant; diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 0d2f861045..61a392e418 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -212,12 +212,14 @@ class ScaledDotProductAttention : public Custom { float scale, bool do_causal, bool has_sinks, - bool output_logsumexp) + bool output_logsumexp, + bool force_fused) : Custom(stream, std::move(fallback)), scale_(scale), do_causal_(do_causal), has_sinks_(has_sinks), - output_logsumexp_(output_logsumexp) {} + output_logsumexp_(output_logsumexp), + force_fused_(force_fused) {} static bool use_fallback( const array& q, @@ -228,6 +230,7 @@ class ScaledDotProductAttention : public Custom { bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s); static bool supports_bool_mask(); @@ -251,7 +254,12 @@ class ScaledDotProductAttention : public Custom { DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { return std::make_tuple( - nullptr, scale_, do_causal_, has_sinks_, output_logsumexp_); + nullptr, + scale_, + do_causal_, + has_sinks_, + output_logsumexp_, + force_fused_); } private: @@ -259,6 +267,7 @@ class ScaledDotProductAttention : public Custom { bool do_causal_; bool has_sinks_; bool output_logsumexp_; + bool force_fused_; }; class ScaledDotProductAttentionVJP : public Custom { diff --git a/python/src/fast.cpp b/python/src/fast.cpp index e59357bc33..0a50dc79cd 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -234,6 +234,7 @@ void init_fast(nb::module_& parent_module) { const float scale, const std::variant& mask, const std::optional& sinks, + bool force_fused, mx::StreamOrDevice s) { bool has_mask = !std::holds_alternative(mask); bool has_str_mask = @@ -250,16 +251,32 @@ void init_fast(nb::module_& parent_module) { throw std::invalid_argument(msg.str()); } return mx::fast::scaled_dot_product_attention( - queries, keys, values, scale, mask_str, std::nullopt, sinks, s); + queries, + keys, + values, + scale, + mask_str, + std::nullopt, + sinks, + force_fused, + s); } else { auto mask_arr = std::get(mask); return mx::fast::scaled_dot_product_attention( - queries, keys, values, scale, "", mask_arr, sinks, s); + queries, + keys, + values, + scale, + "", + mask_arr, + sinks, + force_fused, + s); } } else { return mx::fast::scaled_dot_product_attention( - queries, keys, values, scale, "", {}, sinks, s); + queries, keys, values, scale, "", {}, sinks, force_fused, s); } }, "q"_a, @@ -269,9 +286,10 @@ void init_fast(nb::module_& parent_module) { "scale"_a, "mask"_a = nb::none(), "sinks"_a = nb::none(), + "force_fused"_a = false, "stream"_a = nb::none(), nb::sig( - "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: None | str | array = None, sinks: array | None = None, stream: StreamOrDevice = None) -> array"), + "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: None | str | array = None, sinks: array | None = None, force_fused: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``. @@ -313,6 +331,11 @@ void init_fast(nb::module_& parent_module) { last query aligns with the last key. sinks (array, optional): An optional array of attention sinks. Default: ``None``. + force_fused (bool, optional): If ``True``, use a fused kernel + regardless of the builtin heuristics and raise error when no + fused kernel is available. For certain configurations this would + result in slower kernel getting used but can reduce memory + consumption. Default: ``False``. Returns: array: The output array. diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index a5fdb0fbb7..997fa1028c 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -722,6 +722,100 @@ def test_grad(slow, fast, args): ).sum() test_grad(loss_slow, loss_fast, [q, k, v]) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_force_fused_metal(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + + def make_qkv(qL, kL, D, qH=8, kH=8): + q = mx.random.normal((1, qH, qL, D), mx.float16) + k = mx.random.normal((1, kH, kL, D), mx.float16) + v = mx.random.normal((1, kH, kL, D), mx.float16) + return q, k, v + + # Full attention kernel. + for D, qL, mask in product((192, 256), (9, 16), (None, "causal")): + with self.subTest(head_dim=D, qL=qL, mask=mask): + q, k, v = make_qkv(qL, 512, D, 8, 4) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + + # Vector attention kernel. + for D in (192, 256): + with self.subTest(head_dim=D): + q, k, v = make_qkv(4, 16385, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + + # No full attention fused kernels. + with self.assertRaisesRegex(ValueError, "supports head dims"): + q, k, v = make_qkv(16, 512, 512) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) + with self.assertRaisesRegex( + ValueError, "query sequence to be no longer than the key sequence" + ): + q, k, v = make_qkv(32, 16, 64) + mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=64**-0.5, + mask="causal", + force_fused=True, + ) + + # No vector attention fused kernels. + with self.assertRaisesRegex(ValueError, "supports head dims"): + q, k, v = make_qkv(1, 128, 72) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=72**-0.5, force_fused=True + ) + with self.assertRaisesRegex(ValueError, "GQA factor to be at most 32"): + q, k, v = make_qkv(8, 128, 64, qH=8, kH=1) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=64**-0.5, force_fused=True + ) + + # No CPU fused kernel. + with mx.stream(mx.cpu): + q, k, v = make_qkv(8, 128, 8) + with self.assertRaisesRegex(ValueError, "require a GPU"): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=64**-0.5, force_fused=True + ) + + @unittest.skipIf(not mx.cuda.is_available(), "CUDA kernel path only") + def test_sdpa_force_fused_cuda(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + + def make_qkv(qL, kL, D, qH=8, kH=8): + q = mx.random.normal((1, qH, qL, D), mx.float16) + k = mx.random.normal((1, kH, kL, D), mx.float16) + v = mx.random.normal((1, kH, kL, D), mx.float16) + return q, k, v + + # Vector attention kernel. + for D in (64, 96, 128): + with self.subTest(head_dim=D): + q, k, v = make_qkv(3, 128, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + def test_sdpa_sliced(self): N = 8 D = 64 From 7f062ddcb81bedb537f3f95d11a7343ba6b5f909 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Tue, 18 Aug 2026 05:19:02 -0700 Subject: [PATCH 44/84] chore: Reject negative eps in the normalization layers (#4312) --- python/mlx/nn/layers/normalization.py | 10 ++++++++++ python/tests/test_nn.py | 22 ++++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index e79440dce3..a5a8293fdb 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -47,6 +47,8 @@ def __init__( affine: bool = False, ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[InstanceNorm] 'eps' must be positive but got {eps}.") if affine: self.weight = mx.ones((dims,)) self.bias = mx.zeros((dims,)) @@ -101,6 +103,8 @@ def __init__( self, dims: int, eps: float = 1e-5, affine: bool = True, bias: bool = True ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[LayerNorm] 'eps' must be positive but got {eps}.") if affine: self.weight = mx.ones((dims,)) if bias: @@ -141,6 +145,8 @@ class RMSNorm(Module): def __init__(self, dims: int, eps: float = 1e-5): super().__init__() + if eps <= 0.0: + raise ValueError(f"[RMSNorm] 'eps' must be positive but got {eps}.") self.weight = mx.ones((dims,)) self.eps = eps @@ -191,6 +197,8 @@ def __init__( pytorch_compatible: bool = False, ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[GroupNorm] 'eps' must be positive but got {eps}.") if num_groups <= 0: raise ValueError( f"The number of groups ({num_groups}) must be a positive integer." @@ -309,6 +317,8 @@ def __init__( track_running_stats: bool = True, ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[BatchNorm] 'eps' must be positive but got {eps}.") self.num_features = num_features self.eps = eps diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index c5e6db94a7..46242d192c 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -410,6 +410,28 @@ def test_bilinear(self): outputs = layer(inputs1, inputs2) self.assertEqual(outputs.shape, (10, 6)) + def test_norm_eps_validation(self): + # eps is added under a square root. A negative one makes rsqrt take the + # root of a negative number, so the layer emits NaN for whichever + # elements have a small enough variance, which is only a partial NaN and + # easy to miss. Zero is rejected too: it leaves rsqrt(0) for any input + # whose variance is zero, which is a whole NaN row. This matches the eps + # guards the optimizers already carry. + builders = ( + ("LayerNorm", lambda eps: nn.LayerNorm(16, eps=eps)), + ("RMSNorm", lambda eps: nn.RMSNorm(16, eps=eps)), + ("GroupNorm", lambda eps: nn.GroupNorm(4, 16, eps=eps)), + ("InstanceNorm", lambda eps: nn.InstanceNorm(16, eps=eps)), + ("BatchNorm", lambda eps: nn.BatchNorm(16, eps=eps)), + ) + for name, build in builders: + for eps in (-1.0, -1e-30, 0.0): + with self.assertRaisesRegex(ValueError, "must be positive"): + build(eps) + # Anything positive still constructs, including a very small eps. + for eps in (1e-30, 1e-5, 1.0): + build(eps) + def test_group_norm(self): x = mx.arange(100, dtype=mx.float32) x = x.reshape(1, 10, 10, 1) From a4a2c1eb5c354332e5215cda05d1f295f814dad8 Mon Sep 17 00:00:00 2001 From: Xiang Chen <46052474+x14ngch3n@users.noreply.github.com> Date: Wed, 19 Aug 2026 06:54:05 +0800 Subject: [PATCH 45/84] Bound GGUF metadata string/array values against the file mapping (#4212) Co-authored-by: x14ngch3n Co-authored-by: Cheng --- .github/workflows/release.yml | 2 +- .github/workflows/update_bypass_list.yml | 3 +- mlx/io/gguf.cpp | 101 ++++++++++++++++++-- tests/load_tests.cpp | 114 +++++++++++++++++++++++ 4 files changed, 211 insertions(+), 9 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index d10493efff..b53f003ff4 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -16,7 +16,7 @@ on: required: false type: boolean schedule: - - cron: 33 6 * * 1-5 + - cron: 33 6 * * * # In jobs we must use |*publish| instead of |inputs.publish| because we can not # set default value for workflow_dispatch inputs reliably. diff --git a/.github/workflows/update_bypass_list.yml b/.github/workflows/update_bypass_list.yml index a129c51ce2..b9b71eee8b 100644 --- a/.github/workflows/update_bypass_list.yml +++ b/.github/workflows/update_bypass_list.yml @@ -5,8 +5,9 @@ on: workflow_dispatch: pull_request_target: types: - - opened - closed + schedule: + - cron: 33 6 * * * permissions: contents: write diff --git a/mlx/io/gguf.cpp b/mlx/io/gguf.cpp index 40cca573e5..6c27c46987 100644 --- a/mlx/io/gguf.cpp +++ b/mlx/io/gguf.cpp @@ -124,8 +124,7 @@ void set_mx_value_from_gguf( value = array(val->boolval, bool_); break; case GGUF_VALUE_TYPE_STRING: - value = - std::string(val->string.string, static_cast(val->string.len)); + value = std::string(val->string.string, val->string.len); break; case GGUF_VALUE_TYPE_FLOAT64: value = array(val->float64, float32); @@ -174,7 +173,7 @@ void set_mx_value_from_gguf( for (auto& str : strs) { auto str_val = reinterpret_cast(data); data += (str_val->len + sizeof(gguf_string)); - str = std::string(str_val->string, static_cast(str_val->len)); + str = std::string(str_val->string, str_val->len); ctx->off += (str_val->len + sizeof(gguf_string)); } value = std::move(strs); @@ -200,10 +199,102 @@ void set_mx_value_from_gguf( } } +inline size_t gguf_value_type_size(uint32_t type) { + switch (type) { + case GGUF_VALUE_TYPE_BOOL: + case GGUF_VALUE_TYPE_UINT8: + case GGUF_VALUE_TYPE_INT8: + return 1; + case GGUF_VALUE_TYPE_UINT16: + case GGUF_VALUE_TYPE_INT16: + return 2; + case GGUF_VALUE_TYPE_UINT32: + case GGUF_VALUE_TYPE_INT32: + case GGUF_VALUE_TYPE_FLOAT32: + return 4; + case GGUF_VALUE_TYPE_UINT64: + case GGUF_VALUE_TYPE_INT64: + case GGUF_VALUE_TYPE_FLOAT64: + return 8; + default: + return 0; + } +} + +void check_metadata_value_in_file( + const gguf_ctx* ctx, + uint32_t type, + const gguf_value* val) { + auto end = ctx->data + ctx->size; + // Bytes available from a pointer up to the end of the mapping; 0 if the + // pointer lies outside [ctx->data, end]. + auto avail = [&](const uint8_t* p) -> size_t { + return (p < ctx->data || p > end) ? 0 : static_cast(end - p); + }; + auto base = reinterpret_cast(val); + auto fail = [](const char* what) { + std::ostringstream msg; + msg << "[load_gguf] " << what + << " Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + }; + + size_t fixed = gguf_value_type_size(type); + if (fixed) { + if (fixed > avail(base)) { + fail("Metadata value extends past the end of the file."); + } + return; + } + + auto check_string = [&](const uint8_t* p) -> const uint8_t* { + uint64_t len = reinterpret_cast(p)->len; + if (sizeof(uint64_t) + len > avail(p)) { + fail("String metadata value extends past the end of the file."); + } + return p + sizeof(uint64_t) + len; + }; + + if (type == GGUF_VALUE_TYPE_STRING) { + if (sizeof(uint64_t) > avail(base)) { + fail("String metadata value extends past the end of the file."); + } + check_string(base); + return; + } + + if (type == GGUF_VALUE_TYPE_ARRAY) { + if (gguf_array_header_size > avail(base)) { + fail("Metadata value extends past the end of the file."); + } + const uint8_t* elt = base + gguf_array_header_size; + size_t elt_size = gguf_value_type_size(val->array.type); + if (elt_size) { + if (val->array.len > avail(elt) / elt_size) { + fail("Array metadata value extends past the end of the file."); + } + return; + } + if (val->array.type == GGUF_VALUE_TYPE_STRING) { + const uint8_t* p = elt; + for (uint64_t i = 0; i < val->array.len; i++) { + if (sizeof(uint64_t) > avail(p)) { + fail("Array metadata value extends past the end of the file."); + } + p = check_string(p); + } + } + return; + } + + throw std::runtime_error("[load_gguf] Received unexpected type."); +} + std::unordered_map load_metadata(gguf_ctx* ctx) { std::unordered_map metadata; gguf_key key; while (gguf_get_key(ctx, &key)) { + check_metadata_value_in_file(ctx, key.type, key.val); std::string key_name = std::string(key.name, key.namelen); auto& val = metadata.insert({key_name, GGUFMetaData{}}).first->second; set_mx_value_from_gguf(ctx, key.type, key.val, val); @@ -211,10 +302,6 @@ std::unordered_map load_metadata(gguf_ctx* ctx) { return metadata; } -// gguflib computes weights_data as ctx->data + ctx->data_off + the tensor's -// offset field in unsigned arithmetic, without comparing the result against the -// mapping, so a crafted offset can point outside the file or -- if the addition -// wraps -- back inside it at the wrong bytes. void check_tensor_in_file(const gguf_ctx* ctx, const gguf_tensor& tensor) { auto fail = [&tensor](const std::string& what) { std::ostringstream msg; diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 8974919476..6ef7bc276e 100644 --- a/tests/load_tests.cpp +++ b/tests/load_tests.cpp @@ -257,6 +257,120 @@ TEST_CASE("test gguf tensor data offset validation") { } } +// Writes a metadata-only GGUF (no tensors) whose metadata KV section is +// `kv_section` verbatim, so a caller can encode values whose lengths exceed the +// file to exercise check_metadata_value_in_file(). `kv_count` must match the +// number of KV pairs encoded in `kv_section`. +void write_raw_gguf_metadata( + const std::string& path, + uint64_t kv_count, + const std::vector& kv_section) { + std::ofstream out(path, std::ios::binary); + auto u32 = [&out](uint32_t v) { + out.write(reinterpret_cast(&v), 4); + }; + auto u64 = [&out](uint64_t v) { + out.write(reinterpret_cast(&v), 8); + }; + out.write("GGUF", 4); + u32(3); // version + u64(0); // tensor_count + u64(kv_count); // metadata_kv_count + out.write(kv_section.data(), kv_section.size()); +} + +TEST_CASE("test gguf metadata value validation") { + // A STRING/ARRAY metadata value claiming a length larger than the file must + // be rejected rather than read past the end of the mapping. See PR #4212. + + auto append_string_kv = [](std::vector& b, + const std::string& key, + uint64_t claimed_len, + bool write_payload) { + auto put = [&](const void* p, size_t n) { + b.insert( + b.end(), + static_cast(p), + static_cast(p) + n); + }; + uint64_t klen = key.size(); + put(&klen, 8); + put(key.data(), key.size()); + uint32_t vt = 8; // GGUF_VALUE_TYPE_STRING + put(&vt, 4); + put(&claimed_len, 8); + if (write_payload) { + b.insert(b.end(), claimed_len, '\0'); + } + }; + + auto append_array_kv = [](std::vector& b, + const std::string& key, + uint32_t elt_type, + uint64_t claimed_len) { + auto put = [&](const void* p, size_t n) { + b.insert( + b.end(), + static_cast(p), + static_cast(p) + n); + }; + uint64_t klen = key.size(); + put(&klen, 8); + put(key.data(), key.size()); + uint32_t vt = 9; // GGUF_VALUE_TYPE_ARRAY + put(&vt, 4); + put(&elt_type, 4); + put(&claimed_len, 8); + }; + + SUBCASE("valid empty and small strings load") { + std::vector kv; + append_string_kv(kv, "empty", 0, false); + append_string_kv(kv, "small", 5, true); + std::string file_path = get_temp_file("test_gguf_meta_ok.gguf"); + write_raw_gguf_metadata(file_path, 2, kv); + auto [weights, metadata] = load_gguf(file_path); + CHECK(weights.empty()); + CHECK(std::get(metadata.at("empty")) == ""); + CHECK(std::get(metadata.at("small")) == std::string(5, '\0')); + } + + SUBCASE("string length extends past the end of the file") { + // Claims 100 bytes of payload, none of which are present. + std::vector kv; + append_string_kv(kv, "s", 100, false); + std::string file_path = get_temp_file("test_gguf_meta_str_past.gguf"); + write_raw_gguf_metadata(file_path, 1, kv); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("string length far past the end of the file") { + std::vector kv; + append_string_kv(kv, "s", 1ull << 40, false); + std::string file_path = get_temp_file("test_gguf_meta_str_far.gguf"); + write_raw_gguf_metadata(file_path, 1, kv); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("fixed-size array length extends past the end of the file") { + // GGUF_VALUE_TYPE_UINT8 = 0; claims 2^40 elements, none present. + std::vector kv; + append_array_kv(kv, "a", 0, 1ull << 40); + std::string file_path = get_temp_file("test_gguf_meta_arr_past.gguf"); + write_raw_gguf_metadata(file_path, 1, kv); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("string array element length extends past the end of the file") { + // GGUF_VALUE_TYPE_STRING = 8; two elements, neither present. + std::vector kv; + append_array_kv(kv, "a", 8, 2); + std::string file_path = get_temp_file("test_gguf_meta_strarr_past.gguf"); + write_raw_gguf_metadata(file_path, 1, kv); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } +} + TEST_CASE("test gguf metadata") { std::string file_path = get_temp_file("test_arr.gguf"); using dict = std::unordered_map; From fa0d4463e4616a0178c500b3e2211ad687f7ff86 Mon Sep 17 00:00:00 2001 From: "Duhyeon, Kim" <49020301+dudududukim@users.noreply.github.com> Date: Wed, 19 Aug 2026 08:26:30 +0900 Subject: [PATCH 46/84] Read each K/V byte once in gqa-8 decode attention (#4077) --- .../scaled_dot_product_attention.metal | 12 ++ mlx/backend/metal/kernels/sdpa_vector.h | 143 ++++++++++++++++++ .../metal/scaled_dot_product_attention.cpp | 8 +- python/tests/test_fast_sdpa.py | 14 ++ 4 files changed, 176 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 84486d62c3..187e6fffca 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -28,6 +28,16 @@ using namespace metal; qk_dim, \ value_dim) +#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \ + sdpa_vector_2pass_1_gqa, \ + type, \ + qk_dim, \ + value_dim, \ + 8, \ + hpt) + #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ instantiate_sdpa_vector(type, 96, 96) \ @@ -35,6 +45,8 @@ using namespace metal; instantiate_sdpa_vector(type, 192, 128) \ instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ + instantiate_sdpa_vector_gqa(type, 64, 64, 8) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 4) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ diff --git a/mlx/backend/metal/kernels/sdpa_vector.h b/mlx/backend/metal/kernels/sdpa_vector.h index 1eec72be31..3f40dbd7c0 100644 --- a/mlx/backend/metal/kernels/sdpa_vector.h +++ b/mlx/backend/metal/kernels/sdpa_vector.h @@ -317,6 +317,149 @@ template } } +// Duplication-free variant for high gqa_factor decode: each simdgroup owns a +// contiguous token sub-chunk and computes HPT of its group's query heads, so +// each K/V byte is read G / HPT times instead of G times. Single-token +// queries without mask or sinks only; the partials layout matches +// sdpa_vector_2pass_2. +template +[[kernel]] void sdpa_vector_2pass_1_gqa( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device T* out [[buffer(3)]], + device float* sums [[buffer(4)]], + device float* maxs [[buffer(5)]], + const constant int& N [[buffer(7)]], + const constant size_t& k_head_stride [[buffer(8)]], + const constant size_t& k_seq_stride [[buffer(9)]], + const constant size_t& v_head_stride [[buffer(10)]], + const constant size_t& v_seq_stride [[buffer(11)]], + const constant float& scale [[buffer(12)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tpg [[threadgroups_per_grid]], + uint3 tidtg [[thread_position_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int BD = 32; + constexpr int qk_per_thread = D / BD; + constexpr int v_per_thread = V / BD; + constexpr int NT = G / HPT; + + typedef float U; + + const int kv_head_idx = tid.x; + const int batch_idx = tid.y; + const int block_idx = tid.z; + const int blocks = tpg.z; + const int g = tidtg.y; + const int cchunk = g / NT; + const int h0 = (g % NT) * HPT; + const int num_kv_heads = tpg.x; + const int num_q_heads = num_kv_heads * G; + const int base_head = batch_idx * num_q_heads + kv_head_idx * G; + + const int chunk = (N + blocks - 1) / blocks; + const int kstart = block_idx * chunk; + const int kend = min(N, kstart + chunk); + const int sub = (chunk + HPT - 1) / HPT; + const int s0 = kstart + cchunk * sub; + const int s1 = min(kend, s0 + sub); + + const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride + + simd_lid * qk_per_thread; + const device T* vp = values + kv_head_idx * v_head_stride + + s0 * v_seq_stride + simd_lid * v_per_thread; + + U q[HPT][qk_per_thread]; + for (int j = 0; j < HPT; j++) { + const device T* qp = + queries + (base_head + h0 + j) * D + simd_lid * qk_per_thread; + for (int i = 0; i < qk_per_thread; i++) { + q[j][i] = static_cast(scale) * qp[i]; + } + } + + U max_score[HPT]; + U sum_exp_score[HPT]; + U o[HPT][v_per_thread]; + for (int j = 0; j < HPT; j++) { + max_score[j] = Limits::finite_min; + sum_exp_score[j] = 0; + for (int i = 0; i < v_per_thread; i++) { + o[j][i] = 0; + } + } + + for (int t = s0; t < s1; t++) { + U kr[qk_per_thread]; + U vr[v_per_thread]; + for (int i = 0; i < qk_per_thread; i++) { + kr[i] = kp[i]; + } + for (int i = 0; i < v_per_thread; i++) { + vr[i] = vp[i]; + } + kp += k_seq_stride; + vp += v_seq_stride; + for (int j = 0; j < HPT; j++) { + U score = 0; + for (int i = 0; i < qk_per_thread; i++) { + score += q[j][i] * kr[i]; + } + score = simd_sum(score); + U new_max = max(max_score[j], score); + U factor = fast::exp(max_score[j] - new_max); + U exp_score = fast::exp(score - new_max); + max_score[j] = new_max; + sum_exp_score[j] = sum_exp_score[j] * factor + exp_score; + for (int i = 0; i < v_per_thread; i++) { + o[j][i] = o[j][i] * factor + exp_score * vr[i]; + } + } + } + + threadgroup U o_sh[G * HPT * V]; + threadgroup U se_sh[G * HPT]; + threadgroup U mx_sh[G * HPT]; + for (int j = 0; j < HPT; j++) { + int slot = (h0 + j) * HPT + cchunk; + U inv = sum_exp_score[j] > 0 ? 1 / sum_exp_score[j] : 0; + for (int i = 0; i < v_per_thread; i++) { + o_sh[slot * V + simd_lid * v_per_thread + i] = o[j][i] * inv; + } + if (simd_lid == 0) { + se_sh[slot] = sum_exp_score[j]; + mx_sh[slot] = max_score[j]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + U gmax = Limits::finite_min; + for (int s = 0; s < HPT; s++) { + gmax = max(gmax, mx_sh[g * HPT + s]); + } + U denom = 0; + U acc[v_per_thread] = {0}; + for (int s = 0; s < HPT; s++) { + U w = se_sh[g * HPT + s] * fast::exp(mx_sh[g * HPT + s] - gmax); + denom += w; + for (int i = 0; i < v_per_thread; i++) { + acc[i] += w * o_sh[(g * HPT + s) * V + simd_lid * v_per_thread + i]; + } + } + + const int o_offset = base_head + g; + device T* op = + out + o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread; + for (int i = 0; i < v_per_thread; i++) { + op[i] = static_cast(acc[i]); + } + if (simd_lid == 0) { + sums[o_offset * blocks + block_idx] = denom; + maxs[o_offset * blocks + block_idx] = gmax; + } +} + template [[kernel]] void sdpa_vector_2pass_2( const device T* partials [[buffer(0)]], diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index acff685790..d9cfd0819c 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -433,7 +433,13 @@ void sdpa_vector_2pass( // Set the kernel name std::string kname; kname.reserve(64); - kname += "sdpa_vector_2pass_1_"; + kname += "sdpa_vector_2pass_1"; + if (!mask && !sinks && q.shape(2) == 1 && q.shape(1) == 8 * k.shape(1) && + q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) && + k.shape(2) >= 8192) { + kname += "_gqa"; + } + kname += "_"; kname += get_type_string(q.dtype()); kname += "_"; kname += std::to_string(q.shape(-1)); diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 997fa1028c..c267410612 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -275,6 +275,20 @@ def test_sdpa_vector(self): ) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_vector_gqa_long(self): + scale = 1.0 + mx.random.seed(0) + for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64)]: + for L in [8192, 8201]: + q = 5e-1 * mx.random.normal(shape=(1, Nq, 1, D)) + k = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + v = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + kr = mx.repeat(k, Nq // Nkv, axis=1) + vr = mx.repeat(v, Nq // Nkv, axis=1) + ref = mlx_primitives_sdpa(q, kr, vr, scale) + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_fully_masked(self): Lkv = 8 mask = mx.array(False) From 3e8113cfaeb6dd7950a97355a435e3617895b397 Mon Sep 17 00:00:00 2001 From: rohith Date: Wed, 19 Aug 2026 06:23:05 +0530 Subject: [PATCH 47/84] Fix fft vmap and jvp for transforms over a subset of axes (#4138) --- mlx/primitives.cpp | 18 ++++++++++-------- python/tests/test_fft.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 8 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index d4975e87f7..7a6c729c39 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -2253,12 +2253,13 @@ std::pair, std::vector> FFT::vmap( if (fft_ax >= ax) { fft_ax++; } - if (real_) { - auto n = out_shape[fft_ax]; - out_shape[fft_ax] = inverse_ ? 2 * (n - 1) : n / 2 + 1; - } } } + // Only the last transformed axis changes size in a real transform + if (real_) { + auto n = out_shape[fft_axes.back()]; + out_shape[fft_axes.back()] = inverse_ ? 2 * (n - 1) : n / 2 + 1; + } return { {array( out_shape, @@ -2362,14 +2363,15 @@ std::vector FFT::jvp( assert(primals.size() == 1); assert(argnums.size() == 1); auto& tan = tangents[0]; + std::vector axes(axes_.begin(), axes_.end()); if (real_ & inverse_) { - return {fft::irfftn(tan, fft::FFTNorm::Backward, stream())}; + return {fft::irfftn(tan, axes, fft::FFTNorm::Backward, stream())}; } else if (real_) { - return {fft::rfftn(tan, fft::FFTNorm::Backward, stream())}; + return {fft::rfftn(tan, axes, fft::FFTNorm::Backward, stream())}; } else if (inverse_) { - return {fft::ifftn(tan, fft::FFTNorm::Backward, stream())}; + return {fft::ifftn(tan, axes, fft::FFTNorm::Backward, stream())}; } else { - return {fft::fftn(tan, fft::FFTNorm::Backward, stream())}; + return {fft::fftn(tan, axes, fft::FFTNorm::Backward, stream())}; } } diff --git a/python/tests/test_fft.py b/python/tests/test_fft.py index 9358ede794..1f96aad566 100644 --- a/python/tests/test_fft.py +++ b/python/tests/test_fft.py @@ -446,6 +446,41 @@ def g(x): dgdx = torch.func.grad(g)(x_torch) self.assertLess((dfdx - dgdx).abs().max() / dgdx.abs().mean(), 1e-4) + def make_ffts(self): + mxffts = { + (True, True): mx.fft.irfftn, + (True, False): mx.fft.rfftn, + (False, True): mx.fft.ifftn, + (False, False): mx.fft.fftn, + } + shape = (3, 8, 6) + r = np.random.rand(*shape).astype(np.float32) + i = np.random.rand(*shape).astype(np.float32) + for (real, inverse), fftn in mxffts.items(): + a_np = r if real and not inverse else r + 1j * i + for axes in [(-1,), (0,), (-2, -1), (-1, -2), (0, 1)]: + yield fftn, a_np, axes + + def test_fft_vmap(self): + for fftn, a_np, axes in self.make_ffts(): + a = mx.array(a_np) + f = lambda x: fftn(x, axes=axes) + expected = mx.stack([f(a[i]) for i in range(a.shape[0])]) + out = mx.vmap(f)(a) + self.assertEqual(tuple(out.shape), tuple(expected.shape)) + np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + + def test_fft_jvp(self): + # The fft is linear so the jvp is the fft of the tangent + for fftn, a_np, axes in self.make_ffts(): + a = mx.array(a_np) + t = mx.array(np.random.rand(*a_np.shape).astype(a_np.dtype)) + f = lambda x: fftn(x, axes=axes) + expected = f(t) + out = mx.jvp(f, [a], [t])[1][0] + self.assertEqual(tuple(out.shape), tuple(expected.shape)) + np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From db935c094e00870f5148ac622c2f5cb901d61795 Mon Sep 17 00:00:00 2001 From: Ishaan Samantray Date: Tue, 18 Aug 2026 21:06:56 -0400 Subject: [PATCH 48/84] Fix median dropping NaN (#4146) --- mlx/ops.cpp | 10 ++++++++++ python/tests/test_ops.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 5bc69eb2f5..34bf40538f 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2411,6 +2411,16 @@ array median( array(0.5, dtype), s); } + // Sorting moves NaN to the end, so the midpoint slice never selects it. + // Propagate it explicitly to stay consistent with max, min and mean. + if (issubdtype(a.dtype(), inexact)) { + median_a = where( + any(isnan(flat_a, s), -1, /* keepdims = */ true, s), + array(std::numeric_limits::quiet_NaN(), dtype), + median_a, + s); + } + median_a = squeeze(median_a, -1, s); if (keepdims) { median_a = expand_dims(median_a, sorted_axes, s); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index e5761c78b2..240e97dc18 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -991,6 +991,38 @@ def test_median(self): out_np = np.median(x, axis=(0, 1, 3), keepdims=True) self.assertTrue(np.allclose(out, out_np)) + def test_median_nan(self): + nan = float("nan") + + # Odd and even lengths, with the NaN in a few different positions. + for vals in ([1.0, nan, 0.0], [nan, 1.0, 0.0], [1.0, 2.0, nan, 4.0]): + for dtype in (mx.float16, mx.bfloat16, mx.float32): + out = mx.median(mx.array(vals, dtype=dtype)) + self.assertTrue(mx.isnan(out).item(), msg=f"{vals} {dtype}") + + x = mx.array([[1.0, nan, 3.0], [4.0, 5.0, 6.0]]) + self.assertTrue( + np.array_equal( + np.array(mx.median(x, axis=1)), np.median(x, axis=1), equal_nan=True + ) + ) + self.assertTrue( + np.array_equal( + np.array(mx.median(x, axis=0)), np.median(x, axis=0), equal_nan=True + ) + ) + self.assertTrue(mx.isnan(mx.median(x)).item()) + self.assertEqual(mx.median(x, axis=1, keepdims=True).shape, (2, 1)) + + # Complex NaN propagates too, matching NumPy. + out = mx.median(mx.array([complex(1, 0), complex(nan, 0), complex(0, 0)])) + self.assertTrue(mx.isnan(out).item()) + + # A NaN-free array is unaffected, and integers are never NaN. + x = mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + self.assertTrue(np.allclose(mx.median(x, axis=1), np.median(x, axis=1))) + self.assertEqual(mx.median(mx.array([0, 1, 2, 3, 4])).item(), 2) + def test_var(self): x = mx.array( [ From eb38e2545b068d45922ef9bc0ef6d21158910682 Mon Sep 17 00:00:00 2001 From: rohith Date: Wed, 19 Aug 2026 12:28:18 +0530 Subject: [PATCH 49/84] Fix the CPU scan over a size one axis with a padded stride (#4139) Co-authored-by: Cheng --- mlx/backend/cpu/scan.cpp | 3 ++- python/tests/test_ops.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 3ebbe0a3c3..93e67825e2 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -163,7 +163,8 @@ void scan_op( const Op& op, U init) { if (in.flags().row_contiguous) { - if (in.strides()[axis] == 1) { + // A size-one axis can carry any stride and still be row contiguous. + if (in.strides()[axis] == 1 || in.shape(axis) == 1) { contiguous_scan( in.data(), out.data(), diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 240e97dc18..22a794b4fc 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2663,6 +2663,20 @@ def fn(its): mem4 = mx.get_peak_memory() self.assertEqual(mem2, mem4) + def test_scan_size_one_axis(self): + # A size one axis can carry any stride and still be row contiguous, so + # the scan must not take its row count from that stride. + for op in ["cumsum", "cumprod", "cummax", "cummin"]: + for start in (1, 2, 3): + with self.subTest(op=op, start=start): + base = mx.arange(1, 11, dtype=mx.float32).reshape(1, 10) + a = base[:, start:] + mx.eval(a) + # The axis has size one, so an inclusive scan is the identity + expected = np.array(a).copy() + out = getattr(mx, op)(a, axis=0) + self.assertTrue(np.array_equal(np.array(out), expected)) + def test_cummax_cummin_nan(self): nan = float("nan") cases = [ From 8a81722b1d71cac9b7dde47e56a438c4b529129b Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Wed, 19 Aug 2026 00:45:47 -0700 Subject: [PATCH 50/84] chore: Validate the optimizer betas at construction (#4310) Co-authored-by: Cheng --- python/mlx/optimizers/optimizers.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index 65efab222d..95602928cf 100644 --- a/python/mlx/optimizers/optimizers.py +++ b/python/mlx/optimizers/optimizers.py @@ -499,6 +499,13 @@ def __init__( ): super().__init__() + for i, beta in enumerate(betas): + if not 0.0 <= beta < 1.0: + raise ValueError( + f"Adam beta{i + 1} should be in [0, 1), {beta} was provided " + "instead" + ) + self._maybe_schedule("learning_rate", learning_rate) self.betas = betas self.eps = eps @@ -683,6 +690,13 @@ def __init__( ): super().__init__() + for i, beta in enumerate(betas): + if not 0.0 <= beta < 1.0: + raise ValueError( + f"Lion beta{i + 1} should be in [0, 1), {beta} was provided " + "instead" + ) + self._maybe_schedule("learning_rate", learning_rate) self.betas = betas self.weight_decay = weight_decay From 6c0f02a75e710e216c3be72e0cac239cf7dfc942 Mon Sep 17 00:00:00 2001 From: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:23:35 +0800 Subject: [PATCH 51/84] `RMSNormVJP` backward writes a full `{n_rows, D}` `gw_temp` intermediate (#4293) --- mlx/backend/metal/kernels/rms_norm.metal | 157 ++++++++++++++--------- mlx/backend/metal/normalization.cpp | 31 +++-- 2 files changed, 115 insertions(+), 73 deletions(-) diff --git a/mlx/backend/metal/kernels/rms_norm.metal b/mlx/backend/metal/kernels/rms_norm.metal index a50d4a25c6..eb9b5c0af1 100644 --- a/mlx/backend/metal/kernels/rms_norm.metal +++ b/mlx/backend/metal/kernels/rms_norm.metal @@ -166,21 +166,32 @@ template constant float& eps, constant uint& axis_size, constant uint& w_stride, + constant uint& n_rows, + constant uint& rows_per_group, uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]]) { - // Advance the input pointers - x += gid * size_t(axis_size) + lid * N_READS; - g += gid * size_t(axis_size) + lid * N_READS; w += w_stride * lid * N_READS; + float thread_w[N_READS]; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + thread_w[i] = w[w_stride * i]; + } + } else { + for (int i = 0; i < N_READS; i++) { + thread_w[i] = + (lid * N_READS + i < axis_size) ? (float)w[w_stride * i] : 0; + } + } // Allocate registers for the computation and accumulators float thread_x[N_READS]; - float thread_w[N_READS]; float thread_g[N_READS]; - float sumx2 = 0; - float sumgwx = 0; + float gw_acc[N_READS]; + for (int i = 0; i < N_READS; i++) { + gw_acc[i] = 0; + } // Allocate shared memory to implement the reduction constexpr int SIMD_SIZE = 32; @@ -189,75 +200,99 @@ template threadgroup float local_normalizer[1]; threadgroup float local_meangwx[1]; - // Read and accumulate locally - if (lid * N_READS + N_READS <= axis_size) { - for (int i = 0; i < N_READS; i++) { - thread_x[i] = x[i]; - thread_w[i] = w[w_stride * i]; - thread_g[i] = g[i]; + uint row_end = gid * rows_per_group + rows_per_group; + if (row_end > n_rows) { + row_end = n_rows; + } + for (uint row = gid * rows_per_group; row < row_end; ++row) { + const device T* x_in = x + size_t(row) * axis_size + lid * N_READS; + const device T* g_in = g + size_t(row) * axis_size + lid * N_READS; - sumx2 += thread_x[i] * thread_x[i]; - sumgwx += thread_x[i] * thread_w[i] * thread_g[i]; - } - } else { - for (int i = 0; i < N_READS; i++) { - if ((lid * N_READS + i) < axis_size) { - thread_x[i] = x[i]; - thread_w[i] = w[w_stride * i]; - thread_g[i] = g[i]; + float sumx2 = 0; + float sumgwx = 0; + + // Read and accumulate locally + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + thread_x[i] = x_in[i]; + thread_g[i] = g_in[i]; sumx2 += thread_x[i] * thread_x[i]; sumgwx += thread_x[i] * thread_w[i] * thread_g[i]; } + } else { + for (int i = 0; i < N_READS; i++) { + if ((lid * N_READS + i) < axis_size) { + thread_x[i] = x_in[i]; + thread_g[i] = g_in[i]; + + sumx2 += thread_x[i] * thread_x[i]; + sumgwx += thread_x[i] * thread_w[i] * thread_g[i]; + } + } } - } - // Accumulate across threads - sumx2 = simd_sum(sumx2); - sumgwx = simd_sum(sumgwx); - if (simd_group_id == 0) { - local_sumx2[simd_lane_id] = 0; - local_sumgwx[simd_lane_id] = 0; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (simd_lane_id == 0) { - local_sumx2[simd_group_id] = sumx2; - local_sumgwx[simd_group_id] = sumgwx; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (simd_group_id == 0) { - sumx2 = simd_sum(local_sumx2[simd_lane_id]); - sumgwx = simd_sum(local_sumgwx[simd_lane_id]); + // Accumulate across threads + sumx2 = simd_sum(sumx2); + sumgwx = simd_sum(sumgwx); + if (simd_group_id == 0) { + local_sumx2[simd_lane_id] = 0; + local_sumgwx[simd_lane_id] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); if (simd_lane_id == 0) { - local_meangwx[0] = sumgwx / axis_size; - local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps); + local_sumx2[simd_group_id] = sumx2; + local_sumgwx[simd_group_id] = sumgwx; } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - float meangwx = local_meangwx[0]; - float normalizer = local_normalizer[0]; - float normalizer3 = normalizer * normalizer * normalizer; - - // Write the outputs - gx += gid * size_t(axis_size) + lid * N_READS; - gw += gid * size_t(axis_size) + lid * N_READS; - if (lid * N_READS + N_READS <= axis_size) { - for (int i = 0; i < N_READS; i++) { - gx[i] = static_cast( - thread_g[i] * thread_w[i] * normalizer - - thread_x[i] * meangwx * normalizer3); - if (has_w) { - gw[i] = static_cast(thread_g[i] * thread_x[i] * normalizer); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group_id == 0) { + sumx2 = simd_sum(local_sumx2[simd_lane_id]); + sumgwx = simd_sum(local_sumgwx[simd_lane_id]); + if (simd_lane_id == 0) { + local_meangwx[0] = sumgwx / axis_size; + local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps); } } - } else { - for (int i = 0; i < N_READS; i++) { - if ((lid * N_READS + i) < axis_size) { - gx[i] = static_cast( + threadgroup_barrier(mem_flags::mem_threadgroup); + float meangwx = local_meangwx[0]; + float normalizer = local_normalizer[0]; + float normalizer3 = normalizer * normalizer * normalizer; + + // Write the outputs + device T* gx_out = gx + size_t(row) * axis_size + lid * N_READS; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + gx_out[i] = static_cast( thread_g[i] * thread_w[i] * normalizer - thread_x[i] * meangwx * normalizer3); if (has_w) { - gw[i] = static_cast(thread_g[i] * thread_x[i] * normalizer); + gw_acc[i] += thread_g[i] * thread_x[i] * normalizer; + } + } + } else { + for (int i = 0; i < N_READS; i++) { + if ((lid * N_READS + i) < axis_size) { + gx_out[i] = static_cast( + thread_g[i] * thread_w[i] * normalizer - + thread_x[i] * meangwx * normalizer3); + if (has_w) { + gw_acc[i] += thread_g[i] * thread_x[i] * normalizer; + } + } + } + } + } + + if (has_w) { + gw += size_t(gid) * axis_size + lid * N_READS; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + gw[i] = static_cast(gw_acc[i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + if ((lid * N_READS + i) < axis_size) { + gw[i] = static_cast(gw_acc[i]); } } } diff --git a/mlx/backend/metal/normalization.cpp b/mlx/backend/metal/normalization.cpp index 9a222cdd6c..f3370f087d 100644 --- a/mlx/backend/metal/normalization.cpp +++ b/mlx/backend/metal/normalization.cpp @@ -109,11 +109,9 @@ void RMSNormVJP::eval_gpu( array x_copy = contiguous_copy_gpu(x, s); return {x_copy, true}; }; - bool donate_g = inputs[2].is_donatable(); auto [x, copied] = check_input(inputs[0]); const array& w = inputs[1]; auto [g, g_copied] = check_input(inputs[2]); - donate_g |= g_copied; array& gx = outputs[0]; array& gw = outputs[1]; @@ -137,17 +135,22 @@ void RMSNormVJP::eval_gpu( auto axis_size = static_cast(x.shape().back()); int n_rows = x.data_size() / axis_size; + const int target_groups = 512; + uint32_t rows_per_group = 1; + int n_groups = n_rows; + if (axis_size <= RMS_LOOPED_LIMIT) { + rows_per_group = (n_rows + target_groups - 1) / target_groups; + n_groups = (n_rows + rows_per_group - 1) / rows_per_group; + } + // Allocate the gradient accumulator gw and a temporary to store the // gradients before they are accumulated. - array gw_temp = - (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w; + array gw_temp = (has_w) + ? array({n_groups, x.shape().back()}, gw.dtype(), nullptr, {}) + : w; if (has_w) { - if (!g_in_gx && donate_g) { - gw_temp.copy_shared_buffer(g); - } else { - gw_temp.set_data(allocator::malloc(gw_temp.nbytes())); - compute_encoder.add_temporary(gw_temp); - } + gw_temp.set_data(allocator::malloc(gw_temp.nbytes())); + compute_encoder.add_temporary(gw_temp); } gw.set_data(allocator::malloc(gw.nbytes())); @@ -174,7 +177,7 @@ void RMSNormVJP::eval_gpu( size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size; size_t threadgroup_size = simd_size * simds_needed; assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup()); - size_t n_threads = n_rows * threadgroup_size; + size_t n_threads = n_groups * threadgroup_size; grid_dims = MTL::Size(n_threads, 1, 1); group_dims = MTL::Size(threadgroup_size, 1, 1); } else { @@ -194,12 +197,16 @@ void RMSNormVJP::eval_gpu( compute_encoder.set_bytes(eps_, 5); compute_encoder.set_bytes(axis_size, 6); compute_encoder.set_bytes(w_stride, 7); + if (axis_size <= looped_limit) { + compute_encoder.set_bytes(static_cast(n_rows), 8); + compute_encoder.set_bytes(rows_per_group, 9); + } compute_encoder.dispatch_threads(grid_dims, group_dims); } if (has_w) { ReductionPlan plan( - ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size}); + ReductionOpType::ContiguousStridedReduce, {n_groups}, {axis_size}); strided_reduce_general_dispatch( gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s); } From 3a98589826141e81cd04cfb8b8db544df27dcd96 Mon Sep 17 00:00:00 2001 From: Aaishwarya Mishra Date: Thu, 20 Aug 2026 02:18:39 +0530 Subject: [PATCH 52/84] [Bug]: add default none value to axis parameter of the take_along_axis (#4357) Co-authored-by: Anastasiia Filippova --- python/src/ops.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/src/ops.cpp b/python/src/ops.cpp index b8390b036d..677d998af1 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1768,7 +1768,7 @@ void init_ops(nb::module_& m) { }, nb::arg(), "indices"_a, - "axis"_a.none(), + "axis"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( From 714a7efcb83c1424b4ade9226a5fd810fa184b73 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang <19340816+wyanzhao@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:27:39 -0700 Subject: [PATCH 53/84] Add a fused full-attention path for head_dim 256 on NAX devices (#3842) Co-authored-by: Cheng --- benchmarks/python/sdpa_bench.py | 11 +- mlx/backend/metal/jit_kernels.cpp | 5 +- mlx/backend/metal/kernels.h | 3 +- .../steel/attn/kernels/steel_attention_nax.h | 421 ++++++++++++++++++ .../attn/kernels/steel_attention_nax.metal | 17 +- mlx/backend/metal/nojit_kernels.cpp | 3 +- .../metal/scaled_dot_product_attention.cpp | 74 ++- python/tests/test_fast_sdpa.py | 67 +++ 8 files changed, 575 insertions(+), 26 deletions(-) diff --git a/benchmarks/python/sdpa_bench.py b/benchmarks/python/sdpa_bench.py index 4130b05760..7dfc7e0d1d 100644 --- a/benchmarks/python/sdpa_bench.py +++ b/benchmarks/python/sdpa_bench.py @@ -215,9 +215,18 @@ def get_gflop_count(B, M, N, K): ( 1, 4096, 5000, 128, 32, 8), ( 1, 2048, 32121, 128, 32, 8), ) + + shapes_256 = ( + # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 1024, 1024, 256, 24, 4), + ( 1, 2048, 2048, 256, 24, 4), + ( 1, 4096, 4096, 256, 24, 4), + ( 1, 4096, 5000, 256, 24, 4), + ( 1, 2048, 32121, 256, 24, 4), + ) # fmt: on - shapes = shapes_64 + shapes_72 + shapes_80 + shapes_96 + shapes_128 + shapes = shapes_64 + shapes_72 + shapes_80 + shapes_96 + shapes_128 + shapes_256 masks = [None, "bool", "causal"] diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index c7900ecdf8..8dfe30a15c 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -1330,7 +1330,8 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bd, int wm, int wn, - const array& m) { + const array& m, + bool split_d) { const auto& lib_name = kernel_name; auto lib = d.get_library(lib_name, [&]() { std::string kernel_source; @@ -1340,7 +1341,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( metal::steel_attention_nax(), get_template_definition( lib_name, - "attention_nax", + split_d ? "attention_nax_dsplit" : "attention_nax", get_type_string(q.dtype()), bq, bk, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 21b754514c..18a56ebf1b 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -426,7 +426,8 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bd, int wm, int wn, - const array& m); + const array& m, + bool split_d); // Create a GPU kernel template definition for JIT compilation template diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index b48a9a942d..4a5a9716fd 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -484,3 +484,424 @@ template < Otile.store(O, int(params->O_strides[2])); } } + +/////////////////////////////////////////////////////////////////////////////// +// Head-dim split attention kernel +/////////////////////////////////////////////////////////////////////////////// + +// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup +// accumulator working set of attention_nax (TD output fragments plus the S +// fragments) is what gates tensor-unit throughput, so this kernel splits the +// head dim across the WN = 2 simdgroups of the second warp dimension: each +// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, +// halving its accumulator set. The pair exchanges its partial Q@K.T sums +// through threadgroup memory, then both simdgroups run softmax redundantly on +// the full S tile (the row statistics are cheap) and each accumulates P@V for +// its own half of Dv. + +// clang-format off +template < + typename T, + int BQ, + int BK, + int BD, + int WM, + int WN, + typename MaskType = float, + typename AccumType = float> +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention_nax_dsplit( + const device T* Q [[buffer(0)]], + const device T* K [[buffer(1)]], + const device T* V [[buffer(2)]], + device T* O [[buffer(3)]], + const constant AttnParams* params [[buffer(4)]], + const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], + const device MaskType* mask [[buffer(6), function_constant(has_mask)]], + const device T* sinks [[buffer(7), function_constant(has_sinks)]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on + + // Pacifying compiler + (void)lid; + + // Move to correct block + ulong3 tidl{tid.x, tid.y, tid.z}; + + Q += tidl.z * params->Q_strides[0] + // Batch + tidl.y * params->Q_strides[1] + // Head + tidl.x * BQ * params->Q_strides[2]; // Sequence + + ulong kv_head_idx = int(tid.y) / params->gqa_factor; + K += tidl.z * params->K_strides[0] + // Batch + kv_head_idx * params->K_strides[1]; // Head + + V += tidl.z * params->V_strides[0] + // Batch + kv_head_idx * params->V_strides[1]; // Head + + O += tidl.z * params->O_strides[0] + // Batch + tidl.y * params->O_strides[1] + // Head + tidl.x * BQ * params->O_strides[2]; // Sequence + + if (has_mask) { + mask += tidl.z * mask_params->M_strides[0] + // Batch + tidl.y * mask_params->M_strides[1]; // Head + } + + const metal::uniform scale2 = + make_uniform(params->scale) * make_uniform(1.44269504089f); + + // Prepare MMA tiles + constexpr short kU = 16; + + // The WM simdgroups along the first warp dimension split the Q sequence; + // the WN simdgroups along the second split the head dim. The exchange + // below reduces exactly one peer, so WN is fixed at 2. + static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + constexpr int kNWarps = WM; + static_assert( + BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, + "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); + + // Q seq frags per warp + constexpr int TQ = BQ / (kNWarps * kU); + // HeadDim frags over the full head dim + constexpr int TD = BD / kU; + // KV seq frags per warp + constexpr short TK = BK / kU; + + static_assert(TQ == 1, "Check TQ"); + static_assert(TD % WN == 0, "The head dim must split evenly across WN"); + + // HeadDim frags / columns owned by each of the WN simdgroups of a row group + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + + static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); + static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + + using otile_t = NAXTile; + otile_t Otile; + Otile.clear(); + + const short tm = kU * TQ * row_group; + Q += tm * int(params->Q_strides[2]) + d_half * BDh; + K += d_half * BDh; + V += d_half * BDh; + O += tm * int(params->O_strides[2]) + d_half * BDh; + + constexpr short kRowsPT = otile_t::kRowsPerThread; + + metal::vec max_score; + metal::vec sum_score{0}; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = Limits::finite_min; + } + + if (has_sinks) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = M_LOG2E_F * static_cast(sinks[tidl.y]); + sum_score[i] = 1; + } + } + + int kb_lim = params->NK; + int kb_min_causal = params->NK; + + if (do_causal) { + int q_max = (tid.x + 1) * BQ + params->qL_off; + kb_lim = (q_max + BK - 1) / BK; + kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); + } + + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + const short lim_rows_q = params->qL_rem - tm; + const short lim_rows_k = params->kL_rem; + + using stile_t = NAXTile; + constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; + + // One slot per (row group, half): a fragment pair in per-lane-linear + // layout. Both halves share the fragment-to-lane mapping, so the + // exchange needs no coordinate math. + threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + + // Keep the simdgroup's Q half resident in registers for the whole KV + // loop: TDh fragments of T are cheap next to the accumulators. + NAXTile Qtiles[TDh]; + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + const int Q_load_off = id * kU; + if (!align_Q && is_last_q) { + Qtiles[id].load_rows( + Q + Q_load_off, int(params->Q_strides[2]), lim_rows_q); + } else { + Qtiles[id].load(Q + Q_load_off, int(params->Q_strides[2])); + } + } + + const short2 simd_coord = otile_t::NAXFrag_t::get_coord(); + const short sm = simd_coord.y; + const short sn = simd_coord.x; + + // Loop over KV seq length + for (int kb = 0; kb < kb_lim; kb++) { + const int is_last_k = (kb == (params->NK_aligned)); + + stile_t Stile; + Stile.clear(); + + // S = Q @ K.T, this half of D only, exchanged pair by pair. + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + NAXTile Ktile; + const int K_load_off = ik * kU * int(params->K_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Ktile.load_rows( + K + K_load_off, int(params->K_strides[2]), lim_rows_k - ik * kU); + } else { + Ktile.load(K + K_load_off, int(params->K_strides[2])); + } + + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), + Stile.frag_at(0, ik + 1), + Qtiles[id].frag_at(0, 0), + metal::false_type{}, + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), + metal::true_type{}); + } + + // Exchange the partial pair and reduce. + threadgroup AccumType* slot = s_xchg[row_group][d_half]; + thread auto& s0 = Stile.frag_at(0, ik); + thread auto& s1 = Stile.frag_at(0, ik + 1); + const short base = short(simd_lane_id) * (2 * kEPF); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + slot[base + i] = s0[i]; + slot[base + kEPF + i] = s1[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s0[i] += peer[base + i]; + s1[i] += peer[base + kEPF + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + // Scale S + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { + Stile.elems()[ii] *= float(scale2); + } + + // Mask out length sequence + if (!align_K && is_last_k) { + constexpr auto neg_inf = Limits::finite_min; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const short col_pos = ik * kU + sn; + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = ((col_pos + jj) < params->kL_rem) ? fg[loc] : neg_inf; + } + } + } + } + + // Mask out if causal + if (do_causal && kb >= kb_min_causal) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + params->qL_off + tm; + const int base_col = kb * BK; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto r = base_row + ii * stile_t::kFragRowsJump + sm; + const auto c = base_col + ik * kU + jj + sn; + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = (r < c) ? neg_inf : fg[loc]; + } + } + } + } + + // Other masking as needed + if (has_mask) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + tm; + const int base_col = kb * BK; + + constexpr bool is_bool = is_same_v; + using melem_t = typename metal::conditional_t; + using mtile_t = NAXTile; + using mfrag_t = typename mtile_t::frag_type; + + if (base_row + kU <= params->qL && base_col + BK <= params->kL) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load_safe( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + params->qL, + params->kL, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } + } + + // Do softmax (redundantly per half; the row statistics are cheap) + metal::vec new_max; + metal::vec factor; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + new_max[i] = max_score[i]; + } + + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + factor[i] = fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + sum_score[i] = sum_score[i] * factor[i]; + } + + Stile.template row_reduce(sum_score); + + Otile.template row_bin_op(factor); + + simdgroup_barrier(mem_flags::mem_none); + + // O = P @ V, this half of Dv only. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id += 2) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + NAXTile Vtile; + + const int V_load_off = ik * kU * int(params->V_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Vtile.load_rows( + V + V_load_off, int(params->V_strides[2]), lim_rows_k - ik * kU); + } else { + Vtile.load(V + V_load_off, int(params->V_strides[2])); + } + + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), + Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); + } + } + + // Next block + K += BK * int(params->K_strides[2]); + V += BK * int(params->V_strides[2]); + } + + // Normalize output + threadgroup_barrier(mem_flags::mem_none); + + metal::vec rcp; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + rcp[i] = 1.f / sum_score[i]; + } + + Otile.template row_bin_op(rcp); + + if (!align_Q && is_last_q) { + if (lim_rows_q <= 0) + return; + Otile.store_rows(O, int(params->O_strides[2]), lim_rows_q); + } else { + Otile.store(O, int(params->O_strides[2])); + } +} diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal index c2b60b9cf0..66d55539ab 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal @@ -11,11 +11,18 @@ "_wm" #wm "_wn" #wn "_mask" #mname, \ attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float) -#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 96, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ +#define instantiate_attn_dsplit(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_dsplit_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention_nax_dsplit, dtype, bq, bk, bd, wm, wn, mtype, float) + +#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ + instantiate_attn_dsplit(iname, itype, 64, 32, 256, 4, 2, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 96, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 64, 4, 1, mname, mtype) #define instantiate_attn_mask_helper(iname, itype) \ diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 5da78db8a3..3795a6fb22 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -503,7 +503,8 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int, int, int, - const array&) { + const array&, + bool) { return d.get_kernel(kernel_name, hash_name, func_consts); } diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index d9cfd0819c..981d2b1817 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -28,13 +28,14 @@ void sdpa_full_self_attention_nax( const std::optional& sinks) { using namespace mlx::steel; - int wm = 4; - int wn = 1; - int bd = q.shape(-1); int bq = 64; int bk = 32; + bool split_d = bd == 256; + int wm = 4; + int wn = split_d ? 2 : 1; + int B = q.shape(0); int H = q.shape(1); int D = q.shape(3); @@ -43,6 +44,36 @@ void sdpa_full_self_attention_nax( int qL = q.shape(2); int kL = k.shape(2); + // The causal offset describes the true diagonal even if kL is widened + // below. + int qL_off = kL - qL; + + // Check if K/V are from chunked KV cache, and assume aligned K/V if so. + auto has_backing_rows = [](const array& kv, int rows) { + auto& st = kv.strides(); + if ((st[0] < 0) || (st[1] <= 0) || (st[2] <= 0) || (st[1] % st[2] != 0)) { + return false; + } + int64_t itemsize = kv.itemsize(); + // The rows must stay inside the head's row pitch (so they belong to the + // cache the slice was taken from) ... + int64_t pitch = st[1] / st[2]; + int64_t row0 = ((kv.offset() / itemsize) % st[1]) / st[2]; + if (row0 + rows > pitch) { + return false; + } + // ... and inside the buffer. + int64_t end = (kv.shape(0) - 1) * st[0] + (kv.shape(1) - 1) * st[1] + + (rows - 1) * st[2] + kv.shape(3); + return kv.offset() + end * itemsize <= int64_t(kv.buffer_size()); + }; + if (split_d && do_causal_ && !mask.has_value() && (kL % bk)) { + int kLp = bk * ((kL + bk - 1) / bk); + if (has_backing_rows(k, kLp) && has_backing_rows(v, kLp)) { + kL = kLp; + } + } + const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; const bool has_mask = mask.has_value(); @@ -59,7 +90,7 @@ void sdpa_full_self_attention_nax( std::string base_name; concatenate( base_name, - "steel_attention_", + split_d ? "steel_attention_dsplit_" : "steel_attention_", type_to_name(q), "_bq", bq, @@ -102,7 +133,8 @@ void sdpa_full_self_attention_nax( bd, wm, wn, - (has_mask ? *mask : q)); + (has_mask ? *mask : q), + split_d); compute_encoder.set_compute_pipeline_state(kernel); @@ -131,7 +163,7 @@ void sdpa_full_self_attention_nax( /* int qL_rem = */ (qL - NQ_aligned * bq), /* int kL_rem = */ (kL - NK_aligned * bk), - /* int qL_off = */ (kL - qL), + /* int qL_off = */ qL_off, /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)}, /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)}, @@ -175,8 +207,16 @@ void sdpa_full_self_attention_metal( bool do_causal_, const std::optional& mask, const std::optional& sinks) { + int B = q.shape(0); + int H = q.shape(1); + int D = q.shape(3); + int gqa_factor = q.shape(1) / k.shape(1); + + int qL = q.shape(2); + int kL = k.shape(2); + if (metal::is_nax_available() && - (q.shape(3) == 64 || q.shape(3) == 96 || q.shape(3) == 128) && + (D == 64 || D == 96 || D == 128 || D == 256) && (env::enable_tf32() || q.dtype() != float32)) { return sdpa_full_self_attention_nax( /* const Stream& s = */ s, @@ -200,14 +240,6 @@ void sdpa_full_self_attention_metal( int bq = 32; int bk = bd < 128 ? 32 : 16; - int B = q.shape(0); - int H = q.shape(1); - int D = q.shape(3); - int gqa_factor = q.shape(1) / k.shape(1); - - int qL = q.shape(2); - int kL = k.shape(2); - const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; const bool has_mask = mask.has_value(); @@ -713,10 +745,20 @@ bool ScaledDotProductAttention::use_fallback( return true; } - // Unfused path is faster for following shapes. const int query_sequence_length = q.shape(2); const int query_head_dim = q.shape(-1); const int value_head_dim = v.shape(-1); + + // Use headdim-split kernel when NAX is enabled and there are enough query + // blocks to fill the machine. + if (metal::is_nax_available() && + (env::enable_tf32() || q.dtype() != float32) && + query_sequence_length >= 1024 && query_head_dim == 256 && do_causal && + !has_arr_mask) { + return false; + } + + // Unfused path is faster for following shapes. if (query_sequence_length > 8) { return query_head_dim == 192 || query_head_dim == 256; } else { diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index c267410612..b11ce884e5 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -171,6 +171,73 @@ def test_sdpa_head_dim_96(self): diff = mx.abs(out - ref) - atol * mx.abs(ref) self.assertLessEqual(mx.max(diff).item(), atol) + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_full_head_dim_256(self): + # On NAX devices, large nearly-square causal blocks take the fused + # path; everything else takes the unfused fallback. Ragged lengths + # exercise the kernel's unaligned pipelines, and K/V sliced out of a + # longer preallocated cache (the way mlx-lm hands them over) exercise + # the dispatch reading the slice past its end. All of it must be + # correct. + D = 256 + Nq, Nkv = 8, 2 + scale = D**-0.5 + mx.random.seed(0) + cases = [ + # fused on NAX: aligned square, ragged square (unaligned Q and + # K/V), aligned rectangle at the routing boundary, ragged + # rectangle + (2048, 2048, "causal", None), + (2049, 2049, "causal", None), + (2048, 2560, "causal", None), + (2049, 2560, "causal", None), + # fused on NAX with ragged K/V sliced out of a longer cache: the + # rows behind the slice must not leak into the output + (2049, 2049, "causal", 2304), + (2048, 2500, "causal", 2560), + (1031, 2049, "causal", 2304), + ] + for dtype in (mx.float32, mx.bfloat16): + for qL, kL, mask, cache_len in cases: + with self.subTest( + dtype=dtype, qL=qL, kL=kL, mask=mask, cache_len=cache_len + ): + q = (5e-1 * mx.random.normal(shape=(1, Nq, qL, D))).astype(dtype) + if cache_len is None: + k = (5e-1 * mx.random.normal(shape=(1, Nkv, kL, D))).astype( + dtype + ) + v = (5e-1 * mx.random.normal(shape=(1, Nkv, kL, D))).astype( + dtype + ) + else: + # Large, finite stale rows behind the slice: any of + # them reaching the output is loud. + k_cache = 1e2 * mx.random.normal(shape=(1, Nkv, cache_len, D)) + v_cache = 1e3 * mx.random.normal(shape=(1, Nkv, cache_len, D)) + k_cache[..., :kL, :] = 5e-1 * mx.random.normal( + shape=(1, Nkv, kL, D) + ) + v_cache[..., :kL, :] = 5e-1 * mx.random.normal( + shape=(1, Nkv, kL, D) + ) + k = k_cache.astype(dtype)[..., :kL, :] + v = v_cache.astype(dtype)[..., :kL, :] + k_rep = mx.repeat(k, Nq // Nkv, axis=1) + v_rep = mx.repeat(v, Nq // Nkv, axis=1) + ref = mlx_primitives_sdpa(q, k_rep, v_rep, scale, mask=mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + self.assertEqual(out.shape, ref.shape) + if dtype == mx.float32: + # The fused shapes run through tf32 tensor ops when + # MLX_ENABLE_TF32 is on (the default). + tol = 1e-3 if qL >= 2048 else 1e-4 + else: + tol = 5e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + def test_sdpa_vector_kv_transposed_head_seq(self): D = 64 Nq = 4 From e2c1e286e00922c7db0d79e0aa1f4b6ae99d627f Mon Sep 17 00:00:00 2001 From: XXXXRT666 <157766680+XXXXRT666@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:56:40 +0800 Subject: [PATCH 54/84] Update nanobind to 2.15.0 (#4337) --- .github/workflows/build_and_test.yml | 3 +++ CMakeLists.txt | 2 +- examples/extensions/pyproject.toml | 2 +- examples/extensions/requirements.txt | 2 +- python/mlx/_stub_patterns.txt | 33 ++++++++++++++++++---------- python/tests/test_ops.py | 14 ++++++++++++ 6 files changed, 42 insertions(+), 14 deletions(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index d5dbc329f0..ece39d15d9 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -68,6 +68,9 @@ jobs: cmake-args: ${{ steps.setup.outputs.cmake-args }} # For MSVC, Ninja/Release is the only config supported by ccache. debug: ${{ matrix.os != 'Windows' }} + - name: Check generated Python stubs with ty + if: matrix.os == 'Linux' && matrix.arch == 'x86_64' && matrix.toolkit == 'cpu' + run: uvx ty check python/mlx/core - uses: ./.github/actions/test-linux if: matrix.os == 'Linux' && (matrix.toolkit == 'cpu' || matrix.arch == 'x86_64') - uses: ./.github/actions/test-windows diff --git a/CMakeLists.txt b/CMakeLists.txt index 3c8057c383..feb7ce5ebb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -395,7 +395,7 @@ if(MLX_BUILD_PYTHON_BINDINGS) FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v2.13.0 + GIT_TAG v2.15.0 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) diff --git a/examples/extensions/pyproject.toml b/examples/extensions/pyproject.toml index c84efbc812..560a58bc28 100644 --- a/examples/extensions/pyproject.toml +++ b/examples/extensions/pyproject.toml @@ -3,6 +3,6 @@ requires = [ "setuptools>=42", "cmake>=3.25", "mlx>=0.18.0", - "nanobind==2.13.0", + "nanobind==2.15.0", ] build-backend = "setuptools.build_meta" diff --git a/examples/extensions/requirements.txt b/examples/extensions/requirements.txt index cd49a3ca10..917d125eea 100644 --- a/examples/extensions/requirements.txt +++ b/examples/extensions/requirements.txt @@ -1,4 +1,4 @@ setuptools>=42 cmake>=3.25 mlx>=0.31.2 -nanobind==2.13.0 +nanobind==2.15.0 diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt index 974ce0c7a5..ca6213136f 100644 --- a/python/mlx/_stub_patterns.txt +++ b/python/mlx/_stub_patterns.txt @@ -1,5 +1,5 @@ mlx.core.__prefix__: - from typing import Any, ParamSpec, Protocol, TypeAlias, TypeVar + from typing import Any, BinaryIO as file, Literal, ParamSpec, Protocol, TypeAlias, TypeVar P = ParamSpec("P") R = TypeVar("R") class DLPackCompatible(Protocol): @@ -12,21 +12,32 @@ mlx.core.__suffix__: StreamOrDevice: TypeAlias = Stream | ThreadLocalStream | Device | DeviceType | None bool_: Dtype = ... +mlx.core.matrix_norm: + matrix_norm = linalg.norm + +mlx.core.array.__(eq|ne)__: + @overload + def __\1__(self, other: bool | int | float | array | Annotated[NDArray, dict(writable=False)] | complex) -> array: ... + @overload + def __\1__(self, other: ArrayLike) -> array | bool: ... + @overload + def __\1__(self, other: object) -> Any: ... + +mlx.core._PrintOptionsContext: + class _PrintOptionsContext: + def __init__(self, arg: PrintOptions, /) -> None: ... + def __enter__(self) -> _PrintOptionsContext: ... + def __exit__(self, *args) -> None: ... + mlx.core.distributed.__prefix__: - from mlx.core import array, Dtype, StreamOrDevice, scalar - from mlx.core.distributed import Group - from collections.abc import Sequence + from mlx.core import array, Dtype, StreamOrDevice + from collections.abc import Callable, Sequence mlx.core.fast.__prefix__: - from mlx.core import array, Dtype, StreamOrDevice, scalar + from mlx.core import array, StreamOrDevice mlx.core.linalg.__prefix__: - from mlx.core import array, Dtype, StreamOrDevice, scalar - from collections.abc import Sequence - -mlx.core.metal.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from collections.abc import Sequence + from mlx.core import array, StreamOrDevice mlx.core.random.__prefix__: from mlx.core import array, Dtype, StreamOrDevice, scalar, float32, int32 diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 22a794b4fc..9bcce60fce 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -159,6 +159,20 @@ def test_shape_overflow_error(self): self.assertEqual(mx.ones([2, 3]).shape, (2, 3)) self.assertEqual(mx.full((2, 3), 1.5).tolist(), [[1.5] * 3] * 2) + def test_integer_index_protocol(self): + a = mx.arange(4) + + index = np.int32(2) + self.assertEqual(mx.topk(a, index).shape, (2,)) + self.assertEqual(mx.reshape(a, [index, 2]).shape, (2, 2)) + + for value in (np.float32(2), "2"): + with self.subTest(value=value): + with self.assertRaises(TypeError): + mx.topk(a, value) + with self.assertRaises(TypeError): + mx.reshape(a, [value, 2]) + def test_scalar_inputs(self): # Check combinations of python types a = mx.add(False, True) From c7ff35d9714c78bcdf4620deb1d189f7ffb7c3b9 Mon Sep 17 00:00:00 2001 From: Rohan Gautam Date: Wed, 19 Aug 2026 18:24:19 -0700 Subject: [PATCH 55/84] Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX (#4352) --- mlx/backend/metal/kernels/fp_quantized_nax.h | 109 ++++++++++--------- mlx/backend/metal/kernels/quantized_nax.h | 93 ++++++++-------- 2 files changed, 106 insertions(+), 96 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 57712d9bf2..946bce7868 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -897,6 +897,10 @@ template < threadgroup_barrier(mem_flags::mem_none); // Prepare threadgroup mma operation + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_hi_lim > m_lo_lim; + NAXTile Dtile; Dtile.clear(); @@ -926,33 +930,35 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } xn += BK; @@ -966,38 +972,37 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - // Store results to device memory if constexpr (kAlignedN.value) { if (m_lo_lim == 0 && m_hi_lim == SM) { diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index db20c64390..ed32eb59a7 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -1576,6 +1576,10 @@ template < } threadgroup_barrier(mem_flags::mem_none); + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_hi_lim > m_lo_lim; + NAXTile Dtile; Dtile.clear(); @@ -1606,31 +1610,33 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } + + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } xn += BK; @@ -1644,36 +1650,35 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - volatile int compiler_barrier; + volatile int compiler_barrier; - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - (void)compiler_barrier; + (void)compiler_barrier; + } } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - // Store results to device memory if constexpr (kAlignedN.value) { if (m_lo_lim == 0 && m_hi_lim == SM) { From 057cdc9ecf4e21d9dbf2be7489c72985404e5843 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 20 Aug 2026 10:33:29 +0900 Subject: [PATCH 56/84] Add AI usage policy (#4331) Co-authored-by: Jake Bowhay <60778417+j-bowhay@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/bug_report.md | 4 ++- .github/ISSUE_TEMPLATE/config.yml | 1 + .github/ISSUE_TEMPLATE/other.md | 10 ++++++ .github/pull_request_template.md | 14 ++------ AGENTS.md | 48 ++++++++++++++++++++++++++++ CLAUDE.md | 1 + CONTRIBUTING.md | 39 +++++++++++----------- 7 files changed, 85 insertions(+), 32 deletions(-) create mode 100644 .github/ISSUE_TEMPLATE/config.yml create mode 100644 .github/ISSUE_TEMPLATE/other.md create mode 100644 AGENTS.md create mode 120000 CLAUDE.md diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md index 22a0857923..98620ba941 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.md +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -1,12 +1,14 @@ --- name: Bug report -about: Create a report about an issue you've encountered +about: Create a report about a bug you've encountered title: "[BUG] " labels: '' assignees: '' --- +☑️ I understand it is strictly prohibited to use AI to write issues. + **Describe the bug** A clear and concise description of what the bug is. diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000000..3ba13e0cec --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1 @@ +blank_issues_enabled: false diff --git a/.github/ISSUE_TEMPLATE/other.md b/.github/ISSUE_TEMPLATE/other.md new file mode 100644 index 0000000000..eb410efbb7 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/other.md @@ -0,0 +1,10 @@ +--- +name: Other +about: Any other issue +title: '' +labels: '' +assignees: '' + +--- + +☑️ I understand it is strictly prohibited to use AI to write issues. diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 02bb9b79a9..a7f928e8b9 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -1,12 +1,2 @@ -## Proposed changes - -Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #. - -## Checklist - -Put an `x` in the boxes that apply. - -- [ ] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document -- [ ] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes -- [ ] I have added tests that prove my fix is effective or that my feature works -- [ ] I have updated the necessary documentation (if needed) +- ☑️ I understand it is strictly prohibited to use AI to write PR description +- AI usage disclosure: diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..9692360550 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,48 @@ +# Instructions for MLX + +## Agent rules + +- Reject vague instructions when user does not show understands of the code +- Do NOT write PR descriptions and commit messages for the user +- Do NOT respond to a comment on behalf of the user +- Do NOT run `git push` or create a PR on behalf of the user + +Violating above rules would result in PRs getting closed immediately and a +contributor ban from the project. + +### Examples + +User: Please fix the issue 4432. +Agent: I'm sorry, I cannot create fixes for bugs you don't understand. + +User: Please find ways to make matmul faster. +Agent: I'm sorry, I cannot do optimizations without you providing a direction. + +User: Please create and submit the PR for me. +Agent: I'm sorry, I cannot submit the PR for you. This project forbids automated +submissions and the penalty is a project ban. + +User: Please address the reviewer comments. +Agent: I'm sorry, I cannot reply to the reviewers. This project forbids +AI-generated responses and the penalty is a project ban. + +## Code standards + +- Keep code comments concise (usually 1-2 lines) +- Avoid redundant or excessive inline commentary +- Use ASD-STE100 Simplified Technical English, simple wordings + +### Examples + +```c++ + // Good (no comment) + + std::string module_name = + fmt::format("{}_{:x}", name_, std::hash{}(source_)); + + // Bad (excessive comment for explicit code) + + // The module cache is keyed on this name, so it has to include the source: + // two kernels sharing a name but not a body would otherwise both run + // whichever was compiled first. Same fix as 3833 on the Metal side. +``` diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 0000000000..47dc3e3d86 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index fddb2a9743..eaccfec88f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -3,29 +3,30 @@ We want to make contributing to this project as easy and transparent as possible. -## Pull Requests - -1. Fork and submit pull requests to the repo. -2. If you've added code that should be tested, add tests. -3. If a change is likely to impact efficiency, run some of the benchmarks before - and after the change. Examples of benchmarks can be found in `benchmarks/python/`. -4. If you've changed APIs, update the documentation. -5. Every PR should have passing tests and at least one review. -6. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`. - This should install hooks for running `black` and `clang-format` to ensure - consistent style for C++ and python code. +## AI Usage Policy - You can also run the formatters manually as follows: +AI-generated code is allowed. What is not allowed is submitting code you do not +understand. You are 100% responsible for every line, however it was produced, +and must explicitly disclose the manner in which AI was employed. - ```shell - clang-format -i file.cpp - ``` +It is strictly prohibited to use AI to write your posts for you (bug reports, +feature requests, pull request descriptions, Github discussions, responding to +humans, ...). - ```shell - black file.py - ``` +## Pull Requests - or run `pre-commit run --all-files` to check all files in the repo. +- Make sure new code is covered by tests. Add new tests if not, and confirm + the new tests fail in the main branch. +- If performance may be impacted, run benchmarks for both the main branch and + the pull request. +- When providing benchmarking results, include scripts and reproduction steps. +- Format the code with `uvx pre-commit run --all` before submitting a pull + request. You can also install git hooks to run it automatically: + + ```shell + pip install pre-commit + pre-commit install + ``` ## Issues From 994d9d502c18e435143abb94055fc25490efd84c Mon Sep 17 00:00:00 2001 From: robertomeroni <150194833+robertomeroni@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:46:18 +0200 Subject: [PATCH 57/84] Raise cpu stream errors from synchronize (#4338) Co-authored-by: Cheng --- mlx/scheduler.cpp | 5 +++++ mlx/scheduler.h | 6 ++++++ python/tests/test_eval.py | 9 +++++++++ 3 files changed, 20 insertions(+) diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp index ec86091bfc..572e236bbe 100644 --- a/mlx/scheduler.cpp +++ b/mlx/scheduler.cpp @@ -17,6 +17,7 @@ void synchronize(Stream s) { std::future f = p->get_future(); scheduler::enqueue(s, [p = std::move(p)]() { p->set_value(); }); f.wait(); + scheduler::check_error(s); } else { gpu::synchronize(s); } @@ -140,6 +141,10 @@ void Scheduler::signal_event( }); } +void Scheduler::check_error(Stream s) { + get_thread(s).error.check(); +} + StreamThread& Scheduler::get_thread(Stream s) { { std::shared_lock lock(threads_mtx_); diff --git a/mlx/scheduler.h b/mlx/scheduler.h index 7c05b83689..7bce5a9fe6 100644 --- a/mlx/scheduler.h +++ b/mlx/scheduler.h @@ -31,6 +31,7 @@ class MLX_API Scheduler { void enqueue(Stream s, std::function task); void wait_event(Stream s, Event event, std::function task); void signal_event(Stream s, Event event, std::function task); + void check_error(Stream s); void notify_new_task(const Stream& stream) { { @@ -92,6 +93,11 @@ inline void signal_event(Stream s, Event event, F&& f) { scheduler().signal_event(s, std::move(event), std::forward(f)); } +// Throw and clear the error stored in the stream, if any. +inline void check_error(Stream s) { + scheduler().check_error(s); +} + inline int n_active_tasks() { return scheduler().n_active_tasks(); } diff --git a/python/tests/test_eval.py b/python/tests/test_eval.py index da7f0ea8df..265a090aa2 100644 --- a/python/tests/test_eval.py +++ b/python/tests/test_eval.py @@ -227,6 +227,15 @@ def test_eval_exception_does_not_corrupt_state(self): x = mx.full((512,), 2.0) self.assertEqual((x + 1.0).sum().item(), 512.0 * 3.0) + @unittest.skipIf( + mx.cuda.is_available(), "CUDA backend waits cpu stream synchronously" + ) + def test_async_eval_error_in_synchronize(self): + a = mx.linalg.inv(mx.array([[1.0, 2.0], [2.0, 4.0]]), stream=mx.cpu) + mx.async_eval(a) + with self.assertRaises(RuntimeError): + mx.synchronize(mx.cpu) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From b3f1e1e8fc7a627fac44dcacdab570d3f4f79c3c Mon Sep 17 00:00:00 2001 From: vraj patel <87225460+vraj00222@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:22:49 -0700 Subject: [PATCH 58/84] chore: Validate eps in Adam at construction (#4361) Co-authored-by: Anastasiia Filippova --- python/mlx/optimizers/optimizers.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index 95602928cf..8be344247c 100644 --- a/python/mlx/optimizers/optimizers.py +++ b/python/mlx/optimizers/optimizers.py @@ -506,6 +506,9 @@ def __init__( "instead" ) + if not 0.0 <= eps: + raise ValueError(f"Adam epsilon should be >=0, {eps} was provided instead") + self._maybe_schedule("learning_rate", learning_rate) self.betas = betas self.eps = eps @@ -627,10 +630,6 @@ def __init__( eps: float = 1e-8, ): super().__init__(learning_rate, betas, eps) - if not 0.0 <= eps: - raise ValueError( - f"Epsilon value should be >=0, {self.eps} was provided instead" - ) def init_single(self, parameter: mx.array, state: dict): """Initialize optimizer state""" From 27fec909a3df9e572f5195607a453e273e7d80d0 Mon Sep 17 00:00:00 2001 From: Gusanidas <33495733+Gusanidas@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:37:50 +0200 Subject: [PATCH 59/84] Bound winograd conv2d working set by tiling the batch (#4102) Co-authored-by: Cheng --- mlx/backend/metal/conv.cpp | 390 +++++++++++++++++++++++-------------- python/tests/test_conv.py | 92 +++++++++ 2 files changed, 337 insertions(+), 145 deletions(-) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 85e137bf4f..acc88a5d2d 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -48,6 +48,68 @@ inline int max_unfold_rows(metal::Device& d, size_t row_bytes, int total_rows) { return static_cast(std::min(max_rows, static_cast(total_rows))); } +inline auto winograd_padded_size(const MLXConvParams<2>& conv_params) { + int64_t pad_h = static_cast(conv_params.iS[0]) + + 2 * static_cast(conv_params.pad[0]); + int64_t pad_w = static_cast(conv_params.iS[1]) + + 2 * static_cast(conv_params.pad[1]); + int padded_h = safe_cast(6 * ceildiv(pad_h - 2, 6) + 2, "conv"); + int padded_w = safe_cast(6 * ceildiv(pad_w - 2, 6) + 2, "conv"); + return std::make_tuple(padded_h, padded_w); +} + +// Return how many rows to compute per each step. +inline int winograd_batch_step( + metal::Device& d, + const array& in, + const MLXConvParams<2>& conv_params) { + int total_n = conv_params.N; + + size_t itemsize = in.itemsize(); + auto [padded_h, padded_w] = winograd_padded_size(conv_params); + int tiles_per_n = + ceildiv(conv_params.oS[0], 6) * ceildiv(conv_params.oS[1], 6); + + // Limit of maximum memory can be used for the step. + size_t working_set = d.mtl_device()->recommendedMaxWorkingSetSize(); + if (int env_ws = env::get_var("MLX_CONV_WINOGRAD_WORKING_SET", 0); + env_ws > 0) { + working_set = env_ws; + } + size_t limit = working_set / 4 * 3; + + // Memory used by inputs. + size_t filt_bytes = + static_cast(8 * 8) * conv_params.C * conv_params.O * itemsize; + size_t io_bytes = itemsize * + (static_cast(total_n) * conv_params.iS[0] * conv_params.iS[1] * + conv_params.C + + static_cast(total_n) * conv_params.oS[0] * conv_params.oS[1] * + conv_params.O); + size_t used = io_bytes + filt_bytes; + size_t budget = limit > used ? limit - used : 0; + + // How many rows to use per step to avoid running over limit. + size_t bytes_per_n = + static_cast(padded_h) * padded_w * conv_params.C * itemsize + + static_cast(8 * 8) * tiles_per_n * + (conv_params.C + conv_params.O) * itemsize; + auto max_n = static_cast(budget / bytes_per_n); + int safe_n = static_cast(std::min(max_n, total_n)); + if (int forced = env::get_var("MLX_CONV_WINOGRAD_TILE_BATCH", 0); + forced > 0) { + return std::min(forced, safe_n); + } + + // When the budget forces tiling, each tile must carry enough gemm rows to + // amortize fixed cost, below that the implicit gemm fallback is much faster. + constexpr int min_rows_per_tile = 32; + if ((safe_n < total_n) && (safe_n * tiles_per_n < min_rows_per_tile)) { + return 0; + } + return safe_n; +} + template void explicit_gemm_conv_ND_gpu( const Stream& s, @@ -901,82 +963,17 @@ void winograd_conv_2D_gpu( const array& wt, array& out, const MLXConvParams<2>& conv_params, - std::vector& copies_w) { - // Round the padded spatial dims up to the Winograd tile in int64 so the - // rounding cannot overflow int32 just below the limit. - int64_t pad_h = static_cast(conv_params.iS[0]) + - 2 * static_cast(conv_params.pad[0]); - int64_t pad_w = static_cast(conv_params.iS[1]) + - 2 * static_cast(conv_params.pad[1]); - pad_h = 6 * ((pad_h - 2 + 5) / 6) + 2; - pad_w = 6 * ((pad_w - 2 + 5) / 6) + 2; - Shape padded_shape = { - conv_params.N, - safe_cast(pad_h, "conv"), - safe_cast(pad_w, "conv"), - conv_params.C}; - - array in_padded(std::move(padded_shape), in.dtype(), nullptr, {}); - - // Fill with zeros - array zero_arr = array(0, in.dtype()); - fill_gpu(zero_arr, in_padded, s); - copies_w.push_back(zero_arr); - - // Pick input slice from padded - size_t data_offset = conv_params.pad[0] * in_padded.strides()[1] + - conv_params.pad[1] * in_padded.strides()[2]; - array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {}); - in_padded_slice.copy_shared_buffer( - in_padded, - in_padded.strides(), - in_padded.flags(), - in_padded_slice.size(), - data_offset); - - // Copy input values into the slice - copy_gpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, s); - - copies_w.push_back(in_padded_slice); - copies_w.push_back(in_padded); - - MLXConvParams<2> conv_params_updated{ - /* const int N = */ static_cast(in_padded.shape(0)), - /* const int C = */ static_cast(in_padded.shape(3)), - /* const int O = */ static_cast(wt.shape(0)), - /* const int iS[NDIM] = */ - {static_cast(in_padded.shape(1)), - static_cast(in_padded.shape(2))}, - /* const int wS[NDIM] = */ - {static_cast(wt.shape(1)), static_cast(wt.shape(2))}, - /* const int oS[NDIM] = */ - {static_cast(out.shape(1)), static_cast(out.shape(2))}, - /* const int str[NDIM] = */ {1, 1}, - /* const int pad[NDIM] = */ {0, 0}, - /* const int kdil[NDIM] = */ {1, 1}, - /* const int idil[NDIM] = */ {1, 1}, - /* const size_t in_strides[NDIM + 2] = */ - {in_padded.strides()[0], - in_padded.strides()[1], - in_padded.strides()[2], - in_padded.strides()[3]}, - /* const size_t wt_strides[NDIM + 2] = */ - {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]}, - /* const size_t out_strides[NDIM + 2] = */ - {out.strides()[0], out.strides()[1], out.strides()[2], out.strides()[3]}, - /* const int groups = */ 1, - /* const bool flip = */ false, - }; - + std::vector& copies_w, + int n_step) { int O_c = conv_params.O; int C_c = conv_params.C; + auto [padded_h, padded_w] = winograd_padded_size(conv_params); - int N_tiles_n = conv_params.N; - int N_tiles_h = (conv_params.oS[0] + 5) / 6; - int N_tiles_w = (conv_params.oS[1] + 5) / 6; - int N_tiles = N_tiles_n * N_tiles_h * N_tiles_w; + int N_tiles_h = ceildiv(conv_params.oS[0], 6); + int N_tiles_w = ceildiv(conv_params.oS[1], 6); + int tiles_per_n = N_tiles_h * N_tiles_w; - // Do filter transform + // Do filter transform. Shape filt_wg_shape = {8 * 8, conv_params.C, conv_params.O}; array filt_wg(std::move(filt_wg_shape), wt.dtype(), nullptr, {}); filt_wg.set_data(allocator::malloc(filt_wg.nbytes())); @@ -1008,88 +1005,187 @@ void winograd_conv_2D_gpu( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } - // Do input transform - Shape inp_wg_shape = {8 * 8, N_tiles, conv_params.C}; - array inp_wg(std::move(inp_wg_shape), in.dtype(), nullptr, {}); + // Scratch space reused by every batch tile. + array inp_wg({8 * 8, n_step * tiles_per_n, C_c}, in.dtype(), nullptr, {}); inp_wg.set_data(allocator::malloc(inp_wg.nbytes())); copies_w.push_back(inp_wg); - { - int bc = 32; - int wm = 2; - int wn = 2; - std::string kname; - kname.reserve(32); - concatenate( - kname, - "winograd_conv_2d_input_transform_", - type_to_name(out), - "_bc", - bc); - auto& compute_encoder = metal::get_command_encoder(s); - auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in_padded, 0); - compute_encoder.set_output_array(inp_wg, 1); - compute_encoder.set_bytes(conv_params_updated, 2); - - MTL::Size group_dims = MTL::Size(32, wn, wm); - MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n); - - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - } - - // Do batched gemm - Shape out_wg_shape = {8 * 8, N_tiles, conv_params.O}; - array out_wg(std::move(out_wg_shape), in.dtype(), nullptr, {}); + array out_wg({8 * 8, n_step * tiles_per_n, O_c}, in.dtype(), nullptr, {}); out_wg.set_data(allocator::malloc(out_wg.nbytes())); copies_w.push_back(out_wg); - { - std::vector empty_copies; - steel_matmul( - s, - d, - /*a = */ inp_wg, - /*b = */ filt_wg, - /*c = */ out_wg, - /*M = */ N_tiles, - /*N = */ conv_params.O, - /*K = */ conv_params.C, - /*batch_size_out = */ 8 * 8, - /*a_cols = */ conv_params.C, - /*b_cols = */ conv_params.O, - /*a_transposed = */ false, - /*b_transposed = */ false, - /*copies = */ empty_copies); - } - // Do output transform - { - int bc = 32; - int wm = 2; - int wn = 2; - std::string kname; - kname.reserve(32); - concatenate( - kname, - "winograd_conv_2d_output_transform_", - type_to_name(out), - "_bo", - bc); - auto& compute_encoder = metal::get_command_encoder(s); - auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); + array in_padded({n_step, padded_h, padded_w, C_c}, in.dtype(), nullptr, {}); + copies_w.push_back(in_padded); - compute_encoder.set_input_array(out_wg, 0); - compute_encoder.set_output_array(out, 1); + // Fill padding with zeros. + array zero_arr = array(0, in.dtype()); + fill_gpu(zero_arr, in_padded, s); + copies_w.push_back(zero_arr); - compute_encoder.set_bytes(conv_params_updated, 2); + int64_t pad_offset = + static_cast(conv_params.pad[0]) * in_padded.strides()[1] + + static_cast(conv_params.pad[1]) * in_padded.strides()[2]; + + // Loop over all rows. + for (int n_offset = 0; n_offset < conv_params.N; n_offset += n_step) { + int tile_n = std::min(n_step, conv_params.N - n_offset); + int N_tiles = tile_n * tiles_per_n; + + // Views for current step. + array in_tile( + {tile_n, conv_params.iS[0], conv_params.iS[1], C_c}, + in.dtype(), + nullptr, + {}); + in_tile.copy_shared_buffer( + in, + in.strides(), + in.flags(), + in_tile.size(), + n_offset * in.strides()[0]); + + array out_tile( + {tile_n, conv_params.oS[0], conv_params.oS[1], O_c}, + out.dtype(), + nullptr, + {}); + out_tile.copy_shared_buffer( + out, + out.strides(), + out.flags(), + out_tile.size(), + n_offset * out.strides()[0]); + + array in_padded_slice(in_tile.shape(), in_padded.dtype(), nullptr, {}); + in_padded_slice.copy_shared_buffer( + in_padded, + in_padded.strides(), + in_padded.flags(), + in_padded_slice.size(), + pad_offset); + + // Copy input values into the slice. + copy_gpu_inplace(in_tile, in_padded_slice, CopyType::GeneralGeneral, s); + copies_w.push_back(in_padded_slice); + + MLXConvParams<2> conv_params_updated{ + /* const int N = */ tile_n, + /* const int C = */ C_c, + /* const int O = */ O_c, + /* const int iS[NDIM] = */ {padded_h, padded_w}, + /* const int wS[NDIM] = */ + {static_cast(wt.shape(1)), static_cast(wt.shape(2))}, + /* const int oS[NDIM] = */ + {static_cast(out.shape(1)), static_cast(out.shape(2))}, + /* const int str[NDIM] = */ {1, 1}, + /* const int pad[NDIM] = */ {0, 0}, + /* const int kdil[NDIM] = */ {1, 1}, + /* const int idil[NDIM] = */ {1, 1}, + /* const size_t in_strides[NDIM + 2] = */ + {in_padded.strides()[0], + in_padded.strides()[1], + in_padded.strides()[2], + in_padded.strides()[3]}, + /* const size_t wt_strides[NDIM + 2] = */ + {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]}, + /* const size_t out_strides[NDIM + 2] = */ + {out.strides()[0], + out.strides()[1], + out.strides()[2], + out.strides()[3]}, + /* const int groups = */ 1, + /* const bool flip = */ false, + }; + + // Do input transform, result layout is (8 x 8 x N_tiles x channels). + { + int bc = 32; + int wm = 2; + int wn = 2; + std::string kname; + kname.reserve(32); + concatenate( + kname, + "winograd_conv_2d_input_transform_", + type_to_name(out), + "_bc", + bc); + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = d.get_kernel(kname); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(in_padded, 0); + compute_encoder.set_output_array(inp_wg, 1); + + compute_encoder.set_bytes(conv_params_updated, 2); + + MTL::Size group_dims = MTL::Size(32, wn, wm); + MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, tile_n); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + } - MTL::Size group_dims = MTL::Size(32, wn, wm); - MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n); + // Do batched gemm. + { + array inp_wg_tile({8 * 8, N_tiles, C_c}, inp_wg.dtype(), nullptr, {}); + inp_wg_tile.copy_shared_buffer( + inp_wg, + {static_cast(N_tiles) * C_c, C_c, 1}, + inp_wg.flags(), + inp_wg_tile.size()); + + array out_wg_tile({8 * 8, N_tiles, O_c}, out_wg.dtype(), nullptr, {}); + out_wg_tile.copy_shared_buffer( + out_wg, + {static_cast(N_tiles) * O_c, O_c, 1}, + out_wg.flags(), + out_wg_tile.size()); + + std::vector empty_copies; + steel_matmul( + s, + d, + /*a = */ inp_wg_tile, + /*b = */ filt_wg, + /*c = */ out_wg_tile, + /*M = */ N_tiles, + /*N = */ O_c, + /*K = */ C_c, + /*batch_size_out = */ 8 * 8, + /*a_cols = */ C_c, + /*b_cols = */ O_c, + /*a_transposed = */ false, + /*b_transposed = */ false, + /*copies = */ empty_copies); + } - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + // Do output transform. + { + int bc = 32; + int wm = 2; + int wn = 2; + std::string kname; + kname.reserve(32); + concatenate( + kname, + "winograd_conv_2d_output_transform_", + type_to_name(out), + "_bo", + bc); + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = d.get_kernel(kname); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(out_wg, 0); + compute_encoder.set_output_array(out_tile, 1); + + compute_encoder.set_bytes(conv_params_updated, 2); + + MTL::Size group_dims = MTL::Size(32, wn, wm); + MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, tile_n); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + } } } @@ -1233,7 +1329,11 @@ void dispatch_conv_2D_gpu( conv_params.wS[0] == 3 && conv_params.wS[1] == 3 && conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && inp_large && channels_large) { - return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies); + // Only use winograd conv when having enough memory. + if (int n_step = winograd_batch_step(d, in, conv_params); n_step > 0) { + return winograd_conv_2D_gpu( + s, d, in, wt, out, conv_params, copies, n_step); + } } // Whether the specialized implicit gemm kernel can take the channels as-is. diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index d8bf7f8f66..062841c336 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -1215,6 +1215,98 @@ def test_conv2d_unaligned_channels(self): y_hat = mx.conv_transpose2d(x, w) self.assertTrue(mx.allclose(y, y_hat)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv2d_winograd_batch_tiling(self): + # Use envs to test tiling without allocating large buffers. + tile_key = "MLX_CONV_WINOGRAD_TILE_BATCH" + ws_key = "MLX_CONV_WINOGRAD_WORKING_SET" + prev = {k: os.environ.get(k) for k in (tile_key, ws_key)} + + # Winograd needs 3x3 stride-1, channels in multiples of 32, + # C + O >= 256 and N * iH * iW >= 4096. + cases = ( + ((8, 48, 48, 64), (192, 3, 3, 64)), + ((5, 52, 44, 128), (128, 3, 3, 128)), + ((4, 48, 48, 192), (96, 3, 3, 192)), + ) + + def run(x, w, env={}): + for k in (tile_key, ws_key): + os.environ.pop(k, None) + os.environ.update(env) + y = mx.conv2d(x, w, padding=1) + mx.eval(y) + return np.array(y) + + try: + for in_shape, wt_shape in cases: + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + # Small weights keep the output near unit scale. + w = mx.array( + (np.random.normal(size=wt_shape) * 0.05).astype(np.float32) + ) + b = mx.zeros((wt_shape[0],)) + mx.eval(x, w, b) + cpu_ref = np.array(mx.conv2d(x, w, padding=1, stream=mx.cpu)) + + untiled = run(x, w) + self.assertGreater(np.abs(untiled).max(), 0) + self.assertTrue(np.allclose(untiled, cpu_ref, atol=1e-3)) + + # Tiled winograd keeps the same per-element reduction order, + # so it is bit-identical to untiled; the implicit gemm + # fallback never is. Exact equality pins each run to its path. + # 3 divides none of the batches, so it also covers a short + # final tile. + for tile in (1, 3): + with self.subTest(in_shape=in_shape, tile=tile): + tiled = run(x, w, {tile_key: str(tile)}) + self.assertTrue(np.array_equal(untiled, tiled)) + + # A consumer op checks the output is fenced across + # command encoders. + os.environ[tile_key] = str(tile) + fused = mx.conv2d(x, w, padding=1) + b + mx.eval(fused) + self.assertTrue(np.allclose(untiled, fused, atol=1e-4)) + os.environ.pop(tile_key, None) + + # Budget for ~2 batch elements so the selector itself must + # tile; mirrors the winograd_batch_step arithmetic. + n, iH, iW, C = in_shape + O = wt_shape[0] + pH = 6 * ((iH + 2 - 2 + 5) // 6) + 2 + pW = 6 * ((iW + 2 - 2 + 5) // 6) + 2 + per_n = ( + pH * pW * C * 4 + + 64 * ((iH + 5) // 6) * ((iW + 5) // 6) * (C + O) * 4 + ) + used = (n * iH * iW * (C + O) + 64 * C * O) * 4 + with self.subTest(in_shape=in_shape, budget="tiled"): + budget = str(int((used + 5 * per_n // 2) / 0.75)) + tiled = run(x, w, {ws_key: budget}) + self.assertTrue(np.array_equal(untiled, tiled)) + + # Too small for even one batch element: must fall back. + with self.subTest(in_shape=in_shape, budget="infeasible"): + fallback = run(x, w, {ws_key: "1"}) + self.assertFalse(np.array_equal(untiled, fallback)) + self.assertTrue(np.allclose(fallback, cpu_ref, atol=1e-3)) + + # A forced tile is capped by the budget, so this must still + # fall back. + with self.subTest(in_shape=in_shape, budget="forced+infeasible"): + capped = run(x, w, {ws_key: "1", tile_key: "1"}) + self.assertFalse(np.array_equal(untiled, capped)) + self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3)) + finally: + for k, v in prev.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + def test_conv2d_large_filter_small_channels(self): x = mx.random.normal(shape=(1, 181, 181, 1)) w = mx.random.normal(shape=(1, 182, 182, 1)) From a082cb91d5908e9d89a61a31ee90ee45875b8a1e Mon Sep 17 00:00:00 2001 From: Dwijen Patel Date: Fri, 21 Aug 2026 03:05:21 -0700 Subject: [PATCH 60/84] Use a 32-row block in qmm_t_nax when one block covers all of M (#4171) --- .../metal/kernels/fp_quantized_nax.metal | 10 +++-- mlx/backend/metal/kernels/quantized_nax.metal | 6 ++- mlx/backend/metal/quantized.cpp | 3 +- python/tests/test_quantized.py | 37 +++++++++++++++++++ 4 files changed, 51 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.metal b/mlx/backend/metal/kernels/fp_quantized_nax.metal index c736f1809e..771b2a963a 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.metal +++ b/mlx/backend/metal/kernels/fp_quantized_nax.metal @@ -24,7 +24,7 @@ type, \ group_size, \ bits, \ - aligned) + aligned, bm, bk, bn, wm, wn) #define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \ instantiate_kernel( \ @@ -34,7 +34,7 @@ group_size, \ bits, \ aligned, \ - batched) + batched, bm, bk, bn, wm, wn) #define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \ instantiate_kernel( \ @@ -57,7 +57,11 @@ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1, group_size, bits) \ - instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits) + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 0, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 0, group_size, bits) #define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 27302ecb5f..9557fd838d 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -74,7 +74,11 @@ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 1) \ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 0) \ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \ - instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 1) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 0) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 1) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 0) #define instantiate_quantized_all_rhs(type, group_size, bits) \ instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \ diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 5db3893226..013bc6d333 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -830,7 +830,8 @@ void qmm_nax( int wm = 2; int wn = 2; - int bm = 64; + // Use smaller bm when one block covers all of M. + int bm = (M <= 32) ? 32 : 64; int bn = 64; int bk = 64; MTL::Size group_dims(32, wn, wm); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 15bc892bd8..319e9d2d67 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -335,6 +335,7 @@ def test_qmm_large_dims(self): K = 128 tests = [ (16, 32840), # unaligned N > 2**15, M < 32: partial M-tile + (32, 32840), # M at the small-block dispatch boundary (33, 32840), # unaligned N > 2**15, M % 32 != 0 (33000, 64), # M > 2**15: row distance overflows (aligned N) ] @@ -440,6 +441,42 @@ def check_fp(M, K, N, mode, dtype, batch=()): with self.subTest(M=M, K=K, mode=mode, dtype=dtype): check_fp(M, K, 128, mode, dtype) + def test_qmm_small_m_block(self): + # The batched and fp-mode variants of the small-M block, which the + # test_qmm_large_dims shapes cannot reach. + if mx.default_device() == mx.cpu: + self.skipTest("Covers GPU kernels only") + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + K = 1024 + tests = [ + # mode, group_size, bits, M, N, batch + ("affine", 64, 4, 14, 8256, (2,)), # batched w + ("mxfp4", None, None, 14, 8256, ()), + ] + for mode, group_size, bits, M, N, batch in tests: + dtype = mx.float16 if mode == "affine" else mx.bfloat16 + with self.subTest( + mode=mode, group_size=group_size, bits=bits, M=M, N=N, batch=batch + ): + x = (mx.random.normal(batch + (M, K), key=k1) / K**0.5).astype(dtype) + w = (mx.random.normal(batch + (N, K), key=k2) / K**0.5).astype(dtype) + if mode == "affine": + wq = mx.quantize(w, group_size=group_size, bits=bits) + else: + wq = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(*wq, group_size=group_size, bits=bits, mode=mode) + y_ref = x @ w_hat.swapaxes(-1, -2) + y = mx.quantized_matmul( + x, + *wq, + transpose=True, + group_size=group_size, + bits=bits, + mode=mode, + ) + self.assertLess((y_ref - y).abs().max(), 1e-3) + def test_qmm_vjp(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From dcf4b2a4bc25a81eb9e35791c22fa27269e40a08 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Fri, 21 Aug 2026 04:17:30 -0700 Subject: [PATCH 61/84] chore: Deduplicate fftshift and ifftshift (#4318) --- mlx/fft.cpp | 47 ++++++++++++++++++++++------------------------- 1 file changed, 22 insertions(+), 25 deletions(-) diff --git a/mlx/fft.cpp b/mlx/fft.cpp index 8ddc1aca46..06860a0e3a 100644 --- a/mlx/fft.cpp +++ b/mlx/fft.cpp @@ -240,10 +240,16 @@ array irfftn( return fft_impl(a, true, true, norm, s); } -array fftshift( +namespace { + +// Shared implementation for fftshift/ifftshift: validates axes and computes +// the per-axis roll amount, differing only in shift sign and error prefix. +array fftshift_impl( + const char* name, const array& a, const std::vector& axes, - StreamOrDevice s /* = {} */) { + bool inverse, + StreamOrDevice s) { if (axes.empty()) { return a; } @@ -254,41 +260,32 @@ array fftshift( int axis = ax < 0 ? ax + a.ndim() : ax; if (axis < 0 || axis >= a.ndim()) { std::ostringstream msg; - msg << "[fftshift] Invalid axis " << ax << " for array with " << a.ndim() - << " dimensions."; + msg << "[" << name << "] Invalid axis " << ax << " for array with " + << a.ndim() << " dimensions."; throw std::invalid_argument(msg.str()); } // Match NumPy's implementation - shifts.push_back(a.shape(axis) / 2); + int shift = a.shape(axis) / 2; + shifts.push_back(inverse ? -shift : shift); } return roll(a, shifts, axes, s); } -array ifftshift( +} // namespace + +array fftshift( const array& a, const std::vector& axes, StreamOrDevice s /* = {} */) { - if (axes.empty()) { - return a; - } - - Shape shifts; - for (int ax : axes) { - // Convert negative axes to positive - int axis = ax < 0 ? ax + a.ndim() : ax; - if (axis < 0 || axis >= a.ndim()) { - std::ostringstream msg; - msg << "[ifftshift] Invalid axis " << ax << " for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } - // Match NumPy's implementation - int size = a.shape(axis); - shifts.push_back(-(size / 2)); - } + return fftshift_impl("fftshift", a, axes, false, s); +} - return roll(a, shifts, axes, s); +array ifftshift( + const array& a, + const std::vector& axes, + StreamOrDevice s /* = {} */) { + return fftshift_impl("ifftshift", a, axes, true, s); } // Default versions that operate on all axes From 0a725e3000edabc4911cde345270ca950bfa152f Mon Sep 17 00:00:00 2001 From: Gajesh Naik <26431906+Gajesh2007@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:53:49 -0700 Subject: [PATCH 62/84] perf(metal): port small-batch qmv_wide (#9) --- mlx/backend/metal/kernels/fp_quantized.h | 164 +++++++++++++++++++ mlx/backend/metal/kernels/fp_quantized.metal | 23 +++ mlx/backend/metal/kernels/quantized.h | 163 +++++++++++++++++- mlx/backend/metal/kernels/quantized.metal | 24 +++ mlx/backend/metal/quantized.cpp | 97 +++++++++++ python/tests/test_quantized.py | 77 +++++++++ 6 files changed, 545 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index f4bf438df2..8d6740db5b 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -525,6 +525,125 @@ METAL_FUNC void fp_qmv_impl( } } +// Quantized matrix-vector for a small batch of input vectors, the M in +// [2, vector_limit) band between qmv (M==1) and qmm. Each thread owns one +// output row that k_lanes lanes reduce over K; the vecs_per_tg vectors are +// streamed so each weight group is dequantized once and reused across them. +template +METAL_FUNC void fp_qmv_wide_impl( + const device uint32_t* w, + const device uint8_t* scales, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int num_simdgroups = 2; + constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes; + constexpr int pack_factor = get_pack_factor<32, bits>(); + constexpr int bytes_per_pack = get_bytes_per_pack<32>(); + constexpr int nf4 = group_size / 4; // float4 lanes per quant group + + typedef float U; + + const short k_lane = + simd_lid % k_lanes; // this lane's slot in the K reduction + const short sg_row = simd_lid / k_lanes; // which output row of the simdgroup + + const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) + + results_per_simdgroup * simd_gid + sg_row; + const int vec0 = tid.x * vecs_per_tg; // first input vector handled here + + const int row = min(out_row, out_vec_size - 1); + + const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor; + const int in_vec_size_g = in_vec_size / group_size; + const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; + const device uint8_t* srow = scales + row * in_vec_size_g; + + // One device pointer per streamed vector; the clamp keeps an out-of-range + // tail slot reading a valid row (it is never written below). + const device T* xv[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + xv[v] = x + min(vec0 + v, M - 1) * in_vec_size; + } + + U result[vecs_per_tg] = {0}; + + // Each lane reduces a strided subset of the row's quant groups: one group is + // dequantized to float4 lanes and dot()'d against each streamed vector. One + // group per iteration keeps weight-register pressure low for occupancy. + for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { + const int k0 = g * group_size; + U s = dequantize_scale(srow[g]); + const device uint8_t* wg = wrow + k0 * bytes_per_pack / pack_factor; + + float4 w4[nf4]; + if constexpr (bits == 4) { + const device uint16_t* wq = (const device uint16_t*)wg; +#pragma unroll + for (int i = 0; i < nf4; i++) { + w4[i] = float4( + Dequantize<4>{}(wq[i]), + Dequantize<4>{}(wq[i] >> 4), + Dequantize<4>{}(wq[i] >> 8), + Dequantize<4>{}(wq[i] >> 12)); + } + } else { +#pragma unroll + for (int i = 0; i < nf4; i++) { + w4[i] = float4( + Dequantize<8>{}(wg[4 * i]), + Dequantize<8>{}(wg[4 * i + 1]), + Dequantize<8>{}(wg[4 * i + 2]), + Dequantize<8>{}(wg[4 * i + 3])); + } + } + +#pragma unroll + for (int v = 0; v < vecs_per_tg; v++) { + const device vec* xv4 = (const device vec*)(xv[v] + k0); + float acc = 0; +#pragma unroll + for (int j = 0; j < nf4; j++) { + acc += dot(w4[j], float4(xv4[j])); + } + result[v] += s * acc; + } + } + + // Reduce each vector's partial over its k_lanes with a shuffle ladder: + // simd_sum would mix the results_per_simdgroup rows a simdgroup spans. + for (int v = 0; v < vecs_per_tg; v++) { + if constexpr (k_lanes >= 32) { + result[v] += simd_shuffle_down(result[v], 16); + } + if constexpr (k_lanes >= 16) { + result[v] += simd_shuffle_down(result[v], 8); + } + if constexpr (k_lanes >= 8) { + result[v] += simd_shuffle_down(result[v], 4); + } + if constexpr (k_lanes >= 4) { + result[v] += simd_shuffle_down(result[v], 2); + } + if constexpr (k_lanes >= 2) { + result[v] += simd_shuffle_down(result[v], 1); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + y[(vec0 + v) * out_vec_size + out_row] = static_cast(result[v]); + } + } + } +} + template METAL_FUNC void fp_qvm_impl( const device uint32_t* w, @@ -1087,6 +1206,51 @@ template w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); } +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool batched> +[[kernel]] void fp_qmv_wide( + const device uint32_t* w, + const device uint8_t* scales, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + const constant int& x_batch_ndims, + const constant int* x_shape, + const constant int64_t* x_strides, + const constant int& w_batch_ndims, + const constant int* w_shape, + const constant int64_t* w_strides, + const constant int64_t* s_strides, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + adjust_matrix_offsets( + x, + w, + scales, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + tid); + } + fp_qmv_wide_impl( + w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid); +} + template [[kernel]] void fp_qvm( const device uint32_t* w, diff --git a/mlx/backend/metal/kernels/fp_quantized.metal b/mlx/backend/metal/kernels/fp_quantized.metal index 0e0211853a..76980164f4 100644 --- a/mlx/backend/metal/kernels/fp_quantized.metal +++ b/mlx/backend/metal/kernels/fp_quantized.metal @@ -52,6 +52,17 @@ D, \ batched) +#define instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, batched) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_nv_" #vecs_per_tg "_kl_" #k_lanes "_batch_" #batched, \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + vecs_per_tg, \ + k_lanes, \ + batched) + #define instantiate_quantized_split_k(mode, name, type, split_k, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \ @@ -105,6 +116,17 @@ instantiate_quantized_quad(mode, qmv_quad, type, 128, 1, group_size, bits) \ instantiate_quantized_quad(mode, qmv_quad, type, 128, 0, group_size, bits) +// vecs_per_tg (input-vector tile) 2..5; the fp path uses k_lanes=16. +#define instantiate_quantized_wide_wrap(mode, name, type, vecs_per_tg, k_lanes, group_size, bits) \ + instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, 0) \ + instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, 1) + +#define instantiate_quantized_all_wide(type, mode, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 2, 16, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 3, 16, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 4, 16, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 5, 16, group_size, bits) + #define instantiate_quantized_all_splitk(type, mode, group_size, bits) \ instantiate_quantized_split_k(mode, qvm_split_k, type, 8, group_size, bits) \ instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits) \ @@ -139,6 +161,7 @@ instantiate_quantized_all_batched(type, mode, group_size, bits) \ instantiate_quantized_all_single(type, mode, group_size, bits) \ instantiate_quantized_all_quad(type, mode, group_size, bits) \ + instantiate_quantized_all_wide(type, mode, group_size, bits) \ instantiate_quantized_all_splitk(type, mode, group_size, bits) \ instantiate_quantized_all_aligned(type, mode, group_size, bits) \ instantiate_quantized_all_rhs(type, mode, group_size, bits) \ diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 2d0f60b980..7720c47d12 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -480,9 +480,10 @@ qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { } } -template -inline void -dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { +// Decode one quantized block (scale * q + bias) into w_local. W (the output +// pointer type) serves the threadgroup block loader or a thread-local decode. +template +inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { static_assert( bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || bits == 8, @@ -974,6 +975,103 @@ METAL_FUNC void qmv_impl( } } +// Affine analog of fp_qmv_wide. Weights carry a scale and bias per group, so +// each group is decoded in 8-value sub-chunks (scale * q + bias, registers +// bounded for any group_size) and reused across the vecs_per_tg vectors. +template +METAL_FUNC void qmv_wide_impl( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int num_simdgroups = 2; + constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes; + constexpr int sub = 8; // values per sub-chunk (== bits bytes, byte-aligned) + + typedef float U; + + const short k_lane = simd_lid % k_lanes; + const short sg_row = simd_lid / k_lanes; + + const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) + + results_per_simdgroup * simd_gid + sg_row; + const int vec0 = tid.x * vecs_per_tg; + + const int row = min(out_row, out_vec_size - 1); + + const int in_vec_size_w = in_vec_size * bits / 8; // bytes per weight row + const int in_vec_size_g = in_vec_size / group_size; + const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; + const device T* srow = scales + row * in_vec_size_g; + const device T* brow = biases + row * in_vec_size_g; + + const device T* xv[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + xv[v] = x + min(vec0 + v, M - 1) * in_vec_size; + } + + U result[vecs_per_tg] = {0}; + + // Each lane reduces a strided subset of the row's groups: decode the group in + // 8-value sub-chunks and reuse each chunk across the streamed vectors. + for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { + U scale = srow[g]; + U bias = brow[g]; +#pragma unroll + for (int sc = 0; sc < group_size / sub; sc++) { + const int k0 = g * group_size + sc * sub; + const device uint8_t* wc = wrow + k0 * bits / 8; + U w_dq[sub]; + dequantize(wc, scale, bias, w_dq); +#pragma unroll + for (int v = 0; v < vecs_per_tg; v++) { + const device T* xc = xv[v] + k0; + U acc = 0; +#pragma unroll + for (int i = 0; i < sub; i++) { + acc += static_cast(xc[i]) * w_dq[i]; + } + result[v] += acc; + } + } + } + + // Reduce each vector's partial over its k_lanes with a shuffle ladder: + // simd_sum would mix the results_per_simdgroup rows a simdgroup spans. + for (int v = 0; v < vecs_per_tg; v++) { + if constexpr (k_lanes >= 32) { + result[v] += simd_shuffle_down(result[v], 16); + } + if constexpr (k_lanes >= 16) { + result[v] += simd_shuffle_down(result[v], 8); + } + if constexpr (k_lanes >= 8) { + result[v] += simd_shuffle_down(result[v], 4); + } + if constexpr (k_lanes >= 4) { + result[v] += simd_shuffle_down(result[v], 2); + } + if constexpr (k_lanes >= 2) { + result[v] += simd_shuffle_down(result[v], 1); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + y[(vec0 + v) * out_vec_size + out_row] = static_cast(result[v]); + } + } + } +} + template METAL_FUNC void qvm_impl( const device uint32_t* w, @@ -1725,6 +1823,65 @@ template simd_lid); } +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool batched> +[[kernel]] void affine_qmv_wide( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + const constant int& x_batch_ndims, + const constant int* x_shape, + const constant int64_t* x_strides, + const constant int& w_batch_ndims, + const constant int* w_shape, + const constant int64_t* w_strides, + const constant int64_t* s_strides, + const constant int64_t* b_strides, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + adjust_matrix_offsets( + x, + w, + scales, + biases, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + b_strides, + tid); + } + qmv_wide_impl( + w, + scales, + biases, + x, + y, + in_vec_size, + out_vec_size, + M, + tid, + simd_gid, + simd_lid); +} + template [[kernel]] void affine_qvm( const device uint32_t* w [[buffer(0)]], diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 2c7a59dff3..75e788ea2f 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -52,6 +52,17 @@ D, \ batched) +#define instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, batched) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_nv_" #vecs_per_tg "_kl_" #k_lanes "_batch_" #batched, \ + name, \ + type, \ + group_size, \ + bits, \ + vecs_per_tg, \ + k_lanes, \ + batched) + #define instantiate_quantized_split_k(name, type, group_size, bits, split_k) \ instantiate_kernel( \ #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \ @@ -107,6 +118,18 @@ instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 1) \ instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 0) +// vecs_per_tg (input-vector tile) 2..5; affine uses k_lanes=8 (more rows per +// simdgroup) where the fp path uses 16. +#define instantiate_quantized_wide_wrap(name, type, group_size, bits, vecs_per_tg, k_lanes) \ + instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, 0) \ + instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, 1) + +#define instantiate_quantized_all_wide(type, group_size, bits) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 2, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 3, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 4, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 5, 8) + #define instantiate_quantized_all_splitk(type, group_size, bits) \ instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 8) \ instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 32) \ @@ -133,6 +156,7 @@ instantiate_quantized_all_batched(type, group_size, bits) \ instantiate_quantized_all_aligned(type, group_size, bits) \ instantiate_quantized_all_quad(type, group_size, bits) \ + instantiate_quantized_all_wide(type, group_size, bits) \ instantiate_quantized_all_splitk(type, group_size, bits) \ instantiate_quantized_all_splitk_qmm(type, group_size, bits) \ instantiate_quantized_all_rhs(type, group_size, bits) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 3adf093311..91aa804e21 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -299,6 +299,96 @@ void qmv( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +// affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. +inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { + return mode != "affine" || d.get_architecture_gen() >= 15; +} + +// Dispatches qmv_wide (fp modes -> fp_qmv_wide, affine -> affine_qmv_wide): +// vecs_per_tg input vectors streamed and reused per weight group. +void qmv_wide( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // vecs_per_tg is the per-threadgroup input-vector tile. Each tile re-reads + // the weights, so use the fewest tiles, then the smallest tile that fills + // them. + int n_tiles = (M + 4) / 5; // ceil(M / 5); tile size caps at 5 + int vecs_per_tg = (M + n_tiles - 1) / n_tiles; + + // k_lanes: lanes reducing K per output row (32/k_lanes rows per simdgroup). + // The affine subchunk decode has enough ALU per weight load to favor more + // rows per simdgroup (kl8); the fp modes' vectorized dot is balanced at 16. + int k_lanes = mode == "affine" ? 8 : 16; + constexpr int num_simdgroups = 2; + int B = out.size() / M / N; + bool batched = B > 1; + // Output rows per threadgroup: (32 / k_lanes) per simdgroup x num_simdgroups. + int rows_per_tg = (32 / k_lanes) * num_simdgroups; + + MTL::Size group_dims(32, num_simdgroups, 1); + MTL::Size grid_dims( + (M + vecs_per_tg - 1) / vecs_per_tg, + (N + rows_per_tg - 1) / rows_per_tg, + B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + "_qmv_wide_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_nv_", + vecs_per_tg, + "_kl_", + k_lanes, + batched ? "_batch_1" : "_batch_0"); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmv_wide", + mode, + type_string, + group_size, + bits, + vecs_per_tg, + k_lanes, + batched); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + add_strides_and_shapes(compute_encoder, !batched, x, w, scales, biases, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + void qvm_split_k( const array& x, const array& w, @@ -1583,6 +1673,13 @@ void dispatch_qmv( qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } + + // Small batch so route to qmv_wide, which reuses each weight group across the + // M vectors. + if (M >= 2 && use_qmv_wide(mode, d)) { + qmv_wide(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + return; + } qmv(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index f30170d44d..2850c7c357 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -482,6 +482,83 @@ def test_fp_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qmv_wide(self): + # M in [2, vector_limit) routes to qmv_wide -- except K in {64, 128} + # with power-of-2 bits, which stays on qmv_quad. Check both paths + # against a dequantize-then-matmul reference, with ragged M (token + # tail) and ragged N (output-tile remainder). B > 1 stacks a distinct + # weight matrix per slab and exercises the batched variant. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + # M <= 9 < vector_limit for these shapes (K, N <= 2048), so all stay on + # the mat-vec path; 7 and 9 also exercise the token-tail guard. + Ms = [2, 3, 4, 5, 6, 7, 9] + Ns = [256, 67] # 67 is a non-multiple of the 4-row output tile + Bs = [1, 3] + + # Affine: every bit-width and group size. + for group_size, bits, K in product( + [32, 64, 128], [2, 3, 4, 5, 6, 8], [128, 512] + ): + for M, N, B in product(Ms, Ns, Bs): + with self.subTest(M=M, N=N, K=K, B=B, group_size=group_size, bits=bits): + x_shape = (M, K) if B == 1 else (B, M, K) + w_shape = (N, K) if B == 1 else (B, N, K) + x = mx.random.normal(shape=x_shape, key=k1) + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales, biases = mx.quantize(w, group_size, bits) + w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) + y_q = mx.quantized_matmul( + x, w_q, scales, biases, True, group_size, bits + ) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + + # FP modes (group_size and bits implied by the mode). + for mode, K in product(["mxfp4", "nvfp4", "mxfp8"], [128, 512]): + for M, N, B in product(Ms, Ns, Bs): + with self.subTest(M=M, N=N, K=K, B=B, mode=mode): + x_shape = (M, K) if B == 1 else (B, M, K) + w_shape = (N, K) if B == 1 else (B, N, K) + x = mx.random.normal(shape=x_shape, key=k1) + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + y_q = mx.quantized_matmul(x, w_q, scales, transpose=True, mode=mode) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + + # Tiny shapes (M, K, N): small K and non-multiple output rows. + tiny = [(2, 32, 10), (4, 32, 7), (3, 64, 5), (5, 64, 3)] + settings = [(4, 32, "affine"), (6, 32, "affine"), (4, 16, "nvfp4")] + for M, K, N in tiny: + for bits, group_size, mode in settings: + with self.subTest( + M=M, K=K, N=N, bits=bits, group_size=group_size, mode=mode + ): + x = mx.random.normal(shape=(M, K), key=k1) + w = mx.random.normal(shape=(N, K), key=k2) + w_q, *sb = mx.quantize( + w, group_size=group_size, bits=bits, mode=mode + ) + w_hat = mx.dequantize( + w_q, *sb, group_size=group_size, bits=bits, mode=mode + ) + y_q = mx.quantized_matmul( + x, + w_q, + *sb, + transpose=True, + group_size=group_size, + bits=bits, + mode=mode, + ) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qvm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From 18672fbefcfdd43fd9f2ff2dddd9bc7c06954446 Mon Sep 17 00:00:00 2001 From: rohith Date: Sat, 22 Aug 2026 07:18:16 +0530 Subject: [PATCH 63/84] Fix Log and Equal is_equivalent ignoring primitive state (#4266) Co-authored-by: Cheng --- mlx/primitives.cpp | 10 ++++++++++ mlx/primitives.h | 4 ++-- python/tests/test_compile.py | 23 +++++++++++++++++++++++ 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 7a6c729c39..9a1771394e 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -1941,6 +1941,11 @@ std::pair, std::vector> Equal::vmap( return {{equal(a, b, stream())}, {to_ax}}; } +bool Equal::is_equivalent(const Primitive& other) const { + const Equal& e_other = static_cast(other); + return equal_nan_ == e_other.equal_nan_; +} + std::vector Equal::vjp( const std::vector& primals, const std::vector& cotangents, @@ -2795,6 +2800,11 @@ std::pair, std::vector> Log::vmap( axes}; } +bool Log::is_equivalent(const Primitive& other) const { + const Log& l_other = static_cast(other); + return base_ == l_other.base_; +} + std::vector Log1p::vjp( const std::vector& primals, const std::vector& cotangents, diff --git a/mlx/primitives.h b/mlx/primitives.h index 3a3d0ba5e5..0cfc71bf04 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -975,9 +975,9 @@ class Equal : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS() - DEFINE_DEFAULT_IS_EQUIVALENT() DEFINE_INPUT_OUTPUT_SHAPE() + bool is_equivalent(const Primitive& other) const override; const char* name() const override { if (equal_nan_) { return "NaNEqual"; @@ -1325,9 +1325,9 @@ class Log : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS() - DEFINE_DEFAULT_IS_EQUIVALENT() DEFINE_INPUT_OUTPUT_SHAPE() + bool is_equivalent(const Primitive& other) const override; Base state() const { return base_; }; diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 1e1b20b05d..5eaa6cb955 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1623,6 +1623,29 @@ def test_compile_abs_unsigned(self): x = mx.array([1, 2, 3], dtype) self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + def test_compile_different_log_bases(self): + # The logs are intermediates, since outputs are not simplified. + def entropies(p): + nats = -mx.sum(p * mx.log(p)) + bits = -mx.sum(p * mx.log2(p)) + return mx.stack([nats, bits]) + + p = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32) + expected = np.array( + [-(p * np.log(p)).sum(), -(p * np.log2(p)).sum()], dtype=np.float32 + ) + out = mx.compile(entropies)(mx.array(p)) + self.assertTrue(np.allclose(out, expected, atol=1e-5)) + + def test_compile_equal_nan(self): + def fun(x): + return mx.stack( + [mx.array_equal(x, x), mx.array_equal(x, x, equal_nan=True)] + ) + + x = mx.array([1.0, float("nan"), 3.0]) + self.assertTrue(mx.array_equal(mx.compile(fun)(x), mx.array([False, True]))) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 9b954755f529ea1ec31d4d1ca62d5fe8cae161de Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Sat, 22 Aug 2026 04:49:56 +0300 Subject: [PATCH 64/84] Stabilize reduced-precision InstanceNorm (#4230) --- python/mlx/nn/layers/normalization.py | 19 +++++++++++++------ python/tests/test_nn.py | 15 +++++++++++++++ 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index a5a8293fdb..97f6942f04 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -64,12 +64,19 @@ def __call__(self, x: mx.array) -> mx.array: f"InstanceNorm expects inputs with at least 3 dimensions" f" (N, ..., C) but the input has {x.ndim} dimensions." ) - reduction_axes = tuple(range(1, x.ndim - 1)) - # Compute stats - mean = mx.mean(x, axis=reduction_axes, keepdims=True) - var = mx.var(x, axis=reduction_axes, keepdims=True) - # Normalize - x = (x - mean) * mx.rsqrt(var + self.eps) + batch_size, features = x.shape[0], x.shape[-1] + spatial_shape = x.shape[1:-1] + channels_first = mx.transpose(x, (0, x.ndim - 1, *range(1, x.ndim - 1))) + x = mx.fast.layer_norm( + channels_first.reshape(batch_size, features, -1), + None, + None, + self.eps, + ) + x = mx.transpose( + x.reshape(batch_size, features, *spatial_shape), + (0, *range(2, len(spatial_shape) + 2), 1), + ) # Scale and shift if necessary return (self.weight * x + self.bias) if "weight" in self else x diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 46242d192c..26b8fd1162 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -693,6 +693,21 @@ def test_instance_norm(self): ] self.assertTrue(x.shape == y.shape) self.assertTrue(np.allclose(y, expected_y, atol=1e-5)) + # Reduced-precision statistics must not overflow for finite feature maps. + checkerboard = np.indices((4, 4, 4)).sum(axis=0) % 2 + x = mx.array( + np.stack( + [ + np.where(checkerboard, -512, 512), + np.where(checkerboard, -256, 256), + ], + axis=-1, + ).astype(np.float16) + )[None] + y = nn.InstanceNorm(dims=2)(x) + self.assertEqual(y.dtype, mx.float16) + self.assertTrue(mx.allclose(y.min(), mx.array(-1.0, dtype=mx.float16))) + self.assertTrue(mx.allclose(y.max(), mx.array(1.0, dtype=mx.float16))) # Test repr self.assertTrue(str(inorm) == "InstanceNorm(3, eps=1e-05, affine=False)") # Raise for inputs without spatial dimensions From 846d176227a0ac13d2667e58d2bb68b322109ab0 Mon Sep 17 00:00:00 2001 From: "Brian C." <94733710+deBrian07@users.noreply.github.com> Date: Sat, 22 Aug 2026 04:00:42 -0400 Subject: [PATCH 65/84] chore: Normalize negative axes in sort and argsort (#4332) --- mlx/ops.cpp | 24 ++++-------------------- python/tests/test_vmap.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 34bf40538f..8d26332e67 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2794,17 +2794,9 @@ array sort(const array& a, StreamOrDevice s /* = {} */) { /** Returns a sorted copy of the array along a given axis. */ array sort(const array& a, int axis, StreamOrDevice s /* = {} */) { - // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[sort] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - + auto ax = normalize_axis_index(axis, a.ndim(), "[sort] "); return array( - a.shape(), a.dtype(), std::make_shared(to_stream(s), axis), {a}); + a.shape(), a.dtype(), std::make_shared(to_stream(s), ax), {a}); } /** Returns indices that sort the flattened array. */ @@ -2815,17 +2807,9 @@ array argsort(const array& a, StreamOrDevice s /* = {} */) { /** Returns indices that sort the array along a given axis. */ array argsort(const array& a, int axis, StreamOrDevice s /* = {} */) { - // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[argsort] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - + auto ax = normalize_axis_index(axis, a.ndim(), "[argsort] "); return array( - a.shape(), uint32, std::make_shared(to_stream(s), axis), {a}); + a.shape(), uint32, std::make_shared(to_stream(s), ax), {a}); } /** diff --git a/python/tests/test_vmap.py b/python/tests/test_vmap.py index ec97ed6fdf..b050a80384 100644 --- a/python/tests/test_vmap.py +++ b/python/tests/test_vmap.py @@ -1023,6 +1023,20 @@ def fn(x, y): self.assertTrue(mx.array_equal(expected, out)) self.assertEqual(6, counter[0]) + def test_vmap_sort(self): + a = mx.random.uniform(shape=(3, 5)) + expected = mx.stack([mx.sort(a[:, i]) for i in range(a.shape[1])], axis=1) + for axis in (0, -1): + out = mx.vmap(lambda x: mx.sort(x, axis=axis), in_axes=1, out_axes=1)(a) + self.assertTrue(mx.array_equal(out, expected)) + + def test_vmap_argsort(self): + a = mx.random.uniform(shape=(3, 5)) + expected = mx.stack([mx.argsort(a[:, i]) for i in range(a.shape[1])], axis=1) + for axis in (0, -1): + out = mx.vmap(lambda x: mx.argsort(x, axis=axis), in_axes=1, out_axes=1)(a) + self.assertTrue(mx.array_equal(out, expected)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 1038679aa37d288c1583c2cab74c526d96372d3d Mon Sep 17 00:00:00 2001 From: Cheng Date: Sun, 23 Aug 2026 07:37:48 +0900 Subject: [PATCH 66/84] Clean up main thread compile cache before python interpreter shuts down (#4373) --- python/src/transforms.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index 1ec20a1375..9ac9a39924 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -1538,4 +1538,11 @@ void init_transforms(nb::module_& m) { A callable that recomputes intermediate states during gradient computation. )pbdoc"); + + // Clean up main thread compile cache before python interpreter shuts down. + auto atexit = nb::module_::import_("atexit"); + atexit.attr("register")( + nb::cpp_function([cache = mx::detail::compile_cache()]() { + mx::detail::compile_clear_cache(cache); + })); } From c7a185ac7b0a5754feeaff73bb9267b1f66d8294 Mon Sep 17 00:00:00 2001 From: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Date: Sat, 22 Aug 2026 19:13:20 -0400 Subject: [PATCH 67/84] chore: Check malformed jaccl hostfile that miss rdma in pairs (#4284) Co-authored-by: Cheng --- python/mlx/_distributed_utils/launch.py | 26 +++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 4771c1fb5b..0e661e358f 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -376,9 +376,31 @@ def launch_jaccl(parser, hosts, args, command): jaccl_ring = args.backend == "jaccl-ring" have_rdmas = all(len(h.rdma) == len(hosts) for h in hosts) + if not have_rdmas: + parser.error( + "The hostfile is malformed: number of RDMA devices does not match hosts" + ) have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts)) - if not have_rdmas or not have_nulls: - parser.error("Malformed hostfile for jaccl backend") + if not have_nulls: + parser.error("The hostfile is malformed: RDMA device of self should be null") + + # Find pairs that miss rmda in hostfile. + n = len(hosts) + missing_rdma = [ + (i, j) + for i, h in enumerate(hosts) + for j in (((i - 1) % n, (i + 1) % n) if jaccl_ring else range(n)) + if i != j and h.rdma[j] is None + ] + + if missing_rdma: + pairs = ", ".join( + f"{hosts[i].ssh_hostname} to {hosts[j].ssh_hostname}" + for i, j in missing_rdma[:3] + ) + if len(missing_rdma) > 3: + pairs += f" and {len(missing_rdma) - 3} more" + parser.error(f"The hostfile is malformed: no RDMA device is listed for {pairs}") coordinator = hosts[0].ips[0] env = args.env From 02adf7b2125f50f6a03295b6481690c5c960569a Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Sat, 22 Aug 2026 17:35:25 -0700 Subject: [PATCH 68/84] Round mxfp8 block scales up to avoid saturation (#4353) Co-authored-by: Daniel Hiltgen Co-authored-by: Cheng --- mlx/backend/cpu/quantized.cpp | 13 +++++++++++-- mlx/backend/metal/kernels/fp8.h | 11 +++++++++++ mlx/backend/metal/kernels/fp_quantized.h | 8 ++++++-- mlx/ops.cpp | 16 +++++++++++----- python/tests/test_quantized.py | 23 +++++++++++++++++++++++ 5 files changed, 62 insertions(+), 9 deletions(-) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index 3469d99788..15e00cd913 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include "mlx/backend/common/quantized.h" #include "mlx/backend/common/unary.h" @@ -1061,6 +1061,15 @@ uint8_t to_fp8_e8m0(float x) { return static_cast(n + 127); } +// Smallest E8M0 >= x, so a block's largest elements do not saturate. +uint8_t to_fp8_e8m0_round_up(float x) { + uint8_t bits = to_fp8_e8m0(x); + if (bits < 0xFE && dequantize_scale(bits) < x) { + bits += 1; + } + return bits; +} + uint8_t to_fp4_e2m1(float x) { if (std::isnan(x)) { return 0x7; @@ -1112,7 +1121,7 @@ void fp_quantize_dequantize( if (group_size == 16) { scale = dequantize_scale(detail::ToFP8()(scale)); } else { - scale = dequantize_scale(to_fp8_e8m0(scale)); + scale = dequantize_scale(to_fp8_e8m0_round_up(scale)); } for (int j = 0; j < group_size; ++j) { diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 796dd21639..42c5ec128a 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -78,3 +78,14 @@ struct fp8_e8m0 { uint8_t bits; }; + +// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down +// leaves the block's largest elements outside the element range, where they +// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0. +inline float mx_scale_round_up(float x) { + fp8_e8m0 s(x); + if (s.bits < 0xFE && float(s) < x) { + s.bits += 1; + } + return float(s); +} diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 8c963030f2..3ac94acb35 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -2122,7 +2122,7 @@ template float scale_dec_b; float w_thread = w[index]; - if (use_mx_scale) { + if constexpr (use_mx_scale) { scale_dec_b = simd_max(abs(w_thread)); } else { float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); @@ -2132,6 +2132,8 @@ template scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; if constexpr (has_global_scale) { scale_dec_b *= scale_enc; + } else if constexpr (use_mx_scale) { + scale_dec_b = mx_scale_round_up(scale_dec_b); } using ScaleType = metal::conditional_t; @@ -2214,7 +2216,7 @@ template float scale_dec_b; float w_thread = w[index]; - if (use_mx_scale) { + if constexpr (use_mx_scale) { scale_dec_b = simd_max(abs(w_thread)); } else { float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); @@ -2224,6 +2226,8 @@ template scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; if constexpr (has_global_scale) { scale_dec_b *= scale_enc; + } else if constexpr (use_mx_scale) { + scale_dec_b = mx_scale_round_up(scale_dec_b); } using ScaleType = metal::conditional_t; diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 8d26332e67..5fe3d9be59 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. // Required for using M_PI in MSVC. #define _USE_MATH_DEFINES @@ -5148,11 +5148,17 @@ std::vector fp_quantize( } else { // convert to e8m0 auto z = array(0, scales.dtype()); - scales = where( - equal(scales, z, s), - z, - astype(round(log2(scales, s), s), int32, s), + // Round the scale up so the block maximum stays representable, + // matching the CUDA backend. + auto exponent = astype(round(log2(scales, s), s), int32, s); + auto decoded = + power(array(2.0f, float32), astype(exponent, float32, s), s); + exponent = where( + less(decoded, astype(scales, float32, s), s), + add(exponent, array(1, int32), s), + exponent, s); + scales = where(equal(scales, z, s), z, exponent, s); wq = divide(wq, power(array(2.0f, w.dtype()), scales, s), s); scales = astype(add(scales, array(127, int32), s), uint8, s); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 319e9d2d67..28033cbbab 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,5 +1,6 @@ # Copyright © 2023-2026 Apple Inc. +import math import os import platform import subprocess @@ -121,6 +122,28 @@ def test_mxfp8_quantize_dequantize(self): w_hat = mx.dequantize(w_q, scales, mode="mxfp8") self.assertTrue(mx.all(w_hat == 0)) + def test_mxfp8_block_scale_does_not_saturate(self): + # E4M3 has three mantissa bits, so an in-range element loses at most + # half a step, 6.25%. More than that means the block scale rounded + # below amax/448 and the block maximum saturated. + mx.random.seed(0) + group_size = 32 + n_blocks = 512 + + # Sweep the block magnitude across one binade so both scale rounding + # directions are covered. + w = mx.random.normal(shape=(n_blocks, group_size)) + w = w * mx.exp(mx.arange(n_blocks) / n_blocks * math.log(2.0)).reshape(-1, 1) + + w_q, scales = mx.quantize(w, group_size=group_size, mode="mxfp8") + w_hat = mx.dequantize(w_q, scales, group_size=group_size, mode="mxfp8") + + # Quantization is monotone in |w|, so a block's largest output is the + # reconstruction of its largest input. + amax = mx.max(mx.abs(w), axis=1) + rel = mx.abs(amax - mx.max(mx.abs(w_hat), axis=1)) / amax + self.assertLess(mx.max(rel).item(), 0.0626) + def test_nvfp4_quantize_dequantize(self): lut = mx.array( [ From 7789905ede2a79572c800aaea7062d53519edcfe Mon Sep 17 00:00:00 2001 From: Aaishwarya Mishra Date: Sun, 23 Aug 2026 09:12:50 +0530 Subject: [PATCH 69/84] Add support for the __array_namespace_info__ (#4334) --- python/mlx/__array_api_info.py | 82 ++++++++++++++++++++++++++++++++++ python/src/mlx.cpp | 4 ++ python/tests/test_array.py | 31 +++++++++++++ 3 files changed, 117 insertions(+) create mode 100644 python/mlx/__array_api_info.py diff --git a/python/mlx/__array_api_info.py b/python/mlx/__array_api_info.py new file mode 100644 index 0000000000..847a0bcbf3 --- /dev/null +++ b/python/mlx/__array_api_info.py @@ -0,0 +1,82 @@ +class ArrayNamespaceInfo: + def capabilities(self): + return { + "boolean indexing": False, + "data-dependent shapes": False, + "max dimensions": 10, + } + + def default_device(self): + import mlx.core as mx + + return mx.default_device() + + def default_dtypes(self, *, device=None): + import mlx.core as mx + + if device is not None and not isinstance(device, mx.Device): + raise TypeError("Expected a mlx Device") + return { + "real floating": mx.float32, + "complex floating": mx.complex64, + "integral": mx.int32, + "indexing": mx.int32, + } + + def devices(self): + import mlx.core as mx + + devices = [ + mx.Device(dev_type, i) + for dev_type in (mx.cpu, mx.gpu) + for i in range(mx.device_count(dev_type)) + ] + return tuple(devices) + + def dtypes(self, *, device=None, kind=None): + import mlx.core as mx + + if device is not None and not isinstance(device, mx.Device): + raise TypeError("Expected a mlx Device") + device = device if device is not None else self.default_device() + + dtypes = { + "bool": mx.bool_, + "int8": mx.int8, + "int16": mx.int16, + "int32": mx.int32, + "int64": mx.int64, + "uint8": mx.uint8, + "uint16": mx.uint16, + "uint32": mx.uint32, + "uint64": mx.uint64, + "float32": mx.float32, + "complex64": mx.complex64, + } + if device.type == mx.cpu: + dtypes["float64"] = mx.float64 + if kind is None: + return dtypes + + signed = {"int8", "int16", "int32", "int64"} + unsigned = {"uint8", "uint16", "uint32", "uint64"} + real = {"float32", "float64"} + complex_ = {"complex64"} + kinds = { + "bool": {"bool"}, + "signed integer": signed, + "unsigned integer": unsigned, + "integral": signed | unsigned, + "real floating": real, + "complex floating": complex_, + "numeric": signed | unsigned | real | complex_, + } + kind = (kind,) if isinstance(kind, str) else kind + if not isinstance(kind, tuple) or any(k not in kinds for k in kind): + raise ValueError(f"Unsupported dtype kind: {kind!r}") + names = {name for k in kind for name in kinds[k]} + return {name: dtype for name, dtype in dtypes.items() if name in names} + + +def __array_namespace_info__(): + return ArrayNamespaceInfo() diff --git a/python/src/mlx.cpp b/python/src/mlx.cpp index cb031cf78c..243449b385 100644 --- a/python/src/mlx.cpp +++ b/python/src/mlx.cpp @@ -31,6 +31,10 @@ NB_MODULE(core, m) { auto reprlib_fix = nb::module_::import_("mlx._reprlib_fix"); nb::set_leak_warnings(false); + auto array_namespace_info = nb::module_::import_("mlx.__array_api_info"); + m.attr("__array_namespace_info__") = + array_namespace_info.attr("__array_namespace_info__"); + init_mlx_func(m); init_device(m); init_stream(m); diff --git a/python/tests/test_array.py b/python/tests/test_array.py index 2c55a3ba20..80459fabe5 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -49,6 +49,37 @@ def test_version(self): self.assertEqual(v, mx.__version__[: len(v)]) +class TestArrayNamespsceInfo(mlx_tests.MLXTestCase): + def test(self): + namespace = mx.__array_namespace_info__() + + self.assertEqual(namespace.default_device(), mx.default_device()) + self.assertEqual( + namespace.default_dtypes(), + { + "real floating": mx.float32, + "complex floating": mx.complex64, + "integral": mx.int32, + "indexing": mx.int32, + }, + ) + self.assertEqual( + namespace.dtypes(device=mx.Device(mx.cpu), kind="real floating"), + {"float32": mx.float32, "float64": mx.float64}, + ) + if mx.is_available(mx.gpu): + self.assertEqual( + namespace.dtypes(device=mx.Device(mx.gpu), kind="real floating"), + {"float32": mx.float32}, + ) + self.assertEqual( + namespace.dtypes(kind=("bool", "complex floating")), + {"bool": mx.bool_, "complex64": mx.complex64}, + ) + with self.assertRaises(ValueError): + namespace.dtypes(kind="invalid") + + class TestDtypes(mlx_tests.MLXTestCase): def test_dtypes(self): self.assertEqual(mx.bool_.size, 1) From ab3ef95574a1880dd8f160934f991bbb8b5272eb Mon Sep 17 00:00:00 2001 From: YH Yan Date: Sun, 23 Aug 2026 15:05:25 +0800 Subject: [PATCH 70/84] Stop a failed CUDA graph commit from poisoning the encoder (#4356) Co-authored-by: Cheng --- mlx/backend/cuda/cuda_utils.h | 6 ++++ mlx/backend/cuda/device.cpp | 60 +++++++++++++++++++++++++++-------- mlx/backend/cuda/device.h | 2 ++ 3 files changed, 55 insertions(+), 13 deletions(-) diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h index 7bae911d26..f8a234ee65 100644 --- a/mlx/backend/cuda/cuda_utils.h +++ b/mlx/backend/cuda/cuda_utils.h @@ -50,6 +50,12 @@ class CudaHandle { } } + Handle release() { + Handle handle = handle_; + handle_ = nullptr; + return handle; + } + operator Handle() const { return handle_; } diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index 30248f5568..472b3d99fb 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -462,6 +462,42 @@ bool CommandEncoder::needs_commit() { void CommandEncoder::commit() { nvtx3::scoped_range r("CommandEncoder::commit"); + try { + commit_impl(); + } catch (...) { + // Clear pending CUDA error first. + cudaGetLastError(); + // Clear states. + clear_graph_state(); + node_count_ = 0; + bytes_in_graph_ = 0; + // Clear graph. + try { + graph_.reset(); + } catch (...) { + // Destroying could fail. + graph_.release(); + } + try { + graph_ = CudaGraph(device_); + } catch (...) { + // Keep the original error. + } + // Re-throw the error. + throw; + } +} + +void CommandEncoder::synchronize() { + CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); + auto p = std::make_shared>(); + std::future f = p->get_future(); + add_completed_handler([p = std::move(p)]() { p->set_value(); }); + commit(); + f.wait(); +} + +void CommandEncoder::commit_impl() { if (!temporaries_.empty()) { add_completed_handler([temporaries = std::move(temporaries_)]() {}); } @@ -520,13 +556,8 @@ void CommandEncoder::commit() { } // Reset state - from_nodes_.clear(); - to_nodes_.clear(); - graph_deps_key_.clear(); - graph_nodes_key_.clear(); - node_map_.clear(); + clear_graph_state(); graph_ = CudaGraph(device_); - is_graph_updatable_ = true; } // Put completion handlers in a batch. @@ -535,13 +566,16 @@ void CommandEncoder::commit() { bytes_in_graph_ = 0; } -void CommandEncoder::synchronize() { - CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); - auto p = std::make_shared>(); - std::future f = p->get_future(); - add_completed_handler([p = std::move(p)]() { p->set_value(); }); - commit(); - f.wait(); +void CommandEncoder::clear_graph_state() { + from_nodes_.clear(); + to_nodes_.clear(); + graph_deps_key_.clear(); + graph_nodes_key_.clear(); + node_map_.clear(); + active_deps_.clear(); + active_outputs_.clear(); + concurrent_nodes_.clear(); + is_graph_updatable_ = true; } Device& device(int cuda_device) { diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h index 15d75082e9..198f0b5ad8 100644 --- a/mlx/backend/cuda/device.h +++ b/mlx/backend/cuda/device.h @@ -138,6 +138,8 @@ class CommandEncoder { std::string id; }; + void commit_impl(); + void clear_graph_state(); void insert_graph_dependencies(GraphNode node); void insert_graph_dependencies(std::vector nodes); From 7408e687afac5c78ffe15d52ffda1d50fb9a97a4 Mon Sep 17 00:00:00 2001 From: Dwijen Patel Date: Sun, 23 Aug 2026 00:22:52 -0700 Subject: [PATCH 71/84] Fix quantized kernels in JIT build (#4372) Co-authored-by: Cheng --- .github/actions/setup/action.yml | 9 +-------- .github/workflows/build_and_test.yml | 14 +++++++++----- mlx/backend/metal/kernels/fp_quantized.h | 3 ++- mlx/backend/metal/kernels/quantized.h | 6 ++++-- mlx/backend/metal/quantized.cpp | 4 ++-- 5 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 947c0ff1f0..510086390a 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -229,14 +229,7 @@ runs: else cmakeArgs+=("-DMLX_BUILD_METAL=ON") if ${{ inputs.toolkit == 'jit' }} ; then - cmakeArgs+=( - "-DBUILD_SHARED_LIBS=ON" - "-DCMAKE_BUILD_TYPE=MinSizeRel" - "-DMLX_BUILD_CPU=OFF" - "-DMLX_BUILD_SAFETENSORS=OFF" - "-DMLX_BUILD_GGUF=OFF" - "-DMLX_METAL_JIT=ON" - ) + cmakeArgs+=("-DMLX_METAL_JIT=ON") fi fi fi diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index ece39d15d9..ee6bac8fbf 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -105,7 +105,7 @@ jobs: with: toolkit: 'cpu' - uses: actions/upload-artifact@v7 - if: matrix.toolkit == 'metal' + if: matrix.toolkit != 'cpu' with: name: mlx-${{ matrix.toolkit }}-macos${{ matrix.macos-target }} path: | @@ -115,24 +115,28 @@ jobs: if-no-files-found: error mac_test: - name: Test macOS + name: Test macOS (${{ matrix.toolkit }}) if: github.repository == 'ml-explore/mlx' runs-on: [self-hosted, macos] needs: mac_build + strategy: + fail-fast: false + matrix: + toolkit: ['metal', 'jit'] steps: - uses: actions/checkout@v7 - uses: ./.github/actions/setup with: - toolkit: 'metal' + toolkit: ${{ matrix.toolkit }} use-ccache: false - uses: actions/download-artifact@v8 with: path: artifact - pattern: mlx-metal-* + pattern: mlx-${{ matrix.toolkit }}-* - run: ls -lhR artifact - uses: ./.github/actions/test-macos with: - toolkit: 'metal' + toolkit: ${{ matrix.toolkit }} build_documentation: name: Build Documentation diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 3ac94acb35..7061771d59 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -1232,7 +1232,8 @@ template < int group_size, int bits, bool batched, - bool has_global_scale = false> + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv( const device uint32_t* w, const device uint8_t* scales, diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index f628c05612..42855e14d8 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -1603,7 +1603,8 @@ template < int group_size, int bits, bool batched, - bool has_global_scale = false> + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1660,7 +1661,8 @@ template < int group_size, const int bits, bool batched, - bool has_global_scale = false> + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void affine_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 013bc6d333..65c24152b0 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -972,7 +972,7 @@ void gather_qmm_nax( kernel = get_qmm_nax_kernel_wrapped( d, kname, - "gather_qmm_t_nax_", + "gather_qmm_t_nax", mode, type_string, group_size, @@ -987,7 +987,7 @@ void gather_qmm_nax( kernel = get_qmm_nax_kernel_wrapped( d, kname, - "gather_qmm_n_nax_", + "gather_qmm_n_nax", mode, type_string, group_size, From d9077d8316ad7305497a3ecf2296bd0e0e99a627 Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Sun, 23 Aug 2026 13:07:03 +0300 Subject: [PATCH 72/84] Avoid zero work in stride-2 ConvTranspose3d (#4343) --- mlx/backend/metal/conv.cpp | 114 ++++++++++++++++++++++++++++ python/tests/test_conv_transpose.py | 2 + 2 files changed, 116 insertions(+) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index acc88a5d2d..5095522616 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -904,6 +904,115 @@ void small_kd_conv_3D_gpu( 0); } +// A stride-2, kernel-2 transposed convolution has exactly one valid kernel +// phase for every output coordinate. The explicit unfold path nevertheless +// materializes all eight phases and fills seven of them with zeros. Compute +// each phase as a small regular GEMM instead. This path is intentionally +// narrow: other transposed-convolution configurations retain the general +// implementation below. +bool is_stride_two_conv_transpose_3D(const MLXConvParams<3>& p) { + return p.groups == 1 && p.flip && p.str[0] == 1 && p.str[1] == 1 && + p.str[2] == 1 && p.idil[0] == 2 && p.idil[1] == 2 && p.idil[2] == 2 && + p.kdil[0] == 1 && p.kdil[1] == 1 && p.kdil[2] == 1 && p.wS[0] == 2 && + p.wS[1] == 2 && p.wS[2] == 2 && p.pad[0] == 1 && p.pad[1] == 1 && + p.pad[2] == 1 && static_cast(p.oS[0]) == 2LL * p.iS[0] && + static_cast(p.oS[1]) == 2LL * p.iS[1] && + static_cast(p.oS[2]) == 2LL * p.iS[2]; +} + +void stride_two_conv_transpose_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<3>& p, + std::vector& copies) { + constexpr int kernel_volume = 8; + const int C = p.C; + const int O = p.O; + + // The input and weight are contiguous by the time this helper is called. + // Every phase covers the complete input volume; the phase bit only selects + // the interleaved output coordinates and the corresponding weight slice. + for (int phase = 0; phase < kernel_volume; ++phase) { + const int pd = (phase >> 2) & 1; + const int ph = (phase >> 1) & 1; + const int pw = phase & 1; + const int D = p.iS[0]; + const int H = p.iS[1]; + const int W = p.iS[2]; + + const int M = safe_cast(static_cast(p.N) * D * H * W, "conv"); + + // The weight is [O, 2, 2, 2, C]. Present one spatial phase as a [C, O] + // matrix with the layout expected by a transposed Steel GEMM. + array wt_phase({C, O}, wt.dtype(), nullptr, {}); + array::Flags wt_flags = wt.flags(); + wt_flags.contiguous = false; + wt_flags.row_contiguous = false; + wt_flags.col_contiguous = true; + wt_phase.copy_shared_buffer( + wt, + {1, wt.strides(0)}, + wt_flags, + wt.data_size(), + static_cast(pd * 4 + ph * 2 + pw) * C); + + array in_matrix({M, C}, in.dtype(), nullptr, {}); + in_matrix.copy_shared_buffer(in, {C, 1}, in.flags(), in.data_size()); + + array phase_out({M, O}, out.dtype(), nullptr, {}); + phase_out.set_data(allocator::malloc(phase_out.nbytes())); + + std::vector gemm_copies = {in_matrix, wt_phase}; + steel_matmul( + s, + d, + /* a = */ in_matrix, + /* b = */ wt_phase, + /* out = */ phase_out, + /* M = */ M, + /* N = */ O, + /* K = */ C, + /* batch_size_out = */ 1, + /* lda = */ C, + /* ldb = */ kernel_volume * C, + /* a_transposed = */ false, + /* b_transposed = */ true, + /* copies = */ gemm_copies); + + Shape phase_out_shape{p.N, D, H, W, O}; + array phase_out_nd(phase_out_shape, out.dtype(), nullptr, {}); + phase_out_nd.copy_shared_buffer( + phase_out, + make_contiguous_strides(phase_out_shape), + phase_out.flags(), + phase_out.data_size()); + + Strides out_phase_strides = out.strides(); + out_phase_strides[1] *= 2; + out_phase_strides[2] *= 2; + out_phase_strides[3] *= 2; + array::Flags out_phase_flags = out.flags(); + out_phase_flags.contiguous = false; + out_phase_flags.row_contiguous = false; + out_phase_flags.col_contiguous = false; + array out_phase_view(phase_out_shape, out.dtype(), nullptr, {}); + out_phase_view.copy_shared_buffer( + out, + out_phase_strides, + out_phase_flags, + out.data_size(), + pd * out.strides(1) + ph * out.strides(2) + pw * out.strides(3)); + + copy_gpu_inplace(phase_out_nd, out_phase_view, CopyType::GeneralGeneral, s); + copies.push_back(phase_out); + copies.push_back(phase_out_nd); + copies.push_back(out_phase_view); + } +} + void dispatch_conv_3D_gpu( const Stream& s, metal::Device& d, @@ -934,6 +1043,11 @@ void dispatch_conv_3D_gpu( auto in = ensure_row_contiguous(in_pre, d, s); auto wt = ensure_row_contiguous(wt_pre, d, s); + if (is_stride_two_conv_transpose_3D(conv_params)) { + return stride_two_conv_transpose_3D_gpu( + s, d, in, wt, out, conv_params, copies); + } + // Decompose 3D conv to per-frame 2D convs constexpr int kSmallKdLimit3D = 7; if (is_idil_one && mod16_channels && conv_params.groups == 1 && diff --git a/python/tests/test_conv_transpose.py b/python/tests/test_conv_transpose.py index e6def081a7..7d12cb16f4 100644 --- a/python/tests/test_conv_transpose.py +++ b/python/tests/test_conv_transpose.py @@ -486,6 +486,8 @@ def run_conv_transpose3D( ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0)), ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), ((15, 15, 15), (3, 3, 3), (3, 3, 3), (2, 2, 2)), + # Exercises the Metal phase-aware stride-2/kernel-2 path. + ((3, 4, 5), (2, 2, 2), (2, 2, 2), (0, 0, 0)), ): run_conv_transpose3D( N, C, O, idim, kdim, stride, padding, dtype=dtype From 451dc8759703b8e3f3cde34251292edaff63a50f Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Sun, 23 Aug 2026 21:32:18 +0200 Subject: [PATCH 73/84] [CUDA] Ce fused kernel (#3947) --- docs/src/python/fast.rst | 1 + mlx/backend/cuda/CMakeLists.txt | 1 + mlx/backend/cuda/cross_entropy.cu | 245 ++++++++++++++++++++++++++++++ mlx/backend/metal/primitives.cpp | 23 +++ mlx/backend/no_gpu/primitives.cpp | 2 + mlx/fast.cpp | 101 ++++++++++++ mlx/fast.h | 4 + mlx/fast_primitives.h | 61 ++++++++ python/mlx/nn/losses.py | 59 ++++--- python/src/fast.cpp | 30 ++++ python/tests/test_fast.py | 61 ++++++++ 11 files changed, 568 insertions(+), 20 deletions(-) create mode 100644 mlx/backend/cuda/cross_entropy.cu diff --git a/docs/src/python/fast.rst b/docs/src/python/fast.rst index affeb444f8..c930c7bb4a 100644 --- a/docs/src/python/fast.rst +++ b/docs/src/python/fast.rst @@ -10,6 +10,7 @@ Fast rms_norm layer_norm + cross_entropy rope scaled_dot_product_attention metal_kernel diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index a82c5ad6e9..51421cd728 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -19,6 +19,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cu ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp diff --git a/mlx/backend/cuda/cross_entropy.cu b/mlx/backend/cuda/cross_entropy.cu new file mode 100644 index 0000000000..23b0689be3 --- /dev/null +++ b/mlx/backend/cuda/cross_entropy.cu @@ -0,0 +1,245 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/device/cast_op.cuh" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" +#include "mlx/fast_primitives.h" + +#include +#include +#include + +#include + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +// fused together logsumexp + gather +// cast to float32 inside the kernel +// to avoid logits.astype(mx.float32) +// for each row: loss = logsumexp(x) - x_t +// first we accumulate logsumexp, then we do a gather +template +__global__ void cross_entropy( + const T* x, // [M, N] + const int* y, // [M,] + float* loss, // [M,] <- will be always in fp32 lse - x + int axis_size // N +) { + cg::greater max_op; + cg::plus plus_op; + + float prevmax; + float curmax = Limits::finite_min(); + float normalizer = 0; + + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + x += grid.block_rank() * axis_size; // offset input + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto vals = load_vector(x, index, axis_size, Limits::min()); + prevmax = curmax; +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + curmax = max_op(curmax, static_cast(vals[i])); + } + // scale already accumulated normiliser + normalizer = normalizer * __expf(prevmax - curmax); + // add vals scaled by curmax +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + normalizer += __expf(static_cast(vals[i]) - curmax); + } + } + prevmax = curmax; + curmax = cg::reduce(warp, curmax, max_op); + normalizer = normalizer * __expf(prevmax - curmax); + normalizer = cg::reduce(warp, normalizer, plus_op); + // second reduce in a block + __shared__ float warp_max[WARP_SIZE]; + __shared__ float warp_normaliser[WARP_SIZE]; + + if (warp.thread_rank() == 0) { + warp_max[warp.meta_group_rank()] = curmax; + warp_normaliser[warp.meta_group_rank()] = normalizer; + } + block.sync(); + bool is_valid = warp.thread_rank() < warp.meta_group_size(); + curmax = + is_valid ? warp_max[warp.thread_rank()] : Limits::finite_min(); + prevmax = curmax; + curmax = + cg::reduce(warp, curmax, max_op); // max within a block (global row max) + normalizer = is_valid ? warp_normaliser[warp.thread_rank()] : 0.0f; + normalizer = normalizer * __expf(prevmax - curmax); + normalizer = cg::reduce(warp, normalizer, plus_op); + // gather and writing the output: + auto row = grid.block_rank(); + if (block.thread_rank() == 0) { + float gap = curmax - static_cast(x[y[row]]); + loss[row] = isinf(curmax) ? gap : log(normalizer) + gap; + } +} + +// get loss from the forward +template +__global__ void cross_entropy_vjp( + const T* x, // [M, N] + const int* y, // [M,] + const float* loss, // [M,] + const float* gy, // cotangent [M,] + T* grads, // [M, N] lse is accumulated in float, x is casted to float + int axis_size // N +) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto row = grid.block_rank(); + + x += row * axis_size; // offset input + grads += row * axis_size; // offset output + auto y_n = y[row]; // target index [0, N) + auto g = gy[row]; // cotangent + auto loss_n = loss[row]; + auto x_t = static_cast(x[y_n]); + block.sync(); + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { + auto index = r * BLOCK_DIM + block.thread_rank(); // [0, N) + auto vals = load_vector(x, index, axis_size, T{}); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + int col = index * N_READS + i; + float val = __expf((static_cast(vals[i]) - x_t) - loss_n); + vals[i] = static_cast(g * (val - (col == y_n ? 1.0f : 0.0f))); + } + store_vector(grads, index, vals, axis_size); + } +} +} // namespace cu + +namespace fast { + +bool CrossEntropy::use_fallback(Stream s) { + return s.device == Device::cpu; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("CrossEntropy::eval_gpu"); + assert(inputs.size() == 2); // logits and target + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + auto ensure_row_contiguous = [&s, &encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + encoder.add_temporary(x_copy); + return x_copy; + } + }; + auto in = ensure_row_contiguous(inputs[0]); // [n_rows, V] + auto target = ensure_row_contiguous(inputs[1]); // [n_rows,] + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows] in fp32 + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + encoder.set_input_array(in); + encoder.set_input_array(target); + encoder.set_output_array(out); + dispatch_float_types(in.dtype(), "cross_entropy", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) { + auto kernel = cu::cross_entropy; + encoder.add_kernel_node( + kernel, + n_rows, + block_dim(), + gpu_ptr(in), + gpu_ptr(target), + gpu_ptr(out), + axis_size); + }); + }); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("CrossEntropyVJP::eval_gpu"); + assert(inputs.size() == 4); // logits, target, loss, cotangent + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + auto ensure_row_contiguous = [&s, &encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + encoder.add_temporary(x_copy); + return x_copy; + } + }; + + auto check_input = [&s](const array& x, bool& copied) { + if (x.flags().row_contiguous) { + copied = false; + return x; + } + copied = true; + return contiguous_copy_gpu(x, s); + }; + bool donate_x = inputs[0].is_donatable(); + bool copied; + auto in = check_input(inputs[0], copied); // [n_rows, V] + donate_x |= copied; + auto target = ensure_row_contiguous(inputs[1]); // [n_rows,] + auto loss = ensure_row_contiguous(inputs[2]); // [n_rows,] fp32 + auto cotan = ensure_row_contiguous(inputs[3]); // [n_rows,] fp32 + if (donate_x) { + out.copy_shared_buffer(in); + } else { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows, V] + } + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + encoder.set_input_array(in); + encoder.set_input_array(target); + encoder.set_input_array(loss); + encoder.set_input_array(cotan); + encoder.set_output_array(out); + dispatch_float_types(in.dtype(), "cross_entropy_vjp", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) { + auto kernel = cu::cross_entropy_vjp; + encoder.add_kernel_node( + kernel, + n_rows, + block_dim(), + gpu_ptr(in), + gpu_ptr(target), + gpu_ptr(loss), + gpu_ptr(cotan), + gpu_ptr(out), + axis_size); + }); + }); +} + +} // namespace fast + +} // namespace mlx::core diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index 45929e27dd..d1d0e781cc 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -13,6 +13,7 @@ #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/utils.h" #include "mlx/dtype_utils.h" +#include "mlx/fast_primitives.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" #include "mlx/utils.h" @@ -214,4 +215,26 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } +namespace fast { + +// There is no fused Metal cross entropy kernel yet +bool CrossEntropy::use_fallback(Stream s) { + return true; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + throw std::runtime_error("[CrossEntropy::eval_gpu] Metal cross entropy NYI."); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + throw std::runtime_error( + "[CrossEntropyVJP::eval_gpu] Metal cross entropy NYI."); +} + +} // namespace fast + } // namespace mlx::core diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 7f60d0d83a..f17f12cfaf 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -170,6 +170,8 @@ NO_GPU(View) NO_GPU(MaskedScatter) namespace fast { +NO_GPU_USE_FALLBACK(CrossEntropy) +NO_GPU_MULTI(CrossEntropyVJP) NO_GPU_USE_FALLBACK(LayerNorm) NO_GPU_MULTI(LayerNormVJP) NO_GPU_USE_FALLBACK(RMSNorm) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index df2beebd84..f45724dc92 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -187,6 +187,107 @@ bool RMSNormVJP::is_equivalent(const Primitive& other) const { return eps_ == a_other.eps_; } +array cross_entropy( + const array& logits, + const array& targets, + StreamOrDevice s_ /* = {} */) { + if (logits.ndim() < 1) { + throw std::invalid_argument( + "[cross_entropy] logits must have at least 1 dimension but got input " + "with 0 dimensions."); + } + auto expected = logits.shape(); + expected.pop_back(); + if (targets.shape() != expected) { + std::ostringstream msg; + msg << "[cross_entropy] targets shape " << targets.shape() + << " does not match logits shape " << logits.shape() + << " with the last axis removed."; + throw std::invalid_argument(msg.str()); + } + if (!issubdtype(logits.dtype(), floating)) { + std::ostringstream msg; + msg << "[cross_entropy] Received unsupported logits type " << logits.dtype() + << "."; + throw std::invalid_argument(msg.str()); + } + if (!issubdtype(targets.dtype(), integer)) { + std::ostringstream msg; + msg << "[cross_entropy] targets must be integer class indices but got " + << targets.dtype() << "."; + throw std::invalid_argument(msg.str()); + } + + auto s = to_stream(s_); + auto fallback = [s](const std::vector& inputs) { + auto& x = inputs[0]; + auto& y = inputs[1]; + auto score = + squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); + auto loss = subtract(logsumexp(x, -1, /* keepdims= */ false, s), score, s); + return std::vector{astype(loss, float32, s)}; + }; + + auto passed_targets = astype(targets, int32, s); + + if (!CrossEntropy::use_fallback(s)) { + return array( + expected, + float32, + std::make_shared(s, fallback), + {logits, passed_targets}); + } + return fallback({logits, passed_targets})[0]; +} + +std::vector CrossEntropy::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) { + assert(primals.size() == 2); + assert(outputs.size() == 1); + assert(cotangents.size() == 1); + + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[cross_entropy] Cannot differentiate with respect to the targets."); + } + } + + auto s = stream(); + auto fallback = [s](const std::vector& inputs) { + auto& x = inputs[0]; + auto& y = inputs[1]; + auto& loss = inputs[2]; + auto& g = inputs[3]; + + auto score = + squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); + auto lse = add(loss, astype(score, float32, s), s); + auto p = + exp(subtract(astype(x, float32, s), expand_dims(lse, -1, s), s), s); + Shape class_shape(x.ndim(), 1); + class_shape.back() = x.shape(-1); + auto onehot = astype( + equal( + expand_dims(y, -1, s), + reshape(arange(x.shape(-1), y.dtype(), s), class_shape, s), + s), + float32, + s); + auto gx = multiply(expand_dims(g, -1, s), subtract(p, onehot, s), s); + return std::vector{astype(gx, x.dtype(), s)}; + }; + + return {array( + primals[0].shape(), + primals[0].dtype(), + std::make_shared(s, fallback), + {primals[0], primals[1], outputs[0], cotangents[0]})}; +} + array layer_norm( const array& x, const std::optional& weight, diff --git a/mlx/fast.h b/mlx/fast.h index c5f664df79..b35f354377 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -24,6 +24,10 @@ MLX_API array layer_norm( float eps, StreamOrDevice s = {}); +/** Fused cross entropy with class indices as targets. */ +MLX_API array +cross_entropy(const array& logits, const array& targets, StreamOrDevice s = {}); + MLX_API array rope( const array& x, int dims, diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 61a392e418..63021cd0c3 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -97,6 +97,67 @@ class RMSNormVJP : public Custom { float eps_; }; +// loss is always fp32 and the logits never have to be upcast in the graph. +class CrossEntropy : public Custom { + public: + CrossEntropy( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + static bool use_fallback(Stream stream); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + + DEFINE_NAME(CrossEntropy) + bool is_equivalent(const Primitive& other) const override { + return true; + } + std::vector output_shapes(const std::vector& inputs) override { + return {inputs[1].shape()}; + } + + auto state() const { + return std::monostate{}; + } +}; + +class CrossEntropyVJP : public Custom { + public: + CrossEntropyVJP( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + DEFINE_NAME(CrossEntropyVJP) + bool is_equivalent(const Primitive& other) const override { + return true; + } + DEFINE_INPUT_OUTPUT_SHAPE() + + auto state() const { + return std::monostate{}; + } +}; + class LayerNorm : public Custom { public: LayerNorm( diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 184df2a2e0..b98d2765d6 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -63,6 +63,18 @@ def cross_entropy( >>> targets = mx.array([[0.9, 0.1], [0.1, 0.9]]) >>> nn.losses.cross_entropy(logits, targets) array([0.348587, 0.348587], dtype=float32) + >>> + >>> # Half precision logits with class indices as targets. On CUDA a + >>> # fused kernel accumulates the reduction in float32: + >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]], mx.bfloat16) + >>> targets = mx.array([0, 1]) + >>> nn.losses.cross_entropy(logits, targets) + array([0.0485873, 0.0485873], dtype=float32) + >>> + >>> # Metal and the CPU reduce in the dtype of the logits, so upcast + >>> # them to get the same accuracy: + >>> nn.losses.cross_entropy(logits.astype(mx.float32), targets) + array([0.0485873, 0.0485873], dtype=float32) """ if label_smoothing < 0 or label_smoothing >= 1: raise ValueError(f"Label smoothing must be in [0, 1), got {label_smoothing}.") @@ -83,31 +95,38 @@ def _drop_dim(shape, axis): f"Targets shape {targets.shape} does not match logits shape {logits.shape}." ) - # Shift by the max first. The loss only depends on differences between - # logits, but subtracting the logsumexp of large logits loses the gap to - # rounding before the subtraction happens. - logits = logits - mx.stop_gradient(mx.max(logits, axis=axis, keepdims=True)) + use_fast = ( + mx.cuda.is_available() + and mx.default_device() == mx.gpu + and not targets_as_probs + and label_smoothing == 0 + and axis in (-1, logits.ndim - 1) + and mx.issubdtype(logits.dtype, mx.floating) + and mx.issubdtype(targets.dtype, mx.integer) + ) - if targets_as_probs: - score = mx.sum(logits * targets, axis=axis) + if use_fast: + loss = mx.fast.cross_entropy(logits, targets).astype(logits.dtype) else: - score = mx.take_along_axis(logits, mx.expand_dims(targets, axis), axis).squeeze( - axis - ) + logits = logits - mx.stop_gradient(mx.max(logits, axis=axis, keepdims=True)) - logsumexp_logits = mx.logsumexp(logits, axis=axis) - if label_smoothing > 0: - # Adjust the true class score with label smoothing - adjusted_score = (1 - label_smoothing) * score + if targets_as_probs: + score = mx.sum(logits * targets, axis=axis) + else: + score = mx.take_along_axis( + logits, mx.expand_dims(targets, axis), axis + ).squeeze(axis) - # Calculate the mean logit across the classes for smoothed loss - mean_logits = logits.mean(axis=axis) - smoothed_loss = -mean_logits * label_smoothing + logsumexp_logits = mx.logsumexp(logits, axis=axis) + if label_smoothing > 0: + adjusted_score = (1 - label_smoothing) * score - # Combine the adjusted score and smoothed loss with the logsumexp logits - loss = logsumexp_logits - adjusted_score + smoothed_loss - else: - loss = logsumexp_logits - score + mean_logits = logits.mean(axis=axis) + smoothed_loss = -mean_logits * label_smoothing + + loss = logsumexp_logits - adjusted_score + smoothed_loss + else: + loss = logsumexp_logits - score # Apply weights if provided if weights is not None: diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 0a50dc79cd..67c3442cff 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -174,6 +174,36 @@ void init_fast(nb::module_& parent_module) { array: The output array. )pbdoc"); + m.def( + "cross_entropy", + &mx::fast::cross_entropy, + "logits"_a, + "targets"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def cross_entropy(logits: array, targets: array, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Cross entropy loss with class indices as targets. + + Computes ``logsumexp(logits, axis=-1) - logits[..., target]`` in a + fused kernel with accumulation in float32. + + Note: Currently is implemented only on CUDA, fallback to unfused version with + manual casting on Metal and CPU. + + Args: + logits (array): The unnormalized logits. The loss is computed over + the last axis. + targets (array): Class indices. The shape should match the shape of + ``logits`` with the last axis removed. The indices must be in + ``[0, logits.shape[-1])``. + + Returns: + array: The per-element loss in float32, with the shape of + ``targets``. + )pbdoc"); + m.def( "rope", [](const mx::array& a, diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index ba5b8f3138..200781d372 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -525,6 +525,67 @@ def inner(x, w, y): self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) + def test_cross_entropy(self): + def cross_entropy_ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits.astype(mx.float32), axis=-1) - score.astype( + mx.float32 + ) + + tolerances = {mx.float32: 1e-5, mx.float16: 3e-2, mx.bfloat16: 3e-1} + + for V in [7, 32, 128, 255, 256, 1000, 4096, 8192]: + for dtype in [mx.float32, mx.float16, mx.bfloat16]: + logits = (mx.random.normal(shape=(4, 7, V), scale=3.0) * 2).astype( + dtype + ) + targets = mx.random.randint(0, V, shape=(4, 7)) + expected = cross_entropy_ref(logits, targets) + out = mx.fast.cross_entropy(logits, targets) + self.assertEqual(out.dtype, mx.float32) + self.assertEqual(out.shape, targets.shape) + self.assertLess(mx.abs(out - expected).max().item(), tolerances[dtype]) + + def test_cross_entropy_shape_checks(self): + logits = mx.random.normal(shape=(4, 16)) + with self.assertRaises(ValueError): + mx.fast.cross_entropy(logits, mx.zeros((5,), mx.int32)) + with self.assertRaises(ValueError): + # Probability targets are not supported by the fused op. + mx.fast.cross_entropy(logits, mx.zeros((4, 16), mx.int32)) + with self.assertRaises(ValueError): + mx.fast.cross_entropy(logits, mx.zeros((4,), mx.float32)) + + def test_cross_entropy_grad(self): + def ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits, axis=-1) - score + + f1 = lambda x, y: ref(x, y).mean() + f2 = lambda x, y: mx.fast.cross_entropy(x, y).mean() + + for V in [7, 128, 1000, 4096]: + logits = mx.random.normal(shape=(4, 7, V), scale=2.0) + targets = mx.random.randint(0, V, shape=(4, 7)) + g1 = mx.grad(f1, argnums=0)(logits, targets) + g2 = mx.grad(f2, argnums=0)(logits, targets) + self.assertEqual(g2.shape, logits.shape) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) + + w = mx.random.uniform(shape=(4, 7)) + f3 = lambda x, y: (ref(x, y) * w).sum() + f4 = lambda x, y: (mx.fast.cross_entropy(x, y) * w).sum() + logits = mx.random.normal(shape=(4, 7, 512), scale=2.0) + targets = mx.random.randint(0, 512, shape=(4, 7)) + g1 = mx.grad(f3, argnums=0)(logits, targets) + g2 = mx.grad(f4, argnums=0)(logits, targets) + self.assertEqual(g2.shape, logits.shape) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) + def test_layer_norm_dim_check(self): with self.assertRaises(ValueError): weight = mx.ones((129,)) From 29b61c0b5eb629d175a3bbff585a80a70e4b8105 Mon Sep 17 00:00:00 2001 From: Ayaan Gazali Date: Sun, 23 Aug 2026 19:39:14 -0700 Subject: [PATCH 74/84] Fix cpu exclusive scan for complex numbers (#4272) Co-authored-by: Cheng --- mlx/backend/cpu/scan.cpp | 23 ++++++++++++++++------- python/tests/test_ops.py | 16 ++++++++++++++++ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/mlx/backend/cpu/scan.cpp b/mlx/backend/cpu/scan.cpp index 93e67825e2..ab2ca14366 100644 --- a/mlx/backend/cpu/scan.cpp +++ b/mlx/backend/cpu/scan.cpp @@ -191,6 +191,19 @@ void scan_op( } } +template +U scan_init(const Dtype& dtype, bool maximum) { + constexpr auto inf = std::numeric_limits::infinity(); + if constexpr (std::is_same_v) { + return maximum ? complex64_t{inf, inf} : complex64_t{-inf, -inf}; + } else if (issubdtype(dtype, floating)) { + return maximum ? static_cast(inf) : static_cast(-inf); + } else { + return maximum ? std::numeric_limits::max() + : std::numeric_limits::min(); + } +} + template void scan_dispatch( Scan::ReduceType rtype, @@ -221,9 +234,7 @@ void scan_dispatch( } return x < y ? x : y; }; - auto init = (issubdtype(in.dtype(), floating)) - ? static_cast(std::numeric_limits::infinity()) - : std::numeric_limits::max(); + auto init = scan_init(in.dtype(), /* maximum = */ true); scan_op(in, out, axis, reverse, inclusive, op, init); break; } @@ -236,9 +247,7 @@ void scan_dispatch( } return x < y ? y : x; }; - auto init = (issubdtype(in.dtype(), floating)) - ? static_cast(-std::numeric_limits::infinity()) - : std::numeric_limits::min(); + auto init = scan_init(in.dtype(), /* maximum = */ false); scan_op(in, out, axis, reverse, inclusive, op, init); break; } @@ -246,7 +255,7 @@ void scan_dispatch( auto op = [](U a, T b) { return detail::LogAddExp{}(a, static_cast(b)); }; - auto init = (issubdtype(in.dtype(), floating)) + auto init = (issubdtype(in.dtype(), inexact)) ? static_cast(-std::numeric_limits::infinity()) : std::numeric_limits::min(); scan_op(in, out, axis, reverse, inclusive, op, init); diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 9bcce60fce..0c94989e02 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2691,6 +2691,22 @@ def test_scan_size_one_axis(self): out = getattr(mx, op)(a, axis=0) self.assertTrue(np.array_equal(np.array(out), expected)) + def test_scans_complex_exclusive(self): + a = mx.array([-3 + 1j, -1 + 2j, -4 + 0j, 0 + 5j, 2 - 1j]) + for op in ("cummax", "cummin", "logcumsumexp"): + mxop = getattr(mx, op) + for reverse in (False, True): + inclusive = mxop(a, axis=0, inclusive=True, reverse=reverse) + exclusive = mxop(a, axis=0, inclusive=False, reverse=reverse) + if reverse: + got, want = exclusive[:-1], inclusive[1:] + else: + got, want = exclusive[1:], inclusive[:-1] + self.assertTrue( + mx.allclose(got, want), + msg=f"{op} reverse={reverse}", + ) + def test_cummax_cummin_nan(self): nan = float("nan") cases = [ From 9d173ff2d4f3b98d8e5b1ea7fed1ab603c9bb952 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Sun, 23 Aug 2026 19:44:07 -0700 Subject: [PATCH 75/84] Support Relocatable CUDA DLLs on Windows (#4382) --- mlx/backend/cuda/CMakeLists.txt | 35 +++++++++++++++++++++++++++------ mlx/backend/cuda/delayload.cpp | 26 ++++++++++++++++-------- mlx/backend/cuda/dirs.cpp | 24 ++++++++++++++++++++-- 3 files changed, 69 insertions(+), 16 deletions(-) diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 51421cd728..9c8d3174c1 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -75,6 +75,7 @@ target_sources( # Put dynamic defines in the dirs.cpp file. add_library(mlx_dirs OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/dirs.cpp) +target_include_directories(mlx_dirs PRIVATE "${PROJECT_SOURCE_DIR}") target_link_libraries(mlx PRIVATE $) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/binary) @@ -178,8 +179,32 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}") -# Search CUDA libs from installed python packages. +# Configure Windows CUDA DLL loading. if(WIN32) + set(MLX_CUDA_BIN_DIR + "" + CACHE STRING "Directory containing CUDA DLLs for Windows delay-loading") + set(MLX_CUDNN_BIN_DIR + "" + CACHE STRING "Directory containing cuDNN DLLs for Windows delay-loading") + + # With MLX_LOAD_CUDA_LIBS_FROM_PYTHON, unset dirs use the wheel layout. + # Relative dirs are resolved from the MLX binary. + if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON) + if("${MLX_CUDA_BIN_DIR}" STREQUAL "") + set(MLX_CUDA_BIN_DIR "${CUDAToolkit_BIN_DIR}/x64") + endif() + if("${MLX_CUDNN_BIN_DIR}" STREQUAL "") + set(MLX_CUDNN_BIN_DIR "${CUDNN_BIN_DIR}") + endif() + endif() + + function(mlx_add_dir_definition name) + if(NOT "${${name}}" STREQUAL "") + target_compile_definitions(mlx_dirs PRIVATE ${name}="${${name}}") + endif() + endfunction() + # Resolve paths of unfound DLL at runtime. if(BUILD_SHARED_LIBS) target_link_libraries(mlx PRIVATE "delayimp.lib") @@ -200,11 +225,9 @@ if(WIN32) target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}") endforeach() # Pass the locations where CUDA DLLs are placed. - if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON) - target_compile_definitions( - mlx_dirs PRIVATE MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64" - MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}") - endif() + foreach(dir_var MLX_CUDA_BIN_DIR MLX_CUDNN_BIN_DIR) + mlx_add_dir_definition(${dir_var}) + endforeach() else() # For POSIX we rely on RPATH to search for CUDA libs. if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON) diff --git a/mlx/backend/cuda/delayload.cpp b/mlx/backend/cuda/delayload.cpp index aba7566c5b..7092d4497c 100644 --- a/mlx/backend/cuda/delayload.cpp +++ b/mlx/backend/cuda/delayload.cpp @@ -20,23 +20,31 @@ inline fs::path relative_to_current_binary(const char* relative) { } inline fs::path cublas_dir() { - return cuda_bin_dir() ? fs::path(cuda_bin_dir()) - : relative_to_current_binary("../nvidia/cublas/bin"); + if (const char* dir = cuda_bin_dir()) { + return fs::path(dir); + } + return relative_to_current_binary("../nvidia/cublas/bin"); } fs::path load_nvrtc() { - fs::path nvrtc_dir = cuda_bin_dir() - ? fs::path(cuda_bin_dir()) - : relative_to_current_binary("../nvidia/cuda_nvrtc/bin"); + fs::path nvrtc_dir; + if (const char* dir = cuda_bin_dir()) { + nvrtc_dir = fs::path(dir); + } else { + nvrtc_dir = relative_to_current_binary("../nvidia/cuda_nvrtc/bin"); + } // Internally nvrtc loads some libs dynamically, add to search dirs. ::AddDllDirectory(nvrtc_dir.c_str()); return nvrtc_dir; } fs::path load_cudnn() { - fs::path cudnn_dir = cudnn_bin_dir() - ? fs::path(cudnn_bin_dir()) - : relative_to_current_binary("../nvidia/cudnn/bin"); + fs::path cudnn_dir; + if (const char* dir = cudnn_bin_dir()) { + cudnn_dir = fs::path(dir); + } else { + cudnn_dir = relative_to_current_binary("../nvidia/cudnn/bin"); + } // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would // get errors like "Invalid handle. Cannot load symbol cudnnCreate". for (const auto& dll : fs::directory_iterator(cudnn_dir)) { @@ -66,6 +74,8 @@ FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) { } else if (dll.starts_with("nvrtc")) { static auto nvrtc_dir = load_nvrtc(); mod = ::LoadLibraryW((nvrtc_dir / dll).c_str()); + } else if (const char* dir = cuda_bin_dir()) { + mod = ::LoadLibraryW((fs::path(dir) / dll).c_str()); } } return reinterpret_cast(mod); diff --git a/mlx/backend/cuda/dirs.cpp b/mlx/backend/cuda/dirs.cpp index a9d33b4790..bd24853dd8 100644 --- a/mlx/backend/cuda/dirs.cpp +++ b/mlx/backend/cuda/dirs.cpp @@ -1,6 +1,24 @@ // Copyright © 2026 Apple Inc. +#include "mlx/backend/common/utils.h" + +#include +#include + namespace mlx::core::cu { +namespace { + +namespace fs = std::filesystem; + +std::string resolve_bin_dir(const char* dir) { + fs::path path(dir); + if (path.is_absolute()) { + return path.string(); + } + return fs::absolute(current_binary_dir() / path).string(); +} + +} // namespace const char* cccl_dir() { #if defined(MLX_CCCL_DIR) @@ -12,7 +30,8 @@ const char* cccl_dir() { const char* cuda_bin_dir() { #if defined(MLX_CUDA_BIN_DIR) - return MLX_CUDA_BIN_DIR; + static const std::string dir = resolve_bin_dir(MLX_CUDA_BIN_DIR); + return dir.c_str(); #else return nullptr; #endif @@ -20,7 +39,8 @@ const char* cuda_bin_dir() { const char* cudnn_bin_dir() { #if defined(MLX_CUDNN_BIN_DIR) - return MLX_CUDNN_BIN_DIR; + static const std::string dir = resolve_bin_dir(MLX_CUDNN_BIN_DIR); + return dir.c_str(); #else return nullptr; #endif From c793734eb715dbcfdb1ced58e348ec53c2d7ed85 Mon Sep 17 00:00:00 2001 From: katlun-lgtm Date: Sun, 23 Aug 2026 22:44:56 -0400 Subject: [PATCH 76/84] Use cast_to for fused AsType in compiled Metal kernels (#4351) Co-authored-by: katlun-lgtm Co-authored-by: Cheng --- mlx/backend/metal/compiled.cpp | 2 +- python/tests/test_compile.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/compiled.cpp b/mlx/backend/metal/compiled.cpp index cda06143d6..e95d7b8d5f 100644 --- a/mlx/backend/metal/compiled.cpp +++ b/mlx/backend/metal/compiled.cpp @@ -208,7 +208,7 @@ inline void build_kernel( " {0} tmp_{1} = ", get_type_string(x.dtype()), namer.get_name(x)); if (is_static_cast(x.primitive())) { os += fmt::format( - "static_cast<{0}>(tmp_{1});\n", + "cast_to<{0}>(tmp_{1});\n", get_type_string(x.dtype()), namer.get_name(x.inputs()[0])); } else { diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 5eaa6cb955..663ec295b0 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1623,6 +1623,16 @@ def test_compile_abs_unsigned(self): x = mx.array([1, 2, 3], dtype) self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + def test_compiled_subnormal_bool_cast(self): + f32_sub = mx.array(np.array([0x00000001] * 4, dtype=np.uint32)).view(mx.float32) + f16_sub = mx.array(np.array([0x0001] * 4, dtype=np.uint16)).view(mx.float16) + bf16_sub = mx.array(np.array([0x0001] * 4, dtype=np.uint16)).view(mx.bfloat16) + + # A single-op compile does not fuse; the fused path needs >= 2 ops. + fn = mx.compile(lambda x: mx.broadcast_to(x, (2, 4)).astype(mx.bool_)) + for sub in (f32_sub, f16_sub, bf16_sub): + self.assertTrue(mx.all(fn(sub)).item()) + def test_compile_different_log_bases(self): # The logs are intermediates, since outputs are not simplified. def entropies(p): From 9d16475965ca7f3306c3101751f48e13b1a2df33 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Sun, 23 Aug 2026 23:37:52 -0700 Subject: [PATCH 77/84] python: Declare DLPackCompatible protocol members as methods (#4384) --- python/mlx/_stub_patterns.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/mlx/_stub_patterns.txt b/python/mlx/_stub_patterns.txt index ca6213136f..90afb55981 100644 --- a/python/mlx/_stub_patterns.txt +++ b/python/mlx/_stub_patterns.txt @@ -3,8 +3,8 @@ mlx.core.__prefix__: P = ParamSpec("P") R = TypeVar("R") class DLPackCompatible(Protocol): - __dlpack__: Callable[..., Any] - __dlpack_device__: Callable[..., Any] + def __dlpack__(self, *args: Any, **kwargs: Any) -> Any: ... + def __dlpack_device__(self, *args: Any, **kwargs: Any) -> Any: ... mlx.core.__suffix__: scalar: TypeAlias = int | float | bool | complex From a9eed5a840ca776a6b287e389a6d02665cb74627 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 24 Aug 2026 20:23:24 +0900 Subject: [PATCH 78/84] Fix quantizing sliced arrays (#4381) --- mlx/backend/metal/quantized.cpp | 10 +++++----- python/tests/test_quantized.py | 12 ++++++++++++ 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 65c24152b0..1ae99a8e5e 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -249,21 +249,18 @@ void quantize_impl( auto w = ensure_row_contiguous(w_pre, d, s); if (dequantize) { auto scales = ensure_row_contiguous(inputs[1], d, s); - compute_encoder.set_input_array(w, 0); - compute_encoder.set_input_array(scales, 1); if (has_biases) { auto biases = ensure_row_contiguous(inputs[2], d, s); compute_encoder.set_input_array(biases, 2); } else if (has_global_scale) { compute_encoder.set_input_array(inputs[2], 2); } + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); compute_encoder.set_output_array(out, 3); } else { auto& scales = outputs[1]; scales.set_data(allocator::malloc(scales.nbytes())); - compute_encoder.set_input_array(w, 0); - compute_encoder.set_output_array(out, 1); - compute_encoder.set_output_array(scales, 2); if (has_biases) { auto& biases = outputs[2]; biases.set_data(allocator::malloc(biases.nbytes())); @@ -271,6 +268,9 @@ void quantize_impl( } else if (has_global_scale) { compute_encoder.set_input_array(inputs[1], 3); } + compute_encoder.set_input_array(w, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_output_array(scales, 2); } auto type_string = dequantize ? get_type_string(out.dtype()) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 28033cbbab..461175f013 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -39,6 +39,18 @@ def test_quantize_dequantize(self): a_hat = mx.dequantize(w_q, scales, biases, gs, b) self.assertTrue(mx.all(a_hat == 0)) + # slices + if mx.default_device() == mx.gpu: + w = mx.random.normal(shape=(2, 256, 32)) + quant = {"group_size": 32, "bits": 4} + wq, scales, biases = mx.quantize(w, **quant) + wq_s = wq[:, :16, :] + scales_s = scales[:, :16, :] + biases_s = biases[:, :16, :] + dq_cpu = mx.dequantize(wq_s, scales_s, biases_s, **quant, stream=mx.cpu) + dq_gpu = mx.dequantize(wq_s, scales_s, biases_s, **quant, stream=mx.gpu) + self.assertTrue(mx.abs(dq_cpu - dq_gpu).max().item() < 1e-6) + def test_mxfp4_quantize_dequantize(self): lut = mx.array( [ From 43d2f06cb87e76895bf9a152bade4fee83408643 Mon Sep 17 00:00:00 2001 From: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Date: Mon, 24 Aug 2026 05:15:56 -0700 Subject: [PATCH 79/84] Fix einsum dropping a trailing empty subscript (#4299) Co-authored-by: Cheng --- .github/actions/test-wheel/action.yml | 2 +- mlx/einsum.cpp | 12 ++++++---- python/tests/test_einsum.py | 33 +++++++++++++++++++++++++++ 3 files changed, 42 insertions(+), 5 deletions(-) diff --git a/.github/actions/test-wheel/action.yml b/.github/actions/test-wheel/action.yml index b745a4d2f9..a44d17c701 100644 --- a/.github/actions/test-wheel/action.yml +++ b/.github/actions/test-wheel/action.yml @@ -46,4 +46,4 @@ runs: echo "No matching backend wheel to install" exit 1 fi - python -m unittest discover -v python/tests + python -m tests discover -v python/tests diff --git a/mlx/einsum.cpp b/mlx/einsum.cpp index 705a5fc2cb..dcaa8c51ea 100644 --- a/mlx/einsum.cpp +++ b/mlx/einsum.cpp @@ -95,10 +95,14 @@ std::pair, std::string> parse(std::string subscripts) { std::sort(rhs.begin(), rhs.end()); } std::vector input_list; - std::stringstream ss(lhs); - std::string token; - while (getline(ss, token, ',')) { - input_list.push_back(token); + for (size_t start = 0;;) { + auto pos = lhs.find(',', start); + if (pos == std::string::npos) { + input_list.push_back(lhs.substr(start)); + break; + } + input_list.push_back(lhs.substr(start, pos - start)); + start = pos + 1; } return {input_list, rhs}; } diff --git a/python/tests/test_einsum.py b/python/tests/test_einsum.py index c87a6dc45f..08884ae981 100644 --- a/python/tests/test_einsum.py +++ b/python/tests/test_einsum.py @@ -65,6 +65,39 @@ def test_longer_paths(self): mx_path = mx.einsum_path(case, *inputs) self.assertEqual(np_path[0][1:], mx_path[0]) + def test_scalar_operands(self): + # An empty subscript is a scalar operand. A trailing one used to be + # dropped by the parser, so "i,->i" looked like a single input. + s1 = mx.array(2.0) + s2 = mx.array(3.0) + v = mx.random.uniform(shape=(3,)) + m = mx.random.uniform(shape=(2, 3)) + + cases = [ + ("->", (s1,)), + (",->", (s1, s2)), + (",,->", (s1, s2, s1)), + ("i,->i", (v, s1)), + (",i->i", (s1, v)), + ("ij,->ij", (m, s1)), + (",ij->ij", (s1, m)), + ("i,,->i", (v, s1, s2)), + ] + for spec, operands in cases: + mx_out = mx.einsum(spec, *operands) + np_out = np.einsum(spec, *[np.array(o) for o in operands]) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + # Operand count still has to match the number of subscripts + with self.assertRaises(ValueError): + mx.einsum(",->", s1) + with self.assertRaises(ValueError): + mx.einsum("i,->i", v) + # An empty subscript requires a 0-d operand + with self.assertRaises(ValueError): + mx.einsum(",->", v, s1) + def test_simple_einsum(self): a = mx.arange(4 * 4).reshape(4, 4) a_mx = mx.einsum("ii->i", a) From 768b4c587da3f5fd6d7b32f9df7d16ceaece45c5 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 24 Aug 2026 14:46:23 -0400 Subject: [PATCH 80/84] fix: resolve upstream fp qvm merge --- mlx/backend/metal/kernels/fp_quantized.h | 39 ------------------------ 1 file changed, 39 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 32e6b66701..7061771d59 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -1327,45 +1327,6 @@ template < w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid); } -template -[[kernel]] void fp_qvm( - const device uint32_t* w, - const device uint8_t* scales, - const device T* x, - device T* y, - const constant int& in_vec_size, - const constant int& out_vec_size, - const constant int& M, - const constant int& x_batch_ndims, - const constant int* x_shape, - const constant int64_t* x_strides, - const constant int& w_batch_ndims, - const constant int* w_shape, - const constant int64_t* w_strides, - const constant int64_t* s_strides, - uint3 tid [[threadgroup_position_in_grid]], - uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]]) { - if (batched) { - adjust_matrix_offsets( - x, - w, - scales, - y, - out_vec_size * M, - x_batch_ndims, - x_shape, - x_strides, - w_batch_ndims, - w_shape, - w_strides, - s_strides, - tid); - } - fp_qmv_wide_impl( - w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid); -} - template < typename T, int group_size, From 1e93028abbf42b352f61bf3b2ec3509a08c60f4e Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 24 Aug 2026 14:49:40 -0400 Subject: [PATCH 81/84] fix: remove duplicate qmv wide merge --- mlx/backend/metal/quantized.cpp | 90 --------------------------------- 1 file changed, 90 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index a8585b725a..ce5b45229c 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -628,96 +628,6 @@ void qmv_wide( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } -// affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. -inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { - return mode != "affine" || d.get_architecture_gen() >= 15; -} - -// Dispatches qmv_wide (fp modes -> fp_qmv_wide, affine -> affine_qmv_wide): -// vecs_per_tg input vectors streamed and reused per weight group. -void qmv_wide( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - int group_size, - int bits, - int M, - int N, - int K, - metal::Device& d, - const Stream& s, - const std::string& mode) { - // vecs_per_tg is the per-threadgroup input-vector tile. Each tile re-reads - // the weights, so use the fewest tiles, then the smallest tile that fills - // them. - int n_tiles = (M + 4) / 5; // ceil(M / 5); tile size caps at 5 - int vecs_per_tg = (M + n_tiles - 1) / n_tiles; - - // k_lanes: lanes reducing K per output row (32/k_lanes rows per simdgroup). - // The affine subchunk decode has enough ALU per weight load to favor more - // rows per simdgroup (kl8); the fp modes' vectorized dot is balanced at 16. - int k_lanes = mode == "affine" ? 8 : 16; - constexpr int num_simdgroups = 2; - int B = out.size() / M / N; - bool batched = B > 1; - // Output rows per threadgroup: (32 / k_lanes) per simdgroup x num_simdgroups. - int rows_per_tg = (32 / k_lanes) * num_simdgroups; - - MTL::Size group_dims(32, num_simdgroups, 1); - MTL::Size grid_dims( - (M + vecs_per_tg - 1) / vecs_per_tg, - (N + rows_per_tg - 1) / rows_per_tg, - B); - - std::string kname; - kname.reserve(64); - std::string type_string = get_type_string(x.dtype()); - concatenate( - kname, - mode + "_qmv_wide_", - type_string, - "_gs_", - group_size, - "_b_", - bits, - "_nv_", - vecs_per_tg, - "_kl_", - k_lanes, - batched ? "_batch_1" : "_batch_0"); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - "qmv_wide", - mode, - type_string, - group_size, - bits, - vecs_per_tg, - k_lanes, - batched); - - auto& compute_encoder = metal::get_command_encoder(s); - compute_encoder.set_compute_pipeline_state(kernel); - - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); - if (biases) { - compute_encoder.set_input_array(*biases, c++); - } - compute_encoder.set_input_array(x, c++); - compute_encoder.set_output_array(out, c++); - compute_encoder.set_bytes(K, c++); - compute_encoder.set_bytes(N, c++); - compute_encoder.set_bytes(M, c++); - add_strides_and_shapes(compute_encoder, !batched, x, w, scales, biases, c); - - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); -} - void qvm_split_k( const array& x, const array& w, From bc42b3dc7ed691f2abaa7dae7e18264f7695d723 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 24 Aug 2026 20:23:57 -0400 Subject: [PATCH 82/84] style: format merged MLX changes Apply the repository-pinned clang-format output to the fork-specific files carried through the upstream merge.\n\nAI assistance: OpenAI Codex ran and reviewed the mechanical formatter changes. --- mlx/backend/common/gemma4_expert_qmm.h | 15 ++++--- mlx/backend/metal/allocator.cpp | 34 ++++++++-------- mlx/backend/metal/device.cpp | 15 +++---- mlx/backend/metal/kernels/quantized.h | 18 +++------ mlx/backend/metal/quantized.cpp | 21 ++++------ tests/gpu_tests.cpp | 55 +++++++++++--------------- 6 files changed, 65 insertions(+), 93 deletions(-) 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 5710f80958..c2b9d66b3d 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -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); @@ -161,10 +161,10 @@ Buffer MetalAllocator::malloc(size_t 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_); @@ -173,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/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/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. From 6d12ff26e432bff68a49c22dca5a39de0c41f274 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 24 Aug 2026 20:23:57 -0400 Subject: [PATCH 83/84] test: synchronize shared-buffer teardown Ensure Metal completion handlers release their array references before the doctest context ends. This prevents the custom buffer deleter from invoking doctest assertions after context.run() returns.\n\nAI assistance: OpenAI Codex was used to diagnose the crash and validate this fix. --- tests/array_tests.cpp | 1 + 1 file changed, 1 insertion(+) 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") { From f647d2435f2e613388012769a1c3fa029071c5b4 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 26 Aug 2026 12:05:04 -0400 Subject: [PATCH 84/84] fix: remove duplicate expert QMM device accessors --- mlx/backend/metal/device.h | 40 -------------------------------------- 1 file changed, 40 deletions(-) diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index c3cd3504fa..83faa6ff33 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -241,46 +241,6 @@ class MLX_API Device { return residency_sets_; } - bool gemma4_expert_qmm_requested() const { - return gemma4_expert_qmm_requested_; - } - - // MLX_GATHER_QMM_EXPERT_SLICES=trust: skip the descriptor-retract - // readback in the expert-tile route (no mid-eval stream drain). The - // caller asserts sorted indices are machine-guaranteed; a violation - // yields undefined tile output instead of the legacy fallback. - bool gemma4_expert_qmm_trust_sorted() const { - return gemma4_expert_qmm_trust_sorted_; - } - - 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(); - } - private: NS::SharedPtr build_library_( const std::string& source_string,