perf(cohere2): fuse the parallel-block residual adds with the next norm

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. |

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp

@@ -1650,6 +1650,29 @@ std::unique_ptr<MlxArray> compiled_swiglu_activation(
     return std::make_unique<MlxArray>(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<std::vector<array>(const std::vector<array>&)> get_compiled_add3() {
+        auto fn = [](const std::vector<array>& inputs) -> std::vector<array> {
+            return {mlx::core::add(mlx::core::add(inputs[0], inputs[1]), inputs[2])};
+        };
+        return compile_shapeless_audited("compiled_add3", fn);
+    }
+}
+
+std::unique_ptr<MlxArray> 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<MlxArray>(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)

src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h

@@ -538,6 +538,29 @@ std::unique_ptr<MlxArray> 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<MlxArray>& x_out,
+    std::unique_ptr<MlxArray>& 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<MlxArray> 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

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<T>(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<T>(thread_x[i]) + bp[B_STRIDE * i];
+            }
+        }
+    )";
+
+    struct Add3LayerNormKernelHolder {
+        std::optional<mlx::core::fast::CustomKernelFunction> 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<MlxArray>& x_out,
+    std::unique_ptr<MlxArray>& 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<std::pair<std::string, mlx::core::fast::TemplateArg>> ta = {
+        {"T", T},
+        {"D", D},
+        {"W_STRIDE", 1},
+        {"B_STRIDE", b_stride},
+    };
+    std::vector<array> 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<int>(rows * tg), 1, 1),
+        std::make_tuple(tg, 1, 1),
+        ta, std::nullopt, false, {});
+    x_out = std::make_unique<MlxArray>(std::move(results[0]));
+    h_out = std::make_unique<MlxArray>(std::move(results[1]));
+}
+
 }  // namespace mlx_cxx