Merge upstream MLX and preserve shared-buffer teardown safety - #11
Open
glg2672 wants to merge 89 commits into
Open
Merge upstream MLX and preserve shared-buffer teardown safety#11glg2672 wants to merge 89 commits into
glg2672 wants to merge 89 commits into
Conversation
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.
…t-trim Bound Metal buffer COUNT, not just bytes, in MetalAllocator
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.
…top of 0.32.0 (this branch supersedes it)
Upgrade MLX core 0.31.1 → 0.32.0 (M5 Neural-Accelerator nax fixes) + keep resource-count trim
…ptor builder (Layr-Labs#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<bool> 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.
…rkbloom-base mirror (Layr-Labs#7) * perf(metal): instantiate E=256 expert-tile route for Qwen 3.5/3.6 MoE prefill (mirror of Cmlx/mlx 58fab46) * fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval (mirror) * perf(metal): trust mode skips retract readback (mirror)
Co-authored-by: Cheng <git@zcbenz.com>
…ore#4222) Co-authored-by: codeAnqiang-ma <273298913+codeAnqiang-ma@users.noreply.github.com> Co-authored-by: Cheng <git@zcbenz.com>
…-explore#4125) Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: katlun-lgtm <264247399+katlun-lgtm@users.noreply.github.com> Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Fu Xiaonan <214359569+FU-max-boop@users.noreply.github.com>
…rce (ml-explore#4273) Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Feli <feli@hnu.edu.cn> Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Cheng <git@zcbenz.com>
…ml-explore#4281) Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Cheng <git@zcbenz.com>
…4266) Co-authored-by: Cheng <git@zcbenz.com>
…lore#4284) Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Daniel Hiltgen <daniel.hiltgen@ollama.com> Co-authored-by: Cheng <git@zcbenz.com>
…e#4356) Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: Cheng <git@zcbenz.com>
Co-authored-by: katlun-lgtm <katlun@windyviews.com> Co-authored-by: Cheng <zcbenz@gmail.com>
Co-authored-by: Cheng <git@zcbenz.com>
…tream # Conflicts: # mlx/backend/metal/device.h # mlx/backend/metal/kernels/fp_quantized.h # mlx/backend/metal/kernels/quantized.h # mlx/backend/metal/quantized.cpp # python/tests/test_quantized.py
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.
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.
# Conflicts: # .github/actions/test-linux/action.yml # .github/actions/test-wheel/action.yml # .github/actions/test-windows/action.yml # mlx/backend/common/gemma4_expert_qmm.h # mlx/backend/metal/allocator.cpp # mlx/backend/metal/device.cpp # mlx/backend/metal/kernels/quantized.h # mlx/backend/metal/quantized.cpp # mlx/compile.cpp # mlx/compile_impl.h # tests/gpu_tests.cpp
glg2672
marked this pull request as ready for review
August 26, 2026 16:59
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation
MLX CMake build/tests passed.
Repeated shared-buffer host runs passed.
git diff --checkpassed.☑️ I understand it is strictly prohibited to use AI to write PR description
AI usage disclosure:
AI Was used to help resolve merge conflicts and create tests