From c5e3a13f30c96adad2597998857ae68b4d1e2996 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 21 Sep 2026 20:21:04 +0900 Subject: [PATCH 1/2] perf(cohere2): fuse the parallel-block residual adds into one kernel Cohere2 blocks end with `(attn + ff) + x`, which MLX runs as two dependent elementwise kernels, each one more barrier level on the decode critical path. `compiled_add3` compiles the same expression, with the same association order, into one kernel, so the result is byte-identical and every layer drops one dispatch and one barrier. Validated on c4ai-command-r7b-12-2024 4-bit on M1 Ultra: greedy generations for three prompts at 200 tokens match main exactly, and eight interleaved ABBA pairs (500-token prompt, 128 tokens, command-buffer input budget pinned on both arms) read decode median 110.54 vs 109.35 tok/s, +1.1%. --- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp | 23 ++++++++++++++++++++++ src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 8 ++++++++ src/lib/mlxcel-core/src/lib.rs | 6 ++++++ src/models/cohere2.rs | 6 +++--- 4 files changed, 40 insertions(+), 3 deletions(-) diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp index da77371b1..248da0bba 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp @@ -1650,6 +1650,29 @@ std::unique_ptr compiled_swiglu_activation( return std::make_unique(std::move(result[0])); } +// Compiled three-way add: (a + b) + c as one fused elementwise kernel. +// Same association order as two chained `add` calls, so the result is +// byte-identical; the win is one dispatch and one barrier level fewer per call. +// Used by: Cohere2 +namespace { + static std::function(const std::vector&)> get_compiled_add3() { + auto fn = [](const std::vector& inputs) -> std::vector { + return {mlx::core::add(mlx::core::add(inputs[0], inputs[1]), inputs[2])}; + }; + return compile_shapeless_audited("compiled_add3", fn); + } +} + +std::unique_ptr compiled_add3( + const MlxArray& a, + const MlxArray& b, + const MlxArray& c +) { + static auto compiled_fn = get_compiled_add3(); + auto result = compiled_fn({a.inner, b.inner, c.inner}); + return std::make_unique(std::move(result[0])); +} + // Compiled GptOss SwiGLU activation using the exact mlx-lm formulation: // x_glu = clip(x_glu, max=7) // x_linear = clip(x_linear, min=-7, max=7) diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 79e410679..74647716e 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -538,6 +538,14 @@ std::unique_ptr compiled_swiglu_activation( const MlxArray& x ); +// Three-way add (a + b) + c compiled into one fused kernel (shapeless=true). +// Byte-identical to two chained adds. Used by: Cohere2 +std::unique_ptr compiled_add3( + const MlxArray& a, + const MlxArray& b, + const MlxArray& c +); + // GptOss SwiGLU activation only - compiled with kernel fusion (shapeless=true) // output = clipped_gate * sigmoid(1.702 * clipped_gate) * (clipped_up + 1) // Used by: GptOss diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index e4b881e04..f12d7c33b 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -640,6 +640,12 @@ mod ffi { /// output = silu(gate) * x fn compiled_swiglu_activation(gate: &MlxArray, x: &MlxArray) -> UniquePtr; + /// Compiled three-way add `(a + b) + c` as one fused kernel. + /// Byte-identical to two chained `add` calls (same association order); + /// saves one dispatch and one barrier level per call. + /// Used by: Cohere2 + fn compiled_add3(a: &MlxArray, b: &MlxArray, c: &MlxArray) -> UniquePtr; + /// Compiled GptOss SwiGLU activation with kernel fusion /// Matches mlx-lm gpt_oss.swiglu: clipped gate/up + sigmoid(1.702*gate). /// Used by: GptOss diff --git a/src/models/cohere2.rs b/src/models/cohere2.rs index 623ddeef6..32240571b 100644 --- a/src/models/cohere2.rs +++ b/src/models/cohere2.rs @@ -359,9 +359,9 @@ impl Cohere2TransformerBlock { let attn_h = self.self_attn.forward(&h, cache, mask); let ff_h = self.mlp.forward(&h); - // attn_h + ff_h + x - let sum = mlxcel_core::add(&attn_h, &ff_h); - mlxcel_core::add(&sum, x) + // (attn_h + ff_h) + x as one fused kernel: byte-identical to two adds, + // one barrier level fewer per layer on the decode critical path. + mlxcel_core::compiled_add3(&attn_h, &ff_h, x) } pub fn from_weights( From d392ba8b5917afd0973a8cb4a764c7c71702813c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 21 Sep 2026 21:28:41 +0900 Subject: [PATCH 2/2] perf(cohere2): fuse the residual add with the next LayerNorm A Cohere2 block ends in `(attn + mlp) + x` and the next block starts with a LayerNorm of that sum, so every layer boundary on the decode critical path is two dependent dispatches, a compiled add and MLX's `layer_norm_single_row`. `fused_add3_layer_norm` is one Metal kernel that writes the residual and its normalization. It copies the pinned MLX kernel's threadgroup size, read pattern, two-stage reductions, `precise::rsqrt` and affine step (bias read from memory, as `fast::layer_norm` passes it), so both outputs are byte-identical to the unfused pair rather than close to it. `layers::residual_add3_layer_norm` falls back to the pair off Metal, above a 6656-wide row, or on mixed dtypes, and `MLXCEL_FUSED_ADD_NORM=0` forces it. The Cohere2 layer loop now fuses block i's add with block i+1's input norm, and the last block's add with the final norm. A unit test asserts exact equality across f16 and bf16, with and without bias, several rows and a width off the 8-read boundary; perturbing the kernel's rsqrt by 0.1% makes it fail. Greedy generations for three prompts at 200 tokens match main. command-r7b 4-bit on M1 Ultra, eight ABBA pairs in one binary: decode 112.45 to 113.72 tok/s (+1.1%), prefill unchanged. --- docs/environment-variables.md | 1 + src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 15 ++ src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp | 182 ++++++++++++++++++++ src/lib/mlxcel-core/src/layers.rs | 128 ++++++++++++++ src/lib/mlxcel-core/src/lib.rs | 17 ++ src/models/cohere2.rs | 81 ++++++--- 6 files changed, 398 insertions(+), 26 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 2a344bfaf..65708104c 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -468,6 +468,7 @@ recommended as normal deployment settings. | `MLXCEL_SDPA_PLAN_DEBUG` | `1` enables | off (`0`) | **CUDA only, diagnostic.** Writes one line per cuDNN SDPA call to stderr with the key fields that decide plan reuse: the q shape, `k_len`, the cache buffer extent and row stride, the mask column count and strides, the causal and sinks flags, whether the call took the decode or the bucketed canonicalization, and whether this call built a plan or reused one, plus the resident plan count. This is how the per-shape-class key-field table and the plan-build counts in `docs/benchmark_results/sdpa-plan-cache-bucket-gb10-2026-09-12.md` were produced (issue #1820). It prints per attention call per layer, so it is a diagnostic aid and not something to leave on. | | `MLXCEL_SDPA_FALLBACK_MAX_QUERIES` | non-negative integer | `32` | **CUDA only.** A masked SDPA call with 2 to N query rows over a longer key sequence (the speculative verify shape: a block appended to a KV cache) bypasses cuDNN and takes MLX's own ops fallback (issue #1799). cuDNN caches its execution plan by the exact shapes, and a verify round's key length changes every round, so with cuDNN every such layer class rebuilt a plan on the host each round: about 22 ms per build on GB10, 67 to 76 ms per round on the Laguna DFlash pairing, the whole fixed floor of that round, and the plan cache's lifetime miss counter then aborted the process. The one-row decode step takes the vector kernel and is unaffected; prefill keeps cuDNN (its key length equals its query length, or it has more rows than the bound), except the trailing short chunk of a chunked prefill and a short incremental prefill over a reused prefix-cache prefix, which take the fallback. The fallback's cost grows with the key length (a `[B, heads, q_len, k_len]` score matrix, about 10 ms more per round than cuDNN at block 2 with 350 keys) while the plan build it replaces is constant, so on very long contexts (tens of thousands of keys at block 16, unmeasured) cuDNN could be the cheaper side again. Issue #1820 left this gate's shipped behaviour unchanged: bucketing is off by default, so this gate still claims every array-masked verify block as it did before. It is narrowed only when `MLXCEL_SDPA_PLAN_BUCKET_MAX_QUERIES` is set non-zero, in which case it stops claiming calls whose plan-cache key can be bucketed and covers only a causal-mode block with no array mask, where there is no mask to widen. `0` restores upstream dispatch without a rebuild, which is the kill switch used for the A/B in `docs/benchmark_results/laguna-dflash-verify-cost-gb10-2026-09-11.md`. | | `MLXCEL_PIPELINE_GRANULARITY` | `off`, `layer`, `block:N` | `off` | Inserts layer-boundary async-eval hints for pipeline experiments. | +| `MLXCEL_FUSED_ADD_NORM` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fuses a parallel-residual block's `(attn + mlp) + x` with the LayerNorm that consumes it (the next block's input norm, or the final norm) into one Metal kernel, saving one dispatch and one barrier level per layer during decode. The kernel copies MLX's single-row `layer_norm` reduction, so its outputs are byte-identical to the unfused add and norm (pinned by `residual_add3_layer_norm_matches_the_unfused_pair`); it falls back to the unfused pair off Metal, above a 6656-wide row, or on mixed dtypes. Used by Cohere2: command-r7b 4-bit decode on M1 Ultra +1.1% (eight ABBA pairs, 112.45 to 113.72 tok/s). | | `MLXCEL_FUSED_MOE` | `0`/`false`/`off`/`no` disable; any other value or unset enables | on | Fused single-token decode-MoE kernel (#268), on by default since #282 (Metal) and #319 (CUDA, via `mx.fast.cuda_kernel`); validated on M1 Ultra, M5, and GB10. Set to `0` to force the proven `gather_qmm`/`SwitchGLU` path. Active for afmoe, bailing_moe, cohere2_moe, dbrx, dots.llm1, gemma4, klear, laguna, lfm2, mellum, minimax, mixtral, olmoe, phimoe, qwen2_moe, qwen3_moe, qwen3_next (and Qwen3.5), qwen3_vl_moe, and the qwen3_omni_moe thinker and talker decode. Byte-identical greedy output is checkpoint- and prompt-dependent and was never a general property (#1045): it held on `qwen3-30b-a3b` for the #1045 prompt but not for every prompt (on GB10 one diverges at generated token 39, #1884), and not on Klear. This is not a defect, since the kernel measures roughly 6x closer to an all-f32 ground truth than `gather_qmm` on both, but `gather_qmm` is what mlx-lm mirrors, so set this to `0` when reference-diffing a new MoE port. On the experimental ROCm build the fused kernel has no ROCm port and aborts on affine MoE models, so set this to `0` there until lablup/mlxcel#1803. | | `MLXCEL_FUSED_MOE_SGY` | `1`-`32` | `8` | Simdgroups (Metal) / warps-per-block (CUDA) per threadgroup for the fused decode-MoE kernel; tune per hardware. | | `MLXCEL_FUSED_MOE_MAX_DFF` | positive int | `4096` (Metal) / `8192` (CUDA) | Expert-intermediate (Dff) upper bound for the fused path; above it the caller falls back to `gather_qmm`. The fused path wins only while `gather_qmm` underutilizes the GPU (small experts), so the break-even is backend-dependent and the default is chosen from the live backend: `4096` on Metal (M1 Ultra tuning) and `8192` on CUDA (GB10 re-measured under MLX pin e9463bb, #626; fused wins through Dff 6400 and is break-even at 8192). An explicit value overrides the default on both backends: lower it to force `gather_qmm` sooner, raise it (e.g. `20000`) to force the fused kernel on larger experts such as mixtral (Dff 14336, where it is a slight net loss). Read by every family on the shared `SwitchGLU` fused path, qwen3_moe and qwen3_vl_moe (and the qwen3_omni_moe thinker through it) included since #1884; qwen3_next (and Qwen3.5 and the qwen3_omni_moe talker through it) and gemma4 dispatch their own kernel and do not read it yet. | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 74647716e..813721b4d 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -538,6 +538,21 @@ std::unique_ptr compiled_swiglu_activation( const MlxArray& x ); +// Residual add fused with the next LayerNorm, one Metal launch: +// x_out = (a + b) + x, h_out = layer_norm(x_out, weight, bias). Byte-identical to +// compiled_add3 followed by fast::layer_norm. Metal only, D <= 6656; the Rust +// wrapper (layers::residual_add3_layer_norm) checks that. Used by: Cohere2 +void fused_add3_layer_norm( + const MlxArray& a, + const MlxArray& b, + const MlxArray& x, + const MlxArray& weight, + const MlxArray* bias, + float eps, + std::unique_ptr& x_out, + std::unique_ptr& h_out +); + // Three-way add (a + b) + c compiled into one fused kernel (shapeless=true). // Byte-identical to two chained adds. Used by: Cohere2 std::unique_ptr compiled_add3( diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index 52683d986..6714b764d 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -2531,4 +2531,186 @@ void fused_mamba2_forward( ssm_state_out = std::move(new_ssm_state); } +// ── Residual add fused with the next LayerNorm (Cohere2 parallel block) ───── +// One launch for `x_out = (a + b) + x` and `h_out = layer_norm(x_out, w, bias)`, +// which the unfused graph runs as a compiled add3 kernel followed by MLX's +// `layer_norm_single_row`, two dependent dispatches and two barrier levels per +// layer boundary during decode. +// +// Byte-identical to that pair by construction, not by tolerance: +// - the residual is formed in T with the same association order as +// `compiled_add3`, so `x_out` is the same array element for element; +// - the normalization copies `layer_norm_single_row` from +// mlx/backend/metal/kernels/layer_norm.metal at the pinned MLX commit: the +// same threadgroup size (32 * ceil(ceil(D / 8) / 32)), 8 reads per thread, +// the same two-stage simd/threadgroup reductions for the mean and the +// centred sum of squares, `metal::precise::rsqrt`, and the affine step in T +// with the bias read from memory (a zero scalar with stride 0 when the norm +// has no bias, exactly what `fast::layer_norm` passes), so the compiler +// sees the same expression. +// Covers the single-row kernel only (D <= 6656, MLX's `looped_limit`); the +// caller falls back to the unfused pair above that, off Metal, or on mixed +// dtypes. `residual_add3_layer_norm_matches_the_unfused_pair` pins the identity. +// Used by: Cohere2 +namespace { + static const char* ADD3_LN_METAL_HEADER = R"( + inline void mlxcel_ln_init(threadgroup float* xs, uint lane, uint sg) { + if (sg == 0) { + xs[lane] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + inline void mlxcel_ln_sum(thread float* x, threadgroup float* xs, uint lane, uint sg) { + x[0] = simd_sum(x[0]); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (lane == 0) { + xs[sg] = x[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + x[0] = xs[lane]; + x[0] = simd_sum(x[0]); + } + )"; + + static const char* ADD3_LN_METAL_SOURCE = R"( + constexpr int SIMD_SIZE = 32; + constexpr int N_READS = 8; + uint gid = threadgroup_position_in_grid.x; + uint lid = thread_position_in_threadgroup.x; + uint lane = thread_index_in_simdgroup; + uint sg = simdgroup_index_in_threadgroup; + + float thread_x[N_READS] = {0}; + threadgroup float local_buffer[SIMD_SIZE]; + mlxcel_ln_init(local_buffer, lane, sg); + + size_t off = size_t(gid) * D + lid * N_READS; + const bool safe = lid * N_READS + N_READS <= D; + const int n = int(D) - int(lid * N_READS); + + if (safe) { + for (int i = 0; i < N_READS; i++) { + T s = ra[off + i] + rb[off + i]; + T xn = s + rx[off + i]; + x_out[off + i] = xn; + thread_x[i] = xn; + } + } else { + for (int i = 0; i < n; i++) { + T s = ra[off + i] + rb[off + i]; + T xn = s + rx[off + i]; + x_out[off + i] = xn; + thread_x[i] = xn; + } + } + + float mean = 0; + for (int i = 0; i < N_READS; i++) { + mean += thread_x[i]; + } + mlxcel_ln_sum(&mean, local_buffer, lane, sg); + mean /= D; + + // Upstream starts this loop at `n`, which is negative for threads past + // the end of a narrow row (D < 8 * threadgroup size) and indexes before + // `thread_x`. Its in-range effect is "fill all eight with the mean", + // which clamping the start to 0 reproduces without the out-of-bounds + // write, so the result is unchanged. + float normalizer = 0; + if (!safe) { + for (int i = (n > 0 ? n : 0); i < N_READS; i++) { + thread_x[i] = mean; + } + } + for (int i = 0; i < N_READS; i++) { + thread_x[i] -= mean; + normalizer += thread_x[i] * thread_x[i]; + } + mlxcel_ln_sum(&normalizer, local_buffer, lane, sg); + normalizer = metal::precise::rsqrt(normalizer / D + eps[0]); + + // `auto`: metal_kernel may place a small input in the constant + // address space, so the pointer type follows the input. + auto wp = w + W_STRIDE * lid * N_READS; + auto bp = bias + B_STRIDE * lid * N_READS; + if (safe) { + for (int i = 0; i < N_READS; i++) { + thread_x[i] *= normalizer; + h_out[off + i] = wp[W_STRIDE * i] * static_cast(thread_x[i]) + bp[B_STRIDE * i]; + } + } else { + for (int i = 0; i < n; i++) { + thread_x[i] *= normalizer; + h_out[off + i] = wp[W_STRIDE * i] * static_cast(thread_x[i]) + bp[B_STRIDE * i]; + } + } + )"; + + struct Add3LayerNormKernelHolder { + std::optional kernel; + bool initialized = false; + mlx::core::fast::CustomKernelFunction& get() { + if (!initialized) { + kernel = mlx::core::fast::metal_kernel( + "mlxcel_add3_layer_norm", + {"ra", "rb", "rx", "w", "bias", "eps"}, + {"x_out", "h_out"}, + ADD3_LN_METAL_SOURCE, + ADD3_LN_METAL_HEADER); + initialized = true; + } + return *kernel; + } + }; + static Add3LayerNormKernelHolder& get_add3_layer_norm_kernel() { + static Add3LayerNormKernelHolder holder; + return holder; + } +} + +void fused_add3_layer_norm( + const MlxArray& a, + const MlxArray& b, + const MlxArray& x, + const MlxArray& weight, + const MlxArray* bias, + float eps, + std::unique_ptr& x_out, + std::unique_ptr& h_out +) { + using namespace mlx::core; + auto T = x.inner.dtype(); + const auto& shape = x.inner.shape(); + const int D = shape.back(); + const int64_t rows = x.inner.size() / D; + const int simd = 32; + const int n_reads = 8; + const int tg = simd * (((D + n_reads - 1) / n_reads + simd - 1) / simd); + + // The zero `fast::layer_norm` passes when there is no bias, read through a + // stride-0 pointer as upstream does. One element rather than 0-d, because + // metal_kernel hands a 0-d input to the kernel as a scalar, not a pointer. + array bias_arr = bias ? astype(bias->inner, T) : zeros({1}, T); + const int b_stride = bias && bias->inner.ndim() == 1 ? 1 : 0; + + auto& kernel = get_add3_layer_norm_kernel().get(); + std::vector> ta = { + {"T", T}, + {"D", D}, + {"W_STRIDE", 1}, + {"B_STRIDE", b_stride}, + }; + std::vector inputs = { + a.inner, b.inner, x.inner, astype(weight.inner, T), bias_arr, + full({1}, eps, float32), + }; + auto results = kernel( + inputs, {shape, shape}, {T, T}, + std::make_tuple(static_cast(rows * tg), 1, 1), + std::make_tuple(tg, 1, 1), + ta, std::nullopt, false, {}); + x_out = std::make_unique(std::move(results[0])); + h_out = std::make_unique(std::move(results[1])); +} + } // namespace mlx_cxx diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index fc8ab4b1a..4573b9321 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/src/lib/mlxcel-core/src/layers.rs @@ -977,6 +977,80 @@ impl LayerNorm { } } +/// Largest normalized dimension MLX's single-row `layer_norm` kernel handles +/// (`looped_limit` in `mlx/backend/metal/normalization.cpp`). The fused kernel +/// copies that kernel, so it covers the same range. +const FUSED_ADD3_LAYER_NORM_MAX_DIM: i32 = 6656; + +fn fused_add3_layer_norm_enabled() -> bool { + static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); + *ENABLED.get_or_init(|| { + !matches!( + std::env::var("MLXCEL_FUSED_ADD_NORM").as_deref(), + Ok("0" | "off" | "false" | "no") + ) + }) +} + +/// Residual add fused with the LayerNorm that consumes it: returns +/// `(x_new, norm(x_new))` with `x_new = (a + b) + x`. +/// +/// A parallel-residual block ends in `(attn + mlp) + x` and the next block +/// starts with a LayerNorm of that sum. Unfused that is a compiled add kernel +/// and a norm kernel, two dependent dispatches (two barrier levels) per layer +/// boundary on the decode critical path. On Metal this runs one kernel whose +/// outputs are byte-identical to that pair; elsewhere, or when the shapes and +/// dtypes fall outside what the kernel covers, it runs the pair itself. +/// `MLXCEL_FUSED_ADD_NORM=0` forces the unfused pair. +/// +/// Used by: Cohere2 +pub fn residual_add3_layer_norm( + a: &MlxArray, + b: &MlxArray, + x: &MlxArray, + norm: &LayerNorm, +) -> (UniquePtr, UniquePtr) { + let shape = ffi::array_shape(x); + let dtype = ffi::array_dtype(x); + let dim = shape.last().copied().unwrap_or(0); + let weight = norm.weight.as_ref().unwrap(); + let fusable = fused_add3_layer_norm_enabled() + && ffi::metal_is_available() + && dim > 0 + && dim <= FUSED_ADD3_LAYER_NORM_MAX_DIM + && matches!( + dtype, + crate::dtype::FLOAT16 | crate::dtype::BFLOAT16 | crate::dtype::FLOAT32 + ) + && ffi::array_shape(a) == shape + && ffi::array_shape(b) == shape + && ffi::array_dtype(a) == dtype + && ffi::array_dtype(b) == dtype + && ffi::array_dtype(weight) == dtype + && ffi::array_shape(weight) == [dim] + && norm.bias.as_ref().is_none_or(|bias| { + let bias = bias.as_ref().unwrap(); + ffi::array_dtype(bias) == dtype && ffi::array_shape(bias) == [dim] + }); + if !fusable { + let x_new = ffi::compiled_add3(a, b, x); + let h = norm.forward(&x_new); + return (x_new, h); + } + let bias_ptr = norm + .bias + .as_ref() + .map(|bias| bias.as_ref().unwrap() as *const MlxArray) + .unwrap_or(std::ptr::null()); + let mut x_new = UniquePtr::null(); + let mut h = UniquePtr::null(); + // SAFETY: `bias_ptr` is null or points at `norm.bias`, which outlives the call. + unsafe { + ffi::fused_add3_layer_norm(a, b, x, weight, bias_ptr, norm.eps, &mut x_new, &mut h); + } + (x_new, h) +} + /// Named LoRA weights for runtime on-the-fly application. /// Used by: Phi4MM VLM (language / vision / speech request modes) pub struct LoRAWeights { @@ -8997,3 +9071,57 @@ mod metal4_attention_switch_tests { ); } } + +#[cfg(all(test, feature = "metal"))] +mod residual_add3_layer_norm_tests { + use super::*; + use crate::dtype; + + fn normal(shape: &[i32], dt: i32, seed: u64, scale: f32) -> UniquePtr { + let key = ffi::random_key(seed); + let x = unsafe { ffi::random_normal(shape, dtype::FLOAT32, &*key) }; + let scaled = ffi::multiply(&x, &ffi::full_f32(&[1], scale, dtype::FLOAT32)); + ffi::astype(&scaled, dt) + } + + /// The fused kernel's contract is byte identity with `compiled_add3` + /// followed by `LayerNorm::forward`, the pair it replaces, not closeness. + /// Covers f16 and bf16, with and without a bias, several rows, and a width + /// that is not a multiple of the 8 reads per thread (the kernel's tail + /// branch). A tolerance check here would let a reordered reduction pass. + #[test] + fn residual_add3_layer_norm_matches_the_unfused_pair() { + for (dt, dim, rows, with_bias) in [ + (dtype::FLOAT16, 4096, 1, false), + (dtype::FLOAT16, 4096, 5, false), + (dtype::FLOAT16, 4100, 3, true), + (dtype::BFLOAT16, 4096, 2, true), + (dtype::FLOAT16, 96, 4, false), + ] { + let shape = [1, rows, dim]; + let a = normal(&shape, dt, 1, 1.0); + let b = normal(&shape, dt, 2, 1.0); + let x = normal(&shape, dt, 3, 4.0); + let norm = LayerNorm::new( + normal(&[dim], dt, 4, 0.5), + with_bias.then(|| normal(&[dim], dt, 5, 0.1)), + 1e-5, + ); + + let (x_fused, h_fused) = residual_add3_layer_norm(&a, &b, &x, &norm); + let x_ref = ffi::compiled_add3(&a, &b, &x); + let h_ref = norm.forward(&x_ref); + + let same_x = ffi::array_equal(&x_fused, &x_ref, false); + let same_h = ffi::array_equal(&h_fused, &h_ref, false); + assert!( + ffi::item_bool(&same_x), + "residual differs: dtype {dt} dim {dim} rows {rows} bias {with_bias}" + ); + assert!( + ffi::item_bool(&same_h), + "normalized output differs: dtype {dt} dim {dim} rows {rows} bias {with_bias}" + ); + } + } +} diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index f12d7c33b..913538902 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -640,6 +640,23 @@ mod ffi { /// output = silu(gate) * x fn compiled_swiglu_activation(gate: &MlxArray, x: &MlxArray) -> UniquePtr; + /// Residual add fused with the next LayerNorm in one Metal launch: + /// `x_out = (a + b) + x`, `h_out = layer_norm(x_out, weight, bias)`. + /// Byte-identical to `compiled_add3` + `fast_layer_norm`. Metal only, + /// last dimension <= 6656; call it through + /// [`crate::layers::residual_add3_layer_norm`], which checks both. + /// Used by: Cohere2 + unsafe fn fused_add3_layer_norm( + a: &MlxArray, + b: &MlxArray, + x: &MlxArray, + weight: &MlxArray, + bias: *const MlxArray, + eps: f32, + x_out: &mut UniquePtr, + h_out: &mut UniquePtr, + ); + /// Compiled three-way add `(a + b) + c` as one fused kernel. /// Byte-identical to two chained `add` calls (same association order); /// saves one dispatch and one barrier level per call. diff --git a/src/models/cohere2.rs b/src/models/cohere2.rs index 32240571b..ed267b3d1 100644 --- a/src/models/cohere2.rs +++ b/src/models/cohere2.rs @@ -356,14 +356,25 @@ impl Cohere2TransformerBlock { // h = norm(x) // out = attn(h) + mlp(h) + x let h = self.input_layernorm.forward(x); - let attn_h = self.self_attn.forward(&h, cache, mask); - let ff_h = self.mlp.forward(&h); + let (attn_h, ff_h) = self.attn_and_mlp(&h, cache, mask); // (attn_h + ff_h) + x as one fused kernel: byte-identical to two adds, // one barrier level fewer per layer on the decode critical path. mlxcel_core::compiled_add3(&attn_h, &ff_h, x) } + /// The two parallel branches for an input that is already normalized by + /// this block's `input_layernorm`. The model's layer loop uses this so the + /// residual add of block `i` can be fused with the norm of block `i + 1`. + pub fn attn_and_mlp( + &self, + h: &MlxArray, + cache: &mut KVCache, + mask: Option<&MlxArray>, + ) -> (UniquePtr, UniquePtr) { + (self.self_attn.forward(h, cache, mask), self.mlp.forward(h)) + } + pub fn from_weights( weights: &WeightMap, args: &Cohere2Config, @@ -409,6 +420,43 @@ pub struct Cohere2Model { } impl Cohere2Model { + /// All transformer layers followed by the final norm, returning the + /// normalized hidden state for the LM head. + /// + /// Each block's residual add is fused with the norm that consumes it (the + /// next block's `input_layernorm`, or the model's final `norm` after the + /// last block), which on Metal saves one dispatch and one barrier level per + /// layer and is byte-identical to the unfused add and norm. + fn decoder_stack( + &self, + mut x: UniquePtr, + caches: &mut [KVCache], + full_mask: Option<&MlxArray>, + sliding_mask: Option<&MlxArray>, + ) -> UniquePtr { + let Some(first) = self.layers.first() else { + return self.norm.forward(&x); + }; + let mut h = first.input_layernorm.forward(&x); + for (i, layer) in self.layers.iter().enumerate() { + let mask = if self.config.is_sliding_window_layer(i) { + sliding_mask + } else { + full_mask + }; + let (attn_h, ff_h) = layer.attn_and_mlp(&h, &mut caches[i], mask); + let next_norm = self + .layers + .get(i + 1) + .map_or(&self.norm, |next| &next.input_layernorm); + let (x_new, h_new) = + mlxcel_core::layers::residual_add3_layer_norm(&attn_h, &ff_h, &x, next_norm); + x = x_new; + h = h_new; + } + h + } + /// Forward pass through the entire model pub fn forward_impl( &self, @@ -417,7 +465,7 @@ impl Cohere2Model { _mask: Option<&MlxArray>, ) -> UniquePtr { // Embed tokens - let mut h = self.embed_tokens.forward(input_ids); + let h = self.embed_tokens.forward(input_ids); let shape = mlxcel_core::array_shape(&h); let l = shape[1] as usize; @@ -452,18 +500,8 @@ impl Cohere2Model { (None, None) }; - // Pass through transformer layers - for (i, layer) in self.layers.iter().enumerate() { - let mask = if self.config.is_sliding_window_layer(i) { - sliding_mask.as_ref().map(|m| m.as_ref().unwrap()) - } else { - full_mask.as_ref().map(|m| m.as_ref().unwrap()) - }; - h = layer.forward(&h, &mut caches[i], mask); - } - - // Final norm - let h = self.norm.forward(&h); + // Transformer layers and the final norm. + let h = self.decoder_stack(h, caches, full_mask.as_deref(), sliding_mask.as_deref()); // Output projection let logits = self.lm_head.forward(&h); @@ -487,7 +525,7 @@ impl Cohere2Model { caches: &mut [KVCache], _mask: Option<&MlxArray>, ) -> UniquePtr { - let mut h = if let Some(embeds) = input_embeddings { + let h = if let Some(embeds) = input_embeddings { mlxcel_core::copy(embeds) } else { self.embed_tokens.forward(input_ids) @@ -525,16 +563,7 @@ impl Cohere2Model { (None, None) }; - for (i, layer) in self.layers.iter().enumerate() { - let mask = if self.config.is_sliding_window_layer(i) { - sliding_mask.as_ref().map(|m| m.as_ref().unwrap()) - } else { - full_mask.as_ref().map(|m| m.as_ref().unwrap()) - }; - h = layer.forward(&h, &mut caches[i], mask); - } - - let h = self.norm.forward(&h); + let h = self.decoder_stack(h, caches, full_mask.as_deref(), sliding_mask.as_deref()); let logits = self.lm_head.forward(&h); let scale_arr = mlxcel_core::full_f32(&[1], self.logit_scale, mlxcel_core::array_dtype(&logits));