feat(paged): fused residual-add + RMS norm + weight multiply (patch 0042)

The transformer pre-norm residual chain `h = x + sub_out; n = rms_norm(h) * w`
runs as separate CUDA launches in the paged prefill graph: a k_bin_bcast ADD
(the residual) feeding the existing fused rms_norm+mul. ggml-cuda already fuses
rms_norm+mul (and rms_norm+mul+ADD, where the ADD is a *post*-norm bias) but NOT
the *pre*-norm residual add that feeds the norm. This is the classic add-RMSNorm
fusion (as in vLLM / TensorRT-LLM) that ggml-cuda lacks; it is part of the
unfused-tail prefill gap vs vLLM's torch.compile fusions.

Add it as a CUDA-family graph fusion (paged series owns it; stock stays pure):
- ggml_cuda_can_fuse recognizes { ADD, RMS_NORM, MUL } via ggml_can_fuse_subgraph
  with BOTH the ADD (node_idx) and the MUL (node_idx+2) marked as outputs - the
  residual ADD has a second consumer (the later skip-connection add), so it
  cannot pass the single-use ggml_can_fuse() gate the other rms_norm fusions use.
- New kernel rms_norm_pre_add_mul_f32 computes h = a + b, publishes h to the
  residual buffer (downstream skip add reads it), then sum(h^2) -> scale ->
  dst = scale * h * w in ONE launch, emitting BOTH outputs the graph needs.
- Gated by LLAMA_FUSE_ADD_RMSNORM (default ON) for a clean single-build A/B.

BIT-EXACT (per-path canonical greedy md5, n=48 --temp 0 --seed 1, paged):
  dense q36-27b-nvfp4 : 5951a5b4d624ce891e22ab5fca9bc439 (ON == OFF == canonical)
  MoE   q36-35b-a3b   : 8cb0ce23777bf55f92f63d0292c756b0 (ON == OFF == canonical)
The fused kernel reproduces the exact FP order of the unfused chain: h = a + b
(IEEE add is order-free), the sum(h^2) reduction uses the same block_reduce<SUM>
with the same 256/1024 block-size thresholds, and the same rsqrtf(mean+eps)
scale, so the byte stream is unchanged. test-backend-ops RMS_NORM/ADD/MUL pass
(CUDA0 vs CPU).

PROFILE (dense prefill, nsys --cuda-graph-trace=node, npp512 ntg4 npl8):
  rms_norm_f32<1024>    903 launches / 96.6M ns -> 7   / 0.7M ns
  k_bin_bcast<op_add>  1232 launches / 138.6M ns -> 336 / 1.0M ns
  rms_norm_pre_add_mul (new)               896 launches / 187.2M ns
  -> 896 residual-add + 896 rms_norm launches folded into 896 fused launches;
     the norm+residual slice 233.6M -> 187.2M ns (~20% of that slice, ~1% of
     total prefill GPU time).
S_PP dense (npp512 ntg4 npl32, 3x): 985.5 -> 990.6 t/s (+0.5%, every ON run
beats every OFF run). Modest because the residual tail is a small slice of
prefill; the dominant unfused cost is k_bin_bcast<op_mul> (11%, the GDN
chunked-prefill gating muls) - a separate lever.

Assisted-by: Claude:opus-4.8 [Claude Code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -3764,6 +3764,48 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
         }
     }
 
