From e49b83da01cc7bf39b1b3350521c34cd6e998b6a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 09:17:58 +0900 Subject: [PATCH 1/6] update: advance MLX pin and reconcile backend ownership Advance MLX to 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e, rebase changed CPU/CUDA/Metal overlays against pristine upstream sources, retire the CUDA teardown overlay superseded by upstream #4514, and retain the gathered-row correctness fix while using upstream's Xcode 27 workaround. Retarget the ROCm core overlay and adapt move-only allocation ownership, custom releases and primitive APIs without changing the fork revision. Add host allocation lifetime contracts and preserve reviewed macOS filesystem test portability fixes for candidate CI. Document every retained overlay and hardware validation boundaries. Validation: 27 MLX pin/parser tests, version/dtype/dispatch/llama contracts, 123 Python tooling tests and formatting pass. Native CUDA compilation and candidate Metal execution are coordinated separately; ROCm GPU validation requires an AMD host. Refs #2303 --- Makefile | 1 + ...da-handle-destructor-throws-at-teardown.md | 4 +- .../mlx-pin-243bdf09-reconciliation.md | 66 +++++ .../mlxcelverse/allocator_ownership_test.cpp | 162 +++++++++++++ .../mlxcelverse/allocator_ownership_test.sh | 36 +++ src/lib/mlx-cpp/CMakeLists.txt | 13 +- src/lib/mlx-cpp/patches-cuda/fast.cpp | 65 ++++- src/lib/mlx-cpp/patches-cuda/ops.cpp | 227 +++++++++++++++--- src/lib/mlx-cpp/patches-rocm/CMakeLists.txt | 2 +- .../mlx-cpp/patches-rocm/CORE_RESIDUAL.diff | 52 ++-- src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md | 7 +- src/lib/mlx-cpp/patches-rocm/README.md | 4 +- src/lib/mlx-cpp/patches-rocm/UPSTREAM | 2 +- src/lib/mlx-cpp/patches-rocm/mlx/allocator.h | 129 ++++++++++ .../mlx/backend/common/compiled.cpp | 2 +- .../mlx/backend/common/compiled.h | 7 +- .../mlx/backend/gpu/primitives.cpp | 20 +- .../mlx/backend/rocm/allocator.cpp | 8 +- .../patches-rocm/mlx/backend/rocm/allocator.h | 2 +- .../patches-rocm/mlx/backend/rocm/binary.hip | 1 + .../mlx/backend/rocm/compiled.cpp | 14 +- .../mlx/backend/rocm/device/binary_ops.hpp | 17 +- .../patches-rocm/mlx/backend/rocm/fft.hip | 3 + .../mlx/backend/rocm/gemms/naive_gemm.hip | 6 +- .../mlx/backend/rocm/moe_swiglu.cpp | 9 +- .../mlx/backend/rocm/primitives.cpp | 11 + .../mlx/backend/rocm/reduce/reduce_utils.hpp | 3 +- .../rocm/scaled_dot_product_attention.cpp | 6 +- src/lib/mlx-cpp/patches-rocm/mlx/compile.cpp | 52 ++-- src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp | 63 ++++- .../patches-rocm/mlx/fast_primitives.h | 42 +++- src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp | 225 ++++++++++++++--- .../mlx-cpp/patches-rocm/mlx/primitives.cpp | 187 ++++++++++++--- src/lib/mlx-cpp/patches-rocm/mlx/stream.cpp | 4 + .../patches/mlx/backend/cpu/quantized.cpp | 38 ++- .../mlx/backend/cuda/binary/binary.cuh | 11 +- .../patches/mlx/backend/cuda/cuda_utils.h | 108 --------- .../mlx/backend/cuda/device/qmm_sm80.cuh | 117 ++++++++- .../backend/cuda/quantized/qmm/qmm_sm80.cu | 127 +++++++++- .../mlx/backend/cuda/quantized/quantized.cpp | 34 ++- .../mlx/backend/cuda/reduce/col_reduce.cu | 5 +- .../cuda/scaled_dot_product_attention.cpp | 7 +- .../patches/mlx/backend/metal/device.cpp | 33 ++- .../patches/mlx/backend/metal/kernels/utils.h | 49 +++- .../patches/mlx/backend/metal/quantized.cpp | 106 ++++---- .../mlxcel-surgery/src/ops/split_mtp_tests.rs | 26 +- 46 files changed, 1743 insertions(+), 370 deletions(-) create mode 100644 docs/upstream/mlx-pin-243bdf09-reconciliation.md create mode 100644 scripts/mlxcelverse/allocator_ownership_test.cpp create mode 100644 scripts/mlxcelverse/allocator_ownership_test.sh create mode 100644 src/lib/mlx-cpp/patches-rocm/mlx/allocator.h delete mode 100644 src/lib/mlx-cpp/patches/mlx/backend/cuda/cuda_utils.h diff --git a/Makefile b/Makefile index c4b6e49d3..2e4320a56 100644 --- a/Makefile +++ b/Makefile @@ -874,6 +874,7 @@ verify-rocm-overlay: ## Assert the ROCm overlay records agree with the MLX pin a @echo "$(CYAN)[verify] ROCm overlay records...$(RESET)" @python3 scripts/mlxcelverse/rocm_overlay.py records @bash scripts/mlxcelverse/rocm_overlay_test.sh + @bash scripts/mlxcelverse/allocator_ownership_test.sh .PHONY: verify-binary-assets verify-binary-assets: ## Assert every tracked binary file is declared with a size budget diff --git a/docs/upstream/mlx-cuda-handle-destructor-throws-at-teardown.md b/docs/upstream/mlx-cuda-handle-destructor-throws-at-teardown.md index 8b4c9ff0e..985482aa3 100644 --- a/docs/upstream/mlx-cuda-handle-destructor-throws-at-teardown.md +++ b/docs/upstream/mlx-cuda-handle-destructor-throws-at-teardown.md @@ -1,6 +1,8 @@ # MLX CUDA: `CudaHandle`'s destructor throws while the driver is unloading -A test binary on GB10 sometimes aborts after it has already printed `test result: ok`: +As of the MLX pin update in #2303, the local overlay is retired: upstream [#4514](https://github.com/ml-explore/mlx/pull/4514), commit `9eb3e3ab6fceb4031994f70dfe023872643ed634`, directly destroys CUDA handles without throwing and adds related lifetime fixes. The overlapping [#4531](https://github.com/ml-explore/mlx/pull/4531) was closed unmerged in its favor. The captured target includes #4514; the historical investigation below describes the earlier pin. + +A test binary on GB10 sometimes aborted after it had already printed `test result: ok`: ``` test result: ok. N passed; 0 failed; ... diff --git a/docs/upstream/mlx-pin-243bdf09-reconciliation.md b/docs/upstream/mlx-pin-243bdf09-reconciliation.md new file mode 100644 index 000000000..f42e098bf --- /dev/null +++ b/docs/upstream/mlx-pin-243bdf09-reconciliation.md @@ -0,0 +1,66 @@ +# MLX pin 243bdf09 overlay reconciliation + +Issue: [#2303](https://github.com/lablup/mlxcel/issues/2303). + +The dependency advances from `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` to `243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e`: 111 upstream commits and 312 changed paths. Every comparison below uses pristine `git show COMMIT:path` objects, never the build checkout after overlays were applied. + +## CPU, CUDA and Metal inventory + +The original 32 files comprise 30 upstream-backed overlays and two local extracted helpers. Thirteen upstream-backed files changed relative to their declared bases; twelve were rebased and one was retired. The other seventeen upstream-backed files have zero path-specific commits between their declared base and the target, so their original provenance remains valid. The helpers have no upstream path to replace. There are 31 files after retirement. + +Paths below are relative to `patches/mlx/backend/`, except `dtype.cpp`, `fast.cpp` and `ops.cpp` under `patches-cuda/`. + +| Path | Previous base | Disposition | Retained behavior / evidence | +| --- | --- | --- | --- | +| `cpu/quantized.cpp` | `81ba1c6a` | Rebased | Float accumulation and buffered scalar affine/MXFP dot products (#2248). | +| `cuda/binary/binary.cuh` | `e9463bb` | Rebased | Direct mixed bf16/f32 binary kernels; upstream FloorDivide support retained. | +| `cuda/cuda_utils.h` | `81ba1c6a` | Retired | Retired: upstream #4514 destroys handles without throwing; no other local delta. | +| `cuda/device/binary_ops.cuh` | `81ba1c6a` | Unchanged; zero drift | Mixed-type operators in runtime-compiled graphs. | +| `cuda/device/gemm_sm70.cuh` | `9a795735` | Unchanged; zero drift | Volta tensor-core MMA path. | +| `cuda/device/qmm_naive.cuh` | `59d600b5e` | Unchanged; zero drift | Volta MMA dispatch parameterization. | +| `cuda/device/qmm_sm80.cuh` | `b7c3dd6d` | Rebased | Drain cp.async and synchronize before the aliased shared-memory epilogue (#910). | +| `cuda/gemms/grouped_gemm.h` | `b98831ad` | Unchanged; zero drift | Local grouped GEMM dispatch declarations. | +| `cuda/gemms/grouped_gemm_arch.h` | `9a795735` | Local helper | Local extracted architecture selector with host-test coverage (#1544). | +| `cuda/gemms/grouped_gemm_unaligned.cu` | `c9aa5605` | Unchanged; zero drift | Grouped GEMM architecture selection and local kernel dispatch. | +| `cuda/jit_module.cpp` | `9a795735` | Unchanged; zero drift | Bundled CUDA header discovery, overrides and cold-cache notice. | +| `cuda/jit_module.h` | `9a795735` | Unchanged; zero drift | JIT module interfaces needed by local grouped GEMM dispatch. | +| `cuda/matmul.cpp` | `9a795735` | Unchanged; zero drift | Custom grouped GEMM integration and dispatch. | +| `cuda/quantized/qmm/qmm_naive.cu` | `81ba1c6a` | Unchanged; zero drift | Naive quantized GEMM architecture/tile selection. | +| `cuda/quantized/qmm/qmm_naive_tile.h` | `9a795735` | Local helper | Local extracted tile selector with host-test coverage (#1541). | +| `cuda/quantized/qmm/qmm_sm80.cu` | `b7c3dd6d` | Rebased | Blackwell tile overrides and epilogue-fix JIT revision; upstream gathered GEMM implementation retained. | +| `cuda/quantized/qmm/qmv.cu` | `59d600b5e` | Unchanged; zero drift | Local quantized vector kernels and launch controls. | +| `cuda/quantized/quantized.cpp` | `59d600b5e` | Rebased | Large batch chunking and diagnostics; new FP gather-vector path also uses the chunker. | +| `cuda/reduce/all_reduce.cu` | `59d600b5e` | Unchanged; zero drift | Float accumulation with bf16 output and intermediate selection. | +| `cuda/reduce/col_reduce.cu` | `59d600b5e` | Rebased | Float accumulation/intermediates with bf16 output; consume Data without obsolete free callback. | +| `cuda/reduce/init_reduce.cu` | `9a795735` | Unchanged; zero drift | Initialize reduction output in its actual dtype. | +| `cuda/reduce/reduce_ops.cuh` | `ce45c525` | Unchanged; zero drift | Float accumulation traits for bf16 sum/product. | +| `cuda/reduce/row_reduce.cu` | `59d600b5e` | Unchanged; zero drift | Float accumulation with bf16 output. | +| `cuda/scaled_dot_product_attention.cpp` | `81ba1c6a` | Rebased | cuDNN shape bucketing, bounded graph fallback and determinism controls. | +| `cuda/scaled_dot_product_attention.cu` | `e9463bb` | Unchanged; zero drift | Local fused SDPA dtype behavior. | +| `metal/compiled.cpp` | `81ba1c6a` | Unchanged; zero drift | JIT type promotion behavior. | +| `metal/device.cpp` | `81ba1c6a` | Rebased | Runtime command-buffer input-budget override. | +| `metal/kernels/utils.h` | `81ba1c6a` | Rebased | SIMD-group matrix include and qualified metal::vec types. | +| `metal/quantized.cpp` | `59d600b5e` | Rebased | QMV-wide override and B*M broadcast-row correction in GatherQMM and GatherQQMM (#1599). | +| `dtype.cpp` | `ce45c525` | Unchanged; zero drift | CUDA-only bf16/f32 promotion policy. | +| `fast.cpp` | `59d600b5e` | Rebased | CUDA-only bf16 normalization fallback; upstream explicit stream argument retained. | +| `ops.cpp` | `59d600b5e` | Rebased | CUDA-only reduction/promotion behavior; new upstream operation definitions retained. | + +Three-way conflict review retained the local normalization compute dtype together with upstream's explicit stream argument. CUDA `GatherQMM` preserves the new `fp_gather_qmv` selection while feeding its index/output chunks through the existing batch limiter. The new expert-grouped `gather_qmm_rhs_sm80` route remains upstream's implementation. + +Upstream [#4514](https://github.com/ml-explore/mlx/pull/4514), `9eb3e3ab6fceb4031994f70dfe023872643ed634`, supplies nonthrowing CUDA handle destruction and associated stream/lifetime fixes. The removed overlay had no unrelated local delta; its GLOB registration disappears with the file. The overlapping #4531 was closed unmerged in favor of #4514. + +Upstream [#4594](https://github.com/ml-explore/mlx/pull/4594), `91c83d19b`, provides the Xcode 27 cooperative-tensor workaround. No older `gated_delta_update_nax.h` whole-file overlay is imported from PR #2138. Target `quantized.cpp` still passes the pre-broadcast `x.size() / K` to both gathered matrix routes, so the #1599 correction from #2138 (`7d5459811f16c36d1b05c2a45455f8cb05b99883`) remains necessary and is applied to both as `B * M`. Existing shared-activation MXFP regressions exercise `GatherQMM`; the new NVFP4 matrix route needs appropriate Apple hardware coverage. + +The same PR's two filesystem-portability test corrections are retained for candidate Metal CI: canonical temporary parents on macOS and explicit clean rejection on filesystems that reject non-UTF-8 names. They do not change production surgery behavior. + +## Ownership and integration + +The upstream allocator returns move-only `allocator::Data`, and array donation requires allocator-owned storage. CUDA column-reduction intermediates now consume Data directly; upstream's removed free callback is not reintroduced. The remaining CPU/CUDA/Metal allocations already pass allocator results directly into arrays, preserving the owner and encoder temporary lifetime. The custom C++ bridge's `from_bytes_nocopy` uses the unchanged public raw-pointer/deleter constructor, which upstream now marks non-owned; `kv_inplace_write` uses `copy_shared_buffer`, preserving shared lifetime. `fused_rope_append` delegates allocation to upstream CustomKernel and has no raw allocation handoff to migrate. + +ROCm requires an explicit owned-data friend because the public Data constructor with a deleter intentionally represents borrowed storage. Its retarget and backend compatibility changes are recorded in `patches-rocm/LOCAL_FIXES.md` and `CORE_RESIDUAL.diff`. The fork remains `75915908dfe5028335d318b10340313744fd3a8d`. The host allocator regression uses the actual overlay header and wrapper; it tests moved ownership, borrowed custom release, empty allocations, pending shared ownership and exception cleanup. It does not execute HIP or prove GPU completion ordering. + +The upstream version moves from 0.32.3 to 0.32.4, giving default CUDA PTX caches a new namespace. Validation uses a fresh cache when an explicit `MLX_PTX_CACHE_DIR` override is supplied. + +## Validation scope + +The PR records the exact tested source commit, fetched MLX HEAD, binary hashes and hardware outcomes. CUDA compilation/inference and candidate Apple Silicon CI are coordinated after source reconciliation. Kimi controls require the independent #2296 model fix in an explicitly recorded combined validation tree; that fix is not part of this dependency PR. ROCm compilation, GPU execution and M5-specific execution are not inferred from host tests or another GPU. This checkout's host has no HIP toolchain. diff --git a/scripts/mlxcelverse/allocator_ownership_test.cpp b/scripts/mlxcelverse/allocator_ownership_test.cpp new file mode 100644 index 000000000..e5eb8a8d9 --- /dev/null +++ b/scripts/mlxcelverse/allocator_ownership_test.cpp @@ -0,0 +1,162 @@ +// Copyright 2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Host-only contract test for the ROCm overlay's actual allocator::Data and +// malloc_async handoff. This does not exercise HIP execution or synchronization. +#include "mlx/allocator.h" + +#include +#include +#include +#include +#include +#include + +namespace mlx::core { +class MockAllocator : public allocator::Allocator { + public: + allocator::Buffer malloc(size_t size) override { + auto* ptr = new char[size ? size : 1]; + live.insert(ptr); + ++allocations; + return allocator::Buffer(ptr); + } + allocator::Buffer malloc_async(size_t size, int device, void* stream) { + assert(device == 7 && stream == expected_stream); + return malloc(size); + } + void free(allocator::Buffer buffer) override { + assert(live.erase(buffer.ptr()) == 1); + delete[] static_cast(buffer.ptr()); + ++releases; + } + size_t size(allocator::Buffer) const override { + return 0; + } + std::set live; + int allocations{0}; + int releases{0}; + void* expected_stream{this}; +}; + +MockAllocator mock_allocator; +namespace allocator { +Allocator& allocator() { + return mock_allocator; +} +} // namespace allocator + +namespace rocm { +using hipStream_t = void*; +class CommandEncoder { + public: + struct Device { + int hip_device() const { + return 7; + } + }; + Device device() const { + return {}; + } + void* stream() const { + return mock_allocator.expected_stream; + } +}; +MockAllocator& allocator() { + return mock_allocator; +} +// Extracted verbatim from the overlay, not a duplicate implementation. +#include "rocm_malloc_async.inc" +} // namespace rocm +} // namespace mlx::core + +int main() { + using namespace mlx::core; + using allocator::Data; + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(std::is_nothrow_move_constructible_v); + rocm::CommandEncoder encoder; + + for (size_t size : {size_t{0}, size_t{32}}) { + const int before = mock_allocator.releases; + { + auto original = rocm::malloc_async(size, encoder); + assert(original.is_owned()); + auto buffer = original.buffer(); + Data moved(std::move(original)); + assert(!original.is_owned()); + assert(original.buffer().ptr() == nullptr); + assert(moved.is_owned()); + assert(moved.buffer().ptr() == buffer.ptr()); + assert(mock_allocator.releases == before); + } + assert(mock_allocator.releases == before + 1); + } + + // A borrowed slab shell must retain its custom release and never become + // owned: upstream array::is_donatable requires Data::is_owned(). + int custom_releases = 0; + int slab = 42; + const int before_borrow = mock_allocator.releases; + { + Data borrowed(allocator::Buffer(&slab), [&](allocator::Buffer buffer) { + assert(buffer.ptr() == &slab); + ++custom_releases; + }); + Data moved(std::move(borrowed)); + assert(!moved.is_owned()); + assert(custom_releases == 0); + } + assert(custom_releases == 1); + assert(mock_allocator.releases == before_borrow); + assert(slab == 42); + + // Model the encoder's shared ownership through a pending command batch. + // Destroying the producer must not recycle storage before batch completion. + std::vector> pending; + const int before_batch = mock_allocator.releases; + { + auto producer = std::make_shared(rocm::malloc_async(64, encoder)); + pending.push_back(producer); + } + assert(mock_allocator.releases == before_batch); + pending.clear(); + assert(mock_allocator.releases == before_batch + 1); + + // Failure before handoff releases owned storage, including an empty buffer. + for (size_t size : {size_t{0}, size_t{64}}) { + const int before = mock_allocator.releases; + try { + auto data = rocm::malloc_async(size, encoder); + throw std::runtime_error("handoff failed"); + } catch (const std::runtime_error&) { + } + assert(mock_allocator.releases == before + 1); + } + + // Failure after enqueue leaves the pending batch's reference alive. + const int before_failure = mock_allocator.releases; + try { + auto data = std::make_shared(rocm::malloc_async(64, encoder)); + pending.push_back(data); + throw std::runtime_error("later enqueue failed"); + } catch (const std::runtime_error&) { + } + assert(mock_allocator.releases == before_failure); + pending.clear(); + assert(mock_allocator.releases == before_failure + 1); + assert(mock_allocator.live.empty()); + assert(mock_allocator.allocations == mock_allocator.releases); +} diff --git a/scripts/mlxcelverse/allocator_ownership_test.sh b/scripts/mlxcelverse/allocator_ownership_test.sh new file mode 100644 index 000000000..a12bfdc0d --- /dev/null +++ b/scripts/mlxcelverse/allocator_ownership_test.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Copyright 2026 Lablup Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Exercise the move-only ROCm allocation boundary with a host mock, without HIP. +set -euo pipefail +repo_root=$(cd "$(dirname "$0")/../.." && pwd) +overlay="$repo_root/src/lib/mlx-cpp/patches-rocm" +test_dir=$(mktemp -d "${TMPDIR:-/tmp}/rocm-ownership.XXXXXX") +trap 'rm -rf "$test_dir"' EXIT +mkdir -p "$test_dir/mlx" +# api.h only supplies symbol visibility; the actual allocator header is tested. +printf '#pragma once\n#define MLX_API\n' > "$test_dir/mlx/api.h" +awk ' + /^allocator::Data malloc_async\(size_t size, CommandEncoder& encoder\) \{/ { copying = 1 } + copying { print } + copying && /^}/ { found = 1; exit } + END { if (!found) exit 1 } +' "$overlay/mlx/backend/rocm/allocator.cpp" > "$test_dir/rocm_malloc_async.inc" +"${CXX:-c++}" -std=c++17 -Wall -Wextra -Werror -Wno-unused-parameter \ + -I"$overlay" -I"$test_dir" \ + "$repo_root/scripts/mlxcelverse/allocator_ownership_test.cpp" \ + -o "$test_dir/allocator-ownership-test" +"$test_dir/allocator-ownership-test" +printf '[rocm-ownership] owned, borrowed, empty, move, retained-batch and failure-cleanup contracts passed (host mock; no HIP execution)\n' diff --git a/src/lib/mlx-cpp/CMakeLists.txt b/src/lib/mlx-cpp/CMakeLists.txt index c55f8f14a..7291b6f83 100644 --- a/src/lib/mlx-cpp/CMakeLists.txt +++ b/src/lib/mlx-cpp/CMakeLists.txt @@ -29,10 +29,13 @@ function(mlx_apply_source_overlays mlx_source_dir) # The steel/gemm/mma.h overlay (issue #217: upstream safe-load fix # cherry-pick, PRs #3560/#3565) was retired when the pin moved to # 2026-06-11 upstream main, which carries the fix natively (issue #222). - # quantized.cpp carries one delta, the MLXCEL_QMV_WIDE off-switch on - # `use_qmv_wide` (issue #1187). Upstream has no equivalent knob and the - # predicate is unchanged as of pin 81ba1c6a, so it has to be overlaid here. - # Refresh the file wholesale on a bump and re-apply that one hunk. + # quantized.cpp retains the MLXCEL_QMV_WIDE off-switch (#1187) and the + # broadcast gathered-row count (#1599) in GatherQMM/GatherQQMM. Refresh + # wholesale on a bump and re-apply only those documented local changes. + # Upstream #4594 supplies the Xcode 27 cooperative-tensor workaround; + # no gated_delta_update_nax.h overlay is needed at this pin. + # Upstream #4514 also supplies nonthrowing CUDA handle teardown, so the + # cuda_utils.h overlay is retired. # device.cpp carries a runtime override of the per-command-buffer input # budget (`mlxcel_set_mb_per_buffer_override`), which mlxcel raises around # decode steps and leaves at the device default for prefill. Upstream latches @@ -200,7 +203,7 @@ else() # `)` as ending the declaration, so one there would hide the tag # instead of being read past. # Both parsers fail loudly rather than guess when the shape does not hold. - GIT_TAG 59d600b5e64c238427d0f8d897ab7c682ef4d3d2) + GIT_TAG 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e) # Use FetchContent_Populate + add_subdirectory so we can apply source # overlays before the MLX build system processes the files. diff --git a/src/lib/mlx-cpp/patches-cuda/fast.cpp b/src/lib/mlx-cpp/patches-cuda/fast.cpp index a3dd80aa7..3863f46fd 100644 --- a/src/lib/mlx-cpp/patches-cuda/fast.cpp +++ b/src/lib/mlx-cpp/patches-cuda/fast.cpp @@ -1,8 +1,8 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. // // CUDA patch: bf16 scalar/constant generation in normalization ops // -// Modified from upstream MLX 59d600b5e mlx/fast.cpp (three-way merged at each +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/fast.cpp (three-way merged at each // pin bump; the two fallbacks below are the whole delta) // // Changes: @@ -380,7 +380,7 @@ array layer_norm( auto xc = subtract(x, mu, s); auto v = mean(square(xc, s), /* axis= */ -1, /* keepdims= */ true, s); - x = multiply(xc, rsqrt(add(v, array(eps, compute_type), s), s)); + x = multiply(xc, rsqrt(add(v, array(eps, compute_type), s), s), s); x = astype(x, out_type, s); // If the LN is affine then transform x according to the weight and bias @@ -962,6 +962,7 @@ array scaled_dot_product_attention( has_mask, has_arr_mask, do_causal, + has_sinks, is_training, output_logsumexp, force_fused, @@ -972,8 +973,9 @@ array scaled_dot_product_attention( array& mask = inputs[3]; mask = where( mask, - full_like(mask, 0, final_type, s), - full_like(mask, -inf, final_type, s)); + full_like(mask, 0, final_type, stream), + full_like(mask, -inf, final_type, stream), + stream); } Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)}; auto primitive = std::make_shared( @@ -1005,13 +1007,16 @@ std::vector ScaledDotProductAttention::vjp( const std::vector& outputs) { assert(primals.size() >= 3); assert(cotangents.size() == outputs.size()); - auto s = stream(); if (ScaledDotProductAttentionVJP::use_fallback(primals[0], s)) { assert(outputs.size() == 1); return Custom::vjp(primals, cotangents, argnums, outputs); } + if (outputs.size() != 2) { + return Custom::vjp(primals, cotangents, argnums, outputs); + } + auto fallback = [sdpa = fallback_, s](const std::vector& inputs) { std::vector primals(inputs.begin(), std::prev(inputs.end())); auto [_, vjps] = mlx::core::vjp(sdpa, primals, {inputs.back()}); @@ -1186,6 +1191,54 @@ std::vector gated_delta_update( auto result = fallback({q, k, v, g, beta, h0, mask}); return result; } +std::vector GatedDeltaUpdate::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) { + const int Hk = primals[0].shape(2); + const int Dk = primals[0].shape(3); + const int Hv = primals[2].shape(2); + const int Dv = primals[2].shape(3); + + if (GatedDeltaUpdateVJP::use_fallback(Hk, Dk, Hv, Dv, stream())) { + return Custom::vjp(primals, cotangents, argnums, outputs); + } + + if (primals.size() != 6) { + throw std::runtime_error( + "[GatedDeltaUpdate::vjp] expected 6 primals (q,k,v,g,beta,h0), got " + + std::to_string(primals.size())); + } + if (cotangents.size() != 2) { + throw std::runtime_error( + "[GatedDeltaUpdate::vjp] expected 2 cotangents (out, state), got " + + std::to_string(cotangents.size())); + } + + std::vector inputs = primals; // q,k,v,g,beta,h0 (indices 0..5) + inputs.push_back(cotangents[0]); // cot_o -> index 6 + inputs.push_back(cotangents[1]); // cot_h -> index 7 + + std::vector shapes; + std::vector dtypes; + for (int i = 0; i < 6; ++i) { + shapes.push_back(primals[i].shape()); + dtypes.push_back(i == 5 ? float32 : primals[i].dtype()); // dh is float32 + } + + auto vjps = array::make_arrays( + std::move(shapes), + std::move(dtypes), + std::make_shared(stream(), fallback_), + std::move(inputs)); + + std::vector returned_vjps; + for (int arg : argnums) { + returned_vjps.push_back(std::move(vjps[arg])); + } + return returned_vjps; +} bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); diff --git a/src/lib/mlx-cpp/patches-cuda/ops.cpp b/src/lib/mlx-cpp/patches-cuda/ops.cpp index e74593c9a..4968f43ee 100644 --- a/src/lib/mlx-cpp/patches-cuda/ops.cpp +++ b/src/lib/mlx-cpp/patches-cuda/ops.cpp @@ -2,7 +2,7 @@ // // CUDA patch: Mixed-precision binary operation dispatch for bf16 // -// Modified from upstream MLX 59d600b5e mlx/ops.cpp +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/ops.cpp // // Changes: // - Added bf16_mixed_astype() helper that skips astype conversion when the @@ -511,7 +511,7 @@ array unflatten( std::ostringstream msg; msg << "[unflatten] Invalid axes " << ax << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } size_t size = 1; @@ -755,7 +755,7 @@ array flip(const array& a, StreamOrDevice s /* = {} */) { namespace { inline auto -normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) { +normalize_slice(const Shape& shape, Shape& start, Shape& stop, Shape& strides) { // - Start indices are normalized // - End indices are unchanged as -1 means something different // pre-normalization (the end of the axis) versus post normalization (the @@ -803,7 +803,13 @@ normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) { } // Simplify the stride if it's unused if (out_shape[i] == 1) { - strides[i] = 1; + if (strides[i] < 0 && start[i] > 0) { + strides[i] = -1; + stop[i] = start[i] - 1; + } else { + strides[i] = 1; + stop[i] = start[i] + 1; + } } } @@ -846,7 +852,7 @@ void normalize_dynamic_slice_inputs( std::ostringstream msg; msg << prefix << " Invalid axis " << ax << " for array with dimension " << a.ndim() << "."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } ax = new_ax; } @@ -1799,7 +1805,7 @@ array transpose( std::ostringstream msg; msg << "[transpose] Invalid axis (" << ax << ") for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } if (shape[ax] != 0) { throw std::invalid_argument("[transpose] Repeat axes not allowed."); @@ -1823,6 +1829,17 @@ array transpose(const array& a, StreamOrDevice s /* = {} */) { return transpose(a, std::move(axes), to_stream(s)); } +array matrix_transpose(const array& a, StreamOrDevice s /* = {} */) { + if (a.ndim() < 2) { + throw std::invalid_argument( + "[matrix_transpose] Input array must have at least 2 dimensions."); + } + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + std::swap(axes[a.ndim() - 1], axes[a.ndim() - 2]); + return transpose(a, std::move(axes), to_stream(s)); +} + array broadcast_to( const array& a, const Shape& shape, @@ -2362,7 +2379,7 @@ array mean( std::ostringstream msg; msg << "[mean] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } } auto dtype = at_least_float(a.dtype()); @@ -2396,7 +2413,7 @@ array median( std::ostringstream msg; msg << "[median] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } set_axes.insert(axis < 0 ? axis + ndim : axis); } @@ -2973,6 +2990,148 @@ array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) { return slice(a_partitioned, slice_starts, slice_ends, s); } +std::vector unique( + const array& a, + int size, + bool return_index /* = false */, + bool return_inverse /* = false */, + bool return_counts /* = false */, + const std::optional& fill_value /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + // Validate args + if (size < 0) { + std::ostringstream msg; + msg << "[unique] Received negative size " << size << "."; + throw std::invalid_argument(msg.str()); + } + if (fill_value && fill_value->size() != 1) { + std::ostringstream msg; + msg << "[unique] Fill value must have one element, but got shape " + << fill_value->shape() << "."; + throw std::invalid_argument(msg.str()); + } + const auto flat = flatten(a, s); + const int n = flat.size(); + + // Handle the edge case of an empty array. + if (n == 0) { + // Without an element there is no default fill value. + if (size > 0 && !fill_value) { + throw std::invalid_argument( + "[unique] A fill value is required for an empty input with a" + " non-zero size."); + } + std::vector out; + out.push_back( + size == 0 ? flat + : full( + {size}, + astype(reshape(*fill_value, {}, s), flat.dtype(), s), + flat.dtype(), + s)); + if (return_index) { + out.push_back(zeros({size}, uint32, s)); + } + if (return_inverse) { + out.push_back(zeros(a.shape(), uint32, s)); + } + if (return_counts) { + out.push_back(zeros({size}, int32, s)); + } + return out; + } + + // Sort the array. The values stay differentiable, the permutation does not. + std::optional order; + if (return_index || return_inverse) { + order = stop_gradient(argsort(flat, 0, s), s); + } + const auto sorted = order ? take(flat, *order, 0, s) : sort(flat, 0, s); + + // Edge detection on the sorted array, true where a new unique + // value starts in the sorted array. + const auto boundary = concatenate( + {array({true}), + not_equal( + slice(sorted, {1}, {n}, s), slice(sorted, {0}, {n - 1}, s), s)}, + 0, + s); + + // Cumsum on boundary gives each sorted element the index of the unique + // value it belongs to. It indexes the output, so it is not differentiable. + const auto group = stop_gradient( + subtract( + cumsum(astype(boundary, uint32, s), 0, false, true, s), + array(1, uint32), + s), + s); + + // Use smallest element of the sorted array (index 0) as padding by default. + const auto fill = fill_value + ? astype(reshape(*fill_value, {}, s), flat.dtype(), s) + : slice(sorted, {0}, {1}, s); + + const int buffer_size = std::max(n, size); + // Scatter positions rather than values in the sorted array (64 bit values + // would fail to scatter on the GPU). + const auto slots = stop_gradient( + slice( + scatter_max( + zeros({buffer_size}, uint32, s), + group, + expand_dims( + where( + boundary, + add(arange(n, uint32, s), array(1, uint32), s), + array(0, uint32), + s), + 1, + s), + 0, + s), + {0}, + {size}, + s), + s); + const auto used = greater(slots, array(0, uint32), s); + const auto positions = + subtract(maximum(slots, array(1, uint32), s), array(1, uint32), s); + + // Build the output arrays. + std::vector out; + out.push_back(where(used, take(sorted, positions, 0, s), fill, s)); + if (return_index) { + // The sort is stable, so the first sorted position of a group holds the + // first occurrence in the input. Unused slots point at position zero, + // which pads with the index of the smallest unique value. + out.push_back(take(*order, positions, 0, s)); + } + if (return_inverse) { + // Clamp so that the indices stay inside a truncated output. Without a + // truncation every group index is already smaller than size. + const auto clamped = + minimum(group, array(std::max(size - 1, 0), uint32), s); + const auto inverse = scatter( + zeros({n}, uint32, s), *order, expand_dims(clamped, 1, s), 0, s); + out.push_back(reshape(inverse, a.shape(), s)); + } + // Padding entries count zero, so counts sum to the input size unless the + // output was truncated. + if (return_counts) { + out.push_back(slice( + scatter_add( + zeros({buffer_size}, int32, s), + group, + ones({n, 1}, int32, s), + 0, + s), + {0}, + {size}, + s)); + } + return out; +} + array logsumexp(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) { std::vector axes(a.ndim()); std::iota(axes.begin(), axes.end(), 0); @@ -3181,14 +3340,13 @@ array floor_divide( const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); - if (issubdtype(dtype, inexact)) { - return floor(divide(a, b, s), s); - } - auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); auto shape = inputs[0].shape(); return array( - shape, dtype, std::make_shared(to_stream(s)), std::move(inputs)); + shape, + dtype, + std::make_shared(to_stream(s)), + std::move(inputs)); } array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) { @@ -3712,7 +3870,7 @@ array take( std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -3757,7 +3915,7 @@ array take(const array& a, int index, int axis, StreamOrDevice s /* = {} */) { std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -4284,7 +4442,7 @@ array diff( int ndim = static_cast(a.ndim()); int ax = axis < 0 ? axis + ndim : axis; if (ax < 0 || ax >= ndim) { - throw std::invalid_argument("[diff] Axis is out of bounds for the array."); + throw std::out_of_range("[diff] Axis is out of bounds for the array."); } if (n < 0) { throw std::invalid_argument("[diff] Order `n` must be non-negative."); @@ -4310,12 +4468,13 @@ array logcumsumexp( bool inclusive /* = true*/, StreamOrDevice s /* = {}*/) { axis = normalize_axis_index(axis, a.ndim(), "[logcumsumexp] "); + auto out_type = at_least_float(a.dtype()); return array( a.shape(), - a.dtype(), + out_type, std::make_shared( to_stream(s), Scan::ReduceType::LogAddExp, axis, reverse, inclusive), - {a}); + {astype(a, out_type, s)}); } array logcumsumexp( @@ -5046,10 +5205,10 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { throw std::invalid_argument(msg.str()); } - if (bits < 2 || bits > 8 || bits == 7) { + if (bits < 1 || bits > 8 || bits == 7) { std::ostringstream msg; msg << "[quantize] The requested number of bits " << bits - << " is not supported. The supported bits are 2, 3, 4, 5, 6 and 8."; + << " is not supported. The supported bits are 1, 2, 3, 4, 5, 6 and 8."; throw std::invalid_argument(msg.str()); } @@ -5070,14 +5229,22 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { w_max = astype(w_max, float32, s); w_min = astype(w_min, float32, s); - array mask = greater(abs(w_min, s), abs(w_max, s), s); - array scales = - maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); - scales = where(mask, scales, negative(scales, s), s); - array edge = where(mask, w_min, w_max, s); - array q0 = round(divide(edge, scales, s), s); - scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); - array biases = where(equal(q0, zero, s), zero, edge, s); + array scales(0, float32); + array biases(0, float32); + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scales = maximum(subtract(w_max, w_min, s), eps, s); + biases = w_min; + } else { + array mask = greater(abs(w_min, s), abs(w_max, s), s); + scales = maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); + scales = where(mask, scales, negative(scales, s), s); + array edge = where(mask, w_min, w_max, s); + array q0 = round(divide(edge, scales, s), s); + scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); + biases = where(equal(q0, zero, s), zero, edge, s); + } packed_w = pack_and_quantize(packed_w, scales, biases, bits, s); @@ -5837,7 +6004,7 @@ array vecdot( } int ax = axis < 0 ? axis + a.ndim() : axis; if (ax < 0 || ax >= a.ndim()) { - throw std::invalid_argument("[vecdot] axis is out of bounds."); + throw std::out_of_range("[vecdot] axis is out of bounds."); } if (axis < 0 ? axis + b.ndim() != ax : axis >= b.ndim()) { throw std::invalid_argument("[vecdot] axis is out of bounds."); @@ -6718,7 +6885,7 @@ array roll( std::ostringstream msg; msg << "[roll] Invalid axis " << axes[i] << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } auto sh = shift[i]; diff --git a/src/lib/mlx-cpp/patches-rocm/CMakeLists.txt b/src/lib/mlx-cpp/patches-rocm/CMakeLists.txt index eafa5a166..249905647 100644 --- a/src/lib/mlx-cpp/patches-rocm/CMakeLists.txt +++ b/src/lib/mlx-cpp/patches-rocm/CMakeLists.txt @@ -436,7 +436,7 @@ if(MLX_BUILD_PYTHON_BINDINGS) FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v3.0.1 + GIT_TAG v3.1.0 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) diff --git a/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff b/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff index f8478c216..416df2415 100644 --- a/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff +++ b/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff @@ -15,29 +15,53 @@ # replace: an overlay that silently drops part of an upstream change shows up # here as a new hunk, even if it was wrong before the bump. # -# pin: 59d600b5e64c238427d0f8d897ab7c682ef4d3d2 +# pin: 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e # fork: 75915908dfe5028335d318b10340313744fd3a8d # fork_upstream_merge_base: 39886de461d4f014cd60a3a68ad42fd8d776754a -# overlay-sha256: ad95c895c1bf4d10b3fdece59d617b35c8f6135a135edb5a08f91a8a0a9afb30 CMakeLists.txt +# overlay-sha256: 965377d80af426d42b3e106da236cf63d0d1fc6c19130a9b9b16e1fa51734c74 CMakeLists.txt # overlay-sha256: 9da8043e8823e13c49faa62af44eb944db516da83d4f204fe3cfcb0424bd2c1f mlx/CMakeLists.txt +# overlay-sha256: f18df01b2939c50a5dd1fe9c75d269bef493402c920123939e70b0d968dff211 mlx/allocator.h # overlay-sha256: c255fb0f18fb597cd7ba1708dcc5799ea26ee8541820dedb8997a279d9257496 mlx/backend/common/buffer_cache.h -# overlay-sha256: 0bc33a435855548b05a5d87e300c5b023b3ef618c728e326a5dbef98232e7fd8 mlx/backend/common/compiled.cpp -# overlay-sha256: d53bb3bc38645526fb877ed079c7a2e112bf8f74395a1887f048548bad2a9302 mlx/backend/common/compiled.h -# overlay-sha256: ccf2ae61f566cde4b740b43d62a37fd6c904838b68f8330666073842f290b0f6 mlx/backend/gpu/primitives.cpp -# overlay-sha256: 744670de21dbd5c74e2cee3afe1d542c2b73c45deee7611765fc9b0dcc99c334 mlx/compile.cpp +# overlay-sha256: 6755397b3410b82fe451c1aa0455ba722d4106a67d72fb9ce6a4eecda61ca920 mlx/backend/common/compiled.cpp +# overlay-sha256: 28759d52cd6ec6f55a6ba85ae6e594e24c042fc55bfe139ac1856e556ee42e10 mlx/backend/common/compiled.h +# overlay-sha256: 293cc82e52a6b3d98fb06c66f8c7e471790d908ab29f48a8a33d6e2d77341c81 mlx/backend/gpu/primitives.cpp +# overlay-sha256: 849838a3e3354496e7c4e571b98491e7c80f94f4b84fe9e642093f332069fc8a mlx/compile.cpp # overlay-sha256: 67e68d6e70df6271ed2f54d75957c4762653b6a32fa6c666ea3b014c5d425acd mlx/device.cpp -# overlay-sha256: b45c24910329b5b40085b6c0607bc2c8e514a4fc3bafcce4e7422701a7d0d78c mlx/fast.cpp +# overlay-sha256: 26584a93a3720250aa59c6dd9f091863f8c047691ce55f9a988a01da390ab02a mlx/fast.cpp # overlay-sha256: b1a5366b31412246256ea56f9114bfe7df2fcb93329301ce7ec924f7200c7695 mlx/fast.h -# overlay-sha256: 46684f0d184dc8e1d19c04839ad315e2a3ad38f58e8bd899d6a0b23031391953 mlx/fast_primitives.h +# overlay-sha256: 8839267c3bc469b6a3f7d67c6028701a1a3bd02d0f09b6ac31b15407204d2653 mlx/fast_primitives.h # overlay-sha256: f5fccd165887c2860c04fecf3fcf6073397789081183c2bcaaacad8f8f4d8e9e mlx/io/safetensors.cpp -# overlay-sha256: 92c85296ae688ef1db962a3c1b2132f0b94493d71ee2de350c259c4ce2493581 mlx/ops.cpp -# overlay-sha256: 221cf283f6e15a1ddc4d11451eb79e930f2567828cf5fff45cb9ff020c6431e5 mlx/primitives.cpp -# overlay-sha256: adb49482ec61de1b26887b851ec499b351387d0d030fb240664ded3d3d04e5cd mlx/stream.cpp +# overlay-sha256: 555e31c93e5acd1678862d53820ebafaa0dc4a9074c61501b355ce198a76b81a mlx/ops.cpp +# overlay-sha256: 5af152117800d10f56baf0540418aa8a2d80804106563d4a7190a1a357531c5f mlx/primitives.cpp +# overlay-sha256: 5223221a04503b7f9ab5215fc594e57884b1388a1ca84e6b2ad33edb451d2525 mlx/stream.cpp +# note mlx/allocator.h: LOCAL_FIXES.md item 46 (owned ROCm asynchronous Data allocation friendship). # note mlx/backend/common/compiled.cpp: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's negative_strides tracking and the fork's constant-input skipping. # note mlx/backend/gpu/primitives.cpp: LOCAL_FIXES.md item 24: DynamicSliceUpdate copies through copy_gpu (donates only when is_donatable holds) and drops the fork's graph_active() forced donation. # note mlx/fast_primitives.h: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's compile_options and the fork's output_input_aliases on fast::CustomKernel. # note mlx/io/safetensors.cpp: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's empty-tensor guard (as an early continue) and the fork's ROCm staged_write. +--- a/mlx/allocator.h ++++ b/mlx/allocator.h +@@ -16,6 +16,11 @@ + class CommandEncoder; + allocator::Data malloc_async(size_t size, CommandEncoder& encoder); + } // namespace mlx::core::cu ++ ++namespace mlx::core::rocm { ++class CommandEncoder; ++allocator::Data malloc_async(size_t size, CommandEncoder& encoder); ++} // namespace mlx::core::rocm + + namespace mlx::core::allocator { + +@@ -99,6 +104,7 @@ + explicit Data(Buffer buffer) : buffer_(buffer), owned_(true) {} + friend Data malloc(size_t size); + friend Data cu::malloc_async(size_t size, cu::CommandEncoder& encoder); ++ friend Data rocm::malloc_async(size_t size, rocm::CommandEncoder& encoder); + }; + + inline Data malloc(size_t size) { --- a/mlx/backend/common/compiled.cpp +++ b/mlx/backend/common/compiled.cpp @@ -84,17 +84,13 @@ @@ -131,7 +155,7 @@ compute_dynamic_offset(start_indices, out.strides(), axes_, s); --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h -@@ -479,11 +479,8 @@ +@@ -519,11 +519,8 @@ std::vector scalar_arguments, bool is_precompiled, int shared_memory, @@ -144,7 +168,7 @@ : Primitive(stream), name_(std::move(name)), source_(std::move(source)), -@@ -495,11 +492,8 @@ +@@ -535,11 +532,8 @@ scalar_arguments_(std::move(scalar_arguments)), is_precompiled_(is_precompiled), shared_memory_(shared_memory), @@ -157,7 +181,7 @@ void eval_cpu(const std::vector& inputs, std::vector& outputs) override { -@@ -536,12 +530,9 @@ +@@ -576,12 +570,9 @@ std::vector scalar_arguments_; bool is_precompiled_; int shared_memory_; diff --git a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md index 05910f0f1..342a5dab1 100644 --- a/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md +++ b/src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md @@ -4,7 +4,7 @@ Every change this overlay carries on top of the fork commit in `UPSTREAM`, with Name the files an entry touches: `scripts/mlxcelverse/rocm_overlay.py drift` fails for a backend file that differs from the fork but no entry names, and every core-file residual in `CORE_RESIDUAL.diff` carries a note pointing at its entry. `docs/mlxcelverse/upstream/` holds patch packages made for some of these entries before that decision, kept as historical reference. The procedure for pin bumps and fork syncs is `docs/mlxcelverse/rocm-overlay.md`. -## Retarget onto the MLX pin (59d600b5) +## Retarget onto the MLX pin (243bdf09) The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code onto the pin needed: @@ -69,3 +69,8 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code - `mlx/backend/metal/custom_kernel.cpp` and `mlx/backend/cuda/custom_kernel.cpp` from the fork are not carried: a ROCm build does not compile them, and Metal and CUDA builds never copy this directory. - The fork's Python bindings, tests, docs and benchmark scripts are not needed by mlxcel. + +## Retarget compatibility at 243bdf09 (lablup/mlxcel#2303) + +46. **Move-only allocation ownership.** Upstream changed `allocator::malloc` and `array::set_data` to transfer `allocator::Data`, removed the free wrapper and made donation require owned storage. The new core overlay `mlx/allocator.h` adds only the ROCm forward declaration and friendship corresponding to CUDA's existing asynchronous allocation entry point. `mlx/backend/rocm/allocator.h` and `mlx/backend/rocm/allocator.cpp` return genuinely owned `Data` from the public `rocm::malloc_async`; the underlying `RocmAllocator` still returns raw `Buffer` as required by the allocator interface. This preserves ordinary donation and the existing encoder's shared array retention until batch completion; destructing `Data` reaches the same ROCm cache/deferred-free logic exactly once. `mlx/backend/rocm/gemms/naive_gemm.hip` wraps slab views in borrowed `Data` with the existing shell-only deleter, so they cannot donate externally owned slab storage. `mlx/backend/rocm/moe_swiglu.cpp` and `mlx/backend/rocm/reduce/reduce_utils.hpp` remove the obsolete free/deleter arguments when transferring owned allocations. Retargeting `mlx/backend/common/compiled.{h,cpp}` brings in upstream's `Data`-returning allocation callback and avoids borrowed wrappers. The host mock regression in `scripts/mlxcelverse/allocator_ownership_test.{cpp,sh}` compiles the actual header and async handoff, checks move-only/owned/borrowed/empty allocations, pending-batch shared retention, and cleanup on exceptions before and after enqueue. It does not execute HIP or establish GPU completion correctness. +47. **New primitive and capability contracts.** `mlx/backend/rocm/scaled_dot_product_attention.cpp` accepts upstream's `has_sinks` argument and excludes sinks from WMMA admission, matching the existing execution guard; vector and scalar flash keep their sinks support, and unsupported sinks plus LSE requests take the graph fallback (or reject `force_fused`). `mlx/backend/rocm/primitives.cpp` defines `GatedDeltaUpdateVJP::use_fallback` and `eval_gpu`: the capability check always selects upstream's differentiable graph, like the existing forward fallback, while the explicit unsupported-kernel stub satisfies linking without claiming a native implementation. `mlx/backend/rocm/binary.hip` adds `FloorDivide::eval_gpu` using the existing binary dispatch. Its operation in `mlx/backend/rocm/device/binary_ops.hpp`, also shared by `DivMod`, now floors signed and floating negative quotients instead of truncating them, matching target CUDA semantics. `mlx/backend/rocm/compiled.cpp` makes the same correction for generated kernels and keeps integral division exact rather than rounding 64-bit values through float. `mlx/backend/rocm/fft.hip` documents the audited `odd_out` support: the real transform uses the exact output shape for the hipFFT plan length, cache key and inverse scale, so odd inverse-real lengths are not reconstructed as an even `2 * (input - 1)`. No FFT algorithm change is needed. Actual ROCm compile, graph-fallback execution, negative floor-division numerics and odd-length FFT execution remain to be checked on an AMD host; this retarget host has no HIP toolchain or AMD device. diff --git a/src/lib/mlx-cpp/patches-rocm/README.md b/src/lib/mlx-cpp/patches-rocm/README.md index 7b30b25f0..4140a2c75 100644 --- a/src/lib/mlx-cpp/patches-rocm/README.md +++ b/src/lib/mlx-cpp/patches-rocm/README.md @@ -7,7 +7,7 @@ mlxcelverse is the name for everything mlxcel builds on top of upstream MLX: per The tree mirrors the MLX source tree. Every file is a whole-file overlay, copied over the fetched MLX checkout by `mlx_apply_source_overlays` in `../CMakeLists.txt`, and only when `MLX_BUILD_ROCM` is on. Metal and CUDA builds never see these files. A ROCm build also compiles the one overlay in `../patches/` that every build gets, the CPU backend's `mlx/backend/cpu/quantized.cpp` (float accumulation in the scalar quantized matmul, lablup/mlxcel#2248). That file is not part of this overlay, the fork does not change it, and it is not counted below; its delta and refresh rule are in its header and in `../CMakeLists.txt`. - `mlx/backend/rocm/**`: the ROCm backend (109 files). Upstream MLX has no such directory, so these files never conflict with a pin bump; they only need to keep compiling against the API of the pinned commit. -- 15 MLX core files (`CMakeLists.txt`, `mlx/CMakeLists.txt`, `mlx/backend/common/{buffer_cache.h,compiled.cpp,compiled.h}`, `mlx/backend/gpu/primitives.cpp`, `mlx/{compile,device,fast,ops,primitives,stream}.cpp`, `mlx/{fast,fast_primitives}.h`, `mlx/io/safetensors.cpp`): the upstream file at the pinned commit merged with the fork's ROCm hooks, which are mostly `#ifdef MLX_USE_ROCM` blocks. +- 16 MLX core files (`CMakeLists.txt`, `mlx/CMakeLists.txt`, `mlx/backend/common/{buffer_cache.h,compiled.cpp,compiled.h}`, `mlx/backend/gpu/primitives.cpp`, `mlx/{compile,device,fast,ops,primitives,stream}.cpp`, `mlx/{allocator,fast,fast_primitives}.h`, `mlx/io/safetensors.cpp`): the upstream file at the pinned commit merged with the fork's ROCm hooks, which are mostly `#ifdef MLX_USE_ROCM` blocks. Files in this directory other than `CMakeLists.txt` and `mlx/` (this README, `UPSTREAM`, `LOCAL_FIXES.md`, `CORE_RESIDUAL.diff`) are records and are not copied. `CORE_RESIDUAL.diff` is generated by `scripts/mlxcelverse/rocm_overlay.py drift --write`: it is what the core files carry beyond the MLX pin plus the fork's own change to each file, with a note per file naming the `LOCAL_FIXES.md` entry behind it. `make verify-rocm-overlay` (part of `make verify` and `make verify-rocm`) checks that these records agree with the tree and with the MLX pin. @@ -19,7 +19,7 @@ The backend is vendored from the `rocm-support` branch of [NripeshN/mlx](https:/ The full procedure, and the one for syncing to a newer fork commit, is [docs/mlxcelverse/rocm-overlay.md](../../../../docs/mlxcelverse/rocm-overlay.md). In short: -1. `python3 scripts/mlxcelverse/rocm_overlay.py retarget` 3-way merges each of the 15 core files onto the new pin (`git merge-file overlay old-upstream new-upstream`) and updates `retargeted_to_mlx_pin` in `UPSTREAM`. +1. `python3 scripts/mlxcelverse/rocm_overlay.py retarget` 3-way merges each of the 16 core files onto the new pin (`git merge-file overlay old-upstream new-upstream`) and updates `retargeted_to_mlx_pin` in `UPSTREAM`. 2. `scripts/mlxcelverse/check_api_drift.sh` on an AMD host (or a `--features rocm` build) lists the backend's compile errors against the new API and the primitives without a ROCm `eval_gpu`, which show up only as undefined symbols because `libmlx` is static; give those a `NO_GPU` stub in `mlx/backend/rocm/primitives.cpp` or an implementation. 3. `python3 scripts/mlxcelverse/rocm_overlay.py drift` compares every core file line by line with the new pin plus the fork's own change; review what it reports, then `drift --write`. 4. Record every fix in `LOCAL_FIXES.md`, run `make verify-rocm`, and describe the ROCm part in the PR body. diff --git a/src/lib/mlx-cpp/patches-rocm/UPSTREAM b/src/lib/mlx-cpp/patches-rocm/UPSTREAM index fabbf4485..8e6d61be4 100644 --- a/src/lib/mlx-cpp/patches-rocm/UPSTREAM +++ b/src/lib/mlx-cpp/patches-rocm/UPSTREAM @@ -2,7 +2,7 @@ source: https://github.com/NripeshN/mlx branch: rocm-support commit: 75915908dfe5028335d318b10340313744fd3a8d fork_upstream_merge_base: 39886de461d4f014cd60a3a68ad42fd8d776754a -retargeted_to_mlx_pin: 59d600b5e64c238427d0f8d897ab7c682ef4d3d2 +retargeted_to_mlx_pin: 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e license: MIT (see the MLX LICENSE and NOTICE) copied_when: MLX_BUILD_ROCM=ON (cargo feature `rocm`) local_fixes: LOCAL_FIXES.md diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/allocator.h b/src/lib/mlx-cpp/patches-rocm/mlx/allocator.h new file mode 100644 index 000000000..5753fa9ad --- /dev/null +++ b/src/lib/mlx-cpp/patches-rocm/mlx/allocator.h @@ -0,0 +1,129 @@ +// Copyright © 2023 Apple Inc. + +#pragma once + +#include +#include +#include + +#include "mlx/api.h" + +namespace mlx::core::allocator { +class Data; +} // namespace mlx::core::allocator + +namespace mlx::core::cu { +class CommandEncoder; +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); +} // namespace mlx::core::cu + +namespace mlx::core::rocm { +class CommandEncoder; +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); +} // namespace mlx::core::rocm + +namespace mlx::core::allocator { + +// Simple wrapper around buffer pointers +// WARNING: Only Buffer objects constructed from and those that wrap +// raw pointers from mlx::allocator are supported. +class MLX_API Buffer { + private: + void* ptr_; + + public: + explicit Buffer(void* ptr) : ptr_(ptr) {}; + + // Get the raw data pointer from the buffer + void* raw_ptr(); + + // Get the buffer pointer from the buffer + const void* ptr() const { + return ptr_; + }; + void* ptr() { + return ptr_; + }; +}; + +class MLX_API Allocator { + /** Abstract base class for a memory allocator. */ + public: + virtual Buffer malloc(size_t size) = 0; + virtual void free(Buffer buffer) = 0; + virtual size_t size(Buffer buffer) const = 0; + virtual Buffer make_buffer(void* ptr, size_t size) { + return Buffer{nullptr}; + }; + virtual void release(Buffer buffer) {} + + Allocator() = default; + Allocator(const Allocator& other) = delete; + Allocator(Allocator&& other) = delete; + Allocator& operator=(const Allocator& other) = delete; + Allocator& operator=(Allocator&& other) = delete; + virtual ~Allocator() = default; +}; + +MLX_API Allocator& allocator(); + +using Deleter = std::function; + +// The allocator owns data from malloc +class Data { + public: + Data(Buffer buffer, Deleter d) : buffer_(buffer), deleter_(std::move(d)) {} + Data(const Data& other) = delete; + Data& operator=(const Data& other) = delete; + Data(Data&& other) noexcept + : buffer_(std::exchange(other.buffer_, Buffer(nullptr))), + deleter_(std::exchange(other.deleter_, nullptr)), + owned_(std::exchange(other.owned_, false)) {} + ~Data() { + if (owned_) { + allocator().free(buffer_); + } else if (deleter_) { + deleter_(buffer_); + } + } + + Buffer buffer() const& { + return buffer_; + } + Buffer buffer() const&& = delete; + + bool is_owned() const { + return owned_; + } + + private: + Buffer buffer_; + Deleter deleter_; + bool owned_{false}; + + explicit Data(Buffer buffer) : buffer_(buffer), owned_(true) {} + friend Data malloc(size_t size); + friend Data cu::malloc_async(size_t size, cu::CommandEncoder& encoder); + friend Data rocm::malloc_async(size_t size, rocm::CommandEncoder& encoder); +}; + +inline Data malloc(size_t size) { + return Data(allocator().malloc(size)); +} + +// Make a Buffer from a raw pointer of the given size without a copy. If a +// no-copy conversion is not possible then the returned buffer.ptr() will be +// nullptr. Any buffer created with this function must be released with +// release(buffer) +inline Buffer make_buffer(void* ptr, size_t size) { + return allocator().make_buffer(ptr, size); +}; + +// Release a buffer from the allocator made with make_buffer +inline void release(Buffer buffer) { + allocator().release(buffer); +} + +MLX_API bool can_reuse_alien_buffer(void* ptr); + +} // namespace mlx::core::allocator diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.cpp index 46e356251..e039fc43b 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.cpp @@ -116,7 +116,7 @@ void compiled_allocate_outputs( std::vector& outputs, const std::function& is_constant, bool contiguous, - const std::function& + const std::function& mallocfn /* = allocator::malloc */) { if (contiguous) { int o = 0; diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.h b/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.h index be8d6b073..0ee082488 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.h +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/common/compiled.h @@ -37,9 +37,9 @@ void print_float_constant(std::ostream& os, const array& x) { auto old_precision = os.precision(); if constexpr (std::is_same_v) { - os << std::setprecision(std::numeric_limits::digits10 + 1); + os << std::setprecision(std::numeric_limits::max_digits10); } else { - os << std::setprecision(std::numeric_limits::digits10 + 1); + os << std::setprecision(std::numeric_limits::max_digits10); } os << value << std::setprecision(old_precision); } @@ -81,8 +81,7 @@ void compiled_allocate_outputs( std::vector& outputs, const std::function& is_constant, bool contiguous, - const std::function& mallocfn = - allocator::malloc); + const std::function& mallocfn = allocator::malloc); // Collapse contiguous dims ignoring scalars and constants. std::tuple> diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/gpu/primitives.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/gpu/primitives.cpp index 7795e4959..79e6f3f43 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/gpu/primitives.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/gpu/primitives.cpp @@ -258,15 +258,29 @@ void View::eval_gpu(const std::vector& inputs, array& out) { // - type size is the same // - type size is smaller and the last axis is contiguous // - the entire array is row contiguous - if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) || - in.flags().row_contiguous) { + // and the offset is a multiple of the output type size + if (in.offset() % obytes == 0 && + (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) || + in.flags().row_contiguous)) { auto strides = in.strides(); for (int i = 0; i < static_cast(strides.size()) - 1; ++i) { strides[i] *= ibytes; strides[i] /= obytes; } + // Last axis can have any stride + if (obytes < ibytes) { + strides.back() = 1; + } + auto flags = in.flags(); + if (ibytes != obytes) { + // Last axis size can change contiguity + auto [_, row_contiguous, col_contiguous] = + check_contiguity(out.shape(), strides); + flags.row_contiguous = row_contiguous; + flags.col_contiguous = col_contiguous; + } out.copy_shared_buffer( - in, strides, in.flags(), in.data_size() * ibytes / obytes); + in, strides, flags, in.data_size() * ibytes / obytes); } else { auto tmp = array(in.shape(), in.dtype(), nullptr, {}); tmp.set_data(allocator::malloc(tmp.nbytes())); diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.cpp index 1c49b0341..812ac769b 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.cpp @@ -967,11 +967,13 @@ RocmAllocator& allocator() { return *allocator_; } -Buffer malloc_async(size_t size, CommandEncoder& encoder) { - return allocator().malloc_async( +allocator::Data malloc_async(size_t size, CommandEncoder& encoder) { + // Use Data's owned constructor, just like CUDA: ordinary allocations remain + // donatable and release through the ROCm allocator after encoder retention. + return allocator::Data(allocator().malloc_async( size, encoder.device().hip_device(), - static_cast(encoder.stream())); + static_cast(encoder.stream()))); } } // namespace rocm diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.h b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.h index 4ffc7b87c..4b142a908 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.h +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/allocator.h @@ -189,6 +189,6 @@ RocmAllocator& allocator(); size_t thread_allocated_bytes(); class CommandEncoder; -Buffer malloc_async(size_t size, CommandEncoder& encoder); +allocator::Data malloc_async(size_t size, CommandEncoder& encoder); } // namespace mlx::core::rocm diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/binary.hip b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/binary.hip index fd4de2fe2..382698574 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/binary.hip +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/binary.hip @@ -327,6 +327,7 @@ void binary_op_gpu( BINARY_GPU(Add) BINARY_GPU(ArcTan2) BINARY_GPU(Divide) +BINARY_GPU(FloorDivide) BINARY_GPU(Greater) BINARY_GPU(GreaterEqual) BINARY_GPU(Less) diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/compiled.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/compiled.cpp index 553785525..f6600ed91 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/compiled.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/compiled.cpp @@ -638,7 +638,19 @@ struct Remainder { struct FloorDivide { template __device__ T operator()(T x, T y) { - return T(truncf(static_cast(x) / static_cast(y))); + // HIPRTC has no standard-library type traits here. Clang's integral + // predicate keeps 64-bit quotients exact instead of converting to float. + if constexpr (__is_integral(T)) { + auto q = x / y; + if (x % y != 0 && (x < 0) != (y < 0)) { + --q; + } + return T(q); + } else if constexpr (__is_same(T, double)) { + return floor(x / y); + } else { + return T(floorf(static_cast(x) / static_cast(y))); + } } }; diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device/binary_ops.hpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device/binary_ops.hpp index 7ad184c46..2d6fe9587 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device/binary_ops.hpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/device/binary_ops.hpp @@ -27,14 +27,23 @@ struct FloorDivide { template __device__ T operator()(T x, T y) { if constexpr (std::is_integral_v) { - return x / y; + auto q = x / y; + if constexpr (std::is_signed_v) { + if (x % y != 0 && (x < 0) != (y < 0)) { + --q; + } + } + return q; + } else if constexpr (is_complex_v) { + // Complex floor division is rejected by the frontend. + return hipCdivf(x, y); } else if constexpr (std::is_same_v) { return hip_bfloat16( - truncf(static_cast(x) / static_cast(y))); + floorf(static_cast(x) / static_cast(y))); } else if constexpr (std::is_same_v) { - return __float2half(truncf(__half2float(x) / __half2float(y))); + return __float2half(floorf(__half2float(x) / __half2float(y))); } else { - return truncf(x / y); + return floor(x / y); } } }; diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip index 12a7844f8..1506f7b0c 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/fft.hip @@ -461,6 +461,9 @@ void FFT::eval_gpu(const std::vector& inputs, array& out) { Shape step_shape = prepared.arr.shape(); if (step_real) { + // The output shape already encodes upstream's odd_out_ flag. Preserve + // that exact length for C2R (including odd lengths), not 2 * (in - 1). + // execute_fft keys the hipFFT plan and inverse scaling by this length. step_shape.back() = out.shape(axis); } diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/gemms/naive_gemm.hip b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/gemms/naive_gemm.hip index 1ccef3105..a88ebf521 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/gemms/naive_gemm.hip +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/gemms/naive_gemm.hip @@ -2234,9 +2234,11 @@ void moe_ws_alloc(CommandEncoder& encoder, array& a) { /*host_shadow=*/nullptr, /*host_dirty=*/false, /*alloc_stream=*/encoder.stream()}; - a.set_data(allocator::Buffer{rb}, [](allocator::Buffer b) { + // The slab owns the storage. A custom Data deleter releases only this shell + // and marks the array borrowed, so is_donatable() rejects its storage. + a.set_data(allocator::Data(allocator::Buffer{rb}, [](allocator::Buffer b) { delete static_cast(b.ptr()); - }); + })); } void segmented_mm_device( diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/moe_swiglu.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/moe_swiglu.cpp index 25466d68e..b6f647045 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/moe_swiglu.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/moe_swiglu.cpp @@ -270,20 +270,17 @@ void MoeSwigluSorted::eval_gpu( rocm::malloc_async( static_cast(T) * I * size_of(bfloat16), encoder), mid_shape, - bfloat16, - allocator::free); + bfloat16); array up( rocm::malloc_async( static_cast(T) * I * size_of(bfloat16), encoder), mid_shape, - bfloat16, - allocator::free); + bfloat16); array h( rocm::malloc_async( static_cast(T) * I * size_of(bfloat16), encoder), mid_shape, - bfloat16, - allocator::free); + bfloat16); encoder.add_temporary(gate); encoder.add_temporary(up); encoder.add_temporary(h); diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp index 47a1422fc..c32bebf33 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp @@ -57,6 +57,17 @@ bool GatedDeltaUpdate::use_fallback( return true; } NO_GPU_MULTI(GatedDeltaUpdate) +// The forward primitive is always expanded into its graph on ROCm. Its VJP +// likewise uses upstream's differentiable graph, never an unsupported kernel. +bool GatedDeltaUpdateVJP::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + Stream s) { + return true; +} +NO_GPU_MULTI(GatedDeltaUpdateVJP) } // namespace fast // Note: The following are now implemented in their respective files: diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/reduce/reduce_utils.hpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/reduce/reduce_utils.hpp index 38db204ba..a129ce9ae 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/reduce/reduce_utils.hpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/reduce/reduce_utils.hpp @@ -150,8 +150,7 @@ inline void allocate_same_layout( mlx::core::rocm::malloc_async(out.nbytes(), encoder), data_size, final_strides, - fl, - allocator::free); + fl); } } // namespace mlx::core diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/scaled_dot_product_attention.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/scaled_dot_product_attention.cpp index 58f04bd62..20c349d3d 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/scaled_dot_product_attention.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/scaled_dot_product_attention.cpp @@ -149,6 +149,7 @@ bool ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool /*is_training*/, bool output_logsumexp, bool force_fused, @@ -168,15 +169,14 @@ bool ScaledDotProductAttention::use_fallback( // Vector + scalar flash reject output_logsumexp. WMMA flash can emit LSE and // must be counted here or MLX_SDPA_FLASH_VJP never attaches (frontend only // requests LSE when VJP is non-fallback, and VJP needs LSE forward). - // Note: this is static — no has_sinks_ available. Sinks + LSE is rare; eval - // still refuses sinks on the WMMA path. + // WMMA does not support sinks; keep its admission identical to eval_gpu. bool vector_ok = supports_sdpa_vector( q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp); bool flash_ok = supports_sdpa_flash( q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp); #ifdef MLX_HAS_ROCM_WMMA static const bool wmma_disabled = std::getenv("MLX_SDPA_NO_WMMA") != nullptr; - bool wmma_ok = !wmma_disabled && + bool wmma_ok = !wmma_disabled && !has_sinks && supports_sdpa_flash_wmma(q, k, v, has_arr_mask, output_logsumexp) && s.device == Device::gpu && rocm::device(s.device).has_native_wmma() && sdpa_flash_wmma_smem(q.shape(-1)) <= diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/compile.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/compile.cpp index 8beaab712..7a43d283e 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/compile.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/compile.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -50,17 +51,16 @@ bool is_unary(const Primitive& p) { } bool is_binary(const Primitive& p) { - return ( - typeid(p) == typeid(Add) || typeid(p) == typeid(Divide) || - typeid(p) == typeid(Equal) || typeid(p) == typeid(Greater) || - typeid(p) == typeid(GreaterEqual) || typeid(p) == typeid(Less) || - typeid(p) == typeid(LessEqual) || typeid(p) == typeid(LogicalNot) || - typeid(p) == typeid(LogicalAnd) || typeid(p) == typeid(LogicalOr) || - typeid(p) == typeid(LogAddExp) || typeid(p) == typeid(Maximum) || - typeid(p) == typeid(Minimum) || typeid(p) == typeid(Multiply) || - typeid(p) == typeid(NotEqual) || typeid(p) == typeid(Power) || - typeid(p) == typeid(Subtract) || typeid(p) == typeid(BitwiseBinary) || - typeid(p) == typeid(ArcTan2)); + return typeid(p) == typeid(Add) || typeid(p) == typeid(ArcTan2) || + typeid(p) == typeid(BitwiseBinary) || typeid(p) == typeid(Divide) || + typeid(p) == typeid(Equal) || typeid(p) == typeid(FloorDivide) || + typeid(p) == typeid(Greater) || typeid(p) == typeid(GreaterEqual) || + typeid(p) == typeid(Less) || typeid(p) == typeid(LessEqual) || + typeid(p) == typeid(LogicalNot) || typeid(p) == typeid(LogicalAnd) || + typeid(p) == typeid(LogicalOr) || typeid(p) == typeid(LogAddExp) || + typeid(p) == typeid(Maximum) || typeid(p) == typeid(Minimum) || + typeid(p) == typeid(Multiply) || typeid(p) == typeid(NotEqual) || + typeid(p) == typeid(Power) || typeid(p) == typeid(Subtract); } bool is_ternary(const Primitive& p) { @@ -306,9 +306,9 @@ std::uintptr_t get_function_address(const std::function& fun) { class CompileCache { public: struct CacheEntry { - CacheEntry(Stream stream, bool shapeless) - : stream(stream), shapeless(shapeless) {}; - Stream stream; + CacheEntry(std::optional stream, bool shapeless) + : stream(stream), shapeless(shapeless) {} + std::optional stream; bool shapeless; std::vector inputs; std::vector outputs; @@ -374,7 +374,7 @@ class CompileCache { // Loop over entries and check: // - Default stream and device match the entry's default stream // - Inputs match i.e. shapes and types must be equal. - auto stream = default_stream(default_device()); + auto stream = peek_default_stream(default_device()); for (CacheEntry& entry : entries) { // Check that the default stream and device match if (entry.stream != stream) { @@ -654,11 +654,16 @@ void compile_simplify( return pa.is_equivalent(pb); }; - // Merge scalars + std::unordered_map output_map; + for (auto& o : outputs) { + output_map.insert({o.id(), o}); + } + + // Merge scalars, except outputs std::vector new_tape; for (auto& arr : tape) { // Check if we can merge scalars - if (is_scalar(arr)) { + if (is_scalar(arr) && output_map.find(arr.id()) == output_map.end()) { auto scalar = scalars.find(get_scalar_rep(arr)); if (scalar->second.id() != arr.id()) { merge(scalar->second, arr, parents_map); @@ -672,10 +677,6 @@ void compile_simplify( // Remove no-ops { - std::unordered_map output_map; - for (auto& o : outputs) { - output_map.insert({o.id(), o}); - } for (auto& arr : tape) { if (!arr.has_primitive() || !is_noop(arr.primitive())) { new_tape.push_back(std::move(arr)); @@ -697,11 +698,6 @@ void compile_simplify( tape_order.insert({tape[i].id(), i}); } - std::unordered_set output_set; - for (auto& o : outputs) { - output_set.insert(o.id()); - } - // Multi-pass merge only keeping non-orphaned arrays in the tape for (int pass = 0; pass < passes; ++pass) { for (auto& arr : tape) { @@ -721,7 +717,7 @@ void compile_simplify( auto& src = parents->second[src_idx].first; auto& dst = parents->second[dst_idx].first; if (src.id() != dst.id() && array_equivalent(src, dst) && - output_set.find(src.id()) == output_set.end()) { + output_map.find(src.id()) == output_map.end()) { merge(dst, src, parents_map); mask[src_idx] = true; } @@ -787,7 +783,7 @@ void compile_simplify( } return false; } else { - return output_set.find(a.id()) == output_set.end(); + return output_map.find(a.id()) == output_map.end(); } }; bool discard = maybe_merge_parents(arr); diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp index 5d3a9d6d4..211f13c21 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include #include @@ -352,7 +352,7 @@ array layer_norm( auto xc = subtract(x, mu, s); auto v = mean(square(xc, s), /* axis= */ -1, /* keepdims= */ true, s); - x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s)); + x = multiply(xc, rsqrt(add(v, array(eps, float32), s), s), s); x = astype(x, out_type, s); // If the LN is affine then transform x according to the weight and bias @@ -925,6 +925,7 @@ array scaled_dot_product_attention( has_mask, has_arr_mask, do_causal, + has_sinks, is_training, output_logsumexp, force_fused, @@ -935,8 +936,9 @@ array scaled_dot_product_attention( array& mask = inputs[3]; mask = where( mask, - full_like(mask, 0, final_type, s), - full_like(mask, -inf, final_type, s)); + full_like(mask, 0, final_type, stream), + full_like(mask, -inf, final_type, stream), + stream); } Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)}; auto primitive = std::make_shared( @@ -968,13 +970,16 @@ std::vector ScaledDotProductAttention::vjp( const std::vector& outputs) { assert(primals.size() >= 3); assert(cotangents.size() == outputs.size()); - auto s = stream(); if (ScaledDotProductAttentionVJP::use_fallback(primals[0], s)) { assert(outputs.size() == 1); return Custom::vjp(primals, cotangents, argnums, outputs); } + if (outputs.size() != 2) { + return Custom::vjp(primals, cotangents, argnums, outputs); + } + auto fallback = [sdpa = fallback_, s](const std::vector& inputs) { std::vector primals(inputs.begin(), std::prev(inputs.end())); auto [_, vjps] = mlx::core::vjp(sdpa, primals, {inputs.back()}); @@ -1149,6 +1154,54 @@ std::vector gated_delta_update( auto result = fallback({q, k, v, g, beta, h0, mask}); return result; } +std::vector GatedDeltaUpdate::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) { + const int Hk = primals[0].shape(2); + const int Dk = primals[0].shape(3); + const int Hv = primals[2].shape(2); + const int Dv = primals[2].shape(3); + + if (GatedDeltaUpdateVJP::use_fallback(Hk, Dk, Hv, Dv, stream())) { + return Custom::vjp(primals, cotangents, argnums, outputs); + } + + if (primals.size() != 6) { + throw std::runtime_error( + "[GatedDeltaUpdate::vjp] expected 6 primals (q,k,v,g,beta,h0), got " + + std::to_string(primals.size())); + } + if (cotangents.size() != 2) { + throw std::runtime_error( + "[GatedDeltaUpdate::vjp] expected 2 cotangents (out, state), got " + + std::to_string(cotangents.size())); + } + + std::vector inputs = primals; // q,k,v,g,beta,h0 (indices 0..5) + inputs.push_back(cotangents[0]); // cot_o -> index 6 + inputs.push_back(cotangents[1]); // cot_h -> index 7 + + std::vector shapes; + std::vector dtypes; + for (int i = 0; i < 6; ++i) { + shapes.push_back(primals[i].shape()); + dtypes.push_back(i == 5 ? float32 : primals[i].dtype()); // dh is float32 + } + + auto vjps = array::make_arrays( + std::move(shapes), + std::move(dtypes), + std::make_shared(stream(), fallback_), + std::move(inputs)); + + std::vector returned_vjps; + for (int arg : argnums) { + returned_vjps.push_back(std::move(vjps[arg])); + } + return returned_vjps; +} bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/fast_primitives.h b/src/lib/mlx-cpp/patches-rocm/mlx/fast_primitives.h index 14447db25..02a63e235 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/fast_primitives.h +++ b/src/lib/mlx-cpp/patches-rocm/mlx/fast_primitives.h @@ -1,4 +1,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #include #include @@ -289,6 +289,7 @@ class ScaledDotProductAttention : public Custom { bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, @@ -418,6 +419,12 @@ class GatedDeltaUpdate : public Custom { 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(GatedDeltaUpdate); DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { @@ -425,6 +432,39 @@ class GatedDeltaUpdate : public Custom { } private: + bool is_training_; +}; + +class GatedDeltaUpdateVJP : public Custom { + public: + GatedDeltaUpdateVJP( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + static bool use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + Stream s); + + 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(GatedDeltaUpdateVJP); + DEFINE_INPUT_OUTPUT_SHAPE() + auto state() const { + return std::make_tuple(nullptr); /* TODO */ + } + + private: + bool has_cache_; }; class Quantize : public Custom { diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp index 06ef018de..64462fbe6 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp @@ -482,7 +482,7 @@ array unflatten( std::ostringstream msg; msg << "[unflatten] Invalid axes " << ax << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } size_t size = 1; @@ -726,7 +726,7 @@ array flip(const array& a, StreamOrDevice s /* = {} */) { namespace { inline auto -normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) { +normalize_slice(const Shape& shape, Shape& start, Shape& stop, Shape& strides) { // - Start indices are normalized // - End indices are unchanged as -1 means something different // pre-normalization (the end of the axis) versus post normalization (the @@ -774,7 +774,13 @@ normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) { } // Simplify the stride if it's unused if (out_shape[i] == 1) { - strides[i] = 1; + if (strides[i] < 0 && start[i] > 0) { + strides[i] = -1; + stop[i] = start[i] - 1; + } else { + strides[i] = 1; + stop[i] = start[i] + 1; + } } } @@ -817,7 +823,7 @@ void normalize_dynamic_slice_inputs( std::ostringstream msg; msg << prefix << " Invalid axis " << ax << " for array with dimension " << a.ndim() << "."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } ax = new_ax; } @@ -1770,7 +1776,7 @@ array transpose( std::ostringstream msg; msg << "[transpose] Invalid axis (" << ax << ") for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } if (shape[ax] != 0) { throw std::invalid_argument("[transpose] Repeat axes not allowed."); @@ -1794,6 +1800,17 @@ array transpose(const array& a, StreamOrDevice s /* = {} */) { return transpose(a, std::move(axes), to_stream(s)); } +array matrix_transpose(const array& a, StreamOrDevice s /* = {} */) { + if (a.ndim() < 2) { + throw std::invalid_argument( + "[matrix_transpose] Input array must have at least 2 dimensions."); + } + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + std::swap(axes[a.ndim() - 1], axes[a.ndim() - 2]); + return transpose(a, std::move(axes), to_stream(s)); +} + array broadcast_to( const array& a, const Shape& shape, @@ -2333,7 +2350,7 @@ array mean( std::ostringstream msg; msg << "[mean] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } } auto dtype = at_least_float(a.dtype()); @@ -2367,7 +2384,7 @@ array median( std::ostringstream msg; msg << "[median] axis " << axis << " is out of bounds for array with " << ndim << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } set_axes.insert(axis < 0 ? axis + ndim : axis); } @@ -2944,6 +2961,148 @@ array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) { return slice(a_partitioned, slice_starts, slice_ends, s); } +std::vector unique( + const array& a, + int size, + bool return_index /* = false */, + bool return_inverse /* = false */, + bool return_counts /* = false */, + const std::optional& fill_value /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + // Validate args + if (size < 0) { + std::ostringstream msg; + msg << "[unique] Received negative size " << size << "."; + throw std::invalid_argument(msg.str()); + } + if (fill_value && fill_value->size() != 1) { + std::ostringstream msg; + msg << "[unique] Fill value must have one element, but got shape " + << fill_value->shape() << "."; + throw std::invalid_argument(msg.str()); + } + const auto flat = flatten(a, s); + const int n = flat.size(); + + // Handle the edge case of an empty array. + if (n == 0) { + // Without an element there is no default fill value. + if (size > 0 && !fill_value) { + throw std::invalid_argument( + "[unique] A fill value is required for an empty input with a" + " non-zero size."); + } + std::vector out; + out.push_back( + size == 0 ? flat + : full( + {size}, + astype(reshape(*fill_value, {}, s), flat.dtype(), s), + flat.dtype(), + s)); + if (return_index) { + out.push_back(zeros({size}, uint32, s)); + } + if (return_inverse) { + out.push_back(zeros(a.shape(), uint32, s)); + } + if (return_counts) { + out.push_back(zeros({size}, int32, s)); + } + return out; + } + + // Sort the array. The values stay differentiable, the permutation does not. + std::optional order; + if (return_index || return_inverse) { + order = stop_gradient(argsort(flat, 0, s), s); + } + const auto sorted = order ? take(flat, *order, 0, s) : sort(flat, 0, s); + + // Edge detection on the sorted array, true where a new unique + // value starts in the sorted array. + const auto boundary = concatenate( + {array({true}), + not_equal( + slice(sorted, {1}, {n}, s), slice(sorted, {0}, {n - 1}, s), s)}, + 0, + s); + + // Cumsum on boundary gives each sorted element the index of the unique + // value it belongs to. It indexes the output, so it is not differentiable. + const auto group = stop_gradient( + subtract( + cumsum(astype(boundary, uint32, s), 0, false, true, s), + array(1, uint32), + s), + s); + + // Use smallest element of the sorted array (index 0) as padding by default. + const auto fill = fill_value + ? astype(reshape(*fill_value, {}, s), flat.dtype(), s) + : slice(sorted, {0}, {1}, s); + + const int buffer_size = std::max(n, size); + // Scatter positions rather than values in the sorted array (64 bit values + // would fail to scatter on the GPU). + const auto slots = stop_gradient( + slice( + scatter_max( + zeros({buffer_size}, uint32, s), + group, + expand_dims( + where( + boundary, + add(arange(n, uint32, s), array(1, uint32), s), + array(0, uint32), + s), + 1, + s), + 0, + s), + {0}, + {size}, + s), + s); + const auto used = greater(slots, array(0, uint32), s); + const auto positions = + subtract(maximum(slots, array(1, uint32), s), array(1, uint32), s); + + // Build the output arrays. + std::vector out; + out.push_back(where(used, take(sorted, positions, 0, s), fill, s)); + if (return_index) { + // The sort is stable, so the first sorted position of a group holds the + // first occurrence in the input. Unused slots point at position zero, + // which pads with the index of the smallest unique value. + out.push_back(take(*order, positions, 0, s)); + } + if (return_inverse) { + // Clamp so that the indices stay inside a truncated output. Without a + // truncation every group index is already smaller than size. + const auto clamped = + minimum(group, array(std::max(size - 1, 0), uint32), s); + const auto inverse = scatter( + zeros({n}, uint32, s), *order, expand_dims(clamped, 1, s), 0, s); + out.push_back(reshape(inverse, a.shape(), s)); + } + // Padding entries count zero, so counts sum to the input size unless the + // output was truncated. + if (return_counts) { + out.push_back(slice( + scatter_add( + zeros({buffer_size}, int32, s), + group, + ones({n, 1}, int32, s), + 0, + s), + {0}, + {size}, + s)); + } + return out; +} + array logsumexp(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) { std::vector axes(a.ndim()); std::iota(axes.begin(), axes.end(), 0); @@ -3148,14 +3307,13 @@ array floor_divide( const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); - if (issubdtype(dtype, inexact)) { - return floor(divide(a, b, s), s); - } - auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); auto shape = inputs[0].shape(); return array( - shape, dtype, std::make_shared(to_stream(s)), std::move(inputs)); + shape, + dtype, + std::make_shared(to_stream(s)), + std::move(inputs)); } array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) { @@ -3677,7 +3835,7 @@ array take( std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -3722,7 +3880,7 @@ array take(const array& a, int index, int axis, StreamOrDevice s /* = {} */) { std::ostringstream msg; msg << "[take] Received invalid axis " << axis << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } // Check for valid take @@ -4249,7 +4407,7 @@ array diff( int ndim = static_cast(a.ndim()); int ax = axis < 0 ? axis + ndim : axis; if (ax < 0 || ax >= ndim) { - throw std::invalid_argument("[diff] Axis is out of bounds for the array."); + throw std::out_of_range("[diff] Axis is out of bounds for the array."); } if (n < 0) { throw std::invalid_argument("[diff] Order `n` must be non-negative."); @@ -4275,12 +4433,13 @@ array logcumsumexp( bool inclusive /* = true*/, StreamOrDevice s /* = {}*/) { axis = normalize_axis_index(axis, a.ndim(), "[logcumsumexp] "); + auto out_type = at_least_float(a.dtype()); return array( a.shape(), - a.dtype(), + out_type, std::make_shared( to_stream(s), Scan::ReduceType::LogAddExp, axis, reverse, inclusive), - {a}); + {astype(a, out_type, s)}); } array logcumsumexp( @@ -5063,10 +5222,10 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { throw std::invalid_argument(msg.str()); } - if (bits < 2 || bits > 8 || bits == 7) { + if (bits < 1 || bits > 8 || bits == 7) { std::ostringstream msg; msg << "[quantize] The requested number of bits " << bits - << " is not supported. The supported bits are 2, 3, 4, 5, 6 and 8."; + << " is not supported. The supported bits are 1, 2, 3, 4, 5, 6 and 8."; throw std::invalid_argument(msg.str()); } @@ -5087,14 +5246,22 @@ affine_quantize(const array& w, int group_size, int bits, StreamOrDevice s_) { w_max = astype(w_max, float32, s); w_min = astype(w_min, float32, s); - array mask = greater(abs(w_min, s), abs(w_max, s), s); - array scales = - maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); - scales = where(mask, scales, negative(scales, s), s); - array edge = where(mask, w_min, w_max, s); - array q0 = round(divide(edge, scales, s), s); - scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); - array biases = where(equal(q0, zero, s), zero, edge, s); + array scales(0, float32); + array biases(0, float32); + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scales = maximum(subtract(w_max, w_min, s), eps, s); + biases = w_min; + } else { + array mask = greater(abs(w_min, s), abs(w_max, s), s); + scales = maximum(divide(subtract(w_max, w_min, s), n_bins, s), eps, s); + scales = where(mask, scales, negative(scales, s), s); + array edge = where(mask, w_min, w_max, s); + array q0 = round(divide(edge, scales, s), s); + scales = where(not_equal(q0, zero, s), divide(edge, q0, s), scales); + biases = where(equal(q0, zero, s), zero, edge, s); + } packed_w = pack_and_quantize(packed_w, scales, biases, bits, s); @@ -5875,7 +6042,7 @@ array vecdot( } int ax = axis < 0 ? axis + a.ndim() : axis; if (ax < 0 || ax >= a.ndim()) { - throw std::invalid_argument("[vecdot] axis is out of bounds."); + throw std::out_of_range("[vecdot] axis is out of bounds."); } if (axis < 0 ? axis + b.ndim() != ax : axis >= b.ndim()) { throw std::invalid_argument("[vecdot] axis is out of bounds."); @@ -6768,7 +6935,7 @@ array roll( std::ostringstream msg; msg << "[roll] Invalid axis " << axes[i] << " for array with " << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); + throw std::out_of_range(msg.str()); } auto sh = shift[i]; diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/primitives.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/primitives.cpp index 5bb50a487..7540f1100 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/primitives.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/primitives.cpp @@ -399,7 +399,10 @@ std::vector ArcCos::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcCos::jvp( @@ -427,7 +430,10 @@ std::vector ArcCosh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcCosh::jvp( @@ -437,6 +443,14 @@ std::vector ArcCosh::jvp( assert(primals.size() == 1); assert(argnums.size() == 1); array one = array(1., primals[0].dtype()); + if (issubdtype(primals[0].dtype(), complexfloating)) { + // rsqrt(x^2 - 1) has the wrong sign for complex x with Re(x) < 0. + array d = multiply( + rsqrt(subtract(primals[0], one, stream()), stream()), + rsqrt(add(primals[0], one, stream()), stream()), + stream()); + return {multiply(tangents[0], d, stream())}; + } array t = subtract(square(primals[0], stream()), one, stream()); return {multiply(tangents[0], rsqrt(t, stream()), stream())}; } @@ -454,7 +468,10 @@ std::vector ArcSin::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcSin::jvp( @@ -481,7 +498,10 @@ std::vector ArcSinh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcSinh::jvp( @@ -508,7 +528,10 @@ std::vector ArcTan::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcTan::jvp( @@ -596,7 +619,10 @@ std::vector ArcTanh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector ArcTanh::jvp( @@ -1486,7 +1512,7 @@ std::vector Convolution::vjp( int64_t in_size = dilate_size(in.shape(1 + i), input_dilation_[i]); int64_t out_size = dilate_size(cotan.shape(1 + i), kernel_strides_[i]); - padding_hi[i] = safe_cast(in_size - out_size + padding_hi_[i], "conv"); + padding_hi[i] = safe_cast(in_size - out_size + padding_lo_[i], "conv"); } // Check for negative padding @@ -1498,13 +1524,21 @@ std::vector Convolution::vjp( has_neg_padding |= (pd < 0); } + // Negative padding is cropped from the output below + std::vector conv_padding_lo = padding_lo; + std::vector conv_padding_hi = padding_hi; + for (int i = 0; i < conv_padding_lo.size(); ++i) { + conv_padding_lo[i] = std::max(conv_padding_lo[i], 0); + conv_padding_hi[i] = std::max(conv_padding_hi[i], 0); + } + auto wt_trans = group_transpose(wt, 0, 1, -1); auto grad = conv_general( /* const array& input = */ cotan, /* const array& weight = */ wt_trans, /* std::vector stride = */ input_dilation_, - /* std::vector padding_lo = */ padding_lo, - /* std::vector padding_hi = */ padding_hi, + /* std::vector padding_lo = */ conv_padding_lo, + /* std::vector padding_hi = */ conv_padding_hi, /* std::vector kernel_dilation = */ kernel_dilation_, /* std::vector input_dilation = */ kernel_strides_, /* int groups = */ groups_, @@ -1702,7 +1736,10 @@ std::vector Cos::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return {jvp(primals, cotangents, argnums)}; + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Cos::jvp( @@ -2201,12 +2238,6 @@ std::vector Flatten::output_shapes(const std::vector& inputs) { return {Flatten::output_shape(inputs[0], start_axis_, end_axis_)}; } -bool FFT::is_equivalent(const Primitive& other) const { - const FFT& r_other = static_cast(other); - return axes_ == r_other.axes_ && inverse_ == r_other.inverse_ && - real_ == r_other.real_; -} - std::vector Unflatten::vjp( const std::vector&, const std::vector& cotangents, @@ -2272,13 +2303,13 @@ std::pair, std::vector> FFT::vmap( // 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; + out_shape[fft_axes.back()] = inverse_ ? 2 * (n - 1) + odd_out_ : n / 2 + 1; } return { {array( out_shape, real_ && inverse_ ? float32 : complex64, - std::make_shared(stream(), fft_axes, inverse_, real_), + std::make_shared(stream(), fft_axes, inverse_, real_, odd_out_), {in})}, {ax}}; } @@ -2389,6 +2420,12 @@ std::vector FFT::jvp( } } +bool FFT::is_equivalent(const Primitive& other) const { + const FFT& r_other = static_cast(other); + return axes_ == r_other.axes_ && inverse_ == r_other.inverse_ && + real_ == r_other.real_ && odd_out_ == r_other.odd_out_; +} + std::vector Floor::vjp( const std::vector& primals, const std::vector& cotangents, @@ -2414,6 +2451,32 @@ std::pair, std::vector> Floor::vmap( return {{floor(inputs[0], stream())}, axes}; } +std::vector FloorDivide::vjp( + const std::vector& primals, + const std::vector&, + const std::vector& argnums, + const std::vector&) { + std::vector vjps; + for (auto arg : argnums) { + vjps.push_back(zeros_like(primals[arg], stream())); + } + return vjps; +} + +std::vector FloorDivide::jvp( + const std::vector& primals, + const std::vector&, + const std::vector&) { + return {zeros_like(primals[0], stream())}; +} + +std::pair, std::vector> FloorDivide::vmap( + const std::vector& inputs, + const std::vector& axes) { + auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream()); + return {{floor_divide(a, b, stream())}, {to_ax}}; +} + std::vector Full::vjp( const std::vector& primals, const std::vector& cotangents, @@ -4452,6 +4515,49 @@ std::vector Scan::jvp( if (reduce_type_ == Scan::Sum) { return {cumsum(tangents[0], axis_, reverse_, inclusive_, stream())}; + } else if (reduce_type_ == Scan::LogAddExp) { + if (issubdtype(tangents[0].dtype(), complexfloating)) { + throw std::invalid_argument( + "[logcumsumexp] JVP is not supported for complex inputs."); + } + auto x = primals[0]; + auto t = tangents[0]; + auto y = logcumsumexp(x, axis_, reverse_, inclusive_, stream()); + + auto zero = zeros({1}, t.dtype(), stream()); + auto t_min = array(finfo(t.dtype()).min, t.dtype()); + auto log_abs_t = log(abs(t, stream()), stream()); + auto log_t_positive = + where(greater(t, zero, stream()), log_abs_t, t_min, stream()); + auto log_t_negative = + where(less(t, zero, stream()), log_abs_t, t_min, stream()); + + auto out_pos = exp(subtract( + logcumsumexp( + add(log_t_positive, x, stream()), + axis_, + reverse_, + inclusive_, + stream()), + y, + stream())); + + auto out_neg = exp(subtract( + logcumsumexp( + add(log_t_negative, x, stream()), + axis_, + reverse_, + inclusive_, + stream()), + y, + stream())); + + return {where( + isneginf(y, stream()), + zeros_like(y, stream()), + subtract(out_pos, out_neg, stream()), + stream())}; + } else { throw std::runtime_error( "JVP is not implemented for cumulative prod/min/max"); @@ -4619,14 +4725,19 @@ std::pair, std::vector> Scatter::vmap( // Clone updates along the vmap dimension so they can be applied to each // source tensor in the vmap. auto& updates = inputs.back(); + int idx_ndim = static_cast(inputs[1].ndim()); + int upd_src_ax = idx_ndim + src_ax; if (vmap_axes.back() < 0) { - updates = expand_dims( - updates, {0, static_cast(inputs[1].ndim())}, stream()); + updates = expand_dims(updates, {0, upd_src_ax}, stream()); updates = repeat(updates, vmap_size, 0, stream()); } else { - updates = - expand_dims(updates, static_cast(inputs[1].ndim()), stream()); - updates = moveaxis(updates, vmap_axes.back(), 0, stream()); + int upd_vmap_ax = vmap_axes.back(); + // expand_dims shifts any axis at or after the insertion point + if (upd_vmap_ax >= upd_src_ax) { + upd_vmap_ax++; + } + updates = expand_dims(updates, upd_src_ax, stream()); + updates = moveaxis(updates, upd_vmap_ax, 0, stream()); } } @@ -5610,16 +5721,17 @@ std::vector Sqrt::vjp( assert(primals.size() == 1); assert(cotangents.size() == 1); auto dtype = primals[0].dtype(); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). if (recip_) { auto one_over_x_root_x = divide(outputs[0], primals[0], stream()); return {multiply( multiply(array(-0.5, dtype), cotangents[0], stream()), - one_over_x_root_x, + conjugate(one_over_x_root_x, stream()), stream())}; } else { return {divide( multiply(array(0.5, dtype), cotangents[0], stream()), - outputs[0], + conjugate(outputs[0], stream()), stream())}; } } @@ -5628,10 +5740,21 @@ std::vector Sqrt::jvp( const std::vector& primals, const std::vector& tangents, const std::vector& argnums) { + assert(primals.size() == 1); + assert(tangents.size() == 1); + auto dtype = primals[0].dtype(); if (recip_) { - return vjp(primals, tangents, argnums, {rsqrt(primals[0], stream())}); + auto one_over_x_root_x = + divide(rsqrt(primals[0], stream()), primals[0], stream()); + return {multiply( + multiply(array(-0.5, dtype), tangents[0], stream()), + one_over_x_root_x, + stream())}; } else { - return vjp(primals, tangents, argnums, {sqrt(primals[0], stream())}); + return {divide( + multiply(array(0.5, dtype), tangents[0], stream()), + sqrt(primals[0], stream()), + stream())}; } } @@ -6272,7 +6395,15 @@ std::pair, std::vector> View::vmap( const std::vector& axes) { assert(inputs.size() == 1); assert(axes.size() == 1); - return {{view(inputs[0], dtype_, stream())}, axes}; + auto a = inputs[0]; + int ax = axes[0]; + // Changing the item size resizes the last axis, so move the vmapped axis + if (ax == static_cast(a.ndim()) - 1 && + size_of(a.dtype()) != size_of(dtype_)) { + a = moveaxis(a, ax, 0, stream()); + ax = 0; + } + return {{view(a, dtype_, stream())}, {ax}}; } const char* View::name() const { diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/stream.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/stream.cpp index 2dc959974..71f0109a1 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/stream.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/stream.cpp @@ -50,6 +50,10 @@ Stream default_stream(Device d) { return s.value(); } +std::optional peek_default_stream(Device d) { + return default_stream_storage(d); +} + void set_default_stream(Stream s) { if (!gpu::is_available() && s.device == Device::gpu) { throw std::invalid_argument( diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cpu/quantized.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cpu/quantized.cpp index 56bcda060..cd7be442b 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cpu/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cpu/quantized.cpp @@ -1,7 +1,7 @@ // Copyright © 2023-2026 Apple Inc. // Patched by mlxcel: the scalar quantized matmul loops accumulate in float. -// Modified from upstream MLX 81ba1c6a mlx/backend/cpu/quantized.cpp -// Synced to upstream 81ba1c6a. The delta is the bodies of four functions, +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cpu/quantized.cpp +// Synced to upstream 243bdf09. The delta is the bodies of four functions, // `_qmm`, `_qmm_t`, `fp_qmm` and `fp_qmm_t`, each marked `mlxcel:` below; the // SIMD variants and everything else are upstream verbatim. // @@ -395,6 +395,10 @@ void _qmm_dispatch_typed( int bits, bool transposed_w) { switch (bits) { + case 1: + _qmm_dispatch_group( + result, x, w, scales, biases, M, N, K, group_size, transposed_w); + break; case 2: _qmm_dispatch_group( result, x, w, scales, biases, M, N, K, group_size, transposed_w); @@ -420,7 +424,8 @@ void _qmm_dispatch_typed( result, x, w, scales, biases, M, N, K, group_size, transposed_w); break; default: - throw std::invalid_argument("Quantization bits must be 2, 3, 4, 6 or 8."); + throw std::invalid_argument( + "Quantization bits must be 1, 2, 3, 4, 5, 6 or 8."); } } @@ -1233,15 +1238,24 @@ void quantize( w_min = std::min(w_min, (float)w[w_idx + j]); } bool mask = std::abs(w_min) > std::abs(w_max); - float scale = std::max((w_max - w_min) / n_bins, eps); - scale = mask ? scale : -scale; - - float edge = mask ? w_min : w_max; - float q0 = std::rint(edge / scale); - float bias = 0; - if (q0 != 0) { - scale = edge / q0; - bias = edge; + float scale; + float bias; + + if (bits == 1) { + // Affine 1-bit: bit 0 -> w_min, bit 1 -> w_max + scale = std::max(w_max - w_min, eps); + bias = w_min; + } else { + scale = std::max((w_max - w_min) / n_bins, eps); + scale = mask ? scale : -scale; + + float edge = mask ? w_min : w_max; + float q0 = std::rint(edge / scale); + bias = 0; + if (q0 != 0) { + scale = edge / q0; + bias = edge; + } } size_t out_idx = i * int_per_group; for (int j = 0; j < int_per_group / bytes_per_pack; ++j) { diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/binary/binary.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/binary/binary.cuh index 0b444b022..2d7163512 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/binary/binary.cuh +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/binary/binary.cuh @@ -2,7 +2,7 @@ // // CUDA patch: Mixed-precision binary operation kernels for bf16 // -// Modified from upstream MLX e9463bb mlx/backend/cuda/binary/binary.cuh +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/binary/binary.cuh // // Changes: // - Added mixed-type kernel templates (binary_ss_mixed, binary_sv_mixed, etc.) @@ -19,7 +19,7 @@ // GPU time during bf16 model inference. // // The same-type kernels and dispatch below are byte-identical to upstream -// e9463bb (including the large-grid index_rest fix and +// 243bdf09 (including FloorDivide and the large-grid index_rest fix and // get_launch_args_general); the mixed kernels mirror those conventions. // // This patch is CUDA-only -- applied via the overlay system in CMakeLists.txt @@ -464,9 +464,10 @@ __global__ void binary_g_mixed( template constexpr bool supports_binary_op() { if (std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v || - std::is_same_v || std::is_same_v) { + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v) { return std::is_same_v; } if (std::is_same_v || std::is_same_v || diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/cuda_utils.h b/src/lib/mlx-cpp/patches/mlx/backend/cuda/cuda_utils.h deleted file mode 100644 index c98b50b38..000000000 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/cuda_utils.h +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright © 2025 Apple Inc. -// -// CUDA patch: do not throw from ~CudaHandle while the driver is unloading -// -// Modified from upstream MLX 81ba1c6a mlx/backend/cuda/cuda_utils.h -// -// Changes to ~CudaHandle(): -// - Release the handle directly instead of calling reset(), which throws. -// -// Submitted upstream as inureyes:fix/cuda-handle-destructor-no-throw; drop -// this overlay once it lands and the MLX pin moves past it. Rationale and -// measurements: docs/upstream/mlx-cuda-handle-destructor-throws-at-teardown.md -// (lablup/mlxcel#1422). - -#pragma once - -#include -#include - -namespace mlx::core { - -// Throw exception if the cuda API does not succeed. -void check_cuda_error(const char* name, cudaError_t err); -void check_cuda_error(const char* name, CUresult err); - -// The macro version that prints the command that failed. -#define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd)) - -// Base class for RAII managed CUDA resources. -template -class CudaHandle { - public: - CudaHandle(Handle handle = nullptr) : handle_(handle) {} - - CudaHandle(CudaHandle&& other) : handle_(other.handle_) { - assert(this != &other); - other.handle_ = nullptr; - } - - ~CudaHandle() { - // Skip if there was an error to avoid throwing in the destructors - if (cudaPeekAtLastError() != cudaSuccess) { - return; - } - // Not reset(): it throws via CHECK_CUDA_ERROR, and a throw escaping a - // destructor terminates. At exit the runtime may already be unloading, so - // Destroy fails and the handle can no longer be reclaimed or reported on. - if (handle_ != nullptr) { - Destroy(handle_); - handle_ = nullptr; - } - } - - CudaHandle(const CudaHandle&) = delete; - CudaHandle& operator=(const CudaHandle&) = delete; - - CudaHandle& operator=(CudaHandle&& other) { - assert(this != &other); - reset(); - std::swap(handle_, other.handle_); - return *this; - } - - void reset() { - if (handle_ != nullptr) { - CHECK_CUDA_ERROR(Destroy(handle_)); - handle_ = nullptr; - } - } - - Handle release() { - Handle handle = handle_; - handle_ = nullptr; - return handle; - } - - operator Handle() const { - return handle_; - } - - protected: - Handle handle_; -}; - -namespace cu { -class Device; -}; // namespace cu - -// Wrappers of CUDA resources. -class CudaGraph : public CudaHandle { - public: - using CudaHandle::CudaHandle; - explicit CudaGraph(cu::Device& device); - void end_capture(cudaStream_t stream); -}; - -class CudaGraphExec : public CudaHandle { - public: - void instantiate(cudaGraph_t graph); -}; - -class CudaStream : public CudaHandle { - public: - using CudaHandle::CudaHandle; - explicit CudaStream(cu::Device& device); -}; - -} // namespace mlx::core diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh index 20217d392..c02fd501a 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh @@ -2,7 +2,7 @@ // // CUDA patch: drain in-flight cp.async copies before the epilogue smem store // -// Modified from upstream MLX b7c3dd6d mlx/backend/cuda/device/qmm_sm80.cuh +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/device/qmm_sm80.cuh // // Changes to qmm_sm80_mainloop(): // - Add cp_async_wait<0>() + __syncthreads() at the top of the epilogue, @@ -306,6 +306,9 @@ CUTE_DEVICE void qmm_sm80_mainloop( for (int i = 0; i < size(tCrC_accu); i++) { tCrC(i) = Element(tCrC_accu(i)); } + // The last prefetches may still be landing in the buffers C reuses. + cp_async_wait<0>(); + __syncthreads(); copy(r2s_copy_c, r2s_tCrC, r2s_tCsC); __syncthreads(); copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC); @@ -394,4 +397,116 @@ void qmm_sm80_kernel( thread_idx); } +template +__global__ +__launch_bounds__(decltype(size(make_tiled_mma()))::value) +void qmm_sm80_rhs_kernel( + const Element* A, + const Quant* B, + const Scale* S, + const Element* Z, + const uint32_t* rhs_indices, + Element* C, + int m, + int n, + int k, + int group_count) { + int thread_idx = int(threadIdx.x); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + + auto cta_tiler = CtaTiler{}; + int tile_m = int(size<0>(cta_tiler)); + int y_row = m_coord * tile_m; + int m_max_coord = m - y_row; + int tgp_bm = m_max_coord < tile_m ? m_max_coord : tile_m; + + // Define layouts (mixed). All gathered rows form one matrix; each run of + // equal rhs_indices inside the row tile reuses one expert tile. + auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) + auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL) + auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) + auto S_layout = make_scales_layout(n, k, group_count); + + Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, 1), dA); // (M,K,L) + Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, group_count), dB); // (N,K,L) + Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, 1), dC); // (M,N,L) + + Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) + Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) + + Tensor mA = mA_mkl(_,_,0); // (M,K) + Tensor mC = mC_mnl(_,_,0); // (M,N) + uint32_t index_next = rhs_indices[y_row]; + int offset_next = 0; + while (offset_next < tgp_bm) { + int offset = offset_next; + uint32_t b_batch = index_next; + offset_next = tgp_bm; + for (int row = offset + 1; row < tgp_bm; ++row) { + uint32_t index = rhs_indices[y_row + row]; + if (index != b_batch) { + offset_next = row; + index_next = index; + break; + } + } + + int run_rows = offset_next - offset; + Tensor mA_offset = + domain_offset(make_coord(y_row + offset, 0), mA); + Tensor mC_offset = + domain_offset(make_coord(y_row + offset, 0), mC); + Tensor mB = mB_nkl(_,_,b_batch); // (N,K) + Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + + auto run_qmm = [&](auto run_tiler) { + Tensor gA = local_tile(mA_offset, run_tiler, make_coord(0, n_coord, _), Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC_offset, run_tiler, make_coord(0, n_coord, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) + Tensor gS = local_tile(mS, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gZ = local_tile(mZ, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + + qmm_sm80_mainloop( + run_tiler, + gA, + gB, + gS, + gZ, + gC, + run_rows, + thread_idx); + }; + + constexpr int tile_rows = decltype(size<0>(cta_tiler))::value; + if constexpr (tile_rows > 32) { + if (run_rows <= 16) { + run_qmm(make_shape( + Int<16>{}, size<1>(cta_tiler), size<2>(cta_tiler))); + } else if (run_rows <= 32) { + run_qmm(make_shape( + Int<32>{}, size<1>(cta_tiler), size<2>(cta_tiler))); + } else { + run_qmm(cta_tiler); + } + } else if constexpr (tile_rows > 16) { + if (run_rows <= 16) { + run_qmm(make_shape( + Int<16>{}, size<1>(cta_tiler), size<2>(cta_tiler))); + } else { + run_qmm(cta_tiler); + } + } else { + run_qmm(cta_tiler); + } + // The next run's loads must not overwrite C while it is still being stored. + __syncthreads(); + } +} + } // namespace mlx::core::cu diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu index 3b3a6ac48..e0b26a0b5 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu @@ -1,6 +1,6 @@ // Copyright © 2026 Apple Inc. // Patched by mlxcel. -// Modified from upstream MLX b7c3dd6d mlx/backend/cuda/quantized/qmm/qmm_sm80.cu +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/quantized/qmm/qmm_sm80.cu // Changes: // (1) #637: raise the CTA tile_m cap to 128 on consumer Blackwell // (sm_120/121) for large-M prefill shapes, with MLXCEL_QMM_TILE_{M,N,K} @@ -61,6 +61,10 @@ inline auto make_cta_tiler(int m, int group_size, cu::Device& device) { return cute::make_shape(tile_m, tile_n, tile_k); } +inline int get_quantized_batch_count(const array& w) { + return w.ndim() > 2 ? w.size() / (w.shape(-1) * w.shape(-2)) : 1; +} + } // namespace void qmm_sm80( @@ -157,4 +161,125 @@ void qmm_sm80( broadcast_b); } +bool supports_gather_qmm_rhs_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + int out_m = out.ndim() > 1 ? out.shape(-2) : 1; + if (out_m != 1) { + return false; + } + if (!supports_qmm_sm80( + x, + w, + scales, + biases, + out, + transpose, + bits, + group_size, + mode, + device)) { + return false; + } + + int rows = out.size() / out.shape(-1); + int input_rows = x.size() / x.shape(-1); + int group_count = get_quantized_batch_count(w); + return input_rows == rows && rows >= 16 && group_count > 0 && + rows / group_count >= 4; +} + +void gather_qmm_rhs_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& rhs_indices, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder) { + int m = x.size() / x.shape(-1); + int n = out.shape(-1); + int k = x.shape(-1); + int group_count = get_quantized_batch_count(w); + auto cta_tiler = make_cta_tiler(m, group_size); + + std::string module_name = fmt::format( + "gather_qmm_rhs_sm80_tn_{}_m{}_b{}_g{}_{}", + dtype_to_string(x.dtype()), + cute::size<0>(cta_tiler), + bits, + group_size, + quantization_mode_to_string(mode)); + + auto [ctype_x, ctype_q, ctype_s] = get_qmm_cutlass_types(x, bits, mode); + std::string kernel_name = fmt::format( + "mlx::core::cu::qmm_sm80_rhs_kernel<{}, {}, {}, {}, {}>", + group_size, + ctype_x, + ctype_q, + ctype_s, + cta_tiler_to_string(cta_tiler)); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + return std::make_tuple( + false, jit_source_qmm_sm80, std::vector{kernel_name}); + }); + + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (biases) { + encoder.set_input_array(*biases); + } + encoder.set_input_array(rhs_indices); + encoder.set_output_array(out); + + dim3 num_blocks{ + uint32_t(cute::ceil_div(m, cute::size<0>(cta_tiler))), + uint32_t(cute::ceil_div(n, cute::size<1>(cta_tiler))), + 1}; + dim3 block_dims{uint32_t(cute::size(cu::make_tiled_mma()))}; + + auto [sA_layout, sB_layout, sC_layout] = cu::make_smem_layouts(cta_tiler); + size_t smem_bytes = std::max( + cute::cosize(sA_layout) * x.itemsize() + + cute::cosize(sB_layout) * bits / 8, + cute::cosize(sC_layout) * x.itemsize()); + + auto kernel = mod.get_kernel(kernel_name, [&](CUfunction kernel) { + if (smem_bytes > 48000) { + cuFuncSetAttribute( + kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_bytes); + } + }); + + encoder.add_kernel_node_ex( + kernel, + num_blocks, + block_dims, + {}, + smem_bytes, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + biases ? gpu_ptr(*biases) : nullptr, + gpu_ptr(rhs_indices), + gpu_ptr(out), + m, + n, + k, + group_count); +} + } // namespace mlx::core diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp index ff06174cf..040792d29 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp @@ -12,7 +12,7 @@ // dispatch in decode and MLX exposes no hook for it, so the one-shot trace has // to be emitted from the dispatch site itself. // -// Modified from upstream MLX 59d600b5e mlx/backend/cuda/quantized/quantized.cpp +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/quantized/quantized.cpp // // Between 81ba1c6a and 59d600b5e upstream touched this file in // ml-explore/mlx#4507, plumbing `global_scale` into GatherQMM::eval_gpu for @@ -459,8 +459,13 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { // naive kernel until it does. bool can_use_qmm_sm80 = !global_scale.has_value() && supports(supports_qmm_sm80); + bool can_use_gather_qmm_rhs_sm80 = mode_ != QuantizationMode::Affine && + right_sorted_ && !global_scale.has_value() && + supports(supports_gather_qmm_rhs_sm80); bool can_use_qmm_naive = supports(supports_qmm_naive); - bool can_use_qmv = supports(supports_qmv); + bool can_use_fp_gather_qmv = supports(supports_fp_gather_qmv); + bool can_use_qmv = supports(supports_qmv) || can_use_fp_gather_qmv; + bool prefer_qmv = can_use_fp_gather_qmv || (can_use_qmv && M * B < 8); auto call_qmm_sm80 = [&]() { out.set_data(cu::malloc_async(out.nbytes(), encoder)); @@ -507,6 +512,11 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { run_batch_chunked( out, lhs_indices, rhs_indices, M, N, [&](const array& lc, const array& rc, array& oc) { + if (can_use_fp_gather_qmv) { + fp_gather_qmv( + x, w, scales, global_scale, lc, rc, oc, bits_, group_size_, encoder); + return; + } gather_qmv( x, w, @@ -523,8 +533,24 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { }); }; + if (can_use_gather_qmm_rhs_sm80) { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + gather_qmm_rhs_sm80( + x, + w, + scales, + biases, + rhs_indices, + out, + bits_, + group_size_, + mode_, + encoder); + return; + } + if (can_use_qmm_sm80) { - if (can_use_qmv && (M * B < 8)) { + if (prefer_qmv) { call_qmv(); } else { call_qmm_sm80(); @@ -533,7 +559,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { } if (can_use_qmm_naive) { - if (can_use_qmv && (M * B < 8)) { + if (prefer_qmv) { call_qmv(); } else { call_qmm_naive(); diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/col_reduce.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/col_reduce.cu index ecca3d26e..cb14e06df 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/col_reduce.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/col_reduce.cu @@ -2,7 +2,7 @@ // // CUDA patch: bf16 reduce accumulation optimization // -// Modified from upstream MLX 59d600b5e mlx/backend/cuda/reduce/col_reduce.cu +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/reduce/col_reduce.cu // // Changes: // - Added output type template parameter V (separate from accumulation type U) @@ -413,8 +413,7 @@ void col_reduce_two_pass( cu::malloc_async(intermediate.nbytes(), encoder), data_size, intermediate_strides, - fl, - allocator::free); + fl); // Single dispatch on the original input type handles both passes. // This avoids N^2 template instantiations from nested dispatch_all_types. diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp index 8f852fe38..2a7dc45c2 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. // // Patched by mlxcel. The cuDNN execution plan is cached by an LRU keyed on the // exact q/k/v/mask shapes and strides (build_sdpa_cache_key), and a speculative @@ -8,9 +8,9 @@ // fixed floor of that round), and the LRU's lifetime miss counter then aborted // the process after 2 * MLX_CUDA_SDPA_CACHE_SIZE misses (lablup/mlxcel#1799). // -// Modified from upstream MLX 81ba1c6a mlx/backend/cuda/scaled_dot_product_attention.cpp +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/scaled_dot_product_attention.cpp // -// Three mlxcel changes sit on top of upstream 81ba1c6a: +// Three mlxcel changes sit on top of upstream 243bdf09: // // 1. #1820, the general fix: a small multi-row call carrying an array mask over // a fixed-size KV cache is canonicalized the way upstream already @@ -1218,6 +1218,7 @@ bool ScaledDotProductAttention::use_fallback( bool has_mask, bool has_arr_mask, bool do_causal, + bool has_sinks, bool is_training, bool output_logsumexp, bool force_fused, diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp index efef17f6a..2a0e931f9 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp @@ -1,7 +1,7 @@ // Copyright © 2023-2024 Apple Inc. // Patched by mlxcel: a runtime override for the per-command-buffer input -// budget. Synced to upstream 81ba1c6a. -// Modified from upstream MLX 81ba1c6a mlx/backend/metal/device.cpp The delta is three hunks: the +// budget. Synced to upstream 243bdf09. +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/metal/device.cpp The delta is three hunks: the // include below, the override read in `CommandEncoder::needs_commit`, and the // two bridge entry points `mlxcel_set_mb_per_buffer_override()` / // `mlxcel_mb_per_buffer_override()` at the end of the file. Everything else is @@ -231,7 +231,7 @@ MTL::Library* load_default_library(MTL::Device* device) { return lib; } - NS::Error* error[5]; + NS::Error* error[6]; MTL::Library* lib; // First try the colocated mlx.metallib std::tie(lib, error[0]) = load_colocated_library(device, "mlx"); @@ -244,6 +244,22 @@ MTL::Library* load_default_library(MTL::Device* device) { return lib; } +#ifdef SWIFTPM_BUNDLE + // `swift test` with the default Swift Build build system nests resource + // bundles inside the .xctest and runs tests via the toolchain's + // swiftpm-testing-helper, so the running binary is the toolchain helper: + // neither Bundle.main nor Bundle.allBundles names the test bundle, and the + // SwiftPM bundle search below cannot reach it. In that layout the metallib + // sits at a fixed offset from the binary. + std::tie(lib, error[5]) = load_colocated_library( + device, + std::string("../Resources/") + SWIFTPM_BUNDLE + + ".bundle/Contents/Resources/default"); + if (lib) { + return lib; + } +#endif + // Then try default.metallib in a SwiftPM bundle if we have one std::tie(lib, error[2]) = load_swiftpm_library(device, "default"); if (lib) { @@ -262,7 +278,7 @@ MTL::Library* load_default_library(MTL::Device* device) { if (!lib) { std::ostringstream msg; msg << "Failed to load the default metallib. "; - for (int i = 0; i < 5; i++) { + for (int i = 0; i < 6; i++) { if (error[i] != nullptr) { msg << error[i]->localizedDescription()->utf8String() << " "; } @@ -966,8 +982,13 @@ CommandEncoder& get_command_encoder(Stream s) { } std::unordered_map& get_command_encoders() { - static thread_local std::unordered_map encoders; - return encoders; + if (is_main_thread()) { + static auto* encoders = new std::unordered_map(); + return *encoders; + } else { + static thread_local std::unordered_map encoders; + return encoders; + } } std::unordered_map& get_global_command_encoders() { diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/utils.h b/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/utils.h index 9d7100352..f8305effa 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/utils.h +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/kernels/utils.h @@ -3,7 +3,7 @@ // Patched by mlxcel: include and qualify metal::vec // in elem_to_loc_{2,3}_nd for macOS 26.4 Metal compiler compatibility. // -// Modified from upstream MLX 81ba1c6a mlx/backend/metal/kernels/utils.h +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/metal/kernels/utils.h #pragma once @@ -499,3 +499,50 @@ template <> inline bool cast_to(complex64_t val) { return cast_to(val.real) || cast_to(val.imag); } + +/////////////////////////////////////////////////////////////////////////////// +// Row tile scheduling utils +/////////////////////////////////////////////////////////////////////////////// + +// helper function to identify the tile, called from gather_mm / gather_qmm +template +METAL_FUNC bool schedule_row_tile( + const device int32_t* offsets, + const int num_groups, + const int M, + const int tile, + const uint simd_lane_id, + thread int& row, + thread int& group, + thread short& rows) { + int tiles_before = 0; + // each lane process an expert + for (int e = 0; e < num_groups; e += 32) { + const int g = e + simd_lane_id; // shift by lane id + int start = M; + int end = M; + if (g < num_groups) { + // start of the experts activations + start = offsets[g]; + // end of the experts activations + end = g + 1 < num_groups ? offsets[g + 1] : M; + } + // number of tiles per expert + const int n = (end - start + BM - 1) / BM; + // the total number of tiles up to and including this expert + const int tiles_through = tiles_before + simd_prefix_inclusive_sum(n); + const ushort owner_lane = ushort(simd_sum(int(tiles_through <= tile))); + if (owner_lane < 32) { + // if true, we found an owner + group = e + owner_lane; + // fill the row to start + row = simd_shuffle(start, owner_lane) + + (tile - simd_shuffle(tiles_through - n, owner_lane)) * BM; + // number of valid rows in the tile + rows = short(min(BM, simd_shuffle(end, owner_lane) - row)); + return true; + } + tiles_before = simd_shuffle(tiles_through, 31); + } + return false; +} diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp index 8b16e3563..ad0bea424 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp @@ -1,11 +1,10 @@ // Copyright © 2023-2026 Apple Inc. // Patched by mlxcel: `use_qmv_wide` gains an off-switch, `MLXCEL_QMV_WIDE=0`. -// Modified from upstream MLX 59d600b5e mlx/backend/metal/quantized.cpp -// Synced to upstream 59d600b5e. The delta is three hunks: the includes below, -// `mlxcel_qmv_wide_flag()` with the one call it adds to `use_qmv_wide`, and the -// two bridge entry points `mlxcel_set_qmv_wide()` / `mlxcel_qmv_wide()` at the -// end of the file. Everything else is upstream verbatim, so a bump refreshes -// this file and re-applies those three hunks. +// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/metal/quantized.cpp +// Synced to upstream 243bdf09. Retains the QMV-wide runtime control and +// the #1599 gathered-row correction from PR #2138: gather_qmm_rhs receives +// B * M rows after activation broadcasting, in both GatherQMM and the new +// GatherQQMM matrix route. All other implementation is upstream verbatim. // // Why it exists (lablup/mlxcel#1186, #1187). `use_qmv_wide` sends `M >= 2` // affine quantized matmuls down `qmv_wide` on GPU generation 15 and later @@ -32,6 +31,7 @@ #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/matmul.h" #include "mlx/backend/metal/reduce.h" #include "mlx/backend/metal/unary.h" #include "mlx/backend/metal/utils.h" @@ -92,6 +92,9 @@ inline array ensure_row_contiguous_matrix( const array& x, metal::Device& d, const Stream& s) { + if (x.flags().row_contiguous) { + return x; + } if (x.ndim() < 2) { if (x.strides()[0] == 1) { return x; @@ -143,7 +146,7 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } else if (D <= 4096 && O <= 4096) { return 10; } else { - return 6; + return (arch_gen == 13 && arch_size == 's') ? 7 : 6; } } } else { @@ -507,7 +510,8 @@ void qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + bool fast = (K % qmv_fast_k_alignment(bits) == 0); + bool partial_rows = fast && (N % bn != 0); // A narrower output tile reduces register pressure for large // floating-point quantized matrix-vector products on M5 Max GPUs. bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 && @@ -525,9 +529,11 @@ void qmv( group_size, "_b_", bits, - use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", - global_scale ? "_hgs" : ""); + partial_rows ? "_pr_1" : "_pr_0", + global_scale ? "_hgs_1" : "_hgs_0", + "_r_", + results_per_simdgroup); auto kernel = get_quantized_kernel_wrapped( d, kname, @@ -537,6 +543,7 @@ void qmv( group_size, bits, B > 1, + partial_rows, global_scale.has_value(), results_per_simdgroup); @@ -674,11 +681,11 @@ void qmv_wide( } void qvm_split_k( - const array& x, + array x, const array& w, const array& scales, const std::optional& biases, - array& out, + array out, int group_size, int bits, int M, @@ -692,6 +699,16 @@ void qvm_split_k( int split_k = K > 8192 ? 32 : 8; int split_D = (K + split_k - 1) / split_k; int B = out.size() / M / N; + + // When weights are 2D, collapse leading batch dimensions of x and out into M + // so the kernel and reduction operate on standard 2D matrices. + if (B == 1 && x.ndim() != 2) { + x = reshape_in_eval(x, {M, K}, s); + } + if (B == 1 && out.ndim() != 2) { + out = reshape_in_eval(out, {M, N}, s); + } + B *= split_k; constexpr int num_simdgroups = 2; @@ -1248,15 +1265,11 @@ void qmm_splitk( int k_partition_size = K / split_k; int split_k_partition_stride = M * N; - // Allocate intermediate buffer: insert split_k at the front so that - // partition_stride = M * N matches the leading stride of the buffer. + // Allocate a float32 buffer for the partial sums, so that each partition is + // rounded to the output type only once, after they're summed, as in + // steel_matmul_splitk. auto& compute_encoder = metal::get_command_encoder(s); - auto temp_shape = out.shape(); - if (temp_shape.size() == 1) { - temp_shape.insert(temp_shape.begin(), 1); - } - temp_shape.insert(temp_shape.begin(), split_k); - array intermediate(temp_shape, x.dtype(), nullptr, {}); + array intermediate({split_k, M, N}, float32, nullptr, {}); intermediate.set_data(allocator::malloc(intermediate.nbytes())); compute_encoder.add_temporary(intermediate); @@ -1298,13 +1311,18 @@ void qmm_splitk( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - // Sum across split_k dimension (axis 0) - ReductionPlan plan( - ReductionOpType::ContiguousStridedReduce, - {intermediate.shape(0)}, - {intermediate.strides(0)}); - strided_reduce_general_dispatch( - intermediate, out, "sum", plan, {0}, compute_encoder, d, s); + // Sum the partitions and cast to the output type + auto accum_kname = "steel_gemm_splitk_accum_" + type_to_name(out) + "_" + + type_to_name(intermediate); + auto accum_kernel = get_steel_gemm_splitk_accum_kernel( + d, accum_kname, intermediate, out, false); + compute_encoder.set_compute_pipeline_state(accum_kernel); + compute_encoder.set_input_array(intermediate, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(split_k, 2); + compute_encoder.set_bytes(split_k_partition_stride, 3); + compute_encoder.set_bytes(N, 4); + compute_encoder.dispatch_threads(MTL::Size(N, M, 1), get_block_dims(N, M, 1)); } void gather_qmm( @@ -1614,7 +1632,6 @@ void gather_qmm_rhs_nax( int bn = 64, bk = 64; int wm = 2, wn = 2; - const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -1644,7 +1661,6 @@ void gather_qmm_rhs_nax( global_scale ? "_hgs" : ""); metal::MTLFCList func_consts = { - {&align_M, MTL::DataType::DataTypeBool, 200}, {&align_N, MTL::DataType::DataTypeBool, 201}, {&align_K, MTL::DataType::DataTypeBool, 202}, }; @@ -1655,13 +1671,13 @@ void gather_qmm_rhs_nax( concatenate( hash_name, kname, - "_align_M_", - align_M ? 't' : 'n', "_align_N_", align_N ? 't' : 'n', "_align_K_", align_K ? 't' : 'n'); + array offsets = gather_mm_offsets(indices, E, M, d, s); + // Get and set the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_gather_qmm_nax_kernel( @@ -1683,7 +1699,8 @@ void gather_qmm_rhs_nax( compute_encoder.set_compute_pipeline_state(kernel); MTL::Size group_dims(32, wn, wm); - MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + MTL::Size grid_dims( + (N + bn - 1) / bn, std::min(M, (M + bm - 1) / bm + E - 1), 1); compute_encoder.set_input_array(x, 0); compute_encoder.set_input_array(w, 1); @@ -1694,11 +1711,12 @@ void gather_qmm_rhs_nax( compute_encoder.set_input_array(*gs, 3); } int c = 4; - compute_encoder.set_input_array(indices, c++); + compute_encoder.set_input_array(offsets, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(M, c++); compute_encoder.set_bytes(N, c++); compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(E, c++); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -1775,10 +1793,10 @@ void gather_qmm_rhs( } // TODO: Tune the block sizes + int E = w.size() / w.shape(-1) / w.shape(-2); int bm = 16, bn = 32, bk = 32; int wm = 1, wn = 2; - const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; const bool align_K = (K % bk) == 0; @@ -1807,7 +1825,6 @@ void gather_qmm_rhs( global_scale ? "_hgs" : ""); metal::MTLFCList func_consts = { - {&align_M, MTL::DataType::DataTypeBool, 200}, {&align_N, MTL::DataType::DataTypeBool, 201}, {&align_K, MTL::DataType::DataTypeBool, 202}, }; @@ -1818,13 +1835,13 @@ void gather_qmm_rhs( concatenate( hash_name, kname, - "_align_M_", - align_M ? 't' : 'n', "_align_N_", align_N ? 't' : 'n', "_align_K_", align_K ? 't' : 'n'); + array offsets = gather_mm_offsets(indices, E, M, d, s); + // Get and set the kernel auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_gather_qmm_kernel( @@ -1846,7 +1863,8 @@ void gather_qmm_rhs( compute_encoder.set_compute_pipeline_state(kernel); MTL::Size group_dims(32, wn, wm); - MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, 1); + MTL::Size grid_dims( + (N + bn - 1) / bn, std::min(M, (M + bm - 1) / bm + E - 1), 1); compute_encoder.set_input_array(x, 0); compute_encoder.set_input_array(w, 1); @@ -1857,11 +1875,12 @@ void gather_qmm_rhs( compute_encoder.set_input_array(*gs, 3); } int c = 4; - compute_encoder.set_input_array(indices, c++); + compute_encoder.set_input_array(offsets, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(M, c++); compute_encoder.set_bytes(N, c++); compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(E, c++); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -1882,7 +1901,8 @@ void dispatch_qmv( const Stream& s, const std::string& mode) { // It is a qmv with a small inner dimension so route to qmv_quad kernel - if ((K == 128 || K == 64) && is_power_of_2(bits) && !global_scale) { + if ((K == 128 || (K == 64 && bits >= 2)) && is_power_of_2(bits) && + !global_scale) { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } @@ -2049,7 +2069,8 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { transpose_, group_size_, bits_, - x.size() / K, + // mlxcel #1599: include rows broadcast over the gathered slots. + B * M, N, K, d, @@ -2272,7 +2293,8 @@ void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { true, group_size_, bits_, - x.size() / K, + // mlxcel #1599: include rows broadcast over the gathered slots. + B * M, N, K, d, diff --git a/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs b/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs index 1b8d531df..04b45a4dd 100644 --- a/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs +++ b/src/lib/mlxcel-surgery/src/ops/split_mtp_tests.rs @@ -538,6 +538,7 @@ fn staging_directory_is_private() { #[test] fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { let tmp = tempfile::tempdir().unwrap(); + let resolved_parent = std::fs::canonicalize(tmp.path()).unwrap(); let dot_target = tmp.path().join("dot-target"); let dotdot_child = tmp.path().join("dotdot-target").join("child"); std::fs::create_dir(&dot_target).unwrap(); @@ -545,7 +546,7 @@ fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { for output in [dot_target.join("."), dotdot_child.join("..")] { let staging = StagingDir::create(&output).expect("staging"); - assert_eq!(staging.path().parent(), Some(tmp.path())); + assert_eq!(staging.path().parent(), Some(resolved_parent.as_path())); std::fs::write(staging.path().join("config.json"), b"new").unwrap(); commit_staging_dir(staging, &output, true).expect("commit"); assert_eq!(std::fs::read(output.join("config.json")).unwrap(), b"new"); @@ -554,13 +555,34 @@ fn dot_and_dotdot_outputs_stage_beside_the_resolved_directory() { #[cfg(unix)] #[test] -fn non_utf8_output_name_is_supported() { +fn non_utf8_output_name_is_supported_or_rejected_cleanly() { use std::os::unix::ffi::OsStringExt; let tmp = tempfile::tempdir().unwrap(); let output = tmp.path().join(std::ffi::OsString::from_vec(vec![ b'd', b'r', 0xff, b'f', b't', ])); + + if let Err(probe_error) = std::fs::create_dir(&output) { + let staging_error = match StagingDir::create(&output) { + Ok(_) => panic!("staging must honor the filesystem rejection"), + Err(error) => error, + }; + match staging_error { + crate::SurgeryError::Io(error) => { + assert_eq!(error.kind(), probe_error.kind()); + assert_eq!(error.raw_os_error(), probe_error.raw_os_error()); + } + error => panic!("expected filesystem error, got: {error}"), + } + assert!( + std::fs::read_dir(tmp.path()).unwrap().next().is_none(), + "filesystem rejection must not leave output or staging entries" + ); + return; + } + std::fs::remove_dir(&output).unwrap(); + let staging = StagingDir::create(&output).expect("staging"); std::fs::write(staging.path().join("config.json"), b"new").unwrap(); From 6424eb916e2bdfe7c576cba9f8991d9b614ae9d9 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 09:21:04 +0900 Subject: [PATCH 2/6] fix(cuda): pass device to gathered GEMM tile selection The rebased upstream gathered GEMM route called make_cta_tiler with two arguments, while the retained mlxcel selector needs the active CUDA device for its Blackwell tile cap. Pass encoder.device() as the existing quantized GEMM route does, preserving the row count, quantization group size and runtime tile overrides. The missing argument caused the first candidate CUDA build to fail and the later CuTe layout errors were cascading diagnostics. Validation: inspected both selector call sites and checked the diff. Native CUDA compilation is being rerun on the corrected source. Refs #2303 --- .../mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu index e0b26a0b5..b07dcd2a9 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu @@ -212,7 +212,7 @@ void gather_qmm_rhs_sm80( int n = out.shape(-1); int k = x.shape(-1); int group_count = get_quantized_batch_count(w); - auto cta_tiler = make_cta_tiler(m, group_size); + auto cta_tiler = make_cta_tiler(m, group_size, encoder.device()); std::string module_name = fmt::format( "gather_qmm_rhs_sm80_tn_{}_m{}_b{}_g{}_{}", From 5a0d98303f9830ecf9cd34501dc48c3a6db8306f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 09:42:50 +0900 Subject: [PATCH 3/6] test(cuda): cover aligned sorted MXFP gathered GEMM Add 384-wide sorted-prefill reference cases for MXFP4 and MXFP8 with float16 and bfloat16 activations. The existing 320-wide cases fail the new CUDA gathered RHS kernel's output-alignment admission, so they cannot validate its JIT path or the Blackwell 128-row tile policy. Keep those tail cases and reuse the existing expanded routing fixture, host reference and tolerances. Correct the CUDA ops overlay's stale provenance comment. Formatting and diff checks pass; candidate hardware execution follows in the coordinated validation run. Refs #2303 --- src/lib/mlx-cpp/patches-cuda/ops.cpp | 2 +- src/models/switch_layers_mxfp_tests.rs | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/lib/mlx-cpp/patches-cuda/ops.cpp b/src/lib/mlx-cpp/patches-cuda/ops.cpp index 4968f43ee..9cc608b31 100644 --- a/src/lib/mlx-cpp/patches-cuda/ops.cpp +++ b/src/lib/mlx-cpp/patches-cuda/ops.cpp @@ -10,7 +10,7 @@ // - Modified add, subtract, multiply, divide, maximum, minimum to use the // helper, avoiding copy_v kernel insertion for bf16/fp32 mixed operations. // -// Everything else is byte-identical to upstream 81ba1c6a. +// Everything else is byte-identical to the upstream target declared above. // // This patch is CUDA-only -- applied via the overlay system in CMakeLists.txt // only when MLX_BUILD_CUDA is set. Metal builds use the unmodified upstream file. diff --git a/src/models/switch_layers_mxfp_tests.rs b/src/models/switch_layers_mxfp_tests.rs index b69d60d1f..9a03107f2 100644 --- a/src/models/switch_layers_mxfp_tests.rs +++ b/src/models/switch_layers_mxfp_tests.rs @@ -505,6 +505,22 @@ fn mxfp4_gather_qmm_matches_host_reference() { check_gather(Mode::Mxfp4, OUT_FEATURES, &GATHER_CASES, &ACTIVATIONS); } +/// Unlike the 320-wide tail cases, this width admits CUDA's sorted RHS GEMM. +/// gather_sort expands 32 tokens x 4 experts into 128 physical rows, also +/// exercising the consumer Blackwell 128-row CTA policy. +#[test] +fn mxfp_sorted_prefill_aligned_width_matches_host_reference() { + let _device = lock_default_device(); + for mode in [Mode::Mxfp4, Mode::Mxfp8] { + check_gather( + mode, + 384, + &[GatherCase::PrefillSorted], + &[ACTIVATIONS[0], ACTIVATIONS[1]], + ); + } +} + #[test] fn mxfp8_quantized_matmul_matches_host_reference() { let _device = lock_default_device(); From 3aeccc2b7ac39ea17be3e9959cbb4111cfc02782 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 09:45:56 +0900 Subject: [PATCH 4/6] update(cuda): retire the upstream-covered epilogue barrier overlay Upstream #4554 adds cp_async_wait<0>() and a CTA barrier immediately before the quantized GEMM epilogue overwrites the shared-memory A/B buffers. The former local #910 pair precedes only a register-fragment conversion, so upstream now supplies the same async-copy completion and cross-thread ordering for both ordinary and gathered RHS kernels. Retire the whole device header overlay after verifying that removing its provenance and local barrier hunk reproduces pristine 243bdf09 byte for byte. Record the proof and updated overlay inventory; retain the existing JIT module revision boundary. Validation: pristine-object equality assertion and diff check pass. The coordinated CUDA run will restore the retired cached header from the pinned Git object before rebuilding. Refs #2303 --- .../mlx-pin-243bdf09-reconciliation.md | 6 +- .../mlx/backend/cuda/device/qmm_sm80.cuh | 512 ------------------ 2 files changed, 4 insertions(+), 514 deletions(-) delete mode 100644 src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh diff --git a/docs/upstream/mlx-pin-243bdf09-reconciliation.md b/docs/upstream/mlx-pin-243bdf09-reconciliation.md index f42e098bf..1520fee53 100644 --- a/docs/upstream/mlx-pin-243bdf09-reconciliation.md +++ b/docs/upstream/mlx-pin-243bdf09-reconciliation.md @@ -6,7 +6,7 @@ The dependency advances from `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` to `243b ## CPU, CUDA and Metal inventory -The original 32 files comprise 30 upstream-backed overlays and two local extracted helpers. Thirteen upstream-backed files changed relative to their declared bases; twelve were rebased and one was retired. The other seventeen upstream-backed files have zero path-specific commits between their declared base and the target, so their original provenance remains valid. The helpers have no upstream path to replace. There are 31 files after retirement. +The original 32 files comprise 30 upstream-backed overlays and two local extracted helpers. Thirteen upstream-backed files changed relative to their declared bases; eleven were rebased and two were retired. The other seventeen upstream-backed files have zero path-specific commits between their declared base and the target, so their original provenance remains valid. The helpers have no upstream path to replace. There are 30 files after retirement. Paths below are relative to `patches/mlx/backend/`, except `dtype.cpp`, `fast.cpp` and `ops.cpp` under `patches-cuda/`. @@ -18,7 +18,7 @@ Paths below are relative to `patches/mlx/backend/`, except `dtype.cpp`, `fast.cp | `cuda/device/binary_ops.cuh` | `81ba1c6a` | Unchanged; zero drift | Mixed-type operators in runtime-compiled graphs. | | `cuda/device/gemm_sm70.cuh` | `9a795735` | Unchanged; zero drift | Volta tensor-core MMA path. | | `cuda/device/qmm_naive.cuh` | `59d600b5e` | Unchanged; zero drift | Volta MMA dispatch parameterization. | -| `cuda/device/qmm_sm80.cuh` | `b7c3dd6d` | Rebased | Drain cp.async and synchronize before the aliased shared-memory epilogue (#910). | +| `cuda/device/qmm_sm80.cuh` | `b7c3dd6d` | Retired | Upstream #4554 drains cp.async and synchronizes before the aliased shared-memory epilogue (#910). | | `cuda/gemms/grouped_gemm.h` | `b98831ad` | Unchanged; zero drift | Local grouped GEMM dispatch declarations. | | `cuda/gemms/grouped_gemm_arch.h` | `9a795735` | Local helper | Local extracted architecture selector with host-test coverage (#1544). | | `cuda/gemms/grouped_gemm_unaligned.cu` | `c9aa5605` | Unchanged; zero drift | Grouped GEMM architecture selection and local kernel dispatch. | @@ -49,6 +49,8 @@ Three-way conflict review retained the local normalization compute dtype togethe Upstream [#4514](https://github.com/ml-explore/mlx/pull/4514), `9eb3e3ab6fceb4031994f70dfe023872643ed634`, supplies nonthrowing CUDA handle destruction and associated stream/lifetime fixes. The removed overlay had no unrelated local delta; its GLOB registration disappears with the file. The overlapping #4531 was closed unmerged in favor of #4514. +Upstream [#4554](https://github.com/ml-explore/mlx/pull/4554), `5c89fa11e0beb8b22b79d8aeffccb4c935ff1208`, adds `cp_async_wait<0>()` and `__syncthreads()` immediately before the epilogue writes the C tile into shared storage previously used by A/B. The former local #910 barrier appeared before the preceding register-only accumulator conversion; no shared-memory access or additional asynchronous copy occurs between those positions. The upstream pair therefore preserves both async-copy completion and cross-thread read/write ordering for the ordinary and new gathered RHS kernels. Removing the overlay leaves the device header byte-identical to pristine upstream; its only other delta was explanatory provenance. The existing `qmm_sm80_r2` module name remains to preserve the earlier cache boundary. + Upstream [#4594](https://github.com/ml-explore/mlx/pull/4594), `91c83d19b`, provides the Xcode 27 cooperative-tensor workaround. No older `gated_delta_update_nax.h` whole-file overlay is imported from PR #2138. Target `quantized.cpp` still passes the pre-broadcast `x.size() / K` to both gathered matrix routes, so the #1599 correction from #2138 (`7d5459811f16c36d1b05c2a45455f8cb05b99883`) remains necessary and is applied to both as `B * M`. Existing shared-activation MXFP regressions exercise `GatherQMM`; the new NVFP4 matrix route needs appropriate Apple hardware coverage. The same PR's two filesystem-portability test corrections are retained for candidate Metal CI: canonical temporary parents on macOS and explicit clean rejection on filesystems that reject non-UTF-8 names. They do not change production surgery behavior. diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh deleted file mode 100644 index c02fd501a..000000000 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_sm80.cuh +++ /dev/null @@ -1,512 +0,0 @@ -// Copyright © 2026 Apple Inc. -// -// CUDA patch: drain in-flight cp.async copies before the epilogue smem store -// -// Modified from upstream MLX 243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e mlx/backend/cuda/device/qmm_sm80.cuh -// -// Changes to qmm_sm80_mainloop(): -// - Add cp_async_wait<0>() + __syncthreads() at the top of the epilogue, -// before the C tile is staged into shared memory. -// -// Rationale (lablup/mlxcel#910): SharedStorage is a union, so the epilogue's -// C tile aliases the mainloop's A/B pipeline slots. The mainloop only ever -// waits down to K_PIPE_MAX-2 outstanding cp.async groups, and its final -// iteration issues one more (redundant) fetch_gmem of the last K tile, so up -// to two async copy groups into sA/sB are still in flight when the mainloop -// exits. Upstream then writes the C tile into the same shared memory bytes -// with no cp_async_wait and no barrier, so a late-landing async copy of A/B -// data overwrites freshly staged C values before the smem-to-gmem copy reads -// them. compute-sanitizer racecheck reports this as a shared-memory -// write-write hazard in qmm_sm80_kernel, and the visible effect is bitwise -// non-deterministic (occasionally garbage) quantized-matmul output for every -// M >= 8 dispatch: non-deterministic temp-0 prefill and MoE gather_qmm on -// all 4-bit CUDA models. cp_async_wait<0>() retires all outstanding copy -// groups and the barrier re-converges the CTA (also closing the -// cross-thread write-after-read hazard on the final pipeline slot reads) -// before any thread may overwrite the union with C data. -// -// This patch is CUDA-only, applied via the overlay system in CMakeLists.txt. - -#include "mlx/backend/cuda/device/cute_dequant.cuh" - -#include - -// clang-format off - -namespace mlx::core::cu { - -using namespace cute; - -template -union SharedStorage { - struct { - ArrayEngine> A; - ArrayEngine> B; - } mainloop; - struct { - ArrayEngine> C; - } epilogue; -}; - -template -inline constexpr auto make_smem_layouts(CtaTiler cta_tiler) { - // Note: Kernel launcher assumes cosize being same for all KMajor. - auto [bM, bN, bK] = cta_tiler; - - // Define the A/B smem layouts (static). - auto swizzle_ab = composition(Swizzle<3,3,3>{}, - Layout>, - Stride<_8,Stride<_1,_64>>>{}); - auto bP = Int<3>{}; // pipeline - auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP)); - auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP)); - - // Define the C smem layouts (static). - // TODO: Find a better swizzle. - auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN)); - - return cute::make_tuple(sA_layout, sB_layout, sC_layout); -} - -template -inline constexpr auto make_tiled_mma() { - // Note: Kernel launcher assumes num_threads being same for all parameters. - using Atom = cuda::std::conditional_t< - cuda::std::is_same_v, - SM80_16x8x16_F32F16F16F32_TN, - SM80_16x8x16_F32BF16BF16F32_TN>; - if constexpr (TileM >= 32) { - return make_tiled_mma(Atom{}, Layout>{}, Tile<_32,_32,_16>{}); - } else { - return make_tiled_mma(Atom{}, Layout>{}, Tile<_16,_32,_16>{}); - } -} - -template typename Atom, typename NumThreads> -inline constexpr auto make_tiled_copy(NumThreads num_threads) { - return make_tiled_copy( - Copy_Atom>, T>{}, - make_layout(make_shape(Int{}, Int<8>{}), LayoutRight{}), - make_layout(make_shape(Int<1>{}, Int>{}))); -} - -template -CUTE_DEVICE void qmm_sm80_mainloop( - CtaTiler cta_tiler, - TensorA gA, - TensorB gB, - TensorS gS, - TensorZ gZ, - TensorC gC, - int m_max_coord, - int thread_idx) { - // Get the types of operands. - using Element = typename decltype(gA)::value_type; - using Quant = typename decltype(gB)::value_type; - using Scale = typename decltype(gS)::value_type; - - // Define smem layouts. - auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler); - - // Shared memory buffer. - extern __shared__ char smem_buf[]; - using SharedStorage = SharedStorage; - SharedStorage& smem = *reinterpret_cast(smem_buf); - Tensor sA = make_tensor(make_smem_ptr(smem.mainloop.A.begin()), sA_layout); // (BLK_M,BLK_K) - Tensor sB = make_tensor(make_smem_ptr(smem.mainloop.B.begin()), sB_layout); // (BLK_N,BLK_K) - Tensor sC = make_tensor(make_smem_ptr(smem.epilogue.C.begin()), sC_layout); // (BLK_M,BLK_N) - - // Define MMA. - auto mma = make_tiled_mma(cta_tiler)>(); - auto num_threads = size(mma); - - // Define copy atoms. - constexpr int element_bits = sizeof_bits_v; - constexpr int quant_bits = sizeof_bits_v; - constexpr int qload = 128 / (element_bits / quant_bits); - TiledCopy g2s_copy_a = make_tiled_copy(num_threads); - TiledCopy g2s_copy_b = make_tiled_copy(num_threads); - TiledCopy s2g_copy_c = make_tiled_copy(num_threads); - - Copy_Atom s2r_atom_a; - Copy_Atom>, Quant> s2r_atom_b; - Copy_Atom>, Element> r2s_atom_c; - Copy_Atom, Scale> g2r_atom_s; - - // Partition the copying of A/B/C tiles across the threads. - ThrCopy g2s_thr_copy_a = g2s_copy_a.get_slice(thread_idx); - Tensor tAgA = g2s_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) - Tensor tAsA = g2s_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) - - ThrCopy g2s_thr_copy_b = g2s_copy_b.get_slice(thread_idx); - Tensor tBgB = g2s_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) - Tensor tBsB = g2s_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) - - ThrCopy s2g_thr_copy_c = s2g_copy_c.get_slice(thread_idx); - Tensor s2g_tCsC = s2g_thr_copy_c.partition_S(sC); // (CCPY,CCPY_M,CCPY_N) - Tensor s2g_tCgC = s2g_thr_copy_c.partition_D(gC); // (CCPY,CCPY_M,CCPY_N) - - // MMA. - ThrMMA thr_mma = mma.get_slice(thread_idx); - Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K) - Tensor tCsB = thr_mma.partition_B(sB(_,_,0)); // (MMA,MMA_N,MMA_K) - Tensor tCrB = make_fragment_like(tCsB); // (MMA,MMA_N,MMA_K) - Tensor tCrB_dq = make_fragment_like(tCsB); // (MMA,MMA_N,MMA_K) - Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) - Tensor tCrC_accu = make_fragment_like(tCgC); // (MMA,MMA_M,MMA_N) - Tensor tCrC = make_fragment_like(tCgC); // (MMA,MMA_M,MMA_N) - - Tensor tCgS = thr_mma.partition_B(gS); // (MMA,MMA_N,MMA_K,k) - Tensor tCrS = make_tensor_like(tCgS(_,_,_,0)); // (MMA,MMA_N,MMA_K) - Tensor tCgZ = thr_mma.partition_B(gZ); // (MMA,MMA_N,MMA_K,k) - Tensor tCrZ = make_tensor_like(tCgZ(_,_,_,0)); // (MMA,MMA_N,MMA_K) - - // Copy Atom retiling. - TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_a, mma); - ThrCopy s2r_thr_copy_a = s2r_copy_a.get_slice(thread_idx); - Tensor s2r_tCsA = s2r_thr_copy_a.partition_S(sA); // (ACPY,MMA_M,MMA_K,PIPE) - Tensor s2r_tCrA = s2r_thr_copy_a.retile_D(tCrA); // (ACPY,MMA_M,MMA_K) - - TiledCopy s2r_copy_b = make_tiled_copy_B(s2r_atom_b, mma); - ThrCopy s2r_thr_copy_b = s2r_copy_b.get_slice(thread_idx); - Tensor s2r_tCsB = s2r_thr_copy_b.partition_S(sB); // (BCPY,MMA_N,MMA_K,PIPE) - Tensor s2r_tCrB = s2r_thr_copy_b.retile_D(tCrB); // (BCPY,MMA_N,MMA_K) - - TiledCopy r2s_copy_c = make_tiled_copy_C(r2s_atom_c, mma); - ThrCopy r2s_thr_copy_c = r2s_copy_c.get_slice(thread_idx); - Tensor r2s_tCrC = r2s_thr_copy_c.retile_S(tCrC); // (CCPY,MMA_M,MMA_N) - Tensor r2s_tCsC = r2s_thr_copy_c.partition_D(sC); // (CCPY,MMA_M,MMA_N) - - TiledCopy g2r_copy_s = make_tiled_copy_B(g2r_atom_s, mma); - ThrCopy g2r_thr_copy_s = g2r_copy_s.get_slice(thread_idx); - Tensor g2r_tCgS = g2r_thr_copy_s.partition_S(gS); // (BCPY,MMA_N,MMA_K,k) - Tensor g2r_tCrS = g2r_thr_copy_s.retile_D(tCrS); // (BCPY,MMA_N,MMA_K) - Tensor g2r_tCgZ = g2r_thr_copy_s.partition_S(gZ); // (BCPY,MMA_N,MMA_K,k) - Tensor g2r_tCrZ = g2r_thr_copy_s.retile_D(tCrZ); // (BCPY,MMA_N,MMA_K) - - // Predicates for m bound. - Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K) - Tensor tCpC = make_tensor(make_shape(size<1>(s2g_tCsC), size<2>(s2g_tCsC)), Stride<_1,_0>{}); // (CPY_M,CPY_N) - Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) - Tensor cC = make_identity_tensor(make_shape(size<0>(sC), size<1>(sC))); // (BLK_M,BLK_N) - Tensor tAcA = g2s_thr_copy_a.partition_D(cA); // (CPY,CPY_M,CPY_K) - Tensor tCcC = s2g_thr_copy_c.partition_D(cC); // (CPY,CPY_M,CPY_N) - CUTE_UNROLL - for (int m = 0; m < size<0>(tApA); ++m) { - tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord; - } - CUTE_UNROLL - for (int m = 0; m < size<0>(tCpC); ++m) { - tCpC(m,0) = get<0>(tCcC(0,m,0)) < m_max_coord; - } - - auto K_PIPE_MAX = size<3>(tAsA); - int smem_pipe_read = 0; - int smem_pipe_write = 0; - - // Copy A/B: GMEM => SMEM. - auto fetch_gmem = [&](int tile) { - copy_if(g2s_copy_a, tApA, tAgA(_,_,_,tile), tAsA(_,_,_,smem_pipe_write)); - copy(g2s_copy_b, tBgB(_,_,_,tile), tBsB(_,_,_,smem_pipe_write)); - cp_async_fence(); - smem_pipe_write = (smem_pipe_write + 1) % K_PIPE_MAX; - }; - // Copy S/Z: GMEM => RMEM. - auto fetch_scales = [&](int tile) { - copy(g2r_copy_s, g2r_tCgS(_,_,_,tile), g2r_tCrS); - if constexpr (mlx::core::cu::quant_has_bias_v) { - copy(g2r_copy_s, g2r_tCgZ(_,_,_,tile), g2r_tCrZ); - } - }; - // Copy A/B: SMEM => RMEM. - auto fetch_smem = [&](auto block) { - copy(s2r_atom_a, s2r_tCsA(_,_,block,smem_pipe_read), s2r_tCrA(_,_,block)); - copy(s2r_atom_b, s2r_tCsB(_,_,block,smem_pipe_read), s2r_tCrB(_,_,block)); - CUTE_UNROLL - for (int n = 0; n < size<1>(tCrB); ++n) { - mlx::core::cu::cute_vectorized_dequant( - tCrB(_,n,block), - tCrS(_,n,block), - tCrZ(_,n,block), - tCrB_dq(_,n,block)); - } - }; - - auto K_TILE_MAX = size<3>(tAgA); - auto K_BLOCK_MAX = size<2>(tCrA); - - // Prefetch beginning tiles. - int tile_pipe = 0; - CUTE_UNROLL - for (; tile_pipe < K_PIPE_MAX - 1; ++tile_pipe) { - fetch_gmem(tile_pipe); - } - - // Clear accumulators. - clear(tCrC_accu); - - // Prefetch first block. - if constexpr (K_BLOCK_MAX > 1) { - cp_async_wait(); - __syncthreads(); - fetch_scales(0); - fetch_smem(Int<0>{}); - } - - // Loop over CTA tiles. - for (int tile = 0; tile < K_TILE_MAX; ++tile) { - // Unroll MMA blocks. - CUTE_UNROLL - for (int block = 0; block < K_BLOCK_MAX; ++block) { - // Wait for last tile. - if (block == K_BLOCK_MAX - 1) { - smem_pipe_read = (smem_pipe_read + 1) % K_PIPE_MAX; - cp_async_wait(); - __syncthreads(); - fetch_scales((tile + 1 < K_TILE_MAX) ? tile + 1 : tile); - } - // Prefetch next block. - fetch_smem((block + 1) % K_BLOCK_MAX); - // Prefetch next tile. - if (block == 0) { - fetch_gmem(tile_pipe); - tile_pipe = (tile_pipe + 1 < K_TILE_MAX) ? tile_pipe + 1 : tile_pipe; - } - // MMA. - gemm(mma, tCrA(_,_,block), tCrB_dq(_,_,block), tCrC_accu); - } - } - - // Epilogue. - // [mlxcel #910] The C tile below is staged into the same shared memory - // bytes as the A/B pipeline slots (SharedStorage is a union), but up to - // K_PIPE_MAX-2 cp.async groups into sA/sB are still in flight here (the - // mainloop waits only down to K_PIPE_MAX-2 and its last iteration issues a - // redundant fetch_gmem of the final K tile). Retire them all and - // re-converge the CTA before any thread overwrites the union, otherwise a - // late async copy lands on top of staged C values and the smem-to-gmem - // copy publishes corrupted output (non-deterministic temp-0 decode). The - // barrier also orders every thread's final pipeline-slot smem reads before - // the aliasing C writes. - cp_async_wait<0>(); - __syncthreads(); - CUTE_UNROLL - for (int i = 0; i < size(tCrC_accu); i++) { - tCrC(i) = Element(tCrC_accu(i)); - } - // The last prefetches may still be landing in the buffers C reuses. - cp_async_wait<0>(); - __syncthreads(); - copy(r2s_copy_c, r2s_tCrC, r2s_tCsC); - __syncthreads(); - copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC); -} - -template -inline constexpr auto make_scales_layout(int n, int k, int l) { - auto group_size = Int{}; - return make_layout( - make_shape(n, make_shape(group_size, k / group_size), l), - make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size)); -} - -template -__global__ -__launch_bounds__(decltype(size(make_tiled_mma()))::value) -void qmm_sm80_kernel( - const Element* A, - const Quant* B, - const Scale* S, - const Element* Z, - const uint32_t* lhs_indices, - const uint32_t* rhs_indices, - Element* C, - int m, int n, int k, int l, - bool broadcast_b) { - int thread_idx = int(threadIdx.x); - int m_coord = int(blockIdx.x); - int n_coord = int(blockIdx.y); - int l_coord = int(blockIdx.z); - - // Define layouts (mixed). - auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) - auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL) - auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) - auto S_layout = make_scales_layout(n, k, l); - - // Handle broadcasting. - if (broadcast_b) { - get<2>(dB) = 0; - get<2>(stride(S_layout)) = 0; - } - - // Represent the full tensors. - Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, l), dA); // (M,K,L) - Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, l), dB); // (N,K,L) - Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, l), dC); // (M,N,L) - - Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) - Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) - - // For gather, use index lookup for input batch slicing. - uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord; - uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord; - - // Get batch slice. - Tensor mA = mA_mkl(_,_,a_batch); // (M,K) - Tensor mB = mB_nkl(_,_,b_batch); // (N,K) - Tensor mC = mC_mnl(_,_,l_coord); // (M,N) - - Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) - Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) - - // Get the appropriate blocks for this thread block. - auto cta_tiler = CtaTiler{}; - auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) - Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) - Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) - - Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - - // Compute tile residues for predication. - auto m_max_coord = m - size<0>(gA) * m_coord; // M - BLK_M * m_coord - - qmm_sm80_mainloop( - cta_tiler, - gA, - gB, - gS, - gZ, - gC, - m_max_coord, - thread_idx); -} - -template -__global__ -__launch_bounds__(decltype(size(make_tiled_mma()))::value) -void qmm_sm80_rhs_kernel( - const Element* A, - const Quant* B, - const Scale* S, - const Element* Z, - const uint32_t* rhs_indices, - Element* C, - int m, - int n, - int k, - int group_count) { - int thread_idx = int(threadIdx.x); - int m_coord = int(blockIdx.x); - int n_coord = int(blockIdx.y); - - auto cta_tiler = CtaTiler{}; - int tile_m = int(size<0>(cta_tiler)); - int y_row = m_coord * tile_m; - int m_max_coord = m - y_row; - int tgp_bm = m_max_coord < tile_m ? m_max_coord : tile_m; - - // Define layouts (mixed). All gathered rows form one matrix; each run of - // equal rhs_indices inside the row tile reuses one expert tile. - auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) - auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL) - auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) - auto S_layout = make_scales_layout(n, k, group_count); - - Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, 1), dA); // (M,K,L) - Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, group_count), dB); // (N,K,L) - Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, 1), dC); // (M,N,L) - - Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) - Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) - - Tensor mA = mA_mkl(_,_,0); // (M,K) - Tensor mC = mC_mnl(_,_,0); // (M,N) - uint32_t index_next = rhs_indices[y_row]; - int offset_next = 0; - while (offset_next < tgp_bm) { - int offset = offset_next; - uint32_t b_batch = index_next; - offset_next = tgp_bm; - for (int row = offset + 1; row < tgp_bm; ++row) { - uint32_t index = rhs_indices[y_row + row]; - if (index != b_batch) { - offset_next = row; - index_next = index; - break; - } - } - - int run_rows = offset_next - offset; - Tensor mA_offset = - domain_offset(make_coord(y_row + offset, 0), mA); - Tensor mC_offset = - domain_offset(make_coord(y_row + offset, 0), mC); - Tensor mB = mB_nkl(_,_,b_batch); // (N,K) - Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) - Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) - - auto run_qmm = [&](auto run_tiler) { - Tensor gA = local_tile(mA_offset, run_tiler, make_coord(0, n_coord, _), Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) - Tensor gB = local_tile(mB, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gC = local_tile(mC_offset, run_tiler, make_coord(0, n_coord, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) - Tensor gS = local_tile(mS, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gZ = local_tile(mZ, run_tiler, make_coord(0, n_coord, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - - qmm_sm80_mainloop( - run_tiler, - gA, - gB, - gS, - gZ, - gC, - run_rows, - thread_idx); - }; - - constexpr int tile_rows = decltype(size<0>(cta_tiler))::value; - if constexpr (tile_rows > 32) { - if (run_rows <= 16) { - run_qmm(make_shape( - Int<16>{}, size<1>(cta_tiler), size<2>(cta_tiler))); - } else if (run_rows <= 32) { - run_qmm(make_shape( - Int<32>{}, size<1>(cta_tiler), size<2>(cta_tiler))); - } else { - run_qmm(cta_tiler); - } - } else if constexpr (tile_rows > 16) { - if (run_rows <= 16) { - run_qmm(make_shape( - Int<16>{}, size<1>(cta_tiler), size<2>(cta_tiler))); - } else { - run_qmm(cta_tiler); - } - } else { - run_qmm(cta_tiler); - } - // The next run's loads must not overwrite C while it is still being stored. - __syncthreads(); - } -} - -} // namespace mlx::core::cu From 4485843e48aebd63e352929c7b4b4ce4f2c313d3 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 10:04:19 +0900 Subject: [PATCH 5/6] docs: record MLX pin reconciliation and hardware validation Record final CUDA route, combined-tree Kimi inference, and production-source Metal validation for PR #2304. Preserve source and binary provenance, independent #2296 attribution, and unavailable ROCm/M5 and deferred integration coverage. --- ...4-advance-mlx-pin-ownership-20261014.en.md | 92 +++++++++++++++++++ ...4-advance-mlx-pin-ownership-20261014.ko.md | 92 +++++++++++++++++++ 2 files changed, 184 insertions(+) create mode 100644 TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md create mode 100644 TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md diff --git a/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md new file mode 100644 index 000000000..41dbb612f --- /dev/null +++ b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md @@ -0,0 +1,92 @@ +# Technical Report: PR #2304 - Advance MLX Pin and Reconcile Backend Ownership + +- **Date**: 2026-10-14 +- **Status**: Implementation and available hardware validation completed; ordinary PR CI is pending and the expanded coordinated integration gate is deferred +- **Risk Level**: High (a 111-commit dependency update changes CPU, CUDA, Metal, and ROCm integration, and ROCm hardware validation is unavailable) + +## Executive Summary + +PR #2304 advances the pinned `ml-explore/mlx` source from `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` to `243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e`, a range of 111 commits and 312 changed upstream files. It reconciles every CPU/CUDA/Metal whole-file overlay, adapts the ROCm overlay to move-only allocator ownership and new primitives, preserves required numerical and runtime behavior, and adds ownership and aligned MXFP coverage. CUDA compilation and numerics, explicit RHS-route execution, combined-tree Kimi controls, and production-head Metal CI pass. Actual ROCm GPU and M5-specific validation remain unavailable. + +## 1. Problem Statement + +### 1.1 Background + +mlxcel copies local source overlays over a pinned MLX checkout. A pin-only update can silently replace upstream fixes with stale complete files, so each overlay must be compared with pristine Git objects and either rebased, retained with zero path drift, or retired. This pin also changes allocator ownership and primitive APIs across backends. + +### 1.2 Identified Risks + +- **Overlay replacement**: stale whole-file overlays could discard upstream changes across the 111-commit range. +- **Ownership and lifetime**: upstream allocation now returns move-only `allocator::Data`; incorrect handoff could double-release storage, donate borrowed memory, or release buffers before asynchronous work completes. +- **Backend behavior**: CUDA gathered GEMM, Metal gathered-row broadcasting, ROCm fallbacks, SDPA sinks, floor division, gated-delta VJP, and odd-length inverse FFT all cross the updated dependency boundary. +- **Validation attribution**: Kimi requires the independent PR #2296 source fix, so a combined validation tree must record both source heads and must not attribute that fix to PR #2304. + +## 2. Implementation Details + +### 2.1 CPU, CUDA, and Metal Overlay Reconciliation + +The original inventory contained 30 upstream-backed overlays and two local helpers. Eleven changed files were rebased, 17 files were retained with zero path-specific drift, and the two extracted helpers remained local. Two overlays were retired: merged upstream #4514 supplies the CUDA handle overlay's non-throwing teardown semantics, and upstream #4554 supplies the full #910 async-copy completion and CTA ordering residual from `cuda/device/qmm_sm80.cuh`. The overlapping upstream #4531 was closed without merge because #4514 covered it. Removing the #910 provenance and local barrier hunk reproduces the pristine target file byte for byte. The resulting inventory contains 30 retained files. The reconciliation record in `docs/upstream/mlx-pin-243bdf09-reconciliation.md` documents each disposition and retained behavior. + +The retained behavior includes CPU float accumulation, CUDA batch chunking and dtype/runtime controls, upstream gathered GEMM routes, and the #1599 `B * M` broadcast-row correction in both Metal gathered matrix paths. The Xcode 27 cooperative-tensor workaround comes from upstream #4594, avoiding import of the older complete-file workaround. + +### 2.2 Allocation Ownership and ROCm Retarget + +CUDA reduction intermediates now transfer `allocator::Data` directly. ROCm adds an owned-data bridge because the public deleter constructor represents borrowed storage; borrowed slab shells keep their custom release and remain non-donatable. The retarget keeps fork commit `75915908dfe5028335d318b10340313744fd3a8d`, records 109 backend files, 16 core files, 47 local fixes, and five reviewed core residuals, and adapts compiled callbacks, reduction intermediates, SDPA sinks, gated-delta VJP fallback, floor division, and odd inverse FFT output sizing. + +The host ownership regression exercises move-only transfer, owned and borrowed releases, empty allocation, retained shared ownership, and exception cleanup against the actual overlay header. It models pending command ownership but does not execute HIP or prove GPU completion ordering. + +### 2.3 CUDA Gathered GEMM Follow-ups + +Commit `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9` passes `encoder.device()` to the retained gathered GEMM tile selector, fixing the initial candidate build failure. Commit `5a0d98303f9830ecf9cd34501dc48c3a6db8306f` adds N=384 sorted-prefill MXFP4/MXFP8 reference cases because the existing N=320 cases do not satisfy aligned RHS admission. The default 128-row run selects retained #629 grouped GEMM first, so route validation disables it with the documented `MLXCEL_GATHER_QMM_GROUPED=0` control. Both default tile selection and explicit M=128 selection then pass and generate fresh RHS PTX for f16/bf16 × MXFP4/MXFP8. + +## 3. Review Status + +- Source-integration review found no blocking issue through `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9`. +- Independent overlay and allocation-ownership review found no blocking issue. +- The stale CUDA ops provenance comment was corrected in `5a0d98303f9830ecf9cd34501dc48c3a6db8306f`. +- Security review found no critical, high, or medium findings. Its possible low-severity #910 duplicate barrier was resolved by proving upstream #4554 covers the complete functional residual and retiring the overlay. + +## 4. Verification + +### 4.1 Passed + +1. Pin/parser validation resolves exactly `243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e`; all 27 pin/parser tests pass. +2. CPU/CUDA/Metal pristine-object provenance review reports no unexplained path drift or conflict markers. ROCm records and pristine-object drift checks pass, including 29 overlay tooling checks and the host ownership regression. +3. Contract gates pass: version, kernel dtype keys, kernel port dispatch, llama compatibility, Python tooling with 123 tests, and formatting. +4. CUDA no-run compilation at source `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9` completes in 14m23s after the missing device argument was fixed. +5. Ten focused numerical groups at source `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9` complete with zero failures. The logs contain 94 test executions; five MXFP tests are intentionally repeated by the tile-override group, so 94 is an execution count rather than a unique-test count. The frozen binaries are recorded in `candidate-2303-numerics/provenance.json` with root SHA-256 `8477d85748c60dd1298ab71b4c0e2f8c615f516583641b2ccff562a1e5b60f03`, core `e1734a8aae76a7809e5e9d7bcb3046de585b3833144ee1cfe334a14b1ff9b77e`, and surgery `8b873c576b8696806ff125caf2d61e7c637a09b6562fc9ca478181aa3bb1218e`. +6. Production-head Apple Silicon workflow run `38492190049` passes at `3aeccc2b7ac39ea17be3e9959cbb4111cfc02782`: 12,285 tests passed, 0 failed, 390 ignored, plus 51 video tests passed and 0 failed. The actual step logs are archived in `/tmp/mlxcel-auto-20261012T1250/logs/metal-2303-3aeccc2b-pass.zip`. +7. The retired generated-cache header was restored from pristine pin `243bdf09` before the production-head rebuild. `/tmp/mlxcel-auto-20261012T1250/2303-generated-cache-restoration.json` records the cached file changing from SHA-256 `a9046084ccb93381e366e0d9d6730a4fd1cbd0d49e3dff5a43d2df7e66ffbebf` to pristine SHA-256 `c3c4ba7727db7e67929fa4966d6fffa4d11bf47bb8773213c7d8899ec03b9374`. +8. Combined source `53473340` builds in 3m57s. Its final default numerical run passes 38 switch-layer tests, one explicit aligned M=128 test, and six global-scale tests (45 total, 0 failed), recorded in `/tmp/mlxcel-auto-20261012T1250/2303-final-numerics-validation.json`. +9. Explicit RHS validation passes both N=384 runs with `MLXCEL_GATHER_QMM_GROUPED=0`, once with default tile selection and once with `MLXCEL_QMM_TILE_M=128`. A fresh cache contains four `gather_qmm_rhs_sm80` PTX artifacts covering float16 and bfloat16 with MXFP4 and MXFP8; provenance and binary hashes are recorded in `/tmp/mlxcel-auto-20261012T1250/2303-rhs-validation.json`. +10. Combined-tree Kimi validation passes three focused tests and five exact model controls, including answer `4`, marker `AMBER-RAVEN-7429`, chunked and unchunked prefill, and default SDPA. A separately reviewed 58-token output gives a coherent two-sentence Rayleigh-scattering explanation with no token-zero repetition. Source `5347334039e48d9444be77f2aebeaac17c35c6d1`, pin `243bdf09`, and binary SHA-256 `b52d70388e0420582c4f60f3341d570c2afd736685b37e2c0c6e50ed9eaf5fae` are recorded in `2303-latestpin-model-validation.json` and `2303-latestpin-natural-language-validation.json`; the combined tree's independent PR #2296 fix remains separately attributed. + +### 4.2 Pending or Unavailable + +- **Ordinary PR CI**: seven jobs are queued for runner capacity; their queued state is not treated as a failure or as a pass. +- **ROCm**: this host has neither a ROCm toolchain nor AMD hardware. `check_api_drift.sh`, ROCm compilation, real encoder lifetime execution, `make verify-rocm`, and ROCm inference remain unrun. +- **M5-specific execution**: no M5 result is available; another Apple Silicon result is not substituted. +- **Final combined gate**: the user-requested broad integration run is deferred until all prior expanded follow-ups are ready. This PR's candidate checks are not that final whole-batch gate. + +## 5. Change Summary + +| Item | Value | +| --- | ---: | +| Files changed | 47 | +| Lines added | 1,646 | +| Lines deleted | 767 | +| Upstream range | 111 commits / 312 files | +| CPU/CUDA/Metal overlays after retirement | 30 | +| Focused numerical executions | 94 passed / 0 failed (89 unique because 5 MXFP tests repeat) | +| Production-head Metal workflow | 12,285 passed / 0 failed / 390 ignored, plus 51 video tests passed | +| Combined-tree Kimi validation | 3 focused tests + 5 exact controls + 58-token natural-language check passed | + +## 6. Follow-up Actions + +- [ ] Let the seven queued ordinary PR CI jobs complete and review their results. +- [ ] Run ROCm drift/build/runtime gates on an AMD host, or retain the explicit unavailability record. +- [ ] Run the coordinated full CUDA test/Clippy gate and a fresh main Metal nightly only when the expanded batch is ready. + +## Appendix: Evidence Boundaries + +This report records evidence available while PR #2304 is open and draft. It does not authorize merging, does not claim that PR #2304 closes #1599, #2158, or #1986, and does not treat old-pin Kimi results, host-only ROCm ownership tests, or Metal results from another source revision as current-head coverage beyond the stated source-equivalence boundary. diff --git a/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md new file mode 100644 index 000000000..47c81fec2 --- /dev/null +++ b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md @@ -0,0 +1,92 @@ +# 기술 보고서: PR #2304 - MLX 핀 업데이트 및 백엔드 소유권 조정 + +- **작성일**: 2026-10-14 +- **상태**: 구현 및 사용 가능한 하드웨어 검증 완료; 일반 PR CI는 대기 중이며 expanded 결합 통합 게이트는 연기됨 +- **위험도**: 높음 (111개 커밋 규모의 의존성 업데이트가 CPU, CUDA, Metal, ROCm 통합을 변경하며 ROCm 하드웨어 검증은 수행할 수 없음) + +## 요약 + +PR #2304는 고정된 `ml-explore/mlx` 소스를 `59d600b5e64c238427d0f8d897ab7c682ef4d3d2`에서 `243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e`로 업데이트합니다. 이 구간에는 업스트림 커밋 111개와 변경 파일 312개가 포함됩니다. 모든 CPU/CUDA/Metal 전체 파일 오버레이를 조정하고, ROCm 오버레이를 move-only 할당자 소유권과 신규 프리미티브에 맞추며, 필요한 수치 정확도 및 런타임 동작을 보존하고, 소유권 및 정렬 MXFP 테스트를 추가합니다. CUDA 컴파일 및 수치 테스트, 명시적 RHS 경로 실행, 결합 트리 Kimi 컨트롤, 프로덕션 헤드 Metal CI가 통과했습니다. 실제 ROCm GPU 및 M5 전용 검증은 수행할 수 없습니다. + +## 1. 문제 정의 + +### 1.1 배경 + +mlxcel은 고정된 MLX 체크아웃 위에 로컬 소스 오버레이를 복사합니다. 핀만 업데이트하면 오래된 전체 파일이 업스트림 수정사항을 조용히 덮어쓸 수 있으므로 각 오버레이를 순수 Git 객체와 비교하여 리베이스하거나, 경로 드리프트가 0임을 확인하여 유지하거나, 제거해야 합니다. 이번 핀은 백엔드 전반의 할당자 소유권과 프리미티브 API도 변경합니다. + +### 1.2 확인된 위험 + +- **오버레이 덮어쓰기**: 오래된 전체 파일 오버레이가 111개 커밋 구간의 업스트림 변경을 유실시킬 수 있습니다. +- **소유권과 수명**: 업스트림 할당은 이제 move-only `allocator::Data`를 반환합니다. 잘못 전달하면 이중 해제, 빌린 메모리의 donation, 비동기 작업 종료 전 버퍼 해제가 발생할 수 있습니다. +- **백엔드 동작**: CUDA gathered GEMM, Metal gathered-row 브로드캐스트, ROCm 폴백, SDPA sink, floor division, gated-delta VJP, 홀수 길이 inverse FFT가 모두 갱신된 의존성 경계를 통과합니다. +- **검증 귀속**: Kimi는 독립 PR #2296의 소스 수정이 필요합니다. 결합 검증 트리는 두 소스 헤드를 모두 기록하고 해당 모델 수정을 PR #2304의 변경으로 귀속하지 않아야 합니다. + +## 2. 구현 상세 + +### 2.1 CPU, CUDA, Metal 오버레이 조정 + +기존 인벤토리는 업스트림 기반 오버레이 30개와 로컬 헬퍼 2개였습니다. 변경된 파일 11개를 리베이스하고, 경로별 드리프트가 0인 파일 17개와 추출된 로컬 헬퍼 2개를 유지했습니다. 오버레이 2개를 제거했습니다. 병합된 업스트림 #4514는 CUDA 핸들 오버레이의 non-throwing 종료 의미를 제공하며, 업스트림 #4554는 `cuda/device/qmm_sm80.cuh`에 있던 #910 async-copy 완료 및 CTA 순서 residual 전체를 제공합니다. 겹치는 업스트림 #4531은 #4514가 해당 내용을 포함하여 병합 없이 닫혔습니다. #910 provenance와 로컬 barrier hunk를 제거하면 target의 순수 파일과 바이트 단위로 동일합니다. 최종 인벤토리는 유지된 파일 30개입니다. `docs/upstream/mlx-pin-243bdf09-reconciliation.md`는 각 파일의 처리 방식과 보존된 동작을 기록합니다. + +보존된 동작에는 CPU float 누적, CUDA 배치 청킹 및 dtype/런타임 제어, 업스트림 gathered GEMM 경로, Metal gathered matrix 두 경로의 #1599 `B * M` 브로드캐스트 행 수정이 포함됩니다. Xcode 27 cooperative-tensor 우회는 업스트림 #4594를 사용하므로 과거 전체 파일 우회 구현은 가져오지 않습니다. + +### 2.2 할당 소유권 및 ROCm 리타깃 + +CUDA reduction 중간 버퍼는 이제 `allocator::Data`를 직접 전달합니다. 공개 deleter 생성자는 빌린 저장소를 뜻하므로 ROCm은 owned-data 브리지를 추가합니다. 빌린 slab shell은 커스텀 release를 유지하며 donation 대상이 아닙니다. 리타깃은 포크 커밋 `75915908dfe5028335d318b10340313744fd3a8d`를 유지하고, 백엔드 파일 109개, 코어 파일 16개, 로컬 수정 47개, 검토된 코어 residual 5개를 기록합니다. 또한 compiled callback, reduction 중간 버퍼, SDPA sink, gated-delta VJP 폴백, floor division, 홀수 inverse FFT 출력 길이를 새 API에 맞춥니다. + +호스트 소유권 회귀 테스트는 실제 오버레이 헤더를 사용하여 move-only 전달, owned/borrowed 해제, 빈 할당, 공유 소유권 유지, 예외 정리를 검증합니다. pending command 소유권을 모델링하지만 HIP을 실행하거나 GPU 완료 순서를 증명하지는 않습니다. + +### 2.3 CUDA Gathered GEMM 후속 수정 + +커밋 `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9`는 보존된 gathered GEMM 타일 선택기에 `encoder.device()`를 전달하여 최초 후보 빌드 실패를 수정합니다. 커밋 `5a0d98303f9830ecf9cd34501dc48c3a6db8306f`는 기존 N=320 테스트가 정렬 RHS admission을 만족하지 않으므로 N=384 sorted-prefill MXFP4/MXFP8 참조 테스트를 추가합니다. 기본 128행 실행은 유지된 #629 grouped GEMM을 먼저 선택하므로 경로 검증에서는 문서화된 `MLXCEL_GATHER_QMM_GROUPED=0` 제어를 사용합니다. 기본 타일 선택과 명시적 M=128 선택이 모두 통과하고 f16/bf16 × MXFP4/MXFP8에 대한 새로운 RHS PTX를 생성합니다. + +## 3. 리뷰 상태 + +- 소스 통합 리뷰는 `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9`까지 차단 이슈가 없음을 확인했습니다. +- 독립적인 오버레이 및 할당 소유권 리뷰에서도 차단 이슈가 발견되지 않았습니다. +- 오래된 CUDA ops provenance 주석은 `5a0d98303f9830ecf9cd34501dc48c3a6db8306f`에서 수정되었습니다. +- 보안 리뷰에서 critical, high, medium 발견 사항은 없었습니다. 낮은 심각도일 수 있었던 #910 barrier 중복은 업스트림 #4554가 기능 residual 전체를 포함함을 증명하고 오버레이를 제거하여 해결했습니다. + +## 4. 검증 + +### 4.1 통과 + +1. 핀/파서 검증은 정확히 `243bdf09c0cccbb85ca6a686dab3cc0102cf1f7e`를 해석하며 27개 테스트가 모두 통과했습니다. +2. CPU/CUDA/Metal 순수 객체 provenance 검토에서 설명되지 않은 경로 드리프트나 충돌 마커가 없습니다. ROCm 레코드 및 순수 객체 드리프트 검사도 통과했으며 오버레이 도구 검사 29개와 호스트 소유권 회귀 테스트를 포함합니다. +3. 버전, 커널 dtype 키, 커널 포트 디스패치, llama 호환성, Python 도구 123개 테스트, 포맷 컨트랙트 게이트가 통과했습니다. +4. 소스 `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9`의 CUDA no-run 컴파일은 누락된 device 인자를 수정한 뒤 14분 23초에 완료되었습니다. +5. 소스 `6424eb916e2bdfe7c576cba9f8991d9b614ae9d9`에서 10개 집중 수치 그룹이 실패 없이 완료되었습니다. 로그에는 94회 테스트 실행이 기록되며 타일 override 그룹이 MXFP 테스트 5개를 의도적으로 반복하므로 94는 고유 테스트 수가 아니라 실행 횟수입니다. 동결 바이너리는 `candidate-2303-numerics/provenance.json`에 기록되며 SHA-256은 root `8477d85748c60dd1298ab71b4c0e2f8c615f516583641b2ccff562a1e5b60f03`, core `e1734a8aae76a7809e5e9d7bcb3046de585b3833144ee1cfe334a14b1ff9b77e`, surgery `8b873c576b8696806ff125caf2d61e7c637a09b6562fc9ca478181aa3bb1218e`입니다. +6. 프로덕션 헤드 Apple Silicon 워크플로 실행 `38492190049`는 `3aeccc2b7ac39ea17be3e9959cbb4111cfc02782`에서 통과했습니다. 테스트 12,285개 통과, 0개 실패, 390개 무시와 비디오 테스트 51개 통과, 0개 실패를 기록했습니다. 실제 단계 로그는 `/tmp/mlxcel-auto-20261012T1250/logs/metal-2303-3aeccc2b-pass.zip`에 보관됩니다. +7. 제거된 generated-cache 헤더는 프로덕션 헤드 재빌드 전에 순수 핀 `243bdf09`에서 복원되었습니다. `/tmp/mlxcel-auto-20261012T1250/2303-generated-cache-restoration.json`은 캐시 파일 SHA-256이 `a9046084ccb93381e366e0d9d6730a4fd1cbd0d49e3dff5a43d2df7e66ffbebf`에서 순수 파일 `c3c4ba7727db7e67929fa4966d6fffa4d11bf47bb8773213c7d8899ec03b9374`로 변경되었음을 기록합니다. +8. 결합 소스 `53473340`은 3분 57초에 빌드됩니다. 최종 기본 수치 실행에서 switch-layer 테스트 38개, 명시적 정렬 M=128 테스트 1개, global-scale 테스트 6개가 통과했습니다(총 45개, 실패 0개). 결과는 `/tmp/mlxcel-auto-20261012T1250/2303-final-numerics-validation.json`에 기록됩니다. +9. 명시적 RHS 검증은 `MLXCEL_GATHER_QMM_GROUPED=0`을 사용한 N=384 실행 2개를 모두 통과했습니다. 하나는 기본 타일 선택, 다른 하나는 `MLXCEL_QMM_TILE_M=128`을 사용합니다. 새로운 캐시에 float16 및 bfloat16과 MXFP4 및 MXFP8 조합을 포괄하는 `gather_qmm_rhs_sm80` PTX artifact 4개가 생성되었고 provenance 및 바이너리 해시는 `/tmp/mlxcel-auto-20261012T1250/2303-rhs-validation.json`에 기록됩니다. +10. 결합 트리 Kimi 검증은 집중 테스트 3개와 정확 일치 모델 컨트롤 5개를 통과했습니다. 여기에는 답 `4`, 마커 `AMBER-RAVEN-7429`, 청킹 및 비청킹 prefill, 기본 SDPA가 포함됩니다. 별도로 검토한 58-token 출력은 Rayleigh scattering을 정확하고 일관된 두 문장으로 설명하며 token-zero 반복이 없습니다. 소스 `5347334039e48d9444be77f2aebeaac17c35c6d1`, 핀 `243bdf09`, 바이너리 SHA-256 `b52d70388e0420582c4f60f3341d570c2afd736685b37e2c0c6e50ed9eaf5fae`는 `2303-latestpin-model-validation.json`과 `2303-latestpin-natural-language-validation.json`에 기록되며 결합 트리의 독립 PR #2296 수정은 별도로 귀속됩니다. + +### 4.2 대기 또는 수행 불가 + +- **일반 PR CI**: runner 용량을 기다리는 작업 7개가 queued 상태입니다. queued 상태를 실패나 통과로 간주하지 않습니다. +- **ROCm**: 이 호스트에는 ROCm 도구 체인과 AMD 하드웨어가 없습니다. `check_api_drift.sh`, ROCm 컴파일, 실제 encoder 수명 실행, `make verify-rocm`, ROCm 추론은 수행되지 않았습니다. +- **M5 전용 실행**: M5 결과가 없으며 다른 Apple Silicon 결과로 대체하지 않습니다. +- **최종 결합 게이트**: 사용자가 요청한 전체 통합 실행은 모든 expanded 후속 작업이 준비될 때까지 연기되었습니다. 이 PR의 후보 검증은 최종 전체 배치 게이트가 아닙니다. + +## 5. 변경 요약 + +| 항목 | 값 | +| --- | ---: | +| 변경 파일 | 47 | +| 추가 라인 | 1,646 | +| 삭제 라인 | 767 | +| 업스트림 구간 | 커밋 111개 / 파일 312개 | +| 제거 후 CPU/CUDA/Metal 오버레이 | 30개 | +| 집중 수치 테스트 실행 | 94회 통과 / 0회 실패 (MXFP 5개 반복으로 고유 테스트는 89개) | +| 프로덕션 헤드 Metal 워크플로 | 12,285개 통과 / 0개 실패 / 390개 무시, 비디오 테스트 51개 추가 통과 | +| 결합 트리 Kimi 검증 | 집중 테스트 3개 + 정확 일치 컨트롤 5개 + 58-token 자연어 검사 통과 | + +## 6. 후속 조치 + +- [ ] queued 상태인 일반 PR CI 작업 7개가 완료되면 결과를 검토합니다. +- [ ] AMD 호스트에서 ROCm 드리프트/빌드/런타임 게이트를 실행하거나 수행 불가 기록을 유지합니다. +- [ ] expanded 배치가 준비된 뒤에만 전체 CUDA 테스트/Clippy 게이트와 새로운 main Metal nightly를 실행합니다. + +## 부록: 증거 범위 + +이 보고서는 PR #2304가 열린 초안 상태에서 확보된 증거를 기록합니다. 병합 권한을 부여하지 않으며 PR #2304가 #1599, #2158, #1986을 종결한다고 주장하지 않습니다. 또한 명시된 소스 동등성 범위를 넘어서 이전 핀 Kimi 결과, 호스트 전용 ROCm 소유권 테스트, 다른 소스 리비전의 Metal 결과를 현재 헤드 검증으로 간주하지 않습니다. From 4befdbc7afdcc71d36b80f2ac8fbcc5470051a49 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 23:46:35 +0900 Subject: [PATCH 6/6] fix: update rustls to resolve the final CI security advisories The final PR CI found RUSTSEC-2026-0346 through RUSTSEC-2026-0349 in the existing rustls 0.23.45 lockfile entry. Update rustls alone to 0.23.46; cargo deny check passes without adding advisory exceptions. --- Cargo.lock | 4 ++-- .../2304-advance-mlx-pin-ownership-20261014.en.md | 2 ++ .../2304-advance-mlx-pin-ownership-20261014.ko.md | 2 ++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5d7b2e208..22dbe6c9c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3421,9 +3421,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.45" +version = "0.23.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +checksum = "48e13bd8c0e9365c43cfa5c9e8f9ad49d3c8444926c9aac819e0e4dc503c8fdf" dependencies = [ "log", "once_cell", diff --git a/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md index a8bdece41..afac0afcf 100644 --- a/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md +++ b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.en.md @@ -96,3 +96,5 @@ This report records evidence available while PR #2304 is open and draft. It does Merged main through `59974283`, including the #2305 mlxcelverse directory migration and subsequent ROCm, MoE, and request-memory updates. Keep both retired CUDA overlays deleted; relocate `mlx/allocator.h` into `src/lib/mlxcelverse/overlays/rocm/`; fix the host ownership script path, duplicate local-fix numbering (48 entries), and overlay documentation paths/counts. Main's HIP grid limits, SDPA optimization, dequant-cache default, and MoE changes remain intact. Post-merge validation passes: CUDA test compilation, formatting, 27 pin/parser tests, overlay provenance, ROCm residual and 29 tooling checks, host ownership regression, contract gates, and 123 Python tests. Final GPU-serialized execution passes 25 request-memory tests, 38 switch-layer tests, three Inkling MoE tests, 45 DeepSeek V4 tests, and one forced RHS test (112 executions, including the repeated aligned case). Four fresh RHS PTX artifacts prove f16/bf16 × MXFP4/MXFP8 route execution. Logs and exact tree/binary hashes are in `/tmp/mlx-2304-merge-runtime.json`. Prior Metal and actual Kimi evidence remains attributed to its recorded source; those runs were not repeated for this synchronization. + +Final CI also exposed four new RustSec advisories against the existing rustls 0.23.45 lockfile entry (RUSTSEC-2026-0346 through RUSTSEC-2026-0349). Update only rustls to 0.23.46; local `cargo deny check` passes without adding advisory exceptions. diff --git a/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md index 07787facf..faa4fd9f9 100644 --- a/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md +++ b/TECHNICAL_REPORTS/2304-advance-mlx-pin-ownership-20261014.ko.md @@ -96,3 +96,5 @@ CUDA reduction 중간 버퍼는 이제 `allocator::Data`를 직접 전달합니 #2305의 mlxcelverse 디렉터리 이동 및 이후 ROCm·MoE·요청 메모리 변경을 포함한 main `59974283`까지 병합했다. 폐기한 CUDA 오버레이 두 파일은 삭제를 유지하고 `mlx/allocator.h`는 `src/lib/mlxcelverse/overlays/rocm/`으로 이동했다. 호스트 소유권 검사 경로, 로컬 수정 번호 중복(최종 48개), 오버레이 문서 경로·개수를 수정했다. main의 HIP 그리드 제한, SDPA 최적화, dequant 캐시 기본값 및 MoE 변경은 보존했다. 병합 후 CUDA 테스트 빌드, 포맷, pin/parser 27개, 오버레이 provenance, ROCm residual·도구 검사 29개, 호스트 소유권 회귀, 계약 게이트 및 Python 123개가 통과했다. GPU 잠금 아래 요청 메모리 25개, Switch 38개, Inkling MoE 3개, DeepSeek V4 45개, 강제 RHS 1개가 통과했다(정렬 사례 재실행을 포함한 112회). 새 RHS PTX 네 개로 f16/bf16 × MXFP4/MXFP8 경로 실행을 확인했다. 로그와 정확한 소스 트리·바이너리 해시는 `/tmp/mlx-2304-merge-runtime.json`에 있다. 이전 Metal·실제 Kimi 검증은 기록된 소스에 대한 결과이며 이번 동기화에서 재실행하지 않았다. + +최종 CI에서 기존 rustls 0.23.45에 대한 새 RustSec 권고 네 건(RUSTSEC-2026-0346~0349)이 확인됐다. rustls만 0.23.46으로 갱신했고 예외 추가 없이 로컬 `cargo deny check`를 통과했다.