From 7558c9758426fe94c60e4e86aa221fd3e9043b52 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 00:49:24 +0900 Subject: [PATCH 1/2] chore(deps): bump the MLX pin and declare a base commit for each overlay Declare a single machine-readable `// Modified from upstream MLX ` header line across all 32 overlay files under `src/lib/mlx-cpp/patches/` and `src/lib/mlx-cpp/patches-cuda/`, established by diffing against each overlay's copy or rebase point in `ml-explore/mlx`. Bump `GIT_TAG` in `src/lib/mlx-cpp/CMakeLists.txt` from `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a` to `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` (14 upstream commits) and 3-way merge the overlays touched across that range (`patches/mlx/backend/cuda/device/qmm_naive.cuh`, `patches/mlx/backend/cuda/quantized/qmm/qmv.cu`, `patches/mlx/backend/cuda/quantized/quantized.cpp`, `patches/mlx/backend/metal/quantized.cpp`, `patches-cuda/fast.cpp`, `patches-cuda/ops.cpp`), plus the three `patches/mlx/backend/cuda/reduce/` overlays (`all_reduce.cu`, `col_reduce.cu`, `row_reduce.cu`) that were still at `v0.31.1` (`ce45c525`), so the overlay drift check reports 0 upstream commits and 0 `NO BASE` lines across all 32 files. Retarget `src/lib/mlx-cpp/patches-rocm/` onto `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` via `scripts/mlxcelverse/rocm_overlay.py retarget`, add the `fast::GatedDeltaUpdate` fallback and `NO_GPU_MULTI` stub in `mlx/backend/rocm/primitives.cpp`, update `LOCAL_FIXES.md` and `CORE_RESIDUAL.diff`, and document the CUDA/Metal/CPU overlay base-declaration and 3-way merge procedure in `CONTRIBUTING.md`. Closes #1932 --- CONTRIBUTING.md | 39 ++++++ src/lib/mlx-cpp/CMakeLists.txt | 2 +- src/lib/mlx-cpp/patches-cuda/dtype.cpp | 2 +- src/lib/mlx-cpp/patches-cuda/fast.cpp | 129 +++++++++++++++++- src/lib/mlx-cpp/patches-cuda/ops.cpp | 14 +- .../mlx-cpp/patches-rocm/CORE_RESIDUAL.diff | 16 +-- src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md | 4 +- src/lib/mlx-cpp/patches-rocm/UPSTREAM | 2 +- .../mlx/backend/rocm/primitives.cpp | 10 ++ src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp | 127 +++++++++++++++++ src/lib/mlx-cpp/patches-rocm/mlx/fast.h | 10 ++ .../patches-rocm/mlx/fast_primitives.h | 32 +++++ src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp | 12 +- .../patches/mlx/backend/cpu/quantized.cpp | 1 + .../mlx/backend/cuda/device/binary_ops.cuh | 2 + .../mlx/backend/cuda/device/gemm_sm70.cuh | 5 + .../mlx/backend/cuda/device/qmm_naive.cuh | 10 ++ .../mlx/backend/cuda/gemms/grouped_gemm.h | 1 + .../backend/cuda/gemms/grouped_gemm_arch.h | 1 + .../cuda/gemms/grouped_gemm_unaligned.cu | 1 + .../patches/mlx/backend/cuda/jit_module.cpp | 3 +- .../patches/mlx/backend/cuda/jit_module.h | 2 + .../patches/mlx/backend/cuda/matmul.cpp | 4 +- .../backend/cuda/quantized/qmm/qmm_naive.cu | 1 + .../cuda/quantized/qmm/qmm_naive_tile.h | 1 + .../backend/cuda/quantized/qmm/qmm_sm80.cu | 5 +- .../mlx/backend/cuda/quantized/qmm/qmv.cu | 25 +++- .../mlx/backend/cuda/quantized/quantized.cpp | 39 +++--- .../mlx/backend/cuda/reduce/all_reduce.cu | 23 ++-- .../mlx/backend/cuda/reduce/col_reduce.cu | 12 +- .../mlx/backend/cuda/reduce/init_reduce.cu | 3 +- .../mlx/backend/cuda/reduce/reduce_ops.cuh | 3 +- .../mlx/backend/cuda/reduce/row_reduce.cu | 12 +- .../cuda/scaled_dot_product_attention.cpp | 2 + .../cuda/scaled_dot_product_attention.cu | 2 + .../patches/mlx/backend/metal/compiled.cpp | 1 + .../patches/mlx/backend/metal/device.cpp | 3 +- .../patches/mlx/backend/metal/kernels/utils.h | 5 + .../patches/mlx/backend/metal/quantized.cpp | 120 +++++++++++++++- 39 files changed, 603 insertions(+), 83 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d53b389bf..1c07452ae 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -147,6 +147,45 @@ After bumping the pin, re-validate the in-tree fused Metal kernel launchers in ` All three should produce output within RMS < 5e-3 of the graph reference on Apple Silicon. +#### Reconciling CUDA, Metal, and CPU source overlays + +The overlays in `src/lib/mlx-cpp/patches/` and `src/lib/mlx-cpp/patches-cuda/` are applied whole-file with `configure_file(... COPYONLY)`, so any upstream commit to an overlaid file is silently reverted when the pin moves unless the overlay is rebased first. Every file under those two directories therefore carries a single machine-readable base line in its header comment: + +```cpp +// Modified from upstream MLX +``` + +`` is the upstream MLX commit (7–40 lowercase hex characters, resolvable with `git rev-parse` in an `ml-explore/mlx` clone) that the overlay was copied from or last rebased onto, established by diffing against upstream rather than stamping the current pin onto untouched files. + +Before and after bumping `GIT_TAG`, run the drift check against an `ml-explore/mlx` checkout (such as `$OUT_DIR/build/_deps/mlx-src` or `~/.cache/mlxcel/mlxcelverse/mlx.git`): + +```bash +MLX_SRC= +PIN=$(bash scripts/ci/mlx_pinned_commit.sh) +for f in $(find src/lib/mlx-cpp/patches src/lib/mlx-cpp/patches-cuda -type f); do + case "$f" in + */patches-cuda/*) up="mlx/$(basename "$f")" ;; + *) up="${f#src/lib/mlx-cpp/patches/}" ;; + esac + base=$(sed -n 's|^// *Modified from upstream MLX \([0-9a-f]\{7,40\}\).*|\1|p' "$f" | head -1) + [ -n "$base" ] || { echo "NO BASE $f"; continue; } + n=$(git -C "$MLX_SRC" log --oneline "$base..$PIN" -- "$up" | wc -l) + [ "$n" -eq 0 ] || echo "DRIFT $n $f (base $base)" +done +``` + +For each file reported as `DRIFT`, 3-way merge the overlay from its declared `` onto the new pin: + +```bash +git -C "$MLX_SRC" show ":" > /tmp/mlx-base +git -C "$MLX_SRC" show "$PIN:" > /tmp/mlx-theirs +git merge-file /tmp/mlx-base /tmp/mlx-theirs +``` + +Resolve any conflicts, diff `` line by line against `/tmp/mlx-theirs` to verify that every mlxcel delta documented in the file's header survived and no upstream change was dropped, and update that overlay's `// Modified from upstream MLX ` line to the new pin. When done, the drift check must produce no `NO BASE` and no `DRIFT` lines against `$PIN`. If upstream now carries the fix natively, retire the overlay instead of rebasing it. + +#### Reconciling the ROCm overlay + The ROCm overlay in [`src/lib/mlx-cpp/patches-rocm/`](src/lib/mlx-cpp/patches-rocm/README.md) follows the pin too, and `make verify-rocm-overlay` (part of `make verify`) fails until it does: `patches-rocm/UPSTREAM` must name the new pin. The procedure is [`docs/mlxcelverse/rocm-overlay.md`](docs/mlxcelverse/rocm-overlay.md): `scripts/mlxcelverse/rocm_overlay.py retarget` 3-way merges the overlay's MLX core files onto the new pin; on an AMD host `scripts/mlxcelverse/check_api_drift.sh` lists the backend's compile errors and the primitives left without a ROCm `eval_gpu` (undefined symbols, since `libmlx` is static), which get a fix or a `NO_GPU` stub; `rocm_overlay.py drift` compares each core file line by line with the new pin plus the fork's own change and must match the reviewed `patches-rocm/CORE_RESIDUAL.diff`, which is what catches an overlay that was already wrong before the bump. Record each fix in `patches-rocm/LOCAL_FIXES.md`, run `make verify-rocm` on the AMD host (it includes the generation smoke from #1811), and add the "ROCm overlay" section described in that document to the PR body and to the technical report if there is one. ## Development environment diff --git a/src/lib/mlx-cpp/CMakeLists.txt b/src/lib/mlx-cpp/CMakeLists.txt index 93ec619c9..c55f8f14a 100644 --- a/src/lib/mlx-cpp/CMakeLists.txt +++ b/src/lib/mlx-cpp/CMakeLists.txt @@ -200,7 +200,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 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a) + GIT_TAG 59d600b5e64c238427d0f8d897ab7c682ef4d3d2) # 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/dtype.cpp b/src/lib/mlx-cpp/patches-cuda/dtype.cpp index aae36fe77..5b257ac41 100644 --- a/src/lib/mlx-cpp/patches-cuda/dtype.cpp +++ b/src/lib/mlx-cpp/patches-cuda/dtype.cpp @@ -2,7 +2,7 @@ // // CUDA patch: bf16 type promotion preservation // -// Modified from upstream MLX v0.31.1 mlx/dtype.cpp +// Modified from upstream MLX ce45c525 mlx/dtype.cpp // // Changes to the type_rules promotion table: // - bf16 + float32 → bf16 (was float32) diff --git a/src/lib/mlx-cpp/patches-cuda/fast.cpp b/src/lib/mlx-cpp/patches-cuda/fast.cpp index 7e050475b..a3dd80aa7 100644 --- a/src/lib/mlx-cpp/patches-cuda/fast.cpp +++ b/src/lib/mlx-cpp/patches-cuda/fast.cpp @@ -2,7 +2,7 @@ // // CUDA patch: bf16 scalar/constant generation in normalization ops // -// Modified from upstream MLX 81ba1c6a mlx/fast.cpp (three-way merged at each +// Modified from upstream MLX 59d600b5e mlx/fast.cpp (three-way merged at each // pin bump; the two fallbacks below are the whole delta) // // Changes: @@ -1060,6 +1060,133 @@ bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { has_sinks_ == a_other.has_sinks_; } +std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state, /* = std::nullopt */ + const std::optional& mask_, /* = std::nullopt */ + StreamOrDevice s_ /* = {} */) { + // determine output dtype + auto s = to_stream(s_); + + auto promoted = promote_types(queries.dtype(), keys.dtype()); + auto out_dtype = issubdtype(promoted, floating) + ? promoted + : promote_types(promoted, float32); + + // cast all inputs + auto q = astype(queries, out_dtype, s); + auto k = astype(keys, out_dtype, s); + auto v = astype(values, out_dtype, s); + auto g = astype(gates, out_dtype, s); + auto beta = astype(beta_, out_dtype, s); + + int B = q.shape(0); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); + + auto h0 = initial_state.has_value() ? astype(*initial_state, float32, s) + : zeros({B, Hv, Dv, Dk}, float32, s); + + bool has_mask = mask_.has_value(); + auto mask = has_mask ? astype(*mask_, bool_, s) : array(false); + + auto fallback = [B, T, Hk, Dk, Hv, Dv, has_mask, s]( + std::vector inputs) { + auto out_dtype = inputs[0].dtype(); + auto q = inputs[0]; + auto k = inputs[1]; + auto v = inputs[2]; + auto g = inputs[3]; + auto beta = inputs[4]; + auto state = astype(inputs[5], float32, s); + + if (Hv != Hk) { + int repeat_factor = Hv / Hk; + q = repeat(q, repeat_factor, 2, s); + k = repeat(k, repeat_factor, 2, s); + } + + array mask = has_mask ? astype(inputs[6], bool_, s) : array(false); + const array zero = array(0.0f, float32); + + std::vector outputs; + for (int t = 0; t < T; t++) { + auto get_t = [&](const array& a, int t) { + Shape start(a.ndim(), 0), stop = a.shape(); + start[1] = t; + stop[1] = t + 1; + return squeeze(slice(a, start, stop, s), 1, s); + }; + auto q_t = get_t(q, t); + auto k_t = get_t(k, t); + auto v_t = get_t(v, t); + auto g_t = get_t(g, t); + auto beta_t = get_t(beta, t); + array mask_t = has_mask ? get_t(mask, t) : array(true); + + auto state_prev = state; + + auto decay = (g_t.ndim() == 2) + ? expand_dims(g_t, {-1, -2}, s) // [B,H,1,1] + : expand_dims(g_t, -2, s); // [B,H,1,Dk] + + auto state_next = multiply(state_prev, decay, s); + + auto kv = + sum(multiply(state_next, expand_dims(k_t, -2, s), s), + -1, + false, + s); // [B,H,Dv] + + auto delta = subtract(v_t, kv, s); + delta = multiply(delta, expand_dims(beta_t, -1, s), s); + + state_next = + add(state_next, + multiply(expand_dims(delta, -1, s), expand_dims(k_t, -2, s), s), + s); + + if (has_mask) { + auto state_mask = expand_dims(mask_t, {-1, -2, -3}, s); + state = where(state_mask, state_next, state_prev, s); + } else { + state = state_next; + } + + auto o_t = sum(multiply(state, expand_dims(q_t, -2, s), s), -1, false, s); + + if (has_mask) { + auto out_mask = expand_dims(mask_t, {-1, -2}, s); + o_t = where(out_mask, o_t, zero, s); + } + outputs.push_back(astype(o_t, out_dtype, s)); + } + auto out = stack(outputs, 1, s); + return std::vector{out, state}; + }; + + if (!GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { + auto result = array::make_arrays( + /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, + /* dtypes */ {out_dtype, float32}, + /* primitive */ + std::make_shared(s, fallback), + /* inputs */ {q, k, v, g, beta, h0}); + + return result; + } + + auto result = fallback({q, k, v, g, beta, h0, mask}); + return result; +} + bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); return ( diff --git a/src/lib/mlx-cpp/patches-cuda/ops.cpp b/src/lib/mlx-cpp/patches-cuda/ops.cpp index 4c3b73c75..e74593c9a 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 81ba1c6a mlx/ops.cpp +// Modified from upstream MLX 59d600b5e mlx/ops.cpp // // Changes: // - Added bf16_mixed_astype() helper that skips astype conversion when the @@ -5610,9 +5610,9 @@ array gather_qmm( << "."; throw std::invalid_argument(msg.str()); } - if (to_stream(s).device != Device::gpu || !metal::is_available()) { + if (to_stream(s).device != Device::gpu) { throw std::invalid_argument( - "[gather_qmm] Global scale is only supported on the Metal backend."); + "[gather_qmm] Global scale is only supported on the GPU."); } } if (qmode == QuantizationMode::Affine) { @@ -5715,7 +5715,13 @@ array gather_qqmm( return array( std::move(out_shape), x.dtype(), - std::make_shared(stream, group_size, bits, qmode), + std::make_shared( + stream, + group_size, + bits, + qmode, + sorted_indices && !rhs_indices_, + sorted_indices && !lhs_indices_), std::move(inputs)); } diff --git a/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff b/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff index e6e0c38db..f8478c216 100644 --- a/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff +++ b/src/lib/mlx-cpp/patches-rocm/CORE_RESIDUAL.diff @@ -15,7 +15,7 @@ # 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: 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a +# pin: 59d600b5e64c238427d0f8d897ab7c682ef4d3d2 # fork: 75915908dfe5028335d318b10340313744fd3a8d # fork_upstream_merge_base: 39886de461d4f014cd60a3a68ad42fd8d776754a # overlay-sha256: ad95c895c1bf4d10b3fdece59d617b35c8f6135a135edb5a08f91a8a0a9afb30 CMakeLists.txt @@ -26,11 +26,11 @@ # overlay-sha256: ccf2ae61f566cde4b740b43d62a37fd6c904838b68f8330666073842f290b0f6 mlx/backend/gpu/primitives.cpp # overlay-sha256: 744670de21dbd5c74e2cee3afe1d542c2b73c45deee7611765fc9b0dcc99c334 mlx/compile.cpp # overlay-sha256: 67e68d6e70df6271ed2f54d75957c4762653b6a32fa6c666ea3b014c5d425acd mlx/device.cpp -# overlay-sha256: fb3a1dfce38319f1113732bdd40e070be0886ec48f1a53b26b7377065604deed mlx/fast.cpp -# overlay-sha256: 03638686c779e32639b6887b2ce625e9bda4966f86e6ec757a47045f984361d9 mlx/fast.h -# overlay-sha256: dad142ea136ff6b4c1fd14c5578526f196d0a1724f41495bbed5798ef6c2ff50 mlx/fast_primitives.h +# overlay-sha256: b45c24910329b5b40085b6c0607bc2c8e514a4fc3bafcce4e7422701a7d0d78c mlx/fast.cpp +# overlay-sha256: b1a5366b31412246256ea56f9114bfe7df2fcb93329301ce7ec924f7200c7695 mlx/fast.h +# overlay-sha256: 46684f0d184dc8e1d19c04839ad315e2a3ad38f58e8bd899d6a0b23031391953 mlx/fast_primitives.h # overlay-sha256: f5fccd165887c2860c04fecf3fcf6073397789081183c2bcaaacad8f8f4d8e9e mlx/io/safetensors.cpp -# overlay-sha256: 029ded214e753db3ead896df34d528200abcf21bd7a6e65d96b1e0f0f5e1c041 mlx/ops.cpp +# overlay-sha256: 92c85296ae688ef1db962a3c1b2132f0b94493d71ee2de350c259c4ce2493581 mlx/ops.cpp # overlay-sha256: 221cf283f6e15a1ddc4d11451eb79e930f2567828cf5fff45cb9ff020c6431e5 mlx/primitives.cpp # overlay-sha256: adb49482ec61de1b26887b851ec499b351387d0d030fb240664ded3d3d04e5cd mlx/stream.cpp # 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. @@ -131,7 +131,7 @@ compute_dynamic_offset(start_indices, out.strides(), axes_, s); --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h -@@ -447,11 +447,8 @@ +@@ -479,11 +479,8 @@ std::vector scalar_arguments, bool is_precompiled, int shared_memory, @@ -144,7 +144,7 @@ : Primitive(stream), name_(std::move(name)), source_(std::move(source)), -@@ -463,11 +460,8 @@ +@@ -495,11 +492,8 @@ scalar_arguments_(std::move(scalar_arguments)), is_precompiled_(is_precompiled), shared_memory_(shared_memory), @@ -157,7 +157,7 @@ void eval_cpu(const std::vector& inputs, std::vector& outputs) override { -@@ -504,12 +498,9 @@ +@@ -536,12 +530,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 2f2c9d04a..05910f0f1 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 (81ba1c6a) +## Retarget onto the MLX pin (59d600b5) The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code onto the pin needed: @@ -13,7 +13,7 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code 3. **Compiled kernels.** `compiled_collapse_contiguous_dims` returns a 4-tuple with `negative_strides`; `mlx/backend/rocm/compiled.cpp` unpacks it and forces the large-index kernel for negative strides, as the CUDA backend does. 4. **Custom kernels.** `mlx/backend/rocm/custom_kernel.cpp` passes an empty `CompileOptions::Data` before the aliases argument. `output_input_aliases_` stays out of `CustomKernel::state()` because `mlx/export.cpp` serializes that tuple and has no encoding for it (the fork's `state()` did not include it either). 5. **SDPA.** `ScaledDotProductAttention::use_fallback` takes upstream's new `force_fused` argument, with the CUDA semantics: throw if fused is forced and no fused kernel applies. -6. **New upstream primitives.** `GatherQQMM`, `SearchSorted` and `fast::CrossEntropy` (with its VJP, using the graph fallback) get `NO_GPU` stubs in `mlx/backend/rocm/primitives.cpp`. +6. **New upstream primitives.** `GatherQQMM`, `SearchSorted`, `fast::CrossEntropy` (with its VJP, using the graph fallback) and `fast::GatedDeltaUpdate` (using the graph fallback, as CUDA does in ml-explore/mlx#4020) get `NO_GPU` stubs in `mlx/backend/rocm/primitives.cpp`. 7. **Event errors.** `Event::error()` (ml-explore/mlx#3742) is populated on ROCm since lablup/mlxcel#1804, on the Metal model: each `CommandEncoder` owns an `Error`, and events signaled on a stream point at that stream's `Error` once the stream fails. Measured on gfx1151 with ROCm 7.15 (a probe kernel writing 2^40 bytes past its buffer): the HIP runtime does not abort the process on a queue fault, with or without `HIP_SKIP_ABORT_ON_GPU_ERROR` (the variable made no observable difference, so mlxcel does not set it); it prints the fault and the kernel name to stderr, then every HIP call on every thread returns `hipErrorIllegalAddress` for the rest of the process, `hipDeviceReset` included, and no queued `hipLaunchHostFunc` callback ever runs. That defeated all three of the backend's waits: `HipEvent::wait()` spun while `hipEventQuery` was not `hipSuccess`, `AtomicEvent::wait()` polled a counter that only a host callback lands, and `CommandEncoder::synchronize()` waited on a promise a host callback resolves. `HipEvent::wait()` now ends on any status other than `hipErrorNotReady`; `AtomicEvent` keeps the stream it was signaled from and polls `hipStreamQuery` on it about once a millisecond while waiting; `Worker::commit()` returns the `hipLaunchHostFunc` status and the encoder records a failure; `synchronize()` checks every status and polls the stream while it waits for the handlers. `Event::wait()` calls `check_error()` before and after the wait as Metal and CUDA do, `Event::is_signaled()` reports a failed event as signaled with the error attached (so `array::is_available()` detaches it and throws instead of pending forever), and `Event::wait(Stream)`/`Event::signal(Stream)` go through `scheduler::wait_event`/`signal_event` on CPU streams so the CPU-stream error path is upstream's. Synchronous launch failures: `launch_module_kernel` checks `hipModuleLaunchKernel`, and `gpu::eval` reads the thread's pending HIP error after each `eval_gpu` (kernels launched with `hipLaunchKernelGGL` or `<<<>>>` report a rejected launch no other way) and throws it as that primitive's error; a pending `hipErrorOutOfMemory` is skipped there, because it is what a recovered allocation failure leaves behind (the allocator's cache-release retry and managed fallback, hipBLASLt's workspace fallback) and an unrecovered one has already thrown. Under `MLX_GRAPH_PREFILL_REPLAY=1`, a launch that throws inside `launch_kernel`'s capture ends the capture and falls back to the eager launch, so a now-checked launch failure cannot leave the stream capturing. Destructor-path `hipGraphDestroy`/`hipGraphExecDestroy` calls discard their status explicitly, and `HipHandle::reset()` reports rather than throws, since on a faulted device every destroy fails and a throwing destructor would terminate the process at exit. `MLX_ROCM_GPU_WATCHDOG_SECS` (default off) bounds a single host wait for a kernel that never finishes; it is parsed with `strtol` and must be a whole non-negative integer, otherwise it warns once on stderr and stays off. Not covered: after a queue fault the device is dead for the process and its shutdown may hang, because HIP's teardown waits on the callbacks that never run (measured: the probe process hung at exit whenever a host callback was pending); the process has to be killed and restarted. Files under `mlx/backend/rocm/`: `device.cpp`, `device.h`, `eval.cpp`, `event.h`, `event.hip`, `fence.cpp`, `jit_module.h`, `utils.h`, `worker.cpp`, `worker.h`. Applies to the fork. ## Fixes to the fork's kernels diff --git a/src/lib/mlx-cpp/patches-rocm/UPSTREAM b/src/lib/mlx-cpp/patches-rocm/UPSTREAM index 020ada638..fabbf4485 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: 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a +retargeted_to_mlx_pin: 59d600b5e64c238427d0f8d897ab7c682ef4d3d2 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/backend/rocm/primitives.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp index cb090ca47..47a1422fc 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 @@ -47,6 +47,16 @@ NO_GPU(GatherQQMM) namespace fast { NO_GPU_USE_FALLBACK(CrossEntropy) NO_GPU_MULTI(CrossEntropyVJP) +bool GatedDeltaUpdate::use_fallback( + const int Hk, + const int Dk, + const int Hv, + const int Dv, + const bool has_mask, + Stream s) { + return true; +} +NO_GPU_MULTI(GatedDeltaUpdate) } // namespace fast // Note: The following are now implemented in their respective files: diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp index 568d41b2f..5d3a9d6d4 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/fast.cpp @@ -1023,6 +1023,133 @@ bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { has_sinks_ == a_other.has_sinks_; } +std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state, /* = std::nullopt */ + const std::optional& mask_, /* = std::nullopt */ + StreamOrDevice s_ /* = {} */) { + // determine output dtype + auto s = to_stream(s_); + + auto promoted = promote_types(queries.dtype(), keys.dtype()); + auto out_dtype = issubdtype(promoted, floating) + ? promoted + : promote_types(promoted, float32); + + // cast all inputs + auto q = astype(queries, out_dtype, s); + auto k = astype(keys, out_dtype, s); + auto v = astype(values, out_dtype, s); + auto g = astype(gates, out_dtype, s); + auto beta = astype(beta_, out_dtype, s); + + int B = q.shape(0); + int T = q.shape(1); + int Hk = q.shape(2); + int Dk = q.shape(3); + int Hv = v.shape(2); + int Dv = v.shape(3); + + auto h0 = initial_state.has_value() ? astype(*initial_state, float32, s) + : zeros({B, Hv, Dv, Dk}, float32, s); + + bool has_mask = mask_.has_value(); + auto mask = has_mask ? astype(*mask_, bool_, s) : array(false); + + auto fallback = [B, T, Hk, Dk, Hv, Dv, has_mask, s]( + std::vector inputs) { + auto out_dtype = inputs[0].dtype(); + auto q = inputs[0]; + auto k = inputs[1]; + auto v = inputs[2]; + auto g = inputs[3]; + auto beta = inputs[4]; + auto state = astype(inputs[5], float32, s); + + if (Hv != Hk) { + int repeat_factor = Hv / Hk; + q = repeat(q, repeat_factor, 2, s); + k = repeat(k, repeat_factor, 2, s); + } + + array mask = has_mask ? astype(inputs[6], bool_, s) : array(false); + const array zero = array(0.0f, float32); + + std::vector outputs; + for (int t = 0; t < T; t++) { + auto get_t = [&](const array& a, int t) { + Shape start(a.ndim(), 0), stop = a.shape(); + start[1] = t; + stop[1] = t + 1; + return squeeze(slice(a, start, stop, s), 1, s); + }; + auto q_t = get_t(q, t); + auto k_t = get_t(k, t); + auto v_t = get_t(v, t); + auto g_t = get_t(g, t); + auto beta_t = get_t(beta, t); + array mask_t = has_mask ? get_t(mask, t) : array(true); + + auto state_prev = state; + + auto decay = (g_t.ndim() == 2) + ? expand_dims(g_t, {-1, -2}, s) // [B,H,1,1] + : expand_dims(g_t, -2, s); // [B,H,1,Dk] + + auto state_next = multiply(state_prev, decay, s); + + auto kv = + sum(multiply(state_next, expand_dims(k_t, -2, s), s), + -1, + false, + s); // [B,H,Dv] + + auto delta = subtract(v_t, kv, s); + delta = multiply(delta, expand_dims(beta_t, -1, s), s); + + state_next = + add(state_next, + multiply(expand_dims(delta, -1, s), expand_dims(k_t, -2, s), s), + s); + + if (has_mask) { + auto state_mask = expand_dims(mask_t, {-1, -2, -3}, s); + state = where(state_mask, state_next, state_prev, s); + } else { + state = state_next; + } + + auto o_t = sum(multiply(state, expand_dims(q_t, -2, s), s), -1, false, s); + + if (has_mask) { + auto out_mask = expand_dims(mask_t, {-1, -2}, s); + o_t = where(out_mask, o_t, zero, s); + } + outputs.push_back(astype(o_t, out_dtype, s)); + } + auto out = stack(outputs, 1, s); + return std::vector{out, state}; + }; + + if (!GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { + auto result = array::make_arrays( + /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, + /* dtypes */ {out_dtype, float32}, + /* primitive */ + std::make_shared(s, fallback), + /* inputs */ {q, k, v, g, beta, h0}); + + return result; + } + + auto result = fallback({q, k, v, g, beta, h0, mask}); + return result; +} + bool Quantize::is_equivalent(const Primitive& other) const { const Quantize& p_other = static_cast(other); return ( diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/fast.h b/src/lib/mlx-cpp/patches-rocm/mlx/fast.h index 62ae96687..a28320c73 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/fast.h +++ b/src/lib/mlx-cpp/patches-rocm/mlx/fast.h @@ -60,6 +60,16 @@ MLX_API array scaled_dot_product_attention( bool force_fused = false, StreamOrDevice s = {}); +MLX_API std::vector gated_delta_update( + const array& queries, + const array& keys, + const array& values, + const array& gates, + const array& beta_, + const std::optional& initial_state = std::nullopt, + const std::optional& mask = std::nullopt, + StreamOrDevice s = {}); + using TemplateArg = std::variant; using ScalarArg = std::variant; 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 aecc52ce7..14447db25 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/fast_primitives.h +++ b/src/lib/mlx-cpp/patches-rocm/mlx/fast_primitives.h @@ -395,6 +395,38 @@ class ConvertFP8 : public Primitive { bool to_fp8_; }; +class GatedDeltaUpdate : public Custom { + public: + GatedDeltaUpdate( + 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, + const bool has_mask, + 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(GatedDeltaUpdate); + DEFINE_INPUT_OUTPUT_SHAPE() + auto state() const { + return std::make_tuple(nullptr); /* TODO */ + } + + private: +}; + class Quantize : public Custom { public: explicit Quantize( diff --git a/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp b/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp index b00902267..06ef018de 100644 --- a/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp +++ b/src/lib/mlx-cpp/patches-rocm/mlx/ops.cpp @@ -5648,9 +5648,9 @@ array gather_qmm( << "."; throw std::invalid_argument(msg.str()); } - if (to_stream(s).device != Device::gpu || !metal::is_available()) { + if (to_stream(s).device != Device::gpu) { throw std::invalid_argument( - "[gather_qmm] Global scale is only supported on the Metal backend."); + "[gather_qmm] Global scale is only supported on the GPU."); } } if (qmode == QuantizationMode::Affine) { @@ -5753,7 +5753,13 @@ array gather_qqmm( return array( std::move(out_shape), x.dtype(), - std::make_shared(stream, group_size, bits, qmode), + std::make_shared( + stream, + group_size, + bits, + qmode, + sorted_indices && !rhs_indices_, + sorted_indices && !lhs_indices_), std::move(inputs)); } 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 be566596b..56bcda060 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cpu/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cpu/quantized.cpp @@ -1,5 +1,6 @@ // 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, // `_qmm`, `_qmm_t`, `fp_qmm` and `fp_qmm_t`, each marked `mlxcel:` below; the // SIMD variants and everything else are upstream verbatim. diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/binary_ops.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/binary_ops.cuh index e50e47d9f..c8427bfba 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/binary_ops.cuh +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/binary_ops.cuh @@ -5,6 +5,8 @@ // Without this patch, Gemma 4 and Qwen 3.5 models crash on CUDA because // their computation graphs produce mixed bf16/float intermediates in the // compiled DAG. +// +// Modified from upstream MLX 81ba1c6a mlx/backend/cuda/device/binary_ops.cuh #include "mlx/backend/cuda/device/unary_ops.cuh" diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/gemm_sm70.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/gemm_sm70.cuh index 6ebbe1704..43651e990 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/gemm_sm70.cuh +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/gemm_sm70.cuh @@ -1,4 +1,9 @@ // Copyright © 2026 Apple Inc. +// +// Patched by mlxcel (lablup/mlxcel#1543): add MmaPath::Sm70 Volta 8x8x4 +// quad-pair tensor-core arm to make_tiled_mma for qmm_naive. +// +// Modified from upstream MLX 9a795735 mlx/backend/cuda/device/gemm_sm70.cuh #include diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_naive.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_naive.cuh index 5e2c9aad9..be9dcfe29 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_naive.cuh +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/device/qmm_naive.cuh @@ -1,4 +1,10 @@ // Copyright © 2026 Apple Inc. +// +// Patched by mlxcel (lablup/mlxcel#1543): parameterize qmm_naive_mainloop and +// qmm_naive_kernel on MmaPath instead of bool SM80 so Volta sm_70 can dispatch +// the Sm70 tensor-core MMA arm from gemm_sm70.cuh. +// +// Modified from upstream MLX 59d600b5e mlx/backend/cuda/device/qmm_naive.cuh #include "mlx/backend/cuda/device/cute_dequant.cuh" #include "mlx/backend/cuda/device/gemm_sm70.cuh" @@ -269,6 +275,10 @@ void qmm_naive_kernel( // 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; + // The global scale is per gathered expert; a non-gathered call has one. + if (global_scale && rhs_indices) { + global_scale += b_batch; + } // Get batch slice. Tensor mA = mA_mkl(_,_,a_batch); // (M,K) diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm.h b/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm.h index cf67d3832..45bf78438 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm.h +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm.h @@ -1,5 +1,6 @@ // Copyright © 2025 Apple Inc. // Modified by mlxcel: Added cutlass_gather_mm for general GatherMM support. +// Modified from upstream MLX b98831ad mlx/backend/cuda/gemms/grouped_gemm.h #pragma once diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h b/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h index 93b61cce8..45c21bcfa 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Modified from upstream MLX 9a795735 mlx/backend/cuda/gemms/grouped_gemm_arch.h (factored out of grouped_gemm_unaligned.cu in lablup/mlxcel#1544; not a separate upstream file) // Added by mlxcel (lablup/mlxcel#1544). Not an upstream MLX file. // // The CUTLASS architecture tag selection for the grouped GEMM in diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu index d5b4fac6b..23bd1d6b8 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu @@ -2,6 +2,7 @@ // Modified by mlxcel: Added cutlass_gather_mm for general GatherMM with both // lhs_indices and rhs_indices support. Synced to upstream c9aa5605, which // folds in #3469's Accumulator C-type fix for cutlass half-type matmul. +// Modified from upstream MLX c9aa5605 mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu #include "mlx/backend/cuda/cublas_utils.h" #include "mlx/backend/cuda/cutlass_utils.cuh" diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp index dc4c3013b..bb658eade 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp @@ -1,5 +1,6 @@ // Copyright © 2025 Apple Inc. -// Modified by mlxcel (synced to upstream e9463bb): +// Modified from upstream MLX 9a795735 mlx/backend/cuda/jit_module.cpp +// Modified by mlxcel (rebased onto 9a795735 in #1198): // - CCCL include resolution honors the MLXCEL_CCCL_DIR override and resolves // the executable directory from /proc/self/exe, so the bundled headers are // found regardless of how the process is launched (a relative ./mlxcel from diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.h b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.h index 8aa2dd4d7..56b3b9131 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.h +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.h @@ -3,6 +3,8 @@ // Modified by mlxcel (lablup/mlxcel#1566): serialize the per-kernel lazy // configuration in JitModule::get_kernel_and_dims. // +// Modified from upstream MLX 9a795735 mlx/backend/cuda/jit_module.h +// // Upstream stores each kernel as (CUfunction, configured_flag, max_block_dim) // and fills the last two on first use with no lock, while publishing the flag // *before* the dimension. A second thread that observes the flag already set diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/matmul.cpp b/src/lib/mlx-cpp/patches/mlx/backend/cuda/matmul.cpp index 3432f7a40..89dd4eedd 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/matmul.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/matmul.cpp @@ -1,5 +1,7 @@ // Copyright © 2025 Apple Inc. -// Modified by mlxcel: Added general GatherMM case using CUTLASS grouped GEMM. +// Modified by mlxcel: Added general GatherMM case using CUTLASS grouped GEMM +// (copied at b98831ad and merged with upstream 4652b0086 at 9a795735 in #1198). +// Modified from upstream MLX 9a795735 mlx/backend/cuda/matmul.cpp #include "mlx/backend/common/matmul.h" #include "mlx/backend/cuda/device.h" diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive.cu index 324fbc66f..e5bfd0992 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive.cu @@ -1,4 +1,5 @@ // Copyright © 2026 Apple Inc. +// Modified from upstream MLX 81ba1c6a mlx/backend/cuda/quantized/qmm/qmm_naive.cu // Patched by mlxcel (lablup/mlxcel#1541): size the CTA tile against the // device's real per-block shared-memory budget instead of a shape rule // standing in for one, opt into a larger dynamic shared-memory maximum when diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive_tile.h b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive_tile.h index 70cefafde..bd11553e1 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive_tile.h +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmm_naive_tile.h @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Modified from upstream MLX 9a795735 mlx/backend/cuda/quantized/qmm/qmm_naive_tile.h (factored out of qmm_naive.cu in lablup/mlxcel#1541; not a separate upstream file) // Added by mlxcel (lablup/mlxcel#1541). Not an upstream MLX file. // // The CTA tile selection for `qmm_naive`, factored out of `qmm_naive.cu` as a 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 c3b0a0d6f..3b3a6ac48 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,7 @@ // Copyright © 2026 Apple Inc. -// Patched by mlxcel. Modified from upstream MLX b7c3dd6d -// mlx/backend/cuda/quantized/qmm/qmm_sm80.cu. Changes: +// Patched by mlxcel. +// Modified from upstream MLX b7c3dd6d 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} // env overrides (see make_cta_tiler below); diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu index 51cdb57dd..6452595ed 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cu @@ -1,4 +1,5 @@ // Copyright © 2026 Apple Inc. +// Modified from upstream MLX 59d600b5e mlx/backend/cuda/quantized/qmm/qmv.cu // Patched by mlxcel (lablup/mlxcel#725): add a weight-amortizing multirow qmv // path for small-M/batched decode. The stock kernel maps one input row per // grid.y/grid.z block, so M*B < 8 shapes (batched decode with B in [2,8), @@ -740,6 +741,7 @@ __global__ void gather_qmv_kernel( const Q* w, const S* scales, const T* biases, + const float* global_scale, T* out, const uint32_t* lhs_indices, const uint32_t* rhs_indices, @@ -762,8 +764,12 @@ __global__ void gather_qmv_kernel( x += block.group_index().y * k + m * k * x_idx; out += block.group_index().y * n + m * n * l; + // The global scale is per gathered expert. + if (global_scale) { + global_scale += w_idx; + } qmv_kernel_impl( - x, w, scales, biases, nullptr, out, row, w_idx, n, k); + x, w, scales, biases, global_scale, out, row, w_idx, n, k); } // [mlxcel #725] Kill switch: MLXCEL_QMV_MULTIROW=0 restores the stock per-row @@ -1077,6 +1083,7 @@ void gather_qmv( const Q* w, const S* scales, const T* biases, + const float* global_scale, T* out, const uint32_t* lhs_indices, const uint32_t* rhs_indices, @@ -1093,7 +1100,16 @@ void gather_qmv( uint32_t(cuda::ceil_div(n, rows_per_block)), uint32_t(m), uint32_t(l)}; dim3 block_dims{WARP_SIZE, rows_per_block}; void* args[] = { - &x, &w, &scales, &biases, &out, &lhs_indices, &rhs_indices, &n, &k}; + &x, + &w, + &scales, + &biases, + &global_scale, + &out, + &lhs_indices, + &rhs_indices, + &n, + &k}; dispatch_bool(k % (WARP_SIZE * elems_per_thread), [&](auto has_residue_k) { auto* kernel = &gather_qmv_kernel< @@ -1260,6 +1276,7 @@ void gather_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -1286,6 +1303,9 @@ void gather_qmv( if (biases) { encoder.set_input_array(*biases); } + if (global_scale) { + encoder.set_input_array(*global_scale); + } encoder.set_input_array(lhs_indices); encoder.set_input_array(rhs_indices); encoder.set_output_array(out); @@ -1295,6 +1315,7 @@ void gather_qmv( gpu_ptr(w), gpu_ptr(scales), biases ? gpu_ptr(*biases) : nullptr, + global_scale ? gpu_ptr(*global_scale) : nullptr, gpu_ptr(out), gpu_ptr(lhs_indices), gpu_ptr(rhs_indices), 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 9b3977016..ff06174cf 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 @@ -1,4 +1,5 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. +// // Patched by mlxcel: (1) ensure input contiguity in QuantizedMatmul for // non-contiguous 3D batched weights (e.g. GLM-4 MLA embed_q with // transpose=false); (2) split long-prompt quantized matmuls so no CUDA launch @@ -10,17 +11,14 @@ // qmm_sm80, the naive kernel and qmv is the most architecture-dependent // dispatch in decode and MLX exposes no hook for it, so the one-shot trace has // to be emitted from the dispatch site itself. -// Synced to upstream 81ba1c6a. Between 9a795735 and 81ba1c6a upstream touched -// this file once, ml-explore/mlx#4458, which gives GatherQMM::eval_gpu a -// guard that refuses a global scale outside Metal and keys `biases` on -// `mode_ == Affine` instead of `inputs.size() == 6`. Both land ahead of delta -// (3) and merged without conflict, and the grouped-GEMM gate below already -// reads `biases` only on the affine path. qmm/qmm.h is byte-identical across -// that range, which means every dispatch entry point consumed here -// (qmm_naive, qmm_sm90, qmm_sm80, qmv, gather_qmv, fp_qmv) kept its -// signature, including the `const std::optional& global_scale` -// parameter ml-explore/mlx#3757 added to qmm_naive at position 5 that both -// call sites below pass std::nullopt for. +// +// Modified from upstream MLX 59d600b5e 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 +// nvfp4 and gating qmm_sm80 off when `global_scale` is present. The grouped-GEMM +// fast path already declines nvfp4 above, and `call_qmm_naive` / `call_qmv` +// forward `global_scale` into each batch chunk. // // Re-verify these call sites against qmm.h on the next pin bump; an overlay // that silently reverts an upstream change is how #830 / #831 happened. @@ -328,17 +326,16 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); - if (mode_ != QuantizationMode::Affine && inputs.size() == 6) { - throw std::runtime_error( - "[GatherQMM] Global scale is only supported on the Metal backend."); - } - array x = ensure_row_contiguous(inputs[0], encoder, s); const array& w = inputs[1]; const array& scales = inputs[2]; + // Affine gets biases at index 3, nvfp4 an optional global scale. std::optional biases; + std::optional global_scale; if (mode_ == QuantizationMode::Affine) { biases = inputs[3]; + } else if (inputs.size() == 6) { + global_scale = ensure_row_contiguous(inputs[3], encoder, s); } array lhs_indices = ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s); @@ -458,7 +455,10 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { mode_, encoder.device()); }; - bool can_use_qmm_sm80 = supports(supports_qmm_sm80); + // qmm_sm80 does not apply global scales yet; route such calls to the + // naive kernel until it does. + bool can_use_qmm_sm80 = + !global_scale.has_value() && supports(supports_qmm_sm80); bool can_use_qmm_naive = supports(supports_qmm_naive); bool can_use_qmv = supports(supports_qmv); @@ -491,7 +491,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, - std::nullopt, + global_scale, lc, rc, oc, @@ -512,6 +512,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + global_scale, lc, rc, oc, diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/all_reduce.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/all_reduce.cu index c17f491a6..a96816214 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/all_reduce.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/all_reduce.cu @@ -2,8 +2,7 @@ // // CUDA patch: bf16 reduce accumulation optimization // -// Modified from upstream MLX v0.31.1 -// mlx/backend/cuda/reduce/all_reduce.cu +// Modified from upstream MLX 59d600b5e mlx/backend/cuda/reduce/all_reduce.cu // // Changes: // - Added output type template parameter V (separate from accumulation type U) @@ -80,12 +79,12 @@ void all_reduce( out.set_data(cu::malloc_async(out.nbytes(), encoder)); - auto get_args = [](int size, int N) { - int threads = std::min(512, (size + N - 1) / N); + auto get_args = [](size_t size, size_t N) { + int threads = + static_cast(std::min(size_t{512}, cuda::ceil_div(size, N))); threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; - int reductions_per_step = threads * N; - size_t steps_needed = - (size + reductions_per_step - 1) / reductions_per_step; + size_t reductions_per_step = threads * N; + size_t steps_needed = cuda::ceil_div(size, reductions_per_step); int blocks; if (steps_needed < 32) { @@ -100,7 +99,7 @@ void all_reduce( blocks = 1024; } - size_t steps_per_block = (steps_needed + blocks - 1) / blocks; + size_t steps_per_block = cuda::ceil_div(steps_needed, blocks); size_t block_step = steps_per_block * reductions_per_step; return std::make_tuple(blocks, threads, block_step); @@ -109,9 +108,7 @@ void all_reduce( int blocks, threads; size_t block_step; size_t insize = in.size(); - - // Cub doesn't like const pointers for load (sigh). - void* indata = const_cast(gpu_ptr(in)); + void* indata = gpu_ptr(in); // Large array so allocate an intermediate and accumulate there std::tie(blocks, threads, block_step) = get_args(insize, N_READS); @@ -144,7 +141,7 @@ void all_reduce( kernel_pass1, blocks, threads, - static_cast(indata), + indata, gpu_ptr(intermediate), block_step, insize); @@ -178,7 +175,7 @@ void all_reduce( kernel, blocks, threads, - static_cast(indata), + indata, gpu_ptr(out), block_step, insize); 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 95093d89f..ecca3d26e 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,8 +2,7 @@ // // CUDA patch: bf16 reduce accumulation optimization // -// Modified from upstream MLX v0.31.1 -// mlx/backend/cuda/reduce/col_reduce.cu +// Modified from upstream MLX 59d600b5e mlx/backend/cuda/reduce/col_reduce.cu // // Changes: // - Added output type template parameter V (separate from accumulation type U) @@ -305,8 +304,6 @@ void col_reduce_looped( using OP = MLX_GET_TYPE(reduce_type_tag); using T = cuda_type_t; using U = typename cu::ReduceResult::type; - // Cub doesn't like const pointers for vectorized loads. (sigh) - T* indata = const_cast(gpu_ptr(in)); constexpr int N_READS = 4; constexpr int BM = 32; @@ -320,7 +317,7 @@ void col_reduce_looped( kernel, grid, blocks, - indata, + gpu_ptr(in), gpu_ptr(out), static_cast(args), out.size() / args.reduction_stride); @@ -431,9 +428,6 @@ void col_reduce_two_pass( // --- First pass: T input -> U intermediate --- { - // Cub doesn't like const pointers for vectorized loads. (sigh) - T* indata = const_cast(gpu_ptr(in)); - constexpr int N_READS = 4; constexpr int BM = 32; constexpr int BN = 32; @@ -456,7 +450,7 @@ void col_reduce_two_pass( kernel, grid, blocks, - indata, + gpu_ptr(in), gpu_ptr(intermediate), static_cast(args), out.size() / args.reduction_stride); diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/init_reduce.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/init_reduce.cu index 5f10d5f77..4095a5d6a 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/init_reduce.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/init_reduce.cu @@ -2,8 +2,7 @@ // // CUDA patch: bf16 reduce accumulation optimization // -// Modified from upstream MLX v0.31.1 -// mlx/backend/cuda/reduce/init_reduce.cu +// Modified from upstream MLX 9a795735 mlx/backend/cuda/reduce/init_reduce.cu // // Changes: // - Uses T (input type) for output pointer instead of U (accumulation type) diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/reduce_ops.cuh b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/reduce_ops.cuh index 8d64d44a1..86d97ab6a 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/reduce_ops.cuh +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/reduce_ops.cuh @@ -2,8 +2,7 @@ // // CUDA patch: bf16 reduce accumulation optimization // -// Modified from upstream MLX v0.31.1 -// mlx/backend/cuda/reduce/reduce_ops.cuh +// Modified from upstream MLX ce45c525 mlx/backend/cuda/reduce/reduce_ops.cuh // // Changes: // - Added ReduceResult specialization: accumulate in fp32 diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/row_reduce.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/row_reduce.cu index 1d1e823b3..39c9180f1 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/row_reduce.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/reduce/row_reduce.cu @@ -2,8 +2,7 @@ // // CUDA patch: bf16 reduce accumulation optimization // -// Modified from upstream MLX v0.31.1 -// mlx/backend/cuda/reduce/row_reduce.cu +// Modified from upstream MLX 59d600b5e mlx/backend/cuda/reduce/row_reduce.cu // // Changes: // - Added output type template parameter V (separate from accumulation type U) @@ -296,10 +295,15 @@ void row_reduce_simple( kernel = cu::row_reduce_simple; } - T* indata = const_cast(gpu_ptr(in)); int size = plan.shape.back(); encoder.add_kernel_node( - kernel, grid, block, indata, gpu_ptr(out), out.size(), size); + kernel, + grid, + block, + gpu_ptr(in), + gpu_ptr(out), + out.size(), + size); }); }); } 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 1e8603b98..8f852fe38 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 @@ -8,6 +8,8 @@ // 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 +// // Three mlxcel changes sit on top of upstream 81ba1c6a: // // 1. #1820, the general fix: a small multi-row call carrying an array mask over diff --git a/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu b/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu index a82f50755..39efec20e 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu +++ b/src/lib/mlx-cpp/patches/mlx/backend/cuda/scaled_dot_product_attention.cu @@ -14,6 +14,8 @@ // launches (a no-op for D <= 128 which already fit). The .cpp is unchanged: // cudnn rejects head_dim > 128, so widening this gate alone routes 256/288 // decode to the vector kernels. Synced to upstream e9463bb. +// +// Modified from upstream MLX e9463bb mlx/backend/cuda/scaled_dot_product_attention.cu // Required for using M_LOG2E in MSVC. #define _USE_MATH_DEFINES diff --git a/src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp b/src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp index afe76196a..28cd47732 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. // mlxcel patch: adds static_cast for mixed-dtype compiled ops. +// Modified from upstream MLX 81ba1c6a mlx/backend/metal/compiled.cpp // Synced to upstream 81ba1c6a. The only delta is that cast, in build_kernel's // emission of an op's inputs; everything else is upstream verbatim. Until this // sync the file also emitted `elem_to_loc_1` for 1-D inputs, a leftover 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 824aa4ce8..efef17f6a 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/device.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. // Patched by mlxcel: a runtime override for the per-command-buffer input -// budget. Synced to upstream 81ba1c6a. The delta is three hunks: the +// budget. Synced to upstream 81ba1c6a. +// Modified from upstream MLX 81ba1c6a 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 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 5244ade42..9d7100352 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 @@ -1,4 +1,9 @@ // Copyright © 2023-2024 Apple Inc. +// +// 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 #pragma once 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 b927cf158..8b16e3563 100644 --- a/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp +++ b/src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2026 Apple Inc. // Patched by mlxcel: `use_qmv_wide` gains an off-switch, `MLXCEL_QMV_WIDE=0`. -// Synced to upstream 81ba1c6a. The delta is three hunks: the includes below, +// 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 @@ -864,6 +865,7 @@ void qmm_nax( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, bool transpose, int group_size, @@ -910,7 +912,8 @@ void qmm_nax( "_wn", wn, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); std::string template_def; MTL::ComputePipelineState* kernel; if (transpose) { @@ -928,7 +931,8 @@ void qmm_nax( bk, bn, wm, - wn); + wn, + global_scale.has_value()); } else { kernel = get_qmm_nax_kernel_wrapped( d, @@ -953,6 +957,11 @@ void qmm_nax( compute_encoder.set_input_array(scales, c++); if (biases) { compute_encoder.set_input_array(*biases, c++); + } else if (transpose) { + if (global_scale) { + compute_encoder.set_input_array(*global_scale, c); + } + c++; } compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); @@ -1080,6 +1089,7 @@ void qmm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, bool transpose, int group_size, @@ -1090,6 +1100,10 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { + assert( + (!global_scale || (transpose && !biases)) && + "qmm global scale requires transposed weights without biases"); + bool has_nax_kernel = metal::is_nax_available() && (transpose || mode == "affine"); bool nax_aligned = (K % 64 == 0) && (transpose || N % 64 == 0); @@ -1100,6 +1114,7 @@ void qmm( /* const array& w = */ w, /* const array& scales = */ scales, /* const std::optional& biases = */ biases, + /* const std::optional& global_scale = */ global_scale, /* array& out = */ out, /* bool transpose = */ transpose, /* int group_size = */ group_size, @@ -1135,7 +1150,8 @@ void qmm( "_b_", bits, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); std::string template_def; MTL::ComputePipelineState* kernel; if (transpose) { @@ -1148,7 +1164,8 @@ void qmm( group_size, bits, aligned, - batched); + batched, + global_scale.has_value()); } else { kernel = get_quantized_kernel_wrapped( d, kname, "qmm_n", mode, type_string, group_size, bits, batched); @@ -1161,6 +1178,11 @@ void qmm( compute_encoder.set_input_array(scales, c++); if (biases) { compute_encoder.set_input_array(*biases, c++); + } else if (transpose) { + if (global_scale) { + compute_encoder.set_input_array(*global_scale, c); + } + c++; } compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); @@ -1206,7 +1228,21 @@ void qmm_splitk( } if (split_k <= 1) { return qmm( - x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode); + x, + w, + scales, + biases, + std::nullopt, + out, + true, + group_size, + bits, + M, + N, + K, + d, + s, + mode); } int k_partition_size = K / split_k; @@ -1910,6 +1946,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, out, transpose_, group_size_, @@ -2126,6 +2163,26 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { int K = x.shape(-1); int M = non_batched ? x.size() / K : x.shape(-2); int N = out.shape(-1); + if (has_global_scales && w_quantized && non_batched && + x.dtype() == bfloat16 && K % 32 == 0 && + M >= get_qmv_batch_limit(K, N, d)) { + qmm(x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + true, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } dispatch_qmv( x, w_q, @@ -2196,6 +2253,57 @@ void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { global_scale_w = ensure_row_contiguous(gs_e, d, s); } + int B = out.size() / M / N; + bool use_matrix_kernels = has_global_scales && w_quantized && + x.dtype() == bfloat16 && w_q.ndim() == 3 && + K % (metal::is_nax_available() ? 64 : 32) == 0; + + if (use_matrix_kernels && M == 1 && B >= 16 && right_sorted_) { + int E = w_q.size() / w_q.shape(-1) / w_q.shape(-2); + if (B / E >= 4) { + gather_qmm_rhs( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + rhs_indices, + out, + true, + group_size_, + bits_, + x.size() / K, + N, + K, + d, + s, + mode); + return; + } + } + + if (use_matrix_kernels && M >= get_qmv_batch_limit(K, N, d)) { + gather_qmm( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + true, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); + return; + } + gather_qmv( x, w_q, From a8ae53fb731121a03172ad7456b9b41d6524830f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 14 Oct 2026 01:14:47 +0900 Subject: [PATCH 2/2] docs(reports): add bilingual technical report for PR #2298 Archive the English and Korean technical reports documenting the MLX pin bump to 59d600b5e64c238427d0f8d897ab7c682ef4d3d2, overlay base commit declarations across all 32 patches, 3-way rebasing of drifted overlays, ROCm retargeting, and hardware verification on gfx1151. Refs #1932, #2298 --- ...-bump-mlx-pin-overlay-bases-20261014.en.md | 65 +++++++++++++++++++ ...-bump-mlx-pin-overlay-bases-20261014.ko.md | 65 +++++++++++++++++++ 2 files changed, 130 insertions(+) create mode 100644 TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.en.md create mode 100644 TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.ko.md diff --git a/TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.en.md b/TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.en.md new file mode 100644 index 000000000..bfa256efe --- /dev/null +++ b/TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.en.md @@ -0,0 +1,65 @@ +# Technical Report: PR #2298 - Bump MLX Pin and Declare Overlay Bases + +- **Date**: 2026-10-14 +- **Status**: Completed; implementation reviewed and all contract gates passed +- **Risk Level**: Low (all upstream changes 3-way merged, all 32 overlays verified at 0 drift, local ROCm overlay and unit test suites green) + +## Executive Summary + +PR #2298 updates the pinned upstream `ml-explore/mlx` commit from `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a` to `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` (14 upstream commits) in `src/lib/mlx-cpp/CMakeLists.txt`. It establishes a machine-readable `// Modified from upstream MLX ` base-declaration header across all 32 overlay files in `src/lib/mlx-cpp/patches/` and `src/lib/mlx-cpp/patches-cuda/`, resolving each base commit by diff against upstream history. All 9 conflicting or drifted overlays across that range were 3-way merged onto the new pin while preserving all mlxcel custom kernel and runtime deltas. Upstream fix `dfe17bafb` (#3912) was confirmed intact. The ROCm overlay was retargeted to the new pin with fallback stubs for the new upstream `fast::GatedDeltaUpdate` primitive, and `CONTRIBUTING.md` was extended with documentation for the CUDA/Metal/CPU overlay lifecycle and 3-way merge workflow. + +## 1. Problem Statement + +### 1.1 Background +The `mlxcel` repository applies whole-file source overlays onto upstream `ml-explore/mlx` via `configure_file(... COPYONLY)`. When moving the upstream `GIT_TAG` forward, any upstream modifications to an overlaid file are silently overwritten unless the overlay is explicitly rebased against upstream. Previously, only the ROCm overlay tracked its base commits and drift systematically, while CUDA, Metal, and CPU overlays lacked a machine-readable declaration of their copy or rebase commits. + +### 1.2 Identified Deficiencies +- **Missing Base Declarations**: 29 of 32 overlay files lacked a machine-readable `Modified from upstream MLX` header, preventing automated drift checking. +- **Outdated Reducers**: Three CUDA reduction overlays (`all_reduce.cu`, `col_reduce.cu`, `row_reduce.cu`) were pinned to `v0.31.1` (`ce45c525`), missing upstream const-correctness and 64-bit planning improvements. +- **Upstream Pin Advancement**: Upstream MLX advanced with new features, including global-scale support for quantized operations (#4507), Gated Delta Net kernels (#4020), and non-multiple-of-32 quantized dimension fixes (#3912). +- **ROCm Divergence**: Upstream addition of `fast::GatedDeltaUpdate` required ROCm fallback stubs to prevent undefined symbol errors during static linking. + +## 2. Implementation Details + +### 2.1 Phase 1: Machine-Readable Base Declarations +All 32 overlay files under `src/lib/mlx-cpp/patches/` and `src/lib/mlx-cpp/patches-cuda/` now carry a header comment of the form: +```cpp +// Modified from upstream MLX +``` +Each SHA was established by diffing against upstream `ml-explore/mlx` history. All 32 SHAs resolve in the upstream git repository. + +### 2.2 Phase 2: 3-Way Rebase of Drifted Overlays +Nine overlays touched between their base and `59d600b5e` were 3-way merged onto `59d600b5e`: +1. `patches/mlx/backend/metal/quantized.cpp`: Integrated upstream global-scale loading while preserving the three `MLXCEL_QMV_WIDE` runtime off-switch hunks (#1186, #1187). +2. `patches/mlx/backend/cuda/device/qmm_naive.cuh`: Preserved Volta sm_70 `MmaPath` parameterization (#1543). +3. `patches/mlx/backend/cuda/quantized/qmm/qmv.cu`: Added `global_scale` parameter threading while preserving multirow qmv (#725), row-window caps (#906), and pre-Ampere accumulator optimizations (#1539). +4. `patches/mlx/backend/cuda/quantized/quantized.cpp`: Preserved 65535 grid chunking (#648), input contiguity checks, and sorted MoE prefill grouped-GEMM routes (#629). +5. `patches-cuda/fast.cpp`: Added `gated_delta_update` fallback implementation. +6. `patches-cuda/ops.cpp`: Updated `gather_qmm` GPU device validation and `GatherQQMM` argument list. +7. `patches/mlx/backend/cuda/reduce/{all_reduce,col_reduce,row_reduce}.cu`: Incorporated upstream `const_cast` removal and large-input 64-bit all-reduce planning. + +The drift check confirms 0 upstream commits between declared bases and `59d600b5e` across all 32 overlays. Upstream `dfe17bafb` (#3912) is intact: neither `fp_quantized.h` nor `quantized.h` is overlaid, so both are consumed unmodified from the new pin. + +### 2.3 Phase 3: Pin Bump and ROCm Overlay Retargeting +- `GIT_TAG` in `src/lib/mlx-cpp/CMakeLists.txt` was updated to `59d600b5e64c238427d0f8d897ab7c682ef4d3d2`. +- `patches-rocm/UPSTREAM` `retargeted_to_mlx_pin` was updated to `59d600b5e`. +- `patches-rocm/mlx/backend/rocm/primitives.cpp` added `fast::GatedDeltaUpdate::use_fallback` returning `true` and `NO_GPU_MULTI(GatedDeltaUpdate)`, documented as item 6 in `LOCAL_FIXES.md`. +- `patches-rocm/CORE_RESIDUAL.diff` was refreshed against `59d600b5e`. + +### 2.4 Phase 4: Documentation +`CONTRIBUTING.md` was extended with the `Reconciling CUDA, Metal, and CPU source overlays` subsection, detailing: +- The `// Modified from upstream MLX ` header requirement. +- The automated shell loop for detecting overlay drift. +- The 3-way merge command sequence using `git merge-file`. +- Guidelines for retiring overlays absorbed upstream. + +## 3. Verification + +The following verification steps were executed and confirmed passing: +1. `bash scripts/ci/mlx_pinned_commit.sh` prints `59d600b5e64c238427d0f8d897ab7c682ef4d3d2`. +2. `cargo test -p mlxcel-mlx-pin` passes all 27 tests (26 unit tests + 1 cross-parser integration test). +3. Overlay drift check across all 32 files reports 0 `NO BASE` lines and 0 `DRIFT` lines. +4. `make verify-rocm-overlay` passes with 29 checks OK, validating 109 backend + 15 core files, 45 `LOCAL_FIXES` entries, and 4 noted residuals against pin `59d600b5`. +5. Contract gates passed: `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-python-tooling verify-fmt`. +6. `git diff --check origin/main..HEAD` confirms no whitespace errors or leftover conflict markers. +7. Greedy generation validation on AMD `gfx1151` produces byte-identical tokens before and after the bump. diff --git a/TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.ko.md b/TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.ko.md new file mode 100644 index 000000000..0c4b95fed --- /dev/null +++ b/TECHNICAL_REPORTS/2298-bump-mlx-pin-overlay-bases-20261014.ko.md @@ -0,0 +1,65 @@ +# 기술 보고서: PR #2298 - MLX 핀 업데이트 및 오버레이 베이스 커밋 선언 + +- **날짜**: 2026-10-14 +- **상태**: 완료; 구현 검토 및 전체 컨트랙트 게이트 통과 +- **위험 수준**: 낮음 (모든 업스트림 변경 3-way 병합 완료, 32개 오버레이 드리프트 0 검증, 로컬 ROCm 오버레이 및 단위 테스트 전원 통과) + +## 개요 + +PR #2298은 `src/lib/mlx-cpp/CMakeLists.txt`에서 고정된 업스트림 `ml-explore/mlx` 커밋을 `81ba1c6a0e50a9268b931579c2d4f1158b9aab5a`에서 `59d600b5e64c238427d0f8d897ab7c682ef4d3d2`로 업데이트합니다 (14개 업스트림 커밋). `src/lib/mlx-cpp/patches/` 및 `src/lib/mlx-cpp/patches-cuda/` 하위의 32개 전체 오버레이 파일에 기계 판독 가능한 `// Modified from upstream MLX ` 헤더를 선언하고, 각 베이스 커밋을 업스트림 히스토리와의 diff를 통해 정밀하게 식별하였습니다. 해당 구간에서 충돌하거나 드리프트가 발생한 9개 오버레이는 mlxcel의 커스텀 커널 및 런타임 변경점을 온전히 보존하면서 새 핀으로 3-way 병합되었습니다. 업스트림 수정사항인 `dfe17bafb` (#3912)가 온전히 유지됨을 확인하였습니다. ROCm 오버레이는 새 핀으로 리타깃되었고 신규 업스트림 `fast::GatedDeltaUpdate` 프리미티브에 대한 폴백 스텁이 추가되었으며, `CONTRIBUTING.md`에 CUDA/Metal/CPU 오버레이 수명 주기와 3-way 병합 절차가 문서화되었습니다. + +## 1. 문제 정의 + +### 1.1 배경 +`mlxcel` 저장소는 `configure_file(... COPYONLY)`를 사용하여 업스트림 `ml-explore/mlx`에 파일 단위 소스 오버레이를 적용합니다. 업스트림 `GIT_TAG`를 갱신할 때 오버레이를 명시적으로 리베이스하지 않으면 업스트림의 소스 변경사항이 조용히 덮어씌워져 유실될 위험이 있습니다. 기존에는 ROCm 오버레이만 체계적으로 베이스와 드리프트를 추적하고 있었으며, CUDA, Metal, CPU 오버레이에는 복사 또는 리베이스 시점의 커밋이 명시되지 않았습니다. + +### 1.2 문제점 +- **오버레이 베이스 선언 부재**: 32개 오버레이 중 29개 파일에 기계 판독 가능한 `Modified from upstream MLX` 헤더가 없어 자동 드리프트 감지가 불가능했습니다. +- **오래된 리듀서 오버레이**: 3개의 CUDA reduction 오버레이(`all_reduce.cu`, `col_reduce.cu`, `row_reduce.cu`)가 `v0.31.1` (`ce45c525`)에 머물러 있어 업스트림의 const-correctness 및 64비트 계획 개선사항이 누락되어 있었습니다. +- **업스트림 핀 갱신 필요**: 업스트림 MLX에 양자화 연산의 글로벌 스케일 지원(#4507), Gated Delta Net 커널(#4020), 32의 배수가 아닌 양자화 차원 지원(#3912) 등 주요 개선사항이 추가되었습니다. +- **ROCm 빌드 심볼 정합성**: 신규 추가된 `fast::GatedDeltaUpdate` 프리미티브에 대해 ROCm 백엔드 스텁이 없으면 정적 링킹 시 정의되지 않은 심볼 에러가 발생합니다. + +## 2. 구현 상세 + +### 2.1 1단계: 기계 판독 가능한 베이스 선언 +`src/lib/mlx-cpp/patches/` 및 `src/lib/mlx-cpp/patches-cuda/` 하위의 32개 전체 오버레이 파일에 다음 형태의 헤더 주석을 선언하였습니다: +```cpp +// Modified from upstream MLX +``` +각 SHA는 `ml-explore/mlx` 업스트림 히스토리와의 diff를 통해 확정되었으며, 32개 SHA 모두 업스트림 저장소에서 유효하게 확인됩니다. + +### 2.2 2단계: 드리프트 오버레이 3-way 리베이스 +기존 베이스와 `59d600b5e` 사이에서 변경된 9개 오버레이를 `59d600b5e`로 3-way 병합하였습니다: +1. `patches/mlx/backend/metal/quantized.cpp`: 업스트림 글로벌 스케일 로딩을 통합하고 3개의 `MLXCEL_QMV_WIDE` 런타임 오프스위치 헝크(#1186, #1187)를 유지. +2. `patches/mlx/backend/cuda/device/qmm_naive.cuh`: Volta sm_70 `MmaPath` 템플릿 매개변수화(#1543) 유지. +3. `patches/mlx/backend/cuda/quantized/qmm/qmv.cu`: `global_scale` 파라미터 전달을 추가하고 multirow qmv(#725), 행 윈도우 캡(#906), pre-Ampere 누적 최적화(#1539)를 온전히 유지. +4. `patches/mlx/backend/cuda/quantized/quantized.cpp`: 65535 그리드 청킹(#648), 입력 연속성 검증, 정렬된 MoE 프리필 grouped-GEMM 경로(#629) 유지. +5. `patches-cuda/fast.cpp`: `gated_delta_update` 폴백 구현 추가. +6. `patches-cuda/ops.cpp`: `gather_qmm` GPU 디바이스 유효성 검사 및 `GatherQQMM` 인수 목록 갱신. +7. `patches/mlx/backend/cuda/reduce/{all_reduce,col_reduce,row_reduce}.cu`: 업스트림 `const_cast` 제거 및 대규모 입력에 대한 64비트 all-reduce 계획 반영. + +드리프트 검증 결과 32개 오버레이 전원에 대해 베이스와 `59d600b5e` 간의 커밋 차이가 0임을 확인하였습니다. 업스트림 `dfe17bafb` (#3912)는 `fp_quantized.h` 및 `quantized.h`가 오버레이되지 않아 새 핀에서 온전히 반영됨을 확인했습니다. + +### 2.3 3단계: 핀 갱신 및 ROCm 오버레이 리타깃 +- `src/lib/mlx-cpp/CMakeLists.txt`의 `GIT_TAG`를 `59d600b5e64c238427d0f8d897ab7c682ef4d3d2`로 갱신. +- `patches-rocm/UPSTREAM`의 `retargeted_to_mlx_pin`을 `59d600b5e`로 갱신. +- `patches-rocm/mlx/backend/rocm/primitives.cpp`에 `fast::GatedDeltaUpdate::use_fallback` (`true` 반환) 및 `NO_GPU_MULTI(GatedDeltaUpdate)` 추가 (LOCAL_FIXES.md 항목 6에 기록). +- `patches-rocm/CORE_RESIDUAL.diff` 갱신. + +### 2.4 4단계: 문서화 +`CONTRIBUTING.md`에 `Reconciling CUDA, Metal, and CPU source overlays` 서브섹션을 추가하여 다음 내용을 기술하였습니다: +- `// Modified from upstream MLX ` 헤더 규약. +- 오버레이 드리프트 자동 탐지 셸 루프. +- `git merge-file`을 활용한 3-way 병합 절차. +- 업스트림에 반영된 오버레이의 제거 지침. + +## 3. 검증 결과 + +다음 검증 항목들이 모두 정상 통과되었습니다: +1. `bash scripts/ci/mlx_pinned_commit.sh` 실행 시 `59d600b5e64c238427d0f8d897ab7c682ef4d3d2` 출력. +2. `cargo test -p mlxcel-mlx-pin` 전체 27개 테스트 통과 (26개 단위 테스트 + 1개 크로스 파서 통합 테스트). +3. 32개 오버레이 전체 드리프트 검사 시 0 `NO BASE`, 0 `DRIFT` 출력. +4. `make verify-rocm-overlay` 실행 시 29개 점검 항목 정상 통과 (109개 백엔드 + 15개 코어 파일, 45개 LOCAL_FIXES 항목, 4개 잔차 일치 확인). +5. 컨트랙트 게이트 통과: `make verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-python-tooling verify-fmt`. +6. `git diff --check origin/main..HEAD` 공백 에러 및 충돌 마커 없음 확인. +7. AMD `gfx1151` 환경에서 탐욕적 텍스트 생성 테스트 결과 핀 업데이트 전후 토큰 단위 동일 출력 확인.