+    // Fused residual-add + RMS norm + weight multiply. The transformer residual
+    // ADD feeds the next sublayer's RMS norm but is ALSO consumed by the later
+    // residual add (skip connection), so the ADD node is a graph output too; it
+    // cannot go through the single-use ggml_can_fuse() gate below. Recognize it
+    // here with ggml_can_fuse_subgraph, marking both the ADD (node_idx) and the
+    // final MUL (node_idx + 2) as outputs.
+    std::initializer_list<enum ggml_op> add_rms_norm_mul_ops = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL };
+    if (is_equal(add_rms_norm_mul_ops, ops) &&
+        ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx, node_idx + 2 })) {
+        const ggml_tensor * add      = cgraph->nodes[node_idx];
+        const ggml_tensor * rms_norm = cgraph->nodes[node_idx + 1];
+        const ggml_tensor * mul      = cgraph->nodes[node_idx + 2];
+
+        // RMS norm must consume the residual-add output.
+        if (rms_norm->src[0] != add) {
+            return false;
+        }
+        // All operands F32 (rms norm / fused mul kernel only support F32).
+        if (add->src[0]->type != GGML_TYPE_F32 || add->src[1]->type != GGML_TYPE_F32 ||
+            add->type != GGML_TYPE_F32 || rms_norm->type != GGML_TYPE_F32 ||
+            mul->src[0]->type != GGML_TYPE_F32 || mul->src[1]->type != GGML_TYPE_F32 ||
+            mul->type != GGML_TYPE_F32) {
+            return false;
+        }
+        // The fused kernel computes h = a + b elementwise: same shape, no broadcast.
+        if (!ggml_are_same_shape(add->src[0], add->src[1])) {
+            return false;
+        }
+        // rms_norm kernel assumes contiguous rows for the residual operands and weight.
+        if (!ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous(add->src[1])) {
+            return false;
+        }
+        if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) {
+            return false;
+        }
+        // If rms_norm is the B operand of the mul, broadcast of the A operand is unsupported.
+        if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
+            return false;
+        }
+        return true;
+    }
+
     if (!ggml_can_fuse(cgraph, node_idx, ops)) {
         return false;
     }
@@ -4286,6 +4328,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
         return fused_node_count - 1;
     }
 
+    // Fused residual-add + RMS norm + weight multiply (bit-exact). Default ON;
+    // set LLAMA_FUSE_ADD_RMSNORM=0 for a clean A/B against the unfused path.
+    static const bool fuse_add_rmsnorm = [] {
+        const char * e = getenv("LLAMA_FUSE_ADD_RMSNORM");
+        return e == nullptr || atoi(e) != 0;
+    }();
+    if (fuse_add_rmsnorm &&
+        ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) {
+        ggml_cuda_op_rms_norm_pre_add_mul(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
+        return 2;
+    }
+
     if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) {
         ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
         return 2;

ggml/src/ggml-cuda/norm.cu

@@ -154,6 +154,87 @@ static __global__ void rms_norm_f32(const float * x,
     }
 }
 
+// Fused residual-add + RMS norm + (optional) weight multiply.
+//   h   = a + b                       (the residual stream, written to h_out)
+//   dst = rsqrt(mean(h^2)+eps) * h * mul
+// `a` and `b` are required to be the same shape and contiguous (the transformer
+// residual add), so they share `x`'s strides; `h_out`, `dst` are also contiguous
+// with that shape. `mul` (the RMS weight) broadcasts via the packed-modulo path.
+//
+// Bit-exactness: this reproduces the exact FP order of the unfused chain
+//   k_bin_bcast(add): h[col] = a[col] + b[col]   (f32, elementwise, order-free)
+//   rms_norm:         sumsq over h[col] in column order via block_reduce
+//   mul:              dst[col] = scale * h[col] * mul[col]
+// h is summed from the same f32 values in the same order, so the reduction and
+// the final scale are byte-identical to running the three kernels separately.
+template <int block_size, bool do_multiply = false>
+static __global__ void rms_norm_pre_add_mul_f32(const float * a,
+                                                const float * b,
+                                                float *       h_out,
+                                                float *       dst,
+                                                const int     ncols,
+                                                const int64_t stride_row,
+                                                const int64_t stride_channel,
+                                                const int64_t stride_sample,
+                                                const float   eps,
+                                                const float * mul                  = nullptr,
+                                                const int64_t mul_stride_row       = 0,
+                                                const int64_t mul_stride_channel   = 0,
+                                                const int64_t mul_stride_sample    = 0,
+                                                const uint3   mul_ncols_packed     = make_uint3(0, 0, 0),
+                                                const uint3   mul_nrows_packed     = make_uint3(0, 0, 0),
+                                                const uint3   mul_nchannels_packed = make_uint3(0, 0, 0),
+                                                const uint3   mul_nsamples_packed  = make_uint3(0, 0, 0)) {
+    ggml_cuda_pdl_lc();
+    const int nrows     = gridDim.x;
+    const int nchannels = gridDim.y;
+
+    const int row       = blockIdx.x;
+    const int channel   = blockIdx.y;
+    const int sample    = blockIdx.z;
+    const int tid       = threadIdx.x;
+
+    const int64_t row_offset = sample*stride_sample + channel*stride_channel + row*stride_row;
+    a     += row_offset;
+    b     += row_offset;
+    h_out += row_offset;
+    // dst is laid out contiguously by the scheduler for the MUL output
+    dst   += ((sample*nchannels + channel)*nrows + row)*ncols;
+
+    if constexpr (do_multiply) {
+        const uint32_t mul_row     = fastmodulo(row, mul_nrows_packed);
+        const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
+        const uint32_t mul_sample  = fastmodulo(sample, mul_nsamples_packed);
+        mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
+    }
+
+    float tmp = 0.0f; // partial sum for thread in warp
+
+    ggml_cuda_pdl_sync();
+    for (int col = tid; col < ncols; col += block_size) {
+        const float hi = a[col] + b[col];
+        h_out[col] = hi;          // publish the residual stream for the next add
+        tmp += hi * hi;
+    }
+
+    // sum up partial sums
+    extern __shared__ float s_sum[];
+    tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
+
+    const float mean  = tmp / ncols;
+    const float scale = rsqrtf(mean + eps);
+
+    for (int col = tid; col < ncols; col += block_size) {
+        const float hi = h_out[col];
+        if constexpr (do_multiply) {
+            const int mul_col = fastmodulo(col, mul_ncols_packed);
+            dst[col]          = scale * hi * mul[mul_col];
+        } else {
+            dst[col]          = scale * hi;
+        }
+    }
+}
+
 template <int block_size>
 static __global__ void rms_norm_back_f32(
         const float * grad, const float * xf, float * dst, const int ncols, const float eps) {
@@ -407,6 +488,50 @@ static void rms_norm_mul_f32_cuda(const float *  x,
     }
 }
 
+static void rms_norm_pre_add_mul_f32_cuda(const float *  a,
+                                          const float *  b,
+                                          float *        h_out,
+                                          float *        dst,
+                                          const int      ncols,
+                                          const int      nrows,
+                                          const int      nchannels,
+                                          const int      nsamples,
+                                          const int64_t  stride_row,
+                                          const int64_t  stride_channel,
+                                          const int64_t  stride_sample,
+                                          const float *  mul,
+                                          const int64_t  mul_stride_row,
+                                          const int64_t  mul_stride_channel,
+                                          const int64_t  mul_stride_sample,
+                                          const uint32_t mul_ncols,
+                                          const uint32_t mul_nrows,
+                                          const uint32_t mul_nchannels,
+                                          const uint32_t mul_nsamples,
+                                          const float    eps,
+                                          cudaStream_t   stream) {
+    const dim3 blocks_num(nrows, nchannels, nsamples);
+    GGML_ASSERT(mul != nullptr);
+    const uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
+    const uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
+    const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
+    const uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
+    if (ncols < 1024) {
+        const dim3 block_dims(256, 1, 1);
+        const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float) : 0, stream};
+        ggml_cuda_kernel_launch(rms_norm_pre_add_mul_f32<256, true>, launch_params,
+            a, b, h_out, dst, ncols, stride_row, stride_channel, stride_sample, eps,
+            mul, mul_stride_row, mul_stride_channel, mul_stride_sample,
+            mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed);
+    } else {
+        const dim3 block_dims(1024, 1, 1);
+        const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float) : 0, stream};
+        ggml_cuda_kernel_launch(rms_norm_pre_add_mul_f32<1024, true>, launch_params,
+            a, b, h_out, dst, ncols, stride_row, stride_channel, stride_sample, eps,
+            mul, mul_stride_row, mul_stride_channel, mul_stride_sample,
+            mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed);
+    }
+}
+
 static void rms_norm_back_f32_cuda(const float * grad, const float * xf, float * dst, const int ncols, const int nrows, const float eps, cudaStream_t stream) {
     if (ncols < 1024) {
         const dim3 block_dims(WARP_SIZE, 1, 1);
@@ -647,6 +772,77 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
                           eps, stream);
 }
 
+void ggml_cuda_op_rms_norm_pre_add_mul(ggml_backend_cuda_context & ctx,
+                                       ggml_tensor *               add_tensor,
+                                       ggml_tensor *               rms_norm_tensor,
+                                       ggml_tensor *               mul_tensor) {
+    // The RMS norm consumes the residual-add output.
+    GGML_ASSERT(rms_norm_tensor->src[0] == add_tensor);
+
+    const ggml_tensor * a_src = add_tensor->src[0];
+    const ggml_tensor * b_src = add_tensor->src[1];
+
+    float eps = 0.0f;
+    memcpy(&eps, rms_norm_tensor->op_params, sizeof(float));
+    GGML_ASSERT(eps >= 0.0f);
+
+    const float * a_d = (const float *) a_src->data;
+    const float * b_d = (const float *) b_src->data;
+    float *       h_d = (float *)       add_tensor->data;
+
+    const float *       mul_d   = nullptr;
+    const ggml_tensor * mul_src = nullptr;
+    if (mul_tensor->src[0] == rms_norm_tensor) {
+        mul_d   = (const float *) mul_tensor->src[1]->data;
+        mul_src = mul_tensor->src[1];
+    } else if (mul_tensor->src[1] == rms_norm_tensor) {
+        mul_d   = (const float *) mul_tensor->src[0]->data;
+        mul_src = mul_tensor->src[0];
+    } else {
+        GGML_ASSERT(false);
+    }
+
+    float *      dst_d  = (float *) mul_tensor->data;
+    cudaStream_t stream = ctx.stream();
+
+    GGML_ASSERT(a_src->type == GGML_TYPE_F32);
+    GGML_ASSERT(b_src->type == GGML_TYPE_F32);
+    GGML_ASSERT(add_tensor->type == GGML_TYPE_F32);
+    GGML_ASSERT(rms_norm_tensor->type == GGML_TYPE_F32);
+    GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_are_same_shape(a_src, b_src));
+
+    const int64_t ne00 = add_tensor->ne[0];
+    const int64_t ne01 = add_tensor->ne[1];
+    const int64_t ne02 = add_tensor->ne[2];
+    const int64_t ne03 = add_tensor->ne[3];
+
+    // a and b share the (contiguous) residual layout
+    const size_t ts0 = ggml_type_size(a_src->type);
+    GGML_ASSERT(a_src->nb[0] == ts0 && b_src->nb[0] == ts0);
+    const int64_t s01 = a_src->nb[1] / ts0;
+    const int64_t s02 = a_src->nb[2] / ts0;
+    const int64_t s03 = a_src->nb[3] / ts0;
+
+    const size_t ts_mul = ggml_type_size(mul_src->type);
+    GGML_ASSERT(mul_src->nb[0] == ts_mul);
+    const int64_t mul_s01 = mul_src->nb[1] / ts_mul;
+    const int64_t mul_s02 = mul_src->nb[2] / ts_mul;
+    const int64_t mul_s03 = mul_src->nb[3] / ts_mul;
+
+    const int mul_ncols     = mul_src->ne[0];
+    const int mul_nrows     = mul_src->ne[1];
+    const int mul_nchannels = mul_src->ne[2];
+    const int mul_nsamples  = mul_src->ne[3];
+
+    rms_norm_pre_add_mul_f32_cuda(a_d, b_d, h_d, dst_d,
+                                  ne00, ne01, ne02, ne03,
+                                  /*s00*/ s01, s02, s03,
+                                  mul_d, /*mul_s00*/ mul_s01, mul_s02, mul_s03,
+                                  mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
+                                  eps, stream);
+}
+
 void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * grad  = dst->src[0]; // gradients
     const ggml_tensor * src0f = dst->src[1]; // src0 from forward pass

ggml/src/ggml-cuda/norm.cuh

@@ -13,6 +13,11 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
                                      ggml_tensor *               mul_tensor,
                                      ggml_tensor *               add_tensor);
 
+void ggml_cuda_op_rms_norm_pre_add_mul(ggml_backend_cuda_context & ctx,
+                                       ggml_tensor *               add_tensor,
+                                       ggml_tensor *               rms_norm_tensor,
+                                       ggml_tensor *               mul_tensor);
+
 void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 
 void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst);